fix(coprocessor): host-listener, dependency chain, step 1

Author: rudy-6-4

coprocessor/fhevm-engine/host-listener/src/database/dependence_chains.rs

@@ -0,0 +1,289 @@
+use std::collections::hash_map::Entry;
+use std::collections::{HashMap, HashSet, VecDeque};
+
+use tracing::{error, warn};
+
+use crate::database::tfhe_event_propagate::{
+    tfhe_inputs_handle, tfhe_result_handle,
+};
+use crate::database::tfhe_event_propagate::{
+    ChainCache, ChainDependency, Chains, Handle, LogTfhe, TransactionHash,
+};
+
+type HandleToTransaction = HashMap<Handle, (usize, TransactionHash)>;
+type Transactions = VecDeque<TransactionHash>;
+type TransactionToChain = HashMap<TransactionHash, ChainDependency>;
+
+fn scan_transactions(logs: &[LogTfhe]) -> (HandleToTransaction, Transactions) {
+    let nb_logs = logs.len();
+    let mut prev_tx = None;
+    let mut seen_txs = HashSet::with_capacity(nb_logs);
+    let mut txs = VecDeque::with_capacity(nb_logs);
+    let mut handle_to_tx = HashMap::with_capacity(nb_logs);
+    for (age, log) in logs.iter().enumerate() {
+        let tx = log.transaction_hash.unwrap_or_default();
+        if prev_tx != Some(tx) && !seen_txs.contains(&tx) {
+            prev_tx = Some(tx);
+            seen_txs.insert(tx);
+            txs.push_front(tx);
+        }
+        if let Some(handle) = tfhe_result_handle(&log.event) {
+            handle_to_tx.insert(handle, (age, tx));
+        }
+    }
+    (handle_to_tx, txs)
+}
+
+async fn scan_transactions_dependencies(
+    logs: &[LogTfhe],
+    handle_to_tx: HandleToTransaction,
+    nb_txs: usize,
+    past_chains: &ChainCache,
+) -> (
+    HashMap<TransactionHash, (usize, TransactionHash)>,
+    TransactionToChain,
+) {
+    // tx to its direct dependences or chain
+    let mut tx_direct_tx_dependency: HashMap<
+        TransactionHash,
+        (usize, TransactionHash),
+    > = HashMap::with_capacity(nb_txs);
+    let mut tx_chain_dependency = HashMap::with_capacity(nb_txs);
+    for log in logs {
+        let tx = log.transaction_hash.unwrap_or_default();
+        let log_inputs = tfhe_inputs_handle(&log.event);
+        for input in log_inputs {
+            let same_block_dep = tx_direct_tx_dependency.entry(tx);
+            let already_one = matches!(same_block_dep, Entry::Occupied(_));
+            if let Some((age, dep_tx)) = handle_to_tx.get(&input) {
+                // block local, we keep the most recent tx only, assuming correct order
+                let should_insert = match same_block_dep {
+                    Entry::Vacant(_) => true,
+                    Entry::Occupied(ref e) => *age > e.get().0 && tx != *dep_tx,
+                };
+                if should_insert {
+                    same_block_dep.insert_entry((*age, *dep_tx));
+                }
+                continue;
+            }
+            if already_one {
+                // this tx already has a local dependency, skip checking past ones
+                continue;
+            }
+            let pre_block_dep = tx_chain_dependency
+                .entry(log.transaction_hash.unwrap_or_default());
+            if let Entry::Occupied(_) = pre_block_dep {
+                // this tx already has a chain dependency, skip checking past ones
+                continue;
+            }
+            if let Some(&dep_tx) = past_chains.write().await.get(&input) {
+                // from a previous block
+                pre_block_dep.insert_entry(dep_tx);
+            }
+            // no dependency or only unknown ones
+        }
+    }
+    (tx_direct_tx_dependency, tx_chain_dependency)
+}
+
+fn assign_tx_a_chain_dependency(
+    mut txs: Transactions,
+    tx_direct_tx_dependency: &HashMap<
+        TransactionHash,
+        (usize, TransactionHash),
+    >,
+    tx_chain_dependency: &HashMap<TransactionHash, ChainDependency>,
+) -> (TransactionToChain, Chains) {
+    let nb_txs = txs.len();
+    let mut txn_to_chain_dep = HashMap::with_capacity(nb_txs);
+    let mut chains = HashSet::with_capacity(nb_txs);
+    // tx to its chain dependency
+    while let Some(tx) = txs.pop_front() {
+        let chain_dep = txn_to_chain_dep.entry(tx);
+        if let Entry::Occupied(_) = chain_dep {
+            // already done
+            continue;
+        }
+        // in block dependency we propagate the chain
+        let chain = if let Some(direct_deps) = tx_direct_tx_dependency.get(&tx)
+        {
+            let (_age, tx_dep) = direct_deps;
+            let Some(chain) = txn_to_chain_dep.get(tx_dep) else {
+                warn!(?tx, "Out of order transactions");
+                // only happens if logs are out of order
+                txs.push_back(tx);
+                // let's do its dependency first
+                txs.push_back(*tx_dep);
+                continue;
+            };
+            *chain
+        } else if let Some(chain) = tx_chain_dependency.get(&tx) {
+            *chain
+        } else {
+            // no dependency or unknown ones
+            // createa new chain
+            tx
+        };
+        txn_to_chain_dep.insert(tx, chain);
+        chains.insert(chain);
+    }
+    (txn_to_chain_dep, chains)
+}
+
+pub async fn assign_logs_a_chain_dependency(
+    logs: &mut [LogTfhe],
+    txn_to_chain_dep: &TransactionToChain,
+    past_chains: &ChainCache,
+) {
+    let mut past_chains_write = past_chains.write().await;
+    for log in logs {
+        let tx = log.transaction_hash.unwrap_or_default();
+        let chain_dep = txn_to_chain_dep.get(&tx);
+        if chain_dep.is_none() {
+            error!(?tx, "No chain dependency found for transaction");
+        }
+        let chain_dep = *chain_dep.unwrap_or(&tx);
+        log.dependence_chain = chain_dep;
+        if log.is_allowed == false {
+            // cannot be reused in future blocks
+            continue;
+        }
+        // update past chains cache for next block
+        if let Some(handle) = tfhe_result_handle(&log.event) {
+            past_chains_write.put(handle, chain_dep);
+        }
+    }
+}
+
+pub async fn dependence_chains(
+    logs: &mut [LogTfhe],
+    past_chains: &ChainCache,
+) -> Chains {
+    // handle to transaction
+    let (handle_to_tx, txs) = scan_transactions(logs);
+
+    let nb_txs = txs.len();
+    let (tx_direct_tx_dependency, tx_chain_dependency) =
+        scan_transactions_dependencies(logs, handle_to_tx, nb_txs, past_chains)
+            .await;
+
+    let (txn_to_chain_dep, chains) = assign_tx_a_chain_dependency(
+        txs,
+        &tx_direct_tx_dependency,
+        &tx_chain_dependency,
+    );
+
+    assign_logs_a_chain_dependency(logs, &txn_to_chain_dep, past_chains).await;
+    chains
+}
+
+#[cfg(test)]
+mod tests {
+    // use super::*;
+    use alloy::rpc::types::Log;
+
+    use crate::{
+        contracts::TfheContract::TfheContractEvents as E,
+        database::tfhe_event_propagate::ClearConst,
+    };
+
+    fn push_event(
+        e: E,
+        logs: &mut Vec<Log>,
+        tx: Option<TransactionHash>,
+    ) -> Handle {
+        let Some(tx) = tx else {
+            if let Some(log) = logs.last() {
+                log.transaction_hash
+            } else {
+                TransactionHash::ZERO
+            }
+        };
+        logs.push(Log {
+            inner: e,
+            block_hash: None,
+            block_number: None,
+            block_timestamp: None,
+            transaction_hash: Some(tx),
+            transaction_index: None,
+            log_index: None,
+            removed: true,
+        })
+    }
+
+    fn new_handle() -> Handle {
+        static HANDLE_COUNTER: std::sync::atomic::AtomicU64 =
+            std::sync::atomic::AtomicU64::new(1);
+        let id =
+            HANDLE_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        Handle::from_u64(id)
+    }
+
+    fn input_handle(
+        logs: &mut Vec<Log>,
+        tx: Option<TransactionHash>,
+    ) -> Handle {
+        push_event(
+            E::TrivialEncrypt(C::TrivialEncrypt {
+                caller: alloy::types::Address::ZERO,
+                pt: ClearConst::from_u64(0),
+                toType: 0,
+                resultHandle: new_handle(),
+            }),
+            logs,
+            tx,
+        )
+    }
+
+    fn op1(
+        handle: Handle,
+        logs: &mut Vec<Log>,
+        tx: Option<TransactionHash>,
+    ) -> E {
+        push_event(
+            E::FheAdd(C::FheAdd {
+                lhs: handle,
+                rhs: handle,
+                scalarByte: 1,
+                result: new_handle(),
+                caller: alloy::types::Address::ZERO,
+            }),
+            logs,
+            tx,
+        )
+    }
+
+    fn op2(
+        handle1: Handle,
+        handle2: Handle,
+        logs: &mut Vec<Log>,
+        tx: Option<TransactionHash>,
+    ) -> E {
+        push_event(
+            E::FheAdd(C::FheAdd {
+                lhs: handle1,
+                rhs: handle2,
+                scalarByte: 1,
+                result: new_handle(),
+                caller: alloy::types::Address::ZERO,
+            }),
+            logs,
+            tx,
+        )
+    }
+
+    #[test]
+    fn test_scan_transactions_tx() {
+        let mut logs = vec![];
+        let v0 = input_handle(&mut logs, None);
+        let v1 = op(v0, &mut logs, None);
+        let (handle_to_tx, txs) = scan_transactions(&mut logs);
+        assert_eq!(txs.len(), 1);
+        assert_eq!(handle_to_tx.len(), 2);
+
+        let v2 = input_handle(&mut logs, TransactionHash::with_last_byte(1));
+        let (handle_to_tx, txs) = scan_transactions(&mut logs);
+        assert_eq!(txs.len(), 2);
+        assert_eq!(handle_to_tx.len(), 3);
+    }
+}

