Skip to content

Commit

Permalink
Extend Dump methods to work with Writer&Reader
Browse files Browse the repository at this point in the history
Add methods DumpWriter & ImportDumpReader to allow generating tree
dumps and reading them working with Reader & Writer, in order to write
and read them directly from a file (internally line by line).
  • Loading branch information
arnaucube committed Aug 25, 2022
1 parent 688a2e8 commit 2380214
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 10 deletions.
52 changes: 46 additions & 6 deletions tree.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ the Blake2b hash function, which has much faster computation time.
package arbo

import (
"bufio"
"bytes"
"encoding/binary"
"encoding/hex"
Expand Down Expand Up @@ -1320,12 +1321,26 @@ func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
return t.iterWithStop(rTx, k, 0, f2)
}

// Dump exports all the Tree leafs in a byte array of length:
// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
// Dump exports all the Tree leafs in a byte array
func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
return t.dump(fromRoot, nil)
}

// DumpWriter exports all the Tree leafs writing the bytes in the given Writer
func (t *Tree) DumpWriter(fromRoot []byte, w *bufio.Writer) error {
_, err := t.dump(fromRoot, w)
return err
}

// dump exports all the Tree leafs. If the given w is nil, it will return a
// byte array with the dump, if w contains a *bufio.Writer, it will write the
// dump in w.
// The format of the dump is the following:
// Dump length: [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
// [ len(k) | len(v) | key | value ]
// Where S is the size of the output of the hash function used for the Tree.
func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
func (t *Tree) dump(fromRoot []byte, w *bufio.Writer) ([]byte, error) {
// allow to define which root to use
if fromRoot == nil {
var err error
Expand Down Expand Up @@ -1357,7 +1372,25 @@ func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
kv[1] = byte(len(leafV))
copy(kv[2:2+len(leafK)], leafK)
copy(kv[2+len(leafK):], leafV)
b = append(b, kv...)

if w == nil {
b = append(b, kv...)
} else {
n, err := w.Write(kv)
if err != nil {
callbackErr = fmt.Errorf("dump: w.Write, %s", err)
return true
}
if n != len(kv) {
callbackErr = fmt.Errorf("dump: w.Write n!=len(kv), %s", err)
return true
}
err = w.Flush()
if err != nil {
callbackErr = fmt.Errorf("dump: w.Flush, %s", err)
return true
}
}
return false
})
if callbackErr != nil {
Expand All @@ -1367,8 +1400,16 @@ func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
}

// ImportDump imports the leafs (that have been exported with the Dump method)
// in the Tree.
// in the Tree, reading them from the given byte array.
func (t *Tree) ImportDump(b []byte) error {
bytesReader := bytes.NewReader(b)
r := bufio.NewReader(bytesReader)
return t.ImportDumpReader(r)
}

// ImportDumpReader imports the leafs (that have been exported with the Dump
// method) in the Tree, reading them from the given reader.
func (t *Tree) ImportDumpReader(r *bufio.Reader) error {
if !t.editable() {
return ErrSnapshotNotEditable
}
Expand All @@ -1380,7 +1421,6 @@ func (t *Tree) ImportDump(b []byte) error {
return ErrTreeNotEmpty
}

r := bytes.NewReader(b)
var keys, values [][]byte
for {
l := make([]byte, 2)
Expand Down
39 changes: 35 additions & 4 deletions tree_test.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
package arbo

import (
"bufio"
"encoding/hex"
"math"
"math/big"
"os"
"path/filepath"
"runtime"
"testing"
"time"
Expand Down Expand Up @@ -458,6 +461,14 @@ func TestGenProofAndVerify(t *testing.T) {
}

func TestDumpAndImportDump(t *testing.T) {
testDumpAndImportDump(t, false)
}

func TestDumpAndImportDumpInFile(t *testing.T) {
testDumpAndImportDump(t, true)
}

func testDumpAndImportDump(t *testing.T, inFile bool) {
c := qt.New(t)
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
Expand All @@ -475,17 +486,37 @@ func TestDumpAndImportDump(t *testing.T) {
}
}

e, err := tree1.Dump(nil)
c.Assert(err, qt.IsNil)
var e []byte
filePath := c.TempDir()
fileName := filepath.Join(filePath, "dump.bin")
if inFile {
f, err := os.Create(fileName)
c.Assert(err, qt.IsNil)
w := bufio.NewWriter(f)
err = tree1.DumpWriter(nil, w)
c.Assert(err, qt.IsNil)
} else {
e, err = tree1.Dump(nil)
c.Assert(err, qt.IsNil)
}

database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(Config{Database: database2, MaxLevels: 256,
HashFunction: HashFunctionPoseidon})
c.Assert(err, qt.IsNil)
defer tree2.db.Close() //nolint:errcheck
err = tree2.ImportDump(e)
c.Assert(err, qt.IsNil)

if inFile {
f, err := os.Open(filepath.Clean(fileName))
c.Assert(err, qt.IsNil)
r := bufio.NewReader(f)
err = tree2.ImportDumpReader(r)
c.Assert(err, qt.IsNil)
} else {
err = tree2.ImportDump(e)
c.Assert(err, qt.IsNil)
}

root1, err := tree1.Root()
c.Assert(err, qt.IsNil)
Expand Down

0 comments on commit 2380214

Please sign in to comment.