From 70bf6c0ee168bd98b0db56eb0cb51f786f2d9a52 Mon Sep 17 00:00:00 2001 From: Andrei Borzenkov Date: Tue, 8 Oct 2024 15:38:30 +0400 Subject: [PATCH 1/4] Move case on direction to the outside of the main iteration function --- src/PersistentOrderedMap.mo | 64 +++++++++++++++++++++++++------------ 1 file changed, 43 insertions(+), 21 deletions(-) diff --git a/src/PersistentOrderedMap.mo b/src/PersistentOrderedMap.mo index 18506243..061437dd 100644 --- a/src/PersistentOrderedMap.mo +++ b/src/PersistentOrderedMap.mo @@ -324,27 +324,49 @@ module { /// /// Note: Full map iteration creates `O(n)` temporary objects that will be collected as garbage. public func iter(rbMap : Map, direction : Direction) : I.Iter<(K, V)> { - object { - var trees : IterRep = ?(#tr(rbMap), null); - public func next() : ?(K, V) { - switch (direction, trees) { - case (_, null) { null }; - case (_, ?(#tr(#leaf), ts)) { - trees := ts; - next() - }; - 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))); - next() - } + switch direction { + case (#fwd) { iterForward(rbMap) }; + case (#bwd) { iterBackward(rbMap) } + } + }; + + func iterForward(rbMap : Map) : I.Iter<(K, V)> = object { + var trees : IterRep = ?(#tr(rbMap), null); + public func next() : ?(K, V) { + switch (trees) { + case (null) { null }; + case (?(#tr(#leaf), ts)) { + trees := ts; + next() + }; + case (?(#xy(xy), ts)) { + trees := ts; + ?xy + }; + case (?(#tr(#node(_, l, xy, r)), ts)) { + trees := ?(#tr(l), ?(#xy(xy), ?(#tr(r), ts))); + next() + } + } + } + }; + + func iterBackward(rbMap : Map) : I.Iter<(K, V)> = object { + var trees : IterRep = ?(#tr(rbMap), null); + public func next() : ?(K, V) { + switch (trees) { + case (null) { null }; + case (?(#tr(#leaf), ts)) { + trees := ts; + next() + }; + case (?(#xy(xy), ts)) { + trees := ts; + ?xy + }; + case (?(#tr(#node(_, l, xy, r)), ts)) { + trees := ?(#tr(r), ?(#xy(xy), ?(#tr(l), ts))); + next() } } } From a10a9d63b9dce21e09b75e52ec8154961b00bac1 Mon Sep 17 00:00:00 2001 From: Andrei Borzenkov Date: Wed, 9 Oct 2024 12:53:18 +0400 Subject: [PATCH 2/4] 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 061437dd..407c8cf8 100644 --- a/src/PersistentOrderedMap.mo +++ b/src/PersistentOrderedMap.mo @@ -549,11 +549,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` @@ -589,11 +592,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) + } + } }; From 1488ef2d58eeac70e84f308164b1013612a83e31 Mon Sep 17 00:00:00 2001 From: Andrei Borzenkov Date: Wed, 9 Oct 2024 16:26:50 +0400 Subject: [PATCH 3/4] Implement mapFilter using foldLeft instad of iter --- src/PersistentOrderedMap.mo | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/PersistentOrderedMap.mo b/src/PersistentOrderedMap.mo index 407c8cf8..3d59e13a 100644 --- a/src/PersistentOrderedMap.mo +++ b/src/PersistentOrderedMap.mo @@ -615,19 +615,15 @@ module { }; public func mapFilter(t : Map, compare : (K, K) -> O.Order, f : (K, V1) -> ?V2) : Map{ - var map = #leaf : Map; - for(kv in iter(t, #fwd)) - { - switch(f kv){ - case null {}; - case (?v1) { - // The keys still are monotonic, so we can - // merge trees using `append` and avoid compare here - map := put(map, compare, kv.0, v1); + func combine(key : K, value1 : V1, acc : Map) : Map { + switch (f(key, value1)){ + case null { acc }; + case (?value2) { + put(acc, compare, key, value2) } } }; - map + foldLeft(t, #leaf, combine) }; public func get(t : Map, compare : (K, K) -> O.Order, x : K) : ?V { From 16e4a2a49894475cb0adebe9603bb8fef2b44ee4 Mon Sep 17 00:00:00 2001 From: Andrei Borzenkov Date: Thu, 10 Oct 2024 14:24:49 +0400 Subject: [PATCH 4/4] Remove tuple from Map definition --- src/PersistentOrderedMap.mo | 210 +++++++++++++++--------------- test/PersistentOrderedMap.test.mo | 6 +- 2 files changed, 108 insertions(+), 108 deletions(-) diff --git a/src/PersistentOrderedMap.mo b/src/PersistentOrderedMap.mo index 3d59e13a..d02258db 100644 --- a/src/PersistentOrderedMap.mo +++ b/src/PersistentOrderedMap.mo @@ -45,7 +45,7 @@ module { /// The keys have the generic type `K` and the values the generic type `V`. /// Leaves are considered implicitly black. public type Map = { - #node : (Color, Map, (K, V), Map); + #node : (Color, Map, K, V, Map); #leaf }; @@ -343,8 +343,8 @@ module { trees := ts; ?xy }; - case (?(#tr(#node(_, l, xy, r)), ts)) { - trees := ?(#tr(l), ?(#xy(xy), ?(#tr(r), ts))); + case (?(#tr(#node(_, l, x, y, r)), ts)) { + trees := ?(#tr(l), ?(#xy(x,y), ?(#tr(r), ts))); next() } } @@ -364,8 +364,8 @@ module { trees := ts; ?xy }; - case (?(#tr(#node(_, l, xy, r)), ts)) { - trees := ?(#tr(r), ?(#xy(xy), ?(#tr(l), ts))); + case (?(#tr(#node(_, l, x, y, r)), ts)) { + trees := ?(#tr(r), ?(#xy(x, y), ?(#tr(l), ts))); next() } } @@ -480,8 +480,8 @@ module { func mapRec(m : Map) : Map { switch m { case (#leaf) { #leaf }; - case (#node(c, l, xy, r)) { - #node(c, mapRec l, (xy.0, f xy), mapRec r) // TODO: try destination-passing style to avoid non tail-call recursion + case (#node(c, l, x, y, r)) { + #node(c, mapRec l, x, f(x, y), mapRec r) // TODO: try destination-passing style to avoid non tail-call recursion }; } }; @@ -510,7 +510,7 @@ module { public func size(t : Map) : Nat { switch t { case (#leaf) { 0 }; - case (#node(_, l, _, r)) { + case (#node(_, l, _, _, r)) { size(l) + size(r) + 1 } } @@ -551,7 +551,7 @@ module { { switch (rbMap) { case (#leaf) { base }; - case (#node(_, l, (k, v), r)) { + case (#node(_, l, k, v, r)) { let left = foldLeft(l, base, combine); let middle = combine(k, v, left); foldLeft(r, middle, combine) @@ -594,7 +594,7 @@ module { { switch (rbMap) { case (#leaf) { base }; - case (#node(_, l, (k, v), r)) { + case (#node(_, l, k, v, r)) { let right = foldRight(r, base, combine); let middle = combine(k, v, right); foldRight(l, middle, combine) @@ -626,14 +626,14 @@ module { foldLeft(t, #leaf, combine) }; - public func get(t : Map, compare : (K, K) -> O.Order, x : K) : ?V { + public func get(t : Map, compare : (K, K) -> O.Order, x0 : K) : ?V { switch t { case (#leaf) { null }; - case (#node(_c, l, xy, r)) { - switch (compare(x, xy.0)) { - case (#less) { get(l, compare, x) }; - case (#equal) { ?xy.1 }; - case (#greater) { get(r, compare, x) } + case (#node(_c, l, x1, y, r)) { + switch (compare(x0, x1)) { + case (#less) { get(l, compare, x0) }; + case (#equal) { ?y }; + case (#greater) { get(r, compare, x0) } } } } @@ -641,8 +641,8 @@ module { func redden(t : Map) : Map { switch t { - case (#node (#B, l, xy, r)) { - (#node (#R, l, xy, r)) + case (#node (#B, l, x, y, r)) { + (#node (#R, l, x, y, r)) }; case _ { Debug.trap "RBTree.red" @@ -650,46 +650,46 @@ module { } }; - func lbalance(left : Map, xy : (K,V), right : Map) : Map { + func lbalance(left : Map, x0 : K, y0 : V, right : Map) : Map { switch (left, right) { - case (#node(#R, #node(#R, l1, xy1, r1), xy2, r2), r) { + case (#node(#R, #node(#R, l1, x1, y1, r1), x2, y2, r2), r) { #node( #R, - #node(#B, l1, xy1, r1), - xy2, - #node(#B, r2, xy, r)) + #node(#B, l1, x1, y1, r1), + x2, y2, + #node(#B, r2, x0, y0, r)) }; - case (#node(#R, l1, xy1, #node(#R, l2, xy2, r2)), r) { + case (#node(#R, l1, x1, y1, #node(#R, l2, x2, y2, r2)), r) { #node( #R, - #node(#B, l1, xy1, l2), - xy2, - #node(#B, r2, xy, r)) + #node(#B, l1, x1, y1, l2), + x2, y2, + #node(#B, r2, x0, y0, r)) }; case _ { - #node(#B, left, xy, right) + #node(#B, left, x0, y0, right) } } }; - func rbalance(left : Map, xy : (K,V), right : Map) : Map { + func rbalance(left : Map, x0 : K, y0 : V, right : Map) : Map { switch (left, right) { - case (l, #node(#R, l1, xy1, #node(#R, l2, xy2, r2))) { + case (l, #node(#R, l1, x1, y1, #node(#R, l2, x2, y2, r2))) { #node( #R, - #node(#B, l, xy, l1), - xy1, - #node(#B, l2, xy2, r2)) + #node(#B, l, x0, y0, l1), + x1, y1, + #node(#B, l2, x2, y2, r2)) }; - case (l, #node(#R, #node(#R, l1, xy1, r1), xy2, r2)) { + case (l, #node(#R, #node(#R, l1, x1, y1, r1), x2, y2, r2)) { #node( #R, - #node(#B, l, xy, l1), - xy1, - #node(#B, r1, xy2, r2)) + #node(#B, l, x0, y0, l1), + x1, y1, + #node(#B, r1, x2, y2, r2)) }; case _ { - #node(#B, left, xy, right) + #node(#B, left, x0, y0, right) }; } }; @@ -707,41 +707,41 @@ module { func ins(tree : Map) : Map { switch tree { case (#leaf) { - #node(#R, #leaf, (key,val), #leaf) + #node(#R, #leaf, key,val, #leaf) }; - case (#node(#B, left, xy, right)) { - switch (compare (key, xy.0)) { + case (#node(#B, left, x, y, right)) { + switch (compare (key, x)) { case (#less) { - lbalance(ins left, xy, right) + lbalance(ins left, x, y, right) }; case (#greater) { - rbalance(left, xy, ins right) + rbalance(left, x, y, ins right) }; case (#equal) { - let newVal = onClash({ new = val; old = xy.1 }); - #node(#B, left, (key,newVal), right) + let newVal = onClash({ new = val; old = y }); + #node(#B, left, key, newVal, right) } } }; - case (#node(#R, left, xy, right)) { - switch (compare (key, xy.0)) { + case (#node(#R, left, x, y, right)) { + switch (compare (key, x)) { case (#less) { - #node(#R, ins left, xy, right) + #node(#R, ins left, x, y, right) }; case (#greater) { - #node(#R, left, xy, ins right) + #node(#R, left, x, y, ins right) }; case (#equal) { - let newVal = onClash { new = val; old = xy.1 }; - #node(#R, left, (key,newVal), right) + let newVal = onClash { new = val; old = y }; + #node(#R, left, key, newVal, right) } } } }; }; switch (ins m) { - case (#node(#R, left, xy, right)) { - #node(#B, left, xy, right); + case (#node(#R, left, x, y, right)) { + #node(#B, left, x, y, right); }; case other { other }; }; @@ -772,44 +772,44 @@ module { ) : Map = replace(m, compare, key, val).0; - func balLeft(left : Map, xy : (K,V), right : Map) : Map { + func balLeft(left : Map, x0 : K, y0 : V, right : Map) : Map { switch (left, right) { - case (#node(#R, l1, xy1, r1), r) { + case (#node(#R, l1, x1, y1, r1), r) { #node( #R, - #node(#B, l1, xy1, r1), - xy, + #node(#B, l1, x1, y1, r1), + x0, y0, r) }; - case (_, #node(#B, l2, xy2, r2)) { - rbalance(left, xy, #node(#R, l2, xy2, r2)) + case (_, #node(#B, l2, x2, y2, r2)) { + rbalance(left, x0, y0, #node(#R, l2, x2, y2, r2)) }; - case (_, #node(#R, #node(#B, l2, xy2, r2), xy3, r3)) { + case (_, #node(#R, #node(#B, l2, x2, y2, r2), x3, y3, r3)) { #node(#R, - #node(#B, left, xy, l2), - xy2, - rbalance(r2, xy3, redden r3)) + #node(#B, left, x0, y0, l2), + x2, y2, + rbalance(r2, x3, y3, redden r3)) }; case _ { Debug.trap "balLeft" }; } }; - func balRight(left : Map, xy : (K,V), right : Map) : Map { + func balRight(left : Map, x0 : K, y0 : V, right : Map) : Map { switch (left, right) { - case (l, #node(#R, l1, xy1, r1)) { + case (l, #node(#R, l1, x1, y1, r1)) { #node(#R, l, - xy, - #node(#B, l1, xy1, r1)) + x0, y0, + #node(#B, l1, x1, y1, r1)) }; - case (#node(#B, l1, xy1, r1), r) { - lbalance(#node(#R, l1, xy1, r1), xy, r); + case (#node(#B, l1, x1, y1, r1), r) { + lbalance(#node(#R, l1, x1, y1, r1), x0, y0, r); }; - case (#node(#R, l1, xy1, #node(#B, l2, xy2, r2)), r3) { + case (#node(#R, l1, x1, y1, #node(#B, l2, x2, y2, r2)), r3) { #node(#R, - lbalance(redden l1, xy1, l2), - xy2, - #node(#B, r2, xy, r3)) + lbalance(redden l1, x1, y1, l2), + x2, y2, + #node(#B, r2, x0, y0, r3)) }; case _ { Debug.trap "balRight" }; } @@ -819,40 +819,40 @@ module { switch (left, right) { case (#leaf, _) { right }; case (_, #leaf) { left }; - case (#node (#R, l1, xy1, r1), - #node (#R, l2, xy2, r2)) { + case (#node (#R, l1, x1, y1, r1), + #node (#R, l2, x2, y2, r2)) { switch (append (r1, l2)) { - case (#node (#R, l3, xy3, r3)) { + case (#node (#R, l3, x3, y3, r3)) { #node( #R, - #node(#R, l1, xy1, l3), - xy3, - #node(#R, r3, xy2, r2)) + #node(#R, l1, x1, y1, l3), + x3, y3, + #node(#R, r3, x2, y2, r2)) }; case r1l2 { - #node(#R, l1, xy1, #node(#R, r1l2, xy2, r2)) + #node(#R, l1, x1, y1, #node(#R, r1l2, x2, y2, r2)) } } }; - case (t1, #node(#R, l2, xy2, r2)) { - #node(#R, append(t1, l2), xy2, r2) + case (t1, #node(#R, l2, x2, y2, r2)) { + #node(#R, append(t1, l2), x2, y2, r2) }; - case (#node(#R, l1, xy1, r1), t2) { - #node(#R, l1, xy1, append(r1, t2)) + case (#node(#R, l1, x1, y1, r1), t2) { + #node(#R, l1, x1, y1, append(r1, t2)) }; - case (#node(#B, l1, xy1, r1), #node (#B, l2, xy2, r2)) { + case (#node(#B, l1, x1, y1, r1), #node (#B, l2, x2, y2, r2)) { switch (append (r1, l2)) { - case (#node (#R, l3, xy3, r3)) { + case (#node (#R, l3, x3, y3, r3)) { #node(#R, - #node(#B, l1, xy1, l3), - xy3, - #node(#B, r3, xy2, r2)) + #node(#B, l1, x1, y1, l3), + x3, y3, + #node(#B, r3, x2, y2, r2)) }; case r1l2 { balLeft ( l1, - xy1, - #node(#B, r1l2, xy2, r2) + x1, y1, + #node(#B, r1l2, x2, y2, r2) ) } } @@ -863,34 +863,34 @@ module { public func delete(m : Map, compare : (K, K) -> O.Order, key : K) : Map = remove(m, compare, key).0; - public func remove(tree : Map, compare : (K, K) -> O.Order, x : K) : (Map, ?V) { + public func remove(tree : Map, compare : (K, K) -> O.Order, x0 : K) : (Map, ?V) { var y0 : ?V = null; - func delNode(left : Map, xy : (K, V), right : Map) : Map { - switch (compare (x, xy.0)) { + func delNode(left : Map, x1 : K, y1 : V, right : Map) : Map { + switch (compare (x0, x1)) { case (#less) { let newLeft = del left; switch left { - case (#node(#B, _, _, _)) { - balLeft(newLeft, xy, right) + case (#node(#B, _, _, _, _)) { + balLeft(newLeft, x1, y1, right) }; case _ { - #node(#R, newLeft, xy, right) + #node(#R, newLeft, x1, y1, right) } } }; case (#greater) { let newRight = del right; switch right { - case (#node(#B, _, _, _)) { - balRight(left, xy, newRight) + case (#node(#B, _, _, _, _)) { + balRight(left, x1, y1, newRight) }; case _ { - #node(#R, left, xy, newRight) + #node(#R, left, x1, y1, newRight) } } }; case (#equal) { - y0 := ?xy.1; + y0 := ?y1; append(left, right) }; } @@ -900,14 +900,14 @@ module { case (#leaf) { tree }; - case (#node(_, left, xy, right)) { - delNode(left, xy, right) + case (#node(_, left, x, y, right)) { + delNode(left, x, y, right) } }; }; switch (del(tree)) { - case (#node(#R, left, xy, right)) { - (#node(#B, left, xy, right), y0); + case (#node(#R, left, x, y, right)) { + (#node(#B, left, x, y, right), y0); }; case other { (other, y0) }; }; diff --git a/test/PersistentOrderedMap.test.mo b/test/PersistentOrderedMap.test.mo index 58e7b41d..43263f58 100644 --- a/test/PersistentOrderedMap.test.mo +++ b/test/PersistentOrderedMap.test.mo @@ -33,7 +33,7 @@ func checkMap(rbMap : Map.Map) { func blackDepth(node : Map.Map) : Nat { switch node { case (#leaf) 0; - case (#node(color, left, (key, _), right)) { + case (#node(color, left, key, _, right)) { checkKey(left, func(x) { x < key }); checkKey(right, func(x) { x > key }); let leftBlacks = blackDepth(left); @@ -57,14 +57,14 @@ func blackDepth(node : Map.Map) : Nat { func isRed(node : Map.Map) : Bool { switch node { case (#leaf) false; - case (#node(color, _, _, _)) color == #R + case (#node(color, _, _, _, _)) color == #R } }; func checkKey(node : Map.Map, isValid : Nat -> Bool) { switch node { case (#leaf) {}; - case (#node(_, _, (key, _), _)) { + case (#node(_, _, key, _, _)) { assert (isValid(key)) } }