Merge pull request #263 from hatoo/print-size-all

Support `(print-size)` to print all size of functions
egraphs-good · Oct 20, 2023 · c7405d8 · c7405d8
2 parents 959afb9 + b42eab1
commit c7405d8
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 10 deletions.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -104,7 +104,7 @@ pub enum NCommand {
     Check(Vec<NormFact>),
     CheckProof,
     PrintTable(Symbol, usize),
-    PrintSize(Symbol),
+    PrintSize(Option<Symbol>),
     Output {
         file: String,
         exprs: Vec<Expr>,
@@ -642,8 +642,8 @@ pub enum Command {
     /// prints the first 20 rows of the `Add` function.
     ///
     PrintFunction(Symbol, usize),
-    /// Print out the number of rows in a function.
-    PrintSize(Symbol),
+    /// Print out the number of rows in a function or all functions.
+    PrintSize(Option<Symbol>),
     /// Input a CSV file directly into a function.
     Input {
         name: Symbol,
@@ -699,7 +699,7 @@ impl ToSexp for Command {
             Command::Push(n) => list!("push", n),
             Command::Pop(n) => list!("pop", n),
             Command::PrintFunction(name, n) => list!("print-function", name, n),
-            Command::PrintSize(name) => list!("print-size", name),
+            Command::PrintSize(name) => list!("print-size", ++ name),
             Command::Input { name, file } => list!("input", name, format!("\"{}\"", file)),
             Command::Output { file, exprs } => list!("output", format!("\"{}\"", file), ++ exprs),
             Command::Fail(cmd) => list!("fail", cmd),

diff --git a/src/ast/parse.lalrpop b/src/ast/parse.lalrpop
@@ -82,7 +82,7 @@ Command: Command = {
     LParen "push" <UNum?> RParen => Command::Push(<>.unwrap_or(1)),
     LParen "pop" <UNum?> RParen => Command::Pop(<>.unwrap_or(1)),
     LParen "print-function" <sym:Ident> <n:UNum> RParen => Command::PrintFunction(sym, n),
-    LParen "print-size" <sym:Ident> RParen => Command::PrintSize(sym),
+    LParen "print-size" <sym:Ident?> RParen => Command::PrintSize(sym),
     LParen "input" <name:Ident> <file:String> RParen => Command::Input { name, file },
     LParen "output" <file:String> <exprs:Expr+> RParen => Command::Output { file, exprs },
     LParen "fail" <Command> RParen => Command::Fail(Box::new(<>)),

diff --git a/src/lib.rs b/src/lib.rs
@@ -684,11 +684,36 @@ impl EGraph {
         Ok(())
     }
 
-    pub fn print_size(&mut self, sym: Symbol) -> Result<(), Error> {
-        let f = self.functions.get(&sym).ok_or(TypeError::Unbound(sym))?;
-        log::info!("Function {} has size {}", sym, f.nodes.len());
-        self.print_msg(f.nodes.len().to_string());
-        Ok(())
+    pub fn print_size(&mut self, sym: Option<Symbol>) -> Result<(), Error> {
+        if let Some(sym) = sym {
+            let f = self.functions.get(&sym).ok_or(TypeError::Unbound(sym))?;
+            log::info!("Function {} has size {}", sym, f.nodes.len());
+            self.print_msg(f.nodes.len().to_string());
+            Ok(())
+        } else {
+            // Print size of all functions
+            let mut lens = self
+                .functions
+                .iter()
+                .map(|(sym, f)| (*sym, f.nodes.len()))
+                .collect::<Vec<_>>();
+
+            // Function name's alphabetical order
+            lens.sort_by_key(|(name, _)| name.as_str());
+
+            for (sym, len) in &lens {
+                log::info!("Function {} has size {}", sym, len);
+            }
+
+            self.print_msg(
+                lens.into_iter()
+                    .map(|(name, len)| format!("{}: {}", name, len))
+                    .collect::<Vec<_>>()
+                    .join("\n"),
+            );
+
+            Ok(())
+        }
     }
 
     // returns whether the egraph was updated

diff --git a/tests/math-microbenchmark.egg b/tests/math-microbenchmark.egg
@@ -68,4 +68,6 @@
 (print-size Add)
 (print-size Mul)
 
+(print-size)
+
 (print-stats)