From 0c01a5d9a870ac5a111d290df862435014d25b4a Mon Sep 17 00:00:00 2001
From: zeroqn <zeroqn@pm.me>
Date: Thu, 21 Mar 2024 06:14:16 +0000
Subject: [PATCH] fix(polyexen): invalid resolved polys and lookups toml

---
 tools/polyexen/Cargo.lock  |  12 ++--
 tools/polyexen/Cargo.toml  |   2 +
 tools/polyexen/src/main.rs | 130 +++++++++++++++++++++++++------------
 3 files changed, 96 insertions(+), 48 deletions(-)

diff --git a/tools/polyexen/Cargo.lock b/tools/polyexen/Cargo.lock
index dd02d5e..a7e9582 100644
--- a/tools/polyexen/Cargo.lock
+++ b/tools/polyexen/Cargo.lock
@@ -4066,7 +4066,9 @@ dependencies = [
  "log",
  "num-bigint 0.4.4",
  "polyexen",
+ "serde",
  "summa-solvency",
+ "toml",
 ]
 
 [[package]]
@@ -4289,14 +4291,14 @@ dependencies = [
 
 [[package]]
 name = "toml"
-version = "0.8.10"
+version = "0.8.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9a9aad4a3066010876e8dcf5a8a06e70a558751117a145c6ce2b82c2e2054290"
+checksum = "e9dd1545e8208b4a5af1aa9bbd0b4cf7e9ea08fabc5d0a5c67fcaafa17433aa3"
 dependencies = [
  "serde",
  "serde_spanned",
  "toml_datetime",
- "toml_edit 0.22.6",
+ "toml_edit 0.22.9",
 ]
 
 [[package]]
@@ -4343,9 +4345,9 @@ dependencies = [
 
 [[package]]
 name = "toml_edit"
-version = "0.22.6"
+version = "0.22.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2c1b5fd4128cc8d3e0cb74d4ed9a9cc7c7284becd4df68f5f940e1ad123606f6"
+checksum = "8e40bb779c5187258fd7aad0eb68cb8706a0a81fa712fbea808ab43c4b8374c4"
 dependencies = [
  "indexmap",
  "serde",
diff --git a/tools/polyexen/Cargo.toml b/tools/polyexen/Cargo.toml
index 2198dd1..4715aef 100644
--- a/tools/polyexen/Cargo.toml
+++ b/tools/polyexen/Cargo.toml
@@ -12,3 +12,5 @@ polyexen = { path = "./deps/polyexen" }
 summa-solvency = { path = "./circuit/solvency/zk_prover" }
 log = "0.4.20"
 env_logger = "0.11.2"
+toml = "0.8.12"
+serde = { version = "1.0.197", features = ["derive"] }
diff --git a/tools/polyexen/src/main.rs b/tools/polyexen/src/main.rs
index f68de65..ff71215 100644
--- a/tools/polyexen/src/main.rs
+++ b/tools/polyexen/src/main.rs
@@ -20,6 +20,7 @@ use polyexen::{
         PlafDisplayFixedCSV,
     },
 };
+use serde::Serialize;
 
 mod circuit;
 
@@ -308,8 +309,8 @@ struct AnalysisResult {
 impl AnalysisResult {
     fn write_files(&self, name: &str) -> Result<(), io::Error> {
         if !Path::new("out").exists() {
+            log::debug!("create out dir");
             std::fs::create_dir("out")?;
-            log::debug!("create out dir")
         }
 
         let mut base_file = File::create(format!("out/{}.toml", name))?;
@@ -363,34 +364,56 @@ impl fmt::Display for DisplayPolysBaseTOML<'_> {
             )
         };
 
+        #[derive(Default, Serialize)]
+        struct ResolvedPolys {
+            var_bounds: HashMap<String, String>,
+            polys: Vec<String>,
+        }
+
+        #[derive(Default, Serialize)]
+        struct Constraints {
+            resolved_polys: ResolvedPolys,
+        }
+
+        #[derive(Default, Serialize)]
+        struct Toml {
+            constraints: Constraints,
+        }
+
+        let mut constraints = Constraints::default();
+
         for (c, bound) in self.var_bounds {
-            writeln!(
-                f,
-                "[constraints.resolved_polys.vars.\"{}\"]",
-                CellDisplay {
-                    c,
-                    plaf: &self.plaf,
-                }
-            )?;
-            writeln!(f, "bound = {}", bound)?;
+            let cell_disp = CellDisplay {
+                c,
+                plaf: &self.plaf,
+            };
+            constraints
+                .resolved_polys
+                .var_bounds
+                .insert(cell_disp.to_string(), bound.to_string());
         }
-        writeln!(f)?;
 
         for p in self.polys {
-            writeln!(f, "[constraints.resolved_polys.\"{}\"]", p.name)?;
-            write!(f, "c = \"")?;
-            write!(
-                f,
-                "{}",
-                ExprDisplay {
-                    e: &p.expr,
-                    var_fmt: cell_fmt,
-                }
-            )?;
-            writeln!(f, "\"")?;
+            let expr_disp = ExprDisplay {
+                e: &p.expr,
+                var_fmt: cell_fmt,
+            };
+            constraints
+                .resolved_polys
+                .polys
+                .push(format!("{}{}", p.name, expr_disp));
         }
 
-        Ok(())
+        match toml::to_string_pretty(&Toml { constraints }) {
+            Ok(toml_str) => {
+                write!(f, "{}", toml_str)?;
+                Ok(())
+            }
+            Err(err) => {
+                log::error!("serialize constraints polys to toml failed {}", err);
+                Err(fmt::Error)
+            }
+        }
     }
 }
 
@@ -426,29 +449,50 @@ impl fmt::Display for DisplayLookupsBaseTOML<'_> {
             )
         };
 
+        #[derive(Default, Serialize)]
+        struct ResolvedLookups {
+            lookups: HashMap<String, Vec<String>>,
+        }
+
+        #[derive(Default, Serialize)]
+        struct Constraints {
+            resolved_lookups: ResolvedLookups,
+        }
+
+        #[derive(Default, Serialize)]
+        struct Toml {
+            constraints: Constraints,
+        }
+
+        let mut constraints = Constraints::default();
+
         for l in self.lookups {
-            writeln!(
-                f,
-                "[constraints.resolved_lookups.\"{}\"_\"{}\"]",
-                l.exprs_num.1, l.name
-            )?;
-            write!(f, "l = [")?;
-            for (i, exp) in l.exprs_num.0.iter().enumerate() {
-                if i != 0 {
-                    write!(f, ", ")?;
-                }
-                write!(
-                    f,
-                    "{}",
-                    ExprDisplay {
-                        e: &exp,
-                        var_fmt: cell_fmt
-                    },
-                )?;
+            let exprs_num_name = format!("{}_{}", l.exprs_num.1, l.name);
+            let mut exprs = vec![];
+
+            for exp in l.exprs_num.0.iter() {
+                let expr_disp = ExprDisplay {
+                    e: &exp,
+                    var_fmt: cell_fmt,
+                };
+                exprs.push(expr_disp.to_string());
             }
-            writeln!(f, "]")?;
+
+            constraints
+                .resolved_lookups
+                .lookups
+                .insert(exprs_num_name, exprs);
         }
 
-        Ok(())
+        match toml::to_string_pretty(&Toml { constraints }) {
+            Ok(toml_str) => {
+                write!(f, "{}", toml_str)?;
+                Ok(())
+            }
+            Err(err) => {
+                log::error!("serialize constraints lookups to toml failed {}", err);
+                Err(fmt::Error)
+            }
+        }
     }
 }