balance support

plonky2x/src/frontend/eth/beacon/builder.rs

@@ -1,10 +1,12 @@
 use plonky2::field::extension::Extendable;
 use plonky2::hash::hash_types::RichField;
 
+use super::generators::balance::BeaconValidatorBalanceGenerator;
 use super::generators::validator::BeaconValidatorGenerator;
 use super::vars::{BeaconValidatorVariable, BeaconValidatorsVariable};
 use crate::frontend::builder::CircuitBuilder;
 use crate::frontend::eth::beacon::generators::validators::BeaconValidatorsRootGenerator;
+use crate::frontend::uint::uint256::U256Variable;
 use crate::frontend::vars::Bytes32Variable;
 use crate::prelude::Variable;
 
@@ -59,6 +61,23 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         self.add_simple_generator(&generator);
         generator.validator
     }
+
+    /// Get a validator balance from a given deterministic index.
+    pub fn get_beacon_validator_balance(
+        &mut self,
+        validators: BeaconValidatorsVariable,
+        index: Variable,
+    ) -> U256Variable {
+        let generator = BeaconValidatorBalanceGenerator::new(
+            self,
+            validators.block_root,
+            validators.validators_root,
+            None,
+            Some(index),
+        );
+        self.add_simple_generator(&generator);
+        generator.balance
+    }
 }
 
 #[cfg(test)]
