Base ra-index-map! on ra-slice-for-each-2

This is a bit faster even without special casing the 1-cell.

* mod/newra/lib.scm (ra-index-map!): As stated.

Author: lloda

mod/newra/from.scm

@@ -178,15 +178,15 @@
 ; FIXME going over the args twice, here and in fromb
 
 (define (parse-args A . i)
-  (let ((rank (ra-rank A)))
+  (let ((r (ra-rank A)))
     (let loop ((j 0) (m 0) (ii i)
                (ib '()) (ibi '()) (tb '())
                (iu '()) (iui '()) (tu '()))
       (match ii
         ((i0 . irest)
          (match i0
            (($ <dots> n)
-            (let* ((k (or n (- rank j (count-axes-left irest))))
+            (let* ((k (or n (- r j (count-axes-left irest))))
                    (idest (iota k m)))
               (loop (+ j k) (+ m k) irest
                     (cons (dots k) ib) (fold cons ibi (iota k j)) (fold cons tb idest)
@@ -202,8 +202,8 @@
                       ib ibi tb
                       (cons i0 iu) (cons j iui) (fold cons tu idest)))))))
         (()
-         (when (> j rank)
-           (throw 'too-many-indices-for-rank-of-A j rank))
+         (when (> j r)
+           (throw 'too-many-indices-for-rank-of-A j r))
          (let ((ib (reverse! ib))
                (ibi (reverse! ibi))
                (iu (reverse! iu))

mod/newra/lib.scm

@@ -149,8 +149,6 @@ See also: @code{make-ra}
                             (loop (+ k 1)))))))))
           (make-ra-root (%%ra-root oldra) dims (- ref (ra-offset 0 dims))))))))
 
-; FIXME Depends on traversal order of ra-for-each.
-
 (define (ra->list ra)
   "
 Return a nested list of the elements of array @var{ra}. For example, if @var{ra} is a
@@ -160,32 +158,28 @@ list contains a list for each of the rows of @var{ra}; and so on.
 See also: @code{as-ra}
 "
   (let* ((ra (ra-check ra))
-         (rank (%%ra-rank ra)))
-    (cond
-     ((zero? rank) (ra-ref ra))
-     (else
-      (let ((ra (apply ra-reverse ra (iota rank))))
-        (match (vector-ref (%%ra-dims ra) (- rank 1))
-          (($ <dim> klen klo kstep)
-           (let loop-rank ((ra ra))
-             (cond
-              ((= 1 (%%ra-rank ra))
-               (if (> klen 20)
-                 (ra-fold xcons '() ra)
-                 (let loop-dim ((l '()) (i klo))
-                   (if (> i (dim-hi klen klo))
-                     l
-                     (loop-dim (cons (ra-ref ra i) l) (+ i 1))))))
-              (else
-               (let ((l '()))
-                 (ra-slice-for-each 1 (lambda (x) (set! l (cons (loop-rank x) l))) ra)
-                 l)))))))))))
-
-; Similar to (@ (newra) ra-for-each-slice-1) - since we cannot unroll. It
-; might be cheaper to go Fortran order (building the index lists back to front);
-; should try that. C order and set-cdr! is how oldra does it.
-; This function is provided for compatibility with oldra; generally we shouldn't
-; be building index lists.
+         (rank (%%ra-rank ra))
+         (ra (apply ra-reverse ra (iota rank))))
+    (if (zero? rank)
+      (ra-ref ra)
+      (match (vector-ref (%%ra-dims ra) (- rank 1))
+        (($ <dim> klen klo _)
+         (let loop-rank ((ra ra))
+           (if (= 1 (%%ra-rank ra))
+             (if (> klen 20)
+               (ra-fold xcons '() ra)
+               (let loop-dim ((l '()) (i klo))
+                 (if (> i (dim-hi klen klo))
+                   l
+                   (loop-dim (cons (ra-ref ra i) l) (+ i 1)))))
+             (let ((l '()))
+               (ra-slice-for-each-in-order 1 (lambda (x) (set! l (cons (loop-rank x) l))) ra)
+               l))))))))
+
+; Based on (@ (newra) ra-for-each-slice-2) - we cannot unroll. It might be cheaper to
+; go Fortran order (building the index lists back to front); should try that. C order
+; and set-cdr! is how oldra does it. This function is provided for compatibility with
+; oldra; generally we shouldn't be building index lists.
 
 (define (ra-index-map! ra op)
   "
@@ -204,30 +198,26 @@ x @result{} #%2:2:3(((0 0) (0 1) (0 2)) ((1 0) (1 1) (1 2)))
 
 See also: @code{ra-iota} @code{ra-i}
 "
-  (let* ((kk (ra-rank ra))
-         (ii (make-list kk))
-         (los lens ((@ (newra map) ra-slice-for-each-check) kk ra)))
-      (if (= kk 0)
-        (ra-set! ra (apply op ii))
-        (let loop-rank ((k 0) (ra ra) (endi ii))
-          (let* ((lo (vector-ref los k))
-                 (end (+ lo (vector-ref lens k))))
-            (if (= (+ 1 k) kk)
-              (let loop-dim ((i lo))
-                (if (= i end)
-                  (set-car! endi lo)
-                  (begin
-                    (set-car! endi i)
-                    (ra-set! ra (apply op ii) i)
-                    (loop-dim (+ i 1)))))
-              (let loop-dim ((i lo))
-                (if (= i end)
-                  (set-car! endi lo)
-                  (begin
-                    (set-car! endi i)
-                    (loop-rank (+ k 1) (ra-slice ra i) (cdr endi))
-                    (loop-dim (+ i 1)))))))))
-      ra))
+  (let* ((frame (ra-check ra))
+         (ii (make-list (%%ra-rank frame)))
+         (dims (%%ra-dims frame))
+         (ra (make-ra-root (%%ra-root frame) #() (ra-offset frame))))
+    (let loop-rank ((k 0) (endi ii))
+      (if (null? endi)
+        ((%%ra-vset! ra) (%%ra-root ra) (%%ra-zero ra) (apply op ii))
+        (match (vector-ref dims k)
+          (($ <dim> len lo step)
+           (let loop-dim ((i 0))
+             (if (= i len)
+               (begin
+                 (set-car! endi lo)
+                 (%%ra-zero-set! ra (- (%%ra-zero ra) (* step len))))
+               (begin
+                 (set-car! endi (+ i lo))
+                 (loop-rank (+ k 1) (cdr endi))
+                 (%%ra-zero-set! ra (+ (%%ra-zero ra) step))
+                 (loop-dim (+ i 1)))))))))
+    frame))
 
 
 ; ----------------

mod/newra/map.scm

@@ -149,7 +149,7 @@
              (%let ((ra (frame) (make-ra-root-prefix frame k los)) ...)
 ; since we'll unroll, special case for rank 0
                (if (zero? k)
-; no fresh slice descriptor like in array-slice-for-each. Should be all right b/c the descriptors can be copied.
+; no fresh slice descriptor like in array-slice-for-each. Should be all right bc descriptors can be copied.
                  (%op ra ...)
 ; check early so we can save a step in the loop later.
                  (when (vector-every positive? lens)

mod/test/map-ladder.scm

@@ -43,25 +43,22 @@
          (ra (map (cut make-ra-root-prefix <> kk los) frame)))
     (let loop-rank ((k 0))
       (if (= k kk)
-; no fresh slice descriptor like in array-slice-for-each. Should be all right b/c the descriptors can be copied.
+; no fresh slice descriptor like in array-slice-for-each. Should be all right bc descriptors can be copied.
         (apply op ra)
-        (let  ((lenk (vector-ref lens k)))
+        (let ((len (vector-ref lens k)))
           (let loop-dim ((i 0))
-            (cond
-             ((= i lenk)
-              (for-each
-                  (lambda (ra frame)
-                    (let ((step (dim-step (vector-ref (%%ra-dims frame) k))))
-                      (%%ra-zero-set! ra (- (%%ra-zero ra) (* step lenk)))))
-                ra frame))
-             (else
-              (loop-rank (+ k 1))
-              (for-each
-                  (lambda (ra frame)
-                    (let ((step (dim-step (vector-ref (%%ra-dims frame) k))))
-                      (%%ra-zero-set! ra (+ (%%ra-zero ra) step))))
+            (if (= i len)
+              (for-each (lambda (ra frame)
+                          (let ((step (dim-step (vector-ref (%%ra-dims frame) k))))
+                            (%%ra-zero-set! ra (- (%%ra-zero ra) (* step len)))))
                 ra frame)
-              (loop-dim (+ i 1))))))))))
+              (begin
+                (loop-rank (+ k 1))
+                (for-each (lambda (ra frame)
+                            (let ((step (dim-step (vector-ref (%%ra-dims frame) k))))
+                              (%%ra-zero-set! ra (+ (%%ra-zero ra) step))))
+                  ra frame)
+                (loop-dim (+ i 1))))))))))
 
 ; moving slice with row-major unrolling.
 (define (ra-slice-for-each-3 u op . frame)
@@ -91,7 +88,7 @@
               (if (= k u)
 ; unrolled dimensions.
                 (let loop ((i lenm))
-; no fresh slice descriptor like in array-slice-for-each. Should be all right b/c the descriptors can be copied.
+; no fresh slice descriptor like in array-slice-for-each. Should be all right bc descriptors can be copied.
                   (apply op ra)
                   (cond
                    ((zero? i)

test/everything.scm

@@ -639,6 +639,8 @@
 
 (test-equal "#%2:2:3(((0 0) (0 1) (0 2)) ((1 0) (1 1) (1 2)))"
   (ra->string (let ((x (make-ra 0 2 3))) (ra-index-map! x (lambda x x)))))
+(test-equal "#%2@2:2@3:3(((2 3) (2 4) (2 5)) ((3 3) (3 4) (3 5)))"
+  (ra->string (let ((x (make-ra 0 '(2 3) '(3 5)))) (ra-index-map! x (lambda x x)))))
 (test-equal "#%0(99)"
   (ra->string (let ((x (make-ra 0))) (ra-index-map! x (lambda x 99)))))