From 2a197527501a7d2602b30587e750b98816537f9f Mon Sep 17 00:00:00 2001 From: Ian Alexander Joiner <14581281+iajoiner@users.noreply.github.com> Date: Fri, 6 Sep 2024 16:22:04 -0400 Subject: [PATCH] feat!: allow `ProofPlan::result_evaluate` and `ProofPlan::prover_evaluate` to return columns (#132) # Rationale for this change This is a simple PR to make it possible to compose `ProofPlan`s which facilitates `SliceExec` and `UnionExec`. # What changes are included in this PR? - make `ProofPlan::result_evaluate` and `ProofPlan::prover_evaluate` return `Vec>` # Are these changes tested? Existing tests should pass --- .../proof-of-sql/src/sql/proof/proof_plan.rs | 7 ++-- .../src/sql/proof/query_proof_test.rs | 35 ++++++++++++------- .../sql/proof/verifiable_query_result_test.rs | 18 ++++++---- .../src/sql/proof_plans/dense_filter_exec.rs | 10 +++--- ...dense_filter_exec_test_dishonest_prover.rs | 10 +++--- .../src/sql/proof_plans/dyn_proof_plan.rs | 6 ++-- .../src/sql/proof_plans/filter_exec.rs | 20 ++++++----- .../filter_exec_test_dishonest_prover.rs | 17 ++++++--- .../src/sql/proof_plans/filter_result_expr.rs | 9 +++-- .../src/sql/proof_plans/group_by_exec.rs | 26 ++++++++++---- .../src/sql/proof_plans/projection_exec.rs | 20 ++++++----- 11 files changed, 116 insertions(+), 62 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 index 630f5fa38..de961efea 100644 --- a/crates/proof-of-sql/src/sql/proof/proof_plan.rs +++ b/crates/proof-of-sql/src/sql/proof/proof_plan.rs @@ -2,7 +2,8 @@ use super::{CountBuilder, ProofBuilder, ResultBuilder, VerificationBuilder}; use crate::base::{ commitment::Commitment, database::{ - ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor, OwnedTable, + Column, ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor, + OwnedTable, }, proof::ProofError, scalar::Scalar, @@ -53,7 +54,7 @@ pub trait ProverEvaluate { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ); + ) -> Vec>; /// Evaluate the query and modify `ProofBuilder` to store an intermediate representation /// of the query result and track all the components needed to form the query's proof. @@ -66,7 +67,7 @@ pub trait ProverEvaluate { builder: &mut ProofBuilder<'a, S>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ); + ) -> Vec>; } /// 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_test.rs b/crates/proof-of-sql/src/sql/proof/query_proof_test.rs index da110be83..ede164b20 100644 --- a/crates/proof-of-sql/src/sql/proof/query_proof_test.rs +++ b/crates/proof-of-sql/src/sql/proof/query_proof_test.rs @@ -6,8 +6,9 @@ use crate::{ commitment::{Commitment, InnerProductProof}, database::{ owned_table_utility::{bigint, owned_table}, - ColumnField, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor, MetadataAccessor, - OwnedTable, OwnedTableTestAccessor, TestAccessor, UnimplementedTestAccessor, + Column, ColumnField, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor, + MetadataAccessor, OwnedTable, OwnedTableTestAccessor, TestAccessor, + UnimplementedTestAccessor, }, proof::ProofError, scalar::{Curve25519Scalar, Scalar}, @@ -45,11 +46,12 @@ impl ProverEvaluate for TrivialTestProofPlan { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { let col = alloc.alloc_slice_fill_copy(builder.table_length(), self.column_fill_value); let indexes = Indexes::Sparse(vec![0u64]); builder.set_result_indexes(indexes); builder.produce_result_column(col as &[_]); + vec![Column::BigInt(col)] } fn prover_evaluate<'a>( @@ -57,12 +59,13 @@ impl ProverEvaluate for TrivialTestProofPlan { builder: &mut ProofBuilder<'a, S>, alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { let col = alloc.alloc_slice_fill_copy(builder.table_length(), self.column_fill_value); builder.produce_sumcheck_subpolynomial( SumcheckSubpolynomialType::Identity, vec![(S::ONE, vec![Box::new(col as &[_])])], ); + vec![Column::BigInt(col)] } } impl ProofPlan for TrivialTestProofPlan { @@ -208,11 +211,13 @@ impl ProverEvaluate for SquareTestProofPlan { fn result_evaluate<'a>( &self, builder: &mut ResultBuilder<'a>, - _alloc: &'a Bump, + alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { builder.set_result_indexes(Indexes::Sparse(vec![0, 1])); builder.produce_result_column(self.res); + let res: &[_] = alloc.alloc_slice_copy(&self.res); + vec![Column::BigInt(res)] } fn prover_evaluate<'a>( @@ -220,7 +225,7 @@ impl ProverEvaluate for SquareTestProofPlan { builder: &mut ProofBuilder<'a, S>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { let x = accessor.get_column(ColumnRef::new( "sxt.test".parse().unwrap(), "x".parse().unwrap(), @@ -235,6 +240,7 @@ impl ProverEvaluate for SquareTestProofPlan { (-S::ONE, vec![Box::new(x.clone()), Box::new(x)]), ], ); + vec![Column::BigInt(res)] } } impl ProofPlan for SquareTestProofPlan { @@ -384,11 +390,13 @@ impl ProverEvaluate for DoubleSquareTestProofPlan { fn result_evaluate<'a>( &self, builder: &mut ResultBuilder<'a>, - _alloc: &'a Bump, + alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { builder.set_result_indexes(Indexes::Sparse(vec![0, 1])); builder.produce_result_column(self.res); + let res: &[_] = alloc.alloc_slice_copy(&self.res); + vec![Column::BigInt(res)] } fn prover_evaluate<'a>( @@ -396,7 +404,7 @@ impl ProverEvaluate for DoubleSquareTestProofPlan { builder: &mut ProofBuilder<'a, S>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { let x = accessor.get_column(ColumnRef::new( "sxt.test".parse().unwrap(), "x".parse().unwrap(), @@ -424,6 +432,7 @@ impl ProverEvaluate for DoubleSquareTestProofPlan { (-S::ONE, vec![Box::new(z), Box::new(z)]), ], ); + vec![Column::BigInt(res)] } } impl ProofPlan for DoubleSquareTestProofPlan { @@ -591,10 +600,11 @@ impl ProverEvaluate for ChallengeTestProofPlan { builder: &mut ResultBuilder<'a>, _alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { builder.set_result_indexes(Indexes::Sparse(vec![0, 1])); builder.produce_result_column([9, 25]); builder.request_post_result_challenges(2); + vec![Column::BigInt(&[9, 25])] } fn prover_evaluate<'a>( @@ -602,7 +612,7 @@ impl ProverEvaluate for ChallengeTestProofPlan { builder: &mut ProofBuilder<'a, S>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { let x = accessor.get_column(ColumnRef::new( "sxt.test".parse().unwrap(), "x".parse().unwrap(), @@ -619,6 +629,7 @@ impl ProverEvaluate for ChallengeTestProofPlan { (-alpha, vec![Box::new(x.clone()), Box::new(x)]), ], ); + vec![Column::BigInt(&[9, 25])] } } 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 index e423280c2..430c84e5b 100644 --- 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 @@ -7,8 +7,8 @@ use crate::{ commitment::{Commitment, InnerProductProof}, database::{ owned_table_utility::{bigint, owned_table}, - ColumnField, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor, MetadataAccessor, - OwnedTable, TestAccessor, UnimplementedTestAccessor, + Column, ColumnField, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor, + MetadataAccessor, OwnedTable, TestAccessor, UnimplementedTestAccessor, }, proof::ProofError, scalar::Scalar, @@ -28,16 +28,22 @@ impl ProverEvaluate for EmptyTestQueryExpr { fn result_evaluate<'a>( &self, _builder: &mut ResultBuilder<'a>, - _alloc: &'a Bump, + alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { + let zeros = vec![0; self.length]; + let res: &[_] = alloc.alloc_slice_copy(&zeros); + vec![Column::BigInt(res); self.columns] } fn prover_evaluate<'a>( &self, _builder: &mut ProofBuilder<'a, S>, - _alloc: &'a Bump, + alloc: &'a Bump, _accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { + let zeros = vec![0; self.length]; + let res: &[_] = alloc.alloc_slice_copy(&zeros); + vec![Column::BigInt(res); self.columns] } } impl ProofPlan for EmptyTestQueryExpr { diff --git a/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec.rs index 215d1fe2f..59b6fc993 100644 --- a/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec.rs +++ b/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec.rs @@ -152,7 +152,7 @@ impl ProverEvaluate for DenseFilterExec { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // 1. selection let selection_column: Column<'a, C::Scalar> = self.where_clause @@ -172,10 +172,11 @@ impl ProverEvaluate for DenseFilterExec { // 3. set indexes builder.set_result_indexes(Indexes::Dense(0..(result_len as u64))); // 4. set filtered_columns - for col in filtered_columns { - builder.produce_result_column(col); + for col in &filtered_columns { + builder.produce_result_column(col.clone()); } builder.request_post_result_challenges(2); + filtered_columns } #[tracing::instrument(name = "DenseFilterExec::prover_evaluate", level = "debug", skip_all)] @@ -185,7 +186,7 @@ impl ProverEvaluate for DenseFilterExec { builder: &mut ProofBuilder<'a, C::Scalar>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // 1. selection let selection_column: Column<'a, C::Scalar> = self.where_clause.prover_evaluate(builder, alloc, accessor); @@ -215,6 +216,7 @@ impl ProverEvaluate for DenseFilterExec { &filtered_columns, result_len, ); + filtered_columns } } diff --git a/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec_test_dishonest_prover.rs b/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec_test_dishonest_prover.rs index 265d66b8c..b094d45ac 100644 --- a/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec_test_dishonest_prover.rs +++ b/crates/proof-of-sql/src/sql/proof_plans/dense_filter_exec_test_dishonest_prover.rs @@ -40,7 +40,7 @@ impl ProverEvaluate for DishonestDenseFilterExec, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // 1. selection let selection_column: Column<'a, Curve25519Scalar> = self.where_clause @@ -60,10 +60,11 @@ impl ProverEvaluate for DishonestDenseFilterExec for DishonestDenseFilterExec, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // 1. selection let selection_column: Column<'a, Curve25519Scalar> = self.where_clause.prover_evaluate(builder, alloc, accessor); @@ -107,6 +108,7 @@ impl ProverEvaluate for DishonestDenseFilterExec ProverEvaluate for DynProofPlan { builder: &mut crate::sql::proof::ResultBuilder<'a>, alloc: &'a bumpalo::Bump, accessor: &'a dyn crate::base::database::DataAccessor, - ) { + ) -> Vec> { match self { DynProofPlan::Projection(expr) => expr.result_evaluate(builder, alloc, accessor), DynProofPlan::Filter(expr) => expr.result_evaluate(builder, alloc, accessor), @@ -123,7 +123,7 @@ impl ProverEvaluate for DynProofPlan { builder: &mut crate::sql::proof::ProofBuilder<'a, C::Scalar>, alloc: &'a bumpalo::Bump, accessor: &'a dyn crate::base::database::DataAccessor, - ) { + ) -> Vec> { match self { DynProofPlan::Projection(expr) => expr.prover_evaluate(builder, alloc, accessor), DynProofPlan::Filter(expr) => expr.prover_evaluate(builder, alloc, accessor), 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 index 597e753b3..990b706c4 100644 --- a/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs +++ b/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs @@ -120,7 +120,7 @@ impl ProverEvaluate for FilterExec { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // evaluate where clause let selection_column: Column<'a, C::Scalar> = self.where_clause @@ -139,9 +139,11 @@ impl ProverEvaluate for FilterExec { builder.set_result_indexes(Indexes::Sparse(indexes)); // evaluate result columns - for expr in self.results.iter() { - expr.result_evaluate(builder, accessor); - } + Vec::from_iter( + self.results + .iter() + .map(|expr| expr.result_evaluate(builder, accessor)), + ) } #[tracing::instrument(name = "FilterExec::prover_evaluate", level = "debug", skip_all)] @@ -150,15 +152,17 @@ impl ProverEvaluate for FilterExec { builder: &mut ProofBuilder<'a, C::Scalar>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // evaluate where clause let selection_column: Column<'a, C::Scalar> = self.where_clause.prover_evaluate(builder, alloc, accessor); let selection = selection_column .as_boolean() .expect("selection is not boolean"); - for expr in self.results.iter() { - expr.prover_evaluate(builder, alloc, accessor, selection); - } + Vec::from_iter( + self.results + .iter() + .map(|expr| expr.prover_evaluate(builder, alloc, accessor, selection)), + ) } } 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 index a2afe5ab4..1a1319476 100644 --- 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 @@ -33,7 +33,7 @@ impl ProverEvaluate for DishonestFilterExec { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // evaluate where clause let selection_column: Column<'a, Curve25519Scalar> = self.where_clause @@ -51,6 +51,11 @@ impl ProverEvaluate for DishonestFilterExec { .collect(); indexes[0] += 1; builder.set_result_indexes(Indexes::Sparse(indexes)); + Vec::from_iter( + self.results + .iter() + .map(|expr| expr.result_evaluate(builder, accessor)), + ) } #[tracing::instrument( @@ -63,7 +68,7 @@ impl ProverEvaluate for DishonestFilterExec { builder: &mut ProofBuilder<'a, Curve25519Scalar>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // evaluate where clause let selection_column: Column<'a, Curve25519Scalar> = self.where_clause.prover_evaluate(builder, alloc, accessor); @@ -72,9 +77,11 @@ impl ProverEvaluate for DishonestFilterExec { .expect("selection is not boolean"); // evaluate result columns - for expr in self.results.iter() { - expr.prover_evaluate(builder, alloc, accessor, selection); - } + Vec::from_iter( + self.results + .iter() + .map(|expr| expr.prover_evaluate(builder, alloc, accessor, selection)), + ) } } diff --git a/crates/proof-of-sql/src/sql/proof_plans/filter_result_expr.rs b/crates/proof-of-sql/src/sql/proof_plans/filter_result_expr.rs index cbba90066..3b1b2e00f 100644 --- a/crates/proof-of-sql/src/sql/proof_plans/filter_result_expr.rs +++ b/crates/proof-of-sql/src/sql/proof_plans/filter_result_expr.rs @@ -50,8 +50,10 @@ impl FilterResultExpr { &self, builder: &mut ResultBuilder<'a>, accessor: &'a dyn DataAccessor, - ) { - builder.produce_result_column(accessor.get_column(self.column_ref)); + ) -> Column<'a, S> { + let col = accessor.get_column(self.column_ref); + builder.produce_result_column(col.clone()); + col } /// Given the selected rows (as a slice of booleans), evaluate the filter result expression and @@ -62,7 +64,7 @@ impl FilterResultExpr { alloc: &'a Bump, accessor: &'a dyn DataAccessor, selection: &'a [bool], - ) { + ) -> Column<'a, S> { match accessor.get_column(self.column_ref) { Column::Boolean(col) => prover_evaluate_impl(builder, alloc, selection, col), Column::SmallInt(col) => prover_evaluate_impl(builder, alloc, selection, col), @@ -77,6 +79,7 @@ impl FilterResultExpr { Column::VarChar((_, scals)) => prover_evaluate_impl(builder, alloc, selection, scals), Column::TimestampTZ(_, _, col) => prover_evaluate_impl(builder, alloc, selection, col), }; + accessor.get_column(self.column_ref) } /// Given the evaluation of the selected row's multilinear extension at sumcheck's random point, diff --git a/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs index 6bf651b4e..d81aafd30 100644 --- a/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs +++ b/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs @@ -206,7 +206,7 @@ impl ProverEvaluate for GroupByExec { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // 1. selection let selection_column: Column<'a, C::Scalar> = self.where_clause @@ -238,14 +238,21 @@ impl ProverEvaluate for GroupByExec { // 3. set indexes builder.set_result_indexes(Indexes::Dense(0..(count_column.len() as u64))); // 4. set filtered_columns - for col in group_by_result_columns { - builder.produce_result_column(col); + for col in &group_by_result_columns { + builder.produce_result_column(col.clone()); } - for col in sum_result_columns { - builder.produce_result_column(col); + for col in &sum_result_columns { + builder.produce_result_column(*col); } + let sum_result_columns_iter = sum_result_columns.iter().map(|col| Column::Scalar(col)); builder.produce_result_column(count_column); builder.request_post_result_challenges(2); + Vec::from_iter( + group_by_result_columns + .into_iter() + .chain(sum_result_columns_iter) + .chain(std::iter::once(Column::BigInt(count_column))), + ) } #[tracing::instrument(name = "GroupByExec::prover_evaluate", level = "debug", skip_all)] @@ -255,7 +262,7 @@ impl ProverEvaluate for GroupByExec { builder: &mut ProofBuilder<'a, C::Scalar>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { // 1. selection let selection_column: Column<'a, C::Scalar> = self.where_clause.prover_evaluate(builder, alloc, accessor); @@ -294,6 +301,13 @@ impl ProverEvaluate for GroupByExec { (&group_by_columns, &sum_columns, selection), (&group_by_result_columns, &sum_result_columns, count_column), ); + let sum_result_columns_iter = sum_result_columns.iter().map(|col| Column::Scalar(col)); + Vec::from_iter( + group_by_result_columns + .into_iter() + .chain(sum_result_columns_iter) + .chain(std::iter::once(Column::BigInt(count_column))), + ) } } 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 index e0b332b67..57213b9a3 100644 --- a/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs +++ b/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs @@ -2,7 +2,8 @@ use crate::{ base::{ commitment::Commitment, database::{ - ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor, OwnedTable, + Column, ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor, + OwnedTable, }, proof::ProofError, }, @@ -100,16 +101,17 @@ impl ProverEvaluate for ProjectionExec { builder: &mut ResultBuilder<'a>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { + ) -> Vec> { let columns = Vec::from_iter(self.aliased_results.iter().map(|aliased_expr| { aliased_expr .expr .result_evaluate(builder.table_length(), alloc, accessor) })); builder.set_result_indexes(Indexes::Dense(0..(builder.table_length() as u64))); - for col in columns { - builder.produce_result_column(col); + for col in &columns { + builder.produce_result_column(col.clone()); } + columns } #[tracing::instrument(name = "ProjectionExec::prover_evaluate", level = "debug", skip_all)] @@ -119,9 +121,11 @@ impl ProverEvaluate for ProjectionExec { builder: &mut ProofBuilder<'a, C::Scalar>, alloc: &'a Bump, accessor: &'a dyn DataAccessor, - ) { - self.aliased_results.iter().for_each(|aliased_expr| { - aliased_expr.expr.prover_evaluate(builder, alloc, accessor); - }); + ) -> Vec> { + Vec::from_iter( + self.aliased_results + .iter() + .map(|aliased_expr| aliased_expr.expr.prover_evaluate(builder, alloc, accessor)), + ) } }