1package batchbuilder
 2
 3import (
 4	"math/big"
 5
 6	"github.com/hermeznetwork/hermez-node/common"
 7	"github.com/iden3/go-merkletree/db"
 8)
 9
10// TODO next iteration move the methods of this file into StateDB, which Synchronizer will use in the disk DB, and BatchBuilder will use with the MemoryDB
11
12// GetBalance returns the balance for a given Idx from the DB
13func (bb *BatchBuilder) GetBalance(tx db.Tx, idx common.Idx) (*common.Leaf, error) {
14	idxBytes := idx.Bytes()
15	vBytes, err := tx.Get(idxBytes[:])
16	if err != nil {
17		return nil, err
18	}
19	var b [32 * common.NLEAFELEMS]byte
20	copy(b[:], vBytes)
21	leaf, err := common.LeafFromBytes(b)
22	if err != nil {
23		return nil, err
24	}
25	return leaf, nil
26}
27
28// CreateBalance stores the Leaf into the Idx position in the MerkleTree, also adds db entry for the Leaf value
29func (bb *BatchBuilder) CreateBalance(tx db.Tx, idx common.Idx, leaf common.Leaf) error {
30	// store at the DB the key: v, and value: leaf.Bytes()
31	v, err := leaf.HashValue()
32	if err != nil {
33		return err
34	}
35	leafBytes, err := leaf.Bytes()
36	if err != nil {
37		return err
38	}
39
40	// store the Leaf value
41	tx.Put(v.Bytes(), leafBytes[:])
42	// Add k & v into the MT
43	err = bb.mt.Add(idx.BigInt(), v)
44	if err != nil {
45		return err
46	}
47
48	return nil
49}
50
51// UpdateBalance updates the balance of the leaf of a given Idx.
52// If sending==true: will substract the amount, if sending==false will add the ammount
53func (bb *BatchBuilder) UpdateBalance(tx db.Tx, idx common.Idx, amount *big.Int, sending bool) error {
54	leaf, err := bb.GetBalance(tx, idx)
55	if err != nil {
56		return err
57	}
58
59	// TODO add checks that the numbers are correct and there is no missing value neither impossible values
60	if sending {
61		leaf.Balance = new(big.Int).Sub(leaf.Balance, amount)
62	} else {
63		leaf.Balance = new(big.Int).Add(leaf.Balance, amount)
64	}
65
66	// store at the DB the key: v, and value: leaf.Bytes()
67	v, err := leaf.HashValue()
68	if err != nil {
69		return err
70	}
71	leafBytes, err := leaf.Bytes()
72	if err != nil {
73		return err
74	}
75
76	// store the Leaf value
77	tx.Put(v.Bytes(), leafBytes[:])
78	// Add k & v into the MT
79	err = bb.mt.Update(idx.BigInt(), v)
80	if err != nil {
81		return err
82	}
83
84	return nil
85}