feat: rlc + keccak in compute circuit

sdk-derive/src/lib.rs

@@ -14,7 +14,7 @@ mod input;
 pub fn AxiomComputeInput(_args: TokenStream, input: TokenStream) -> TokenStream {
     let input_clone = input.clone();
     let ast = parse_macro_input!(input_clone as ItemStruct);
-    let name = ast.ident.clone();
+    // let name = ast.ident.clone();
     let new_struct = impl_new_struct(&ast);
     if let Err(err) = new_struct {
         return err.into();
@@ -28,11 +28,6 @@ pub fn AxiomComputeInput(_args: TokenStream, input: TokenStream) -> TokenStream
         #ast
         #new_struct
         #flatten
-
-        fn main() {
-            env_logger::init();
-            axiom_sdk::cmd::run_cli::<#name>();
-        }
     }
     .into()
 }

sdk/data/keccak_input.json

@@ -0,0 +1,5 @@
+{
+    "a": 2,
+    "b": 3,
+    "c": 4
+}

sdk/examples/account_age.rs

@@ -2,6 +2,7 @@ use std::fmt::Debug;
 
 use axiom_sdk::{
     axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    cmd::run_cli,
     halo2_base::{
         gates::{GateChip, GateInstructions, RangeInstructions},
         AssignedValue,
@@ -44,3 +45,8 @@ impl AxiomComputeFn for AccountAgeInput {
         ]
     }
 }
+
+fn main() {
+    env_logger::init();
+    run_cli::<AccountAgeInput>();
+}

sdk/examples/keccak.rs

@@ -0,0 +1,36 @@
+use std::{fmt::Debug, str::FromStr};
+
+use axiom_sdk::{
+    axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    cmd::run_cli,
+    halo2_base::AssignedValue,
+    subquery::{AccountField, HeaderField, TxField},
+    AxiomComputeInput, Fr, HiLo,
+};
+use ethers::types::{Address, H256};
+
+#[AxiomComputeInput]
+pub struct KeccakInput {
+    pub a: u64,
+    pub b: u64,
+    pub c: u64,
+}
+
+impl AxiomComputeFn for KeccakInput {
+    fn compute(
+        api: &mut AxiomAPI,
+        assigned_inputs: KeccakCircuitInput<AssignedValue<Fr>>,
+    ) -> Vec<AxiomResult> {
+        let res = api.keccak_fix_len(vec![
+            assigned_inputs.a,
+            assigned_inputs.b,
+            assigned_inputs.c,
+        ]);
+        vec![res.into()]
+    }
+}
+
+fn main() {
+    env_logger::init();
+    run_cli::<KeccakInput>();
+}

sdk/examples/quickstart.rs

@@ -2,6 +2,7 @@ use std::{fmt::Debug, str::FromStr};
 
 use axiom_sdk::{
     axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    cmd::run_cli,
     halo2_base::AssignedValue,
     subquery::{AccountField, HeaderField, TxField},
     AxiomComputeInput, Fr, HiLo,
@@ -99,3 +100,8 @@ impl AxiomComputeFn for QuickstartInput {
         vec![]
     }
 }
+
+fn main() {
+    env_logger::init();
+    run_cli::<QuickstartInput>();
+}

sdk/src/api.rs

@@ -3,7 +3,8 @@ use std::sync::{Arc, Mutex};
 use axiom_circuit::{
     axiom_codec::HiLo,
     axiom_eth::{
-        halo2_base::{gates::RangeChip, AssignedValue, Context},
+        halo2_base::{gates::RangeChip, safe_types::SafeTypeChip, AssignedValue, Context},
+        keccak::promise::{KeccakFixLenCall, KeccakVarLenCall},
         rlc::circuit::builder::RlcCircuitBuilder,
     },
     subquery::caller::SubqueryCaller,
@@ -112,4 +113,61 @@ impl<'a> AxiomAPI<'a> {
         let ctx = self.builder.base.main(0);
         get_tx(ctx, self.subquery_caller.clone(), block_number, tx_idx)
     }
+
+    pub fn keccak_fix_len(&mut self, bytes: Vec<AssignedValue<Fr>>) -> HiLo<AssignedValue<Fr>> {
+        let ctx = self.builder.base.main(0);
+        let subquery_caller = self.subquery_caller.clone();
+        let mut subquery_caller = subquery_caller.lock().unwrap();
+
+        let safe_type_chip = SafeTypeChip::new(self.range);
+        let len = bytes.len();
+        let bytes = safe_type_chip.raw_to_fix_len_bytes_vec(ctx, bytes, len);
+
+        subquery_caller.keccak(ctx, KeccakFixLenCall::new(bytes))
+    }
+
+    pub fn keccak_var_len(
+        &mut self,
+        bytes: Vec<AssignedValue<Fr>>,
+        len: AssignedValue<Fr>,
+        max_len: usize,
+    ) -> HiLo<AssignedValue<Fr>> {
+        let ctx = self.builder.base.main(0);
+        let subquery_caller = self.subquery_caller.clone();
+        let mut subquery_caller = subquery_caller.lock().unwrap();
+
+        let safe_type_chip = SafeTypeChip::new(self.range);
+        let bytes = safe_type_chip.raw_to_var_len_bytes_vec(ctx, bytes, len, max_len);
+
+        subquery_caller.keccak(ctx, KeccakVarLenCall::new(bytes, 0))
+    }
+
+    pub fn keccak_fix_len_unsafe(
+        &mut self,
+        bytes: Vec<AssignedValue<Fr>>,
+    ) -> HiLo<AssignedValue<Fr>> {
+        let ctx = self.builder.base.main(0);
+        let subquery_caller = self.subquery_caller.clone();
+        let mut subquery_caller = subquery_caller.lock().unwrap();
+
+        let len = bytes.len();
+        let bytes = SafeTypeChip::unsafe_to_fix_len_bytes_vec(bytes, len);
+
+        subquery_caller.keccak(ctx, KeccakFixLenCall::new(bytes))
+    }
+
+    pub fn keccak_var_len_unsafe(
+        &mut self,
+        bytes: Vec<AssignedValue<Fr>>,
+        len: AssignedValue<Fr>,
+        max_len: usize,
+    ) -> HiLo<AssignedValue<Fr>> {
+        let ctx = self.builder.base.main(0);
+        let subquery_caller = self.subquery_caller.clone();
+        let mut subquery_caller = subquery_caller.lock().unwrap();
+
+        let bytes = SafeTypeChip::unsafe_to_var_len_bytes_vec(bytes, len, max_len);
+
+        subquery_caller.keccak(ctx, KeccakVarLenCall::new(bytes, 0))
+    }
 }

