From 1e2831da47cca6df7e97f3d936fedbb5269f5f02 Mon Sep 17 00:00:00 2001
From: Jacob Lifshay <programmerjake@gmail.com>
Date: Mon, 30 Sep 2024 21:20:35 -0700
Subject: [PATCH] add validation of connects and matches when validating module

this is useful for catching errors in transformation passes
---
 crates/fayalite/src/module.rs | 87 ++++++++++++++++++++++-------------
 1 file changed, 56 insertions(+), 31 deletions(-)

diff --git a/crates/fayalite/src/module.rs b/crates/fayalite/src/module.rs
index 2b7534c..e88b37a 100644
--- a/crates/fayalite/src/module.rs
+++ b/crates/fayalite/src/module.rs
@@ -185,6 +185,40 @@ pub struct StmtConnect {
     pub source_location: SourceLocation,
 }
 
+impl StmtConnect {
+    #[track_caller]
+    fn assert_validity_with_original_types(&self, lhs_orig_ty: impl Type, rhs_orig_ty: impl Type) {
+        let Self {
+            lhs,
+            rhs,
+            source_location,
+        } = *self;
+        assert!(
+            Expr::ty(lhs).can_connect(Expr::ty(rhs)),
+            "can't connect types that are not equivalent:\nlhs type:\n{lhs_orig_ty:?}\nrhs type:\n{rhs_orig_ty:?}\nat: {source_location}",
+        );
+        assert!(
+            matches!(Expr::flow(lhs), Flow::Sink | Flow::Duplex),
+            "can't connect to source, connect lhs must have sink or duplex flow\nat: {source_location}"
+        );
+        assert!(
+            lhs.target().is_some(),
+            "can't connect to non-target\nat: {source_location}"
+        );
+        match Expr::flow(rhs) {
+            Flow::Source | Flow::Duplex => {}
+            Flow::Sink => assert!(
+                Expr::ty(rhs).is_passive(),
+                "can't connect from sink with non-passive type\nat: {source_location}"
+            ),
+        }
+    }
+    #[track_caller]
+    fn assert_validity(&self) {
+        self.assert_validity_with_original_types(Expr::ty(self.lhs), Expr::ty(self.rhs));
+    }
+}
+
 impl fmt::Debug for StmtConnect {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let Self {
@@ -283,6 +317,13 @@ pub struct StmtMatch<S: ModuleBuildingStatus = ModuleBuilt> {
     pub blocks: Interned<[S::Block]>,
 }
 
+impl StmtMatch {
+    #[track_caller]
+    fn assert_validity(&self) {
+        assert_eq!(Expr::ty(self.expr).variants().len(), self.blocks.len());
+    }
+}
+
 impl<S: ModuleBuildingStatus> fmt::Debug for StmtMatch<S> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let Self {
@@ -1657,11 +1698,13 @@ impl AssertValidityState {
         }
         for stmt in stmts {
             match stmt {
-                Stmt::Connect(StmtConnect {
-                    lhs,
-                    rhs,
-                    source_location,
-                }) => {
+                Stmt::Connect(connect) => {
+                    connect.assert_validity();
+                    let StmtConnect {
+                        lhs,
+                        rhs,
+                        source_location,
+                    } = connect;
                     self.set_connect_side_written(lhs, source_location, true, block);
                     self.set_connect_side_written(rhs, source_location, false, block);
                 }
@@ -1671,6 +1714,7 @@ impl AssertValidityState {
                     self.process_conditional_sub_blocks(block, sub_blocks)
                 }
                 Stmt::Match(match_stmt) => {
+                    match_stmt.assert_validity();
                     let sub_blocks = Vec::from_iter(
                         match_stmt
                             .blocks
@@ -2517,24 +2561,12 @@ pub fn connect_any_with_loc<Lhs: ToExpr, Rhs: ToExpr>(
     let rhs_orig = rhs.to_expr();
     let lhs = Expr::canonical(lhs_orig);
     let rhs = Expr::canonical(rhs_orig);
-    assert!(
-        Expr::ty(lhs).can_connect(Expr::ty(rhs)),
-        "can't connect types that are not equivalent:\nlhs type:\n{:?}\nrhs type:\n{:?}",
-        Expr::ty(lhs_orig),
-        Expr::ty(rhs_orig)
-    );
-    assert!(
-        matches!(Expr::flow(lhs), Flow::Sink | Flow::Duplex),
-        "can't connect to source, connect lhs must have sink or duplex flow"
-    );
-    assert!(lhs.target().is_some(), "can't connect to non-target");
-    match Expr::flow(rhs) {
-        Flow::Source | Flow::Duplex => {}
-        Flow::Sink => assert!(
-            Expr::ty(rhs).is_passive(),
-            "can't connect from sink with non-passive type"
-        ),
-    }
+    let connect = StmtConnect {
+        lhs,
+        rhs,
+        source_location,
+    };
+    connect.assert_validity_with_original_types(Expr::ty(lhs_orig), Expr::ty(rhs_orig));
     ModuleBuilder::with(|m| {
         m.impl_
             .borrow_mut()
@@ -2542,14 +2574,7 @@ pub fn connect_any_with_loc<Lhs: ToExpr, Rhs: ToExpr>(
             .builder_normal_body()
             .block(m.block_stack.top())
             .stmts
-            .push(
-                StmtConnect {
-                    lhs,
-                    rhs,
-                    source_location,
-                }
-                .into(),
-            );
+            .push(connect.into());
     });
 }