feat: add postprocessing and PostprocessingLevel (#20)

.gitignore

@@ -1,4 +1,5 @@
 .env
+**/*.json
 
 # Rust
 # Generated by Cargo

scripts/count-ethereum-core/src/main.rs

@@ -36,7 +36,7 @@ async fn count_table(
     table_ref: &str,
 ) -> Result<i64, Box<dyn core::error::Error>> {
     let uppercased_table_ref = table_ref.to_uppercase();
-    let query = format!("SELECT * FROM {uppercased_table_ref}");
+    let query = format!("SELECT COUNT(*) FROM {uppercased_table_ref}");
     let table = client
         .query_and_verify(&query, &uppercased_table_ref)
         .await?;

sdk/src/args.rs

@@ -1,5 +1,5 @@
-use crate::SxTClient;
-use clap::Parser;
+use crate::{PostprocessingLevel, SxTClient};
+use clap::{Parser, ValueEnum};
 
 /// Struct to define and parse command-line arguments for Proof of SQL Client.
 ///
@@ -79,6 +79,39 @@ pub struct SdkArgs {
         default_value = "verifier_setup.bin"
     )]
     pub verifier_setup: String,
+
+    /// Level of postprocessing allowed. Default is `Cheap`.
+    #[arg(
+        long,
+        value_name = "POSTPROCESSING_LEVEL",
+        default_value = "cheap",
+        value_enum,
+        help = "Level of postprocessing allowed, default is `Cheap`"
+    )]
+    pub postprocessing_level: PostprocessingLevelArg,
+}
+
+/// Level of postprocessing allowed.
+///
+/// We have it as a separate enum from `PostprocessingLevel` so that the core of the lib won't depend on `clap`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
+pub enum PostprocessingLevelArg {
+    /// No postprocessing allowed
+    None,
+    /// Only cheap postprocessing allowed
+    Cheap,
+    /// All postprocessing allowed
+    All,
+}
+
+impl From<PostprocessingLevelArg> for PostprocessingLevel {
+    fn from(arg: PostprocessingLevelArg) -> Self {
+        match arg {
+            PostprocessingLevelArg::None => PostprocessingLevel::None,
+            PostprocessingLevelArg::Cheap => PostprocessingLevel::Cheap,
+            PostprocessingLevelArg::All => PostprocessingLevel::All,
+        }
+    }
 }
 
 impl From<&SdkArgs> for SxTClient {
@@ -90,5 +123,6 @@ impl From<&SdkArgs> for SxTClient {
             args.sxt_api_key.clone(),
             args.verifier_setup.clone(),
         )
+        .with_postprocessing(args.postprocessing_level.into())
     }
 }

sdk/src/client.rs

@@ -4,7 +4,11 @@ use proof_of_sql::{
     proof_primitive::dory::{
         DoryScalar, DynamicDoryCommitment, DynamicDoryEvaluationProof, VerifierSetup,
     },
-    sql::{parse::QueryExpr, proof::VerifiableQueryResult},
+    sql::{
+        parse::QueryExpr,
+        postprocessing::{apply_postprocessing_steps, OwnedTablePostprocessing},
+        proof::VerifiableQueryResult,
+    },
 };
 use prover::{ProverContextRange, ProverQuery, ProverResponse};
 use reqwest::Client;
@@ -14,6 +18,19 @@ mod prover {
     tonic::include_proto!("sxt.core");
 }
 
+/// Level of postprocessing allowed
+///
+/// Some postprocessing steps are expensive so we allow the user to control the level of postprocessing.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum PostprocessingLevel {
+    /// No postprocessing allowed
+    None,
+    /// Only cheap postprocessing allowed
+    Cheap,
+    /// All postprocessing allowed
+    All,
+}
+
 /// Space and Time (SxT) client
 #[derive(Debug, Clone)]
 pub struct SxTClient {
@@ -34,6 +51,9 @@ pub struct SxTClient {
 
     /// Path to the verifier setup binary file
     pub verifier_setup: String,
+
+    /// Level of postprocessing allowed. Default is [`PostprocessingLevel::Cheap`].
+    pub postprocessing_level: PostprocessingLevel,
 }
 
 impl SxTClient {
@@ -51,9 +71,16 @@ impl SxTClient {
             substrate_node_url,
             sxt_api_key,
             verifier_setup,
+            postprocessing_level: PostprocessingLevel::Cheap,
         }
     }
 
+    /// Set the level of postprocessing allowed
+    pub fn with_postprocessing(mut self, postprocessing_level: PostprocessingLevel) -> Self {
+        self.postprocessing_level = postprocessing_level;
+        self
+    }
+
     /// Query and verify a SQL query
     ///
     /// Run a SQL query and verify the result using Dynamic Dory.
@@ -117,6 +144,22 @@ impl SxTClient {
         let owned_table_result = proof
             .verify(proof_plan, &accessor, &serialized_result, &&verifier_setup)?
             .table;
-        Ok(owned_table_result)
+        // Apply postprocessing steps
+        let postprocessing = query_expr.postprocessing();
+        let is_postprocessing_expensive = postprocessing.iter().any(|step| {
+            matches!(
+                step,
+                OwnedTablePostprocessing::Slice(_) | OwnedTablePostprocessing::GroupBy(_)
+            )
+        });
+        match (self.postprocessing_level, postprocessing.len(), is_postprocessing_expensive) {
+            (_, 0, false) => Ok(owned_table_result),
+            (PostprocessingLevel::All, _, _) | (PostprocessingLevel::Cheap, _, false) => {
+                let transformed_result: OwnedTable<DoryScalar> =
+                    apply_postprocessing_steps(owned_table_result, postprocessing)?;
+                Ok(transformed_result)
+            }
+            _ => Err("Required postprocessing is not allowed. Please examine your query or change `PostprocessingLevel` using `SxTClient::with_postprocessing`".into()),
+        }
     }
 }

sdk/src/lib.rs

@@ -1,9 +1,9 @@
 mod args;
-pub use args::SdkArgs;
+pub use args::{PostprocessingLevelArg, SdkArgs};
 mod auth;
 pub use auth::get_access_token;
 mod client;
-pub use client::SxTClient;
+pub use client::{PostprocessingLevel, SxTClient};
 mod substrate;
 pub use substrate::query_commitments;
 mod sxt_chain_runtime;