add 4byte_traces (#99)

crates/freeze/src/datasets/four_byte_traces.rs

@@ -0,0 +1,104 @@
+use crate::*;
+use polars::prelude::*;
+use std::collections::BTreeMap;
+
+/// columns for transactions
+#[cryo_to_df::to_df(Datatype::FourByteTraces)]
+#[derive(Default)]
+pub struct FourByteTraces {
+    pub(crate) n_rows: u64,
+    pub(crate) block_number: Vec<Option<u32>>,
+    pub(crate) transaction_index: Vec<Option<u32>>,
+    pub(crate) transaction_hash: Vec<Option<Vec<u8>>>,
+    pub(crate) signature: Vec<Vec<u8>>,
+    pub(crate) size: Vec<u64>,
+    pub(crate) count: Vec<u64>,
+    pub(crate) chain_id: Vec<u64>,
+}
+
+#[async_trait::async_trait]
+impl Dataset for FourByteTraces {
+    fn aliases() -> Vec<&'static str> {
+        vec!["4byte_traces"]
+    }
+}
+
+type BlockTxsTraces = (Option<u32>, Vec<Option<Vec<u8>>>, Vec<BTreeMap<String, u64>>);
+
+#[async_trait::async_trait]
+impl CollectByBlock for FourByteTraces {
+    type Response = BlockTxsTraces;
+
+    async fn extract(request: Params, source: Arc<Source>, query: Arc<Query>) -> R<Self::Response> {
+        let schema =
+            query.schemas.get(&Datatype::FourByteTraces).ok_or(err("schema not provided"))?;
+        let include_txs = schema.has_column("transaction_hash");
+        source
+            .fetcher
+            .geth_debug_trace_block_4byte_traces(request.block_number()? as u32, include_txs)
+            .await
+    }
+
+    fn transform(response: Self::Response, columns: &mut Self, query: &Arc<Query>) -> R<()> {
+        process_storage_reads(&response, columns, &query.schemas)
+    }
+}
+
+#[async_trait::async_trait]
+impl CollectByTransaction for FourByteTraces {
+    type Response = BlockTxsTraces;
+
+    async fn extract(request: Params, source: Arc<Source>, query: Arc<Query>) -> R<Self::Response> {
+        let schema =
+            query.schemas.get(&Datatype::FourByteTraces).ok_or(err("schema not provided"))?;
+        let include_block_number = schema.has_column("block_number");
+        let tx = request.transaction_hash()?;
+        source.fetcher.geth_debug_trace_transaction_4byte_traces(tx, include_block_number).await
+    }
+
+    fn transform(response: Self::Response, columns: &mut Self, query: &Arc<Query>) -> R<()> {
+        process_storage_reads(&response, columns, &query.schemas)
+    }
+}
+
+pub(crate) fn process_storage_reads(
+    response: &BlockTxsTraces,
+    columns: &mut FourByteTraces,
+    schemas: &Schemas,
+) -> R<()> {
+    let schema = schemas.get(&Datatype::FourByteTraces).ok_or(err("schema not provided"))?;
+    let (block_number, txs, traces) = response;
+    for (index, (trace, tx)) in traces.iter().zip(txs).enumerate() {
+        for (signature_size, count) in trace.iter() {
+            let (signature, size) = parse_signature_size(signature_size)?;
+            columns.n_rows += 1;
+            store!(schema, columns, block_number, *block_number);
+            store!(schema, columns, transaction_index, Some(index as u32));
+            store!(schema, columns, transaction_hash, tx.clone());
+            store!(schema, columns, signature, signature.clone());
+            store!(schema, columns, size, size);
+            store!(schema, columns, count, *count);
+        }
+    }
+    Ok(())
+}
+
+fn parse_signature_size(signature_size: &str) -> R<(Vec<u8>, u64)> {
+    let parts: Vec<&str> = signature_size.splitn(2, '-').collect();
+    if parts.len() != 2 {
+        return Err(err("could not parse 4byte-size pair"))
+    }
+
+    // Parse the hexadecimal part (assuming there's no "0x" prefix as in the example given)
+    let hex_part = parts[0].trim_start_matches("0x");
+    let bytes = (0..hex_part.len())
+        .step_by(2)
+        .map(|i| u8::from_str_radix(&hex_part[i..i + 2], 16))
+        .collect::<Result<Vec<u8>, _>>()
+        .map_err(|_| err("could not parse signature bytes"))?;
+
+    // Parse the number as u64
+    let number = parts[1].parse::<u64>().map_err(|_| err("could not parse call data size"))?;
+
+    Ok((bytes, number))
+}

