fast-path vector-map on two vectors

ashinn · Nov 13, 2024 · 3716d99 · 3716d99
1 parent 49072eb
commit 3716d99
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/lib/scheme/extras.scm b/lib/scheme/extras.scm
@@ -187,12 +187,22 @@
             (lp (cdr ls) (+ i v-len)))))))
 
 (define (vector-map proc vec . lov)
-  (if (null? lov)
+  (cond
+   ((null? lov)
+    (let lp ((i (vector-length vec)) (res '()))
+      (if (zero? i)
+          (list->vector res)
+          (lp (- i 1) (cons (proc (vector-ref vec (- i 1))) res)))))
+   ((null? (cdr lov))
+    (let ((vec2 (car lov)))
       (let lp ((i (vector-length vec)) (res '()))
         (if (zero? i)
             (list->vector res)
-            (lp (- i 1) (cons (proc (vector-ref vec (- i 1))) res))))
-      (list->vector (apply map proc (map vector->list (cons vec lov))))))
+            (lp (- i 1)
+                (cons (proc (vector-ref vec (- i 1)) (vector-ref vec2 (- i 1)))
+                      res))))))
+   (else
+    (list->vector (apply map proc (map vector->list (cons vec lov)))))))
 
 (define (vector-for-each proc vec . lov)
   (if (null? lov)