Trie data structure in Swift lang

Question

This is an implementation of a generic Trie DS, in Swift-lang, that is using however, for the sake of a easy example, strings. It implements two operations, insert and search. So it can insert a text, and look for substrings of that text, for example. I'd be interested in recommendations as to how correct the code is, but also how idiomatic, relative to the Swift language.

The Node<T>struct represents a node in the Trie<T> tree. Each node holds a dictionary that maps a key of type T, towards another Node<T>.

This is the code:

import Foundation

struct Node<T: Hashable> {
    var children: [T : Node<T>] = [:]
    
    var isLeaf: Bool {
        children.isEmpty
    }
    
    mutating func addChild(withKey key: T) {
        guard children[key] == nil else {  return   }
        children[key] = Node()
    }
    
    func hasChild(withKey key: T) -> Bool {
        children[key] != nil
    }
}

struct Trie<T: Hashable> {
    public var root = Node<T>()
    
    mutating func addContent(of content: some Collection<T>) {
       var iter = content.makeIterator()
       addCotentHelper(&root, &iter)
        
        func addCotentHelper(_ node: inout Node<T>, _ iter: inout some IteratorProtocol<T>) {
            guard let key = iter.next() else { return }
            node.addChild(withKey: key)
            addCotentHelper(&node.children[key]!, &iter)
        }
    }
    
   

    private func containsPrefix(of content: some Collection<T>, at root: Node<T>) -> Bool {
        guard let first = content.first else { return true }
        
        if root.hasChild(withKey: first) {
            return containsPrefix(of: content.dropFirst(), at: root.children[first]!)
        }
        return false
    }
    
    public func hasContent(of word: some Collection<T>) -> Bool {
        return hasContentHelper(of: word, at: root)
        
        func hasContentHelper(of word: some Collection<T>, at root: Node<T>) -> Bool {
            guard !word.isEmpty else { return true }
            
            guard root.isLeaf == false else {
                return false
            }
            
            if containsPrefix(of: word, at: root) {
                return  true
            }
            
            return root.children.values.contains(where: { child in hasContentHelper(of: word, at: child) })
        }
    }
    
    func printTrie() {
        printTrieHelper(node: root, indent: 1)
        
        func printTrieHelper(node: Node<T>, indent: Int) {
            let leading_indent = "| "
            let last_indent = "|-"
            for (k, v) in node.children {
                print(
                    String(
                        repeating: leading_indent,
                        count: indent - 1) + last_indent,
                    terminator: "")
                print(k, terminator: "\n")
                printTrieHelper(node: v, indent: indent + 1)
            }
        }
    }
}

Answer 1

Access control

Some types and methods are explicitly made public, but in order to make the code compilable and usable as an external library, all types and methods which are meant to be called from outside must be made public. Also Trie needs a public init method. So the public interface should look like this:

public struct Node<T> where T : Hashable {
}

public struct Trie<T> where T : Hashable {

    public init()
    public mutating func addContent(of content: some Collection<T>)
    public func hasContent(of word: some Collection<T>) -> Bool
    public func printTrie()
}

Another option is to make Node private to the Trie type:

public struct Trie<T: Hashable> {

    private struct Node<T: Hashable> { ... }

    private var root = Node<T>()

    // ...
}

Minor remarks

The generic type placeholder can be omitted if it is given from the context, i.e. in struct Node<T: Hashable> can

var children: [T : Node<T>] = [:]
    

be simplified to

var children: [T : Node] = [:]
    

There is a typo in addCotentHelper().

To guard or not to guard

The guard statement is useful to avoid the “if-let pyramid of doom” and to handle exceptional situations. But it should (in my opinion) not be used as a general “if not” statement. So it is perfectly fine here

guard let key = iter.next() else { return }

But these statements

guard children[key] == nil else {  return   }
children[key] = Node()

guard !word.isEmpty else { return true }
        
guard root.isLeaf == false else {
    return false
}

are easier to read and to understand as simple if statements:

if children[key] == nil {
    children[key] = Node()
}

if word.isEmpty {
    return true
}

if root.isLeaf {
    return false
}

Simplifying the code

The addContent() method uses recursion and passes Node instances as inout parameters around a lot. This gets much easier if we make Node a class (so that pointers to instances can be passed around) and let the addChild() return the child node:

class Node<T: Hashable> {
    var children: [T : Node] = [:]
    
    func addChild(withKey key: T) -> Node {
        if let child = children[key] {
            return child
        } else {
            let child = Node()
            children[key] = child
            return child
        }
    }
    
    // ...
}

public mutating func addContent(of content: some Collection<T>) {
    var node = root
    for key in content {
        node = node.addChild(withKey: key)
    }
}

Similarly, containsPrefix() can now use iteration of recursion, and creating slices of the content is no longer necessary:

private func containsPrefix(of content: some Collection<T>, at root: Node<T>) -> Bool {
    var node = root
    for key in content {
        guard let next = node.children[key] else {
            return false
        }
        node = next
    }
    return true
}

In hasContentHelper(), the check for word.isEmpty is not needed because containsPrefix(of: word, at: root) will return true in that case. The check for root.isLeaf is also not needed, because root.children.values.contains will return false if there are not children.

It is a matter of taste, but I find an explicit loop easier to read than calling root.children.values.contains() with a closure:

for child in root.children.values {
    if hasContentHelper(of: word, at: child) {
        return true
    }
}
return false

With the above changes the isLeaf property and the hasChild() method of Node are no no longer needed.

Naming

Again a matter of personal taste, but I would use

private func containsPrefix(_ content: some Collection<T>, at root: Node<T>)
public func contains(_ word: some Collection<T>) -> Bool
func containsHelper(_ word: some Collection<T>, at root: Node<T>) -> Bool
public func print()

The first three method resemble the contains() method of collections. And “Trie” in printTrie() is a “needless word” because it just repeats the type, which is clear from the context.

As an alternative to a print() method one can also implement the CustomStringConvertible or CustomDebugStringConvertible protocol, e.g.

extension Trie: CustomDebugStringConvertible {
    public var debugDescription: String {
        // add code here ...
    }
}