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:
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.