Embed the mempool into the Blockchain

coordinator/tributary/src/blockchain.rs

@@ -2,33 +2,47 @@ use std::collections::{HashSet, HashMap};
 
 use ciphersuite::{Ciphersuite, Ristretto};
 
-use crate::{Signed, TransactionKind, Transaction, ProvidedTransactions, BlockError, Block};
+use crate::{Signed, TransactionKind, Transaction, ProvidedTransactions, BlockError, Block, Mempool};
 
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Blockchain<T: Transaction> {
   genesis: [u8; 32],
   // TODO: db
   tip: [u8; 32],
-  provided: ProvidedTransactions<T>,
-  // TODO: Mempool
   next_nonces: HashMap<<Ristretto as Ciphersuite>::G, u32>,
+
+  provided: ProvidedTransactions<T>,
+  mempool: Mempool<T>,
 }
 
 impl<T: Transaction> Blockchain<T> {
   pub fn new(genesis: [u8; 32], participants: &[<Ristretto as Ciphersuite>::G]) -> Self {
-    // TODO: Reload provided/nonces
+    // TODO: Reload next_nonces/provided/mempool
 
     let mut next_nonces = HashMap::new();
     for participant in participants {
       next_nonces.insert(*participant, 0);
     }
-    Self { genesis, tip: genesis, provided: ProvidedTransactions::new(), next_nonces }
+
+    Self {
+      genesis,
+
+      tip: genesis,
+      next_nonces,
+
+      provided: ProvidedTransactions::new(),
+      mempool: Mempool::new(genesis),
+    }
   }
 
   pub fn tip(&self) -> [u8; 32] {
     self.tip
   }
 
+  pub fn add_transaction(&mut self, tx: T) -> bool {
+    self.mempool.add(&self.next_nonces, tx)
+  }
+
   pub fn provide_transaction(&mut self, tx: T) {
     self.provided.provide(tx)
   }
@@ -38,9 +52,8 @@ impl<T: Transaction> Blockchain<T> {
     self.next_nonces.get(&key).cloned()
   }
 
-  // TODO: Embed mempool
-  pub fn build_block(&self, txs: HashMap<[u8; 32], T>) -> Block<T> {
-    let block = Block::new(self.tip, &self.provided, txs);
+  pub fn build_block(&mut self) -> Block<T> {
+    let block = Block::new(self.tip, &self.provided, self.mempool.block(&self.next_nonces));
     // build_block should not return invalid blocks
     self.verify_block(&block).unwrap();
     block

coordinator/tributary/src/mempool.rs

@@ -24,13 +24,21 @@ impl<T: Transaction> Mempool<T> {
   ) -> bool {
     match tx.kind() {
       TransactionKind::Signed(Signed { signer, nonce, .. }) => {
-        // If the mempool doesn't have a nonce tracked, grab it from the blockchain
-        if !self.next_nonces.contains_key(signer) {
-          let Some(blockchain_next_nonces) = blockchain_next_nonces.get(signer).cloned() else {
-            // Not a participant
-            return false;
-          };
-          self.next_nonces.insert(*signer, blockchain_next_nonces);
+        // Get the nonce from the blockchain
+        let Some(blockchain_next_nonce) = blockchain_next_nonces.get(signer).cloned() else {
+          // Not a participant
+          return false;
+        };
+
+        // If the blockchain's nonce is greater than the mempool's, use it
+        // Default to true so if the mempool hasn't tracked this nonce yet, it'll be inserted
+        let mut blockchain_is_greater = true;
+        if let Some(mempool_next_nonce) = self.next_nonces.get(signer) {
+          blockchain_is_greater = blockchain_next_nonce > *mempool_next_nonce;
+        }
+
+        if blockchain_is_greater {
+          self.next_nonces.insert(*signer, blockchain_next_nonce);
         }
 
         if verify_transaction(&tx, self.genesis, &mut HashSet::new(), &mut self.next_nonces)

coordinator/tributary/src/tests/blockchain.rs

@@ -8,7 +8,7 @@ use blake2::{Digest, Blake2s256};
 use ciphersuite::{group::ff::Field, Ciphersuite, Ristretto};
 
 use crate::{
-  merkle, Transaction, ProvidedTransactions, Block, Blockchain,
+  merkle, Signed, TransactionKind, Transaction, ProvidedTransactions, Block, Blockchain,
   tests::{ProvidedTransaction, SignedTransaction, random_provided_transaction},
 };
 
@@ -31,7 +31,7 @@ fn new_blockchain<T: Transaction>(
 fn block_addition() {
   let genesis = new_genesis();
   let mut blockchain = new_blockchain::<SignedTransaction>(genesis, &[]);
-  let block = blockchain.build_block(HashMap::new());
+  let block = blockchain.build_block();
   assert_eq!(block.header.parent, genesis);
   assert_eq!(block.header.transactions, [0; 32]);
   blockchain.verify_block(&block).unwrap();
@@ -42,9 +42,9 @@ fn block_addition() {
 #[test]
 fn invalid_block() {
   let genesis = new_genesis();
-  let blockchain = new_blockchain::<SignedTransaction>(genesis, &[]);
+  let mut blockchain = new_blockchain::<SignedTransaction>(genesis, &[]);
 
-  let block = blockchain.build_block(HashMap::new());
+  let block = blockchain.build_block();
 
   // Mutate parent
   {
@@ -92,7 +92,9 @@ fn invalid_block() {
 
   {
     // Add a valid transaction
-    let mut block = blockchain.build_block(HashMap::from([(tx.hash(), tx.clone())]));
+    let mut blockchain = blockchain.clone();
+    assert!(blockchain.add_transaction(tx.clone()));
+    let mut block = blockchain.build_block();
     assert_eq!(block.header.transactions, merkle(&[tx.hash()]));
     blockchain.verify_block(&block).unwrap();
 
@@ -112,7 +114,9 @@ fn invalid_block() {
 
   {
     // Invalid signature
-    let mut block = blockchain.build_block(HashMap::from([(tx.hash(), tx)]));
+    let mut blockchain = blockchain;
+    assert!(blockchain.add_transaction(tx));
+    let mut block = blockchain.build_block();
     blockchain.verify_block(&block).unwrap();
     block.transactions[0].1.signature.s += <Ristretto as Ciphersuite>::F::ONE;
     assert!(blockchain.verify_block(&block).is_err());
@@ -134,11 +138,36 @@ fn signed_transaction() {
   let mut blockchain = new_blockchain::<SignedTransaction>(genesis, &[signer]);
   assert_eq!(blockchain.next_nonce(signer), Some(0));
 
-  let test = |blockchain: &mut Blockchain<SignedTransaction>, mempool: HashMap<_, _>| {
+  let test = |blockchain: &mut Blockchain<SignedTransaction>,
+              mempool: HashMap<[u8; 32], SignedTransaction>| {
     let mut hashes = mempool.keys().cloned().collect::<HashSet<_>>();
 
+    // These transactions do need to be added, in-order, to the mempool for the blockchain to
+    // build a block off them
+    {
+      let mut ordered = HashMap::new();
+      for (_, tx) in mempool.clone().drain() {
+        let nonce = if let TransactionKind::Signed(Signed { nonce, .. }) = tx.kind() {
+          *nonce
+        } else {
+          panic!("non-signed TX in test mempool");
+        };
+        ordered.insert(nonce, tx);
+      }
+
+      let mut i = 0;
+      while !ordered.contains_key(&i) {
+        i += 1;
+      }
+      for i in i .. (i + u32::try_from(ordered.len()).unwrap()) {
+        assert!(blockchain.add_transaction(ordered.remove(&i).unwrap()));
+      }
+    }
+
     let tip = blockchain.tip();
-    let block = blockchain.build_block(mempool);
+    let block = blockchain.build_block();
+    // The Block constructor should sort these these, and build_block should've called Block::new
+    assert_eq!(block, Block::new(blockchain.tip(), &ProvidedTransactions::new(), mempool));
     assert_eq!(blockchain.tip(), tip);
     assert_eq!(block.header.parent, tip);