Import Upstream version 2c0upstream/2c0

1 files changed, 8 insertions, 6 deletions

diff --git a/arraymap.scm b/arraymap.scm
index 18ee64a..d3dedba 100644
--- a/arraymap.scm
+++ b/arraymap.scm
@@ -52,25 +52,27 @@
 	    (rafe crshape (cons i inds))))))
   (rafe (array-shape (car ras)) '()))
 
-(define (shape->indexes shape)
-  (define ra0 (apply make-array '() shape))
+(define (array-index-map! ra fun)
   (define (ramap rshape inds)
     (if (null? (cdr rshape))
 	(do ((i (cadar rshape) (+ -1 i))
 	     (is (cons (cadar rshape) inds)
 		 (cons (+ -1 i) inds)))
 	    ((< i (caar rshape)))
-	  (apply array-set! ra0 is is))
+	  (apply array-set! ra (apply fun is) is))
 	(let ((crshape (cdr rshape))
 	      (ll (caar rshape)))
 	  (do ((i (cadar rshape) (+ -1 i)))
 	      ((< i ll))
 	    (ramap crshape (cons i inds))))))
-  (ramap (reverse shape) '())
-  ra0)
+  (if (zero? (array-rank ra))
+      (array-set! ra (fun))
+      (ramap (reverse (array-shape ra)) '())))
 
 (define (array-indexes ra)
-  (shape->indexes (array-shape ra)))
+  (let ((ra0 (apply make-array '() (array-shape ra))))
+    (array-index-map! ra0 list)
+    ra0))
 
 (define (array-copy! source dest)
   (array-map! dest identity source))