From a4169f1641b5013cb6715ef913466f57fea6e5d6 Mon Sep 17 00:00:00 2001
From: Anton Reinhard <anton.reinhard@proton.me>
Date: Wed, 8 May 2024 13:27:41 +0200
Subject: [PATCH] Add trie workaround

---
 src/trie.jl | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/trie.jl b/src/trie.jl
index b3babca..441890a 100644
--- a/src/trie.jl
+++ b/src/trie.jl
@@ -48,7 +48,26 @@ function insert_helper!(
     trie::NodeIdTrie{NodeType},
     node::NodeType,
     depth::Int,
-) where {TaskType <: AbstractTask, NodeType <: Union{DataTaskNode{TaskType}, ComputeTaskNode{TaskType}}}
+) where {TaskType <: AbstractDataTask, NodeType <: DataTaskNode{TaskType}}
+    if (length(children(node)) == depth)
+        push!(trie.value, node)
+        return nothing
+    end
+
+    depth = depth + 1
+    id = node.children[depth].id
+
+    if (!haskey(trie.children, id))
+        trie.children[id] = NodeIdTrie{NodeType}()
+    end
+    return insert_helper!(trie.children[id], node, depth)
+end
+# TODO: Remove this workaround once https://github.com/JuliaLang/julia/issues/54404 is fixed in julia 1.10+
+function insert_helper!(
+    trie::NodeIdTrie{NodeType},
+    node::NodeType,
+    depth::Int,
+) where {TaskType <: AbstractComputeTask, NodeType <: ComputeTaskNode{TaskType}}
     if (length(children(node)) == depth)
         push!(trie.value, node)
         return nothing
@@ -63,6 +82,7 @@ function insert_helper!(
     return insert_helper!(trie.children[id], node, depth)
 end
 
+
 """
     insert!(trie::NodeTrie, node::Node)
 
@@ -71,7 +91,17 @@ Insert the given node into the trie. It's sorted by its type in the first layer,
 function insert!(
     trie::NodeTrie,
     node::NodeType,
-) where {TaskType <: AbstractTask, NodeType <: Union{DataTaskNode{TaskType}, ComputeTaskNode{TaskType}}}
+) where {TaskType <: AbstractDataTask, NodeType <: DataTaskNode{TaskType}}
+    if (!haskey(trie.children, NodeType))
+        trie.children[NodeType] = NodeIdTrie{NodeType}()
+    end
+    return insert_helper!(trie.children[NodeType], node, 0)
+end
+# TODO: Remove this workaround once https://github.com/JuliaLang/julia/issues/54404 is fixed in julia 1.10+
+function insert!(
+    trie::NodeTrie,
+    node::NodeType,
+) where {TaskType <: AbstractComputeTask, NodeType <: ComputeTaskNode{TaskType}}
     if (!haskey(trie.children, NodeType))
         trie.children[NodeType] = NodeIdTrie{NodeType}()
     end