From 4f53745a84b4f93b008e337eb28d6a7851574b8a Mon Sep 17 00:00:00 2001
From: Martin Ivanov <xeartybg@gmail.com>
Date: Mon, 24 Jun 2024 11:33:03 +0300
Subject: [PATCH] refactor(balance_verification_final_layer): Don't register
 public inputs manually

---
 .../final_layer.rs                               | 10 ++++++++--
 .../final_layer.rs                               | 16 +++++++++-------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/beacon-light-client/plonky2/crates/circuits/src/deposit_accumulator_balance_aggregator_diva/final_layer.rs b/beacon-light-client/plonky2/crates/circuits/src/deposit_accumulator_balance_aggregator_diva/final_layer.rs
index 61978c5a2..6e7c8feda 100644
--- a/beacon-light-client/plonky2/crates/circuits/src/deposit_accumulator_balance_aggregator_diva/final_layer.rs
+++ b/beacon-light-client/plonky2/crates/circuits/src/deposit_accumulator_balance_aggregator_diva/final_layer.rs
@@ -24,6 +24,7 @@ use circuit_derive::{CircuitTarget, SerdeCircuitTarget};
 use plonky2::{
     field::{extension::Extendable, goldilocks_field::GoldilocksField},
     hash::hash_types::{HashOut, RichField},
+    iop::target::Target,
     plonk::{
         circuit_builder::CircuitBuilder,
         circuit_data::{CircuitConfig, CircuitData},
@@ -80,6 +81,9 @@ pub struct DepositAccumulatorBalanceAggregatorDivaFinalLayerTarget {
     #[target(in)]
     #[serde(with = "serde_bool_array_to_hex_string_nested")]
     pub slot_branch: Sha256MerkleBranchTarget<5>,
+
+    #[target(out)]
+    pub public_inputs_hash_bytes: [Target; 32],
 }
 
 pub struct DepositAccumulatorBalanceAggregatorDivaFinalLayer;
@@ -179,8 +183,9 @@ impl Circuit for DepositAccumulatorBalanceAggregatorDivaFinalLayer {
         public_inputs_hash[1] = builder._false();
         public_inputs_hash[2] = builder._false();
 
-        let public_inputs_hash_bytes = bits_to_bytes_target(builder, &public_inputs_hash);
-        builder.register_public_inputs(&public_inputs_hash_bytes);
+        let public_inputs_hash_bytes = bits_to_bytes_target(builder, &public_inputs_hash)
+            .try_into()
+            .unwrap();
 
         Self::Target {
             balance_aggregation_proof,
@@ -196,6 +201,7 @@ impl Circuit for DepositAccumulatorBalanceAggregatorDivaFinalLayer {
             execution_block_number_branch: input.execution_block_number_branch,
             slot: input.slot,
             slot_branch: input.slot_branch,
+            public_inputs_hash_bytes,
         }
     }
 }
diff --git a/beacon-light-client/plonky2/crates/circuits/src/withdrawal_credentials_balance_aggregator/final_layer.rs b/beacon-light-client/plonky2/crates/circuits/src/withdrawal_credentials_balance_aggregator/final_layer.rs
index c9cb430b5..c67093455 100644
--- a/beacon-light-client/plonky2/crates/circuits/src/withdrawal_credentials_balance_aggregator/final_layer.rs
+++ b/beacon-light-client/plonky2/crates/circuits/src/withdrawal_credentials_balance_aggregator/final_layer.rs
@@ -6,7 +6,7 @@ use crate::{
     },
     utils::circuit::{
         assert_slot_is_in_epoch::assert_slot_is_in_epoch,
-        biguint_to_bits_target,
+        biguint_to_bits_target, bits_to_bytes_target,
         hashing::{
             merkle::{sha256::assert_merkle_proof_is_valid_const_sha256, ssz::ssz_num_to_bits},
             sha256::sha256,
@@ -22,6 +22,7 @@ use itertools::Itertools;
 use plonky2::{
     field::{extension::Extendable, goldilocks_field::GoldilocksField},
     hash::hash_types::RichField,
+    iop::target::Target,
     plonk::{
         circuit_builder::CircuitBuilder,
         circuit_data::{CircuitConfig, CircuitData},
@@ -64,6 +65,9 @@ pub struct FinalCircuitTargets<const WITHDRAWAL_CREDENTIALS_COUNT: usize> {
     #[serde(with = "serde_bool_array_to_hex_string_nested")]
     pub balances_branch: Sha256MerkleBranchTarget<6>,
 
+    #[target(out)]
+    pub public_inputs_hash_bytes: [Target; 32],
+
     pub balance_verification_proof: ProofWithPublicInputsTarget<D>,
     pub validators_commitment_mapper_proof: ProofWithPublicInputsTarget<D>,
 }
@@ -164,12 +168,9 @@ impl<const WITHDRAWAL_CREDENTIALS_COUNT: usize> Circuit
         public_inputs_hash[1] = builder._false();
         public_inputs_hash[2] = builder._false();
 
-        let public_inputs_hash_bytes = public_inputs_hash
-            .chunks(8)
-            .map(|x| builder.le_sum(x.iter().rev()))
-            .collect_vec();
-
-        builder.register_public_inputs(&public_inputs_hash_bytes);
+        let public_inputs_hash_bytes = bits_to_bytes_target(builder, &public_inputs_hash)
+            .try_into()
+            .unwrap();
 
         Self::Target {
             block_root: input.block_root,
@@ -181,6 +182,7 @@ impl<const WITHDRAWAL_CREDENTIALS_COUNT: usize> Circuit
             validators_branch: input.validators_branch,
             balance_verification_proof,
             validators_commitment_mapper_proof,
+            public_inputs_hash_bytes,
         }
     }
 }