Added support for column-major md cspan.

API change: the create function cspan_md(order, array, d1, d2, ...) has the new first argument order, which must be either 'C' or 'F' (C: row-major or Fortran: column-major). The representation of strides was changed.

1 files changed, 93 insertions, 71 deletions

diff --git a/include/stc/cspan.h b/include/stc/cspan.h
index d7a72267..35f5029f 100644
--- a/include/stc/cspan.h
+++ b/include/stc/cspan.h
@@ -30,7 +30,7 @@ using_cspan(Intspan, int, 1);
 
 int demo1() {
     float raw[4*5];
-    Span2f ms = cspan_md(raw, 4, 5);
+    Span2f ms = cspan_md('C', raw, 4, 5);
     
     for (int i=0; i<ms.shape[0]; i++)
         for (int j=0; j<ms.shape[1]; j++)
@@ -72,7 +72,7 @@ int demo2() {
     typedef struct { \
         Self##_value *data; \
         int32_t shape[RANK]; \
-        cspan_idx##RANK stride; \
+        cspan_tup##RANK stride; \
     } Self; \
     \
     typedef struct { Self##_value *ref; int32_t pos[RANK]; const Self *_s; } Self##_iter; \
@@ -96,27 +96,34 @@ int demo2() {
         return it; \
     } \
     STC_INLINE void Self##_next(Self##_iter* it) { \
-        it->ref += _cspan_next##RANK(RANK, it->pos, it->_s->shape, it->_s->stride.d); \
-        if (it->pos[0] == it->_s->shape[0]) it->ref = NULL; \
+        static const struct { int8_t i, j, inc; } t[2] = {{RANK - 1, 0, -1}, {0, RANK - 1, 1}}; \
+        const int order = (it->_s->stride.d[0] < it->_s->stride.d[RANK - 1]); /* 0='C', 1='F' */ \
+        it->ref += _cspan_next##RANK(it->pos, it->_s->shape, it->_s->stride.d, RANK, t[order].i, t[order].inc); \
+        if (it->pos[t[order].j] == it->_s->shape[t[order].j]) \
+            it->ref = NULL; \
     } \
     struct stc_nostruct
 
 #define using_cspan2(Self, T) using_cspan_3(Self, T, 1); using_cspan_3(Self##2, T, 2)
 #define using_cspan3(Self, T) using_cspan2(Self, T); using_cspan_3(Self##3, T, 3)
 #define using_cspan4(Self, T) using_cspan3(Self, T); using_cspan_3(Self##4, T, 4)
-typedef struct { int32_t d[1]; } cspan_idx1;
-typedef struct { int32_t d[2]; } cspan_idx2;
-typedef struct { int32_t d[3]; } cspan_idx3;
-typedef struct { int32_t d[4]; } cspan_idx4;
-typedef struct { int32_t d[5]; } cspan_idx5;
-typedef struct { int32_t d[6]; } cspan_idx6;
+typedef struct { int32_t d[1]; } cspan_tup1;
+typedef struct { int32_t d[2]; } cspan_tup2;
+typedef struct { int32_t d[3]; } cspan_tup3;
+typedef struct { int32_t d[4]; } cspan_tup4;
+typedef struct { int32_t d[5]; } cspan_tup5;
+typedef struct { int32_t d[6]; } cspan_tup6;
 #define c_END -1
 #define c_ALL 0,-1
 
-#define cspan_md(array, ...) \
-    {.data=array, .shape={__VA_ARGS__}, .stride={.d={__VA_ARGS__}}}
+#define cspan_md(order, array, ...) \
+    {.data=array, .shape={__VA_ARGS__}, \
+     .stride=*(c_PASTE(cspan_tup, c_NUMARGS(__VA_ARGS__))*)_cspan_shape2stride(order, ((int32_t[]){__VA_ARGS__}), c_NUMARGS(__VA_ARGS__))}
 
-/* For static initialization, use cspan_init(). c_init() for non-static only. */
+#define cspan_transpose(self) \
+    _cspan_transpose((self)->shape, (self)->stride.d, cspan_rank(self))
+
+/* Use cspan_init() for static initialization only. c_init() for non-static init. */
 #define cspan_init(SpanType, ...) \
     {.data=(SpanType##_value[])__VA_ARGS__, .shape={sizeof((SpanType##_value[])__VA_ARGS__)/sizeof(SpanType##_value)}}
 
@@ -134,17 +141,18 @@ typedef struct { int32_t d[6]; } cspan_idx6;
 
 #define cspan_size(self) _cspan_size((self)->shape, cspan_rank(self))
 #define cspan_rank(self) c_arraylen((self)->shape)
+#define cspan_order(self) ((self)->stride.d[0] < (self)->stride.d[cspan_rank(self) - 1])
 
-#define cspan_index(self, ...) c_PASTE(cspan_idx_, c_NUMARGS(__VA_ARGS__))(self, __VA_ARGS__)
-#define cspan_idx_1 cspan_idx_4
-#define cspan_idx_2 cspan_idx_4
-#define cspan_idx_3 cspan_idx_4
-#define cspan_idx_4(self, ...) \
+#define cspan_index(self, ...) c_PASTE(cspan_tup_, c_NUMARGS(__VA_ARGS__))(self, __VA_ARGS__)
+#define cspan_tup_1 cspan_tup_3
+#define cspan_tup_2 cspan_tup_3
+#define cspan_tup_3(self, ...) \
     c_PASTE(_cspan_idx, c_NUMARGS(__VA_ARGS__))((self)->shape, (self)->stride, __VA_ARGS__) // small/fast
-#define cspan_idx_5(self, ...) \
+#define cspan_tup_4(self, ...) \
     (_cspan_idxN(c_NUMARGS(__VA_ARGS__), (self)->shape, (self)->stride.d, (int32_t[]){__VA_ARGS__}) + \
      c_static_assert(cspan_rank(self) == c_NUMARGS(__VA_ARGS__))) // general
-#define cspan_idx_6 cspan_idx_5
+#define cspan_tup_5 cspan_tup_4
+#define cspan_tup_6 cspan_tup_4
 
 #define cspan_at(self, ...) ((self)->data + cspan_index(self, __VA_ARGS__))
 #define cspan_front(self) ((self)->data)
@@ -162,19 +170,20 @@ typedef struct { int32_t d[6]; } cspan_idx6;
 #define cspan_submd4(...) c_MACRO_OVERLOAD(cspan_submd4, __VA_ARGS__)
 #define cspan_submd3(...) c_MACRO_OVERLOAD(cspan_submd3, __VA_ARGS__)
 #define cspan_submd2(self, x) \
-    {.data=cspan_at(self, x, 0), .shape={(self)->shape[1]}}
+    {.data=cspan_at(self, x, 0), .shape={(self)->shape[1]}, .stride={.d={(self)->stride.d[1]}}}
 #define cspan_submd3_2(self, x) \
     {.data=cspan_at(self, x, 0, 0), .shape={(self)->shape[1], (self)->shape[2]}, \
-                                    .stride={.d={0, (self)->stride.d[2]}}}
+                                    .stride={.d={(self)->stride.d[1], (self)->stride.d[2]}}}
 #define cspan_submd3_3(self, x, y) \
-    {.data=cspan_at(self, x, y, 0), .shape={(self)->shape[2]}}
+    {.data=cspan_at(self, x, y, 0), .shape={(self)->shape[2]}, .stride={.d={(self)->stride.d[2]}}}
 #define cspan_submd4_2(self, x) \
     {.data=cspan_at(self, x, 0, 0, 0), .shape={(self)->shape[1], (self)->shape[2], (self)->shape[3]}, \
-                                       .stride={.d={0, (self)->stride.d[2], (self)->stride.d[3]}}}
+                                       .stride={.d={(self)->stride.d[1], (self)->stride.d[2], (self)->stride.d[3]}}}
 #define cspan_submd4_3(self, x, y) \
-    {.data=cspan_at(self, x, y, 0, 0), .shape={(self)->shape[2], (self)->shape[3]}, .stride={.d={0, (self)->stride.d[3]}}}
+    {.data=cspan_at(self, x, y, 0, 0), .shape={(self)->shape[2], (self)->shape[3]}, \
+                                       .stride={.d={(self)->stride.d[2], (self)->stride.d[3]}}}
 #define cspan_submd4_4(self, x, y, z) \
-    {.data=cspan_at(self, x, y, z, 0), .shape={(self)->shape[3]}}
+    {.data=cspan_at(self, x, y, z, 0), .shape={(self)->shape[3]}, .stride={.d={(self)->stride.d[3]}}}
 
 // private definitions:
 
@@ -184,83 +193,96 @@ STC_INLINE intptr_t _cspan_size(const int32_t shape[], int rank) {
     return sz;
 }
 
-STC_INLINE intptr_t _cspan_idx1(const int32_t shape[1], const cspan_idx1 stri, int32_t x)
+STC_INLINE void _cspan_transpose(int32_t shape[], int32_t stride[], int rank) {
+    for (int i = 0; i < --rank; ++i) {
+        c_swap(int32_t, shape + i, shape + rank);
+        c_swap(int32_t, stride + i, stride + rank);
+    }
+}
+
+STC_INLINE intptr_t _cspan_idx1(const int32_t shape[1], const cspan_tup1 stri, int32_t x)
     { c_ASSERT(c_LTu(x, shape[0])); return x; }
 
-STC_INLINE intptr_t _cspan_idx2(const int32_t shape[2], const cspan_idx2 stri, int32_t x, int32_t y)
-    { c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1])); return (intptr_t)stri.d[1]*x + y; }
+STC_INLINE intptr_t _cspan_idx2(const int32_t shape[2], const cspan_tup2 stri, int32_t x, int32_t y)
+    { c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1])); return (intptr_t)stri.d[0]*x + stri.d[1]*y; }
 
