Don't realize lazy seq more than once in reductions

squint-cljs · Dec 7, 2023 · aaac8f4 · aaac8f4
1 parent 4ca081f
commit aaac8f4
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 9 deletions.
diff --git a/src/squint/core.js b/src/squint/core.js
@@ -583,20 +583,36 @@ export function reduce(f, arg1, arg2) {
 }
 
 function _reductions2(f, coll) {
-  const s = seq(coll);
+  let s = iterable(coll)[Symbol.iterator]();
+  let fst, rst;
+  const vd = s.next();
+  if (vd.done) {
+    s = null;
+  } else {
+    fst = vd.value;
+    rst = s;
+  }
   return new LazySeq(function () {
-    return s ? _reductions3(f, first(s), rest(s)) : list(f());
+    return s ? _reductions3(f, fst, rst) : list(f());
   });
 }
 
 function _reductions3(f, init, coll) {
-  const s = seq(coll);
-    if (reduced_QMARK_(init)) {
-      return list(init.value);
-    }
-    return cons(init, new LazySeq(function () {
+  if (reduced_QMARK_(init)) {
+    return list(init.value);
+  }
+  let s = iterable(coll)[Symbol.iterator]();
+  let fst, rst;
+  const vd = s.next();
+  if (vd.done) {
+    s = null;
+  } else {
+    fst = vd.value;
+    rst = s;
+  }
+  return cons(init, new LazySeq(function () {
       if (s) {
-        return _reductions3(f, f(init, first(s)), rest(s));
+        return _reductions3(f, f(init, fst), rst);
       }
     }));
 }

diff --git a/test/squint/compiler_test.cljs b/test/squint/compiler_test.cljs
@@ -1019,7 +1019,12 @@
             (jsv! '(vec (take 3 (as-> (repeat {:height 50}) posts
                                   (map #(assoc %1 :offset %2)
                                        posts
-                                       (reductions + 0 (map :height posts)))))))))))
+                                       (reductions + 0 (map :height posts))))))))))
+  (testing "lazy-reusage"
+    (is (= 100 (jsv! '(do (def a (atom []))
+                          (defn spy [x] (swap! a conj x) x)
+                          (vec (reductions + (map spy (range 100))))
+                          (count @a)))))))
 
 (deftest seq-test
   (is (= "abc" (jsv! '(seq "abc"))))