Define Base.delete! for Tries

src/trie.jl

@@ -47,6 +47,27 @@ function Base.getindex(t::Trie, key)
     throw(KeyError("key not found: $key"))
 end
 
+function Base.delete!(t::Trie, key; at_key=true)
+    if isempty(key) || (at_key && !haskey(t, key))
+        return t
+    end
+    node = subtrie(t, key)
+    isnothing(node) && return t # no subtrie in t at key, do nothing
+    if at_key # only set is_key to false on the inner-most node
+        node.is_key = false
+    end
+    all_dangling = true # assume all children are dangling
+    for (k,v) in node.children # clean up dangling nodes
+        if isempty(v.children) && !v.is_key # dangling
+            delete!(node.children, k)
+        else
+            all_dangling = false # at least one not dangling
+        end
+    end
+    # only recurse if all children were dangling
+    return all_dangling ? delete!(t, @view(key[1:end-1]); at_key=false) : t
+end
+
 function subtrie(t::Trie, prefix)
     node = t
     for char in prefix

test/test_trie.jl

@@ -13,6 +13,22 @@
         @test sort(keys(t)) == ["amy", "ann", "emma", "kevin", "rob", "roger"]
         @test t["rob"] == 27
         @test sort(keys_with_prefix(t,"ro")) == ["rob", "roger"]
+
+        delete!(t, "roger")
+        @test !haskey(t, "roger")
+        @test haskey(t, "rob")
+        @test get(t,"rob",nothing) == 27
+        @test sort(keys(t)) == ["amy", "ann", "emma", "kevin", "rob"]
+        @test t["rob"] == 27
+        @test sort(keys_with_prefix(t,"ro")) == ["rob"]
+
+        # Ensure deletion cleans up dangling nodes
+        t = Trie(["A", "ABC"])
+        delete!(t, "ABC")
+        t_no_abc = Trie(["A"])
+        partial_no_abc = collect(partial_path(t_no_abc, "ABC"))
+        partial_original = collect(partial_path(t, "ABC"))
+        @test length(partial_no_abc) == length(partial_original)
     end
 
     @testset "Constructors" begin