-STC_INLINE intptr_t _cspan_idx3(const int32_t shape[3], const cspan_idx3 stri, int32_t x, int32_t y, int32_t z) {
+STC_INLINE intptr_t _cspan_idx3(const int32_t shape[3], const cspan_tup3 stri, int32_t x, int32_t y, int32_t z) {
     c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1]) && c_LTu(z, shape[2]));
-    return (intptr_t)stri.d[2]*(stri.d[1]*x + y) + z;
+    return (intptr_t)stri.d[0]*x + stri.d[1]*y + stri.d[2]*z;
 }
-STC_INLINE intptr_t _cspan_idx4(const int32_t shape[4], const cspan_idx4 stri, int32_t x, int32_t y,
-                                                                               int32_t z, int32_t w) {
-    c_ASSERT(c_LTu(x, shape[0]) && c_LTu(y, shape[1]) && c_LTu(z, shape[2]) && c_LTu(w, shape[3]));
-    return (intptr_t)stri.d[3]*(stri.d[2]*(stri.d[1]*x + y) + z) + w;
+
+STC_INLINE intptr_t _cspan_idxN(int rank, const int32_t shape[], const int32_t stride[], const int32_t a[]) {
+    intptr_t off = 0;
+    while (rank--) {
+        c_ASSERT(c_LTu(a[rank], shape[rank]));
+        off += stride[rank]*a[rank];
+    }
+    return off;
 }
 
-STC_API intptr_t _cspan_idxN(int rank, const int32_t shape[], const int32_t stri[], const int32_t a[]);
-STC_API intptr_t _cspan_next2(int rank, int32_t pos[], const int32_t shape[], const int32_t stride[]);
-#define _cspan_next1(r, pos, d, s) (++pos[0], 1)
+#define _cspan_next1(pos, d, s, r, i, inc) (++pos[0], 1)
+STC_API intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int rank, int i, int inc);
 #define _cspan_next3 _cspan_next2
 #define _cspan_next4 _cspan_next2
 #define _cspan_next5 _cspan_next2
 #define _cspan_next6 _cspan_next2
 
