fix(polyexen): invalid resolved polys and lookups toml

zBlock-2 · Mar 21, 2024 · 0c01a5d · 0c01a5d
1 parent 06a6b8a
commit 0c01a5d
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 48 deletions.
diff --git a/tools/polyexen/Cargo.lock b/tools/polyexen/Cargo.lock
diff --git 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
@@ -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)
+            }
+        }
     }
 }