add --txs support for state_diffs and vm_traces

paradigmxyz · Aug 8, 2023 · f81cd69 · f81cd69
1 parent d8352b5
commit f81cd69
Show file tree

Hide file tree

Showing 6 changed files with 325 additions and 40 deletions.
diff --git a/crates/freeze/src/datasets/balance_diffs.rs b/crates/freeze/src/datasets/balance_diffs.rs
@@ -4,7 +4,8 @@ use polars::prelude::*;
 
 use super::state_diffs;
 use crate::types::{
-    BalanceDiffs, BlockChunk, CollectError, ColumnType, Dataset, Datatype, RowFilter, Source, Table,
+    BalanceDiffs, BlockChunk, CollectError, ColumnType, Dataset, Datatype, RowFilter, Source,
+    Table, TransactionChunk,
 };
 
 #[async_trait::async_trait]
@@ -51,6 +52,30 @@ impl Dataset for BalanceDiffs {
         schema: &Table,
         filter: Option<&RowFilter>,
     ) -> Result<DataFrame, CollectError> {
-        state_diffs::collect_single(&Datatype::BalanceDiffs, chunk, source, schema, filter).await
+        state_diffs::collect_block_state_diffs(
+            &Datatype::BalanceDiffs,
+            chunk,
+            source,
+            schema,
+            filter,
+        )
+        .await
+    }
+
+    async fn collect_transaction_chunk(
+        &self,
+        chunk: &TransactionChunk,
+        source: &Source,
+        schema: &Table,
+        filter: Option<&RowFilter>,
+    ) -> Result<DataFrame, CollectError> {
+        state_diffs::collect_transaction_state_diffs(
+            &Datatype::BalanceDiffs,
+            chunk,
+            source,
+            schema,
+            filter,
+        )
+        .await
     }
 }
diff --git a/crates/freeze/src/datasets/code_diffs.rs b/crates/freeze/src/datasets/code_diffs.rs
@@ -5,6 +5,7 @@ use polars::prelude::*;
 use super::state_diffs;
 use crate::types::{
     BlockChunk, CodeDiffs, CollectError, ColumnType, Dataset, Datatype, RowFilter, Source, Table,
+    TransactionChunk,
 };
 
 #[async_trait::async_trait]
@@ -51,6 +52,24 @@ impl Dataset for CodeDiffs {
         schema: &Table,
         filter: Option<&RowFilter>,
     ) -> Result<DataFrame, CollectError> {
-        state_diffs::collect_single(&Datatype::CodeDiffs, chunk, source, schema, filter).await
+        state_diffs::collect_block_state_diffs(&Datatype::CodeDiffs, chunk, source, schema, filter)
+            .await
+    }
+
+    async fn collect_transaction_chunk(
+        &self,
+        chunk: &TransactionChunk,
+        source: &Source,
+        schema: &Table,
+        filter: Option<&RowFilter>,
+    ) -> Result<DataFrame, CollectError> {
+        state_diffs::collect_transaction_state_diffs(
+            &Datatype::CodeDiffs,
+            chunk,
+            source,
+            schema,
+            filter,
+        )
+        .await
     }
 }
diff --git a/crates/freeze/src/datasets/nonce_diffs.rs b/crates/freeze/src/datasets/nonce_diffs.rs
@@ -5,6 +5,7 @@ use polars::prelude::*;
 use super::state_diffs;
 use crate::types::{
     BlockChunk, CollectError, ColumnType, Dataset, Datatype, NonceDiffs, RowFilter, Source, Table,
+    TransactionChunk,
 };
 
 #[async_trait::async_trait]
@@ -51,6 +52,24 @@ impl Dataset for NonceDiffs {
         schema: &Table,
         filter: Option<&RowFilter>,
     ) -> Result<DataFrame, CollectError> {
-        state_diffs::collect_single(&Datatype::NonceDiffs, chunk, source, schema, filter).await
+        state_diffs::collect_block_state_diffs(&Datatype::NonceDiffs, chunk, source, schema, filter)
+            .await
+    }
+
+    async fn collect_transaction_chunk(
+        &self,
+        chunk: &TransactionChunk,
+        source: &Source,
+        schema: &Table,
+        filter: Option<&RowFilter>,
+    ) -> Result<DataFrame, CollectError> {
+        state_diffs::collect_transaction_state_diffs(
+            &Datatype::NonceDiffs,
+            chunk,
+            source,
+            schema,
+            filter,
+        )
+        .await
     }
 }