You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1554 lines
42 KiB

3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
  1. /*
  2. Package arbo implements a Merkle Tree compatible with the circomlib
  3. implementation of the MerkleTree, following the specification from
  4. https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
  5. https://eprint.iacr.org/2018/955.
  6. Allows to define which hash function to use. So for example, when working with
  7. zkSnarks the Poseidon hash function can be used, but when not, it can be used
  8. the Blake2b hash function, which has much faster computation time.
  9. */
  10. package arbo
  11. import (
  12. "bufio"
  13. "bytes"
  14. "encoding/binary"
  15. "encoding/hex"
  16. "fmt"
  17. "io"
  18. "math"
  19. "runtime"
  20. "sync"
  21. "go.vocdoni.io/dvote/db"
  22. )
  23. const (
  24. // PrefixValueLen defines the bytes-prefix length used for the Value
  25. // bytes representation stored in the db
  26. PrefixValueLen = 2
  27. // PrefixValueEmpty is used for the first byte of a Value to indicate
  28. // that is an Empty value
  29. PrefixValueEmpty = 0
  30. // PrefixValueLeaf is used for the first byte of a Value to indicate
  31. // that is a Leaf value
  32. PrefixValueLeaf = 1
  33. // PrefixValueIntermediate is used for the first byte of a Value to
  34. // indicate that is a Intermediate value
  35. PrefixValueIntermediate = 2
  36. // nChars is used to crop the Graphviz nodes labels
  37. nChars = 4
  38. maxUint8 = int(^uint8(0)) // 2**8 -1
  39. maxUint16 = int(^uint16(0)) // 2**16 -1
  40. )
  41. var (
  42. // DefaultThresholdNLeafs defines the threshold number of leafs in the
  43. // tree that determines if AddBatch will work in memory or in disk. It
  44. // is defined when calling NewTree, and if set to 0 it will work always
  45. // in disk.
  46. DefaultThresholdNLeafs = 65536
  47. dbKeyRoot = []byte("root")
  48. dbKeyNLeafs = []byte("nleafs")
  49. emptyValue = []byte{0}
  50. // ErrKeyNotFound is used when a key is not found in the db neither in
  51. // the current db Batch.
  52. ErrKeyNotFound = fmt.Errorf("key not found")
  53. // ErrKeyAlreadyExists is used when trying to add a key as leaf to the
  54. // tree that already exists.
  55. ErrKeyAlreadyExists = fmt.Errorf("key already exists")
  56. // ErrInvalidValuePrefix is used when going down into the tree, a value
  57. // is read from the db and has an unrecognized prefix.
  58. ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
  59. // ErrDBNoTx is used when trying to use Tree.dbPut but Tree.dbBatch==nil
  60. ErrDBNoTx = fmt.Errorf("dbPut error: no db Batch")
  61. // ErrMaxLevel indicates when going down into the tree, the max level is
  62. // reached
  63. ErrMaxLevel = fmt.Errorf("max level reached")
  64. // ErrMaxVirtualLevel indicates when going down into the tree, the max
  65. // virtual level is reached
  66. ErrMaxVirtualLevel = fmt.Errorf("max virtual level reached")
  67. // ErrSnapshotNotEditable indicates when the tree is a non writable
  68. // snapshot, thus can not be modified
  69. ErrSnapshotNotEditable = fmt.Errorf("snapshot tree can not be edited")
  70. // ErrTreeNotEmpty indicates when the tree was expected to be empty and
  71. // it is not
  72. ErrTreeNotEmpty = fmt.Errorf("tree is not empty")
  73. )
  74. // Tree defines the struct that implements the MerkleTree functionalities
  75. type Tree struct {
  76. sync.Mutex
  77. db db.Database
  78. maxLevels int
  79. // thresholdNLeafs defines the threshold number of leafs in the tree
  80. // that determines if AddBatch will work in memory or in disk. It is
  81. // defined when calling NewTree, and if set to 0 it will work always in
  82. // disk.
  83. thresholdNLeafs int
  84. snapshotRoot []byte
  85. hashFunction HashFunction
  86. // TODO in the methods that use it, check if emptyHash param is len>0
  87. // (check if it has been initialized)
  88. emptyHash []byte
  89. dbg *dbgStats
  90. }
  91. // Config defines the configuration for calling NewTree & NewTreeWithTx methods
  92. type Config struct {
  93. Database db.Database
  94. MaxLevels int
  95. ThresholdNLeafs int
  96. HashFunction HashFunction
  97. }
  98. // NewTree returns a new Tree, if there is a Tree still in the given database, it
  99. // will load it.
  100. func NewTree(cfg Config) (*Tree, error) {
  101. wTx := cfg.Database.WriteTx()
  102. defer wTx.Discard()
  103. t, err := NewTreeWithTx(wTx, cfg)
  104. if err != nil {
  105. return nil, err
  106. }
  107. if err = wTx.Commit(); err != nil {
  108. return nil, err
  109. }
  110. return t, nil
  111. }
  112. // NewTreeWithTx returns a new Tree using the given db.WriteTx, which will not
  113. // be ccommited inside this method, if there is a Tree still in the given
  114. // database, it will load it.
  115. func NewTreeWithTx(wTx db.WriteTx, cfg Config) (*Tree, error) {
  116. // if thresholdNLeafs is set to 0, use the DefaultThresholdNLeafs
  117. if cfg.ThresholdNLeafs == 0 {
  118. cfg.ThresholdNLeafs = DefaultThresholdNLeafs
  119. }
  120. t := Tree{db: cfg.Database, maxLevels: cfg.MaxLevels,
  121. thresholdNLeafs: cfg.ThresholdNLeafs, hashFunction: cfg.HashFunction}
  122. t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
  123. _, err := wTx.Get(dbKeyRoot)
  124. if err == db.ErrKeyNotFound {
  125. // store new root 0 (empty)
  126. if err = wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
  127. return nil, err
  128. }
  129. if err = t.setNLeafs(wTx, 0); err != nil {
  130. return nil, err
  131. }
  132. return &t, nil
  133. } else if err != nil {
  134. return nil, err
  135. }
  136. return &t, nil
  137. }
  138. // Root returns the root of the Tree
  139. func (t *Tree) Root() ([]byte, error) {
  140. rTx := t.db.ReadTx()
  141. defer rTx.Discard()
  142. return t.RootWithTx(rTx)
  143. }
  144. // RootWithTx returns the root of the Tree using the given db.ReadTx
  145. func (t *Tree) RootWithTx(rTx db.ReadTx) ([]byte, error) {
  146. // if snapshotRoot is defined, means that the tree is a snapshot, and
  147. // the root is not obtained from the db, but from the snapshotRoot
  148. // parameter
  149. if t.snapshotRoot != nil {
  150. return t.snapshotRoot, nil
  151. }
  152. // get db root
  153. return rTx.Get(dbKeyRoot)
  154. }
  155. func (t *Tree) setRoot(wTx db.WriteTx, root []byte) error {
  156. return wTx.Set(dbKeyRoot, root)
  157. }
  158. // HashFunction returns Tree.hashFunction
  159. func (t *Tree) HashFunction() HashFunction {
  160. return t.hashFunction
  161. }
  162. // editable returns true if the tree is editable, and false when is not
  163. // editable (because is a snapshot tree)
  164. func (t *Tree) editable() bool {
  165. return t.snapshotRoot == nil
  166. }
  167. // Invalid is used when a key-value can not be added trough AddBatch, and
  168. // contains the index of the key-value and the error.
  169. type Invalid struct {
  170. Index int
  171. Error error
  172. }
  173. // AddBatch adds a batch of key-values to the Tree. Returns an array containing
  174. // the indexes of the keys failed to add. Supports empty values as input
  175. // parameters, which is equivalent to 0 valued byte array.
  176. func (t *Tree) AddBatch(keys, values [][]byte) ([]Invalid, error) {
  177. wTx := t.db.WriteTx()
  178. defer wTx.Discard()
  179. invalids, err := t.AddBatchWithTx(wTx, keys, values)
  180. if err != nil {
  181. return invalids, err
  182. }
  183. return invalids, wTx.Commit()
  184. }
  185. // AddBatchWithTx does the same than the AddBatch method, but allowing to pass
  186. // the db.WriteTx that is used. The db.WriteTx will not be committed inside
  187. // this method.
  188. func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
  189. t.Lock()
  190. defer t.Unlock()
  191. if !t.editable() {
  192. return nil, ErrSnapshotNotEditable
  193. }
  194. e := []byte{}
  195. // equal the number of keys & values
  196. if len(keys) > len(values) {
  197. // add missing values
  198. for i := len(values); i < len(keys); i++ {
  199. values = append(values, e)
  200. }
  201. } else if len(keys) < len(values) {
  202. // crop extra values
  203. values = values[:len(keys)]
  204. }
  205. nLeafs, err := t.GetNLeafsWithTx(wTx)
  206. if err != nil {
  207. return nil, err
  208. }
  209. if nLeafs > t.thresholdNLeafs {
  210. return t.addBatchInDisk(wTx, keys, values)
  211. }
  212. return t.addBatchInMemory(wTx, keys, values)
  213. }
  214. func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
  215. nCPU := flp2(runtime.NumCPU())
  216. if nCPU == 1 || len(keys) < nCPU {
  217. var invalids []Invalid
  218. for i := 0; i < len(keys); i++ {
  219. if err := t.addWithTx(wTx, keys[i], values[i]); err != nil {
  220. invalids = append(invalids, Invalid{i, err})
  221. }
  222. }
  223. return invalids, nil
  224. }
  225. kvs, invalids, err := keysValuesToKvs(t.maxLevels, keys, values)
  226. if err != nil {
  227. return nil, err
  228. }
  229. buckets := splitInBuckets(kvs, nCPU)
  230. root, err := t.RootWithTx(wTx)
  231. if err != nil {
  232. return nil, err
  233. }
  234. l := int(math.Log2(float64(nCPU)))
  235. subRoots, err := t.getSubRootsAtLevel(wTx, root, l+1)
  236. if err != nil {
  237. return nil, err
  238. }
  239. if len(subRoots) != nCPU {
  240. // Already populated Tree but Unbalanced.
  241. // add one key at each bucket, and then continue with the flow
  242. for i := 0; i < len(buckets); i++ {
  243. // add one leaf of the bucket, if there is an error when
  244. // adding the k-v, try to add the next one of the bucket
  245. // (until one is added)
  246. inserted := -1
  247. for j := 0; j < len(buckets[i]); j++ {
  248. if newRoot, err := t.add(wTx, root, 0,
  249. buckets[i][j].k, buckets[i][j].v); err == nil {
  250. inserted = j
  251. root = newRoot
  252. break
  253. }
  254. }
  255. // remove the inserted element from buckets[i]
  256. if inserted != -1 {
  257. buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
  258. }
  259. }
  260. subRoots, err = t.getSubRootsAtLevel(wTx, root, l+1)
  261. if err != nil {
  262. return nil, err
  263. }
  264. }
  265. if len(subRoots) != nCPU {
  266. return nil, fmt.Errorf("This error should not be reached."+
  267. " len(subRoots) != nCPU, len(subRoots)=%d, nCPU=%d."+
  268. " Please report it in a new issue:"+
  269. " https://github.com/vocdoni/arbo/issues/new", len(subRoots), nCPU)
  270. }
  271. invalidsInBucket := make([][]Invalid, nCPU)
  272. txs := make([]db.WriteTx, nCPU)
  273. for i := 0; i < nCPU; i++ {
  274. txs[i] = t.db.WriteTx()
  275. err := txs[i].Apply(wTx)
  276. if err != nil {
  277. return nil, err
  278. }
  279. }
  280. var wg sync.WaitGroup
  281. wg.Add(nCPU)
  282. for i := 0; i < nCPU; i++ {
  283. go func(cpu int) {
  284. // use different wTx for each cpu, after once all
  285. // are done, iter over the cpuWTxs and copy their
  286. // content into the main wTx
  287. for j := 0; j < len(buckets[cpu]); j++ {
  288. newSubRoot, err := t.add(txs[cpu], subRoots[cpu],
  289. l, buckets[cpu][j].k, buckets[cpu][j].v)
  290. if err != nil {
  291. invalidsInBucket[cpu] = append(invalidsInBucket[cpu],
  292. Invalid{buckets[cpu][j].pos, err})
  293. continue
  294. }
  295. // if there has not been errors, set the new subRoots[cpu]
  296. subRoots[cpu] = newSubRoot
  297. }
  298. wg.Done()
  299. }(i)
  300. }
  301. wg.Wait()
  302. for i := 0; i < nCPU; i++ {
  303. if err := wTx.Apply(txs[i]); err != nil {
  304. return nil, err
  305. }
  306. txs[i].Discard()
  307. }
  308. for i := 0; i < len(invalidsInBucket); i++ {
  309. invalids = append(invalids, invalidsInBucket[i]...)
  310. }
  311. newRoot, err := t.upFromSubRoots(wTx, subRoots)
  312. if err != nil {
  313. return nil, err
  314. }
  315. // update dbKeyNLeafs
  316. if err := t.SetRootWithTx(wTx, newRoot); err != nil {
  317. return nil, err
  318. }
  319. // update nLeafs
  320. if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
  321. return nil, err
  322. }
  323. return invalids, nil
  324. }
  325. func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) {
  326. // is a method of Tree just to get access to t.hashFunction and
  327. // t.emptyHash.
  328. // go up from subRoots to up, storing nodes in the given WriteTx
  329. // once up at the root, store it in the WriteTx using the dbKeyRoot
  330. if len(subRoots) == 1 {
  331. return subRoots[0], nil
  332. }
  333. // get the subRoots values to know the node types of each subRoot
  334. nodeTypes := make([]byte, len(subRoots))
  335. for i := 0; i < len(subRoots); i++ {
  336. if bytes.Equal(subRoots[i], t.emptyHash) {
  337. nodeTypes[i] = PrefixValueEmpty
  338. continue
  339. }
  340. v, err := wTx.Get(subRoots[i])
  341. if err != nil {
  342. return nil, err
  343. }
  344. nodeTypes[i] = v[0]
  345. }
  346. var newSubRoots [][]byte
  347. for i := 0; i < len(subRoots); i += 2 {
  348. if (bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash)) ||
  349. (nodeTypes[i] == PrefixValueLeaf && bytes.Equal(subRoots[i+1], t.emptyHash)) {
  350. // when both sub nodes are empty, the parent is also empty
  351. // or
  352. // when 1st sub node is a leaf but the 2nd is empty, the
  353. // leaf is used as 'parent'
  354. newSubRoots = append(newSubRoots, subRoots[i])
  355. continue
  356. }
  357. if bytes.Equal(subRoots[i], t.emptyHash) && nodeTypes[i+1] == PrefixValueLeaf {
  358. // when 2nd sub node is a leaf but the 1st is empty,
  359. // the leaf is used as 'parent'
  360. newSubRoots = append(newSubRoots, subRoots[i+1])
  361. continue
  362. }
  363. k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
  364. if err != nil {
  365. return nil, err
  366. }
  367. // store k-v to db
  368. if err = wTx.Set(k, v); err != nil {
  369. return nil, err
  370. }
  371. newSubRoots = append(newSubRoots, k)
  372. }
  373. return t.upFromSubRoots(wTx, newSubRoots)
  374. }
  375. func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
  376. // go at level l and return each node key, where each node key is the
  377. // subRoot of the subTree that starts there
  378. var subRoots [][]byte
  379. err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
  380. if currLvl == l && !bytes.Equal(k, t.emptyHash) {
  381. subRoots = append(subRoots, k)
  382. }
  383. if currLvl >= l {
  384. return true // to stop the iter from going down
  385. }
  386. return false
  387. })
  388. return subRoots, err
  389. }
  390. func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
  391. vt, err := t.loadVT(wTx)
  392. if err != nil {
  393. return nil, err
  394. }
  395. invalids, err := vt.addBatch(keys, values)
  396. if err != nil {
  397. return nil, err
  398. }
  399. // once the VirtualTree is build, compute the hashes
  400. pairs, err := vt.computeHashes()
  401. if err != nil {
  402. // currently invalids in computeHashes are not counted,
  403. // but should not be needed, as if there is an error there is
  404. // nothing stored in the db and the error is returned
  405. return nil, err
  406. }
  407. // store pairs in db
  408. for i := 0; i < len(pairs); i++ {
  409. if err := wTx.Set(pairs[i][0], pairs[i][1]); err != nil {
  410. return nil, err
  411. }
  412. }
  413. // store root (from the vt) to db
  414. if vt.root != nil {
  415. if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil {
  416. return nil, err
  417. }
  418. }
  419. // update nLeafs
  420. if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
  421. return nil, err
  422. }
  423. return invalids, nil
  424. }
  425. // loadVT loads a new virtual tree (vt) from the current Tree, which contains
  426. // the same leafs.
  427. func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
  428. vt := newVT(t.maxLevels, t.hashFunction)
  429. vt.params.dbg = t.dbg
  430. var callbackErr error
  431. err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
  432. if v[0] != PrefixValueLeaf {
  433. return false
  434. }
  435. leafK, leafV := ReadLeafValue(v)
  436. if err := vt.add(0, leafK, leafV); err != nil {
  437. callbackErr = err
  438. return true
  439. }
  440. return false
  441. })
  442. if callbackErr != nil {
  443. return vt, callbackErr
  444. }
  445. return vt, err
  446. }
  447. // Add inserts the key-value into the Tree. If the inputs come from a
  448. // *big.Int, is expected that are represented by a Little-Endian byte array
  449. // (for circom compatibility).
  450. func (t *Tree) Add(k, v []byte) error {
  451. wTx := t.db.WriteTx()
  452. defer wTx.Discard()
  453. if err := t.AddWithTx(wTx, k, v); err != nil {
  454. return err
  455. }
  456. return wTx.Commit()
  457. }
  458. // AddWithTx does the same than the Add method, but allowing to pass the
  459. // db.WriteTx that is used. The db.WriteTx will not be committed inside this
  460. // method.
  461. func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
  462. t.Lock()
  463. defer t.Unlock()
  464. if !t.editable() {
  465. return ErrSnapshotNotEditable
  466. }
  467. return t.addWithTx(wTx, k, v)
  468. }
  469. // warning: addWithTx does not use the Tree mutex, the mutex is responsibility
  470. // of the methods calling this method, and same with t.editable().
  471. func (t *Tree) addWithTx(wTx db.WriteTx, k, v []byte) error {
  472. root, err := t.RootWithTx(wTx)
  473. if err != nil {
  474. return err
  475. }
  476. root, err = t.add(wTx, root, 0, k, v) // add from level 0
  477. if err != nil {
  478. return err
  479. }
  480. // store root to db
  481. if err := t.setRoot(wTx, root); err != nil {
  482. return err
  483. }
  484. // update nLeafs
  485. if err = t.incNLeafs(wTx, 1); err != nil {
  486. return err
  487. }
  488. return nil
  489. }
  490. // keyPathFromKey returns the keyPath and checks that the key is not bigger
  491. // than maximum key length for the tree maxLevels size.
  492. // This is because if the key bits length is bigger than the maxLevels of the
  493. // tree, two different keys that their difference is at the end, will collision
  494. // in the same leaf of the tree (at the max depth).
  495. func keyPathFromKey(maxLevels int, k []byte) ([]byte, error) {
  496. maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
  497. if len(k) > maxKeyLen {
  498. return nil, fmt.Errorf("len(k) can not be bigger than ceil(maxLevels/8), where"+
  499. " len(k): %d, maxLevels: %d, max key len=ceil(maxLevels/8): %d. Might need"+
  500. " a bigger tree depth (maxLevels>=%d) in order to input keys of length %d",
  501. len(k), maxLevels, maxKeyLen, len(k)*8, len(k)) //nolint:gomnd
  502. }
  503. keyPath := make([]byte, maxKeyLen) //nolint:gomnd
  504. copy(keyPath[:], k)
  505. return keyPath, nil
  506. }
  507. func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
  508. keyPath, err := keyPathFromKey(t.maxLevels, k)
  509. if err != nil {
  510. return nil, err
  511. }
  512. path := getPath(t.maxLevels, keyPath)
  513. // go down to the leaf
  514. var siblings [][]byte
  515. _, _, siblings, err = t.down(wTx, k, root, siblings, path, fromLvl, false)
  516. if err != nil {
  517. return nil, err
  518. }
  519. leafKey, leafValue, err := t.newLeafValue(k, v)
  520. if err != nil {
  521. return nil, err
  522. }
  523. if err := wTx.Set(leafKey, leafValue); err != nil {
  524. return nil, err
  525. }
  526. // go up to the root
  527. if len(siblings) == 0 {
  528. // return the leafKey as root
  529. return leafKey, nil
  530. }
  531. root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
  532. if err != nil {
  533. return nil, err
  534. }
  535. return root, nil
  536. }
  537. // down goes down to the leaf recursively
  538. func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
  539. path []bool, currLvl int, getLeaf bool) (
  540. []byte, []byte, [][]byte, error) {
  541. if currLvl > t.maxLevels {
  542. return nil, nil, nil, ErrMaxLevel
  543. }
  544. var err error
  545. var currValue []byte
  546. if bytes.Equal(currKey, t.emptyHash) {
  547. // empty value
  548. return currKey, emptyValue, siblings, nil
  549. }
  550. currValue, err = rTx.Get(currKey)
  551. if err != nil {
  552. return nil, nil, nil, err
  553. }
  554. switch currValue[0] {
  555. case PrefixValueEmpty: // empty
  556. fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
  557. hex.EncodeToString(newKey), hex.EncodeToString(currKey),
  558. currLvl, hex.EncodeToString(currValue))
  559. panic("This point should not be reached, as the 'if currKey==t.emptyHash'" +
  560. " above should avoid reaching this point. This panic is temporary" +
  561. " for reporting purposes, will be deleted in future versions." +
  562. " Please paste this log (including the previous log lines) in a" +
  563. " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
  564. case PrefixValueLeaf: // leaf
  565. if !bytes.Equal(currValue, emptyValue) {
  566. if getLeaf {
  567. return currKey, currValue, siblings, nil
  568. }
  569. oldLeafKey, _ := ReadLeafValue(currValue)
  570. if bytes.Equal(newKey, oldLeafKey) {
  571. return nil, nil, nil, ErrKeyAlreadyExists
  572. }
  573. oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey)
  574. if err != nil {
  575. return nil, nil, nil, err
  576. }
  577. // if currKey is already used, go down until paths diverge
  578. oldPath := getPath(t.maxLevels, oldLeafKeyFull)
  579. siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
  580. if err != nil {
  581. return nil, nil, nil, err
  582. }
  583. }
  584. return currKey, currValue, siblings, nil
  585. case PrefixValueIntermediate: // intermediate
  586. if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
  587. return nil, nil, nil,
  588. fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
  589. PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
  590. }
  591. // collect siblings while going down
  592. if path[currLvl] {
  593. // right
  594. lChild, rChild := ReadIntermediateChilds(currValue)
  595. siblings = append(siblings, lChild)
  596. return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
  597. }
  598. // left
  599. lChild, rChild := ReadIntermediateChilds(currValue)
  600. siblings = append(siblings, rChild)
  601. return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
  602. default:
  603. return nil, nil, nil, ErrInvalidValuePrefix
  604. }
  605. }
  606. // downVirtually is used when in a leaf already exists, and a new leaf which
  607. // shares the path until the existing leaf is being added
  608. func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
  609. newPath []bool, currLvl int) ([][]byte, error) {
  610. var err error
  611. if currLvl > t.maxLevels-1 {
  612. return nil, ErrMaxVirtualLevel
  613. }
  614. if oldPath[currLvl] == newPath[currLvl] {
  615. siblings = append(siblings, t.emptyHash)
  616. siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
  617. if err != nil {
  618. return nil, err
  619. }
  620. return siblings, nil
  621. }
  622. // reached the divergence
  623. siblings = append(siblings, oldKey)
  624. return siblings, nil
  625. }
  626. // up goes up recursively updating the intermediate nodes
  627. func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
  628. currLvl, toLvl int) ([]byte, error) {
  629. var k, v []byte
  630. var err error
  631. if path[currLvl+toLvl] {
  632. k, v, err = t.newIntermediate(siblings[currLvl], key)
  633. if err != nil {
  634. return nil, err
  635. }
  636. } else {
  637. k, v, err = t.newIntermediate(key, siblings[currLvl])
  638. if err != nil {
  639. return nil, err
  640. }
  641. }
  642. // store k-v to db
  643. if err = wTx.Set(k, v); err != nil {
  644. return nil, err
  645. }
  646. if currLvl == 0 {
  647. // reached the root
  648. return k, nil
  649. }
  650. return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
  651. }
  652. func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
  653. t.dbg.incHash()
  654. return newLeafValue(t.hashFunction, k, v)
  655. }
  656. // newLeafValue takes a key & value from a leaf, and computes the leaf hash,
  657. // which is used as the leaf key. And the value is the concatenation of the
  658. // inputed key & value. The output of this function is used as key-value to
  659. // store the leaf in the DB.
  660. // [ 1 byte | 1 byte | N bytes | M bytes ]
  661. // [ type of node | length of key | key | value ]
  662. func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
  663. leafKey, err := hashFunc.Hash(k, v, []byte{1})
  664. if err != nil {
  665. return nil, nil, err
  666. }
  667. var leafValue []byte
  668. leafValue = append(leafValue, byte(PrefixValueLeaf))
  669. if len(k) > maxUint8 {
  670. return nil, nil, fmt.Errorf("newLeafValue: len(k) > %v", maxUint8)
  671. }
  672. leafValue = append(leafValue, byte(len(k)))
  673. leafValue = append(leafValue, k...)
  674. leafValue = append(leafValue, v...)
  675. return leafKey, leafValue, nil
  676. }
  677. // ReadLeafValue reads from a byte array the leaf key & value
  678. func ReadLeafValue(b []byte) ([]byte, []byte) {
  679. if len(b) < PrefixValueLen {
  680. return []byte{}, []byte{}
  681. }
  682. kLen := b[1]
  683. if len(b) < PrefixValueLen+int(kLen) {
  684. return []byte{}, []byte{}
  685. }
  686. k := b[PrefixValueLen : PrefixValueLen+kLen]
  687. v := b[PrefixValueLen+kLen:]
  688. return k, v
  689. }
  690. func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
  691. t.dbg.incHash()
  692. return newIntermediate(t.hashFunction, l, r)
  693. }
  694. // newIntermediate takes the left & right keys of a intermediate node, and
  695. // computes its hash. Returns the hash of the node, which is the node key, and a
  696. // byte array that contains the value (which contains the left & right child
  697. // keys) to store in the DB.
  698. // [ 1 byte | 1 byte | N bytes | N bytes ]
  699. // [ type of node | length of left key | left key | right key ]
  700. func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
  701. b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
  702. b[0] = PrefixValueIntermediate
  703. if len(l) > maxUint8 {
  704. return nil, nil, fmt.Errorf("newIntermediate: len(l) > %v", maxUint8)
  705. }
  706. b[1] = byte(len(l))
  707. copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
  708. copy(b[PrefixValueLen+hashFunc.Len():], r)
  709. key, err := hashFunc.Hash(l, r)
  710. if err != nil {
  711. return nil, nil, err
  712. }
  713. return key, b, nil
  714. }
  715. // ReadIntermediateChilds reads from a byte array the two childs keys
  716. func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
  717. if len(b) < PrefixValueLen {
  718. return []byte{}, []byte{}
  719. }
  720. lLen := b[1]
  721. if len(b) < PrefixValueLen+int(lLen) {
  722. return []byte{}, []byte{}
  723. }
  724. l := b[PrefixValueLen : PrefixValueLen+lLen]
  725. r := b[PrefixValueLen+lLen:]
  726. return l, r
  727. }
  728. func getPath(numLevels int, k []byte) []bool {
  729. path := make([]bool, numLevels)
  730. for n := 0; n < numLevels; n++ {
  731. path[n] = k[n/8]&(1<<(n%8)) != 0
  732. }
  733. return path
  734. }
  735. // Update updates the value for a given existing key. If the given key does not
  736. // exist, returns an error.
  737. func (t *Tree) Update(k, v []byte) error {
  738. wTx := t.db.WriteTx()
  739. defer wTx.Discard()
  740. if err := t.UpdateWithTx(wTx, k, v); err != nil {
  741. return err
  742. }
  743. return wTx.Commit()
  744. }
  745. // UpdateWithTx does the same than the Update method, but allowing to pass the
  746. // db.WriteTx that is used. The db.WriteTx will not be committed inside this
  747. // method.
  748. func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
  749. t.Lock()
  750. defer t.Unlock()
  751. if !t.editable() {
  752. return ErrSnapshotNotEditable
  753. }
  754. keyPath, err := keyPathFromKey(t.maxLevels, k)
  755. if err != nil {
  756. return err
  757. }
  758. path := getPath(t.maxLevels, keyPath)
  759. root, err := t.RootWithTx(wTx)
  760. if err != nil {
  761. return err
  762. }
  763. var siblings [][]byte
  764. _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
  765. if err != nil {
  766. return err
  767. }
  768. oldKey, _ := ReadLeafValue(valueAtBottom)
  769. if !bytes.Equal(oldKey, k) {
  770. return ErrKeyNotFound
  771. }
  772. leafKey, leafValue, err := t.newLeafValue(k, v)
  773. if err != nil {
  774. return err
  775. }
  776. if err := wTx.Set(leafKey, leafValue); err != nil {
  777. return err
  778. }
  779. // go up to the root
  780. if len(siblings) == 0 {
  781. return t.setRoot(wTx, leafKey)
  782. }
  783. root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
  784. if err != nil {
  785. return err
  786. }
  787. // store root to db
  788. if err := t.setRoot(wTx, root); err != nil {
  789. return err
  790. }
  791. return nil
  792. }
  793. // GenProof generates a MerkleTree proof for the given key. The leaf value is
  794. // returned, together with the packed siblings of the proof, and a boolean
  795. // parameter that indicates if the proof is of existence (true) or not (false).
  796. func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
  797. rTx := t.db.ReadTx()
  798. defer rTx.Discard()
  799. return t.GenProofWithTx(rTx, k)
  800. }
  801. // GenProofWithTx does the same than the GenProof method, but allowing to pass
  802. // the db.ReadTx that is used.
  803. func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
  804. keyPath, err := keyPathFromKey(t.maxLevels, k)
  805. if err != nil {
  806. return nil, nil, nil, false, err
  807. }
  808. path := getPath(t.maxLevels, keyPath)
  809. root, err := t.RootWithTx(rTx)
  810. if err != nil {
  811. return nil, nil, nil, false, err
  812. }
  813. // go down to the leaf
  814. var siblings [][]byte
  815. _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
  816. if err != nil {
  817. return nil, nil, nil, false, err
  818. }
  819. s, err := PackSiblings(t.hashFunction, siblings)
  820. if err != nil {
  821. return nil, nil, nil, false, err
  822. }
  823. leafK, leafV := ReadLeafValue(value)
  824. if !bytes.Equal(k, leafK) {
  825. // key not in tree, proof of non-existence
  826. return leafK, leafV, s, false, nil
  827. }
  828. return leafK, leafV, s, true, nil
  829. }
  830. // PackSiblings packs the siblings into a byte array.
  831. // [ 2 byte | 2 byte | L bytes | S * N bytes ]
  832. // [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
  833. // Where the bitmap indicates if the sibling is 0 or a value from the siblings
  834. // array. And S is the size of the output of the hash function used for the
  835. // Tree. The 2 2-byte that define the full length and bitmap length, are
  836. // encoded in little-endian.
  837. func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
  838. var b []byte
  839. var bitmap []bool
  840. emptySibling := make([]byte, hashFunc.Len())
  841. for i := 0; i < len(siblings); i++ {
  842. if bytes.Equal(siblings[i], emptySibling) {
  843. bitmap = append(bitmap, false)
  844. } else {
  845. bitmap = append(bitmap, true)
  846. b = append(b, siblings[i]...)
  847. }
  848. }
  849. bitmapBytes := bitmapToBytes(bitmap)
  850. l := len(bitmapBytes)
  851. if l > maxUint16 {
  852. return nil, fmt.Errorf("PackSiblings: bitmapBytes length > %v", maxUint16)
  853. }
  854. fullLen := 4 + l + len(b) //nolint:gomnd
  855. if fullLen > maxUint16 {
  856. return nil, fmt.Errorf("PackSiblings: fullLen > %v", maxUint16)
  857. }
  858. res := make([]byte, fullLen)
  859. binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
  860. binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
  861. copy(res[4:4+l], bitmapBytes)
  862. copy(res[4+l:], b)
  863. return res, nil
  864. }
  865. // UnpackSiblings unpacks the siblings from a byte array.
  866. func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
  867. fullLen := binary.LittleEndian.Uint16(b[0:2])
  868. l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
  869. if len(b) != int(fullLen) {
  870. return nil,
  871. fmt.Errorf("expected len: %d, current len: %d",
  872. fullLen, len(b))
  873. }
  874. bitmapBytes := b[4 : 4+l]
  875. bitmap := bytesToBitmap(bitmapBytes)
  876. siblingsBytes := b[4+l:]
  877. iSibl := 0
  878. emptySibl := make([]byte, hashFunc.Len())
  879. var siblings [][]byte
  880. for i := 0; i < len(bitmap); i++ {
  881. if iSibl >= len(siblingsBytes) {
  882. break
  883. }
  884. if bitmap[i] {
  885. siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
  886. iSibl += hashFunc.Len()
  887. } else {
  888. siblings = append(siblings, emptySibl)
  889. }
  890. }
  891. return siblings, nil
  892. }
  893. func bitmapToBytes(bitmap []bool) []byte {
  894. bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
  895. b := make([]byte, bitmapBytesLen)
  896. for i := 0; i < len(bitmap); i++ {
  897. if bitmap[i] {
  898. b[i/8] |= 1 << (i % 8)
  899. }
  900. }
  901. return b
  902. }
  903. func bytesToBitmap(b []byte) []bool {
  904. var bitmap []bool
  905. for i := 0; i < len(b); i++ {
  906. for j := 0; j < 8; j++ {
  907. bitmap = append(bitmap, b[i]&(1<<j) > 0)
  908. }
  909. }
  910. return bitmap
  911. }
  912. // Get returns the value in the Tree for a given key. If the key is not found,
  913. // will return the error ErrKeyNotFound, and in the leafK & leafV parameters
  914. // will be placed the data found in the tree in the leaf that was on the path
  915. // going to the input key.
  916. func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
  917. rTx := t.db.ReadTx()
  918. defer rTx.Discard()
  919. return t.GetWithTx(rTx, k)
  920. }
  921. // GetWithTx does the same than the Get method, but allowing to pass the
  922. // db.ReadTx that is used. If the key is not found, will return the error
  923. // ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
  924. // found in the tree in the leaf that was on the path going to the input key.
  925. func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
  926. keyPath, err := keyPathFromKey(t.maxLevels, k)
  927. if err != nil {
  928. return nil, nil, err
  929. }
  930. path := getPath(t.maxLevels, keyPath)
  931. root, err := t.RootWithTx(rTx)
  932. if err != nil {
  933. return nil, nil, err
  934. }
  935. // go down to the leaf
  936. var siblings [][]byte
  937. _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
  938. if err != nil {
  939. return nil, nil, err
  940. }
  941. leafK, leafV := ReadLeafValue(value)
  942. if !bytes.Equal(k, leafK) {
  943. return leafK, leafV, ErrKeyNotFound
  944. }
  945. return leafK, leafV, nil
  946. }
  947. // CheckProof verifies the given proof. The proof verification depends on the
  948. // HashFunction passed as parameter.
  949. func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
  950. siblings, err := UnpackSiblings(hashFunc, packedSiblings)
  951. if err != nil {
  952. return false, err
  953. }
  954. keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd
  955. copy(keyPath[:], k)
  956. key, _, err := newLeafValue(hashFunc, k, v)
  957. if err != nil {
  958. return false, err
  959. }
  960. path := getPath(len(siblings), keyPath)
  961. for i := len(siblings) - 1; i >= 0; i-- {
  962. if path[i] {
  963. key, _, err = newIntermediate(hashFunc, siblings[i], key)
  964. if err != nil {
  965. return false, err
  966. }
  967. } else {
  968. key, _, err = newIntermediate(hashFunc, key, siblings[i])
  969. if err != nil {
  970. return false, err
  971. }
  972. }
  973. }
  974. if bytes.Equal(key[:], root) {
  975. return true, nil
  976. }
  977. return false, nil
  978. }
  979. func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
  980. oldNLeafs, err := t.GetNLeafsWithTx(wTx)
  981. if err != nil {
  982. return err
  983. }
  984. newNLeafs := oldNLeafs + nLeafs
  985. return t.setNLeafs(wTx, newNLeafs)
  986. }
  987. func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
  988. b := make([]byte, 8)
  989. binary.LittleEndian.PutUint64(b, uint64(nLeafs))
  990. if err := wTx.Set(dbKeyNLeafs, b); err != nil {
  991. return err
  992. }
  993. return nil
  994. }
  995. // GetNLeafs returns the number of Leafs of the Tree.
  996. func (t *Tree) GetNLeafs() (int, error) {
  997. rTx := t.db.ReadTx()
  998. defer rTx.Discard()
  999. return t.GetNLeafsWithTx(rTx)
  1000. }
  1001. // GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
  1002. // pass the db.ReadTx that is used.
  1003. func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
  1004. b, err := rTx.Get(dbKeyNLeafs)
  1005. if err != nil {
  1006. return 0, err
  1007. }
  1008. nLeafs := binary.LittleEndian.Uint64(b)
  1009. return int(nLeafs), nil
  1010. }
  1011. // SetRoot sets the root to the given root
  1012. func (t *Tree) SetRoot(root []byte) error {
  1013. wTx := t.db.WriteTx()
  1014. defer wTx.Discard()
  1015. if err := t.SetRootWithTx(wTx, root); err != nil {
  1016. return err
  1017. }
  1018. return wTx.Commit()
  1019. }
  1020. // SetRootWithTx sets the root to the given root using the given db.WriteTx
  1021. func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
  1022. if !t.editable() {
  1023. return ErrSnapshotNotEditable
  1024. }
  1025. if root == nil {
  1026. return fmt.Errorf("can not SetRoot with nil root")
  1027. }
  1028. // check that the root exists in the db
  1029. if !bytes.Equal(root, t.emptyHash) {
  1030. if _, err := wTx.Get(root); err == ErrKeyNotFound {
  1031. return fmt.Errorf("can not SetRoot with root %x, as it"+
  1032. " does not exist in the db", root)
  1033. } else if err != nil {
  1034. return err
  1035. }
  1036. }
  1037. return wTx.Set(dbKeyRoot, root)
  1038. }
  1039. // Snapshot returns a read-only copy of the Tree from the given root
  1040. func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
  1041. // allow to define which root to use
  1042. if fromRoot == nil {
  1043. var err error
  1044. fromRoot, err = t.Root()
  1045. if err != nil {
  1046. return nil, err
  1047. }
  1048. }
  1049. rTx := t.db.ReadTx()
  1050. defer rTx.Discard()
  1051. // check that the root exists in the db
  1052. if !bytes.Equal(fromRoot, t.emptyHash) {
  1053. if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
  1054. return nil,
  1055. fmt.Errorf("can not do a Snapshot with root %x,"+
  1056. " as it does not exist in the db", fromRoot)
  1057. } else if err != nil {
  1058. return nil, err
  1059. }
  1060. }
  1061. return &Tree{
  1062. db: t.db,
  1063. maxLevels: t.maxLevels,
  1064. snapshotRoot: fromRoot,
  1065. emptyHash: t.emptyHash,
  1066. hashFunction: t.hashFunction,
  1067. dbg: t.dbg,
  1068. }, nil
  1069. }
  1070. // Iterate iterates through the full Tree, executing the given function on each
  1071. // node of the Tree.
  1072. func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
  1073. rTx := t.db.ReadTx()
  1074. defer rTx.Discard()
  1075. return t.IterateWithTx(rTx, fromRoot, f)
  1076. }
  1077. // IterateWithTx does the same than the Iterate method, but allowing to pass
  1078. // the db.ReadTx that is used.
  1079. func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
  1080. // allow to define which root to use
  1081. if fromRoot == nil {
  1082. var err error
  1083. fromRoot, err = t.RootWithTx(rTx)
  1084. if err != nil {
  1085. return err
  1086. }
  1087. }
  1088. return t.iter(rTx, fromRoot, f)
  1089. }
  1090. // IterateWithStop does the same than Iterate, but with int for the current
  1091. // level, and a boolean parameter used by the passed function, is to indicate to
  1092. // stop iterating on the branch when the method returns 'true'.
  1093. func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
  1094. rTx := t.db.ReadTx()
  1095. defer rTx.Discard()
  1096. // allow to define which root to use
  1097. if fromRoot == nil {
  1098. var err error
  1099. fromRoot, err = t.RootWithTx(rTx)
  1100. if err != nil {
  1101. return err
  1102. }
  1103. }
  1104. return t.iterWithStop(rTx, fromRoot, 0, f)
  1105. }
  1106. // IterateWithStopWithTx does the same than the IterateWithStop method, but
  1107. // allowing to pass the db.ReadTx that is used.
  1108. func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
  1109. f func(int, []byte, []byte) bool) error {
  1110. // allow to define which root to use
  1111. if fromRoot == nil {
  1112. var err error
  1113. fromRoot, err = t.RootWithTx(rTx)
  1114. if err != nil {
  1115. return err
  1116. }
  1117. }
  1118. return t.iterWithStop(rTx, fromRoot, 0, f)
  1119. }
  1120. func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
  1121. f func(int, []byte, []byte) bool) error {
  1122. var v []byte
  1123. var err error
  1124. if bytes.Equal(k, t.emptyHash) {
  1125. v = t.emptyHash
  1126. } else {
  1127. v, err = rTx.Get(k)
  1128. if err != nil {
  1129. return err
  1130. }
  1131. }
  1132. currLevel++
  1133. switch v[0] {
  1134. case PrefixValueEmpty:
  1135. f(currLevel, k, v)
  1136. case PrefixValueLeaf:
  1137. f(currLevel, k, v)
  1138. case PrefixValueIntermediate:
  1139. stop := f(currLevel, k, v)
  1140. if stop {
  1141. return nil
  1142. }
  1143. l, r := ReadIntermediateChilds(v)
  1144. if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
  1145. return err
  1146. }
  1147. if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
  1148. return err
  1149. }
  1150. default:
  1151. return ErrInvalidValuePrefix
  1152. }
  1153. return nil
  1154. }
  1155. func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
  1156. f2 := func(currLvl int, k, v []byte) bool {
  1157. f(k, v)
  1158. return false
  1159. }
  1160. return t.iterWithStop(rTx, k, 0, f2)
  1161. }
  1162. // Dump exports all the Tree leafs in a byte array
  1163. func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
  1164. return t.dump(fromRoot, nil)
  1165. }
  1166. // DumpWriter exports all the Tree leafs writing the bytes in the given Writer
  1167. func (t *Tree) DumpWriter(fromRoot []byte, w io.Writer) error {
  1168. _, err := t.dump(fromRoot, w)
  1169. return err
  1170. }
  1171. // dump exports all the Tree leafs. If the given w is nil, it will return a
  1172. // byte array with the dump, if w contains a *bufio.Writer, it will write the
  1173. // dump in w.
  1174. // The format of the dump is the following:
  1175. // Dump length: [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
  1176. // [ 1 byte | 1 byte | S bytes | len(v) bytes ]
  1177. // [ len(k) | len(v) | key | value ]
  1178. // Where S is the size of the output of the hash function used for the Tree.
  1179. func (t *Tree) dump(fromRoot []byte, w io.Writer) ([]byte, error) {
  1180. // allow to define which root to use
  1181. if fromRoot == nil {
  1182. var err error
  1183. fromRoot, err = t.Root()
  1184. if err != nil {
  1185. return nil, err
  1186. }
  1187. }
  1188. // WARNING current encoding only supports key & values of 255 bytes each
  1189. // (due using only 1 byte for the length headers).
  1190. var b []byte
  1191. var callbackErr error
  1192. err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
  1193. if v[0] != PrefixValueLeaf {
  1194. return false
  1195. }
  1196. leafK, leafV := ReadLeafValue(v)
  1197. kv := make([]byte, 2+len(leafK)+len(leafV))
  1198. if len(leafK) > maxUint8 {
  1199. callbackErr = fmt.Errorf("len(leafK) > %v", maxUint8)
  1200. return true
  1201. }
  1202. kv[0] = byte(len(leafK))
  1203. if len(leafV) > maxUint8 {
  1204. callbackErr = fmt.Errorf("len(leafV) > %v", maxUint8)
  1205. return true
  1206. }
  1207. kv[1] = byte(len(leafV))
  1208. copy(kv[2:2+len(leafK)], leafK)
  1209. copy(kv[2+len(leafK):], leafV)
  1210. if w == nil {
  1211. b = append(b, kv...)
  1212. } else {
  1213. n, err := w.Write(kv)
  1214. if err != nil {
  1215. callbackErr = fmt.Errorf("dump: w.Write, %s", err)
  1216. return true
  1217. }
  1218. if n != len(kv) {
  1219. callbackErr = fmt.Errorf("dump: w.Write n!=len(kv), %s", err)
  1220. return true
  1221. }
  1222. }
  1223. return false
  1224. })
  1225. if callbackErr != nil {
  1226. return nil, callbackErr
  1227. }
  1228. return b, err
  1229. }
  1230. // ImportDump imports the leafs (that have been exported with the Dump method)
  1231. // in the Tree, reading them from the given byte array.
  1232. func (t *Tree) ImportDump(b []byte) error {
  1233. bytesReader := bytes.NewReader(b)
  1234. r := bufio.NewReader(bytesReader)
  1235. return t.ImportDumpReader(r)
  1236. }
  1237. // ImportDumpReader imports the leafs (that have been exported with the Dump
  1238. // method) in the Tree, reading them from the given reader.
  1239. func (t *Tree) ImportDumpReader(r io.Reader) error {
  1240. if !t.editable() {
  1241. return ErrSnapshotNotEditable
  1242. }
  1243. root, err := t.Root()
  1244. if err != nil {
  1245. return err
  1246. }
  1247. if !bytes.Equal(root, t.emptyHash) {
  1248. return ErrTreeNotEmpty
  1249. }
  1250. var keys, values [][]byte
  1251. for {
  1252. l := make([]byte, 2)
  1253. _, err = io.ReadFull(r, l)
  1254. if err == io.EOF {
  1255. break
  1256. } else if err != nil {
  1257. return err
  1258. }
  1259. k := make([]byte, l[0])
  1260. _, err = io.ReadFull(r, k)
  1261. if err != nil {
  1262. return err
  1263. }
  1264. v := make([]byte, l[1])
  1265. _, err = io.ReadFull(r, v)
  1266. if err != nil {
  1267. return err
  1268. }
  1269. keys = append(keys, k)
  1270. values = append(values, v)
  1271. }
  1272. if _, err = t.AddBatch(keys, values); err != nil {
  1273. return err
  1274. }
  1275. return nil
  1276. }
  1277. // Graphviz iterates across the full tree to generate a string Graphviz
  1278. // representation of the tree and writes it to w
  1279. func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
  1280. return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
  1281. }
  1282. // GraphvizFirstNLevels iterates across the first NLevels of the tree to
  1283. // generate a string Graphviz representation of the first NLevels of the tree
  1284. // and writes it to w
  1285. func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
  1286. fmt.Fprintf(w, `digraph hierarchy {
  1287. node [fontname=Monospace,fontsize=10,shape=box]
  1288. `)
  1289. rTx := t.db.ReadTx()
  1290. defer rTx.Discard()
  1291. if fromRoot == nil {
  1292. var err error
  1293. fromRoot, err = t.RootWithTx(rTx)
  1294. if err != nil {
  1295. return err
  1296. }
  1297. }
  1298. nEmpties := 0
  1299. err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
  1300. if currLvl == untilLvl {
  1301. return true // to stop the iter from going down
  1302. }
  1303. switch v[0] {
  1304. case PrefixValueEmpty:
  1305. case PrefixValueLeaf:
  1306. fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
  1307. // key & value from the leaf
  1308. kB, vB := ReadLeafValue(v)
  1309. fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
  1310. hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
  1311. hex.EncodeToString(vB[:nChars]))
  1312. fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
  1313. hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
  1314. case PrefixValueIntermediate:
  1315. l, r := ReadIntermediateChilds(v)
  1316. lStr := hex.EncodeToString(l[:nChars])
  1317. rStr := hex.EncodeToString(r[:nChars])
  1318. eStr := ""
  1319. if bytes.Equal(l, t.emptyHash) {
  1320. lStr = fmt.Sprintf("empty%v", nEmpties)
  1321. eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
  1322. lStr)
  1323. nEmpties++
  1324. }
  1325. if bytes.Equal(r, t.emptyHash) {
  1326. rStr = fmt.Sprintf("empty%v", nEmpties)
  1327. eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
  1328. rStr)
  1329. nEmpties++
  1330. }
  1331. fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
  1332. lStr, rStr)
  1333. fmt.Fprint(w, eStr)
  1334. default:
  1335. }
  1336. return false
  1337. })
  1338. fmt.Fprintf(w, "}\n")
  1339. return err
  1340. }
  1341. // PrintGraphviz prints the output of Tree.Graphviz
  1342. func (t *Tree) PrintGraphviz(fromRoot []byte) error {
  1343. if fromRoot == nil {
  1344. var err error
  1345. fromRoot, err = t.Root()
  1346. if err != nil {
  1347. return err
  1348. }
  1349. }
  1350. return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
  1351. }
  1352. // PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
  1353. func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
  1354. if fromRoot == nil {
  1355. var err error
  1356. fromRoot, err = t.Root()
  1357. if err != nil {
  1358. return err
  1359. }
  1360. }
  1361. w := bytes.NewBufferString("")
  1362. fmt.Fprintf(w,
  1363. "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
  1364. err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
  1365. if err != nil {
  1366. fmt.Println(w)
  1367. return err
  1368. }
  1369. fmt.Fprintf(w,
  1370. "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
  1371. fmt.Println(w)
  1372. return nil
  1373. }
  1374. // TODO circom proofs
  1375. // TODO data structure for proofs (including root, key, value, siblings,
  1376. // hashFunction) + method to verify that data structure