fix recursisve node checks for custom inverses

uiua-lang · Nov 11, 2024 · df674ae · df674ae
1 parent efecea4
commit df674ae
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -9,7 +9,7 @@ name = "uiua"
 readme = "readme.md"
 repository = "https://github.com/uiua-lang/uiua"
 rust-version = "1.78"
-version = "0.14.0-dev.1"
+version = "0.14.0-dev.2"
 
 [dependencies]
 # Core dependencies

diff --git a/src/tree.rs b/src/tree.rs
@@ -232,6 +232,13 @@ impl CustomInverse {
             }
         })
     }
+    /// Iterate over all nodes
+    pub fn nodes(&self) -> impl Iterator<Item = &SigNode> {
+        (self.normal.as_ref().into_iter())
+            .chain(self.un.as_ref())
+            .chain(self.anti.as_ref())
+            .chain(self.under.as_ref().into_iter().flat_map(|(b, a)| [a, b]))
+    }
 }
 
 impl Default for Node {
@@ -611,6 +618,9 @@ impl Node {
                 Node::Switch { branches, .. } => branches
                     .iter()
                     .all(|br| recurse(&br.node, purity, asm, visited)),
+                Node::CustomInverse(cust, _) => (cust.normal.as_ref().ok())
+                    .or(cust.un.as_ref())
+                    .is_some_and(|sn| recurse(&sn.node, purity, asm, visited)),
                 _ => true,
             };
             visited.truncate(len);
@@ -651,6 +661,9 @@ impl Node {
                 Node::Switch { branches, .. } => {
                     branches.iter().all(|br| recurse(&br.node, asm, visited))
                 }
+                Node::CustomInverse(cust, _) => (cust.normal.as_ref().ok())
+                    .or(cust.un.as_ref())
+                    .is_some_and(|sn| recurse(&sn.node, asm, visited)),
                 _ => true,
             };
             visited.truncate(len);
@@ -675,6 +688,9 @@ impl Node {
                 Node::Switch { branches, .. } => {
                     branches.iter().any(|br| recurse(&br.node, asm, visited))
                 }
+                Node::CustomInverse(cust, _) => (cust.normal.as_ref().ok())
+                    .or(cust.un.as_ref())
+                    .is_some_and(|sn| recurse(&sn.node, asm, visited)),
                 _ => false,
             };
             visited.truncate(len);