More work for demo

uwsampl · Apr 27, 2024 · b3ef865 · b3ef865
1 parent dd5cacf
commit b3ef865
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/churchroad-js/package-lock.json b/churchroad-js/package-lock.json
diff --git a/web-demo/Cargo.toml b/web-demo/Cargo.toml
@@ -21,6 +21,11 @@ rev = "325814fd90767b5e43c72bc2eb65e14ff0b8746c"
 wee_alloc = "0.4.5"
 churchroad = {path = "../churchroad"}
 
+egraph-serialize = "0.1"
+indexmap = "2.0.0"
+rustc-hash = "1.1.0"
+ordered-float = "3"
+
 log = "0.4.19"
 wasm-logger = "0.2"
 

diff --git a/web-demo/src/lib.rs b/web-demo/src/lib.rs
@@ -1,4 +1,8 @@
-#![allow(clippy::unused_unit)] // weird clippy bug with wasm-bindgen
+#![allow(clippy::unused_unit)]
+use egglog::SerializeConfig;
+use egraph_serialize::{ClassId, NodeId};
+use indexmap::IndexMap;
+// weird clippy bug with wasm-bindgen
 use log::{Level, Log, Metadata, Record};
 use wasm_bindgen::prelude::*;
 use web_sys::console;
@@ -18,8 +22,17 @@ pub fn run_program(input: &str) -> Result {
     match egraph.parse_and_run_program(input) {
         Ok(outputs) => {
             let serialized = egraph.serialize_for_graphviz(false);
+            let serialized2 = egraph.serialize(SerializeConfig::default());
+            let choices = GreedyDagExtractor::default().extract(&serialized2, &[]);
+            let mut out = churchroad::to_verilog_egraph_serialize(&serialized2, &choices, "");
+
+            out.push_str("\n\nOutputs:\n");
+            for (i, output) in outputs.iter().enumerate() {
+                out.push_str(&format!("Output {}: {}\n", i, output));
+            }
+
             Result {
-                text: outputs.join("<br>"),
+                text: out,
                 dot: serialized.to_dot(),
             }
         }
@@ -111,3 +124,82 @@ pub fn init() {
         Err(e) => console::error_1(&JsValue::from(e.to_string())),
     }
 }
+
+
+use rustc_hash::FxHashMap;
+
+pub type Cost = ordered_float::NotNan<f64>;
+
+struct CostSet {
+    costs: FxHashMap<ClassId, Cost>,
+    total: Cost,
+    choice: NodeId,
+}
+
+#[derive(Default)]
+pub struct GreedyDagExtractor;
+impl GreedyDagExtractor {
+    fn extract(&self, egraph: &egraph_serialize::EGraph, _roots: &[egraph_serialize::ClassId]) -> IndexMap<ClassId, NodeId>
+    {
+        let mut costs = FxHashMap::<ClassId, CostSet>::with_capacity_and_hasher(
+            egraph.classes().len(),
+            Default::default(),
+        );
+
+        let mut keep_going = true;
+
+        let mut i = 0;
+        while keep_going {
+            i += 1;
+            println!("iteration {}", i);
+            keep_going = false;
+
+            'node_loop: for (node_id, node) in &egraph.nodes {
+                let cid = egraph.nid_to_cid(node_id);
+                let mut cost_set = CostSet {
+                    costs: Default::default(),
+                    total: Cost::default(),
+                    choice: node_id.clone(),
+                };
+
+                // compute the cost set from the children
+                for child in &node.children {
+                    let child_cid = egraph.nid_to_cid(child);
+                    if let Some(child_cost_set) = costs.get(child_cid) {
+                        // prevent a cycle
+                        if child_cost_set.costs.contains_key(cid) {
+                            continue 'node_loop;
+                        }
+                        cost_set.costs.extend(child_cost_set.costs.clone());
+                    } else {
+                        continue 'node_loop;
+                    }
+                }
+
+                // add this node
+                cost_set.costs.insert(cid.clone(), node.cost);
+
+                dbg!(node.cost);
+
+                cost_set.total = cost_set.costs.values().sum();
+
+                // if the cost set is better than the current one, update it
+                if let Some(old_cost_set) = costs.get(cid) {
+                    if cost_set.total < old_cost_set.total {
+                        costs.insert(cid.clone(), cost_set);
+                        keep_going = true;
+                    }
+                } else {
+                    costs.insert(cid.clone(), cost_set);
+                    keep_going = true;
+                }
+            }
+        }
+
+        let mut result = IndexMap::default();
+        for (cid, cost_set) in costs {
+            result.insert(cid, cost_set.choice);
+        }
+        result
+    }
+}