Make the BTree size aware, and add BTree.size() function

canscale · Feb 4, 2023 · dc07d30 · dc07d30
1 parent cc6367a
commit dc07d30
Show file tree

Hide file tree

Showing 6 changed files with 85 additions and 8 deletions.
diff --git a/src/BTree.mo b/src/BTree.mo
@@ -35,10 +35,10 @@ module {
           var count = 0;
         };
       }); 
+      var size = 0;
       order = btreeOrder;
     }
   };
-
 
   /// Allows one to quickly create a BTree using an array of key value pairs
   public func createBTreeWithKVPairs<K, V>(order: Nat, compare: (K, K) -> O.Order, kvPairs: [(K, V)]): BTree<K, V> {
@@ -49,6 +49,9 @@ module {
     t;
   };
 
+  /// Get the current count of key-value pairs present in the BTree
+  public func size<K, V>(tree: BTree<K, V>): Nat { tree.size };
+
 
   /// Retrieves the value corresponding to the key of BTree if it exists
   public func get<K, V>(tree: BTree<K, V>, compare: (K, K) -> O.Order, key: K): ?V {
@@ -67,7 +70,14 @@ module {
     };
 
     switch(insertResult) {
-      case (#insert(ov)) { ov };
+      case (#insert(ov)) {
+        switch(ov) {
+          // if inserted a value that was not previously there, increment the tree size counter
+          case null { tree.size += 1 };
+          case _ {};
+        };
+        ov
+      };
       case (#promote({ kv; leftChild; rightChild; })) {
         tree.root := #internal({
           data = {
@@ -83,6 +93,8 @@ module {
             else { null }
           });
         });
+        // promotion always comes from inserting a new element, so increment the tree size counter
+        tree.size += 1;
 
         null
       }
@@ -99,14 +111,14 @@ module {
           case (#keyFound(deleteIndex)) { 
             leafNode.data.count -= 1;
             let (_, deletedValue) = AU.deleteAndShiftValuesOver<(K, V)>(leafNode.data.kvs, deleteIndex);
+            tree.size -= 1;
             ?deletedValue
           };
           case _ { null }
         }
-
       };
       case (#internal(internalNode)) { 
-        switch(internalDeleteHelper(internalNode, tree.order, compare, key, false)) {
+        let deletedValueResult = switch(internalDeleteHelper(internalNode, tree.order, compare, key, false)) {
           case (#delete(value)) { value };
           case (#mergeChild({ internalChild; deletedValue })) {
             if (internalChild.data.count > 0) {
@@ -122,9 +134,15 @@ module {
             };
             deletedValue
           }
-        }
-      }
+        };
 
+        switch(deletedValueResult) {
+          // if deleted a value from the BTree, decrement the size
+          case (?deletedValue) { tree.size -= 1 };
+          case null {}
+        };
+        deletedValueResult
+      }
     }
   };
 
@@ -1176,6 +1194,7 @@ module {
   public func toText<K, V>(t: BTree<K, V>, keyToText: K -> Text, valueToText: V -> Text): Text {
     var textOutput = "BTree={";
     textOutput #= "root=" # rootToText<K, V>(t.root, keyToText, valueToText) # "; ";
+    textOutput #= "size=" # Nat.toText(t.size) # "; ";
     textOutput #= "order=" # Nat.toText(t.order) # "; ";
     textOutput # "}";
   };
@@ -1188,7 +1207,7 @@ module {
     keyEquals: (K, K) -> Bool,
     valueEquals: (V, V) -> Bool
   ): Bool {
-    if (t1.order != t2.order) return false;
+    if (t1.order != t2.order or t1.size != t2.size) return false;
 
     nodeEquals(t1.root, t2.root, keyEquals, valueEquals);
   };

diff --git a/src/Types.mo b/src/Types.mo
@@ -21,6 +21,7 @@ module {
 
   public type BTree<K, V> = {
     var root: Node<K, V>;
+    var size: Nat;
     order: Nat;
   };
 }
diff --git a/test/BTreeDeleteOrder4PropertyTest.mo b/test/BTreeDeleteOrder4PropertyTest.mo
@@ -49,6 +49,7 @@ let btreePropertyTests = S.suite("check validity of mass insertion/deletion", [
             var count = 0;
           }
         });
+        var size = 0;
         order = 4;
       }))
     ),
@@ -87,6 +88,7 @@ let btreePropertyTests = S.suite("check validity of mass insertion/deletion", [
             var count = 0;
           }
         });
+        var size = 0;
         order = 4;
       }))
     ),
@@ -121,6 +123,7 @@ let btreePropertyTests = S.suite("check validity of mass insertion/deletion", [
             var count = 0;
           }
         });
+        var size = 0;
         order = 4;
       }))
     ),

diff --git a/test/BTreeDeleteOrder6PropertyTest.mo b/test/BTreeDeleteOrder6PropertyTest.mo
@@ -49,6 +49,7 @@ let btreePropertyTests = S.suite("check validity of mass insertion/deletion", [
             var count = 0;
           }
         });
+        var size = 0;
         order = 6;
       }))
     ),
@@ -87,6 +88,7 @@ let btreePropertyTests = S.suite("check validity of mass insertion/deletion", [
             var count = 0;
           }
         });
+        var size = 0;
         order = 6;
       }))
     ),
@@ -121,6 +123,7 @@ let btreePropertyTests = S.suite("check validity of mass insertion/deletion", [
             var count = 0;
           }
         });
+        var size = 0;
         order = 6;
       }))
     ),