diff --git a/addbatch_test.go b/addbatch_test.go index 431c3bd..afaf99c 100644 --- a/addbatch_test.go +++ b/addbatch_test.go @@ -1,6 +1,7 @@ package arbo import ( + "bytes" "crypto/rand" "encoding/hex" "fmt" @@ -1091,6 +1092,118 @@ func TestAddBatchThresholdInDisk(t *testing.T) { c.Check(len(invalids), qt.Equals, 0) // check that both trees roots are equal checkRoots(c, tree1, tree3) + + // now add one leaf more to the trees to ensure that the previous + // actions did not left the tree in an invalid state + k := randomBytes(32) + v := randomBytes(32) + err = tree1.Add(k, v) + c.Assert(err, qt.IsNil) + err = tree2.Add(k, v) + c.Assert(err, qt.IsNil) + err = tree3.Add(k, v) + c.Assert(err, qt.IsNil) +} + +func initTestUpFromSubRoots(c *qt.C) (*Tree, *Tree) { + database1, err := badgerdb.New(db.Options{Path: c.TempDir()}) + c.Assert(err, qt.IsNil) + tree1, err := NewTree(Config{database1, 256, DefaultThresholdNLeafs, + HashFunctionBlake2b}) + c.Assert(err, qt.IsNil) + + database2, err := badgerdb.New(db.Options{Path: c.TempDir()}) + c.Assert(err, qt.IsNil) + tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs, + HashFunctionBlake2b}) + c.Assert(err, qt.IsNil) + return tree1, tree2 +} + +func testUpFromSubRoots(c *qt.C, tree1, tree2 *Tree, preSubRoots [][]byte) { + // add the preSubRoots to the tree1 + for i := 0; i < len(preSubRoots); i++ { + if bytes.Equal(preSubRoots[i], tree1.emptyHash) { + continue + } + err := tree1.Add(preSubRoots[i], nil) + c.Assert(err, qt.IsNil) + } + root1, err := tree1.Root() + c.Assert(err, qt.IsNil) + + wTx := tree2.db.WriteTx() + subRoots := make([][]byte, len(preSubRoots)) + for i := 0; i < len(preSubRoots); i++ { + if preSubRoots[i] == nil || bytes.Equal(preSubRoots[i], tree1.emptyHash) { + subRoots[i] = tree1.emptyHash + continue + } + leafKey, leafValue, err := tree2.newLeafValue(preSubRoots[i], nil) + c.Assert(err, qt.IsNil) + subRoots[i] = leafKey + + err = wTx.Set(leafKey, leafValue) + c.Assert(err, qt.IsNil) + } + // first fill the leaf nodes + // then call upFromSubRoots + root2FromUp, err := tree2.upFromSubRoots(wTx, subRoots) + c.Assert(err, qt.IsNil) + + err = tree2.SetRootWithTx(wTx, root2FromUp) + c.Assert(err, qt.IsNil) + err = wTx.Commit() + c.Assert(err, qt.IsNil) + + root2, err := tree2.Root() + c.Assert(err, qt.IsNil) + + c.Assert(root1, qt.DeepEquals, root2) +} + +func testUpFromSubRootsWithEmpties(c *qt.C, preSubRoots [][]byte, indexEmpties []int) { + tree1, tree2 := initTestUpFromSubRoots(c) + defer tree1.db.Close() //nolint:errcheck + defer tree2.db.Close() //nolint:errcheck + + testPreSubRoots := make([][]byte, len(preSubRoots)) + copy(testPreSubRoots[:], preSubRoots[:]) + for i := 0; i < len(indexEmpties); i++ { + testPreSubRoots[indexEmpties[i]] = tree1.emptyHash + } + testUpFromSubRoots(c, tree1, tree2, testPreSubRoots) +} + +func TestUpFromSubRoots(t *testing.T) { + c := qt.New(t) + + // prepare preSubRoots + preSubRoots := [][]byte{ + BigIntToBytes(32, big.NewInt(4)), + BigIntToBytes(32, big.NewInt(2)), + BigIntToBytes(32, big.NewInt(1)), + BigIntToBytes(32, big.NewInt(3)), + } + + // test using the full 4 leafs as subRoots + testUpFromSubRootsWithEmpties(c, preSubRoots, nil) + // 1st subRoot empty + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0}) + // 2nd subRoot empty + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{1}) + // 3rd subRoot empty + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{2}) + // 4th subRoot empty + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{3}) + + // other combinations of empty SubRoots + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0, 1, 2, 3}) + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0, 1}) + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{1, 2}) + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{1, 3}) + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{2, 3}) + testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0, 2, 3}) } // TODO test adding batch with multiple invalid keys diff --git a/tree.go b/tree.go index 20216c0..5836f68 100644 --- a/tree.go +++ b/tree.go @@ -376,24 +376,38 @@ func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) if len(subRoots) == 1 { return subRoots[0], nil } + // get the subRoots values to know the node types of each subRoot + nodeTypes := make([]byte, len(subRoots)) + for i := 0; i < len(subRoots); i++ { + if bytes.Equal(subRoots[i], t.emptyHash) { + nodeTypes[i] = PrefixValueEmpty + continue + } + v, err := wTx.Get(subRoots[i]) + if err != nil { + return nil, err + } + nodeTypes[i] = v[0] + } var newSubRoots [][]byte - // TODO store the nodes key-values into the wTx for i := 0; i < len(subRoots); i += 2 { - if bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash) { - // TODO: || (ns[i].typ() == vtLeaf && ns[i+1].typ() == vtEmpty) { - + if (bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash)) || + (nodeTypes[i] == PrefixValueLeaf && bytes.Equal(subRoots[i+1], t.emptyHash)) { // when both sub nodes are empty, the parent is also empty // or - // (TODO WIP) when 1st sub node is a leaf but the 2nd is empty, the - // leaf is used as parent + // when 1st sub node is a leaf but the 2nd is empty, the + // leaf is used as 'parent' newSubRoots = append(newSubRoots, subRoots[i]) continue } - // TODO if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtLeaf { - // when 2nd sub node is a leaf but the 1st is empty, the - // leaf is used as 'parent' + if bytes.Equal(subRoots[i], t.emptyHash) && nodeTypes[i+1] == PrefixValueLeaf { + // when 2nd sub node is a leaf but the 1st is empty, + // the leaf is used as 'parent' + newSubRoots = append(newSubRoots, subRoots[i+1]) + continue + } k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1]) if err != nil { @@ -408,6 +422,7 @@ func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) return t.upFromSubRoots(wTx, newSubRoots) } + func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) { // go at level l and return each node key, where each node key is the // subRoot of the subTree that starts there diff --git a/vt.go b/vt.go index 1285c66..8d4728d 100644 --- a/vt.go +++ b/vt.go @@ -271,7 +271,7 @@ func upFromNodes(ns []*node) (*node, error) { // when both sub nodes are empty, the parent is also empty // or // when 1st sub node is a leaf but the 2nd is empty, the - // leaf is used as parent + // leaf is used as 'parent' res = append(res, ns[i]) continue }