[Pitch] Add easy <Sequence> conformance for arbitrary homogeneous trees

Overview:

If you’re using a tree data structure in Swift, there’s no easy way to visit all the nodes of the tree.

For instance, if you use SceneKit, each SCNNode has:

var childNodes: [SCNNode]

If you want to, say, set a constant pieceIndex on each node in the tree, you’d probably write something recursive, like:

let pieceIndex = 2
func updatePieceIndexForGeometryUnderNode(_ node: SCNNode) {
    node.geometry?.setValue(pieceIndex, forKey: "pieceIndex")
    node.childNodes.forEach { updatePieceIndexForGeometryUnderNode($0) }
}
updatePieceIndexForGeometryUnderNode(node)

If, however, SCNNode conformed to Sequence, you could write this much more simply:

let pieceIndex = 2
node.forEach { $0.geometry?.setValue(pieceIndex, forKey: "pieceIndex”) }

Implementation:

I propose adding a simple protocol (currently called RegularTree) that requires just one extra property, children, and turns any tree into a Sequence.

Here’s the code I have right now:

public protocol RegularTree : Sequence {
    associatedtype ChildrenCollection where ChildrenCollection : Collection, ChildrenCollection.Element == Self
    var children: ChildrenCollection { get }
}

public extension RegularTree {
    var isLeaf: Bool { children.isEmpty }

    func makeIterator() -> RegularTreeIterator<Self> {
        RegularTreeIterator<Self>(self)
    }
}

public struct RegularTreeIterator<TreeType> : Sequence, IteratorProtocol where TreeType : RegularTree {
    // MARK: init
    public init(_ root: TreeType) {
        self.root = root
    }

    // MARK: <IteratorProtocol>
    public mutating func next() -> TreeType? {
        guard rootWasUsed else {
            rootWasUsed = true
            iteratorStack.append(root.children.makeIterator())
            return root
        }

        guard var lastIterator = iteratorStack.last else { return nil }
        iteratorStack.removeLast()

        if let nextItem = lastIterator.next() {
            iteratorStack.append(lastIterator)
            if !nextItem.isLeaf {
                iteratorStack.append(nextItem.children.makeIterator())
            }
            return nextItem
        } else {
            return next() // go up a level, if there are any levels above us
        }
    }

    // MARK: private properties
    private let root: TreeType
    private var rootWasUsed = false
    private var iteratorStack: [TreeType.ChildrenCollection.Iterator] = []
}

Here’s how it’d be applied to SceneKit’s SCNNode:

extension SCNNode : RegularTree {
    var children: [SCNNode] { childNodes }
}

And that’s all that’d be needed to be able to use an SCNNode as a sequence.

Flaws and Possible Improvements:

• It’d be good if there were both depth-first and breadth-first iterators.
• The name RegularTree is...not good.
• The implementation above is probably not the fastest.
• I haven’t verified that this implementation won’t blow up the stack if traversing a very very large tree.

2 Likes

Would that suggest instead, rather, that it's not the tree that should conform to Sequence, but instead a depth-first traversal view and a breadth-first traversal view of that tree?

11 Likes

I’m not an expert in views but it seems like an Iterator already is a view onto something, so having another view would be kinda extra.

Basically, if you don’t care what order you see the children in, you can just use the standard sequence methods on the tree (eg, forEach, map, etc), but if you do care you can just instantiate the appropriate iterator yourself and use that as the view onto the tree.

A tree shouldn’t itself conform to Sequence because there are multiple ways to flatten a tree into a sequence (depth-first, breadth-first, etc.). You wouldn’t be able to iterate through a tree in different orders if the type itself conformed to Sequence. Instead, an intermediary type can be used as a lens/view onto that tree. I recently built something that does this. I started to think about how it might fit into the swift-algorithms package. However, it's still too early to put in tree algorithms into the Swift standard library or swift-algorithms because there isn’t yet enough direction for how tree/graph structures would be represented in Swift.

4 Likes

My point is that if you have multiple Iterator types, you can always manually create an iterator, so you don’t need a lens/view: the iterator is that already.

To wit, nothing keeps you from instantiating a custom Iterator type on an existing Sequence (of any type) and using that iterator instead of the default one.

But also I think it’s very handy to be able to treat a tree, by default, as a sequence, because it’s very common to want to get at all its nodes.

I don’t know about “too early,” I’m just saying this is something I’ve found eliminates code (and in one case eliminated a crasher, in a place where I was recursing into a tree 300 frames deep).

-Wil