sdk/src/cmd.rs

@@ -11,9 +11,11 @@ use axiom_circuit::{
         halo2_base::{gates::circuit::BaseCircuitParams, AssignedValue},
         halo2_proofs::{plonk::ProvingKey, SerdeFormat},
         halo2curves::bn256::G1Affine,
+        rlc::circuit::RlcCircuitParams,
+        utils::keccak::decorator::RlcKeccakCircuitParams,
     },
     scaffold::AxiomCircuit,
-    types::AxiomCircuitPinning,
+    types::{AxiomCircuitParams, AxiomCircuitPinning},
 };
 pub use clap::Parser;
 use clap::Subcommand;
@@ -62,6 +64,11 @@ pub struct Cli {
     pub input_path: Option<PathBuf>,
     #[arg(short, long = "data-path")]
     pub data_path: Option<PathBuf>,
+    //Advanced options
+    #[arg(long = "keccak-rows")]
+    pub keccak_rows_per_round: Option<u32>,
+    #[arg(long = "rlc-columns")]
+    pub rlc_columns: Option<u32>,
 }
 
 pub fn run_cli<A: AxiomComputeFn>()
@@ -102,33 +109,46 @@ where
     let provider = Provider::<Http>::try_from(provider_uri).unwrap();
     let data_path = cli.data_path.unwrap_or_else(|| PathBuf::from("data"));
 
