From 0b8c90493b91e9bc878ed613158e794f78ec878a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20L=C3=A4ufer?= <laeufer@cornell.edu>
Date: Tue, 3 Dec 2024 10:39:01 -0500
Subject: [PATCH] fix simulator

---
 examples/sim.rs                 | 21 +++++++---------
 src/system/transition_system.rs | 44 +++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/examples/sim.rs b/examples/sim.rs
index 0344582..75f109f 100644
--- a/examples/sim.rs
+++ b/examples/sim.rs
@@ -8,7 +8,7 @@ use patronus::expr::*;
 use patronus::sim::*;
 use patronus::system::*;
 use patronus::*;
-use std::collections::{HashMap, HashSet};
+use rustc_hash::FxHashMap;
 use std::io::BufRead;
 
 #[derive(Parser, Debug)]
@@ -91,7 +91,7 @@ fn main() {
         std::fs::File::open(testbench_file).expect("Failed to load testbench file"),
     );
 
-    let name_to_ref: HashMap<String, ExprRef> = todo!(); // sys.generate_name_to_ref(&ctx);
+    let name_to_ref = sys.get_name_map(&ctx);
     let (inputs, outputs) =
         read_header(&mut tb, &name_to_ref, &sys, &ctx).expect("Failed to read testbench header");
     if args.verbose {
@@ -131,7 +131,7 @@ type IOInfo = Vec<(usize, ExprRef, String, WidthInt)>;
 /// Correlates the header with the inputs and outputs of the system.
 fn read_header(
     input: &mut impl BufRead,
-    name_to_ref: &HashMap<String, ExprRef>,
+    name_to_ref: &FxHashMap<String, ExprRef>,
     sys: &TransitionSystem,
     ctx: &Context,
 ) -> std::io::Result<(IOInfo, IOInfo)> {
@@ -143,14 +143,11 @@ fn read_header(
         let name = cell.trim();
         if let Some(signal_ref) = name_to_ref.get(name) {
             let width = signal_ref.get_bv_type(ctx).unwrap();
-            todo!();
-            // let signal = sys.get_signal(*signal_ref).unwrap();
-            // if signal.is_input() {
-            //     inputs.push((cell_id, *signal_ref, name.to_string(), width));
-            // }
-            // if signal.is_output() {
-            //     outputs.push((cell_id, *signal_ref, name.to_string(), width));
-            // }
+            if let Some(signal_ref) = sys.lookup_input(ctx, name) {
+                inputs.push((cell_id, signal_ref, name.to_string(), width));
+            } else if let Some(signal_ref) = sys.lookup_output(ctx, name) {
+                outputs.push((cell_id, signal_ref, name.to_string(), width));
+            }
         }
     }
     Ok((inputs, outputs))
@@ -191,7 +188,7 @@ fn do_step(
     // calculate the output values
     sim.update();
 
-    // print values if the option is enables
+    // print values if the option is enabled
     if !signal_to_print.is_empty() {
         println!();
         for (name, expr) in signal_to_print.iter() {
diff --git a/src/system/transition_system.rs b/src/system/transition_system.rs
index 46a1e03..f663a91 100644
--- a/src/system/transition_system.rs
+++ b/src/system/transition_system.rs
@@ -196,4 +196,48 @@ impl TransitionSystem {
 
         out
     }
+
+    /// Creates a map from signal name to expression
+    pub fn get_name_map(&self, ctx: &Context) -> FxHashMap<String, ExprRef> {
+        let mut m = FxHashMap::default();
+        for out in self.outputs.iter() {
+            m.insert(ctx[out.name].to_string(), out.expr);
+        }
+        for &e in self
+            .bad_states
+            .iter()
+            .chain(self.constraints.iter())
+            .chain(self.inputs.iter())
+            .chain(self.states.iter().map(|s| &s.symbol))
+        {
+            if let Some(name) = ctx[e].get_symbol_name(ctx) {
+                m.insert(name.to_string(), e);
+            }
+            if let Some(name) = self.names[e] {
+                m.insert(ctx[name].to_string(), e);
+            }
+        }
+        m
+    }
+
+    /// Returns input by name.
+    pub fn lookup_input(&self, ctx: &Context, name: &str) -> Option<ExprRef> {
+        self.inputs
+            .iter()
+            .find(|&&i| {
+                ctx[i]
+                    .get_symbol_name(ctx)
+                    .map(|i_name| i_name == name)
+                    .unwrap_or(false)
+            })
+            .cloned()
+    }
+
+    /// Returns output by name.
+    pub fn lookup_output(&self, ctx: &Context, name: &str) -> Option<ExprRef> {
+        self.outputs
+            .iter()
+            .find(|&&o| ctx[o.name] == name)
+            .map(|o| o.expr)
+    }
 }