-STC_API intptr_t _cspan_slice(int32_t odim[], int32_t ostri[], int* orank, 
-                              const int32_t shape[], const int32_t stri[], 
+STC_API intptr_t _cspan_slice(int32_t oshape[], int32_t ostride[], int* orank,
+                              const int32_t shape[], const int32_t stride[],
                               int rank, const int32_t a[][2]);
+
+STC_API int32_t* _cspan_shape2stride(char order, int32_t shape[], int rank);
 #endif // STC_CSPAN_H_INCLUDED
 
 /* -------------------------- IMPLEMENTATION ------------------------- */
 #if defined(i_implement) || defined(i_static)
 
-STC_DEF intptr_t _cspan_idxN(int rank, const int32_t shape[], const int32_t stri[], const int32_t a[]) {
-    intptr_t off = a[0];
-    c_ASSERT(c_LTu(a[0], shape[0]));
-    for (int i = 1; i < rank; ++i) {
-        off *= stri[i];
-        off += a[i];
-        c_ASSERT(c_LTu(a[i], shape[i]));
+STC_DEF int32_t* _cspan_shape2stride(char order, int32_t shape[], int rank) {
+    int32_t k = 1, i, j, inc, s1, s2;
+    if (order == 'F') i = 0, j = rank, inc = 1; 
+    else  /* 'C' */   i = rank - 1, j = -1, inc = -1;
+    s1 = shape[i]; shape[i] = 1;
+
+    for (i += inc; i != j; i += inc) {
+        s2 = shape[i];
+        shape[i] = (k *= s1);
+        s1 = s2;
     }
-    return off;
+    return shape;
 }
 
-STC_DEF intptr_t _cspan_next2(int rank, int32_t pos[], const int32_t shape[], const int32_t stride[]) {
-    intptr_t off = 1, rs = 1;
-    ++pos[rank - 1];
-    while (--rank && pos[rank] == shape[rank]) {
-        pos[rank] = 0, ++pos[rank - 1];
-        const intptr_t ds = rs*shape[rank];
-        rs *= stride[rank];
-        off += rs - ds;
+STC_DEF intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int rank, int i, int inc) {
+    intptr_t off = stride[i];
+    ++pos[i];
+    for (; --rank && pos[i] == shape[i]; i += inc) {
+        pos[i] = 0; ++pos[i + inc];
+        off += stride[i + inc] - stride[i]*shape[i];
     }
     return off;
 }
 
-STC_DEF intptr_t _cspan_slice(int32_t odim[], int32_t ostri[], int* orank, 
-                              const int32_t shape[], const int32_t stri[], 
+STC_DEF intptr_t _cspan_slice(int32_t oshape[], int32_t ostride[], int* orank, 
+                              const int32_t shape[], const int32_t stride[], 
                               int rank, const int32_t a[][2]) {
     intptr_t off = 0;
-    int i = 0, j = 0;
-    int32_t t, s = 1;
+    int i = 0, oi = 0;
+    int32_t end;
     for (; i < rank; ++i) {
-        off *= stri[i];
-        off += a[i][0];
-        switch (a[i][1]) { 
-            case 0: s *= stri[i]; c_ASSERT(c_LTu(a[i][0], shape[i])); continue;
-            case -1: t = shape[i]; break;
-            default: t = a[i][1]; break; 
+        off += stride[i]*a[i][0];
+        switch (a[i][1]) {
+            case 0: c_ASSERT(c_LTu(a[i][0], shape[i])); continue;
+            case -1: end = shape[i]; break;
+            default: end = a[i][1];
         }
-        odim[j] = t - a[i][0];
-        ostri[j] = s*stri[i];
-        c_ASSERT(c_LTu(0, odim[j]) & !c_LTu(shape[i], t));
-        s = 1; ++j;
+        oshape[oi] = end - a[i][0];
+        ostride[oi] = stride[i];
+        c_ASSERT(c_LTu(0, oshape[oi]) & !c_LTu(shape[i], end));
+        ++oi;
     }
-    *orank = j;
+    *orank = oi;
     return off;
 }
+
 #endif
 #undef i_opt
 #undef i_header