def test_update(self):
"""
Verify that items in the tree can be updated
"""
tree = MerkleTree(2)
tree.append(FQ.random())
tree.append(FQ.random())
proof_0_before = tree.proof(0)
proof_1_before = tree.proof(1)
root_before = tree.root
self.assertTrue(proof_0_before.verify(tree.root))
self.assertTrue(proof_1_before.verify(tree.root))
leaf_0_after = FQ.random()
tree.update(0, leaf_0_after)
root_after_0 = tree.root
proof_0_after = tree.proof(0)
self.assertTrue(proof_0_after.verify(tree.root))
self.assertNotEqual(root_before, root_after_0)
leaf_1_after = FQ.random()
tree.update(1, leaf_1_after)
root_after_1 = tree.root
proof_1_after = tree.proof(1)
self.assertTrue(proof_1_after.verify(tree.root))
self.assertNotEqual(root_before, root_after_1)
self.assertNotEqual(root_after_0, root_after_1)
class AccountManager(object):
def __init__(self, tree_size):
self._accounts = []
self._key2idx = dict()
self._tree = MerkleTree(tree_size)
def lookup_accounts(self, *args):
return [self.lookup_account(_) for _ in args]
def lookup_account(self, index):
if isinstance(index, AccountState):
assert index.index is not None
index = index.index
elif isinstance(index, Point):
index = self._key2idx[index]
return self._accounts[index]
def new_account(self, balance=0):
secret, pubkey = eddsa_random_keypair()
return secret, self.add_account(pubkey, balance)
def add_account(self, pubkey, balance=0, nonce=0):
assert isinstance(pubkey, Point)
state = AccountState(pubkey, balance, nonce)
state.index = self._tree.append(state.hash())
self._accounts.append(state)
self._key2idx[pubkey] = state.index
return state
def new_transaction(self, from_account, to_account, amount):
from_account, to_account = self.lookup_accounts(from_account, to_account)
return OnchainTransaction(from_account.index, to_account.index, amount)
def apply_transaction(self, stx):
"""
Records the state transition of the transaction being applied to the tree
"""
assert isinstance(stx, SignedTransaction)
tx = stx.tx
from_account, to_account = self.lookup_accounts(tx.from_idx, tx.to_idx)
if from_account.balance < tx.amount:
raise RuntimeError("Balance not sufficient to perform transfer")
merkle_root = self._tree.root
# Update `from` leaf, recording its state before modification
state_from = deepcopy(from_account)
from_account.nonce += 1
from_account.balance -= tx.amount
proof_before_from = self._tree.proof(tx.from_idx)
self._tree.update(tx.from_idx, from_account.hash())
# Update `to` leaf, recording its state before modification
state_to = deepcopy(to_account)
to_account.balance += tx.amount
proof_before_to = self._tree.proof(tx.to_idx)
self._tree.update(tx.to_idx, to_account.hash())
return TransactionProof(merkle_root, stx, state_from, state_to, proof_before_from, proof_before_to)
