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.

857 lines
22 KiB

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 (when using the Poseidon hash function),
  4. following the specification from
  5. https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
  6. https://eprint.iacr.org/2018/955.
  7. Also allows to define which hash function to use. So for example, when working
  8. with zkSnarks the Poseidon hash function can be used, but when not, it can be
  9. used the Blake3 hash function, which improves the computation time.
  10. */
  11. package arbo
  12. import (
  13. "bytes"
  14. "encoding/binary"
  15. "encoding/hex"
  16. "fmt"
  17. "io"
  18. "math"
  19. "sync"
  20. "sync/atomic"
  21. "time"
  22. "github.com/iden3/go-merkletree/db"
  23. )
  24. const (
  25. // PrefixValueLen defines the bytes-prefix length used for the Value
  26. // bytes representation stored in the db
  27. PrefixValueLen = 2
  28. // PrefixValueEmpty is used for the first byte of a Value to indicate
  29. // that is an Empty value
  30. PrefixValueEmpty = 0
  31. // PrefixValueLeaf is used for the first byte of a Value to indicate
  32. // that is a Leaf value
  33. PrefixValueLeaf = 1
  34. // PrefixValueIntermediate is used for the first byte of a Value to
  35. // indicate that is a Intermediate value
  36. PrefixValueIntermediate = 2
  37. )
  38. var (
  39. dbKeyRoot = []byte("root")
  40. dbKeyNLeafs = []byte("nleafs")
  41. emptyValue = []byte{0}
  42. )
  43. // Tree defines the struct that implements the MerkleTree functionalities
  44. type Tree struct {
  45. sync.RWMutex
  46. tx db.Tx
  47. db db.Storage
  48. lastAccess int64 // in unix time // TODO delete, is a feature of a upper abstraction level
  49. maxLevels int
  50. root []byte
  51. hashFunction HashFunction
  52. // TODO in the methods that use it, check if emptyHash param is len>0
  53. // (check if it has been initialized)
  54. emptyHash []byte
  55. dbg *dbgStats
  56. }
  57. // NewTree returns a new Tree, if there is a Tree still in the given storage, it
  58. // will load it.
  59. func NewTree(storage db.Storage, maxLevels int, hash HashFunction) (*Tree, error) {
  60. t := Tree{db: storage, maxLevels: maxLevels, hashFunction: hash}
  61. t.updateAccessTime()
  62. t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
  63. root, err := t.dbGet(dbKeyRoot)
  64. if err == db.ErrNotFound {
  65. // store new root 0
  66. t.tx, err = t.db.NewTx()
  67. if err != nil {
  68. return nil, err
  69. }
  70. t.root = t.emptyHash
  71. if err = t.dbPut(dbKeyRoot, t.root); err != nil {
  72. return nil, err
  73. }
  74. if err = t.setNLeafs(0); err != nil {
  75. return nil, err
  76. }
  77. if err = t.tx.Commit(); err != nil {
  78. return nil, err
  79. }
  80. return &t, err
  81. } else if err != nil {
  82. return nil, err
  83. }
  84. t.root = root
  85. return &t, nil
  86. }
  87. func (t *Tree) updateAccessTime() {
  88. atomic.StoreInt64(&t.lastAccess, time.Now().Unix())
  89. }
  90. // LastAccess returns the last access timestamp in Unixtime
  91. func (t *Tree) LastAccess() int64 {
  92. return atomic.LoadInt64(&t.lastAccess)
  93. }
  94. // Root returns the root of the Tree
  95. func (t *Tree) Root() []byte {
  96. return t.root
  97. }
  98. // HashFunction returns Tree.hashFunction
  99. func (t *Tree) HashFunction() HashFunction {
  100. return t.hashFunction
  101. }
  102. // Add inserts the key-value into the Tree. If the inputs come from a *big.Int,
  103. // is expected that are represented by a Little-Endian byte array (for circom
  104. // compatibility).
  105. func (t *Tree) Add(k, v []byte) error {
  106. t.updateAccessTime()
  107. t.Lock()
  108. defer t.Unlock()
  109. var err error
  110. t.tx, err = t.db.NewTx()
  111. if err != nil {
  112. return err
  113. }
  114. err = t.add(0, k, v) // add from level 0
  115. if err != nil {
  116. return err
  117. }
  118. // store root to db
  119. if err := t.dbPut(dbKeyRoot, t.root); err != nil {
  120. return err
  121. }
  122. // update nLeafs
  123. if err = t.incNLeafs(1); err != nil {
  124. return err
  125. }
  126. return t.tx.Commit()
  127. }
  128. func (t *Tree) add(fromLvl int, k, v []byte) error {
  129. // TODO check validity of key & value (for the Tree.HashFunction type)
  130. keyPath := make([]byte, t.hashFunction.Len())
  131. copy(keyPath[:], k)
  132. path := getPath(t.maxLevels, keyPath)
  133. // go down to the leaf
  134. var siblings [][]byte
  135. _, _, siblings, err := t.down(k, t.root, siblings, path, fromLvl, false)
  136. if err != nil {
  137. return err
  138. }
  139. leafKey, leafValue, err := t.newLeafValue(k, v)
  140. if err != nil {
  141. return err
  142. }
  143. if err := t.dbPut(leafKey, leafValue); err != nil {
  144. return err
  145. }
  146. // go up to the root
  147. if len(siblings) == 0 {
  148. t.root = leafKey
  149. return nil
  150. }
  151. root, err := t.up(leafKey, siblings, path, len(siblings)-1, fromLvl)
  152. if err != nil {
  153. return err
  154. }
  155. t.root = root
  156. return nil
  157. }
  158. // down goes down to the leaf recursively
  159. func (t *Tree) down(newKey, currKey []byte, siblings [][]byte,
  160. path []bool, currLvl int, getLeaf bool) (
  161. []byte, []byte, [][]byte, error) {
  162. if currLvl > t.maxLevels-1 {
  163. return nil, nil, nil, fmt.Errorf("max level")
  164. }
  165. var err error
  166. var currValue []byte
  167. if bytes.Equal(currKey, t.emptyHash) {
  168. // empty value
  169. return currKey, emptyValue, siblings, nil
  170. }
  171. currValue, err = t.dbGet(currKey)
  172. if err != nil {
  173. return nil, nil, nil, err
  174. }
  175. switch currValue[0] {
  176. case PrefixValueEmpty: // empty
  177. // TODO WIP WARNING should not be reached, as the 'if' above should avoid
  178. // reaching this point
  179. // return currKey, empty, siblings, nil
  180. panic("should not be reached, as the 'if' above should avoid reaching this point") // TMP
  181. case PrefixValueLeaf: // leaf
  182. if bytes.Equal(newKey, currKey) {
  183. // TODO move this error msg to const & add test that
  184. // checks that adding a repeated key this error is
  185. // returned
  186. return nil, nil, nil, fmt.Errorf("key already exists")
  187. }
  188. if !bytes.Equal(currValue, emptyValue) {
  189. if getLeaf {
  190. return currKey, currValue, siblings, nil
  191. }
  192. oldLeafKey, _ := ReadLeafValue(currValue)
  193. oldLeafKeyFull := make([]byte, t.hashFunction.Len())
  194. copy(oldLeafKeyFull[:], oldLeafKey)
  195. // if currKey is already used, go down until paths diverge
  196. oldPath := getPath(t.maxLevels, oldLeafKeyFull)
  197. siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
  198. if err != nil {
  199. return nil, nil, nil, err
  200. }
  201. }
  202. return currKey, currValue, siblings, nil
  203. case PrefixValueIntermediate: // intermediate
  204. if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
  205. return nil, nil, nil,
  206. fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
  207. PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
  208. }
  209. // collect siblings while going down
  210. if path[currLvl] {
  211. // right
  212. lChild, rChild := ReadIntermediateChilds(currValue)
  213. siblings = append(siblings, lChild)
  214. return t.down(newKey, rChild, siblings, path, currLvl+1, getLeaf)
  215. }
  216. // left
  217. lChild, rChild := ReadIntermediateChilds(currValue)
  218. siblings = append(siblings, rChild)
  219. return t.down(newKey, lChild, siblings, path, currLvl+1, getLeaf)
  220. default:
  221. return nil, nil, nil, fmt.Errorf("invalid value")
  222. }
  223. }
  224. // downVirtually is used when in a leaf already exists, and a new leaf which
  225. // shares the path until the existing leaf is being added
  226. func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
  227. newPath []bool, currLvl int) ([][]byte, error) {
  228. var err error
  229. if currLvl > t.maxLevels-1 {
  230. return nil, fmt.Errorf("max virtual level %d", currLvl)
  231. }
  232. if oldPath[currLvl] == newPath[currLvl] {
  233. siblings = append(siblings, t.emptyHash)
  234. siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
  235. if err != nil {
  236. return nil, err
  237. }
  238. return siblings, nil
  239. }
  240. // reached the divergence
  241. siblings = append(siblings, oldKey)
  242. return siblings, nil
  243. }
  244. // up goes up recursively updating the intermediate nodes
  245. func (t *Tree) up(key []byte, siblings [][]byte, path []bool, currLvl, toLvl int) ([]byte, error) {
  246. var k, v []byte
  247. var err error
  248. if path[currLvl+toLvl] {
  249. k, v, err = t.newIntermediate(siblings[currLvl], key)
  250. if err != nil {
  251. return nil, err
  252. }
  253. } else {
  254. k, v, err = t.newIntermediate(key, siblings[currLvl])
  255. if err != nil {
  256. return nil, err
  257. }
  258. }
  259. // store k-v to db
  260. if err = t.dbPut(k, v); err != nil {
  261. return nil, err
  262. }
  263. if currLvl == 0 {
  264. // reached the root
  265. return k, nil
  266. }
  267. return t.up(k, siblings, path, currLvl-1, toLvl)
  268. }
  269. func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
  270. t.dbg.incHash()
  271. return newLeafValue(t.hashFunction, k, v)
  272. }
  273. func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
  274. leafKey, err := hashFunc.Hash(k, v, []byte{1})
  275. if err != nil {
  276. return nil, nil, err
  277. }
  278. var leafValue []byte
  279. leafValue = append(leafValue, byte(1))
  280. leafValue = append(leafValue, byte(len(k)))
  281. leafValue = append(leafValue, k...)
  282. leafValue = append(leafValue, v...)
  283. return leafKey, leafValue, nil
  284. }
  285. // ReadLeafValue reads from a byte array the leaf key & value
  286. func ReadLeafValue(b []byte) ([]byte, []byte) {
  287. if len(b) < PrefixValueLen {
  288. return []byte{}, []byte{}
  289. }
  290. kLen := b[1]
  291. if len(b) < PrefixValueLen+int(kLen) {
  292. return []byte{}, []byte{}
  293. }
  294. k := b[PrefixValueLen : PrefixValueLen+kLen]
  295. v := b[PrefixValueLen+kLen:]
  296. return k, v
  297. }
  298. func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
  299. t.dbg.incHash()
  300. return newIntermediate(t.hashFunction, l, r)
  301. }
  302. func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
  303. b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
  304. b[0] = 2
  305. b[1] = byte(len(l))
  306. copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
  307. copy(b[PrefixValueLen+hashFunc.Len():], r)
  308. key, err := hashFunc.Hash(l, r)
  309. if err != nil {
  310. return nil, nil, err
  311. }
  312. return key, b, nil
  313. }
  314. // ReadIntermediateChilds reads from a byte array the two childs keys
  315. func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
  316. if len(b) < PrefixValueLen {
  317. return []byte{}, []byte{}
  318. }
  319. lLen := b[1]
  320. if len(b) < PrefixValueLen+int(lLen) {
  321. return []byte{}, []byte{}
  322. }
  323. l := b[PrefixValueLen : PrefixValueLen+lLen]
  324. r := b[PrefixValueLen+lLen:]
  325. return l, r
  326. }
  327. func getPath(numLevels int, k []byte) []bool {
  328. path := make([]bool, numLevels)
  329. for n := 0; n < numLevels; n++ {
  330. path[n] = k[n/8]&(1<<(n%8)) != 0
  331. }
  332. return path
  333. }
  334. // Update updates the value for a given existing key. If the given key does not
  335. // exist, returns an error.
  336. func (t *Tree) Update(k, v []byte) error {
  337. t.updateAccessTime()
  338. t.Lock()
  339. defer t.Unlock()
  340. var err error
  341. t.tx, err = t.db.NewTx()
  342. if err != nil {
  343. return err
  344. }
  345. keyPath := make([]byte, t.hashFunction.Len())
  346. copy(keyPath[:], k)
  347. path := getPath(t.maxLevels, keyPath)
  348. var siblings [][]byte
  349. _, valueAtBottom, siblings, err := t.down(k, t.root, siblings, path, 0, true)
  350. if err != nil {
  351. return err
  352. }
  353. oldKey, _ := ReadLeafValue(valueAtBottom)
  354. if !bytes.Equal(oldKey, k) {
  355. return fmt.Errorf("key %s does not exist", hex.EncodeToString(k))
  356. }
  357. leafKey, leafValue, err := t.newLeafValue(k, v)
  358. if err != nil {
  359. return err
  360. }
  361. if err := t.dbPut(leafKey, leafValue); err != nil {
  362. return err
  363. }
  364. // go up to the root
  365. if len(siblings) == 0 {
  366. t.root = leafKey
  367. return t.tx.Commit()
  368. }
  369. root, err := t.up(leafKey, siblings, path, len(siblings)-1, 0)
  370. if err != nil {
  371. return err
  372. }
  373. t.root = root
  374. // store root to db
  375. if err := t.dbPut(dbKeyRoot, t.root); err != nil {
  376. return err
  377. }
  378. return t.tx.Commit()
  379. }
  380. // GenProof generates a MerkleTree proof for the given key. If the key exists in
  381. // the Tree, the proof will be of existence, if the key does not exist in the
  382. // tree, the proof will be of non-existence.
  383. func (t *Tree) GenProof(k []byte) ([]byte, []byte, error) {
  384. t.updateAccessTime()
  385. keyPath := make([]byte, t.hashFunction.Len())
  386. copy(keyPath[:], k)
  387. path := getPath(t.maxLevels, keyPath)
  388. // go down to the leaf
  389. var siblings [][]byte
  390. _, value, siblings, err := t.down(k, t.root, siblings, path, 0, true)
  391. if err != nil {
  392. return nil, nil, err
  393. }
  394. leafK, leafV := ReadLeafValue(value)
  395. if !bytes.Equal(k, leafK) {
  396. fmt.Println("key not in Tree")
  397. fmt.Println(leafK)
  398. fmt.Println(leafV)
  399. // TODO proof of non-existence
  400. panic(fmt.Errorf("unimplemented"))
  401. }
  402. s := PackSiblings(t.hashFunction, siblings)
  403. return leafV, s, nil
  404. }
  405. // PackSiblings packs the siblings into a byte array.
  406. // [ 1 byte | L bytes | S * N bytes ]
  407. // [ bitmap length (L) | bitmap | N non-zero siblings ]
  408. // Where the bitmap indicates if the sibling is 0 or a value from the siblings
  409. // array. And S is the size of the output of the hash function used for the
  410. // Tree.
  411. func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
  412. var b []byte
  413. var bitmap []bool
  414. emptySibling := make([]byte, hashFunc.Len())
  415. for i := 0; i < len(siblings); i++ {
  416. if bytes.Equal(siblings[i], emptySibling) {
  417. bitmap = append(bitmap, false)
  418. } else {
  419. bitmap = append(bitmap, true)
  420. b = append(b, siblings[i]...)
  421. }
  422. }
  423. bitmapBytes := bitmapToBytes(bitmap)
  424. l := len(bitmapBytes)
  425. res := make([]byte, l+1+len(b))
  426. res[0] = byte(l) // set the bitmapBytes length
  427. copy(res[1:1+l], bitmapBytes)
  428. copy(res[1+l:], b)
  429. return res
  430. }
  431. // UnpackSiblings unpacks the siblings from a byte array.
  432. func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
  433. l := b[0]
  434. bitmapBytes := b[1 : 1+l]
  435. bitmap := bytesToBitmap(bitmapBytes)
  436. siblingsBytes := b[1+l:]
  437. iSibl := 0
  438. emptySibl := make([]byte, hashFunc.Len())
  439. var siblings [][]byte
  440. for i := 0; i < len(bitmap); i++ {
  441. if iSibl >= len(siblingsBytes) {
  442. break
  443. }
  444. if bitmap[i] {
  445. siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
  446. iSibl += hashFunc.Len()
  447. } else {
  448. siblings = append(siblings, emptySibl)
  449. }
  450. }
  451. return siblings, nil
  452. }
  453. func bitmapToBytes(bitmap []bool) []byte {
  454. bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
  455. b := make([]byte, bitmapBytesLen)
  456. for i := 0; i < len(bitmap); i++ {
  457. if bitmap[i] {
  458. b[i/8] |= 1 << (i % 8)
  459. }
  460. }
  461. return b
  462. }
  463. func bytesToBitmap(b []byte) []bool {
  464. var bitmap []bool
  465. for i := 0; i < len(b); i++ {
  466. for j := 0; j < 8; j++ {
  467. bitmap = append(bitmap, b[i]&(1<<j) > 0)
  468. }
  469. }
  470. return bitmap
  471. }
  472. // Get returns the value for a given key
  473. func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
  474. keyPath := make([]byte, t.hashFunction.Len())
  475. copy(keyPath[:], k)
  476. path := getPath(t.maxLevels, keyPath)
  477. // go down to the leaf
  478. var siblings [][]byte
  479. _, value, _, err := t.down(k, t.root, siblings, path, 0, true)
  480. if err != nil {
  481. return nil, nil, err
  482. }
  483. leafK, leafV := ReadLeafValue(value)
  484. if !bytes.Equal(k, leafK) {
  485. panic(fmt.Errorf("Tree.Get error: keys doesn't match, %s != %s",
  486. BytesToBigInt(k), BytesToBigInt(leafK)))
  487. }
  488. return leafK, leafV, nil
  489. }
  490. // CheckProof verifies the given proof. The proof verification depends on the
  491. // HashFunction passed as parameter.
  492. func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
  493. siblings, err := UnpackSiblings(hashFunc, packedSiblings)
  494. if err != nil {
  495. return false, err
  496. }
  497. keyPath := make([]byte, hashFunc.Len())
  498. copy(keyPath[:], k)
  499. key, _, err := newLeafValue(hashFunc, k, v)
  500. if err != nil {
  501. return false, err
  502. }
  503. path := getPath(len(siblings), keyPath)
  504. for i := len(siblings) - 1; i >= 0; i-- {
  505. if path[i] {
  506. key, _, err = newIntermediate(hashFunc, siblings[i], key)
  507. if err != nil {
  508. return false, err
  509. }
  510. } else {
  511. key, _, err = newIntermediate(hashFunc, key, siblings[i])
  512. if err != nil {
  513. return false, err
  514. }
  515. }
  516. }
  517. if bytes.Equal(key[:], root) {
  518. return true, nil
  519. }
  520. return false, nil
  521. }
  522. func (t *Tree) dbPut(k, v []byte) error {
  523. if t.tx == nil {
  524. return fmt.Errorf("dbPut error: no db Tx")
  525. }
  526. t.dbg.incDbPut()
  527. return t.tx.Put(k, v)
  528. }
  529. func (t *Tree) dbGet(k []byte) ([]byte, error) {
  530. // if key is empty, return empty as value
  531. if bytes.Equal(k, t.emptyHash) {
  532. return t.emptyHash, nil
  533. }
  534. t.dbg.incDbGet()
  535. v, err := t.db.Get(k)
  536. if err == nil {
  537. return v, nil
  538. }
  539. if t.tx != nil {
  540. return t.tx.Get(k)
  541. }
  542. return nil, db.ErrNotFound
  543. }
  544. // Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
  545. // after the setNLeafs call.
  546. func (t *Tree) incNLeafs(nLeafs int) error {
  547. oldNLeafs, err := t.GetNLeafs()
  548. if err != nil {
  549. return err
  550. }
  551. newNLeafs := oldNLeafs + nLeafs
  552. return t.setNLeafs(newNLeafs)
  553. }
  554. // Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
  555. // after the setNLeafs call.
  556. func (t *Tree) setNLeafs(nLeafs int) error {
  557. b := make([]byte, 8)
  558. binary.LittleEndian.PutUint64(b, uint64(nLeafs))
  559. if err := t.dbPut(dbKeyNLeafs, b); err != nil {
  560. return err
  561. }
  562. return nil
  563. }
  564. // GetNLeafs returns the number of Leafs of the Tree.
  565. func (t *Tree) GetNLeafs() (int, error) {
  566. b, err := t.dbGet(dbKeyNLeafs)
  567. if err != nil {
  568. return 0, err
  569. }
  570. nLeafs := binary.LittleEndian.Uint64(b)
  571. return int(nLeafs), nil
  572. }
  573. // Iterate iterates through the full Tree, executing the given function on each
  574. // node of the Tree.
  575. func (t *Tree) Iterate(f func([]byte, []byte)) error {
  576. // TODO allow to define which root to use
  577. t.updateAccessTime()
  578. return t.iter(t.root, f)
  579. }
  580. // IterateWithStop does the same than Iterate, but with int for the current
  581. // level, and a boolean parameter used by the passed function, is to indicate to
  582. // stop iterating on the branch when the method returns 'true'.
  583. func (t *Tree) IterateWithStop(f func(int, []byte, []byte) bool) error {
  584. t.updateAccessTime()
  585. return t.iterWithStop(t.root, 0, f)
  586. }
  587. func (t *Tree) iterWithStop(k []byte, currLevel int, f func(int, []byte, []byte) bool) error {
  588. v, err := t.dbGet(k)
  589. if err != nil {
  590. return err
  591. }
  592. currLevel++
  593. switch v[0] {
  594. case PrefixValueEmpty:
  595. f(currLevel, k, v)
  596. case PrefixValueLeaf:
  597. f(currLevel, k, v)
  598. case PrefixValueIntermediate:
  599. stop := f(currLevel, k, v)
  600. if stop {
  601. return nil
  602. }
  603. l, r := ReadIntermediateChilds(v)
  604. if err = t.iterWithStop(l, currLevel, f); err != nil {
  605. return err
  606. }
  607. if err = t.iterWithStop(r, currLevel, f); err != nil {
  608. return err
  609. }
  610. default:
  611. return fmt.Errorf("invalid value")
  612. }
  613. return nil
  614. }
  615. func (t *Tree) iter(k []byte, f func([]byte, []byte)) error {
  616. f2 := func(currLvl int, k, v []byte) bool {
  617. f(k, v)
  618. return false
  619. }
  620. return t.iterWithStop(k, 0, f2)
  621. }
  622. // Dump exports all the Tree leafs in a byte array of length:
  623. // [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
  624. // [ 1 byte | 1 byte | S bytes | len(v) bytes ]
  625. // [ len(k) | len(v) | key | value ]
  626. // Where S is the size of the output of the hash function used for the Tree.
  627. func (t *Tree) Dump() ([]byte, error) {
  628. t.updateAccessTime()
  629. // TODO allow to define which root to use
  630. // WARNING current encoding only supports key & values of 255 bytes each
  631. // (due using only 1 byte for the length headers).
  632. var b []byte
  633. err := t.Iterate(func(k, v []byte) {
  634. if v[0] != PrefixValueLeaf {
  635. return
  636. }
  637. leafK, leafV := ReadLeafValue(v)
  638. kv := make([]byte, 2+len(leafK)+len(leafV))
  639. kv[0] = byte(len(leafK))
  640. kv[1] = byte(len(leafV))
  641. copy(kv[2:2+len(leafK)], leafK)
  642. copy(kv[2+len(leafK):], leafV)
  643. b = append(b, kv...)
  644. })
  645. return b, err
  646. }
  647. // ImportDump imports the leafs (that have been exported with the ExportLeafs
  648. // method) in the Tree.
  649. func (t *Tree) ImportDump(b []byte) error {
  650. t.updateAccessTime()
  651. r := bytes.NewReader(b)
  652. var err error
  653. var keys, values [][]byte
  654. for {
  655. l := make([]byte, 2)
  656. _, err = io.ReadFull(r, l)
  657. if err == io.EOF {
  658. break
  659. } else if err != nil {
  660. return err
  661. }
  662. k := make([]byte, l[0])
  663. _, err = io.ReadFull(r, k)
  664. if err != nil {
  665. return err
  666. }
  667. v := make([]byte, l[1])
  668. _, err = io.ReadFull(r, v)
  669. if err != nil {
  670. return err
  671. }
  672. keys = append(keys, k)
  673. values = append(values, v)
  674. }
  675. if _, err = t.AddBatch(keys, values); err != nil {
  676. return err
  677. }
  678. return nil
  679. }
  680. // Graphviz iterates across the full tree to generate a string Graphviz
  681. // representation of the tree and writes it to w
  682. func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
  683. return t.GraphvizFirstNLevels(w, rootKey, t.maxLevels)
  684. }
  685. // GraphvizFirstNLevels iterates across the first NLevels of the tree to
  686. // generate a string Graphviz representation of the first NLevels of the tree
  687. // and writes it to w
  688. func (t *Tree) GraphvizFirstNLevels(w io.Writer, rootKey []byte, untilLvl int) error {
  689. fmt.Fprintf(w, `digraph hierarchy {
  690. node [fontname=Monospace,fontsize=10,shape=box]
  691. `)
  692. nChars := 4 // TODO move to global constant
  693. nEmpties := 0
  694. err := t.iterWithStop(t.root, 0, func(currLvl int, k, v []byte) bool {
  695. if currLvl == untilLvl {
  696. return true // to stop the iter from going down
  697. }
  698. switch v[0] {
  699. case PrefixValueEmpty:
  700. case PrefixValueLeaf:
  701. fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
  702. // key & value from the leaf
  703. kB, vB := ReadLeafValue(v)
  704. fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
  705. hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
  706. hex.EncodeToString(vB[:nChars]))
  707. fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
  708. hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
  709. case PrefixValueIntermediate:
  710. l, r := ReadIntermediateChilds(v)
  711. lStr := hex.EncodeToString(l[:nChars])
  712. rStr := hex.EncodeToString(r[:nChars])
  713. eStr := ""
  714. if bytes.Equal(l, t.emptyHash) {
  715. lStr = fmt.Sprintf("empty%v", nEmpties)
  716. eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
  717. lStr)
  718. nEmpties++
  719. }
  720. if bytes.Equal(r, t.emptyHash) {
  721. rStr = fmt.Sprintf("empty%v", nEmpties)
  722. eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
  723. rStr)
  724. nEmpties++
  725. }
  726. fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
  727. lStr, rStr)
  728. fmt.Fprint(w, eStr)
  729. default:
  730. }
  731. return false
  732. })
  733. fmt.Fprintf(w, "}\n")
  734. return err
  735. }
  736. // PrintGraphviz prints the output of Tree.Graphviz
  737. func (t *Tree) PrintGraphviz(rootKey []byte) error {
  738. return t.PrintGraphvizFirstNLevels(rootKey, t.maxLevels)
  739. }
  740. // PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
  741. func (t *Tree) PrintGraphvizFirstNLevels(rootKey []byte, untilLvl int) error {
  742. if rootKey == nil {
  743. rootKey = t.Root()
  744. }
  745. w := bytes.NewBufferString("")
  746. fmt.Fprintf(w,
  747. "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
  748. err := t.GraphvizFirstNLevels(w, nil, untilLvl)
  749. if err != nil {
  750. fmt.Println(w)
  751. return err
  752. }
  753. fmt.Fprintf(w,
  754. "End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
  755. fmt.Println(w)
  756. return nil
  757. }
  758. // Purge WIP: unimplemented
  759. func (t *Tree) Purge(keys [][]byte) error {
  760. return nil
  761. }