Optimise reduce.

backend-v2/runtime/PersistentList.c

@@ -250,14 +250,34 @@ RTValue PersistentList_reduce(PersistentList *self, RTValue f, RTValue start) {
   RTValue args[2];
 
   while (current && !RT_isNull(current->first)) {
-    args[0] = acc;
-    args[1] = current->first;
-    retain(args[1]);
-    acc = RT_invokeMethodWithFrame(frame, f, method, args, 2);
-    current = current->rest;
+    PersistentList *next = current->rest;
+    RTValue first = current->first;
+
+    if (Ptr_isReusable(current)) {
+      // Sole owner: steal children and destroy node immediately
+      current->first = RT_boxNull();
+      current->rest = NULL;
+      Ptr_release(current);
+
+      args[0] = acc;
+      args[1] = first;
+      // No retain(first) because we stole the reference from the destroyed node
+      acc = RT_invokeMethodWithFrame(frame, f, method, args, 2);
+    } else {
+      // Shared: standard retain/release
+      retain(first);
+      if (next) Ptr_retain(next);
+
+      args[0] = acc;
+      args[1] = first;
+      acc = RT_invokeMethodWithFrame(frame, f, method, args, 2);
+
+      Ptr_release(current);
+    }
+    current = next;
   }
 
-  Ptr_release(self);
+  if (current) Ptr_release(current);
   release(f);
   return acc;
 }
@@ -267,11 +287,23 @@ RTValue PersistentList_reduce2(PersistentList *self, RTValue f) {
     Ptr_release(self);
     return RT_invokeDynamic(f, NULL, 0);
   }
