Skip to content

Commit

Permalink
feat: axiom-client-sdk cli
Browse files Browse the repository at this point in the history
  • Loading branch information
rpalakkal committed Jan 10, 2024
1 parent b4d84ba commit 392142a
Show file tree
Hide file tree
Showing 22 changed files with 272 additions and 219 deletions.
4 changes: 2 additions & 2 deletions axiom-client-derive/src/input.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ pub fn impl_new_struct(ast: &ItemStruct) -> Result<TokenStream, TokenStream> {
.zip(field_types.iter())
.map(|(name, field_type)| {
quote! {
#name: <#field_type as axiom_client::input::raw_input::RawInput<crate::Fr>>::FEType<T>,
#name: <#field_type as axiom_client::input::raw_input::RawInput<axiom_client_sdk::Fr>>::FEType<T>,
}
})
.collect();
Expand Down Expand Up @@ -218,7 +218,7 @@ pub fn impl_flatten_and_raw_input(ast: &DeriveInput) -> TokenStream {
}
}

impl #old_impl_generics crate::compute::AxiomComputeInput for #raw_circuit_name_ident #old_ty_generics {
impl #old_impl_generics axiom_client_sdk::compute::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;
}
Expand Down
5 changes: 5 additions & 0 deletions axiom-client-derive/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +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 new_struct = impl_new_struct(&ast);
if let Err(err) = new_struct {
return err.into();
Expand All @@ -27,6 +28,10 @@ pub fn AxiomComputeInput(_args: TokenStream, input: TokenStream) -> TokenStream
#ast
#new_struct
#flatten

fn main() {
axiom_client_sdk::cmd::run_cli::<#name>();
}
}
.into()
}
Expand Down
3 changes: 3 additions & 0 deletions axiom-client-sdk/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,3 +12,6 @@ ethers = "2.0"
anyhow = "1.0.75"
serde = { version = "1.0", features = ["derive"] }
lazy_static = "1.4.0"
clap = {version = "4.4.7", features = ["derive"]}
clap-num = "=1.0.2"
serde_json = "1.0.111"
4 changes: 4 additions & 0 deletions axiom-client-sdk/data/account_age_input.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"addr": "0x897dDbe14c9C7736EbfDC58461355697FbF70048",
"claimed_block_number": 9173677
}
8 changes: 8 additions & 0 deletions axiom-client-sdk/data/quickstart_input.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
{
"addr": "0x8dde5d4a8384f403f888e1419672d94c570440c9",
"block": 9730000,
"tx_block_number": 9728956,
"tx_idx": 10,
"slot": "0x0000000000000000000000000000000000000000000000000000000000000002",
"mapping_slot": "0x0000000000000000000000000000000000000000000000000000000000000001"
}
Original file line number Diff line number Diff line change
@@ -1,21 +1,16 @@
use std::fmt::Debug;

