Fix copying reshape.

Closes #986.
ashinn · May 31, 2024 · 0b55c0b · 0b55c0b
1 parent 4b5ab83
commit 0b55c0b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/lib/srfi/231/test.sld b/lib/srfi/231/test.sld
@@ -3767,12 +3767,13 @@
            '#(1 1 2 1))
           (make-interval '#(4))))
 
-        '(test #t
-            (specialized-array-reshape
-             (array-sample (array-copy (make-array (make-interval '#(3 4)) list))
-                           '#(2 1))
-             (make-interval '#(8))
-             #t))
+        (test '((0 0) (0 1) (0 2) (0 3) (2 0) (2 1) (2 2) (2 3))
+            (array->list*
+             (specialized-array-reshape
+              (array-sample (array-copy (make-array (make-interval '#(3 4)) list))
+                            '#(2 1))
+              (make-interval '#(8))
+              #t)))
         (test '(() ())
             (array->list*
              (specialized-array-reshape

diff --git a/lib/srfi/231/transforms.scm b/lib/srfi/231/transforms.scm
@@ -514,6 +514,24 @@
        (array-domain source))
       destination)))
 
+(define (array-assign/reshape! destination source)
+  (let ((dest-domain (array-domain destination))
+        (source-domain (array-domain source)))
+    (assert (and (mutable-array? destination) (array? source)
+                 (= (interval-volume dest-domain)
+                    (interval-volume source-domain))))
+    (let ((getter (array-getter source))
+          (setter (array-setter destination)))
+      (let lp ((source-ivc (interval-cursor source-domain))
+               (dest-ivc (interval-cursor dest-domain)))
+        (apply setter
+               (apply getter (interval-cursor-get source-ivc))
+               (interval-cursor-get dest-ivc))
+        (when (and (interval-cursor-next! source-ivc)
+                   (interval-cursor-next! dest-ivc))
+          (lp source-ivc dest-ivc)))
+      destination)))
+
 (define (reshape-without-copy array new-domain)
   (let* ((domain (array-domain array))
          (orig-indexer (array-indexer array))
@@ -568,7 +586,7 @@
                   new-domain
                   (array-storage-class array)
                   (array-safe? array))))
-        (array-assign! res array)
+        (array-assign/reshape! res array)
         res))
      (else
       (error "can't reshape" array new-domain)))))