diff options
Diffstat (limited to 'include/stc/cspan.h')
| -rw-r--r-- | include/stc/cspan.h | 19 |
1 files changed, 12 insertions, 7 deletions
diff --git a/include/stc/cspan.h b/include/stc/cspan.h index b8b191f1..32921390 100644 --- a/include/stc/cspan.h +++ b/include/stc/cspan.h @@ -244,24 +244,29 @@ STC_API int32_t* _cspan_shape2stride(char order, int32_t shape[], int rank); STC_DEF intptr_t _cspan_next2(int32_t pos[], const int32_t shape[], const int32_t stride[], int rank, int* done) { int i, inc; - if (stride[0] < stride[rank - 1]) i = rank - 1, inc = -1; else i = 0, inc = 1; + if (stride[0] < stride[rank - 1]) i = rank - 1, inc = -1; + else /* order 'C' */ i = 0, inc = 1; + intptr_t off = stride[i]; ++pos[i]; - for (; --rank && pos[i] == shape[i]; i += inc) { + while (--rank && pos[i] == shape[i]) { pos[i] = 0; ++pos[i + inc]; off += stride[i + inc] - stride[i]*shape[i]; + i += inc; } *done = pos[i] == shape[i]; return off; } 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; + int i, inc; + if (order == 'F') i = 0, inc = 1; + else i = rank - 1, inc = -1; + int32_t k = 1, s1 = shape[i], s2; - for (i += inc; i != j; i += inc) { + shape[i] = 1; + while (--rank) { + i += inc; s2 = shape[i]; shape[i] = (k *= s1); s1 = s2; |