add dataflow algorithm

Author: UnsignedByte

crates/ir/src/control.rs

@@ -1,8 +1,8 @@
-use crate::{Expr, Prop, Time};
+use crate::{Direction, Expr, PortOwner, Prop, Time};
 
 use super::{
     Access, CompIdx, Component, Ctx, Event, ExprIdx, Fact, Foreign, InfoIdx,
-    InstIdx, InvIdx, ParamIdx, PortIdx, PropIdx, Range, TimeIdx, TimeSub,
+    InstIdx, InvIdx, ParamIdx, Port, PortIdx, PropIdx, Range, TimeIdx, TimeSub,
 };
 
 #[derive(Clone, PartialEq, Eq)]
@@ -171,6 +171,52 @@ impl InvIdx {
         let inst = self.inst(ctx);
         inst.comp(ctx)
     }
+
+    // Input ports to the iterator
+    pub fn inputs<C>(self, ctx: &C) -> impl Iterator<Item = PortIdx> + '_
+    where
+        C: Ctx<Invoke> + Ctx<Port>,
+    {
+        let inv = ctx.get(self);
+        inv.ports.iter().filter_map(|&port| {
+            let Port {
+                owner: PortOwner::Inv { dir, .. },
+                ..
+            } = ctx.get(port)
+            else {
+                unreachable!("Port is not an invocation port")
+            };
+
+            if *dir == Direction::In {
+                Some(port)
+            } else {
+                None
+            }
+        })
+    }
+
+    // Output ports to the iterator
+    pub fn outputs<C>(self, ctx: &C) -> impl Iterator<Item = PortIdx> + '_
+    where
+        C: Ctx<Invoke> + Ctx<Port>,
+    {
+        let inv = ctx.get(self);
+        inv.ports.iter().filter_map(|&port| {
+            let Port {
+                owner: PortOwner::Inv { dir, .. },
+                ..
+            } = ctx.get(port)
+            else {
+                unreachable!("Port is not an invocation port")
+            };
+
+            if *dir == Direction::Out {
+                Some(port)
+            } else {
+                None
+            }
+        })
+    }
 }
 
 #[derive(Clone, PartialEq, Eq)]

crates/ir/src/utils/dataflow.rs

@@ -1,38 +1,52 @@
-use std::collections::HashMap;
+use itertools::Itertools;
 
 use super::DenseIndexInfo;
 use crate::{self as ir, Ctx};
 
 /// Dataflow graph of a component
 #[derive(Clone)]
 pub struct Dataflow {
-    preds: HashMap<ir::Port, Vec<ir::Port>>,
-    succs: HashMap<ir::Port, Vec<ir::Port>>,
+    pub preds: DenseIndexInfo<ir::Port, Vec<ir::PortIdx>>,
+    pub succs: DenseIndexInfo<ir::Port, Vec<ir::PortIdx>>,
 }
 
 impl From<&ir::Component> for Dataflow {
     fn from(comp: &ir::Component) -> Self {
-        let mut preds = DenseIndexInfo::with_default(comp.ports().len());
-        let mut succs = DenseIndexInfo::with_default(comp.ports().len());
+        let mut preds: DenseIndexInfo<ir::Port, Vec<_>> =
+            DenseIndexInfo::with_default(comp.ports().len());
+        let mut succs: DenseIndexInfo<ir::Port, Vec<_>> =
+            DenseIndexInfo::with_default(comp.ports().len());
 
         for cmd in &comp.cmds {
             match cmd {
                 crate::Command::Invoke(idx) => {
-                    let inv = comp.get(*idx);
-                    for port in &inv.ports {
-                        for &pred in &port.pred {
-                            preds.insert(*port, pred);
+                    let inputs = idx.inputs(comp);
+                    let outputs = idx.outputs(comp).collect_vec();
+
+                    for pred in inputs {
+                        for succ in &outputs {
+                            preds.get_mut(*succ).push(pred);
+                            succs.get_mut(pred).push(*succ);
                         }
                     }
                 }
-                crate::Command::Instance(idx) => todo!(),
-                crate::Command::BundleDef(idx) => todo!(),
-                crate::Command::Connect(connect) => todo!(),
-                crate::Command::Let(_) => todo!(),
-                crate::Command::ForLoop(_) => todo!(),
-                crate::Command::If(_) => todo!(),
-                crate::Command::Fact(fact) => todo!(),
-                crate::Command::Exists(exists) => todo!(),
+                crate::Command::Connect(ir::Connect { src, dst, .. }) => {
+                    assert!(src.is_port(comp) && dst.is_port(comp), "Bundles should be resolved before constructing dataflow");
+
+                    preds.get_mut(dst.port).push(src.port);
+                    succs.get_mut(src.port).push(dst.port);
+                }
+                crate::Command::BundleDef(_)
+                | crate::Command::ForLoop(_)
+                | crate::Command::If(_) => {
+                    unreachable!(
+                        "Components should be monomorphic and bundle-free"
+                    )
+                }
+                crate::Command::Instance(_)
+                | crate::Command::Let(_)
+                | crate::Command::Fact(_)
+                | crate::Command::Exists(_) => {}
             }
         }