refactor!: let ProofPlan::result_evaluate and `final_round_evaluate…

…` return `Table`
spaceandtimelabs · Nov 13, 2024 · 2550185 · 2550185
1 parent 36a8eba
commit 2550185
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 34 deletions.
diff --git a/crates/proof-of-sql/src/sql/proof/proof_plan.rs b/crates/proof-of-sql/src/sql/proof/proof_plan.rs
@@ -1,6 +1,6 @@
 use super::{CountBuilder, FinalRoundBuilder, FirstRoundBuilder, VerificationBuilder};
 use crate::base::{
-    database::{Column, ColumnField, ColumnRef, DataAccessor, OwnedTable, TableRef},
+    database::{ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
     map::{IndexMap, IndexSet},
     proof::ProofError,
     scalar::Scalar,
@@ -40,7 +40,7 @@ pub trait ProverEvaluate {
         &self,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>>;
+    ) -> Table<'a, S>;
 
     /// Evaluate the query and modify `FirstRoundBuilder` to form the query's proof.
     fn first_round_evaluate(&self, builder: &mut FirstRoundBuilder);
@@ -56,7 +56,7 @@ pub trait ProverEvaluate {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>>;
+    ) -> Table<'a, S>;
 }
 
 /// Marker used as a trait bound for generic [`ProofPlan`] types to indicate the honesty of their implementation.

diff --git a/crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs
@@ -1,6 +1,6 @@
 use crate::{
     base::{
-        database::{Column, ColumnField, ColumnRef, DataAccessor, OwnedTable, TableRef},
+        database::{ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
         map::{IndexMap, IndexSet},
         proof::ProofError,
         scalar::Scalar,
@@ -67,8 +67,8 @@ impl ProverEvaluate for EmptyExec {
         &self,
         _alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
-        Vec::new()
+    ) -> Table<'a, S> {
+        Table::<'a, S>::try_new(IndexMap::default()).unwrap()
     }
 
     fn first_round_evaluate(&self, _builder: &mut FirstRoundBuilder) {}
@@ -80,7 +80,7 @@ impl ProverEvaluate for EmptyExec {
         _builder: &mut FinalRoundBuilder<'a, S>,
         _alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
-        Vec::new()
+    ) -> Table<'a, S> {
+        Table::<'a, S>::try_new(IndexMap::default()).unwrap()
     }
 }
diff --git a/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs
@@ -141,7 +141,7 @@ impl ProverEvaluate for FilterExec {
         &self,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let column_refs = self.get_column_references();
         let used_table =
             Table::<'a, S>::from_columns(&column_refs, self.table.table_ref, accessor, alloc);
@@ -160,7 +160,13 @@ impl ProverEvaluate for FilterExec {
 
         // Compute filtered_columns and indexes
         let (filtered_columns, _) = filter_columns(alloc, &columns, selection);
-        filtered_columns
+        Table::<'a, S>::try_from_iter(
+            self.aliased_expr
+                .iter()
+                .map(|expr| expr.alias)
+                .zip(filtered_columns),
+        )
+        .expect("Failed to create table from iterator")
     }
 
     fn first_round_evaluate(&self, builder: &mut FirstRoundBuilder) {
@@ -174,7 +180,7 @@ impl ProverEvaluate for FilterExec {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let column_refs = self.get_column_references();
         let used_table =
             Table::<'a, S>::from_columns(&column_refs, self.table.table_ref, accessor, alloc);
@@ -216,7 +222,13 @@ impl ProverEvaluate for FilterExec {
             &filtered_columns,
             result_len,
         );
-        filtered_columns
+        Table::<'a, S>::try_from_iter(
+            self.aliased_expr
+                .iter()
+                .map(|expr| expr.alias)
+                .zip(filtered_columns),
+        )
+        .expect("Failed to create table from iterator")
     }
 }
 

diff --git a/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test_dishonest_prover.rs b/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test_dishonest_prover.rs
@@ -37,7 +37,7 @@ impl ProverEvaluate for DishonestFilterExec {
         &self,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let column_refs = self.get_column_references();
         let used_table =
             Table::<'a, S>::from_columns(&column_refs, self.table.table_ref, accessor, alloc);
@@ -55,7 +55,13 @@ impl ProverEvaluate for DishonestFilterExec {
         // Compute filtered_columns
         let (filtered_columns, _) = filter_columns(alloc, &columns, selection);
         let filtered_columns = tamper_column(alloc, filtered_columns);
-        filtered_columns
+        Table::<'a, S>::try_from_iter(
+            self.aliased_expr
+                .iter()
+                .map(|expr| expr.alias)
+                .zip(filtered_columns),
+        )
+        .expect("Failed to create table from iterator")
     }
 
     fn first_round_evaluate(&self, builder: &mut FirstRoundBuilder) {
@@ -73,7 +79,7 @@ impl ProverEvaluate for DishonestFilterExec {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let column_refs = self.get_column_references();
         let used_table =
             Table::<'a, S>::from_columns(&column_refs, self.table.table_ref, accessor, alloc);
@@ -115,7 +121,13 @@ impl ProverEvaluate for DishonestFilterExec {
             &filtered_columns,
             result_len,
         );
-        filtered_columns
+        Table::<'a, S>::try_from_iter(
+            self.aliased_expr
+                .iter()
+                .map(|expr| expr.alias)
+                .zip(filtered_columns),
+        )
+        .expect("Failed to create table from iterator")
     }
 }
 

diff --git a/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs
@@ -1,6 +1,6 @@
 use crate::{
     base::{
-        database::{Column, ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
+        database::{ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
         map::{IndexMap, IndexSet},
         proof::ProofError,
         scalar::Scalar,
@@ -89,16 +89,17 @@ impl ProverEvaluate for ProjectionExec {
         &self,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let column_refs = self.get_column_references();
         let used_table =
             Table::<'a, S>::from_columns(&column_refs, self.table.table_ref, accessor, alloc);
-        let columns: Vec<_> = self
-            .aliased_results
-            .iter()
-            .map(|aliased_expr| aliased_expr.expr.result_evaluate(alloc, &used_table))
-            .collect();
-        columns
+        Table::<'a, S>::try_from_iter(self.aliased_results.iter().map(|aliased_expr| {
+            (
+                aliased_expr.alias,
+                aliased_expr.expr.result_evaluate(alloc, &used_table),
+            )
+        }))
+        .expect("Failed to create table from iterator")
     }
 
     fn first_round_evaluate(&self, _builder: &mut FirstRoundBuilder) {}
@@ -119,17 +120,15 @@ impl ProverEvaluate for ProjectionExec {
         let used_table =
             Table::<'a, S>::from_columns(&column_refs, self.table.table_ref, accessor, alloc);
         // 1. Evaluate result expressions
-        let res: Vec<_> = self
-            .aliased_results
-            .iter()
-            .map(|aliased_expr| {
-                aliased_expr
-                    .expr
-                    .prover_evaluate(builder, alloc, &used_table)
-            })
-            .collect();
+        let res = Table::<'a, S>::try_from_iter(self.aliased_results.iter().map(|aliased_expr| {
+            (
+                aliased_expr.alias,
+                aliased_expr.expr.result_evaluate(alloc, &used_table),
+            )
+        }))
+        .expect("Failed to create table from iterator");
         // 2. Produce MLEs
-        res.clone().into_iter().for_each(|column| {
+        res.inner_table().values().for_each(|column| {
             builder.produce_intermediate_mle(column.as_scalar(alloc));
         });
         res