test(blend): test interaction between core PoQ and Merkle logic (#1878)

nomos-core/chain-defs/src/mantle/keys.rs

@@ -39,6 +39,11 @@ impl SecretKey {
         &self.0
     }
 
+    #[must_use]
+    pub fn into_inner(self) -> Fr {
+        self.0
+    }
+
     #[must_use]
     pub fn to_public_key(&self) -> PublicKey {
         PublicKey(Poseidon2Bn254Hasher::digest(&[*NOMOS_KDF, self.0]))

nomos-services/blend/src/merkle.rs

@@ -197,10 +197,11 @@ fn compute_selectors(
     let mut result = [false; CORE_MERKLE_TREE_HEIGHT];
     let mut idx = *index;
 
-    // The selector at each level is determined by the bit at that position
-    // in the binary representation of the index
-    // Bit 0 (LSB) determines position at level 0, bit 1 at level 1, etc.
-    for result_entry in result.iter_mut().take(CORE_MERKLE_TREE_HEIGHT) {
+    // The selector at each level is determined by the corresponding bit of the leaf
+    // index. Iterating from the last element to the first (leaf → root):
+    // result[CORE_MERKLE_TREE_HEIGHT-1] (leaf level) = bit 0 (LSB) of index
+    // result[0] (root level) = MSB of index
+    for result_entry in result.iter_mut().take(CORE_MERKLE_TREE_HEIGHT).rev() {
         *result_entry = (idx & 1) == 1;
         idx >>= 1u8;
     }
@@ -213,7 +214,21 @@ mod tests {
     use core::iter::repeat_n;
 
     use groth16::{Field as _, fr_from_bytes_unchecked};
-    use nomos_core::{crypto::ZkHash, mantle::keys::PublicKey};
+    use nomos_blend_message::crypto::{
+        keys::Ed25519PublicKey,
+        proofs::quota::{
+            ProofOfQuota,
+            inputs::prove::{
+                PrivateInputs, PublicInputs,
+                private::ProofOfCoreQuotaInputs,
+                public::{CoreInputs, LeaderInputs},
+            },
+        },
+    };
+    use nomos_core::{
+        crypto::ZkHash,
+        mantle::keys::{PublicKey, SecretKey},
+    };
     use num_bigint::BigUint;
 
     use crate::merkle::{Error, MerkleTree, TOTAL_MERKLE_LEAVES};
@@ -255,11 +270,12 @@ mod tests {
             .verify_proof_for_key(&proof_for_key_one, &key_one)
             .unwrap();
         // We check that the first key is the right child of the bottom sub-tree...
-        assert!(proof_for_key_one.first().unwrap().1);
+        assert!(proof_for_key_one.last().unwrap().1);
         // ...but the left of all sub-trees above that.
         assert!(
             !proof_for_key_one
                 .iter()
+                .rev()
                 .skip(1)
                 .any(|(_, selector)| *selector)
         );
@@ -289,10 +305,11 @@ mod tests {
         merkle_tree
             .verify_proof_for_key(&proof_for_key_one, &key_one)
             .unwrap();
-        assert!(proof_for_key_one.first().unwrap().1);
+        assert!(proof_for_key_one.last().unwrap().1);
         assert!(
             !proof_for_key_one
                 .iter()
+                .rev()
                 .skip(1)
                 .any(|(_, selector)| *selector)
         );
@@ -307,15 +324,16 @@ mod tests {
         merkle_tree
             .verify_proof_for_key(&proof_for_key_three, &key_three)
             .unwrap();
-        // First selector is `true` because it's the left child...
-        assert!(proof_for_key_one[0].1);
-        // Second selector is `false` because it's already in the right sub-tree at this
-        // level (first sub-tree are keys 1 and 2).
-        assert!(!proof_for_key_one[1].1);
+        // Last selector is `false` because it's the left child...
+        assert!(!proof_for_key_three.last().unwrap().1);
+        // Second-to-last selector is `true` because it's already in the right sub-tree
+        // at this level (first sub-tree are keys 1 and 2).
+        assert!(proof_for_key_three[proof_for_key_three.len() - 2].1);
         // It's in the left-most sub-tree going above.
         assert!(
-            !proof_for_key_one
+            !proof_for_key_three
                 .iter()
+                .rev()
                 .skip(2)
                 .any(|(_, selector)| *selector)
         );
@@ -354,4 +372,139 @@ mod tests {
         let key = PublicKey::new(ZkHash::ONE);
         assert_eq!(MerkleTree::new(vec![key, key]), Err(Error::DuplicateKey));
     }
+
+    struct PoQInputs<const INPUTS: usize> {
+        public_inputs: PublicInputs,
+        secret_inputs: [ProofOfCoreQuotaInputs; INPUTS],
+    }
+
+    fn generate_inputs<const INPUTS: usize>() -> PoQInputs<INPUTS> {
+        let keys: [_; INPUTS] = (1..=INPUTS as u64)
+            .map(|i| {
+                let sk = SecretKey::new(ZkHash::from(i));
+                let pk = sk.to_public_key();
+                (sk, pk)
+            })
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap();
+        let merkle_tree = MerkleTree::new(keys.clone().map(|(_, pk)| pk).to_vec()).unwrap();
+        let public_inputs = {
+            let core_inputs = CoreInputs {
+                quota: 1,
+                zk_root: merkle_tree.root(),
+            };
+            let leader_inputs = LeaderInputs {
+                message_quota: 1,
+                pol_epoch_nonce: ZkHash::ZERO,
+                pol_ledger_aged: ZkHash::ZERO,
+                total_stake: 1,
+            };
+            let session = 1;
+            let signing_key: Ed25519PublicKey = [10; 32].try_into().unwrap();
+            PublicInputs {
+                core: core_inputs,
+                leader: leader_inputs,
+                session,
+                signing_key,
+            }
+        };
+        let secret_inputs = keys.map(|(sk, pk)| {
+            let proof = merkle_tree.get_proof_for_key(&pk).unwrap();
+            ProofOfCoreQuotaInputs {
+                core_sk: sk.into_inner(),
+                core_path_and_selectors: proof,
+            }
+        });
+
+        PoQInputs {
+            public_inputs,
+            secret_inputs,
+        }
+    }
+
+    #[test]
+    fn poq_interaction_single_key() {
+        let PoQInputs {
+            public_inputs,
+            secret_inputs,
+        } = generate_inputs::<1>();
+
+        for secret_input in secret_inputs {
+            let (poq, _) = ProofOfQuota::new(
+                &public_inputs,
+                PrivateInputs::new_proof_of_core_quota_inputs(0, secret_input),
+            )
+            .unwrap();
+            poq.verify(&public_inputs).unwrap();
+        }
+    }
+
+    #[test]
+    fn poq_interaction_two_keys() {
+        let PoQInputs {
+            public_inputs,
+            secret_inputs,
+        } = generate_inputs::<2>();
+
+        for secret_input in secret_inputs {
+            let (poq, _) = ProofOfQuota::new(
+                &public_inputs,
+                PrivateInputs::new_proof_of_core_quota_inputs(0, secret_input),
+            )
+            .unwrap();
+            poq.verify(&public_inputs).unwrap();
+        }
+    }
+
+    #[test]
+    fn poq_interaction_three_keys() {
+        let PoQInputs {
+            public_inputs,
+            secret_inputs,
+        } = generate_inputs::<3>();
+
+        for secret_input in secret_inputs {
+            let (poq, _) = ProofOfQuota::new(
+                &public_inputs,
+                PrivateInputs::new_proof_of_core_quota_inputs(0, secret_input),
+            )
+            .unwrap();
+            poq.verify(&public_inputs).unwrap();
+        }
+    }
+
+    #[test]
+    fn poq_interaction_four_keys() {
+        let PoQInputs {
+            public_inputs,
+            secret_inputs,
+        } = generate_inputs::<3>();
+
+        for secret_input in secret_inputs {
+            let (poq, _) = ProofOfQuota::new(
+                &public_inputs,
+                PrivateInputs::new_proof_of_core_quota_inputs(0, secret_input),
+            )
+            .unwrap();
+            poq.verify(&public_inputs).unwrap();
+        }
+    }
+
+    #[test]
+    fn poq_interaction_one_hundred_keys() {
+        let PoQInputs {
+            public_inputs,
+            secret_inputs,
+        } = generate_inputs::<100>();
+
+        for secret_input in secret_inputs {
+            let (poq, _) = ProofOfQuota::new(
+                &public_inputs,
+                PrivateInputs::new_proof_of_core_quota_inputs(0, secret_input),
+            )
+            .unwrap();
+            poq.verify(&public_inputs).unwrap();
+        }
+    }
 }