cspan flat-iterator now prints c_COLMAJOR matrices as in Python. Simpler code, although it now doesn't print column-major spans in column-by-column order, but rather in row-by-row, like row-major spans.

2 files changed, 11 insertions, 15 deletions

diff --git a/docs/cspan_api.md b/docs/cspan_api.md
index c3556dc3..0f345a7e 100644
--- a/docs/cspan_api.md
+++ b/docs/cspan_api.md
@@ -147,7 +147,6 @@ if __name__ == '__main__':
 
     for i in a.flat: print(" {}".format(i), end='')
     print('')
-    # prints row-wise
     for i in b.flat: print(" {}".format(i), end='')
 '''
  13 14 15 16
@@ -198,7 +197,6 @@ int main(void) {
 
     c_foreach (i, myspan2, a) printf(" %d", *i.ref);
     puts("");
-    // prints in storage order, same as a.
     c_foreach (i, myspan2, b) printf(" %d", *i.ref);
 }
 ```
diff --git a/include/stc/cspan.h b/include/stc/cspan.h
index f9f3b02a..8a422dad 100644
--- a/include/stc/cspan.h
+++ b/include/stc/cspan.h
@@ -97,9 +97,8 @@ int demo2() {
         return it; \
     } \
     STC_INLINE void Self##_next(Self##_iter* it) { \
-        int i, inc, done; \
-        if (cspan_is_colmajor(it->_s)) i=0, inc=1; else i=RANK-1, inc=-1; \
-        it->ref += _cspan_next##RANK(it->pos, it->_s->shape, it->_s->stride.d, RANK, i, inc, &done); \
+        int done; \
+        it->ref += _cspan_next##RANK(it->pos, it->_s->shape, it->_s->stride.d, RANK, &done); \
         if (done) it->ref = NULL; \
     } \
     struct stc_nostruct
@@ -227,8 +226,8 @@ STC_INLINE intptr_t _cspan_idxN(int rank, const int32_t shape[], const int32_t s
     return off;
 }
 
-STC_API intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int rank, int i, int inc, int* done);
-#define _cspan_next1(pos, shape, stride, rank, i, inc, done) (*done = ++pos[0]==shape[0], (void)(i|inc), stride[0])
+STC_API intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int rank, int* done);
+#define _cspan_next1(pos, shape, stride, rank, done) (*done = ++pos[0]==shape[0], stride[0])
 #define _cspan_next3 _cspan_next2
 #define _cspan_next4 _cspan_next2
 #define _cspan_next5 _cspan_next2
@@ -246,16 +245,15 @@ STC_API int32_t* _cspan_shape2stride(cspan_layout layout, int32_t shape[], int r
 /* --------------------- IMPLEMENTATION --------------------- */
 #if defined(i_implement) || defined(i_static)
 
-STC_DEF intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int rank, int i, int inc, int* done) {
-    intptr_t off = stride[i];
-    ++pos[i];
+STC_DEF intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int r, int* done) {
+    intptr_t off = stride[--r];
+    ++pos[r];
 
-    while (--rank && pos[i] == shape[i]) {
-        pos[i] = 0; ++pos[i + inc];
-        off += stride[i + inc] - stride[i]*shape[i];
-        i += inc;
+    for (; r && pos[r] == shape[r]; --r) {
+        pos[r] = 0; ++pos[r - 1];
+        off += stride[r - 1] - stride[r]*shape[r];
     }
-    *done = pos[i] == shape[i];
+    *done = pos[r] == shape[r];
     return off;
 }