From c8b0d6747e705458e515e5d6f3a1d08b6a544dbb Mon Sep 17 00:00:00 2001
From: thomaslavaur <thomas19.lavaur@gmail.com>
Date: Wed, 13 May 2026 09:45:18 +0200
Subject: [PATCH] update pol circuits for arbitrary merkle depth

---
 .../mantle/generate_inputs_for_pol.py         | 32 ++++++++++---------
 circom_circuits/mantle/pol.circom             |  2 +-
 circom_circuits/mantle/pol_lib.circom         | 32 +++++++++----------
 3 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/circom_circuits/mantle/generate_inputs_for_pol.py b/circom_circuits/mantle/generate_inputs_for_pol.py
index 34a8f88..be5ffb3 100755
--- a/circom_circuits/mantle/generate_inputs_for_pol.py
+++ b/circom_circuits/mantle/generate_inputs_for_pol.py
@@ -204,13 +204,14 @@ def PoseidonSponge(data, capacity, output_len):
     return output
 
 
-if len(sys.argv) != Integer(4):
-    print("Usage: <script> <epoch_nonce> <slot_number> <total_stake>")
+if len(sys.argv) != Integer(5):
+    print("Usage: <script> <merkle_depth> <epoch_nonce> <slot_number> <total_stake>")
     exit()
 
-epoch_nonce = int(sys.argv[Integer(1)])
-slot_number = int(sys.argv[Integer(2)])
-total_stake = int(sys.argv[Integer(3)])
+merkle_depth = int(sys.argv[Integer(1)])
+epoch_nonce = int(sys.argv[Integer(2)])
+slot_number = int(sys.argv[Integer(3)])
+total_stake = int(sys.argv[Integer(4)])
 
 if epoch_nonce >= p:
     print("epoch nonce must be less than p")
@@ -241,23 +242,24 @@ while(ticket > threshold):
     note_id = poseidon2_hash([F(232989242343357190262606),tx_hash,output_number,value,pk])
     ticket = poseidon2_hash([F(13887241025832268),F(epoch_nonce),F(slot_number),note_id,sk])
  
-aged_nodes = [F(randrange(0,p,1)) for i in range(32)]
-aged_selectors = randrange(0,2**32,1)
-aged_selectors = format(aged_selectors,'032b')
+aged_nodes = [F(randrange(0,p,1)) for i in range(merkle_depth)]
+aged_selectors = randrange(0,2**merkle_depth,1)
+string = '0'+str(merkle_depth)+'b'
+aged_selectors = format(aged_selectors,string)
 aged_root = note_id
-for i in range(32):
-    if int(aged_selectors[31-i]) == 0:
+for i in range(merkle_depth):
+    if int(aged_selectors[merkle_depth-1-i]) == 0:
         aged_root = Compression([aged_root,aged_nodes[i]])
     else:
         aged_root = Compression([aged_nodes[i],aged_root])
 
-unspent_nodes = [F(randrange(0,p,1)) for i in range(32)]
-unspent_selectors = randrange(0,2**32,1)
-unspent_selectors = format(unspent_selectors,'032b')
+unspent_nodes = [F(randrange(0,p,1)) for i in range(merkle_depth)]
+unspent_selectors = randrange(0,2**merkle_depth,1)
+unspent_selectors = format(unspent_selectors,string)
 
 latest_root = note_id
-for i in range(32):
-    if int(unspent_selectors[31-i]) == 0:
+for i in range(merkle_depth):
+    if int(unspent_selectors[merkle_depth-1-i]) == 0:
         latest_root = Compression([latest_root,unspent_nodes[i]])
     else:
         latest_root = Compression([unspent_nodes[i],latest_root])
diff --git a/circom_circuits/mantle/pol.circom b/circom_circuits/mantle/pol.circom
index a362a6a..14387fd 100644
--- a/circom_circuits/mantle/pol.circom
+++ b/circom_circuits/mantle/pol.circom
@@ -3,4 +3,4 @@ pragma circom 2.1.9;
 
 include "pol_lib.circom";
 
