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

…` return `Table`
spaceandtimelabs · Nov 14, 2024 · 513d374 · 513d374
1 parent cecf821
commit 513d374
Show file tree

Hide file tree

Showing 13 changed files with 215 additions and 103 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/query_proof.rs b/crates/proof-of-sql/src/sql/proof/query_proof.rs
@@ -6,7 +6,7 @@ use crate::{
     base::{
         bit::BitDistribution,
         commitment::CommitmentEvaluationProof,
-        database::{Column, CommitmentAccessor, DataAccessor, MetadataAccessor, TableRef},
+        database::{CommitmentAccessor, DataAccessor, MetadataAccessor, TableRef},
         map::IndexMap,
         math::log2_up,
         polynomial::{compute_evaluation_vector, CompositePolynomialInfo},
@@ -76,9 +76,14 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         let alloc = Bump::new();
 
         // Evaluate query result
-        let result_cols = expr.result_evaluate(&alloc, accessor);
-        let output_length = result_cols.first().map_or(0, Column::len);
-        let provable_result = ProvableQueryResult::new(output_length as u64, &result_cols);
+        let result_table = expr.result_evaluate(&alloc, accessor);
+        let result_cols = result_table
+            .inner_table()
+            .values()
+            .cloned()
+            .collect::<Vec<_>>();
+        let provable_result =
+            ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols);
 
         // Prover First Round
         let mut first_round_builder = FirstRoundBuilder::new();

diff --git a/crates/proof-of-sql/src/sql/proof/query_proof_test.rs b/crates/proof-of-sql/src/sql/proof/query_proof_test.rs
@@ -7,9 +7,9 @@ use crate::{
         database::{
             owned_table_utility::{bigint, owned_table},
             Column, ColumnField, ColumnRef, ColumnType, DataAccessor, OwnedTable,
-            OwnedTableTestAccessor, TableRef,
+            OwnedTableTestAccessor, Table, TableRef,
         },
-        map::{indexset, IndexMap, IndexSet},
+        map::{indexmap, indexset, IndexMap, IndexSet},
         proof::ProofError,
         scalar::{Curve25519Scalar, Scalar},
     },
@@ -44,9 +44,12 @@ impl ProverEvaluate for TrivialTestProofPlan {
         &self,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let col = alloc.alloc_slice_fill_copy(self.length, self.column_fill_value);
-        vec![Column::BigInt(col)]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(col),
+        })
+        .unwrap()
     }
 
     fn first_round_evaluate(&self, _builder: &mut FirstRoundBuilder) {}
@@ -56,14 +59,17 @@ impl ProverEvaluate for TrivialTestProofPlan {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let col = alloc.alloc_slice_fill_copy(self.length, self.column_fill_value);
         builder.produce_intermediate_mle(col as &[_]);
         builder.produce_sumcheck_subpolynomial(
             SumcheckSubpolynomialType::Identity,
             vec![(S::ONE, vec![Box::new(col as &[_])])],
         );
-        vec![Column::BigInt(col)]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(col),
+        })
+        .unwrap()
     }
 }
 impl ProofPlan for TrivialTestProofPlan {
@@ -213,9 +219,12 @@ impl ProverEvaluate for SquareTestProofPlan {
         &self,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let res: &[_] = alloc.alloc_slice_copy(&self.res);
-        vec![Column::BigInt(res)]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(res),
+        })
+        .unwrap()
     }
 
     fn first_round_evaluate(&self, _builder: &mut FirstRoundBuilder) {}
@@ -225,7 +234,7 @@ impl ProverEvaluate for SquareTestProofPlan {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let x = accessor.get_column(ColumnRef::new(
             "sxt.test".parse().unwrap(),
             "x".parse().unwrap(),
@@ -240,7 +249,10 @@ impl ProverEvaluate for SquareTestProofPlan {
                 (-S::ONE, vec![Box::new(x), Box::new(x)]),
             ],
         );
-        vec![Column::BigInt(res)]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(&[9, 25]),
+        })
+        .unwrap()
     }
 }
 impl ProofPlan for SquareTestProofPlan {
@@ -390,9 +402,12 @@ impl ProverEvaluate for DoubleSquareTestProofPlan {
         &self,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let res: &[_] = alloc.alloc_slice_copy(&self.res);
-        vec![Column::BigInt(res)]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(res),
+        })
+        .unwrap()
     }
 
     fn first_round_evaluate(&self, _builder: &mut FirstRoundBuilder) {}
@@ -402,7 +417,7 @@ impl ProverEvaluate for DoubleSquareTestProofPlan {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let x = accessor.get_column(ColumnRef::new(
             "sxt.test".parse().unwrap(),
             "x".parse().unwrap(),
@@ -430,7 +445,10 @@ impl ProverEvaluate for DoubleSquareTestProofPlan {
             ],
         );
         builder.produce_intermediate_mle(res);
-        vec![Column::BigInt(res)]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(res),
+        })
+        .unwrap()
     }
 }
 impl ProofPlan for DoubleSquareTestProofPlan {
@@ -597,8 +615,11 @@ impl ProverEvaluate for ChallengeTestProofPlan {
         &self,
         _alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
-        vec![Column::BigInt(&[9, 25])]
+    ) -> Table<'a, S> {
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(&[9, 25]),
+        })
+        .unwrap()
     }
 
     fn first_round_evaluate(&self, builder: &mut FirstRoundBuilder) {
@@ -610,7 +631,7 @@ impl ProverEvaluate for ChallengeTestProofPlan {
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let x = accessor.get_column(ColumnRef::new(
             "sxt.test".parse().unwrap(),
             "x".parse().unwrap(),
@@ -627,7 +648,10 @@ impl ProverEvaluate for ChallengeTestProofPlan {
                 (-alpha, vec![Box::new(x), Box::new(x)]),
             ],
         );
-        vec![Column::BigInt(&[9, 25])]
+        Table::<'a, S>::try_new(indexmap! {
+            "a1".parse().unwrap() => Column::BigInt(&[9, 25]),
+        })
+        .unwrap()
     }
 }
 impl ProofPlan for ChallengeTestProofPlan {

diff --git a/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs b/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs
@@ -8,7 +8,7 @@ use crate::{
         database::{
             owned_table_utility::{bigint, owned_table},
             Column, ColumnField, ColumnRef, ColumnType, DataAccessor, OwnedTable,
-            OwnedTableTestAccessor, TableRef,
+            OwnedTableTestAccessor, Table, TableRef,
         },
         map::{indexset, IndexMap, IndexSet},
         proof::ProofError,
@@ -29,24 +29,30 @@ impl ProverEvaluate for EmptyTestQueryExpr {
         &self,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let zeros = vec![0; self.length];
         let res: &[_] = alloc.alloc_slice_copy(&zeros);
-        vec![Column::BigInt(res); self.columns]
+        Table::<'a, S>::try_from_iter(
+            (1..=self.columns).map(|i| (format!("a{i}").parse().unwrap(), Column::BigInt(res))),
+        )
+        .unwrap()
     }
     fn first_round_evaluate(&self, _builder: &mut FirstRoundBuilder) {}
     fn final_round_evaluate<'a, S: Scalar>(
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
-    ) -> Vec<Column<'a, S>> {
+    ) -> Table<'a, S> {
         let zeros = vec![0; self.length];
         let res: &[_] = alloc.alloc_slice_copy(&zeros);
         let _ = std::iter::repeat_with(|| builder.produce_intermediate_mle(res))
             .take(self.columns)
             .collect::<Vec<_>>();
-        vec![Column::BigInt(res); self.columns]
+        Table::<'a, S>::try_from_iter(
+            (1..=self.columns).map(|i| (format!("a{i}").parse().unwrap(), Column::BigInt(res))),
+        )
+        .unwrap()
     }
 }
 impl ProofPlan for EmptyTestQueryExpr {

diff --git a/crates/proof-of-sql/src/sql/proof_plans/dyn_proof_plan.rs b/crates/proof-of-sql/src/sql/proof_plans/dyn_proof_plan.rs
@@ -1,7 +1,7 @@
 use super::{EmptyExec, FilterExec, GroupByExec, ProjectionExec, TableExec};
 use crate::{
     base::{
-        database::{Column, ColumnField, ColumnRef, DataAccessor, OwnedTable, TableRef},
+        database::{ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
         map::{IndexMap, IndexSet},
         proof::ProofError,
         scalar::Scalar,

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 = accessor.get_table(self.table.table_ref, &column_refs);
         // 1. selection
@@ -159,7 +159,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_results
+                .iter()
+                .map(|expr| expr.alias)
+                .zip(filtered_columns),
+        )
+        .expect("Failed to create table from iterator")
     }
 
     fn first_round_evaluate(&self, builder: &mut FirstRoundBuilder) {
@@ -173,7 +179,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 = accessor.get_table(self.table.table_ref, &column_refs);
         // 1. selection
@@ -214,7 +220,13 @@ impl ProverEvaluate for FilterExec {
             &filtered_columns,
             result_len,
         );
-        filtered_columns
+        Table::<'a, S>::try_from_iter(
+            self.aliased_results
+                .iter()
+                .map(|expr| expr.alias)
+                .zip(filtered_columns),
+        )
+        .expect("Failed to create table from iterator")
     }
 }