M addbatch.go → addbatch.go

@@ -8,7 +8,6 @@ )
 
 /*
 
-
 AddBatch design
 ===============
 
@@ -18,7 +17,7 @@ =================================================
 - Build the full tree from bottom to top (from all the leaf to the root)
 
 
-CASE B: ALMOST CASE A, Almost empty Tree --> if Tree has numLeafs < numBuckets
+CASE B: ALMOST CASE A, Almost empty Tree --> if Tree has numLeafs < minLeafsThreshold
 ==============================================================================
 - Get the Leafs (key & value) (iterate the tree from the current root getting
 the leafs)
@@ -33,7 +32,7 @@ / \
       B   C
 
 
-CASE C: ALMOST CASE B --> if Tree has few Leafs (but numLeafs>=numBuckets)
+CASE C: ALMOST CASE B --> if Tree has few Leafs (but numLeafs>=minLeafsThreshold)
 ==============================================================================
 - Use A, B, G, F as Roots of subtrees
 - Do CASE B for each subtree
@@ -112,8 +111,8 @@
 Algorithm decision
 ==================
 - if nLeafs==0 (root==0): CASE A
-- if nLeafs<nBuckets: CASE B
-- if nLeafs>=nBuckets && nLeafs < minLeafsThreshold: CASE C
+- if nLeafs<minLeafsThreshold: CASE B
+- if nLeafs>=minLeafsThreshold && (nLeafs/nBuckets) < minLeafsThreshold: CASE C
 - else: CASE D & CASE E
 
 
@@ -148,14 +147,35 @@ return nil, err
 	}
 
 	// if nLeafs==0 (root==0): CASE A
-	e := make([]byte, t.hashFunction.Len())
-	if bytes.Equal(t.root, e) {
-		// CASE A
+	if bytes.Equal(t.root, t.emptyHash) {
 		// sort keys & values by path
 		sortKvs(kvs)
 		return t.buildTreeBottomUp(kvs)
 	}
 
+	// if nLeafs<nBuckets: CASE B
+	nLeafs, err := t.GetNLeafs()
+	if err != nil {
+		return nil, err
+	}
+	minLeafsThreshold := uint64(100) // nolint:gomnd // TMP WIP
+	if nLeafs < minLeafsThreshold {
+		// get already existing keys
+		aKs, aVs, err := t.getLeafs()
+		if err != nil {
+			return nil, err
+		}
+		aKvs, err := t.keysValuesToKvs(aKs, aVs)
+		if err != nil {
+			return nil, err
+		}
+		// add already existing key-values to the inputted key-values
+		kvs = append(kvs, aKvs...)
+		// proceed with CASE A
+		sortKvs(kvs)
+		return t.buildTreeBottomUp(kvs)
+	}
+
 	return nil, fmt.Errorf("UNIMPLEMENTED")
 }
 
@@ -209,6 +229,18 @@
 	return kvs, nil
 }
 
+/*
+func (t *Tree) kvsToKeysValues(kvs []kv) ([][]byte, [][]byte) {
+	ks := make([][]byte, len(kvs))
+	vs := make([][]byte, len(kvs))
+	for i := 0; i < len(kvs); i++ {
+		ks[i] = kvs[i].k
+		vs[i] = kvs[i].v
+	}
+	return ks, vs
+}
+*/
+
 // keys & values must be sorted by path, and must be length multiple of 2
 // TODO return index of failed keyvaules
 func (t *Tree) buildTreeBottomUp(kvs []kv) ([]int, error) {
@@ -254,3 +286,16 @@ rKs = append(rKs, k)
 	}
 	return t.upFromKeys(rKs)
 }
+
+func (t *Tree) getLeafs() ([][]byte, [][]byte, error) {
+	var ks, vs [][]byte
+	err := t.Iterate(func(k, v []byte) {
+		if v[0] != PrefixValueLeaf {
+			return
+		}
+		leafK, leafV := readLeafValue(v)
+		ks = append(ks, leafK)
+		vs = append(vs, leafV)
+	})
+	return ks, vs, err
+}

M addbatch_test.go → addbatch_test.go

@@ -49,3 +49,53 @@
 	// check that both trees roots are equal
 	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
 }
+
+func TestAddBatchCaseB(t *testing.T) {
+	c := qt.New(t)
+
+	nLeafs := 1024
+
+	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree.db.Close()
+
+	start := time.Now()
+	for i := 0; i < nLeafs; i++ {
+		k := BigIntToBytes(big.NewInt(int64(i)))
+		v := BigIntToBytes(big.NewInt(int64(i * 2)))
+		if err := tree.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+	fmt.Println(time.Since(start))
+
+	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close()
+
+	// add the initial leafs to fill a bit the tree before calling the
+	// AddBatch method
+	for i := 0; i < 99; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
+		k := BigIntToBytes(big.NewInt(int64(i)))
+		v := BigIntToBytes(big.NewInt(int64(i * 2)))
+		if err := tree2.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	var keys, values [][]byte
+	for i := 99; i < nLeafs; i++ {
+		k := BigIntToBytes(big.NewInt(int64(i)))
+		v := BigIntToBytes(big.NewInt(int64(i * 2)))
+		keys = append(keys, k)
+		values = append(values, v)
+	}
+	start = time.Now()
+	indexes, err := tree2.AddBatchOpt(keys, values)
+	c.Assert(err, qt.IsNil)
+	fmt.Println(time.Since(start))
+	c.Check(len(indexes), qt.Equals, 0)
+
+	// check that both trees roots are equal
+	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
+}