use axiom_client::{
axiom_eth::halo2_base::{
use axiom_client_sdk::{
axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
halo2_base::{
gates::{GateChip, GateInstructions, RangeInstructions},
AssignedValue,
},
subquery::account::AccountField,
subquery::AccountField,
AxiomComputeInput, Fr,
};
use axiom_client_derive::AxiomComputeInput;
use ethers::types::Address;

use crate::{
api::AxiomAPI,
compute::{AxiomComputeFn, AxiomResult},
Fr,
};

#[AxiomComputeInput]
pub struct AccountAgeInput {
pub addr: Address,
Expand All @@ -24,7 +19,7 @@ pub struct AccountAgeInput {

impl AxiomComputeFn for AccountAgeInput {
fn compute(
api: &mut AxiomAPI<Self::Provider>,
api: &mut AxiomAPI,
assigned_inputs: AccountAgeCircuitInput<AssignedValue<Fr>>,
) -> Vec<AxiomResult> {
let gate = GateChip::new();
Expand All @@ -49,20 +44,3 @@ impl AxiomComputeFn for AccountAgeInput {
]
}
}

#[cfg(test)]
mod tests {
use std::str::FromStr;

use super::*;
use crate::axiom_compute_tests;

fn inputs() -> AccountAgeInput {
AccountAgeInput {
addr: Address::from_str("0x897dDbe14c9C7736EbfDC58461355697FbF70048").unwrap(),
claimed_block_number: 9173677,
}
}

axiom_compute_tests!(AccountAgeInput, inputs, 12);
}
Original file line number Diff line number Diff line change
@@ -1,19 +1,13 @@
use std::{fmt::Debug, str::FromStr};

use axiom_client::{
axiom_codec::HiLo,
axiom_eth::halo2_base::AssignedValue,
subquery::{account::AccountField, header::HeaderField, tx::TxField},
use axiom_client_sdk::{
axiom::{AxiomAPI, AxiomComputeFn, AxiomResult},
halo2_base::AssignedValue,
subquery::{AccountField, HeaderField, TxField},
AxiomComputeInput, Fr, HiLo,
};
use axiom_client_derive::AxiomComputeInput;
use ethers::types::{Address, H256};

use crate::{
api::AxiomAPI,
compute::{AxiomComputeFn, AxiomResult},
Fr,
};

#[AxiomComputeInput]
pub struct QuickstartInput {
pub block: u64,
Expand All @@ -26,7 +20,7 @@ pub struct QuickstartInput {

impl AxiomComputeFn for QuickstartInput {
fn compute(
api: &mut AxiomAPI<Self::Provider>,
api: &mut AxiomAPI,
assigned_inputs: QuickstartCircuitInput<AssignedValue<Fr>>,
) -> Vec<AxiomResult> {
// fetch block header data
Expand Down Expand Up @@ -105,24 +99,3 @@ impl AxiomComputeFn for QuickstartInput {
vec![]
}
}

#[cfg(test)]
mod tests {
use std::str::FromStr;

use super::*;
use crate::axiom_compute_tests;

fn inputs() -> QuickstartInput {
QuickstartInput {
addr: Address::from_str("0x8dde5d4a8384f403f888e1419672d94c570440c9").unwrap(),
block: 9730000,
tx_block_number: 9728956,
tx_idx: 10,
slot: H256::from_low_u64_be(2),
mapping_slot: H256::from_low_u64_be(1),
}
}

axiom_compute_tests!(QuickstartInput, inputs, 12);
}
11 changes: 11 additions & 0 deletions axiom-client-sdk/readme.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
# axiom-client-sdk

## Usage

See [./examples/account_age.rs] for an example Axiom compute circuit. To run the `account_age` circuit:

```
cargo run --example account_age -- --input data/quickstart_input.json -k 12 -p <PROVIDER_URI> <CMD>
```

where `PROVIDER_URI` is a JSON-RPC URL, and `CMD` is `mock`, `prove`, `keygen`, or `run`.
24 changes: 12 additions & 12 deletions axiom-client-sdk/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use axiom_client::{
subquery::caller::SubqueryCaller,
utils::{from_hi_lo, to_hi_lo},
};
use ethers::providers::JsonRpcClient;
use ethers::providers::Http;

use crate::{
subquery::{
Expand All @@ -23,17 +23,17 @@ use crate::{
Fr,
};

pub struct AxiomAPI<'a, P: JsonRpcClient> {
pub struct AxiomAPI<'a> {
pub builder: &'a mut RlcCircuitBuilder<Fr>,
pub range: &'a RangeChip<Fr>,
pub subquery_caller: Arc<Mutex<SubqueryCaller<P, Fr>>>,
pub subquery_caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
}

impl<'a, P: JsonRpcClient> AxiomAPI<'a, P> {
impl<'a> AxiomAPI<'a> {
pub fn new(
builder: &'a mut RlcCircuitBuilder<Fr>,
range: &'a RangeChip<Fr>,
subquery_caller: Arc<Mutex<SubqueryCaller<P, Fr>>>,
subquery_caller: Arc<Mutex<SubqueryCaller<Http, Fr>>>,
) -> Self {
Self {
builder,
Expand All @@ -42,7 +42,7 @@ impl<'a, P: JsonRpcClient> AxiomAPI<'a, P> {
}
}

pub fn subquery_caller(&self) -> Arc<Mutex<SubqueryCaller<P, Fr>>> {
pub fn subquery_caller(&self) -> Arc<Mutex<SubqueryCaller<Http, Fr>>> {
self.subquery_caller.clone()
}

Expand All @@ -64,12 +64,12 @@ impl<'a, P: JsonRpcClient> AxiomAPI<'a, P> {
&mut self,
block_number: AssignedValue<Fr>,
addr: AssignedValue<Fr>,
) -> Account<P> {
) -> Account {
let ctx = self.builder.base.main(0);
get_account(ctx, self.subquery_caller.clone(), block_number, addr)
}

pub fn get_header(&mut self, block_number: AssignedValue<Fr>) -> Header<P> {
pub fn get_header(&mut self, block_number: AssignedValue<Fr>) -> Header {
let ctx = self.builder.base.main(0);
get_header(ctx, self.subquery_caller.clone(), block_number)
}
Expand All @@ -79,7 +79,7 @@ impl<'a, P: JsonRpcClient> AxiomAPI<'a, P> {
block_number: AssignedValue<Fr>,
addr: AssignedValue<Fr>,
mapping_slot: HiLo<AssignedValue<Fr>>,
) -> SolidityMapping<P> {
) -> SolidityMapping {
let ctx = self.builder.base.main(0);
get_mapping(
ctx,
Expand All @@ -94,7 +94,7 @@ impl<'a, P: JsonRpcClient> AxiomAPI<'a, P> {
&mut self,
block_number: AssignedValue<Fr>,
tx_idx: AssignedValue<Fr>,
) -> Receipt<P> {
) -> Receipt {
let ctx = self.builder.base.main(0);
get_receipt(ctx, self.subquery_caller.clone(), block_number, tx_idx)
}
Expand All @@ -103,12 +103,12 @@ impl<'a, P: JsonRpcClient> AxiomAPI<'a, P> {
&mut self,
block_number: AssignedValue<Fr>,
addr: AssignedValue<Fr>,
) -> Storage<P> {
) -> Storage {
let ctx = self.builder.base.main(0);
get_storage(ctx, self.subquery_caller.clone(), block_number, addr)
}

pub fn get_tx(&mut self, block_number: AssignedValue<Fr>, tx_idx: AssignedValue<Fr>) -> Tx<P> {
pub fn get_tx(&mut self, block_number: AssignedValue<Fr>, tx_idx: AssignedValue<Fr>) -> Tx {
let ctx = self.builder.base.main(0);
get_tx(ctx, self.subquery_caller.clone(), block_number, tx_idx)
}
Expand Down
112 changes: 112 additions & 0 deletions axiom-client-sdk/src/cmd.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
use std::{env, fmt::Debug, fs, path::PathBuf};

use axiom_client::axiom_eth::halo2_base::{gates::circuit::BaseCircuitParams, AssignedValue};
pub use clap::Parser;
use clap::Subcommand;
use ethers::providers::{Http, Provider};

use crate::{
compute::{AxiomCompute, AxiomComputeFn},
Fr,
};

#[derive(Clone, Copy, Debug, Subcommand)]
pub enum SnarkCmd {
/// Run the mock prover
Mock,
/// Generate new proving & verifying keys
Keygen,
/// Generate a new proof
Prove,
/// Generate an Axiom compute query
Run,
}

impl std::fmt::Display for SnarkCmd {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Mock => write!(f, "mock"),
Self::Keygen => write!(f, "keygen"),
Self::Prove => write!(f, "prove"),
Self::Run => write!(f, "run"),
}
}
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
/// Command-line helper for building Axiom compute circuits
pub struct Cli {
#[command(subcommand)]
pub command: SnarkCmd,
#[arg(short = 'k', long = "degree")]
pub degree: u32,
#[arg(short = 'p', long = "provider")]
pub provider: Option<String>,
#[arg(short, long = "input")]
pub input_path: Option<PathBuf>,
}

pub fn run_cli<A: AxiomComputeFn>()
where
A::Input<Fr>: Default + Debug,
A::Input<AssignedValue<Fr>>: Debug,
{
let cli = Cli::parse();
match cli.command {
SnarkCmd::Mock | SnarkCmd::Prove | SnarkCmd::Run => {
if cli.input_path.is_none() {
panic!("The `input_path` argument is required for the selected command.");
}
}
_ => {}
}
let input_path = cli.input_path.unwrap();
let json_str = fs::read_to_string(input_path).expect("Unable to read file");
let input: A::LogicInput = serde_json::from_str(&json_str).expect("Unable to parse JSON");
if cli.provider.is_none() && env::var("PROVIDER_URI").is_err() {
panic!("The `provider` argument is required for the selected command. Either pass it as an argument or set the `PROVIDER_URI` environment variable.");
}
let provider_uri = cli
.provider
.unwrap_or_else(|| env::var("PROVIDER_URI").unwrap());
let provider = Provider::<Http>::try_from(provider_uri).unwrap();

let params = BaseCircuitParams {
k: 12,
num_advice_per_phase: vec![4],
num_fixed: 1,
num_lookup_advice_per_phase: vec![1],
lookup_bits: Some(11),
num_instance_columns: 1,
};
match cli.command {
SnarkCmd::Mock => {
AxiomCompute::<A>::new()
.use_inputs(input)
.use_params(params)
.use_provider(provider)
.mock();
}
SnarkCmd::Keygen => {
AxiomCompute::<A>::new()
.use_params(params)
.use_provider(provider)
.keygen();
}
SnarkCmd::Prove => {
let compute = AxiomCompute::<A>::new()
.use_params(params)
.use_provider(provider);
let (_vk, pk) = compute.keygen();
compute.use_inputs(input).prove(pk);
}
SnarkCmd::Run => {
let compute = AxiomCompute::<A>::new()
.use_params(params)
.use_provider(provider);
let (_vk, pk) = compute.keygen();
compute.use_inputs(input).run(pk);
}
}
}
Loading

0 comments on commit 392142a

Please sign in to comment.