-  RTValue first = self->first;
-  retain(first);
-  PersistentList *rest = self->rest;
-  if (rest) Ptr_retain(rest);
-  Ptr_release(self);
+
+  RTValue first;
+  PersistentList *rest;
+
+  if (Ptr_isReusable(self)) {
+    first = self->first;
+    rest = self->rest;
+    self->first = RT_boxNull();
+    self->rest = NULL;
+    Ptr_release(self);
+  } else {
+    first = self->first;
+    retain(first);
+    rest = self->rest;
+    if (rest) Ptr_retain(rest);
+    Ptr_release(self);
+  }
 
   if (!rest) {
     release(f);

backend-v2/runtime/PersistentVector.c

@@ -106,12 +106,14 @@ String *PersistentVector_toString(PersistentVector *restrict self) {
 
 /* outside refcount system */
 void PersistentVector_destroy(PersistentVector *restrict self,
-                              bool deallocateChildren) {
-  Ptr_release(self->tail);
-  if (self->root)
-    Ptr_release(self->root);
+                               bool deallocateChildren) {
+  if (deallocateChildren) {
+    if (self->root)
+      Ptr_release(self->root);
+    if (self->tail)
+      Ptr_release(self->tail);
+  }
 }
-
 void PersistentVector_promoteToShared(PersistentVector *self, uword_t current) {
   if (current & SHARED_BIT)
     return;
@@ -686,7 +688,48 @@ bool PersistentVector_equals_managed(PersistentVector *self, RTValue other) {
   return res;
 }
 
-RTValue PersistentVector_reduce(PersistentVector *self, RTValue f, RTValue start) {
+
+static RTValue PersistentVectorNode_reduce(PersistentVectorNode *node,
+                                           uword_t level, RTValue acc,
+                                           Frame *frame, RTValue f,
+                                           FunctionMethod *method,
+                                           bool reusable) {
+  RTValue args[2];
+  if (node->type == leafNode) {
+    uword_t cnt = node->count;
+    for (uword_t i = 0; i < cnt; i++) {
+      args[0] = acc;
+      args[1] = node->array[i];
+      if (!reusable) {
+        retain(args[1]);
+      }
+      acc = RT_invokeMethodWithFrame(frame, f, method, args, 2);
+    }
+    if (reusable) {
+      node->count = 0;
+      Ptr_release(node);
+    }
+  } else {
+    uword_t cnt = node->count;
+    for (uword_t i = 0; i < cnt; i++) {
+      PersistentVectorNode *child =
+          (PersistentVectorNode *)RT_unboxPtr(node->array[i]);
+      bool childReusable = reusable && Ptr_isReusable(child);
+      acc = PersistentVectorNode_reduce(child, level - RRB_BITS, acc, frame, f,
+                                        method, childReusable);
+      if (reusable && childReusable) {
+        node->array[i] = RT_boxNull();
+      }
+    }
+    if (reusable) {
+      Ptr_release(node);
+    }
+  }
+  return acc;
+}
+
+RTValue PersistentVector_reduce(PersistentVector *self, RTValue f,
+                                RTValue start) {
   RTValue acc = start;
   FunctionMethod *method = Function_extractMethod(f, 2);
 
@@ -695,32 +738,37 @@ RTValue PersistentVector_reduce(PersistentVector *self, RTValue f, RTValue start
   frame->leafFrame = NULL;
   frame->bailoutEntryIndex = -1;
 
-  RTValue args[2];
-  uword_t totalCount = self->count;
-  uword_t tailCount = self->tail ? self->tail->count : 0;
-  uword_t tailStart = totalCount - tailCount;
-
-  // 1. Process Trie-resident blocks (32 elements at a time)
-  for (uword_t i = 0; i < tailStart; i += 32) {
-    PersistentVectorNode *node = PersistentVector_nthBlock(self, i);
-    // nthBlock for i < tailStart is guaranteed to be a leaf in the trie
-    for (uword_t j = 0; j < 32; j++) {
-      args[0] = acc;
-      args[1] = node->array[j];
-      retain(args[1]);
-      acc = RT_invokeMethodWithFrame(frame, f, method, args, 2);
+  bool reusable = Ptr_isReusable(self);
+  PersistentVectorNode *root = self->root;
+  PersistentVectorNode *tail = self->tail;
+  uword_t shift = self->shift;
+
+  if (root) {
+    bool rootReusable = reusable && Ptr_isReusable(root);
+    acc = PersistentVectorNode_reduce(root, shift, acc, frame, f, method,
+                                      rootReusable);
+    if (reusable && rootReusable) {
+      self->root = NULL;
     }
   }
 
-  // 2. Process the Tail
-  if (self->tail) {
-    PersistentVectorNode *tail = self->tail;
-    for (uword_t j = 0; j < tail->count; j++) {
+  if (tail) {
+    bool tailReusable = reusable && Ptr_isReusable(tail);
+    uword_t cnt = tail->count;
+    RTValue args[2];
+    for (uword_t i = 0; i < cnt; i++) {
       args[0] = acc;
-      args[1] = tail->array[j];
-      retain(args[1]);
+      args[1] = tail->array[i];
+      if (!tailReusable) {
+        retain(args[1]);
+      }
       acc = RT_invokeMethodWithFrame(frame, f, method, args, 2);
     }
+    if (tailReusable) {
+      tail->count = 0;
+      self->tail = NULL;
+      Ptr_release(tail);
+    }
   }
 
   Ptr_release(self);
@@ -734,8 +782,10 @@ RTValue PersistentVector_reduce2(PersistentVector *self, RTValue f) {
     return RT_invokeDynamic(f, NULL, 0);
   }
 
+  RTValue first;
   Ptr_retain(self);
-  RTValue first = PersistentVector_nth(self, 0);
+  first = PersistentVector_nth(self, 0);
+
   if (self->count == 1) {
     Ptr_release(self);
     release(f);
@@ -744,7 +794,16 @@ RTValue PersistentVector_reduce2(PersistentVector *self, RTValue f) {
 
   PersistentVectorIterator it = PersistentVector_iterator(self);
   it.index = 1;
+  if (it.blockIndex < it.block->count - 1) {
+    it.blockIndex++;
+  } else {
+    it.block = PersistentVector_nthBlock(self, 1);
+    it.blockIndex = 0;
+  }
+
   PersistentVectorChunkedSeq *seq = PersistentVectorChunkedSeq_create(it);
+  Ptr_release(self);
+
   return PersistentVectorChunkedSeq_reduce(seq, f, first);
 }
 

backend-v2/runtime/tests/PersistentList_test.c

@@ -4,6 +4,10 @@
 #include <stdio.h>
 
 #include "TestTools.h"
+#include "../BigInteger.h"
+#include "../Function.h"
+#include "../Integer.h"
+#include "../PersistentList.h"
 
 #define CONJ_TO_LOOP_FACTOR 100
 
@@ -129,12 +133,79 @@ static void testListPromotionStop(void **state) {
   });
 }
 
+static RTValue MockBigIntAddition(Frame *frame, RTValue a, RTValue b, RTValue a2,
+                                   RTValue a3, RTValue a4) {
+  BigInteger *valA = (BigInteger *)RT_unboxPtr(a);
+  BigInteger *valB = (BigInteger *)RT_unboxPtr(b);
+  // retain because BigInteger_add will consume them, 
+  // but RT_invokeMethodWithFrame will ALSO release them later.
+  Ptr_retain(valA);
+  Ptr_retain(valB);
+  BigInteger *res = BigInteger_add(valA, valB);
+  return RT_boxPtr(res);
+}
+
+static RTValue create_bigint_add_fn() {
+  ClojureFunction *f = Function_create(1, 2, false);
+  Function_fillMethod(f, 0, 0, 2, false, MockBigIntAddition, "add", 0);
+  return RT_boxPtr(f);
+}
+
+static RTValue MockAddition(Frame *frame, RTValue a, RTValue b, RTValue a2,
+                            RTValue a3, RTValue a4) {
+  int32_t valA = RT_unboxInt32(a);
+  int32_t valB = RT_unboxInt32(b);
+  release(a);
+  release(b);
+  return RT_boxInt32(valA + valB);
+}
+
+static RTValue create_mock_add_fn() {
+  ClojureFunction *f = Function_create(1, 2, false);
+  Function_fillMethod(f, 0, 0, 2, false, MockAddition, "add", 0);
+  return RT_boxPtr(f);
+}
+
+static void testListReduce(void **state) {
+  (void)state;
+  ASSERT_MEMORY_ALL_BALANCED({
+    PersistentList *l = PersistentList_empty();
+    int count = 1000;
+    for (int i = 0; i < count; i++) {
+      l = PersistentList_conj(l, RT_boxPtr(BigInteger_createFromInt(i)));
+    }
+
+    // 1. reduce with start value
+    Ptr_retain(l);
+    RTValue addFn = create_bigint_add_fn();
+    RTValue res = PersistentList_reduce(l, addFn, RT_boxPtr(BigInteger_createFromInt(0)));
+
+    BigInteger *expected = BigInteger_createFromInt((count * (count - 1)) / 2);
+    assert_true(BigInteger_equals((BigInteger *)RT_unboxPtr(res), expected));
+    release(res);
+    Ptr_release(expected);
+
+    // 2. reduce2 (without start value)
+    Ptr_retain(l);
+    RTValue addFn2 = create_bigint_add_fn();
+    RTValue res2 = PersistentList_reduce2(l, addFn2);
+
+    BigInteger *expected2 = BigInteger_createFromInt((count * (count - 1)) / 2);
+    assert_true(BigInteger_equals((BigInteger *)RT_unboxPtr(res2), expected2));
+    release(res2);
+    Ptr_release(expected2);
+
+    Ptr_release(l);
+  });
+}
+
 int main(void) {
   const struct CMUnitTest tests[] = {
       cmocka_unit_test_prestate(testScalingBehavior,
                                 listConjunctionAndPerformance),
       cmocka_unit_test(testListPromotion),
       cmocka_unit_test(testListPromotionStop),
+      cmocka_unit_test(testListReduce),
   };
   initialise_memory();
   RuntimeInterface_initialise();