refactor!: replace DataAccessor in ProofPlan with IndexMap

crates/proof-of-sql/src/base/database/accessor.rs

@@ -90,7 +90,7 @@ pub trait DataAccessor<S: Scalar>: MetadataAccessor {
     /// Creates a new [`Table`] from a [`TableRef`] and [`ColumnRef`]s.
     ///
     /// Columns are retrieved from the [`DataAccessor`] using the provided [`TableRef`] and [`ColumnRef`]s.
-    /// The only reason why [`table_ref` is needed is because `column_refs` can be empty.
+    /// The only reason why [`table_ref`] is needed is because [`column_refs`] can be empty.
     /// # Panics
     /// Column length mismatches can occur in theory. In practice, this should not happen.
     fn get_table(&self, table_ref: TableRef, column_refs: &IndexSet<ColumnRef>) -> Table<S> {

crates/proof-of-sql/src/sql/proof/proof_plan.rs

@@ -1,6 +1,6 @@
 use super::{CountBuilder, FinalRoundBuilder, FirstRoundBuilder, VerificationBuilder};
 use crate::base::{
-    database::{ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
+    database::{ColumnField, ColumnRef, OwnedTable, Table, TableRef},
     map::{IndexMap, IndexSet},
     proof::ProofError,
     scalar::Scalar,
@@ -39,7 +39,7 @@ pub trait ProverEvaluate {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S>;
 
     /// Evaluate the query and modify `FirstRoundBuilder` to form the query's proof.
@@ -55,7 +55,7 @@ pub trait ProverEvaluate {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S>;
 }
 

crates/proof-of-sql/src/sql/proof/query_proof.rs

@@ -6,8 +6,10 @@ use crate::{
     base::{
         bit::BitDistribution,
         commitment::CommitmentEvaluationProof,
-        database::{CommitmentAccessor, DataAccessor, MetadataAccessor, TableRef},
-        map::IndexMap,
+        database::{
+            ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor, Table, TableRef,
+        },
+        map::{IndexMap, IndexSet},
         math::log2_up,
         polynomial::{compute_evaluation_vector, CompositePolynomialInfo},
         proof::{Keccak256Transcript, ProofError, Transcript},
@@ -75,8 +77,22 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
 
         let alloc = Bump::new();
 
+        let total_col_refs = expr.get_column_references();
+        let table_map: IndexMap<TableRef, Table<CP::Scalar>> = expr
+            .get_table_references()
+            .into_iter()
+            .map(|table_ref| {
+                let col_refs: IndexSet<ColumnRef> = total_col_refs
+                    .iter()
+                    .filter(|col_ref| col_ref.table_ref() == table_ref)
+                    .copied()
+                    .collect();
+                (table_ref, accessor.get_table(table_ref, &col_refs))
+            })
+            .collect();
+
         // Evaluate query result
-        let provable_result = expr.result_evaluate(&alloc, accessor).into();
+        let provable_result = expr.result_evaluate(&alloc, &table_map).into();
 
         // Prover First Round
         let mut first_round_builder = FirstRoundBuilder::new();
@@ -99,11 +115,11 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         let mut builder =
             FinalRoundBuilder::new(range_length, num_sumcheck_variables, post_result_challenges);
 
-        expr.get_column_references().into_iter().for_each(|col| {
-            builder.produce_anchored_mle(accessor.get_column(col));
-        });
+        for col_ref in total_col_refs {
+            builder.produce_anchored_mle(accessor.get_column(col_ref));
+        }
 
-        expr.final_round_evaluate(&mut builder, &alloc, accessor);
+        expr.final_round_evaluate(&mut builder, &alloc, &table_map);
 
         let num_sumcheck_variables = builder.num_sumcheck_variables();
 

crates/proof-of-sql/src/sql/proof/query_proof_test.rs

@@ -7,8 +7,8 @@ use crate::{
         database::{
             owned_table_utility::{bigint, owned_table},
             table_utility::*,
-            ColumnField, ColumnRef, ColumnType, DataAccessor, OwnedTable, OwnedTableTestAccessor,
-            Table, TableRef,
+            ColumnField, ColumnRef, ColumnType, OwnedTable, OwnedTableTestAccessor, Table,
+            TableRef,
         },
         map::{indexset, IndexMap, IndexSet},
         proof::ProofError,
@@ -17,6 +17,7 @@ use crate::{
     sql::proof::{FirstRoundBuilder, QueryData, SumcheckSubpolynomialType},
 };
 use bumpalo::Bump;
+use proof_of_sql_parser::Identifier;
 use serde::Serialize;
 
 /// Type to allow us to prove and verify an artificial polynomial where we prove
@@ -44,7 +45,7 @@ impl ProverEvaluate for TrivialTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         let col = vec![self.column_fill_value; self.length];
         table([borrowed_bigint("a1", col, alloc)])
@@ -56,7 +57,7 @@ impl ProverEvaluate for TrivialTestProofPlan {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         let col = alloc.alloc_slice_fill_copy(self.length, self.column_fill_value);
         builder.produce_intermediate_mle(col as &[_]);
@@ -217,7 +218,7 @@ impl ProverEvaluate for SquareTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         table([borrowed_bigint("a1", self.res, alloc)])
     }
@@ -228,13 +229,14 @@ impl ProverEvaluate for SquareTestProofPlan {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        let x = accessor.get_column(ColumnRef::new(
-            "sxt.test".parse().unwrap(),
-            "x".parse().unwrap(),
-            ColumnType::BigInt,
-        ));
+        let x = *table_map
+            .get(&TableRef::new("sxt.test".parse().unwrap()))
+            .unwrap()
+            .inner_table()
+            .get(&"x".parse::<Identifier>().unwrap())
+            .unwrap();
         let res: &[_] = alloc.alloc_slice_copy(&self.res);
         builder.produce_intermediate_mle(res);
         builder.produce_sumcheck_subpolynomial(
@@ -393,7 +395,7 @@ impl ProverEvaluate for DoubleSquareTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         table([borrowed_bigint("a1", self.res, alloc)])
     }
@@ -404,13 +406,14 @@ impl ProverEvaluate for DoubleSquareTestProofPlan {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        let x = accessor.get_column(ColumnRef::new(
-            "sxt.test".parse().unwrap(),
-            "x".parse().unwrap(),
-            ColumnType::BigInt,
-        ));
+        let x = *table_map
+            .get(&TableRef::new("sxt.test".parse().unwrap()))
+            .unwrap()
+            .inner_table()
+            .get(&"x".parse::<Identifier>().unwrap())
+            .unwrap();
         let res: &[_] = alloc.alloc_slice_copy(&self.res);
         let z: &[_] = alloc.alloc_slice_copy(&self.z);
         builder.produce_intermediate_mle(z);
@@ -599,7 +602,7 @@ impl ProverEvaluate for ChallengeTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         table([borrowed_bigint("a1", [9, 25], alloc)])
     }
@@ -612,13 +615,14 @@ impl ProverEvaluate for ChallengeTestProofPlan {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        let x = accessor.get_column(ColumnRef::new(
-            "sxt.test".parse().unwrap(),
-            "x".parse().unwrap(),
-            ColumnType::BigInt,
-        ));
+        let x = *table_map
+            .get(&TableRef::new("sxt.test".parse().unwrap()))
+            .unwrap()
+            .inner_table()
+            .get(&"x".parse::<Identifier>().unwrap())
+            .unwrap();
         let res: &[_] = alloc.alloc_slice_copy(&[9, 25]);
         let alpha = builder.consume_post_result_challenge();
         let _beta = builder.consume_post_result_challenge();

crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs

@@ -8,8 +8,8 @@ use crate::{
         database::{
             owned_table_utility::{bigint, owned_table},
             table_utility::*,
-            ColumnField, ColumnRef, ColumnType, DataAccessor, OwnedTable, OwnedTableTestAccessor,
-            Table, TableRef,
+            ColumnField, ColumnRef, ColumnType, OwnedTable, OwnedTableTestAccessor, Table,
+            TableRef,
         },
         map::{indexset, IndexMap, IndexSet},
         proof::ProofError,
@@ -29,7 +29,7 @@ impl ProverEvaluate for EmptyTestQueryExpr {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         let zeros = vec![0_i64; self.length];
         table_with_row_count(
@@ -42,7 +42,7 @@ impl ProverEvaluate for EmptyTestQueryExpr {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         let zeros = vec![0_i64; self.length];
         let res: &[_] = alloc.alloc_slice_copy(&zeros);

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::{ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableRef},
+        database::{ColumnField, ColumnRef, OwnedTable, Table, TableRef},
         map::{IndexMap, IndexSet},
         proof::ProofError,
         scalar::Scalar,

crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs

@@ -1,8 +1,6 @@
 use crate::{
     base::{
-        database::{
-            ColumnField, ColumnRef, DataAccessor, OwnedTable, Table, TableOptions, TableRef,
-        },
+        database::{ColumnField, ColumnRef, OwnedTable, Table, TableOptions, TableRef},
         map::{IndexMap, IndexSet},
         proof::ProofError,
         scalar::Scalar,
@@ -68,7 +66,7 @@ impl ProverEvaluate for EmptyExec {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         _alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         // Create an empty table with one row
         Table::<'a, S>::try_new_with_options(IndexMap::default(), TableOptions::new(Some(1)))
@@ -83,7 +81,7 @@ impl ProverEvaluate for EmptyExec {
         &self,
         _builder: &mut FinalRoundBuilder<'a, S>,
         _alloc: &'a Bump,
-        _accessor: &'a dyn DataAccessor<S>,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
         // Create an empty table with one row
         Table::<'a, S>::try_new_with_options(IndexMap::default(), TableOptions::new(Some(1)))

crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs

@@ -2,8 +2,8 @@ use super::{fold_columns, fold_vals};
 use crate::{
     base::{
         database::{
-            filter_util::filter_columns, Column, ColumnField, ColumnRef, DataAccessor, OwnedTable,
-            Table, TableOptions, TableRef,
+            filter_util::filter_columns, Column, ColumnField, ColumnRef, OwnedTable, Table,
+            TableOptions, TableRef,
         },
         map::{IndexMap, IndexSet},
         proof::ProofError,
@@ -140,12 +140,13 @@ impl ProverEvaluate for FilterExec {
     fn result_evaluate<'a, S: Scalar>(
         &self,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        let column_refs = self.get_column_references();
-        let used_table = accessor.get_table(self.table.table_ref, &column_refs);
+        let table = table_map
+            .get(&self.table.table_ref)
+            .expect("Table not found");
         // 1. selection
-        let selection_column: Column<'a, S> = self.where_clause.result_evaluate(alloc, &used_table);
+        let selection_column: Column<'a, S> = self.where_clause.result_evaluate(alloc, table);
         let selection = selection_column
             .as_boolean()
             .expect("selection is not boolean");
@@ -155,7 +156,7 @@ impl ProverEvaluate for FilterExec {
         let columns: Vec<_> = self
             .aliased_results
             .iter()
-            .map(|aliased_expr| aliased_expr.expr.result_evaluate(alloc, &used_table))
+            .map(|aliased_expr| aliased_expr.expr.result_evaluate(alloc, table))
             .collect();
 
         // Compute filtered_columns and indexes
@@ -180,14 +181,14 @@ impl ProverEvaluate for FilterExec {
         &self,
         builder: &mut FinalRoundBuilder<'a, S>,
         alloc: &'a Bump,
-        accessor: &'a dyn DataAccessor<S>,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        let column_refs = self.get_column_references();
-        let used_table = accessor.get_table(self.table.table_ref, &column_refs);
+        let table = table_map
+            .get(&self.table.table_ref)
+            .expect("Table not found");
         // 1. selection
         let selection_column: Column<'a, S> =
-            self.where_clause
-                .prover_evaluate(builder, alloc, &used_table);
+            self.where_clause.prover_evaluate(builder, alloc, table);
         let selection = selection_column
             .as_boolean()
             .expect("selection is not boolean");
@@ -197,11 +198,7 @@ impl ProverEvaluate for FilterExec {
         let columns: Vec<_> = self
             .aliased_results
             .iter()
-            .map(|aliased_expr| {
-                aliased_expr
-                    .expr
-                    .prover_evaluate(builder, alloc, &used_table)
-            })
+            .map(|aliased_expr| aliased_expr.expr.prover_evaluate(builder, alloc, table))
             .collect();
         // Compute filtered_columns
         let (filtered_columns, result_len) = filter_columns(alloc, &columns, selection);