multi_iterator.h

Go to the documentation of this file.

#ifndef WXM_MULTI_ITERATOR_H
#define WXM_MULTI_ITERATOR_H
 
#include "meta.h"
 
#include <array>
#include <type_traits>
 
namespace wxm
{
template<size_t max_dims> class multi_iterator
{
    std::array<size_t, max_dims> idcs;
 
    std::array<size_t, max_dims> shape_;
 
    size_t ndims_;
 
    size_t idx = 0;
 
    size_t size_;
 
public:
    template<class... Idcs,
             class = typename std::enable_if<
                 meta::all<std::is_integral<Idcs>::value...>::value>::type>
    multi_iterator(Idcs... lengths)
        : idcs{0}, shape_{static_cast<size_t>(lengths)...}, ndims_(sizeof...(Idcs))
    {
        static_assert(
            sizeof...(Idcs) < max_dims,
            "Number of dimensions for the multi-iterator cannot exceed max_dims");
        size_ = 1;
        for (size_t i = 0; i < ndims_; ++i)
        {
            size_ *= shape_[i];
        }
    }
 
    template<class C> multi_iterator(const C& cont) : idcs{0}
    {
        ndims_ = cont.size();
        std::copy(cont.begin(), cont.end(), shape_.begin());
        size_ = 1;
        for (size_t i = 0; i < ndims_; ++i)
        {
            size_ *= shape_[i];
        }
    }
 
    const std::array<size_t, max_dims>& indices() const
    {
        return idcs;
    }
 
    size_t index() const
    {
        return idx;
    }
 
    size_t size() const
    {
        return size_;
    }
 
    size_t ndims() const
    {
        return ndims_;
    }
 
    const std::array<size_t, max_dims>& shape() const
    {
        return shape_;
    }
 
    multi_iterator& operator++()
    {
        ++idx;
        ++idcs[ndims_ - 1];
        for (size_t i = ndims_ - 1; i > 0; --i)
        {
            if (idcs[i] >= shape_[i])
            {
                idcs[i] = 0;
                ++idcs[i - 1];
                continue;
            }
            break;
        }
        return *this;
    }
 
    multi_iterator& operator+=(size_t val)
    {
        idx += val;
        idcs[ndims_ - 1] += val;
        bool no_carry = false;
        for (size_t i = ndims_ - 1; i > 0 && !no_carry; --i)
        {
            no_carry = true;
            while (idcs[i] >= shape_[i])
            {
                idcs[i] -= shape_[i];
                ++idcs[i - 1];
                no_carry = false;
            }
        }
        return *this;
    }
 
    multi_iterator& advance(size_t dim)
    {
        ++idcs[dim];
 
        size_t val = 1;
        for (size_t i = dim + 1; i < ndims_; ++i)
        {
            val *= shape_[i];
        }
        idx += val;
 
        for (size_t i = dim; i > 0; --i)
        {
            if (idcs[i] >= shape_[i])
            {
                idcs[i] = 0;
                ++idcs[i - 1];
                continue;
            }
            break;
        }
        return *this;
    }
 
    multi_iterator& advance(size_t dim, size_t val)
    {
        idcs[dim] += val;
 
        for (size_t i = dim + 1; i < ndims_; ++i)
        {
            val *= shape_[i];
        }
        idx += val;
 
        bool no_carry = false;
        for (size_t i = dim; i > 0 && !no_carry; --i)
        {
            no_carry = true;
            while (idcs[i] >= shape_[i])
            {
                idcs[i] -= shape_[i];
                ++idcs[i - 1];
                no_carry = false;
            }
        }
        return *this;
    }
};
} // namespace wxm
 
#endif