+    let base_params = BaseCircuitParams {
+        k: cli.degree.unwrap() as usize,
+        num_advice_per_phase: vec![4],
+        num_fixed: 1,
+        num_lookup_advice_per_phase: vec![1],
+        lookup_bits: Some(11),
+        num_instance_columns: 1,
+    };
+
+    let keccak_rows_per_round = cli.keccak_rows_per_round.unwrap_or(0) as usize;
+    let rlc_columns = cli.rlc_columns.unwrap_or(0) as usize;
+
+    let params = if keccak_rows_per_round > 0 {
+        AxiomCircuitParams::Keccak(RlcKeccakCircuitParams {
+            keccak_rows_per_round,
+            rlc: RlcCircuitParams {
+                base: base_params,
+                num_rlc_columns: rlc_columns,
+            },
+        })
+    } else if rlc_columns > 0 {
+        AxiomCircuitParams::Rlc(RlcCircuitParams {
+            base: base_params,
+            num_rlc_columns: rlc_columns,
+        })
+    } else {
+        AxiomCircuitParams::Base(base_params)
+    };
+
     match cli.command {
         SnarkCmd::Mock => {
-            let params = BaseCircuitParams {
-                k: cli.degree.unwrap() as usize,
-                num_advice_per_phase: vec![4],
-                num_fixed: 1,
-                num_lookup_advice_per_phase: vec![1],
-                lookup_bits: Some(11),
-                num_instance_columns: 1,
-            };
             AxiomCompute::<A>::new()
                 .use_inputs(input)
                 .use_params(params)
                 .use_provider(provider)
                 .mock();
         }
         SnarkCmd::Keygen => {
-            let params = BaseCircuitParams {
-                k: cli.degree.unwrap() as usize,
-                num_advice_per_phase: vec![4],
-                num_fixed: 1,
-                num_lookup_advice_per_phase: vec![1],
-                lookup_bits: Some(11),
-                num_instance_columns: 1,
-            };
             let circuit = AxiomCompute::<A>::new()
-                .use_params(params.clone())
+                .use_params(params)
                 .use_provider(provider);
             let (_, pkey, pinning) = circuit.keygen();
             let pk_path = data_path.join(PathBuf::from("pk.bin"));

sdk/src/compute.rs

@@ -5,10 +5,7 @@ use std::{
 
 use axiom_circuit::{
     axiom_eth::{
-        halo2_base::{
-            gates::{circuit::BaseCircuitParams, RangeChip},
-            AssignedValue,
-        },
+        halo2_base::{gates::RangeChip, AssignedValue},
         halo2_proofs::plonk::{ProvingKey, VerifyingKey},
         halo2curves::bn256::G1Affine,
         rlc::circuit::builder::RlcCircuitBuilder,
@@ -34,17 +31,33 @@ pub trait AxiomComputeInput: Clone + Default + Debug {
 }
 
 pub trait AxiomComputeFn: AxiomComputeInput {
-    // type Provider: JsonRpcClient + Clone = Self::ProviderType;
+    type FirstPhasePayload: Clone + Default = ();
+
     fn compute(
         api: &mut AxiomAPI,
         assigned_inputs: Self::Input<AssignedValue<Fr>>,
     ) -> Vec<AxiomResult>;
+
+    fn compute_phase0(
+        api: &mut AxiomAPI,
+        assigned_inputs: Self::Input<AssignedValue<Fr>>,
+    ) -> (Vec<AxiomResult>, Self::FirstPhasePayload) {
+        (Self::compute(api, assigned_inputs), Default::default())
+    }
+
+    #[allow(unused_variables)]
+    fn compute_phase1(
+        builder: &mut RlcCircuitBuilder<Fr>,
+        range: &RangeChip<Fr>,
+        payload: Self::FirstPhasePayload,
+    ) {
+    }
 }
 
 #[derive(Debug, Clone)]
 pub struct AxiomCompute<A: AxiomComputeFn> {
     provider: Option<Provider<Http>>,
-    params: Option<BaseCircuitParams>,
+    params: Option<AxiomCircuitParams>,
     pinning: Option<AxiomCircuitPinning>,
     input: Option<A::LogicInput>,
 }
@@ -67,16 +80,17 @@ where
 {
     type InputValue = A::Input<Fr>;
     type InputWitness = A::Input<AssignedValue<Fr>>;
+    type FirstPhasePayload = A::FirstPhasePayload;
 
     fn virtual_assign_phase0(
         builder: &mut RlcCircuitBuilder<Fr>,
         range: &RangeChip<Fr>,
         subquery_caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
         callback: &mut Vec<HiLo<AssignedValue<Fr>>>,
         assigned_inputs: Self::InputWitness,
-    ) {
+    ) -> <A as AxiomComputeFn>::FirstPhasePayload {
         let mut api = AxiomAPI::new(builder, range, subquery_caller);
-        let result = A::compute(&mut api, assigned_inputs);
+        let (result, payload) = A::compute_phase0(&mut api, assigned_inputs);
         let hilo_output = result
             .into_iter()
             .map(|result| match result {
@@ -85,6 +99,15 @@ where
             })
             .collect::<Vec<_>>();
         callback.extend(hilo_output);
+        payload
+    }
+
+    fn virtual_assign_phase1(
+        builder: &mut RlcCircuitBuilder<Fr>,
+        range: &RangeChip<Fr>,
+        payload: Self::FirstPhasePayload,
+    ) {
+        A::compute_phase1(builder, range, payload);
     }
 }
 
@@ -101,7 +124,7 @@ where
         self.provider = Some(provider);
     }
 
-    pub fn set_params(&mut self, params: BaseCircuitParams) {
+    pub fn set_params(&mut self, params: AxiomCircuitParams) {
         self.params = Some(params);
     }
 
@@ -118,7 +141,7 @@ where
         self
     }
 
-    pub fn use_params(mut self, params: BaseCircuitParams) -> Self {
+    pub fn use_params(mut self, params: AxiomCircuitParams) -> Self {
         self.set_params(params);
         self
     }
@@ -149,7 +172,7 @@ where
         let provider = self.provider.clone().unwrap();
         let params = self.params.clone().unwrap();
         let converted_input = self.input.clone().map(|input| input.into());
-        mock::<Http, Self>(provider, AxiomCircuitParams::Base(params), converted_input);
+        mock::<Http, Self>(provider, params, converted_input);
     }
 
     pub fn keygen(
@@ -162,7 +185,7 @@ where
         self.check_provider_and_params_set();
         let provider = self.provider.clone().unwrap();
         let params = self.params.clone().unwrap();
-        keygen::<Http, Self>(provider, AxiomCircuitParams::Base(params), None)
+        keygen::<Http, Self>(provider, params, None)
     }
 
     pub fn prove(&self, pk: ProvingKey<G1Affine>) -> Snark {
@@ -183,7 +206,7 @@ where
         self.check_provider_and_params_set();
         let provider = self.provider.clone().unwrap();
         let params = self.params.clone().unwrap();
-        AxiomCircuit::new(provider, AxiomCircuitParams::Base(params))
+        AxiomCircuit::new(provider, params)
     }
 }