Files
arbo/tree_test.go
arnaucube 5f6c35e435 Update Snapshot & Root approach
Update Snapshot & Root approach to get the root always from the db,
except in the cases that the tree is a snapshot, in which the root will
be in memory.
In this way, when a snapshot is performed and the original tree gets
modifyed, the snapshot will still point to the old root. Also, the root
obtained from the db, uses also the db.ReadTx, so if the root is being
modifyied in the current tx (db.WriteTx), when getting the root it will
be return the lastest version that is in the tx but not yet in the db.
2021-08-13 12:38:05 +02:00

539 lines
16 KiB
Go

package arbo
import (
"encoding/hex"
"math/big"
"testing"
"time"
qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
)
func checkRootBIString(c *qt.C, tree *Tree, expected string) {
root, err := tree.Root()
c.Assert(err, qt.IsNil)
rootBI := BytesToBigInt(root)
c.Check(rootBI.String(), qt.Equals, expected)
}
func TestDBTx(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
wTx := database.WriteTx()
_, err = wTx.Get([]byte("a"))
c.Assert(err, qt.Equals, db.ErrKeyNotFound)
err = wTx.Set([]byte("a"), []byte("b"))
c.Assert(err, qt.IsNil)
v, err := wTx.Get([]byte("a"))
c.Assert(err, qt.IsNil)
c.Assert(v, qt.DeepEquals, []byte("b"))
}
func TestAddTestVectors(t *testing.T) {
c := qt.New(t)
// Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
testVectorsPoseidon := []string{
"0000000000000000000000000000000000000000000000000000000000000000",
"13578938674299138072471463694055224830892726234048532520316387704878000008795",
"5412393676474193513566895793055462193090331607895808993925969873307089394741",
"14204494359367183802864593755198662203838502594566452929175967972147978322084",
}
testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
testVectorsSha256 := []string{
"0000000000000000000000000000000000000000000000000000000000000000",
"46910109172468462938850740851377282682950237270676610513794735904325820156367",
"59481735341404520835410489183267411392292882901306595567679529387376287440550",
"20573794434149960984975763118181266662429997821552560184909083010514790081771",
}
testAdd(c, HashFunctionSha256, testVectorsSha256)
}
func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 10, hashFunc)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
root, err := tree.Root()
c.Assert(err, qt.IsNil)
c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0])
bLen := hashFunc.Len()
err = tree.Add(
BigIntToBytes(bLen, big.NewInt(1)),
BigIntToBytes(bLen, big.NewInt(2)))
c.Assert(err, qt.IsNil)
checkRootBIString(c, tree, testVectors[1])
err = tree.Add(
BigIntToBytes(bLen, big.NewInt(33)),
BigIntToBytes(bLen, big.NewInt(44)))
c.Assert(err, qt.IsNil)
checkRootBIString(c, tree, testVectors[2])
err = tree.Add(
BigIntToBytes(bLen, big.NewInt(1234)),
BigIntToBytes(bLen, big.NewInt(9876)))
c.Assert(err, qt.IsNil)
checkRootBIString(c, tree, testVectors[3])
}
func TestAddBatch(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
for i := 0; i < 1000; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(0))
if err := tree.Add(k, v); err != nil {
t.Fatal(err)
}
}
checkRootBIString(c, tree,
"296519252211642170490407814696803112091039265640052570497930797516015811235")
database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree2.db.Close() //nolint:errcheck
var keys, values [][]byte
for i := 0; i < 1000; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(0))
keys = append(keys, k)
values = append(values, v)
}
indexes, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
checkRootBIString(c, tree2,
"296519252211642170490407814696803112091039265640052570497930797516015811235")
}
func TestAddDifferentOrder(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
bLen := tree1.HashFunction().Len()
for i := 0; i < 16; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(0))
if err := tree1.Add(k, v); err != nil {
t.Fatal(err)
}
}
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree2.db.Close() //nolint:errcheck
for i := 16 - 1; i >= 0; i-- {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(0))
if err := tree2.Add(k, v); err != nil {
t.Fatal(err)
}
}
root1, err := tree1.Root()
c.Assert(err, qt.IsNil)
root2, err := tree2.Root()
c.Assert(err, qt.IsNil)
c.Check(hex.EncodeToString(root2), qt.Equals, hex.EncodeToString(root1))
c.Check(hex.EncodeToString(root1), qt.Equals,
"3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
}
func TestAddRepeatedIndex(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
k := BigIntToBytes(bLen, big.NewInt(int64(3)))
v := BigIntToBytes(bLen, big.NewInt(int64(12)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
err = tree.Add(k, v) // repeating same key-value
c.Check(err, qt.Equals, ErrKeyAlreadyExists)
}
func TestUpdate(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
k := BigIntToBytes(bLen, big.NewInt(int64(20)))
v := BigIntToBytes(bLen, big.NewInt(int64(12)))
if err := tree.Add(k, v); err != nil {
t.Fatal(err)
}
v = BigIntToBytes(bLen, big.NewInt(int64(11)))
err = tree.Update(k, v)
c.Assert(err, qt.IsNil)
gettedKey, gettedValue, err := tree.Get(k)
c.Assert(err, qt.IsNil)
c.Check(gettedKey, qt.DeepEquals, k)
c.Check(gettedValue, qt.DeepEquals, v)
// add more leafs to the tree to do another test
for i := 0; i < 16; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
if err := tree.Add(k, v); err != nil {
t.Fatal(err)
}
}
k = BigIntToBytes(bLen, big.NewInt(int64(3)))
v = BigIntToBytes(bLen, big.NewInt(int64(11)))
// check that before the Update, value for 3 is !=11
gettedKey, gettedValue, err = tree.Get(k)
c.Assert(err, qt.IsNil)
c.Check(gettedKey, qt.DeepEquals, k)
c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
err = tree.Update(k, v)
c.Assert(err, qt.IsNil)
// check that after Update, the value for 3 is ==11
gettedKey, gettedValue, err = tree.Get(k)
c.Assert(err, qt.IsNil)
c.Check(gettedKey, qt.DeepEquals, k)
c.Check(gettedValue, qt.DeepEquals, v)
c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
}
func TestAux(t *testing.T) { // TODO split in proper tests
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
k := BigIntToBytes(bLen, big.NewInt(int64(1)))
v := BigIntToBytes(bLen, big.NewInt(int64(0)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
k = BigIntToBytes(bLen, big.NewInt(int64(256)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
k = BigIntToBytes(bLen, big.NewInt(int64(257)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
k = BigIntToBytes(bLen, big.NewInt(int64(515)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
k = BigIntToBytes(bLen, big.NewInt(int64(770)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
k = BigIntToBytes(bLen, big.NewInt(int64(388)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
k = BigIntToBytes(bLen, big.NewInt(int64(900)))
err = tree.Add(k, v)
c.Assert(err, qt.IsNil)
//
// err = tree.PrintGraphviz(nil)
// c.Assert(err, qt.IsNil)
}
func TestGet(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
for i := 0; i < 10; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
if err := tree.Add(k, v); err != nil {
t.Fatal(err)
}
}
k := BigIntToBytes(bLen, big.NewInt(int64(7)))
gettedKey, gettedValue, err := tree.Get(k)
c.Assert(err, qt.IsNil)
c.Check(gettedKey, qt.DeepEquals, k)
c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
}
func TestGenProofAndVerify(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
for i := 0; i < 10; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
if err := tree.Add(k, v); err != nil {
t.Fatal(err)
}
}
k := BigIntToBytes(bLen, big.NewInt(int64(7)))
v := BigIntToBytes(bLen, big.NewInt(int64(14)))
kAux, proofV, siblings, existence, err := tree.GenProof(k)
c.Assert(err, qt.IsNil)
c.Assert(proofV, qt.DeepEquals, v)
c.Assert(k, qt.DeepEquals, kAux)
c.Assert(existence, qt.IsTrue)
root, err := tree.Root()
c.Assert(err, qt.IsNil)
verif, err := CheckProof(tree.hashFunction, k, v, root, siblings)
c.Assert(err, qt.IsNil)
c.Check(verif, qt.IsTrue)
}
func TestDumpAndImportDump(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
bLen := tree1.HashFunction().Len()
for i := 0; i < 16; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
if err := tree1.Add(k, v); err != nil {
t.Fatal(err)
}
}
e, err := tree1.Dump(nil)
c.Assert(err, qt.IsNil)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree2.db.Close() //nolint:errcheck
err = tree2.ImportDump(e)
c.Assert(err, qt.IsNil)
root1, err := tree1.Root()
c.Assert(err, qt.IsNil)
root2, err := tree2.Root()
c.Assert(err, qt.IsNil)
c.Check(root2, qt.DeepEquals, root1)
c.Check(hex.EncodeToString(root2), qt.Equals,
"0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
}
func TestRWMutex(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
bLen := tree.HashFunction().Len()
var keys, values [][]byte
for i := 0; i < 1000; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(0))
keys = append(keys, k)
values = append(values, v)
}
go func() {
_, err = tree.AddBatch(keys, values)
if err != nil {
panic(err)
}
}()
time.Sleep(500 * time.Millisecond)
k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
if err := tree.Add(k, v); err != nil {
t.Fatal(err)
}
}
// TODO UPDATE
// func TestSetGetNLeafs(t *testing.T) {
// c := qt.New(t)
// database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
// c.Assert(err, qt.IsNil)
// tree, err := NewTree(database, 100, HashFunctionPoseidon)
// c.Assert(err, qt.IsNil)
//
// // 0
// tree.dbBatch = tree.db.NewBatch()
//
// err = tree.setNLeafs(0)
// c.Assert(err, qt.IsNil)
//
// err = tree.dbBatch.Write()
// c.Assert(err, qt.IsNil)
//
// n, err := tree.GetNLeafs()
// c.Assert(err, qt.IsNil)
// c.Assert(n, qt.Equals, 0)
//
// // 1024
// tree.dbBatch = tree.db.NewBatch()
//
// err = tree.setNLeafs(1024)
// c.Assert(err, qt.IsNil)
//
// err = tree.dbBatch.Write()
// c.Assert(err, qt.IsNil)
//
// n, err = tree.GetNLeafs()
// c.Assert(err, qt.IsNil)
// c.Assert(n, qt.Equals, 1024)
//
// // 2**64 -1
// tree.dbBatch = tree.db.NewBatch()
//
// maxUint := ^uint(0)
// maxInt := int(maxUint >> 1)
//
// err = tree.setNLeafs(maxInt)
// c.Assert(err, qt.IsNil)
//
// err = tree.dbBatch.Write()
// c.Assert(err, qt.IsNil)
//
// n, err = tree.GetNLeafs()
// c.Assert(err, qt.IsNil)
// c.Assert(n, qt.Equals, maxInt)
// }
func TestSnapshot(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 100, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
// fill the tree
bLen := tree.HashFunction().Len()
var keys, values [][]byte
for i := 0; i < 1000; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(int64(i)))
keys = append(keys, k)
values = append(values, v)
}
indexes, err := tree.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
checkRootBIString(c, tree,
"13742386369878513332697380582061714160370929283209286127733983161245560237407")
// do a snapshot, and expect the same root than the original tree
snapshotTree, err := tree.Snapshot(nil)
c.Assert(err, qt.IsNil)
checkRootBIString(c, snapshotTree,
"13742386369878513332697380582061714160370929283209286127733983161245560237407")
// check that the snapshotTree can not be updated
_, err = snapshotTree.AddBatch(keys, values)
c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
err = snapshotTree.Add([]byte("test"), []byte("test"))
c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
err = snapshotTree.Update([]byte("test"), []byte("test"))
c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
err = snapshotTree.ImportDump(nil)
c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
// update the original tree by adding a new key-value, and check that
// snapshotTree still has the old root, but the original tree has a new
// root
err = tree.Add([]byte("test"), []byte("test"))
c.Assert(err, qt.IsNil)
checkRootBIString(c, snapshotTree,
"13742386369878513332697380582061714160370929283209286127733983161245560237407")
checkRootBIString(c, tree,
"1025190963769001718196479367844646783678188389989148142691917685159698888868")
}
func BenchmarkAdd(b *testing.B) {
bLen := 32 // for both Poseidon & Sha256
// prepare inputs
var ks, vs [][]byte
for i := 0; i < 1000; i++ {
k := BigIntToBytes(bLen, big.NewInt(int64(i)))
v := BigIntToBytes(bLen, big.NewInt(int64(i)))
ks = append(ks, k)
vs = append(vs, v)
}
b.Run("Poseidon", func(b *testing.B) {
benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
})
b.Run("Sha256", func(b *testing.B) {
benchmarkAdd(b, HashFunctionSha256, ks, vs)
})
}
func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
c := qt.New(b)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 140, hashFunc)
c.Assert(err, qt.IsNil)
defer tree.db.Close() //nolint:errcheck
for i := 0; i < len(ks); i++ {
if err := tree.Add(ks[i], vs[i]); err != nil {
b.Fatal(err)
}
}
}