node: WIP cleanup and port for client and state.

src/node/client.rs

@@ -1,16 +1,10 @@
 use async_std::sync::{Arc, Mutex};
-
 use incrementalmerkletree::{bridgetree::BridgeTree, Tree};
-use log::{debug, info, warn};
-use smol::Executor;
-use url::Url;
+use lazy_init::Lazy;
+use log::{debug, error, info};
 
-use super::{
-    service::GatewayClient,
-    state::{state_transition, State, StateUpdate},
-};
+use super::state::{state_transition, State};
 use crate::{
-    blockchain::{rocks::columns, Rocks, RocksColumn, Slab},
     crypto::{
         address::Address,
         coin::Coin,
@@ -20,149 +14,130 @@ use crate::{
         types::DrkTokenId,
         OwnCoin,
     },
-    tx,
-    util::serial::{Decodable, Encodable},
-    wallet::{
-        cashierdb::CashierDbPtr,
-        walletdb::{Balances, WalletPtr},
+    tx::{
+        builder::{
+            TransactionBuilder, TransactionBuilderClearInputInfo, TransactionBuilderInputInfo,
+            TransactionBuilderOutputInfo,
+        },
+        Transaction,
     },
-    zk::circuit::{MintContract, SpendContract},
+    util::serial::Encodable,
+    wallet::walletdb::{Balances, WalletPtr},
+    zk::circuit::MintContract,
     Result,
 };
 
+pub type ClientResult<T> = std::result::Result<T, ClientFailed>;
+
 #[derive(Debug, Clone, thiserror::Error)]
 pub enum ClientFailed {
-    #[error("Here is not enough value {0}")]
+    #[error("Not enough value: {0}")]
     NotEnoughValue(u64),
-    #[error("Invalid Address  {0}")]
+
+    #[error("Invalid address: {0}")]
     InvalidAddress(String),
-    #[error("Invalid Amount {0}")]
+
+    #[error("Invalid amount: {0}")]
     InvalidAmount(u64),
-    #[error("Unable to get deposit address")]
-    UnableToGetDepositAddress,
-    #[error("Unable to get withdraw address")]
-    UnableToGetWithdrawAddress,
-    #[error("Does not have cashier public key")]
-    DoesNotHaveCashierPublicKey,
-    #[error("Does not have keypair")]
-    DoesNotHaveKeypair,
-    #[error("Password is empty. Cannot create database")]
-    EmptyPassword,
-    #[error("Wallet already initialized")]
-    WalletInitialized,
-    #[error("Keypair already exists")]
-    KeyExists,
-    #[error("{0}")]
+
+    #[error("Client error: {0}")]
     ClientError(String),
+
     #[error("Verify error: {0}")]
     VerifyError(String),
-    #[error("Merkle tree already exists")]
-    TreeExists,
 }
 
-pub type ClientResult<T> = std::result::Result<T, ClientFailed>;
-
 impl From<crate::error::Error> for ClientFailed {
-    fn from(err: crate::error::Error) -> ClientFailed {
-        ClientFailed::ClientError(err.to_string())
+    fn from(err: crate::error::Error) -> Self {
+        Self::ClientError(err.to_string())
     }
 }
 
 impl From<super::state::VerifyFailed> for ClientFailed {
-    fn from(err: super::state::VerifyFailed) -> ClientFailed {
-        ClientFailed::VerifyError(err.to_string())
+    fn from(err: super::state::VerifyFailed) -> Self {
+        Self::VerifyError(err.to_string())
     }
 }
 
+/// The Client structure, used for transaction operations.
+/// This includes, receiving, broadcasting, and building.
 pub struct Client {
     pub main_keypair: Keypair,
-    gateway: GatewayClient,
     wallet: WalletPtr,
-    mint_pk: ProvingKey,
-    spend_pk: ProvingKey,
+    mint_pk: Lazy<ProvingKey>,
+    burn_pk: Lazy<ProvingKey>,
 }
 
 impl Client {
-    pub async fn new(
-        rocks: Arc<Rocks>,
-        gateway_addrs: (Url, Url),
-        wallet: WalletPtr,
-    ) -> Result<Self> {
+    pub async fn new(wallet: WalletPtr) -> Result<Self> {
+        // Initialize or load the wallet
         wallet.init_db().await?;
 
-        // Check if there is a default keypair
+        // Check if there is a default keypair and generate one in
+        // case we don't have any.
         if wallet.get_default_keypair().await.is_err() {
-            // Generate a new keypair if we don't have any.
+            // TODO: Clean this up with Option<T> to have less calls.
             if wallet.get_keypairs().await?.is_empty() {
                 wallet.keygen().await?;
             }
-            // set the first keypair as the default one
-            wallet.set_default_keypair(&wallet.get_keypairs().await?[0].public).await?;
         }
 
+        wallet.set_default_keypair(&wallet.get_keypairs().await?[0].public).await?;
+
         // Generate merkle tree if we don't have one.
+        // TODO: See what to do about this
         if wallet.get_tree().await.is_err() {
             wallet.tree_gen().await?;
         }
 
         let main_keypair = wallet.get_default_keypair().await?;
-        info!("Main keypair: {}", Address::from(main_keypair.public).to_string());
-
-        debug!("Creating GatewayClient");
-        let slabstore = RocksColumn::<columns::Slabs>::new(rocks);
-        let gateway = GatewayClient::new(gateway_addrs.0, gateway_addrs.1, slabstore)?;
-
-        // TODO: These should go to a better place.
-        debug!("Building proving key for the mint contract...");
-        let mint_pk = ProvingKey::build(11, &MintContract::default());
-        debug!("Building proving key for the spend contract...");
-        let spend_pk = ProvingKey::build(11, &SpendContract::default());
-
-        let client = Client { main_keypair, gateway, wallet, mint_pk, spend_pk };
-        Ok(client)
-    }
-
-    pub async fn start(&mut self) -> Result<()> {
-        self.gateway.start().await
+        info!(target: "client", "Main keypair: {}", Address::from(main_keypair.public));
+
+        Ok(Self {
+            //main_keypair: Mutex::new(main_keypair),
+            main_keypair,
+            wallet,
+            mint_pk: Lazy::new(),
+            burn_pk: Lazy::new(),
+        })
     }
 
     async fn build_slab_from_tx(
-        &mut self,
+        &self,
         pubkey: PublicKey,
         value: u64,
         token_id: DrkTokenId,
         clear_input: bool,
         state: Arc<Mutex<State>>,
     ) -> ClientResult<Vec<Coin>> {
-        debug!("Begin building slab from tx");
-        let mut clear_inputs: Vec<tx::TransactionBuilderClearInputInfo> = vec![];
-        let mut inputs: Vec<tx::TransactionBuilderInputInfo> = vec![];
-        let mut outputs: Vec<tx::TransactionBuilderOutputInfo> = vec![];
-        let mut coins: Vec<Coin> = vec![];
+        debug!("build_slab_from_tx(): Begin building slab from tx");
+        let mut clear_inputs = vec![];
+        let mut inputs = vec![];
+        let mut outputs = vec![];
+        let mut coins = vec![];
 
         if clear_input {
-            // TODO: FIXME:
+            debug!("build_slab_from_tx(): Building clear input");
             let signature_secret = self.main_keypair.secret;
-            let input = tx::TransactionBuilderClearInputInfo { value, token_id, signature_secret };
+            let input = TransactionBuilderClearInputInfo { value, token_id, signature_secret };
             clear_inputs.push(input);
         } else {
-            debug!("Start building tx inputs");
-            let mut inputs_value = 0_u64;
+            debug!("build_slab_from_tx(): Building tx inputs");
+            let mut inputs_value = 0;
             let state_m = state.lock().await;
             let own_coins = self.wallet.get_own_coins().await?;
 
             for own_coin in own_coins.iter() {
                 if inputs_value >= value {
+                    debug!("build_slab_from_tx(): inputs_value >= value");
                     break
                 }
 
-                let node = MerkleNode(own_coin.coin.0);
-                let (leaf_position, merkle_path) = state_m.tree.authentication_path(&node).unwrap();
-                // TODO: What is this counting? Is it everything or does it know to separate
-                // different tokens?
+                let leaf_position = own_coin.leaf_position;
+                let merkle_path = state_m.tree.authentication_path(leaf_position).unwrap();
                 inputs_value += own_coin.note.value;
 
-                let input = tx::TransactionBuilderInputInfo {
+                let input = TransactionBuilderInputInfo {
                     leaf_position,
                     merkle_path,
                     secret: own_coin.secret,
@@ -172,56 +147,60 @@ impl Client {
                 inputs.push(input);
                 coins.push(own_coin.coin);
             }
+            // Release state lock
+            drop(state_m);
 
             if inputs_value < value {
+                error!("build_slab_from_tx(): Not enough value to build tx inputs");
                 return Err(ClientFailed::NotEnoughValue(inputs_value))
             }
 
             if inputs_value > value {
-                let return_value: u64 = inputs_value - value;
-
-                outputs.push(tx::TransactionBuilderOutputInfo {
+                let return_value = inputs_value - value;
+                outputs.push(TransactionBuilderOutputInfo {
                     value: return_value,
                     token_id,
                     public: self.main_keypair.public,
                 });
             }
 
-            debug!("Finish building inputs");
+            debug!("build_slab_from_tx(): Finished building inputs");
         }
 
-        outputs.push(tx::TransactionBuilderOutputInfo { value, token_id, public: pubkey });
-
-        let builder = tx::TransactionBuilder { clear_inputs, inputs, outputs };
-
+        outputs.push(TransactionBuilderOutputInfo { value, token_id, public: pubkey });
+        let builder = TransactionBuilder { clear_inputs, inputs, outputs };
         let mut tx_data = vec![];
 
-        let tx: tx::Transaction = builder.build(&self.mint_pk, &self.spend_pk)?;
-        tx.encode(&mut tx_data).expect("encode tx");
-
-        let slab = Slab::new(tx_data);
-        debug!("Finish building slab from tx");
+        let mint_pk = self.mint_pk.get_or_create(Client::build_mint_pk);
+        let burn_pk = self.burn_pk.get_or_create(Client::build_burn_pk);
+        let tx = builder.build(mint_pk, burn_pk)?;
+        tx.encode(&mut tx_data)?;
 
-        // Check if it's valid before sending to gateway
+        // Check if state transition is valid before broadcasting
+        debug!("build_slab_from_tx(): Checking if state transition is valid");
         let state = &*state.lock().await;
+        debug!("build_slab_from_tx(): Got state lock");
         state_transition(state, tx)?;
+        debug!("build_slab_from_tx(): Successful state transition");
+
+        debug!("build_slab_from_tx(): Broadcasting transaction");
+        // TODO: Send to some channel, let's not p2p here
+        //self.p2p.broadcast(Tx(Transaction)).await?;
+        debug!("build_slab_from_tx(): Broadcasted successfully");
 
-        debug!("Sending slab to gateway");
-        self.gateway.put_slab(slab).await?;
-        debug!("Slab sent to gateway successfully");
         Ok(coins)
     }
 
     pub async fn send(
-        &mut self,
+        &self,
         pubkey: PublicKey,
         amount: u64,
         token_id: DrkTokenId,
         clear_input: bool,
         state: Arc<Mutex<State>>,
     ) -> ClientResult<()> {
-        // TODO: TOKEN debug
-        debug!("Sending {}", amount);
+        // TODO: Token id debug
+        debug!("send(): Sending {}", amount);
 
         if amount == 0 {
             return Err(ClientFailed::InvalidAmount(0))
@@ -229,147 +208,57 @@ impl Client {
 
         let coins = self.build_slab_from_tx(pubkey, amount, token_id, clear_input, state).await?;
         for coin in coins.iter() {
+            // TODO: This should be more robust. In case our transaction is denied,
+            // we want to revert to be able to send again.
             self.wallet.confirm_spend_coin(coin).await?;
         }
 
-        debug!("Sent {}", amount);
+        debug!("send(): Sent {}", amount);
         Ok(())
     }
 
     pub async fn transfer(
-        &mut self,
+        &self,
         token_id: DrkTokenId,
         pubkey: PublicKey,
         amount: u64,
         state: Arc<Mutex<State>>,
     ) -> ClientResult<()> {
-        debug!("Start transfer {}", amount);
-        let token_id_exists = self.wallet.token_id_exists(token_id).await?;
-
-        if token_id_exists {
+        debug!("transfer(): Start transfer {}", amount);
+        if self.wallet.token_id_exists(token_id).await? {
             self.send(pubkey, amount, token_id, false, state).await?;
-        } else {
-            return Err(ClientFailed::NotEnoughValue(amount))
+            debug!("transfer(): Finish transfer {}", amount);
+            return Ok(())
         }
 
-        debug!("Finish transfer {}", amount);
-        Ok(())
+        Err(ClientFailed::NotEnoughValue(amount))
     }
 
+    // TODO: Should this function run on finalized blocks and iterate over its transactions?
     async fn update_state(
         secret_keys: Vec<SecretKey>,
-        slab: &Slab,
+        tx: Transaction,
         state: Arc<Mutex<State>>,
         wallet: WalletPtr,
         notify: Option<async_channel::Sender<(PublicKey, u64)>>,
     ) -> Result<()> {
-        debug!("Building tx from slab and updating the state");
-        let payload = slab.get_payload();
-        /*
-        use std::io::Write;
-        let mut file = std::fs::File::create("/tmp/payload.txt")?;
-        file.write_all(&payload)?;
-        */
-        debug!("Decoding payload");
-        let tx = tx::Transaction::decode(&payload[..])?;
-
-        let update: StateUpdate;
-
-        // This is separate because otherwise the mutex is never unlocked.
+        debug!("update_state(): Begin state update");
+        debug!("update_state(): Acquiring state lock");
+        let update;
         {
-            debug!("Acquiring state lock");
             let state = &*state.lock().await;
             update = state_transition(state, tx)?;
-            debug!("Successfully passed state_transition");
         }
 
-        debug!("Acquiring state lock");
+        debug!("update_state(): Trying to apply the new state");
         let mut state = state.lock().await;
-        debug!("Trying to apply the new state");
         state.apply(update, secret_keys, notify, wallet).await?;
-        debug!("Successfully passed state.apply");
+        drop(state);
+        debug!("update_state(): Successfully updated state");
 
         Ok(())
     }
 
-    pub async fn connect_to_subscriber_from_cashier(
-        &self,
-        state: Arc<Mutex<State>>,
-        cashier_wallet: CashierDbPtr,
-        notify: async_channel::Sender<(PublicKey, u64)>,
-        executor: Arc<Executor<'_>>,
-    ) -> Result<()> {
-        debug!("Start subscriber for cashier");
-        let gateway_slabs_sub = self.gateway.start_subscriber(executor.clone()).await?;
-
-        let secret_key = self.main_keypair.secret;
-        let wallet = self.wallet.clone();
-
-        let task: smol::Task<Result<()>> = executor.spawn(async move {
-            loop {
-                let slab = gateway_slabs_sub.recv().await?;
-                debug!("Received new slab");
-
-                let mut secret_keys = vec![secret_key];
-                let mut withdraw_keys = cashier_wallet.get_withdraw_private_keys().await?;
-                secret_keys.append(&mut withdraw_keys);
-
-                let update_state = Self::update_state(
-                    secret_keys,
-                    &slab,
-                    state.clone(),
-                    wallet.clone(),
-                    Some(notify.clone()),
-                )
-                .await;
-
-                if let Err(e) = update_state {
-                    warn!("Update state: {}", e);
-                    continue
-                }
-            }
-        });
-
-        task.detach();
-        Ok(())
-    }
-
-    pub async fn connect_to_subscriber(
-        &self,
-        state: Arc<Mutex<State>>,
-        executor: Arc<Executor<'_>>,
-    ) -> Result<()> {
-        debug!("Start subscriber for darkfid");
-        let gateway_slabs_sub = self.gateway.start_subscriber(executor.clone()).await?;
-
-        let secret_key = self.main_keypair.secret;
-        let wallet = self.wallet.clone();
-
-        let task: smol::Task<Result<()>> = executor.spawn(async move {
-            loop {
-                let slab = gateway_slabs_sub.recv().await?;
-                debug!("Received new slab");
-
-                let update_state = Self::update_state(
-                    vec![secret_key],
-                    &slab,
-                    state.clone(),
-                    wallet.clone(),
-                    None,
-                )
-                .await;
-
-                if let Err(e) = update_state {
-                    warn!("Update state: {}", e);
-                    continue
-                }
-            }
-        });
-
-        task.detach();
-        Ok(())
-    }
-
     pub async fn init_db(&self) -> Result<()> {
         self.wallet.init_db().await
     }
@@ -408,4 +297,12 @@ impl Client {
     pub async fn get_tree(&self) -> Result<BridgeTree<MerkleNode, 32>> {
         self.wallet.get_tree().await
     }
+
+    fn build_mint_pk() -> ProvingKey {
+        ProvingKey::build(11, &MintContract::default())
+    }
+
+    fn build_burn_pk() -> ProvingKey {
+        ProvingKey::build(11, &MintContract::default())
+    }
 }

src/node/mod.rs

@@ -1,3 +1,2 @@
 pub mod client;
-pub mod service;
 pub mod state;

src/node/state.rs

@@ -2,7 +2,7 @@ use incrementalmerkletree::{bridgetree::BridgeTree, Frontier, Tree};
 use log::{debug, error};
 
 use crate::{
-    blockchain::{rocks::columns, RocksColumn},
+    blockchain2::{nfstore::NullifierStore, rootstore::RootStore},
     crypto::{
         coin::Coin,
         keypair::{PublicKey, SecretKey},
@@ -12,103 +12,125 @@ use crate::{
         proof::VerifyingKey,
         OwnCoin,
     },
-    error,
     tx::Transaction,
     wallet::walletdb::WalletPtr,
     Result,
 };
 
-pub trait ProgramState {
-    fn is_valid_cashier_public_key(&self, public: &PublicKey) -> bool;
-    fn is_valid_merkle(&self, merkle: &MerkleNode) -> bool;
-    fn nullifier_exists(&self, nullifier: &Nullifier) -> bool;
-    fn mint_vk(&self) -> &VerifyingKey;
-    fn spend_vk(&self) -> &VerifyingKey;
-}
-
-pub struct StateUpdate {
-    pub nullifiers: Vec<Nullifier>,
-    pub coins: Vec<Coin>,
-    pub enc_notes: Vec<EncryptedNote>,
-}
-
 pub type VerifyResult<T> = std::result::Result<T, VerifyFailed>;
 
 #[derive(Debug, Clone, thiserror::Error)]
 pub enum VerifyFailed {
-    #[error("Invalid cashier public key for clear input {0}")]
-    InvalidCashierKey(usize),
+    #[error("Invalid cashier/faucet public key for clear input {0}")]
+    InvalidCashierOrFaucetKey(usize),
+
     #[error("Invalid merkle root for input {0}")]
     InvalidMerkle(usize),
-    #[error("Duplicate nullifier for input {0}")]
-    DuplicateNullifier(usize),
-    #[error("Spend proof for input {0}")]
-    SpendProof(usize),
-    #[error("Mint proof for input {0}")]
-    MintProof(usize),
-    #[error("Invalid signature for clear input {0}")]
-    ClearInputSignature(usize),
+
     #[error("Invalid signature for input {0}")]
     InputSignature(usize),
-    #[error("Money in does not match money out (value commits)")]
-    MissingFunds,
-    #[error("Assets don't match some inputs or outputs (token commits)")]
+
+    #[error("Invalid signature for clear input {0}")]
+    ClearInputSignature(usize),
+
+    #[error("Token commitments in inputs or outputs do not match")]
     AssetMismatch,
-    #[error("Inetrnal error: {0}")]
+
+    #[error("Money in does not match money out (value commitments")]
+    MissingFunds,
+
+    #[error("Mint proof verification failure for input {0}")]
+    MintProof(usize),
+
+    #[error("Burn proof verification failure for input {0}")]
+    BurnProof(usize),
+
+    #[error("Internal error: {0}")]
     InternalError(String),
 }
 
-impl From<error::Error> for VerifyFailed {
-    fn from(err: error::Error) -> Self {
-        VerifyFailed::InternalError(err.to_string())
+impl From<crate::error::Error> for VerifyFailed {
+    fn from(err: crate::error::Error) -> Self {
+        Self::InternalError(err.to_string())
     }
 }
 
-pub fn state_transition<S: ProgramState>(state: &S, tx: Transaction) -> VerifyResult<StateUpdate> {
-    // Check deposits are legit
-    debug!(target: "STATE TRANSITION", "iterate clear_inputs");
+/// Trait implementing the state functions used by the state transition.
+pub trait ProgramState {
+    /// Check if the public key is coming from a trusted cashier
+    fn is_valid_cashier_public_key(&self, public: &PublicKey) -> bool;
+    /// Check if the public key is coming from a trusted faucet
+    fn is_valid_faucet_public_key(&self, public: &PublicKey) -> bool;
+    /// Check if a merkle root is valid in this context
+    fn is_valid_merkle(&self, merkle: &MerkleNode) -> bool;
+    /// Check if the nullifier has been seen already
+    fn nullifier_exists(&self, nullifier: &Nullifier) -> bool;
+    /// Mint proof verification key
+    fn mint_vk(&self) -> &VerifyingKey;
+    /// Burn proof verification key
+    fn burn_vk(&self) -> &VerifyingKey;
+}
+
+/// A struct representing a state update.
+/// This gets applied on top of an existing state.
+pub struct StateUpdate {
+    /// All nullifiers in a transaction
+    pub nullifiers: Vec<Nullifier>,
+    /// All coins in a transaction
+    pub coins: Vec<Coin>,
+    /// All encrypted notes in a transaction
+    pub enc_notes: Vec<EncryptedNote>,
+}
 
+/// State transition function
+pub fn state_transition<S: ProgramState>(state: &S, tx: Transaction) -> VerifyResult<StateUpdate> {
+    // Check the public keys in the clear inputs to see if they're coming
+    // from a valid cashier or faucet.
+    debug!(target: "state_transition", "Iterate clear_inputs");
     for (i, input) in tx.clear_inputs.iter().enumerate() {
-        // Check the public key in the clear inputs
-        // It should be a valid public key for the cashier
-        if !state.is_valid_cashier_public_key(&input.signature_public) {
-            error!(target: "STATE TRANSITION", "Invalid cashier public key");
-            return Err(VerifyFailed::InvalidCashierKey(i))
+        let pk = &input.signature_public;
+        if !state.is_valid_cashier_public_key(pk) || !state.is_valid_faucet_public_key(pk) {
+            error!(target: "state_transition", "Invalid pubkey for clear input: {:?}", pk);
+            return Err(VerifyFailed::InvalidCashierOrFaucetKey(i))
         }
     }
 
-    debug!(target: "STATE TRANSITION", "iterate inputs");
-
+    debug!(target: "state_transition", "Iterate inputs");
     for (i, input) in tx.inputs.iter().enumerate() {
         let merkle = &input.revealed.merkle_root;
 
-        // Merkle is used to know whether this is a coin that existed
-        // in a previous state.
+        // The Merkle root is used to know whether this is a coin that
+        // existed in a previous state.
         if !state.is_valid_merkle(merkle) {
+            error!(target: "state_transition", "Invalid Merkle root (input {})", i);
+            debug!(target: "state_transition", "root: {:?}", merkle);
             return Err(VerifyFailed::InvalidMerkle(i))
         }
 
-        // The nullifiers should not already exist
-        // It is double spend protection.
+        // The nullifiers should not already exist.
+        // It is the double-spend protection.
         let nullifier = &input.revealed.nullifier;
-
         if state.nullifier_exists(nullifier) {
-            return Err(VerifyFailed::DuplicateNullifier(i))
+            error!(target: "state_transition", "Duplicate nullifier found (input {})", i);
+            debug!(target: "state_transition", "nullifier: {:?}", nullifier);
         }
     }
 
-    debug!(target: "STATE TRANSITION", "Check the tx verifies correctly");
-    tx.verify(state.mint_vk(), state.spend_vk())?;
-    debug!(target: "STATE TRANSITION", "Verified successfully");
+    debug!(target: "state_transition", "Verifying zk proofs");
+    match tx.verify(state.mint_vk(), state.burn_vk()) {
+        Ok(()) => debug!(target: "state_transition", "Verified successfully"),
+        Err(e) => error!(target: "state_transition", "Failed verifying zk proofs: {}", e),
+    }
 
-    let mut nullifiers = vec![];
+    // Gather all the nullifiers
+    let mut nullifiers = Vec::with_capacity(tx.inputs.len());
     for input in tx.inputs {
         nullifiers.push(input.revealed.nullifier);
     }
 
-    // Newly created coins for this tx
-    let mut coins = vec![];
-    let mut enc_notes = vec![];
+    // Newly created coins for this transaction
+    let mut coins = Vec::with_capacity(tx.outputs.len());
+    let mut enc_notes = Vec::with_capacity(tx.outputs.len());
     for output in tx.outputs {
         // Gather all the coins
         coins.push(output.revealed.coin);
@@ -118,23 +140,27 @@ pub fn state_transition<S: ProgramState>(state: &S, tx: Transaction) -> VerifyRe
     Ok(StateUpdate { nullifiers, coins, enc_notes })
 }
 
+/// Struct holding the state which we can apply a [`StateUpdate`] onto.
 pub struct State {
     /// The entire Merkle tree state
     pub tree: BridgeTree<MerkleNode, 32>,
     /// List of all previous and the current merkle roots.
     /// This is the hashed value of all the children.
-    pub merkle_roots: RocksColumn<columns::MerkleRoots>,
+    pub merkle_roots: RootStore,
     /// Nullifiers prevent double-spending
-    pub nullifiers: RocksColumn<columns::Nullifiers>,
+    pub nullifiers: NullifierStore,
     /// List of Cashier public keys
-    pub public_keys: Vec<PublicKey>,
-    /// Verifying key for the Mint contract
+    pub cashier_pubkeys: Vec<PublicKey>,
+    /// List of Faucet public keys
+    pub faucet_pubkeys: Vec<PublicKey>,
+    /// Verifying key for the Mint ZK proof
     pub mint_vk: VerifyingKey,
-    /// Verifying key for the Spend contract
-    pub spend_vk: VerifyingKey,
+    /// Verifying key for the Burn ZK proof
+    pub burn_vk: VerifyingKey,
 }
 
 impl State {
+    /// Apply a [`StateUpdate`] to some state.
     pub async fn apply(
         &mut self,
         update: StateUpdate,
@@ -142,44 +168,41 @@ impl State {
         notify: Option<async_channel::Sender<(PublicKey, u64)>>,
         wallet: WalletPtr,
     ) -> Result<()> {
-        // Extend our list of nullifiers with the ones from the update.
-        debug!("Extend nullifiers");
-        for nullifier in update.nullifiers {
-            self.nullifiers.put(nullifier, vec![] as Vec<u8>)?;
-        }
+        debug!(target: "state_apply", "Extend nullifier set");
+        self.nullifiers.insert(&update.nullifiers)?;
 
-        debug!("Update Merkle tree and witness");
+        debug!(target: "state_apply", "Update Merkle tree and witnesses");
         for (coin, enc_note) in update.coins.into_iter().zip(update.enc_notes.iter()) {
             // Add the new coins to the Merkle tree
             let node = MerkleNode(coin.0);
             self.tree.append(&node);
 
             // Keep track of all Merkle roots that have existed
-            self.merkle_roots.put(self.tree.root(), vec![] as Vec<u8>)?;
+            self.merkle_roots.insert(&[self.tree.root()])?;
 
             for secret in secret_keys.iter() {
                 if let Some(note) = State::try_decrypt_note(enc_note, *secret) {
-                    self.tree.witness();
+                    debug!(target: "state_apply", "Received a coin: amount {}", note.value);
+                    let leaf_position = self.tree.witness().unwrap();
                     let nullifier = Nullifier::new(*secret, note.serial);
+                    let own_coin =
+                        OwnCoin { coin, note, secret: *secret, nullifier, leaf_position };
 
-                    let own_coin = OwnCoin { coin, note, secret: *secret, nullifier };
-
-                    wallet.put_own_coins(own_coin).await?;
+                    wallet.put_own_coin(own_coin).await?;
 
                     let pubkey = PublicKey::from_secret(*secret);
-
-                    debug!("Received a coin: amount {}", note.value);
-                    debug!("Send a notification");
+                    debug!(target: "state_apply", "Send a notification");
                     if let Some(ch) = notify.clone() {
                         ch.send((pubkey, note.value)).await?;
                     }
                 }
             }
-            // Save updated merkle tree into wallet.
+
+            // Save updated merkle tree into the wallet.
             wallet.put_tree(&self.tree).await?;
         }
 
-        debug!("apply() exiting successfully");
+        debug!(target: "state_apply", "Finished apply() successfully.");
         Ok(())
     }
 
@@ -193,23 +216,30 @@ impl State {
 
 impl ProgramState for State {
     fn is_valid_cashier_public_key(&self, public: &PublicKey) -> bool {
-        debug!("Check if it is a valid cashier public key");
-        self.public_keys.contains(public)
+        debug!(target: "state_transition", "Checking if pubkey is a valid cashier");
+        self.cashier_pubkeys.contains(public)
+    }
+
+    fn is_valid_faucet_public_key(&self, public: &PublicKey) -> bool {
+        debug!(target: "state_transition", "Checking if pubkey is a valid faucet");
+        self.faucet_pubkeys.contains(public)
     }
 
     fn is_valid_merkle(&self, merkle_root: &MerkleNode) -> bool {
-        debug!("Check if it is valid merkle");
-        if let Ok(mr) = self.merkle_roots.key_exist(*merkle_root) {
+        debug!(target: "state_transition", "Checking if Merkle root is valid");
+        if let Ok(mr) = self.merkle_roots.contains(merkle_root) {
             return mr
         }
+        // FIXME: An error here means a db issue
         false
     }
 
     fn nullifier_exists(&self, nullifier: &Nullifier) -> bool {
-        debug!("Check if nullifier exists");
-        if let Ok(nl) = self.nullifiers.key_exist(nullifier.to_bytes()) {
-            return nl
+        debug!(target: "state_transition", "Checking if Nullifier exists");
+        if let Ok(nf) = self.nullifiers.contains(nullifier) {
+            return nf
         }
+        // FIXME: An error here means a db issue
         false
     }
 
@@ -217,7 +247,7 @@ impl ProgramState for State {
         &self.mint_vk
     }
 
-    fn spend_vk(&self) -> &VerifyingKey {
-        &self.spend_vk
+    fn burn_vk(&self) -> &VerifyingKey {
+        &self.burn_vk
     }
 }