TOOLS/NestedLoopIterator.h at master · TNWorkers/TOOLS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
#ifndef NESTEDLOOPITERATOR
#define NESTEDLOOPITERATOR

//#include <Eigen/Dense>
#include <vector>
#include <initializer_list>
#include <algorithm>
#include <functional>
#include <assert.h>
#include <numeric>

class NestedLoopIterator
{
public:

	NestedLoopIterator(){};
	NestedLoopIterator (std::size_t dim_input, std::initializer_list<std::size_t> ranges_input);
	NestedLoopIterator (std::size_t dim_input, std::vector<std::size_t> ranges_input);
	NestedLoopIterator (std::size_t dim_input, std::size_t const_range);

	void operator++() {++curr_index; make_reverseTensorIndex();};

	std::size_t operator*() {return index();}

	NestedLoopIterator& operator= (const std::size_t &cmp) {curr_index=cmp; make_tensorIndex(); return *this;}
	bool operator<  (const std::size_t &cmp) {return curr_index< cmp;}
	bool operator<= (const std::size_t &cmp) {return curr_index<=cmp;}
	bool operator!= (const std::size_t &cmp) {return curr_index!=cmp;}
	bool operator== (const std::size_t &cmp) {return curr_index==cmp;}

	std::size_t begin();
	std::size_t end();

	inline std::size_t index() {return curr_index;}
	inline std::size_t index_sum() {return accumulate(tensor_index.begin(), tensor_index.end(), 0);}
	inline std::size_t operator() (std::size_t index) const {return tensor_index[index];}
	inline std::vector<std::size_t> operator() () const {return tensor_index;}

private:

	std::size_t dim;
	std::size_t curr_index;
	std::size_t total_range;

	void make_tensorIndex();
	void make_reverseTensorIndex();

	std::vector<std::size_t> tensor_index;
	std::vector<std::size_t> ranges;
};

NestedLoopIterator::
NestedLoopIterator (std::size_t dim_input, std::initializer_list<std::size_t> ranges_input)
:dim(dim_input)
{
	tensor_index.resize(dim);
	ranges.resize(dim);
	assert(ranges_input.size()==dim);

	std::size_t i=0;
	for (auto range=ranges_input.begin(); range!=ranges_input.end(); ++range)
	{
		assert(*range>=0);
		ranges[i] = *range;
		++i;
	}
	total_range = accumulate(ranges.begin(), ranges.end(), 1, std::multiplies<std::size_t>());

	std::fill(tensor_index.begin(), tensor_index.end(), 0);
	curr_index = 0;
}

NestedLoopIterator::
NestedLoopIterator (std::size_t dim_input, std::vector<std::size_t> ranges_input)
:dim(dim_input)
{
	assert(ranges_input.size() == dim);
	tensor_index.resize(dim);
	ranges.resize(dim);

	ranges = ranges_input;
	total_range = accumulate(ranges.begin(), ranges.end(), 1, std::multiplies<std::size_t>());

	std::fill(tensor_index.begin(), tensor_index.end(), 0);
	curr_index = 0;
}

NestedLoopIterator::
NestedLoopIterator (std::size_t dim_input, std::size_t const_range)
:dim(dim_input)
{
	tensor_index.resize(dim);
	ranges.resize(dim);

	std::fill(ranges.begin(), ranges.end(), const_range);
	total_range = accumulate(ranges.begin(), ranges.end(), 1, std::multiplies<std::size_t>());

	std::fill(tensor_index.begin(), tensor_index.end(), 0);
	curr_index = 0;
}

std::size_t NestedLoopIterator::
begin()
{
	return 0;
}

std::size_t NestedLoopIterator::
end()
{
	return total_range;
}

void NestedLoopIterator::
make_tensorIndex()
{
	std::size_t r = curr_index;
	for (int s=dim-1; s>=0; --s) // must be int!
	{
		tensor_index[s] = r%ranges[s];
		r /= ranges[s];
	}
}

void NestedLoopIterator::
make_reverseTensorIndex()
{
	std::size_t r = curr_index;
	for (int s=0; s<dim; s++)
	{
		tensor_index[s] = r%ranges[s];
		r /= ranges[s];
	}
}

#endif