use BTreeSet and BTreeMap for containers

examples/signed_dict.rs

@@ -2,7 +2,7 @@
 //! Simple example of building a signed dict and verifying it
 //!
 //! Run: `cargo run --release --example signed_dict`
-use std::collections::HashSet;
+use std::collections::BTreeSet;
 
 use pod2::{
     backends::plonky2::{primitives::ec::schnorr::SecretKey, signer::Signer},
@@ -26,7 +26,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     builder.insert("name", "Alice");
     builder.insert("lucky_number", 42);
     builder.insert("human", true);
-    let friends_set: HashSet<Value> = ["Bob", "Charlie", "Dave"]
+    let friends_set: BTreeSet<Value> = ["Bob", "Charlie", "Dave"]
         .into_iter()
         .map(Value::from)
         .collect();

src/backends/plonky2/mainpod/mod.rs

@@ -817,7 +817,7 @@ impl Pod for MainPod {
 
 #[cfg(test)]
 pub mod tests {
-    use std::{any::Any, collections::HashSet};
+    use std::{any::Any, collections::BTreeSet};
 
     use num::{BigUint, One};
 
@@ -1133,7 +1133,7 @@ pub mod tests {
     fn test_set_contains() -> frontend::Result<()> {
         let params = Params::default();
         let mut builder = MainPodBuilder::new(&params, &DEFAULT_VD_SET);
-        let set: HashSet<_> = [1, 2, 3].into_iter().map(|n| n.into()).collect();
+        let set: BTreeSet<_> = [1, 2, 3].into_iter().map(|n| n.into()).collect();
         let set = Set::new(set);
         builder.pub_op(frontend::Operation::set_contains(set, 1))?;
 

src/backends/plonky2/primitives/merkletree/circuit.rs

@@ -671,7 +671,7 @@ impl MerkleTreeStateTransitionProofTarget {
 
 #[cfg(test)]
 pub mod tests {
-    use std::collections::HashMap;
+    use std::collections::BTreeMap;
 
     use plonky2::plonk::{circuit_builder::CircuitBuilder, circuit_data::CircuitConfig};
 
@@ -777,7 +777,7 @@ pub mod tests {
 
     // test logic to be reused both by the existence & nonexistence tests
     fn test_merkleproof_verify_opt(max_depth: usize, existence: bool) -> Result<()> {
-        let mut kvs: HashMap<RawValue, RawValue> = HashMap::new();
+        let mut kvs: BTreeMap<RawValue, RawValue> = BTreeMap::new();
         for i in 0..10 {
             kvs.insert(
                 RawValue::from(hash_value(&RawValue::from(i))),
@@ -834,7 +834,7 @@ pub mod tests {
     }
 
     fn test_merkleproof_only_existence_verify_opt(max_depth: usize) -> Result<()> {
-        let mut kvs: HashMap<RawValue, RawValue> = HashMap::new();
+        let mut kvs: BTreeMap<RawValue, RawValue> = BTreeMap::new();
         for i in 0..10 {
             kvs.insert(
                 RawValue::from(hash_value(&RawValue::from(i))),
@@ -886,7 +886,7 @@ pub mod tests {
         //        /\
         //       5  13
 
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         kvs.insert(RawValue::from(0), RawValue::from(1000));
         kvs.insert(RawValue::from(2), RawValue::from(1002));
         kvs.insert(RawValue::from(5), RawValue::from(1005));
@@ -950,7 +950,7 @@ pub mod tests {
 
     #[test]
     fn test_wrong_witness() -> Result<()> {
-        let mut kvs: HashMap<RawValue, RawValue> = HashMap::new();
+        let mut kvs: BTreeMap<RawValue, RawValue> = BTreeMap::new();
         for i in 0..10 {
             kvs.insert(RawValue::from(i), RawValue::from(i));
         }
@@ -1046,7 +1046,7 @@ pub mod tests {
     #[test]
     fn test_state_transition_gadget() -> Result<()> {
         let max_depth: usize = 32;
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         for i in 0..8 {
             kvs.insert(RawValue::from(i), RawValue::from(1000 + i));
         }
@@ -1105,7 +1105,7 @@ pub mod tests {
         // the last level (max_depth-1)
 
         let max_depth: usize = 32;
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         for i in 0..8 {
             kvs.insert(RawValue::from(i), RawValue::from(1000 + i));
         }
@@ -1169,7 +1169,7 @@ pub mod tests {
     #[test]
     fn test_state_transition_gadget_with_alteration() -> Result<()> {
         let max_depth: usize = 32;
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         for i in 0..8 {
             kvs.insert(RawValue::from(i), RawValue::from(1000 + i));
         }
@@ -1231,7 +1231,7 @@ pub mod tests {
     #[test]
     fn test_state_transition_gadget_disabled() -> Result<()> {
         let max_depth: usize = 32;
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         for i in 0..8 {
             kvs.insert(RawValue::from(i), RawValue::from(1000 + i));
         }

src/backends/plonky2/primitives/merkletree/mod.rs

@@ -1,6 +1,6 @@
 //! Module that implements the MerkleTree specified at
 //! <https://0xparc.github.io/pod2/merkletree.html> .
-use std::{collections::HashMap, fmt, iter::IntoIterator};
+use std::{collections::BTreeMap, fmt, iter::IntoIterator};
 
 use itertools::zip_eq;
 use plonky2::field::types::Field;
@@ -32,7 +32,7 @@ impl Eq for MerkleTree {}
 
 impl MerkleTree {
     /// builds a new `MerkleTree` where the leaves contain the given key-values
-    pub fn new(kvs: &HashMap<RawValue, RawValue>) -> Self {
+    pub fn new(kvs: &BTreeMap<RawValue, RawValue>) -> Self {
         // Start with an empty node as root.
         let mut root = Node::None;
 
@@ -968,7 +968,7 @@ pub mod tests {
 
     #[test]
     fn test_merkletree() -> TreeResult<()> {
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         for i in 0..8 {
             if i == 1 {
                 continue;
@@ -1043,7 +1043,7 @@ pub mod tests {
 
     #[test]
     fn test_state_transition() -> TreeResult<()> {
-        let mut kvs = HashMap::new();
+        let mut kvs = BTreeMap::new();
         for i in 0..8 {
             kvs.insert(RawValue::from(i), RawValue::from(1000 + i));
         }

src/examples/mod.rs

@@ -1,6 +1,6 @@
 pub mod custom;
 
-use std::{collections::HashSet, sync::LazyLock};
+use std::{collections::BTreeSet, sync::LazyLock};
 
 use custom::eth_dos_batch;
 use num::BigUint;
@@ -46,7 +46,7 @@ pub fn zu_kyc_pod_builder(
 ) -> Result<MainPodBuilder> {
     let now_minus_18y = ZU_KYC_NOW_MINUS_18Y;
     let now_minus_1y = ZU_KYC_NOW_MINUS_1Y;
-    let sanctions_values: HashSet<Value> = ZU_KYC_SANCTION_LIST
+    let sanctions_values: BTreeSet<Value> = ZU_KYC_SANCTION_LIST
         .iter()
         .map(|s| Value::from(*s))
         .collect();
@@ -71,7 +71,7 @@ pub fn zu_kyc_pod_builder(
 }
 
 pub fn zu_kyc_pod_request(gov_signer: &Value, pay_signer: &Value) -> Result<PodRequest> {
-    let sanctions_values: HashSet<Value> = ZU_KYC_SANCTION_LIST
+    let sanctions_values: BTreeSet<Value> = ZU_KYC_SANCTION_LIST
         .iter()
         .map(|s| Value::from(*s))
         .collect();
@@ -427,6 +427,6 @@ pub fn tickets_pod_full_flow(params: &Params, vd_set: &VDSet) -> Result<MainPodB
         &signed_dict,
         123,
         true,
-        &Set::new(HashSet::new()),
+        &Set::new(BTreeSet::new()),
     )
 }

src/frontend/custom.rs

@@ -262,7 +262,7 @@ fn resolve_wildcard(args: &[&str], priv_args: &[&str], s: &str) -> Result<Wildca
 
 #[cfg(test)]
 mod tests {
-    use std::collections::HashSet;
+    use std::collections::BTreeSet;
 
     use super::*;
     use crate::{
@@ -361,7 +361,7 @@ mod tests {
 
         let mut mp_builder = MainPodBuilder::new(&params, vd_set);
 
-        let set_values: HashSet<Value> = [1, 2, 3].iter().map(|i| Value::from(*i)).collect();
+        let set_values: BTreeSet<Value> = [1, 2, 3].iter().map(|i| Value::from(*i)).collect();
         let s1 = Set::new(set_values);
         let s2 = 1;
 

src/frontend/mod.rs

@@ -2,7 +2,7 @@
 //! with Pods.
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashSet},
     convert::From,
     fmt,
 };
@@ -36,7 +36,7 @@ pub use pod_request::*;
 #[derive(Clone, Debug)]
 pub struct SignedDictBuilder {
     pub params: Params,
-    pub kvs: HashMap<Key, Value>,
+    pub kvs: BTreeMap<Key, Value>,
 }
 
 impl fmt::Display for SignedDictBuilder {
@@ -53,7 +53,7 @@ impl SignedDictBuilder {
     pub fn new(params: &Params) -> Self {
         Self {
             params: params.clone(),
-            kvs: HashMap::new(),
+            kvs: BTreeMap::new(),
         }
     }
 
@@ -105,7 +105,7 @@ impl SignedDict {
             .then_some(())
             .ok_or(Error::custom("Invalid signature!"))
     }
-    pub fn kvs(&self) -> &HashMap<Key, Value> {
+    pub fn kvs(&self) -> &BTreeMap<Key, Value> {
         self.dict.kvs()
     }
     pub fn get(&self, key: impl Into<Key>) -> Option<&Value> {

src/frontend/multi_pod/mod.rs

@@ -411,6 +411,7 @@ impl MultiPodBuilder {
             .map_err(|e| Error::Frontend(e.to_string()))?;
 
         self.statements.push(stmt.clone());
+        println!("DBG op {:#?}", op);
         self.operations.push(op);
 
         Ok(stmt)
@@ -484,11 +485,20 @@ impl MultiPodBuilder {
         // Build statement content groups for deduplication.
         // Statements with identical content share a single slot in the POD.
         // Group statement indices by their content.
-        let mut content_to_indices: HashMap<&Statement, Vec<usize>> = HashMap::new();
+        let mut statement_content_group_index: HashMap<&Statement, usize> = HashMap::new();
+        let mut statement_content_groups = Vec::new();
         for (idx, stmt) in self.statements.iter().enumerate() {
-            content_to_indices.entry(stmt).or_default().push(idx);
+            use std::collections::hash_map::Entry;
+            match statement_content_group_index.entry(stmt) {
+                Entry::Vacant(e) => {
+                    e.insert(statement_content_groups.len());
+                    statement_content_groups.push(vec![idx]);
+                }
+                Entry::Occupied(e) => {
+                    statement_content_groups[*e.get()].push(idx);
+                }
+            }
         }
-        let statement_content_groups: Vec<Vec<usize>> = content_to_indices.into_values().collect();
 
         // Run solver
         let input = solver::SolverInput {

src/frontend/multi_pod/solver.rs

@@ -73,6 +73,7 @@ pub struct MultiPodSolution {
 }
 
 /// Input to the MILP solver.
+#[derive(Debug)]
 pub struct SolverInput<'a> {
     /// Number of statements.
     pub num_statements: usize,
@@ -170,7 +171,11 @@ pub fn solve(input: &SolverInput) -> Result<MultiPodSolution> {
     // Incremental approach: try solving with increasing POD counts
     // Start with min_pods and increment until we find a feasible solution
     for target_pods in min_pods..=input.max_pods {
+        println!("DBG try solve for {} pods", target_pods);
         if let Some(solution) = try_solve_with_pods(input, target_pods, &all_batches)? {
+            println!("DBG solution found!");
+            dbg!(&input);
+            dbg!(&solution);
             return Ok(solution);
         }
         // Infeasible with target_pods, try more

src/frontend/serialization.rs

@@ -62,7 +62,7 @@ impl TryFrom<SerializedMainPod> for MainPod {
 
 #[cfg(test)]
 mod tests {
-    use std::collections::{HashMap, HashSet};
+    use std::collections::{BTreeMap, BTreeSet};
 
     use pretty_assertions::assert_eq;
     use schemars::schema_for;
@@ -100,7 +100,7 @@ mod tests {
             ),
             (
                 TypedValue::Dictionary(
-                    Dictionary::new(HashMap::from([
+                    Dictionary::new(BTreeMap::from([
                         // The set of valid keys is equal to the set of valid JSON keys
                         ("foo".into(), 123.into()),
                         // Empty strings are valid JSON keys
@@ -118,7 +118,7 @@ mod tests {
                 "{\"kvs\":{\"\":\"baz\",\"\\u0000\":\"\",\"    hi\":false,\"!@£$%^&&*()\":\"\",\"foo\":{\"Int\":\"123\"},\"🥳\":\"party time!\"}}",
             ),
             (
-                TypedValue::Set(Set::new(HashSet::from(["foo".into(), "bar".into()]))),
+                TypedValue::Set(Set::new(BTreeSet::from(["foo".into(), "bar".into()]))),
                 "{\"set\":[\"bar\",\"foo\"]}",
             ),
         ];
@@ -145,15 +145,15 @@ mod tests {
         builder.insert("very_large_int", 1152921504606846976);
         builder.insert(
             "a_dict_containing_one_key",
-            Dictionary::new(HashMap::from([
+            Dictionary::new(BTreeMap::from([
                 ("foo".into(), 123.into()),
                 (
                     "an_array_containing_three_ints".into(),
                     Array::new(vec![1.into(), 2.into(), 3.into()]).into(),
                 ),
                 (
                     "a_set_containing_two_strings".into(),
-                    Set::new(HashSet::from([
+                    Set::new(BTreeSet::from([
                         Array::new(vec!["foo".into(), "bar".into()]).into(),
                         "baz".into(),
                     ]))

src/lang/frontend_ast_lower.rs

@@ -4,7 +4,7 @@
 //! Supports automatic predicate splitting and multi-batch packing.
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     str::FromStr,
     sync::Arc,
 };
@@ -174,13 +174,13 @@ pub fn lower_literal(lit: &LiteralValue) -> Value {
             Value::from(array)
         }
         LiteralValue::Set(s) => {
-            let elements: std::collections::HashSet<_> =
+            let elements: std::collections::BTreeSet<_> =
                 s.elements.iter().map(lower_literal).collect();
             let set = containers::Set::new(elements);
             Value::from(set)
         }
         LiteralValue::Dict(d) => {
-            let pairs: HashMap<_, _> = d
+            let pairs: BTreeMap<_, _> = d
                 .pairs
                 .iter()
                 .map(|pair| {