From 03e9ab5fa7a3b85656b79e0520ec1c9da456af27 Mon Sep 17 00:00:00 2001
From: Andrei Borzenkov <root@sandwitch.dev>
Date: Fri, 11 Oct 2024 11:41:53 +0400
Subject: [PATCH 1/2] Reuse iteration code and parametrize it with map walking
 function

---
 src/PersistentOrderedMap.mo | 37 +++++++++++++++++++++++--------------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/PersistentOrderedMap.mo b/src/PersistentOrderedMap.mo
index 18506243..ecf51626 100644
--- a/src/PersistentOrderedMap.mo
+++ b/src/PersistentOrderedMap.mo
@@ -324,30 +324,39 @@ module {
   ///
   /// Note: Full map iteration creates `O(n)` temporary objects that will be collected as garbage.
   public func iter<K, V>(rbMap : Map<K, V>, direction : Direction) : I.Iter<(K, V)> {
-    object {
-      var trees : IterRep<K, V> = ?(#tr(rbMap), null);
-      public func next() : ?(K, V) {
-        switch (direction, trees) {
-          case (_, null) { null };
-          case (_, ?(#tr(#leaf), ts)) {
+    let turnLeftFirst : MapTraverser<K, V>
+     = func (l, xy, r, ts) { ?(#tr(l), ?(#xy(xy), ?(#tr(r), ts))) };
+
+    let turnRightFirst : MapTraverser<K, V>
+     = func (l, xy, r, ts) { ?(#tr(r), ?(#xy(xy), ?(#tr(l), ts))) };
+
+    switch direction {
+      case (#fwd) IterMap(rbMap, turnLeftFirst);
+      case (#bwd) IterMap(rbMap, turnRightFirst)
+    }
+  };
+
+  type MapTraverser<K, V> = (Map<K, V>, (K, V), Map<K, V>, IterRep<K, V>) -> IterRep<K, V>;
+
+  class IterMap<K, V>(rbMap : Map<K, V>, mapTraverser : MapTraverser<K, V>) {
+    var trees : IterRep<K, V> = ?(#tr(rbMap), null);
+    public func next() : ?(K, V) {
+        switch (trees) {
+          case (null) { null };
+          case (?(#tr(#leaf), ts)) {
             trees := ts;
             next()
           };
-          case (_, ?(#xy(xy), ts)) {
+          case (?(#xy(xy), ts)) {
             trees := ts;
             ?xy
-          }; // TODO: Let's float-out case on direction
-          case (#fwd, ?(#tr(#node(_, l, xy, r)), ts)) {
-            trees := ?(#tr(l), ?(#xy(xy), ?(#tr(r), ts)));
-            next()
           };
-          case (#bwd, ?(#tr(#node(_, l, xy, r)), ts)) {
-            trees := ?(#tr(r), ?(#xy(xy), ?(#tr(l), ts)));
+          case (?(#tr(#node(_, l, xy, r)), ts)) {
+            trees := mapTraverser(l, xy, r, ts);
             next()
           }
         }
       }
-    }
   };
 
   /// Returns an Iterator (`Iter`) over the key-value pairs in the map.

From 11613a584983088aa6da1453830fdd05f937c6e8 Mon Sep 17 00:00:00 2001
From: Andrei Borzenkov <root@sandwitch.dev>
Date: Wed, 9 Oct 2024 12:53:18 +0400
Subject: [PATCH 2/2] Use direct recursion in folding functions

---
 src/PersistentOrderedMap.mo | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/PersistentOrderedMap.mo b/src/PersistentOrderedMap.mo
index ecf51626..f93e4a61 100644
--- a/src/PersistentOrderedMap.mo
+++ b/src/PersistentOrderedMap.mo
@@ -536,11 +536,14 @@ module {
     combine : (Key, Value, Accum) -> Accum
   ) : Accum
   {
-    var acc = base;
-    for(val in iter(rbMap, #fwd)){
-      acc := combine(val.0, val.1, acc);
-    };
-    acc
+    switch (rbMap) {
+      case (#leaf) { base };
+      case (#node(_, l, (k, v), r)) {
+        let left = foldLeft(l, base, combine);
+        let middle = combine(k, v, left);
+        foldLeft(r, middle, combine)
+      }
+    }
   };
 
   /// Collapses the elements in `rbMap` into a single value by starting with `base`
@@ -576,11 +579,14 @@ module {
     combine : (Key, Value, Accum) -> Accum
   ) : Accum
   {
-    var acc = base;
-    for(val in iter(rbMap, #bwd)){
-      acc := combine(val.0, val.1, acc);
-    };
-    acc
+    switch (rbMap) {
+      case (#leaf) { base };
+      case (#node(_, l, (k, v), r)) {
+        let right = foldRight(r, base, combine);
+        let middle = combine(k, v, right);
+        foldRight(l, middle, combine)
+      }
+    }
   };