working on deduce_resets pass

crates/fayalite/src/module/transform/deduce_resets.rs

@@ -10,16 +10,13 @@ use petgraph::{
 use crate::{
     bundle::{BundleField, BundleType},
     enum_::{EnumType, EnumVariant},
-    expr::{
-        ops,
-        target::{TargetBase, TargetPathArrayElement, TargetPathBundleField, TargetPathElement},
-        Flow,
-    },
+    expr::{ops, ExprEnum, Flow},
     intern::{Intern, Interned, Memoize},
     module::{
-        AnnotatedModuleIO, Block, ExprInInstantiatedModule, ExternModuleBody, InstantiatedModule,
-        ModuleBody, ModuleIO, NormalModuleBody, Stmt, StmtConnect, StmtDeclaration, StmtFormal,
-        StmtIf, StmtInstance, StmtMatch, StmtReg, StmtWire, TargetInInstantiatedModule,
+        transform::visit::{Visit, Visitor},
+        AnnotatedModuleIO, Block, ExternModuleBody, InstantiatedModule, ModuleBody, ModuleIO,
+        NormalModuleBody, Stmt, StmtConnect, StmtDeclaration, StmtFormal, StmtIf, StmtInstance,
+        StmtMatch, StmtReg, StmtWire, TargetInInstantiatedModule,
     },
     prelude::*,
 };
@@ -51,7 +48,7 @@ impl From<DeduceResetsError> for std::io::Error {
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 enum ResetTarget {
     Base {
-        base: Interned<TargetBase>,
+        base: Expr<CanonicalType>,
     },
     BundleField {
         parent: Interned<ResetTarget>,
@@ -75,7 +72,7 @@ enum ResetTarget {
 impl ResetTarget {
     fn canonical_ty(self) -> CanonicalType {
         match self {
-            ResetTarget::Base { base } => base.canonical_ty(),
+            ResetTarget::Base { base } => Expr::ty(base),
             ResetTarget::BundleField { field_ty, .. } => field_ty,
             ResetTarget::EnumVariant { variant_ty, .. } => variant_ty,
             ResetTarget::ArraysElements { element_ty, .. } => element_ty,
@@ -89,7 +86,7 @@ impl ResetTarget {
             | ResetTarget::ArraysElements { parent, .. } => Some(parent),
         }
     }
-    fn base(mut self) -> Interned<TargetBase> {
+    fn base(mut self) -> Expr<CanonicalType> {
         loop {
             match self {
                 ResetTarget::Base { base } => break base,
@@ -150,14 +147,6 @@ impl ResetTarget {
     }
 }
 
-impl<T: Into<TargetBase>> From<T> for ResetTarget {
-    fn from(base: T) -> Self {
-        ResetTarget::Base {
-            base: TargetBase::intern_sized(base.into()),
-        }
-    }
-}
-
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 struct ResetTargetInInstantiatedModule {
     instantiated_module: InstantiatedModule,
@@ -289,6 +278,23 @@ impl<C: FromIterator<A>, A, Step: ProcessStep> FromIterator<Processed<A, Step>>
     }
 }
 
+fn for_each_subexpr_exclusive<E>(
+    expr: Expr<CanonicalType>,
+    f: impl FnMut(Expr<CanonicalType>) -> Result<(), E>,
+) -> Result<(), E> {
+    struct MyVisitor<F> {
+        f: F,
+    }
+    impl<F: FnMut(Expr<CanonicalType>) -> Result<(), E>, E> Visitor for MyVisitor<F> {
+        type Error = E;
+        fn visit_expr_enum(&mut self, v: &ExprEnum) -> Result<(), Self::Error> {
+            (self.f)(v.to_expr())?;
+            v.default_visit(self)
+        }
+    }
+    Expr::expr_enum(expr).default_visit(&mut MyVisitor { f })
+}
+
 struct AddNodesToGraphStep;
 
 impl ProcessStep for AddNodesToGraphStep {
@@ -342,6 +348,41 @@ impl RunProcessStep<AddNodesToGraphStep> for ResetTarget {
     }
 }
 
+impl RunProcessStep<AddEdgesToGraphStep> for ResetTarget {
+    fn run_process_step(
+        self,
+        _instantiated_module: InstantiatedModule,
+        _state: &mut State,
+        _step: &mut AddEdgesToGraphStep,
+    ) -> Result<Processed<Self, AddEdgesToGraphStep>, Infallible> {
+        Ok(Processed(()))
+    }
+}
+
+impl<Step: ProcessStep> RunProcessStep<Step> for Expr<CanonicalType>
+where
+    ResetTarget: RunProcessStep<Step>,
+{
+    fn run_process_step(
+        self,
+        instantiated_module: InstantiatedModule,
+        state: &mut State,
+        step: &mut Step,
+    ) -> Result<Processed<Self, Step>, Step::Error> {
+        let retval = ResetTarget::Base { base: self }
+            .run_process_step(instantiated_module, state, step)?
+            .map(|target| match target {
+                ResetTarget::Base { base } => base,
+                _ => unreachable!(),
+            });
+        for_each_subexpr_exclusive(self, |base| {
+            ResetTarget::Base { base }.run_process_step(instantiated_module, state, step)?;
+            Ok(())
+        })?;
+        Ok(retval)
+    }
+}
+
 struct ConnectAndLhsInstantiatedModule {
     lhs_instantiated_module: InstantiatedModule,
     lhs: Expr<CanonicalType>,
@@ -349,24 +390,49 @@ struct ConnectAndLhsInstantiatedModule {
     source_location: SourceLocation,
 }
 
-impl<Step: ProcessStep> RunProcessStep<Step> for ConnectAndLhsInstantiatedModule {
+impl RunProcessStep<AddNodesToGraphStep> for ConnectAndLhsInstantiatedModule {
+    fn run_process_step(
+        self,
+        _instantiated_module: InstantiatedModule,
+        _state: &mut State,
+        _step: &mut AddNodesToGraphStep,
+    ) -> Result<Processed<Self, AddNodesToGraphStep>, Infallible> {
+        Ok(Processed(()))
+    }
+}
+
+impl RunProcessStep<AddEdgesToGraphStep> for ConnectAndLhsInstantiatedModule {
     fn run_process_step(
         self,
         rhs_instantiated_module: InstantiatedModule,
         state: &mut State,
-        step: &mut Step,
-    ) -> Result<Processed<Self, Step>, Step::Error> {
+        step: &mut AddEdgesToGraphStep,
+    ) -> Result<Processed<Self, AddEdgesToGraphStep>, Infallible> {
         let Self {
             lhs_instantiated_module,
             lhs,
             rhs,
             source_location,
         } = self;
-        todo!();
+        match Expr::ty(lhs) {
+            CanonicalType::UInt(ty) => todo!(),
+            CanonicalType::SInt(ty) => todo!(),
+            CanonicalType::Bool(ty) => todo!(),
+            CanonicalType::AsyncReset(ty) => todo!(),
+            CanonicalType::SyncReset(ty) => todo!(),
+            CanonicalType::Reset(ty) => todo!(),
+            CanonicalType::Clock(ty) => todo!(),
+            CanonicalType::Array(ty) => todo!(),
+            CanonicalType::Enum(ty) => todo!(),
+            CanonicalType::Bundle(ty) => todo!(),
+        }
     }
 }
 
-impl<Step: ProcessStep> RunProcessStep<Step> for StmtConnect {
+impl<Step: ProcessStep> RunProcessStep<Step> for StmtConnect
+where
+    ConnectAndLhsInstantiatedModule: RunProcessStep<Step>,
+{
     fn run_process_step(
         self,
         instantiated_module: InstantiatedModule,
@@ -444,16 +510,12 @@ where
         state: &mut State,
         step: &mut Step,
     ) -> Result<Processed<Self, Step>, Step::Error> {
-        Ok(ResetTarget::from(self)
+        Ok(self
+            .to_expr()
             .run_process_step(instantiated_module, state, step)?
-            .map(|target| {
-                let ResetTarget::Base { base } = target else {
-                    unreachable!();
-                };
-                let TargetBase::ModuleIO(module_io) = *base else {
-                    unreachable!();
-                };
-                module_io
+            .map(|expr| match *Expr::expr_enum(expr) {
+                ExprEnum::ModuleIO(module_io) => module_io,
+                _ => unreachable!(),
             }))
     }
 }
@@ -469,16 +531,12 @@ where
         step: &mut Step,
     ) -> Result<Processed<Self, Step>, Step::Error> {
         let Self { annotations, wire } = self;
-        Ok(ResetTarget::from(wire)
+        Ok(wire
+            .to_expr()
             .run_process_step(instantiated_module, state, step)?
-            .map(|target| {
-                let ResetTarget::Base { base } = target else {
-                    unreachable!();
-                };
-                let TargetBase::Wire(wire) = *base else {
-                    unreachable!();
-                };
-                Self { annotations, wire }
+            .map(|expr| match *Expr::expr_enum(expr) {
+                ExprEnum::Wire(wire) => Self { annotations, wire },
+                _ => unreachable!(),
             }))
     }
 }
@@ -494,16 +552,12 @@ where
         step: &mut Step,
     ) -> Result<Processed<Self, Step>, Step::Error> {
         let Self { annotations, reg } = self;
-        Ok(ResetTarget::from(reg)
+        Ok(reg
+            .to_expr()
             .run_process_step(instantiated_module, state, step)?
-            .map(|target| {
-                let ResetTarget::Base { base } = target else {
-                    unreachable!();
-                };
-                let TargetBase::Reg(reg) = *base else {
-                    unreachable!();
-                };
-                Self { annotations, reg }
+            .map(|expr| match *Expr::expr_enum(expr) {
+                ExprEnum::Reg(reg) => Self { annotations, reg },
+                _ => unreachable!(),
             }))
     }
 }
@@ -511,6 +565,7 @@ where
 impl<Step: ProcessStep> RunProcessStep<Step> for StmtInstance
 where
     ResetTarget: RunProcessStep<Step>,
+    ConnectAndLhsInstantiatedModule: RunProcessStep<Step>,
 {
     fn run_process_step(
         self,
@@ -563,19 +618,14 @@ where
             }
             .run_process_step(rhs_instantiated_module, state, step)?;
         }
-        Ok(ResetTarget::from(instance)
+        Ok(Expr::canonical(instance.to_expr())
             .run_process_step(instantiated_module, state, step)?
-            .map(|target| {
-                let ResetTarget::Base { base } = target else {
-                    unreachable!();
-                };
-                let TargetBase::Instance(instance) = *base else {
-                    unreachable!();
-                };
-                Self {
+            .map(|expr| match *Expr::expr_enum(expr) {
+                ExprEnum::Instance(instance) => Self {
                     annotations,
                     instance,
-                }
+                },
+                _ => unreachable!(),
             }))
     }
 }
@@ -583,6 +633,7 @@ where
 impl<Step: ProcessStep> RunProcessStep<Step> for StmtDeclaration
 where
     ResetTarget: RunProcessStep<Step>,
+    ConnectAndLhsInstantiatedModule: RunProcessStep<Step>,
 {
     fn run_process_step(
         self,
@@ -607,6 +658,7 @@ where
 impl<Step: ProcessStep> RunProcessStep<Step> for Stmt
 where
     ResetTarget: RunProcessStep<Step>,
+    ConnectAndLhsInstantiatedModule: RunProcessStep<Step>,
 {
     fn run_process_step(
         self,
@@ -637,6 +689,7 @@ where
 impl<Step: ProcessStep> RunProcessStep<Step> for Block
 where
     ResetTarget: RunProcessStep<Step>,
+    ConnectAndLhsInstantiatedModule: RunProcessStep<Step>,
 {
     fn run_process_step(
         self,
@@ -649,7 +702,11 @@ where
         // so always just use the old `memories` value. we add the ports to the graph anyway to make the other logic easier.
         for memory in memories {
             for port in memory.ports() {
-                ResetTarget::from(port).run_process_step(instantiated_module, state, step)?;
+                Expr::canonical(port.to_expr()).run_process_step(
+                    instantiated_module,
+                    state,
+                    step,
+                )?;
             }
         }
         let stmts = Result::<Processed<_, _>, _>::from_iter(
@@ -664,6 +721,7 @@ where
 impl<Step: ProcessStep> RunProcessStep<Step> for Module<Bundle>
 where
     ResetTarget: RunProcessStep<Step>,
+    ConnectAndLhsInstantiatedModule: RunProcessStep<Step>,
 {
     fn run_process_step(
         self,
@@ -722,7 +780,12 @@ pub fn deduce_resets(
         &mut state,
         &mut AddNodesToGraphStep,
     );
-    todo!("add edges");
+    let mut step = AddEdgesToGraphStep {
+        union_find: UnionFind::new(state.graph.node_count()),
+    };
+    let Ok(Processed(())) =
+        module.run_process_step(InstantiatedModule::Base(module), &mut state, &mut step);
+    let AddEdgesToGraphStep { union_find } = step;
     todo!("deduce types");
     Ok(todo!("transform module"))
 }