feat: add ProofPlan::get_table_references (#323)

# Rationale for this change

An upstream crates needs to know what tables are being queried against
from the proof plan.

# What changes are included in this PR?

`ProofPlan::get_table_references` is added

# Are these changes tested?

Yes

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

@@ -3,7 +3,7 @@ use crate::base::{
     commitment::Commitment,
     database::{
         Column, ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor,
-        OwnedTable,
+        OwnedTable, TableRef,
     },
     map::IndexSet,
     proof::ProofError,
@@ -46,6 +46,9 @@ pub trait ProofPlan<C: Commitment>: Debug + Send + Sync + ProverEvaluate<C::Scal
 
     /// Return all the columns referenced in the Query
     fn get_column_references(&self) -> IndexSet<ColumnRef>;
+
+    /// Return all the tables referenced in the Query
+    fn get_table_references(&self) -> IndexSet<TableRef>;
 }
 
 pub trait ProverEvaluate<S: Scalar> {

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

@@ -7,7 +7,7 @@ use crate::{
         database::{
             owned_table_utility::{bigint, owned_table},
             Column, ColumnField, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor,
-            MetadataAccessor, OwnedTable, OwnedTableTestAccessor, TestAccessor,
+            MetadataAccessor, OwnedTable, OwnedTableTestAccessor, TableRef, TestAccessor,
             UnimplementedTestAccessor,
         },
         map::IndexSet,
@@ -109,6 +109,9 @@ impl<C: Commitment> ProofPlan<C> for TrivialTestProofPlan {
     fn get_column_references(&self) -> IndexSet<ColumnRef> {
         unimplemented!("no real usage for this function yet")
     }
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        unimplemented!("no real usage for this function yet")
+    }
 }
 
 fn verify_a_trivial_query_proof_with_given_offset(n: usize, offset_generators: usize) {
@@ -278,6 +281,9 @@ impl<C: Commitment> ProofPlan<C> for SquareTestProofPlan {
     fn get_column_references(&self) -> IndexSet<ColumnRef> {
         unimplemented!("no real usage for this function yet")
     }
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        unimplemented!("no real usage for this function yet")
+    }
 }
 
 fn verify_a_proof_with_an_anchored_commitment_and_given_offset(offset_generators: usize) {
@@ -481,6 +487,9 @@ impl<C: Commitment> ProofPlan<C> for DoubleSquareTestProofPlan {
     fn get_column_references(&self) -> IndexSet<ColumnRef> {
         unimplemented!("no real usage for this function yet")
     }
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        unimplemented!("no real usage for this function yet")
+    }
 }
 
 fn verify_a_proof_with_an_intermediate_commitment_and_given_offset(offset_generators: usize) {
@@ -677,6 +686,9 @@ impl<C: Commitment> ProofPlan<C> for ChallengeTestProofPlan {
     fn get_column_references(&self) -> IndexSet<ColumnRef> {
         unimplemented!("no real usage for this function yet")
     }
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        unimplemented!("no real usage for this function yet")
+    }
 }
 
 fn verify_a_proof_with_a_post_result_challenge_and_given_offset(offset_generators: usize) {

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, CommitmentAccessor, DataAccessor,
-            MetadataAccessor, OwnedTable, TestAccessor, UnimplementedTestAccessor,
+            MetadataAccessor, OwnedTable, TableRef, TestAccessor, UnimplementedTestAccessor,
         },
         map::IndexSet,
         proof::ProofError,
@@ -88,6 +88,10 @@ impl<C: Commitment> ProofPlan<C> for EmptyTestQueryExpr {
     fn get_column_references(&self) -> IndexSet<ColumnRef> {
         unimplemented!("no real usage for this function yet")
     }
+
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        unimplemented!("no real usage for this function yet")
+    }
 }
 
 #[test]

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

@@ -1,9 +1,21 @@
 use super::{FilterExec, GroupByExec, ProjectionExec};
 use crate::{
-    base::{commitment::Commitment, database::Column, map::IndexSet},
-    sql::proof::{ProofPlan, ProverEvaluate},
+    base::{
+        commitment::Commitment,
+        database::{
+            Column, ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor,
+            OwnedTable, TableRef,
+        },
+        map::IndexSet,
+        proof::ProofError,
+    },
+    sql::proof::{
+        CountBuilder, FinalRoundBuilder, FirstRoundBuilder, ProofPlan, ProverEvaluate,
+        VerificationBuilder,
+    },
 };
 use alloc::vec::Vec;
+use bumpalo::Bump;
 use serde::{Deserialize, Serialize};
 
 /// The query plan for proving a query
@@ -34,25 +46,25 @@ pub enum DynProofPlan<C: Commitment> {
 impl<C: Commitment> ProofPlan<C> for DynProofPlan<C> {
     fn count(
         &self,
-        builder: &mut crate::sql::proof::CountBuilder,
-        accessor: &dyn crate::base::database::MetadataAccessor,
-    ) -> Result<(), crate::base::proof::ProofError> {
+        builder: &mut CountBuilder,
+        accessor: &dyn MetadataAccessor,
+    ) -> Result<(), ProofError> {
         match self {
             DynProofPlan::Projection(expr) => expr.count(builder, accessor),
             DynProofPlan::GroupBy(expr) => expr.count(builder, accessor),
             DynProofPlan::Filter(expr) => expr.count(builder, accessor),
         }
     }
 
-    fn get_length(&self, accessor: &dyn crate::base::database::MetadataAccessor) -> usize {
+    fn get_length(&self, accessor: &dyn MetadataAccessor) -> usize {
         match self {
             DynProofPlan::Projection(expr) => expr.get_length(accessor),
             DynProofPlan::GroupBy(expr) => expr.get_length(accessor),
             DynProofPlan::Filter(expr) => expr.get_length(accessor),
         }
     }
 
-    fn get_offset(&self, accessor: &dyn crate::base::database::MetadataAccessor) -> usize {
+    fn get_offset(&self, accessor: &dyn MetadataAccessor) -> usize {
         match self {
             DynProofPlan::Projection(expr) => expr.get_offset(accessor),
             DynProofPlan::GroupBy(expr) => expr.get_offset(accessor),
@@ -63,41 +75,49 @@ impl<C: Commitment> ProofPlan<C> for DynProofPlan<C> {
     #[tracing::instrument(name = "DynProofPlan::verifier_evaluate", level = "debug", skip_all)]
     fn verifier_evaluate(
         &self,
-        builder: &mut crate::sql::proof::VerificationBuilder<C>,
-        accessor: &dyn crate::base::database::CommitmentAccessor<C>,
-        result: Option<&crate::base::database::OwnedTable<C::Scalar>>,
-    ) -> Result<Vec<C::Scalar>, crate::base::proof::ProofError> {
+        builder: &mut VerificationBuilder<C>,
+        accessor: &dyn CommitmentAccessor<C>,
+        result: Option<&OwnedTable<C::Scalar>>,
+    ) -> Result<Vec<C::Scalar>, ProofError> {
         match self {
             DynProofPlan::Projection(expr) => expr.verifier_evaluate(builder, accessor, result),
             DynProofPlan::GroupBy(expr) => expr.verifier_evaluate(builder, accessor, result),
             DynProofPlan::Filter(expr) => expr.verifier_evaluate(builder, accessor, result),
         }
     }
 
-    fn get_column_result_fields(&self) -> Vec<crate::base::database::ColumnField> {
+    fn get_column_result_fields(&self) -> Vec<ColumnField> {
         match self {
             DynProofPlan::Projection(expr) => expr.get_column_result_fields(),
             DynProofPlan::GroupBy(expr) => expr.get_column_result_fields(),
             DynProofPlan::Filter(expr) => expr.get_column_result_fields(),
         }
     }
 
-    fn get_column_references(&self) -> IndexSet<crate::base::database::ColumnRef> {
+    fn get_column_references(&self) -> IndexSet<ColumnRef> {
         match self {
             DynProofPlan::Projection(expr) => expr.get_column_references(),
             DynProofPlan::GroupBy(expr) => expr.get_column_references(),
             DynProofPlan::Filter(expr) => expr.get_column_references(),
         }
     }
+
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        match self {
+            DynProofPlan::Projection(expr) => expr.get_table_references(),
+            DynProofPlan::GroupBy(expr) => expr.get_table_references(),
+            DynProofPlan::Filter(expr) => expr.get_table_references(),
+        }
+    }
 }
 
 impl<C: Commitment> ProverEvaluate<C::Scalar> for DynProofPlan<C> {
     #[tracing::instrument(name = "DynProofPlan::result_evaluate", level = "debug", skip_all)]
     fn result_evaluate<'a>(
         &self,
         input_length: usize,
-        alloc: &'a bumpalo::Bump,
-        accessor: &'a dyn crate::base::database::DataAccessor<C::Scalar>,
+        alloc: &'a Bump,
+        accessor: &'a dyn DataAccessor<C::Scalar>,
     ) -> Vec<Column<'a, C::Scalar>> {
         match self {
             DynProofPlan::Projection(expr) => expr.result_evaluate(input_length, alloc, accessor),
@@ -106,7 +126,7 @@ impl<C: Commitment> ProverEvaluate<C::Scalar> for DynProofPlan<C> {
         }
     }
 
-    fn first_round_evaluate(&self, builder: &mut crate::sql::proof::FirstRoundBuilder) {
+    fn first_round_evaluate(&self, builder: &mut FirstRoundBuilder) {
         match self {
             DynProofPlan::Projection(expr) => expr.first_round_evaluate(builder),
             DynProofPlan::GroupBy(expr) => expr.first_round_evaluate(builder),
@@ -117,9 +137,9 @@ impl<C: Commitment> ProverEvaluate<C::Scalar> for DynProofPlan<C> {
     #[tracing::instrument(name = "DynProofPlan::final_round_evaluate", level = "debug", skip_all)]
     fn final_round_evaluate<'a>(
         &self,
-        builder: &mut crate::sql::proof::FinalRoundBuilder<'a, C::Scalar>,
-        alloc: &'a bumpalo::Bump,
-        accessor: &'a dyn crate::base::database::DataAccessor<C::Scalar>,
+        builder: &mut FinalRoundBuilder<'a, C::Scalar>,
+        alloc: &'a Bump,
+        accessor: &'a dyn DataAccessor<C::Scalar>,
     ) -> Vec<Column<'a, C::Scalar>> {
         match self {
             DynProofPlan::Projection(expr) => expr.final_round_evaluate(builder, alloc, accessor),

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

@@ -4,7 +4,7 @@ use crate::{
         commitment::Commitment,
         database::{
             filter_util::filter_columns, Column, ColumnField, ColumnRef, CommitmentAccessor,
-            DataAccessor, MetadataAccessor, OwnedTable,
+            DataAccessor, MetadataAccessor, OwnedTable, TableRef,
         },
         map::IndexSet,
         proof::ProofError,
@@ -139,6 +139,10 @@ where
 
         columns
     }
+
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        IndexSet::from_iter([self.table.table_ref])
+    }
 }
 
 /// Alias for a filter expression with a honest prover.

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

@@ -153,6 +153,10 @@ fn we_can_correctly_fetch_all_the_referenced_columns() {
             )
         ])
     );
+
+    let ref_tables = provable_ast.get_table_references();
+
+    assert_eq!(ref_tables, IndexSet::from_iter([table_ref]));
 }
 
 #[test]

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

@@ -7,7 +7,7 @@ use crate::{
                 aggregate_columns, compare_indexes_by_owned_columns, AggregatedColumns,
             },
             Column, ColumnField, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor,
-            MetadataAccessor, OwnedTable,
+            MetadataAccessor, OwnedTable, TableRef,
         },
         map::IndexSet,
         proof::ProofError,
@@ -202,6 +202,10 @@ impl<C: Commitment> ProofPlan<C> for GroupByExec<C> {
 
         columns
     }
+
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        IndexSet::from_iter([self.table.table_ref])
+    }
 }
 
 impl<C: Commitment> ProverEvaluate<C::Scalar> for GroupByExec<C> {

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

@@ -3,7 +3,7 @@ use crate::{
         commitment::Commitment,
         database::{
             Column, ColumnField, ColumnRef, CommitmentAccessor, DataAccessor, MetadataAccessor,
-            OwnedTable,
+            OwnedTable, TableRef,
         },
         map::IndexSet,
         proof::ProofError,
@@ -92,6 +92,10 @@ impl<C: Commitment> ProofPlan<C> for ProjectionExec<C> {
         });
         columns
     }
+
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        IndexSet::from_iter([self.table.table_ref])
+    }
 }
 
 impl<C: Commitment> ProverEvaluate<C::Scalar> for ProjectionExec<C> {

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

@@ -102,6 +102,10 @@ fn we_can_correctly_fetch_all_the_referenced_columns() {
             ),
         ])
     );
+
+    let ref_tables = provable_ast.get_table_references();
+
+    assert_eq!(ref_tables, IndexSet::from_iter([table_ref]));
 }
 
 #[test]