@@ -106,6 +125,8 @@ pub(crate) mod tests {
         (0..1).for_each(|i| {
             let idx = builder.constant::<Variable>(F::from_canonical_u64(i));
             let validator = builder.get_beacon_validator(validators, idx);
+            let balance = builder.get_beacon_validator_balance(validators, idx);
+            builder.watch(&balance, "balance");
             builder.watch(&validator.effective_balance, "hi");
         });
 

plonky2x/src/frontend/eth/beacon/generators/balance.rs

@@ -0,0 +1,143 @@
+use core::marker::PhantomData;
+
+use curta::math::prelude::PrimeField64;
+use plonky2::field::extension::Extendable;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::generator::{GeneratedValues, SimpleGenerator};
+use plonky2::iop::target::Target;
+use plonky2::iop::witness::PartitionWitness;
+use plonky2::plonk::circuit_data::CommonCircuitData;
+use plonky2::util::serialization::{Buffer, IoResult};
+use tokio::runtime::Runtime;
+
+use crate::frontend::builder::CircuitBuilder;
+use crate::frontend::uint::uint256::U256Variable;
+use crate::frontend::vars::{Bytes32Variable, CircuitVariable};
+use crate::prelude::Variable;
+use crate::utils::eth::beacon::BeaconClient;
+use crate::utils::hex;
+
+#[derive(Debug, Clone)]
+pub struct BeaconValidatorBalanceGenerator<F: RichField + Extendable<D>, const D: usize> {
+    client: BeaconClient,
+    block_root: Bytes32Variable,
+    validators_root: Bytes32Variable,
+    deterministic_idx: Option<u64>,
+    dynamic_idx: Option<Variable>,
+    pub balance: U256Variable,
+    _phantom: PhantomData<F>,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> BeaconValidatorBalanceGenerator<F, D> {
+    pub fn new(
+        builder: &mut CircuitBuilder<F, D>,
+        block_root: Bytes32Variable,
+        validators_root: Bytes32Variable,
+        deterministic_idx: Option<u64>,
+        dynamic_idx: Option<Variable>,
+    ) -> Self {
+        Self {
+            client: builder.beacon_client.clone().unwrap(),
+            block_root,
+            validators_root,
+            deterministic_idx,
+            dynamic_idx,
+            balance: builder.init::<U256Variable>(),
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D>
+    for BeaconValidatorBalanceGenerator<F, D>
+{
+    fn id(&self) -> String {
+        "BeaconValidatorBalanceGenerator".to_string()
+    }
+
+    fn dependencies(&self) -> Vec<Target> {
+        let mut targets = Vec::new();
+        targets.extend(self.block_root.targets());
+        targets.extend(self.validators_root.targets());
+        targets
+    }
+
+    fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
+        let block_root = self.block_root.get(witness);
+        let rt = Runtime::new().expect("failed to create tokio runtime");
+        let result = rt.block_on(async {
+            if self.deterministic_idx.is_some() {
+                self.client
+                    .get_validator_balance(hex!(block_root), self.deterministic_idx.unwrap())
+                    .await
+                    .expect("failed to get validator")
+            } else {
+                let idx = self.dynamic_idx.unwrap().get(witness).as_canonical_u64();
+                self.client
+                    .get_validator_balance(hex!(block_root), idx)
+                    .await
+                    .expect("failed to get validator")
+            }
+        });
+        self.balance.set(out_buffer, result);
+    }
+
+    #[allow(unused_variables)]
+    fn serialize(&self, dst: &mut Vec<u8>, common_data: &CommonCircuitData<F, D>) -> IoResult<()> {
+        todo!()
+    }
+
+    #[allow(unused_variables)]
+    fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData<F, D>) -> IoResult<Self> {
+        todo!()
+    }
+}
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use std::env;
+
+    use plonky2::field::goldilocks_field::GoldilocksField;
+    use plonky2::iop::witness::PartialWitness;
+    use plonky2::plonk::config::PoseidonGoldilocksConfig;
+
+    use crate::frontend::builder::CircuitBuilder;
+    use crate::frontend::eth::beacon::generators::validator::BeaconValidatorGenerator;
+    use crate::frontend::vars::Bytes32Variable;
+    use crate::utils::bytes32;
+    use crate::utils::eth::beacon::BeaconClient;
+
+    #[test]
+    #[cfg_attr(feature = "ci", ignore)]
+    fn test_get_validator_generator() {
+        dotenv::dotenv().ok();
+
+        type F = GoldilocksField;
+        type C = PoseidonGoldilocksConfig;
+        const D: usize = 2;
+
+        let consensus_rpc = env::var("CONSENSUS_RPC_URL").unwrap();
+        let client = BeaconClient::new(consensus_rpc);
+
+        let mut builder = CircuitBuilder::<F, D>::new();
+        builder.set_beacon_client(client);
+
+        let block_root = builder.constant::<Bytes32Variable>(bytes32!(
+            "0xe6d6e23b8e07e15b98811579e5f6c36a916b749fd7146d009196beeddc4a6670"
+        ));
+        let validators = builder.get_beacon_validators(block_root);
+        let generator = BeaconValidatorGenerator::new(
+            &mut builder,
+            validators.block_root,
+            validators.validators_root,
+            Some(0),
+            None,
+        );
+        builder.add_simple_generator(&generator);
+
+        let circuit = builder.build::<C>();
+        let pw = PartialWitness::new();
+        let proof = circuit.data.prove(pw).unwrap();
+        circuit.data.verify(proof).unwrap();
+    }
+}

plonky2x/src/frontend/eth/beacon/generators/mod.rs

@@ -1,2 +1,3 @@
+pub mod balance;
 pub mod validator;
 pub mod validators;

plonky2x/src/utils/eth/beacon/mod.rs

@@ -1,4 +1,5 @@
 use anyhow::Result;
+use ethers::types::U256;
 use num::BigInt;
 use reqwest::Client;
 use serde::Deserialize;
@@ -55,12 +56,43 @@ pub struct GetBeaconValidatorsRoot {
     pub description: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct InnerData {
+    index: String,
+    balance: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct Data {
+    data: Vec<InnerData>,
+    execution_optimistic: bool,
+    finalized: bool,
+}
+
+#[derive(Debug, Deserialize)]
+struct Wrapper {
+    data: Data,
+}
+
 impl BeaconClient {
     /// Creates a new BeaconClient based on a rpc url.
     pub fn new(rpc_url: String) -> Self {
         Self { rpc_url }
     }
 
+    pub async fn get_validator_balance(&self, _: String, validator_idx: u64) -> Result<U256> {
+        let endpoint = format!(
+            "{}/eth/v1/beacon/states/head/validator_balances?id={}",
+            self.rpc_url, validator_idx
+        );
+        println!("{}", endpoint);
+        let client = Client::new();
+        let response = client.get(endpoint).send().await?;
+        let response: Wrapper = response.json().await?;
+        let balance = response.data.data[0].balance.parse::<u64>()?;
+        Ok(U256::from(balance))
+    }
+
     /// Gets the validators root based on a beacon_id and the SSZ proof from
     /// `stateRoot -> validatorsRoot`.
     pub async fn get_validators_root(&self, beacon_id: String) -> Result<GetBeaconValidatorsRoot> {