-component main {public [sl,epoch_nonce,t0,t1,ledger_aged,ledger_latest,P_lead_part_one,P_lead_part_two]}= proof_of_leadership(25);
\ No newline at end of file
+component main {public [sl,epoch_nonce,t0,t1,ledger_aged,ledger_latest,P_lead_part_one,P_lead_part_two]}= proof_of_leadership(64);
\ No newline at end of file
diff --git a/circom_circuits/mantle/pol_lib.circom b/circom_circuits/mantle/pol_lib.circom
index cfe8fc8..38dd6c6 100644
--- a/circom_circuits/mantle/pol_lib.circom
+++ b/circom_circuits/mantle/pol_lib.circom
@@ -43,15 +43,15 @@ template derive_entropy(){
     out <== hash.out;
 }
 
-template would_win_leadership(secret_depth){
+template would_win_leadership(merkle_depth){
     signal input slot;
     signal input epoch_nonce;
     signal input t0;
     signal input t1;
 
     //Part of the note id proof of membership to prove aged
-    signal input aged_nodes[32];
-    signal input aged_selectors[32];         // must be bits
+    signal input aged_nodes[merkle_depth];
+    signal input aged_selectors[merkle_depth];         // must be bits
     signal input aged_root;
 
     //Used to derive the note identifier
@@ -83,12 +83,12 @@ template would_win_leadership(secret_depth){
 
     // Check the note ID is aged enough
             //First check selectors are indeed bits
-    for(var i = 0; i < 32; i++){
+    for(var i = 0; i < merkle_depth; i++){
         aged_selectors[i] * (1 - aged_selectors[i]) === 0;
     }
             //Then check the proof of membership
-    component aged_membership = proof_of_membership(32);
-    for(var i = 0; i < 32; i++){
+    component aged_membership = proof_of_membership(merkle_depth);
+    for(var i = 0; i < merkle_depth; i++){
         aged_membership.nodes[i] <== aged_nodes[i];
         aged_membership.selector[i] <== aged_selectors[i];
     }
@@ -123,15 +123,15 @@ template would_win_leadership(secret_depth){
 } 
 
 
-template proof_of_leadership(secret_depth){
+template proof_of_leadership(merkle_depth){
     signal input sl;
     signal input epoch_nonce;  // the epoch nonce eta
     signal input t0;
     signal input t1;
 
     //Part of the note id proof of membership to prove aged
-    signal input noteid_aged_path[32];
-    signal input noteid_aged_selectors[32];         // must be bits
+    signal input noteid_aged_path[merkle_depth];
+    signal input noteid_aged_selectors[merkle_depth];         // must be bits
     signal input ledger_aged;
 
     //Used to derive the note identifier
@@ -140,8 +140,8 @@ template proof_of_leadership(secret_depth){
     signal input note_output_number;
     
     //Part of the note id proof of membership to prove it's unspent
-    signal input noteid_latest_path[32];
-    signal input noteid_latest_selectors[32];         // must be bits
+    signal input noteid_latest_path[merkle_depth];
+    signal input noteid_latest_selectors[merkle_depth];         // must be bits
     signal input ledger_latest;
 
     // The winning note.
@@ -149,12 +149,12 @@ template proof_of_leadership(secret_depth){
 
 
     // Verify the note is winning the lottery
-    component lottery_checker = would_win_leadership(secret_depth);
+    component lottery_checker = would_win_leadership(merkle_depth);
     lottery_checker.slot <== sl;
     lottery_checker.epoch_nonce <== epoch_nonce;
     lottery_checker.t0 <== t0;
     lottery_checker.t1 <== t1;
-    for(var i = 0; i < 32; i++){
+    for(var i = 0; i < merkle_depth; i++){
         lottery_checker.aged_nodes[i] <== noteid_aged_path[i];
         lottery_checker.aged_selectors[i] <== noteid_aged_selectors[i];
     }
@@ -181,12 +181,12 @@ template proof_of_leadership(secret_depth){
 
     // Check that the note is unspent
             //First check selectors are indeed bits
-    for(var i = 0; i < 32; i++){
+    for(var i = 0; i < merkle_depth; i++){
         noteid_latest_selectors[i] * (1 - noteid_latest_selectors[i]) === 0;
     }
             //Then check the note id is in the latest ledger state
-    component unspent_membership = proof_of_membership(32);
-    for(var i = 0; i < 32; i++){
+    component unspent_membership = proof_of_membership(merkle_depth);
+    for(var i = 0; i < merkle_depth; i++){
         unspent_membership.nodes[i] <== noteid_latest_path[i];
         unspent_membership.selector[i] <== noteid_latest_selectors[i];
     }