From df674ae1bc41077fce1bb1d1cee5867745d41a67 Mon Sep 17 00:00:00 2001
From: Kai Schmidt <kaikaliischmidt@gmail.com>
Date: Mon, 11 Nov 2024 07:20:33 -0800
Subject: [PATCH] fix recursisve node checks for custom inverses

---
 Cargo.lock  |  2 +-
 Cargo.toml  |  2 +-
 src/tree.rs | 16 ++++++++++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 86a58b2d8..16b6daa08 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4077,7 +4077,7 @@ checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9"
 
 [[package]]
 name = "uiua"
-version = "0.14.0-dev.1"
+version = "0.14.0-dev.2"
 dependencies = [
  "arboard",
  "bitflags 2.5.0",
diff --git a/Cargo.toml b/Cargo.toml
index 1ea11308e..3aeebf406 100644
--- 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
index dae8fbebe..2ba0222ba 100644
--- 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);