crates/freeze/src/datasets/mod.rs

@@ -30,6 +30,8 @@ pub mod erc721_metadata;
 pub mod erc721_transfers;
 /// eth calls
 pub mod eth_calls;
+/// four_byte traces
+pub mod four_byte_traces;
 /// geth balance diffs
 pub mod geth_balance_diffs;
 /// geth code diffs
@@ -81,6 +83,7 @@ pub use erc20_transfers::*;
 pub use erc721_metadata::*;
 pub use erc721_transfers::*;
 pub use eth_calls::*;
+pub use four_byte_traces::*;
 pub use geth_balance_diffs::*;
 pub use geth_code_diffs::*;
 pub use geth_nonce_diffs::*;

crates/freeze/src/types/datatypes/scalar.rs

@@ -19,6 +19,7 @@ define_datatypes!(
     Erc721Metadata,
     Erc721Transfers,
     EthCalls,
+    FourByteTraces,
     GethCodeDiffs,
     GethBalanceDiffs,
     GethStorageDiffs,

crates/freeze/src/types/sources.rs

@@ -418,6 +418,29 @@ impl<P: JsonRpcClient> Fetcher<P> {
         Ok((Some(block_number), txs, traces))
     }
 
+    /// get geth debug block call traces
+    pub async fn geth_debug_trace_block_4byte_traces(
+        &self,
+        block_number: u32,
+        include_transaction_hashes: bool,
+    ) -> Result<(Option<u32>, Vec<Option<Vec<u8>>>, Vec<BTreeMap<String, u64>>)> {
+        let tracer = GethDebugTracerType::BuiltInTracer(GethDebugBuiltInTracerType::FourByteTracer);
+        let options = GethDebugTracingOptions { tracer: Some(tracer), ..Default::default() };
+        let (block, txs, traces) =
+            self.geth_debug_trace_block(block_number, options, include_transaction_hashes).await?;
+
+        let mut calls = Vec::new();
+        for trace in traces.into_iter() {
+            match trace {
+                GethTrace::Known(GethTraceFrame::FourByteTracer(FourByteFrame(frame))) => {
+                    calls.push(frame)
+                }
+                _ => return Err(CollectError::CollectError("invalid trace result".to_string())),
+            }
+        }
+        Ok((block, txs, calls))
+    }
+
     /// get geth debug block call traces
     pub async fn geth_debug_trace_block_prestate(
         &self,
@@ -540,6 +563,30 @@ impl<P: JsonRpcClient> Fetcher<P> {
         Ok((block_number, vec![Some(transaction_hash)], traces))
     }
 
+    /// get geth debug block 4byte traces
+    pub async fn geth_debug_trace_transaction_4byte_traces(
+        &self,
+        transaction_hash: Vec<u8>,
+        include_block_number: bool,
+    ) -> Result<(Option<u32>, Vec<Option<Vec<u8>>>, Vec<BTreeMap<String, u64>>)> {
+        let tracer = GethDebugTracerType::BuiltInTracer(GethDebugBuiltInTracerType::FourByteTracer);
+        let options = GethDebugTracingOptions { tracer: Some(tracer), ..Default::default() };
+        let (block, txs, traces) = self
+            .geth_debug_trace_transaction(transaction_hash, options, include_block_number)
+            .await?;
+
+        let mut calls = Vec::new();
+        for trace in traces.into_iter() {
+            match trace {
+                GethTrace::Known(GethTraceFrame::FourByteTracer(FourByteFrame(frame))) => {
+                    calls.push(frame)
+                }
+                _ => return Err(CollectError::CollectError("invalid trace result".to_string())),
+            }
+        }
+        Ok((block, txs, calls))
+    }
+
     /// get geth debug block call traces
     pub async fn geth_debug_trace_transaction_prestate(
         &self,