Add PredicateTxGenerator

fuzzamoto-ir/src/compiler.rs

@@ -71,6 +71,8 @@ pub type ConnectionId = usize;
 pub struct CompiledMetadata {
     // Map from blockhash to (block variable index, list of transaction variable indices)
     block_tx_var_map: HashMap<bitcoin::BlockHash, (usize, usize, Vec<usize>)>,
+    // Map from txid to tx variable index
+    txo_var_map: HashMap<Txid, VariableIndex>,
     // Map from connection ids to connection variable indices.
     connection_map: HashMap<ConnectionId, VariableIndex>,
     // List of instruction indices that correspond to actions in the compiled program (does not include probe operation)
@@ -86,6 +88,7 @@ impl CompiledMetadata {
         Self {
             block_tx_var_map: HashMap::new(),
             connection_map: HashMap::new(),
+            txo_var_map: HashMap::new(),
             action_indices: Vec::new(),
             variable_indices: Vec::new(),
             instructions: 0,
@@ -102,16 +105,20 @@ impl CompiledMetadata {
             .map(|(header_var, block_var, tx_vars)| (*header_var, *block_var, tx_vars.as_slice()))
     }
 
-    // Get the list of instruction indices that correspond to actions in the compiled program
-    pub fn instruction_indices(&self) -> &[InstructionIndex] {
-        &self.action_indices
+    pub fn txo_variables(&self, txid: Txid) -> Option<&VariableIndex> {
+        self.txo_var_map.get(&txid)
     }
 
     // Get the list of instruction indices that correspond to variables in the compiled program
     pub fn variable_indices(&self) -> &[InstructionIndex] {
         &self.variable_indices
     }
 
+    // Get the list of instruction indices that correspond to actions in the compiled program
+    pub fn instruction_indices(&self) -> &[InstructionIndex] {
+        &self.action_indices
+    }
+
     pub fn connection_map(&self) -> &HashMap<ConnectionId, VariableIndex> {
         &self.connection_map
     }
@@ -1211,12 +1218,20 @@ impl Compiler {
             }
             Operation::TakeTxo | Operation::TakeCoinbaseTxo => {
                 let tx_var = self.get_input_mut::<Tx>(&instruction.inputs, 0)?;
+                let txid = tx_var.id;
                 let num_txos = tx_var.txos.len();
                 let mut txo = Txo::new();
                 if num_txos != 0 {
                     txo = tx_var.txos[tx_var.output_selector % num_txos].clone();
                     tx_var.output_selector += 1;
                 }
+                let txo_index = self.variables.len();
+
+                self.output
+                    .metadata
+                    .txo_var_map
+                    .entry(txid)
+                    .or_insert(txo_index);
 
                 self.append_variable(txo);
             }

fuzzamoto-ir/src/generators/block.rs

@@ -229,7 +229,7 @@ impl<R: RngCore> Generator<R> for TipBlockGenerator {
         &self,
         program: &crate::Program,
         rng: &mut R,
-        meta: Option<&PerTestcaseMetadata>,
+        meta: Option<&mut PerTestcaseMetadata>,
     ) -> Option<usize> {
         if let Some(meta) = meta.as_ref()
             && let Some(nth) = meta.recent_blocks.iter().max()
@@ -291,7 +291,7 @@ impl<R: RngCore> Generator<R> for ReorgBlockGenerator {
         &self,
         program: &crate::Program,
         rng: &mut R,
-        meta: Option<&PerTestcaseMetadata>,
+        meta: Option<&mut PerTestcaseMetadata>,
     ) -> Option<usize> {
         if let Some(meta) = meta.as_ref()
             && let Some(max) = meta.recent_blocks.iter().max_by_key(|i| i.defining_block.1)

fuzzamoto-ir/src/generators/block_txn.rs

@@ -123,7 +123,7 @@ impl<R: RngCore> Generator<R> for BlockTxnGenerator {
         &self,
         program: &crate::Program,
         rng: &mut R,
-        meta: Option<&PerTestcaseMetadata>,
+        meta: Option<&mut PerTestcaseMetadata>,
     ) -> Option<usize> {
         if let Some(meta) = meta
             && !meta.block_txn_request().is_empty()

fuzzamoto-ir/src/generators/mod.rs

@@ -66,7 +66,7 @@ pub trait Generator<R: RngCore> {
         &self,
         program: &Program,
         rng: &mut R,
-        _meta: Option<&PerTestcaseMetadata>,
+        _meta: Option<&mut PerTestcaseMetadata>,
     ) -> Option<usize> {
         program.get_random_instruction_index(rng, self.requested_context())
     }

fuzzamoto-ir/src/generators/tx.rs

@@ -1,5 +1,7 @@
+use std::marker::PhantomData;
+
 use crate::{
-    IndexedVariable, Operation, PerTestcaseMetadata, TaprootLeafSpec,
+    IndexedVariable, MempoolTxo, Operation, PerTestcaseMetadata, TaprootLeafSpec,
     generators::{Generator, ProgramBuilder},
 };
 use bitcoin::{
@@ -133,12 +135,23 @@ fn build_outputs<R: RngCore>(
     Ok(())
 }
 
-fn build_tx<R: RngCore>(
+fn build_tx_from_txos<R: RngCore>(
     builder: &mut ProgramBuilder,
     rng: &mut R,
     funding_txos: &[IndexedVariable],
     tx_version: u32,
     output_amounts: &[(u64, OutputType)],
+) -> Result<(IndexedVariable, Vec<IndexedVariable>), GeneratorError> {
+    let txos: Vec<usize> = funding_txos.iter().map(|txo| txo.index).collect();
+    build_tx(builder, rng, &txos, tx_version, output_amounts)
+}
+
+fn build_tx<R: RngCore>(
+    builder: &mut ProgramBuilder,
+    rng: &mut R,
+    funding_txos: &[usize],
+    tx_version: u32,
+    output_amounts: &[(u64, OutputType)],
 ) -> Result<(IndexedVariable, Vec<IndexedVariable>), GeneratorError> {
     let tx_version_var =
         builder.force_append_expect_output(vec![], Operation::LoadTxVersion(tx_version));
@@ -154,7 +167,7 @@ fn build_tx<R: RngCore>(
         let sequence_var =
             builder.force_append_expect_output(vec![], Operation::LoadSequence(0xffffffff));
         builder.force_append(
-            vec![mut_inputs_var.index, funding_txo.index, sequence_var.index],
+            vec![mut_inputs_var.index, *funding_txo, sequence_var.index],
             Operation::AddTxInput,
         );
     }
@@ -198,6 +211,135 @@ fn build_tx<R: RngCore>(
     Ok((const_tx_var, outputs))
 }
 
+/// `PredicateTxGenerator` generates transactions that spends the transactions in mempool depending on the given predicate.
+pub struct PredicateTxGenerator<F> {
+    predicate: F,
+    phantom: PhantomData<F>,
+}
+
+impl<F, R: RngCore> Generator<R> for PredicateTxGenerator<F>
+where
+    F: Fn(&MempoolTxo) -> bool,
+{
+    fn generate(
+        &self,
+        builder: &mut ProgramBuilder,
+        rng: &mut R,
+        meta: Option<&PerTestcaseMetadata>,
+    ) -> GeneratorResult {
+        if let Some(meta) = meta
+            && !meta.txo_metadata().txo_entry.is_empty()
+        {
+            let chosen = meta
+                .txo_metadata
+                .choice
+                .expect("We should've chosen a txo to spend by now");
+            let txo = meta
+                .txo_metadata
+                .txo_entry
+                .get(chosen)
+                .expect("Getting the chosen spent txo should always succeed")
+                .definition
+                .0;
+
+            let tx_version = *[1, 2, 3].choose(rng).unwrap(); // 3 = TRUC-violation
+            let amount = (
+                rng.gen_range(5000..100_000_000),
+                get_random_output_type(rng),
+            );
+            let (const_tx_var, _) = build_tx(builder, rng, &[txo], tx_version, &[amount])?;
+
+            let conn_var = builder.get_or_create_random_connection(rng);
+
+            let mut_inventory_var =
+                builder.force_append_expect_output(vec![], Operation::BeginBuildInventory);
+            builder.force_append(
+                vec![mut_inventory_var.index, const_tx_var.index],
+                Operation::AddWtxidInv,
+            );
+            let const_inventory_var = builder.force_append_expect_output(
+                vec![mut_inventory_var.index],
+                Operation::EndBuildInventory,
+            );
+
+            builder.force_append(
+                vec![conn_var.index, const_inventory_var.index],
+                Operation::SendInv,
+            );
+            builder.force_append(vec![conn_var.index, const_tx_var.index], Operation::SendTx);
+
+            Ok(())
+        } else {
+            Ok(())
+        }
+    }
+
+    fn choose_index(
+        &self,
+        program: &crate::Program,
+        rng: &mut R,
+        meta: Option<&mut PerTestcaseMetadata>,
+    ) -> Option<usize> {
+        let meta = meta?;
+        let txo_meta = meta.txo_metadata();
+        let filtered = txo_meta
+            .txo_entry
+            .iter()
+            .enumerate()
+            .filter(|(_, x)| (self.predicate)(*x))
+            .collect::<Vec<_>>();
+
+        let chosen = filtered.choose(rng);
+        if let Some((idx, txo)) = chosen {
+            let (_, inst) = txo.definition;
+            meta.txo_metadata_mut().choice = Some(*idx);
+            program.get_random_instruction_index_from(
+                rng,
+                <Self as Generator<R>>::requested_context(self),
+                inst + 1,
+            )
+        } else {
+            None
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        "PredicateTxGenerator"
+    }
+}
+
+impl<F> PredicateTxGenerator<F> {
+    pub fn new(predicate: F) -> Self {
+        Self {
+            predicate,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl PredicateTxGenerator<fn(&MempoolTxo) -> bool> {
+    pub fn double_spend() -> Self {
+        Self {
+            predicate: |x: &MempoolTxo| !x.spentby.is_empty(),
+            phantom: PhantomData,
+        }
+    }
+
+    pub fn chain_spend() -> Self {
+        Self {
+            predicate: |x: &MempoolTxo| x.depends.is_empty(),
+            phantom: PhantomData,
+        }
+    }
+
+    pub fn any() -> Self {
+        Self {
+            predicate: |_: &MempoolTxo| true,
+            phantom: PhantomData,
+        }
+    }
+}
+
 /// `SingleTxGenerator` generates instructions for a single new transaction into a program
 pub struct SingleTxGenerator;
 
@@ -225,7 +367,8 @@ impl<R: RngCore> Generator<R> for SingleTxGenerator {
             }
             amounts
         };
-        let (const_tx_var, _) = build_tx(builder, rng, &funding_txos, tx_version, &output_amounts)?;
+        let (const_tx_var, _) =
+            build_tx_from_txos(builder, rng, &funding_txos, tx_version, &output_amounts)?;
 
         if rng.gen_bool(0.5) {
             let conn_var = builder.get_or_create_random_connection(rng);
@@ -277,7 +420,7 @@ impl<R: RngCore> Generator<R> for OneParentOneChildGenerator {
             return Err(GeneratorError::MissingVariables);
         };
 
-        let (parent_tx_var, parent_output_vars) = build_tx(
+        let (parent_tx_var, parent_output_vars) = build_tx_from_txos(
             builder,
             rng,
             &funding_txos,
@@ -287,7 +430,7 @@ impl<R: RngCore> Generator<R> for OneParentOneChildGenerator {
                 (10000, OutputType::PayToAnchor),
             ],
         )?;
-        let (child_tx_var, _) = build_tx(
+        let (child_tx_var, _) = build_tx_from_txos(
             builder,
             rng,
             &[parent_output_vars.last().unwrap().clone()],
@@ -354,7 +497,7 @@ impl<R: RngCore> Generator<R> for LongChainGenerator {
         // transaction spends the output of the previous transaction
         let mut tx_vars = Vec::new();
         for i in 0..25 {
-            let (tx_var, outputs) = build_tx(
+            let (tx_var, outputs) = build_tx_from_txos(
                 builder,
                 rng,
                 &funding_txos,
@@ -423,7 +566,7 @@ impl<R: RngCore> Generator<R> for LargeTxGenerator {
         let conn_var = builder.get_or_create_random_connection(rng);
 
         for utxo in funding_txos {
-            let (tx_var, _) = build_tx(
+            let (tx_var, _) = build_tx_from_txos(
                 builder,
                 rng,
                 &[utxo.clone()],

fuzzamoto-ir/src/lib.rs

@@ -20,11 +20,11 @@ pub use minimizers::*;
 pub use mutators::*;
 pub use operation::*;
 
+use bitcoin::Txid;
 pub use fuzzamoto::taproot::*;
 use rand::{RngCore, seq::IteratorRandom};
-pub use variable::*;
-
 use std::{collections::HashMap, fmt, hash::Hash};
+pub use variable::*;
 
 /// Program represent a sequence of operations to perform on target nodes.
 #[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Hash)]
@@ -334,11 +334,39 @@ pub struct GetBlockTxn {
     pub tx_indices_variables: Vec<usize>,
 }
 
+/// The metadata holds the txid of transactions in mempool which is already spent by another tx in the mempool
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct MempoolTxo {
+    pub txid: Txid,
+    pub definition: (usize, usize),
+    pub spentby: Vec<Txid>,
+    pub depends: Vec<Txid>,
+}
+
+/// The metadata holds the `MempoolTxo` list and the next txo used to mutate.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct TxoMetadata {
+    pub txo_entry: Vec<MempoolTxo>,
+    pub choice: Option<usize>,
+}
+
+impl Default for TxoMetadata {
+    fn default() -> Self {
+        Self {
+            txo_entry: Vec::new(),
+            choice: None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub enum ProbeResult {
     GetBlockTxn {
         get_block_txn: GetBlockTxn,
     },
+    Mempool {
+        txo_entry: Vec<MempoolTxo>,
+    },
     Failure {
         /// The command that failed to be decoded
         command: String,