From 2380214412bb0f5b18a743b99bff35d1b73fbf7a Mon Sep 17 00:00:00 2001 From: arnaucube Date: Thu, 25 Aug 2022 10:04:09 +0200 Subject: [PATCH] Extend Dump methods to work with Writer&Reader 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). --- tree.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++------ tree_test.go | 39 +++++++++++++++++++++++++++++++++++---- 2 files changed, 81 insertions(+), 10 deletions(-) diff --git a/tree.go b/tree.go index ef1b830..8eb646e 100644 --- a/tree.go +++ b/tree.go @@ -11,6 +11,7 @@ the Blake2b hash function, which has much faster computation time. package arbo import ( + "bufio" "bytes" "encoding/binary" "encoding/hex" @@ -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 @@ -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 { @@ -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 } @@ -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) diff --git a/tree_test.go b/tree_test.go index d3bd48f..00a0dce 100644 --- a/tree_test.go +++ b/tree_test.go @@ -1,9 +1,12 @@ package arbo import ( + "bufio" "encoding/hex" "math" "math/big" + "os" + "path/filepath" "runtime" "testing" "time" @@ -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) @@ -475,8 +486,19 @@ 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) @@ -484,8 +506,17 @@ func TestDumpAndImportDump(t *testing.T) { 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)