fix: docs + re-exports

sdk-derive/src/input.rs

@@ -218,7 +218,7 @@ pub fn impl_flatten_and_raw_input(ast: &DeriveInput) -> TokenStream {
             }
         }
 
-        impl #old_impl_generics axiom_sdk::compute::AxiomComputeInput for #raw_circuit_name_ident #old_ty_generics {
+        impl #old_impl_generics axiom_sdk::axiom::AxiomComputeInput for #raw_circuit_name_ident #old_ty_generics {
             type LogicInput = #raw_circuit_name_ident #old_ty_generics;
             type Input<T: Copy> = #name #ty_generics;
         }

sdk-derive/src/lib.rs

@@ -11,6 +11,10 @@ mod input;
 
 #[proc_macro_attribute]
 #[allow(non_snake_case)]
+/// Derive the `AxiomComputeInput` trait for a struct.
+/// The struct must be named `_Input`. Ex: `ExampleInput`.
+/// All the fields of the struct must implement `RawInput` from `axiom_circuit::input::raw_input`,
+/// which has already been implemented for most primitive Rust types and Ethers types.
 pub fn AxiomComputeInput(_args: TokenStream, input: TokenStream) -> TokenStream {
     let input_clone = input.clone();
     let ast = parse_macro_input!(input_clone as ItemStruct);

sdk/examples/account_age.rs

@@ -1,15 +1,15 @@
 use std::fmt::Debug;
 
 use axiom_sdk::{
-    axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
     cmd::run_cli,
     ethers::types::Address,
     halo2_base::{
-        gates::{GateChip, GateInstructions, RangeInstructions},
+        gates::{GateInstructions, RangeInstructions},
         AssignedValue,
     },
     subquery::AccountField,
-    AxiomComputeInput, Fr,
+    Fr,
 };
 
 #[AxiomComputeInput]
@@ -23,10 +23,12 @@ impl AxiomComputeFn for AccountAgeInput {
         api: &mut AxiomAPI,
         assigned_inputs: AccountAgeCircuitInput<AssignedValue<Fr>>,
     ) -> Vec<AxiomResult> {
-        let gate = GateChip::new();
         let zero = api.ctx().load_zero();
         let one = api.ctx().load_constant(Fr::one());
-        let prev_block = gate.sub(api.ctx(), assigned_inputs.claimed_block_number, one);
+        let prev_block = api
+            .range
+            .gate()
+            .sub(api.ctx(), assigned_inputs.claimed_block_number, one);
 
         let account_prev_block = api.get_account(prev_block, assigned_inputs.addr);
         let prev_nonce = account_prev_block.call(AccountField::Nonce);

sdk/examples/keccak.rs

@@ -1,10 +1,10 @@
 use std::fmt::Debug;
 
 use axiom_sdk::{
-    axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
     cmd::run_cli,
     halo2_base::AssignedValue,
-    AxiomComputeInput, Fr,
+    Fr,
 };
 
 #[AxiomComputeInput]

sdk/examples/quickstart.rs

@@ -1,12 +1,12 @@
 use std::{fmt::Debug, str::FromStr};
 
 use axiom_sdk::{
-    axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
     cmd::run_cli,
     ethers::types::{Address, H256},
     halo2_base::AssignedValue,
     subquery::{AccountField, HeaderField, TxField},
-    AxiomComputeInput, Fr, HiLo,
+    Fr, HiLo,
 };
 
 #[AxiomComputeInput]

sdk/examples/rlc.rs

@@ -2,13 +2,13 @@ use std::fmt::Debug;
 
 use axiom_circuit::axiom_eth::rlc::circuit::builder::RlcCircuitBuilder;
 use axiom_sdk::{
-    axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
+    axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
     cmd::run_cli,
     halo2_base::{
         gates::{GateInstructions, RangeChip, RangeInstructions},
         AssignedValue,
     },
-    AxiomComputeInput, Fr,
+    Fr,
 };
 
 #[AxiomComputeInput]
@@ -43,9 +43,7 @@ impl AxiomComputeFn for RlcInput {
         let gate = range.gate();
         let rlc_chip = builder.rlc_chip(gate);
         let (ctx, rlc_ctx) = builder.rlc_ctx_pair();
-
         gate.add(ctx, payload[0], payload[1]);
-        // builder.base.main(1).debug_assert_false();
         let x = vec![
             ctx.load_constant(Fr::from(1)),
             ctx.load_constant(Fr::from(2)),

sdk/src/api.rs

@@ -24,11 +24,14 @@ use crate::{
     Fr,
 };
 
-/// Axiom Circuit API
+/// Axiom Circuit API for making both subquery calls (e.g. `get_account`, `get_header`, etc.) and for more general ZK primitives (e.g. `add`, `mul`, etc.).
 pub struct AxiomAPI<'a> {
+    /// The `halo2-lib` struct used to construct the circuit
     pub builder: &'a mut RlcCircuitBuilder<Fr>,
+    /// The main chip for ZK primitives
     pub range: &'a RangeChip<Fr>,
-    pub subquery_caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
+    /// The struct that manages all subquery calls
+    subquery_caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
 }
 
 impl<'a> AxiomAPI<'a> {

sdk/src/compute.rs

@@ -24,7 +24,7 @@ use serde::{de::DeserializeOwned, Serialize};
 
 use crate::{api::AxiomAPI, Fr};
 
-/// A trait for the input to an Axiom Compute function
+/// A trait for specifying the input to an Axiom Compute function
 pub trait AxiomComputeInput: Clone + Default + Debug {
     /// The type of the native input (ie. Rust types) to the compute function
     type LogicInput: Clone + Debug + Serialize + DeserializeOwned + Into<Self::Input<Fr>>;
@@ -34,6 +34,7 @@ pub trait AxiomComputeInput: Clone + Default + Debug {
 
 /// A trait for specifying an Axiom Compute function
 pub trait AxiomComputeFn: AxiomComputeInput {
+    /// An optional type for the first phase payload -- only needed if you are using `compute_phase1`
     type FirstPhasePayload: Clone + Default = ();
 
     /// Axiom Compute function
@@ -63,6 +64,7 @@ pub trait AxiomComputeFn: AxiomComputeInput {
 }
 
 #[derive(Debug, Clone)]
+/// Helper struct that contains all the necessary metadata and inputs to run an Axiom Compute function
 pub struct AxiomCompute<A: AxiomComputeFn> {
     provider: Option<Provider<Http>>,
     params: Option<AxiomCircuitParams>,
@@ -124,57 +126,69 @@ where
     A::Input<Fr>: Default + Debug,
     A::Input<AssignedValue<Fr>>: Debug,
 {
+    /// Create a new AxiomCompute instance
     pub fn new() -> Self {
         Self::default()
     }
 
+    /// Set the provider for the AxiomCompute instance
     pub fn set_provider(&mut self, provider: Provider<Http>) {
         self.provider = Some(provider);
     }
 
+    /// Set the params for the AxiomCompute instance
     pub fn set_params(&mut self, params: AxiomCircuitParams) {
         self.params = Some(params);
     }
 
+    /// Set the inputs for the AxiomCompute instance
     pub fn set_inputs(&mut self, input: A::LogicInput) {
         self.input = Some(input);
     }
 
+    /// Set the pinning for the AxiomCompute instance
     pub fn set_pinning(&mut self, pinning: AxiomCircuitPinning) {
         self.pinning = Some(pinning);
     }
 
+    /// Use the given provider for the AxiomCompute instance
     pub fn use_provider(mut self, provider: Provider<Http>) -> Self {
         self.set_provider(provider);
         self
     }
 
+    /// Use the given params for the AxiomCompute instance
     pub fn use_params(mut self, params: AxiomCircuitParams) -> Self {
         self.set_params(params);
         self
     }
 
+    /// Use the given inputs for the AxiomCompute instance
     pub fn use_inputs(mut self, input: A::LogicInput) -> Self {
         self.set_inputs(input);
         self
     }
 
+    /// Use the given pinning for the AxiomCompute instance
     pub fn use_pinning(mut self, pinning: AxiomCircuitPinning) -> Self {
         self.set_pinning(pinning);
         self
     }
 
+    /// Check that all the necessary configurations are set
     fn check_all_set(&self) {
         assert!(self.provider.is_some());
         assert!(self.pinning.is_some());
         assert!(self.input.is_some());
     }
 
+    /// Check that the provider and params are set
     fn check_provider_and_params_set(&self) {
         assert!(self.provider.is_some());
         assert!(self.params.is_some());
     }
 
+    /// Run the mock prover
     pub fn mock(&self) {
         self.check_provider_and_params_set();
         let provider = self.provider.clone().unwrap();
@@ -183,6 +197,7 @@ where
         mock::<Http, Self>(provider, params, converted_input);
     }
 
+    /// Run key generation and return the proving and verifying keys, and the circuit pinning
     pub fn keygen(
         &self,
     ) -> (
@@ -196,20 +211,23 @@ where
         keygen::<Http, Self>(provider, params, None)
     }
 
+    /// Run the prover and return the resulting snark
     pub fn prove(&self, pk: ProvingKey<G1Affine>) -> Snark {
         self.check_all_set();
         let provider = self.provider.clone().unwrap();
         let converted_input = self.input.clone().map(|input| input.into());
         prove::<Http, Self>(provider, self.pinning.clone().unwrap(), converted_input, pk)
     }
 
+    /// Run the prover and return the outputs needed to make an on-chain compute query
     pub fn run(&self, pk: ProvingKey<G1Affine>) -> AxiomV2CircuitOutput {
         self.check_all_set();
         let provider = self.provider.clone().unwrap();
         let converted_input = self.input.clone().map(|input| input.into());
         run::<Http, Self>(provider, self.pinning.clone().unwrap(), converted_input, pk)
     }
 
+    /// Returns an [AxiomCircuit] instance, for functions that expect the halo2 circuit trait
     pub fn circuit(&self) -> AxiomCircuit<Fr, Http, Self> {
         self.check_provider_and_params_set();
         let provider = self.provider.clone().unwrap();