coprocessor/fhevm-engine/host-listener/src/database/ingest.rs

@@ -9,6 +9,7 @@ use tracing::{error, info};
 
 use crate::cmd::block_history::BlockSummary;
 use crate::contracts::{AclContract, TfheContract};
+use crate::database::dependence_chains::dependence_chains;
 use crate::database::tfhe_event_propagate::{
     acl_result_handles, tfhe_result_handle, Database, LogTfhe,
 };
@@ -100,9 +101,11 @@ pub async fn ingest_block_logs(
                 let log = LogTfhe {
                     event,
                     transaction_hash: log.transaction_hash,
-                    is_allowed: false, // updated in the next loop
                     block_number,
                     block_timestamp,
+                    // updated in the next loop and dependence_chains
+                    is_allowed: false,
+                    dependence_chain: Default::default(),
                 };
                 tfhe_event_log.push(log);
                 continue;
@@ -119,18 +122,19 @@ pub async fn ingest_block_logs(
             );
         }
     }
-
-    for tfhe_log in tfhe_event_log {
-        let is_allowed =
+    for tfhe_log in tfhe_event_log.iter_mut() {
+        tfhe_log.is_allowed =
             if let Some(result_handle) = tfhe_result_handle(&tfhe_log.event) {
                 is_allowed.contains(&result_handle.to_vec())
             } else {
                 false
             };
-        let tfhe_log = LogTfhe {
-            is_allowed,
-            ..tfhe_log
-        };
+    }
+
+    let chains =
+        dependence_chains(&mut tfhe_event_log, &db.dependence_chain).await;
+
+    for tfhe_log in tfhe_event_log {
         let inserted = db.insert_tfhe_event(&mut tx, &tfhe_log).await?;
         if block_logs.catchup && inserted {
             info!(tfhe_log = ?tfhe_log, "TFHE event missed before");
@@ -147,5 +151,6 @@ pub async fn ingest_block_logs(
     }
 
     db.mark_block_as_valid(&mut tx, &block_logs.summary).await?;
+    db.update_dependence_chain(&mut tx, chains).await?;
     tx.commit().await
 }

coprocessor/fhevm-engine/host-listener/src/database/mod.rs

@@ -1,2 +1,3 @@
+pub mod dependence_chains;
 pub mod ingest;
 pub mod tfhe_event_propagate;