flatten-/expand-array

borodust · Feb 7, 2017 · 2ab72e8 · 2ab72e8
1 parent ffead82
commit 2ab72e8
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 24 deletions.
diff --git a/cl-bodge.asd b/cl-bodge.asd
@@ -15,7 +15,8 @@
   :pathname "utils"
   :serial t
   :components ((:file "packages")
-               (:file "utils")))
+               (:file "utils")
+               (:file "arrays")))
 
 
 (defsystem cl-bodge/engine

diff --git a/utils/arrays.lisp b/utils/arrays.lisp
@@ -0,0 +1,63 @@
+(in-package :cl-bodge.utils)
+
+
+(defclass array-slice ()
+  ((dimensions :initform nil)
+   (slice :initform nil)))
+
+
+(defmethod initialize-instance :after ((this array-slice) &key dimensions)
+  (with-slots ((dims dimensions) slice) this
+    (setf slice (mapcar (constantly 0) dimensions)
+          dims (reverse dimensions))))
+
+
+(defun %next-slice (dimensions slice)
+  (when slice
+    (if (= (1- (first dimensions)) (first slice))
+        (when (%next-slice (rest dimensions) (rest slice))
+          (setf (first slice) 0))
+        (incf (first slice)))))
+
+
+(defun next-slice (obj)
+  (with-slots (dimensions slice) obj
+    (when (%next-slice dimensions slice)
+      (reverse slice))))
+
+
+(defun slice (obj)
+  (with-slots (slice) obj
+    slice))
+
+
+(defun flat-index (obj)
+  (with-slots (dimensions slice) obj
+    (labels ((total-offset (offsets sizes prev)
+               (if offsets
+                   (+ (* (first offsets) (apply #'* prev))
+                      (total-offset (rest offsets) (rest sizes) (cons (first sizes) prev)))
+                   0)))
+      (total-offset slice dimensions nil))))
+
+
+(defun flatten-array (array)
+  (let* ((dimensions (array-dimensions array))
+         (slice (make-instance 'array-slice :dimensions dimensions))
+         (result (make-array (apply #'* dimensions) :element-type (array-element-type array))))
+    (loop for slice-indexes = (slice slice) then (next-slice slice) while slice-indexes
+       do (setf (aref result (flat-index slice))
+                (apply #'aref array slice-indexes)))
+    result))
+
+
+
+(defun expand-array (array dimensions)
+  (unless (= (apply #'* (ensure-list dimensions)) (length array))
+    (error "incorrect dimensions"))
+  (let* ((slice (make-instance 'array-slice :dimensions dimensions))
+         (result (make-array dimensions :element-type (array-element-type array))))
+    (loop for slice-indexes = (slice slice) then (next-slice slice) while slice-indexes
+       do (setf (apply #'aref result slice-indexes)
+                (aref array (flat-index slice))))
+    result))
diff --git a/utils/packages.lisp b/utils/packages.lisp
@@ -59,6 +59,7 @@
            reexporting
            when-debugging
            flatten-array
+           expand-array
            split-sequence
            foreign-function-pointer
            stringify))
diff --git a/utils/utils.lisp b/utils/utils.lisp
@@ -229,29 +229,6 @@
                    :initial-contents list)))))
 
 
-(defun flatten-array (array)
-  (let* ((dims (array-dimensions array)))
-    (labels ((total-offset (offsets sizes)
-               (if offsets
-                   (+ (* (first offsets) (reduce #'* sizes))
-                      (total-offset (rest offsets) (rest sizes)))
-                   0))
-             (copy-dimension (dst src offset rest-dims)
-               (if (null rest-dims)
-                   (let ((dst-offset (total-offset offset (rest dims))))
-                     (setf (aref dst dst-offset) (apply #'aref src offset)))
-                   (loop with size = (first rest-dims)
-                      for i from 0 below size
-                      do (copy-dimension dst src (append offset (list i)) (rest rest-dims))))))
-
-      (if (null (cdr dims))
-          array
-          (let* ((result (make-array (reduce #'* dims)
-                                     :element-type (array-element-type array))))
-            (copy-dimension result array '() dims)
-            result)))))
-
-
 (defun foreign-function-pointer (function-name)
   (when-let* ((fn (autowrap:find-function function-name)))
     (let ((name (autowrap:foreign-symbol-c-symbol fn)))