[cider/flat] Print the flat env port states (#1580)

* let the env hold a reference to the context

* pretty print all the env ports

* fix the printing

* add an explanatory comment for the iterator

interp/src/flatten/flat_ir/base.rs

@@ -1,4 +1,7 @@
-use std::{num::NonZeroU32, ops::Add};
+use std::{
+    num::NonZeroU32,
+    ops::{Add, Sub},
+};
 
 use crate::flatten::structures::index_trait::{
     impl_index, impl_index_nonzero, IndexRange, IndexRef,
@@ -379,3 +382,67 @@ impl Add<&LocalRefCellOffset> for &BaseIndices {
         GlobalRefCellId::new(self.ref_cell_base.index() + rhs.index())
     }
 }
+
+impl Sub<&BaseIndices> for GlobalPortId {
+    type Output = LocalPortOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalPortOffset::new(self.index() - rhs.port_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for GlobalRefPortId {
+    type Output = LocalRefPortOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalRefPortOffset::new(self.index() - rhs.ref_port_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for GlobalCellId {
+    type Output = LocalCellOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalCellOffset::new(self.index() - rhs.cell_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for GlobalRefCellId {
+    type Output = LocalRefCellOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalRefCellOffset::new(self.index() - rhs.ref_cell_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for &GlobalPortId {
+    type Output = LocalPortOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalPortOffset::new(self.index() - rhs.port_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for &GlobalRefPortId {
+    type Output = LocalRefPortOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalRefPortOffset::new(self.index() - rhs.ref_port_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for &GlobalCellId {
+    type Output = LocalCellOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalCellOffset::new(self.index() - rhs.cell_base.index())
+    }
+}
+
+impl Sub<&BaseIndices> for &GlobalRefCellId {
+    type Output = LocalRefCellOffset;
+
+    fn sub(self, rhs: &BaseIndices) -> Self::Output {
+        LocalRefCellOffset::new(self.index() - rhs.ref_cell_base.index())
+    }
+}

interp/src/flatten/mod.rs

@@ -12,4 +12,5 @@ pub fn flat_main(ctx: &calyx_ir::Context) {
 
     let env = Environment::new(&i_ctx);
     env.print_env_stats();
+    env.print_env();
 }

interp/src/flatten/structures/environment.rs

@@ -1,3 +1,5 @@
+use itertools::Itertools;
+
 use super::{context::Context, indexed_map::IndexedMap};
 use crate::{
     flatten::{
@@ -73,7 +75,7 @@ pub(crate) struct ProgramCounter {
 }
 
 #[derive(Debug)]
-pub struct Environment {
+pub struct Environment<'a> {
     /// A map from global port IDs to their current values.
     ports: PortMap,
     /// A map from global cell IDs to their current state and execution info.
@@ -85,10 +87,13 @@ pub struct Environment {
 
     /// The program counter for the whole program execution.
     pcs: ProgramCounter,
+
+    /// The immutable context. This is retained for ease of use.
+    ctx: &'a Context,
 }
 
-impl Environment {
-    pub fn new(ctx: &Context) -> Self {
+impl<'a> Environment<'a> {
+    pub fn new(ctx: &'a Context) -> Self {
         let root = ctx.entry_point;
         let aux = &ctx.secondary[root];
 
@@ -102,29 +107,30 @@ impl Environment {
                 aux.ref_port_offset_map.count(),
             ),
             pcs: ProgramCounter {},
+            ctx,
         };
 
         let root_node = CellLedger::comp(root, &env);
         let root = env.cells.push(root_node);
-        env.layout_component(root, ctx);
+        env.layout_component(root);
 
         env
     }
 
-    fn layout_component(&mut self, comp: GlobalCellId, ctx: &Context) {
+    fn layout_component(&mut self, comp: GlobalCellId) {
         let ComponentLedger {
             index_bases,
             comp_id,
         } = self.cells[comp]
             .as_comp()
             .expect("Called layout component with a non-component cell.");
-        let comp_aux = &ctx.secondary[*comp_id];
+        let comp_aux = &self.ctx.secondary[*comp_id];
 
         let comp_id = *comp_id;
 
         // first layout the signature
         for sig_port in comp_aux.signature.iter() {
-            let width = ctx.lookup_port_def(&comp_id, sig_port).width;
+            let width = self.ctx.lookup_port_def(&comp_id, sig_port).width;
             let idx = self.ports.push(Value::zeroes(width));
             debug_assert_eq!(index_bases + sig_port, idx);
         }
@@ -136,18 +142,21 @@ impl Environment {
 
             //go
             let go = self.ports.push(Value::bit_low());
-            debug_assert_eq!(go, index_bases + ctx.primary[group_idx].go);
+            debug_assert_eq!(go, index_bases + self.ctx.primary[group_idx].go);
 
             //done
             let done = self.ports.push(Value::bit_low());
-            debug_assert_eq!(done, index_bases + ctx.primary[group_idx].done);
+            debug_assert_eq!(
+                done,
+                index_bases + self.ctx.primary[group_idx].done
+            );
         }
 
         for (cell_off, def_idx) in comp_aux.cell_offset_map.iter() {
-            let info = &ctx.secondary[*def_idx];
+            let info = &self.ctx.secondary[*def_idx];
             if !info.prototype.is_component() {
                 for port in info.ports.iter() {
-                    let width = ctx.lookup_port_def(&comp_id, port).width;
+                    let width = self.ctx.lookup_port_def(&comp_id, port).width;
                     let idx = self.ports.push(Value::zeroes(width));
                     debug_assert_eq!(
                         &self.cells[comp].as_comp().unwrap().index_bases + port,
@@ -171,13 +180,13 @@ impl Environment {
                     cell
                 );
 
-                self.layout_component(cell, ctx);
+                self.layout_component(cell);
             }
         }
 
         // ref cells and ports are initialized to None
         for (ref_cell, def_idx) in comp_aux.ref_cell_offset_map.iter() {
-            let info = &ctx.secondary[*def_idx];
+            let info = &self.ctx.secondary[*def_idx];
             for port_idx in info.ports.iter() {
                 let port_actual = self.ref_ports.push(None);
                 debug_assert_eq!(
@@ -193,6 +202,72 @@ impl Environment {
         }
     }
 
+    pub fn print_env(&self) {
+        let root_idx = GlobalCellId::new(0);
+        let mut hierarchy = Vec::new();
+        self.print_component(root_idx, &mut hierarchy)
+    }
+
+    fn print_component(
+        &self,
+        target: GlobalCellId,
+        hierarchy: &mut Vec<GlobalCellId>,
+    ) {
+        let info = self.cells[target].as_comp().unwrap();
+        let comp = &self.ctx.secondary[info.comp_id];
+        hierarchy.push(target);
+
+        // This funky iterator chain first pulls the first element (the
+        // entrypoint) and extracts its name. Subsequent element are pairs of
+        // global offsets produced by a staggered iteration, yielding `(root,
+        // child)` then `(child, grandchild)` and so on. All the strings are
+        // finally collected and concatenated with a `.` separator to produce
+        // the fully qualified name prefix for the given component instance.
+        let name_prefix = hierarchy
+            .first()
+            .iter()
+            .map(|x| {
+                let info = self.cells[**x].as_comp().unwrap();
+                let prior_comp = &self.ctx.secondary[info.comp_id];
+                &self.ctx.secondary[prior_comp.name]
+            })
+            .chain(hierarchy.iter().zip(hierarchy.iter().skip(1)).map(
+                |(l, r)| {
+                    let info = self.cells[*l].as_comp().unwrap();
+                    let prior_comp = &self.ctx.secondary[info.comp_id];
+                    let local_target = r - (&info.index_bases);
+
+                    let def_idx = &prior_comp.cell_offset_map[local_target];
+
+                    let id = &self.ctx.secondary[*def_idx];
+                    &self.ctx.secondary[id.name]
+                },
+            ))
+            .join(".");
+
+        for (cell_off, def_idx) in comp.cell_offset_map.iter() {
+            let definition = &self.ctx.secondary[*def_idx];
+
+            println!("{}.{}", name_prefix, self.ctx.secondary[definition.name]);
+            for port in definition.ports.iter() {
+                let definition =
+                    &self.ctx.secondary[comp.port_offset_map[port]];
+                println!(
+                    "    {}: {}",
+                    self.ctx.secondary[definition.name],
+                    self.ports[&info.index_bases + port]
+                );
+            }
+
+            if definition.prototype.is_component() {
+                let child_target = &info.index_bases + cell_off;
+                self.print_component(child_target, hierarchy);
+            }
+        }
+
+        hierarchy.pop();
+    }
+
     pub fn print_env_stats(&self) {
         println!("Environment Stats:");
         println!("  Ports: {}", self.ports.len());

interp/src/flatten/structures/indexed_map.rs

@@ -106,6 +106,10 @@ where
     pub fn capacity(&self) -> usize {
         self.data.capacity()
     }
+
+    pub fn first(&self) -> Option<&D> {
+        self.data.first()
+    }
 }
 
 impl<T, K> Default for IndexedMap<K, T>