diff --git a/lib/apl/runtime.sx b/lib/apl/runtime.sx index 30081175..8f53df38 100644 --- a/lib/apl/runtime.sx +++ b/lib/apl/runtime.sx @@ -606,3 +606,58 @@ ((old-c-last (apl-safe-mod (+ c-last rot) last-dim))) (nth ravel (+ flat (- old-c-last c-last)))))) (range 0 n)))))))) + +(define + apl-catenate + (fn + (a b) + (let + ((a-s (if (scalar? a) (list 1) (get a :shape))) + (b-s (if (scalar? b) (list 1) (get b :shape))) + (a-r (get a :ravel)) + (b-r (get b :ravel))) + (let + ((a-last (last a-s)) (prefix (take a-s (- (len a-s) 1)))) + (let + ((new-shape (append prefix (list (+ a-last (last b-s))))) + (a-strides (apl-strides a-s)) + (b-strides (apl-strides b-s))) + (let + ((new-size (reduce * 1 new-shape)) + (new-strides (apl-strides new-shape))) + (make-array + new-shape + (map + (fn + (new-flat) + (let + ((new-coords (apl-flat->multi new-flat new-shape new-strides))) + (let + ((last-c (last new-coords)) + (prefix-c (take new-coords (- (len new-coords) 1)))) + (if + (< last-c a-last) + (nth + a-r + (apl-multi->flat + (append prefix-c (list last-c)) + a-strides)) + (nth + b-r + (apl-multi->flat + (append prefix-c (list (- last-c a-last))) + b-strides)))))) + (range 0 new-size))))))))) + +(define + apl-catenate-first + (fn + (a b) + (let + ((a-s (if (scalar? a) (list 1) (get a :shape))) + (b-s (if (scalar? b) (list 1) (get b :shape))) + (a-r (get a :ravel)) + (b-r (get b :ravel))) + (make-array + (cons (+ (first a-s) (first b-s)) (rest a-s)) + (append a-r b-r))))) diff --git a/lib/apl/tests/structural.sx b/lib/apl/tests/structural.sx index 72b1e961..4a6dddd8 100644 --- a/lib/apl/tests/structural.sx +++ b/lib/apl/tests/structural.sx @@ -322,4 +322,70 @@ (apl-rotate-first (apl-scalar 1) (make-array (list 2 3) (list 1 2 3 4 5 6)))) - (list 4 5 6 1 2 3)) \ No newline at end of file + (list 4 5 6 1 2 3)) + +(apl-test + "cat v,v ravel" + (rv + (apl-catenate + (make-array (list 3) (list 1 2 3)) + (make-array (list 2) (list 4 5)))) + (list 1 2 3 4 5)) + +(apl-test + "cat v,v shape" + (sh + (apl-catenate + (make-array (list 3) (list 1 2 3)) + (make-array (list 2) (list 4 5)))) + (list 5)) + +(apl-test + "cat scalar,v" + (rv (apl-catenate (apl-scalar 99) (make-array (list 3) (list 1 2 3)))) + (list 99 1 2 3)) + +(apl-test + "cat v,scalar" + (rv (apl-catenate (make-array (list 3) (list 1 2 3)) (apl-scalar 99))) + (list 1 2 3 99)) + +(apl-test + "cat matrix last-axis shape" + (sh + (apl-catenate + (make-array (list 2 3) (list 1 2 3 4 5 6)) + (make-array (list 2 2) (list 7 8 9 10)))) + (list 2 5)) + +(apl-test + "cat matrix last-axis ravel" + (rv + (apl-catenate + (make-array (list 2 3) (list 1 2 3 4 5 6)) + (make-array (list 2 2) (list 7 8 9 10)))) + (list 1 2 3 7 8 4 5 6 9 10)) + +(apl-test + "cat-first v,v shape" + (sh + (apl-catenate-first + (make-array (list 3) (list 1 2 3)) + (make-array (list 2) (list 4 5)))) + (list 5)) + +(apl-test + "cat-first matrix shape" + (sh + (apl-catenate-first + (make-array (list 2 3) (list 1 2 3 4 5 6)) + (make-array (list 3 3) (list 11 12 13 14 15 16 17 18 19)))) + (list 5 3)) + +(apl-test + "cat-first matrix ravel" + (rv + (apl-catenate-first + (make-array (list 2 3) (list 1 2 3 4 5 6)) + (make-array (list 3 3) (list 11 12 13 14 15 16 17 18 19)))) + (list 1 2 3 4 5 6 11 12 13 14 15 16 17 18 19)) \ No newline at end of file