diff --git a/tree.go b/tree.go index 6e8affa..c97d0b7 100644 --- a/tree.go +++ b/tree.go @@ -434,8 +434,22 @@ func bytesToBitmap(b []byte) []bool { // Get returns the value for a given key func (t *Tree) Get(k []byte) ([]byte, []byte, error) { - // unimplemented - return nil, nil, fmt.Errorf("unimplemented") + keyPath := make([]byte, t.hashFunction.Len()) + copy(keyPath[:], k) + + path := getPath(t.maxLevels, keyPath) + // go down to the leaf + var siblings [][]byte + _, value, _, err := t.down(k, t.root, siblings, path, 0, true) + if err != nil { + return nil, nil, err + } + leafK, leafV := readLeafValue(value) + if !bytes.Equal(k, leafK) { + panic(fmt.Errorf("%s != %s", BytesToBigInt(k), BytesToBigInt(leafK))) + } + + return leafK, leafV, nil } // CheckProof verifies the given proof. The proof verification depends on the diff --git a/tree_test.go b/tree_test.go index 0a00326..bbab864 100644 --- a/tree_test.go +++ b/tree_test.go @@ -145,6 +145,26 @@ func TestAux(t *testing.T) { assert.Nil(t, err) } +func TestGet(t *testing.T) { + tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon) + require.Nil(t, err) + + defer tree.db.Close() + for i := 0; i < 10; 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) + } + } + + k := BigIntToBytes(big.NewInt(int64(7))) + gettedKey, gettedValue, err := tree.Get(k) + assert.Nil(t, err) + assert.Equal(t, k, gettedKey) + assert.Equal(t, BigIntToBytes(big.NewInt(int64(7*2))), gettedValue) +} + func TestGenProofAndVerify(t *testing.T) { tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon) require.Nil(t, err)