Lazy Recursive Drill-Down of Tree Structures in Swift 2.1

I really do like the Swift programming language, because of features like multiple return types, nice support for optionals, value types (i.e., structs), enums with associated values, method overloading just with the return type, operator overloading, and many more.

But there are some language decision that are quite alien to me and I want to write them down (and my workarounds to them, which might help you), starting with this post. Perhaps you (yes I mean you 🙂 ) have a better answer for me, then please comment.

I stumbled over the obstacle I will mention in a few lines while I had a really precisely-defined task:

I develop the iOS game Fusionate (the site is work in progress). It is a SpriteKit based game and completely written in Swift. You can find it already in the App Store. I have written tests using XCTest that play the game from beginning to the end automatically and check whether the game behaves the same as it did, when I recorded the tests. This approach is called differential testing and it is quite powerful. Perhaps I will write about it in another blog post, what do you think?

Besides the behavioral check my tests inspect whether the visual state of the scene graph (i.e., the SKNodes of the SKScene) is the same in each turn (its a turn-based game) both playing the game and loading each turn separately from storage. For this I added in my test classes an extension to SKNode, which extracts the necessary data for comparing values like position, rotation, and color in a small struct NodeData. This additional information should only be collected for my own node types, since there are other nodes like the game menu, which will have different values each time the test is executed and would lead to failing tests.

In order to solve this problem we need at first the extension to SKNode:

extension SKNode {
  var nodeData: NodeData {
    return NodeData(
      name: name ?? self.dynamicType.description(),
      position: (x: Double(self.position.x), y: Double(self.position.y)),
      zPosition: Double(self.zPosition),
      zRotation: Double(self.zRotation),
      alpha: Double(self.alpha),
      scale: (xScale: Double(self.xScale), yScale: Double(self.yScale)),
      children: // the creation of `childNodeDatas` has to be implemented
    )
  }
}

I will omit the details of NodeData here as this is out of scope for this post. The following marker interface is used to identify my own nodes:

private protocol FusionateNode {
    var nodeData: NodeData { get }
}

With this in mind we can now do a first approach to separate the wheat from the chaff:

let childNodeDatas = self.children
  .filter() { $0 is FusionateNode}
  .map() { ($0 as! FusionateNode).nodeData }

This implementation retrieves all children deeply that implement my marker protocol FusionateNode. So will they live happily ever after? No! If we have a construct where a fusionate node is hidden by one or more hierarchy levels of SKNodes, they won’t show up in the target tree.

What can we do about it? We can build a recursive lambda expression using a little trick:

var inner: SKNode -> [NodeData] = { node in
  return []
}

let recurse: SKNode -> [NodeData] = { node in
  return Array(
    (node.children.map { node in
      if let node = node as? FusionateNode {
        return [node.nodeData]
      }
      return inner(node)
    } as [[NodeData]]).flatten()
  )
}
inner = recurse
let childNodeDatas = recurse(self)

We created a recursive call that drills-down each non-FusionateNode and searches deeply for FusionateNodes. It collects all fusionate nodes and puts them together (i.e., flattens them) to one big child array filled just with NodeDatas.

The first thing I do not like about this solution is, that I cannot use the lambda expression recurse() directly (that results in Variable used within its own initial value), but have to define an empty lambda expression inner (which cannot use fatalError() as implementation what is really annoying) which I have to reassign with the lambda expression recurse in a second step. The interface of the function is well-defined. Therefore, it should be possible to use it for recursion (the compiler may use the same trick that I do, but I don’t want to see it and be able to make mistakes/misuse it). We can omit this by creating a method in the SKNode extension:

private func recurse() -> [NodeData] {
  return Array((self.children.lazy.map { node in
    if let node = node as? FusionateNode {
      return [node.nodeData]
    }
    else {
      return node.recurse()
    }
  } as [[NodeData]]).flatten())
}

Second, I have to cast the result of map() to as [[NodeData]] or else I get an Cannot invoke 'map' with an argument list of type '(@noescape (SKNode) throws -> _)'. You could say now, that there is a recursive call in this expression and therefore the result of the call to map is hard to infer. But on the one hand this is not true, as the return type of inner does not have to be inferred as it is given directly, and on the other hand returning an empty array of NodeDatas (instead of calling recurse()) has the same effect. The problem seems to be introduced by the two return statements. The Swift compiler is not capable of following the control flow of one if-statement to infer the result type. If I remove the if-clause everything compiles well (even using the recursive call).

Third, this recursive call is not lazy. That means on every hierarchy level during the search for fusionate nodes a 2-dimensional array is created, flattened and collected into an array again. Hence, this solution is not very efficient. Adding a call to the property self.children.lazy will not help as every call to the constructor of Array collects the elements of the given lazy collection, eagerly.

Although performance is not so important for test code I was curious how to solve such an issue in Swift. Hence, I started the following odyssey, in order to fix the aforementioned third nuisance. Since I used a lot of program languages with static typing (like Java, Scala, Xtend, …) I first tried to use the most general (I intendedly do not use the term generic since it is a very overloaded notion) collection type as a return type of my recursive function in order to be able to collect all my intermediate collection views lazily until I do a terminal step of collecting all of them into my target collection (like in the Java 8 Stream API).

But right from the beginning an obstacle stood into my way: type aliases. Apple decided to support generics for classes and structs, but not for protocols (in Java you would say interface). Protocols can have associated types which are introduced via the keyword typealias. Unfortunately, as soon as you use a type alias in a protocol (it becomes an protocol with associated types (PAT)) you can use them as a generic constraint only. You can find more information on that in this talk from Alexis Gallager. I ended up with this “solution”:

private func recurse<R: CollectionType where R.Generator.Element == NodeData>() -> R {
  return (self.children.lazy.map { node in
    if let node = node as? FusionateNode {
      return [node.nodeData] as! R
    }
    else {
      return node.recurse()
    }
  } as LazyMapCollection<[SKNode], R>).flatten() as! R
}

I have to use the collection type as a generic constraint but, then I have to downcast every intermediate result and with calling this function the result type R is parameterized and sifts up through the complete call. No matter which (concrete) result type I choose, since all concrete collection types are structs which can inherit from protocols but not from each other or classes, it will be the wrong type, since [node.nodeData] and node.recurse() will never have the same type (if you proof me wrong, please write a comment; but still this is not coder-friendly isn’t it?). Hence, this solution compiles but will lead to nasty runtime errors for the forced downcasts.

Next, I tried to find a solution using the struct implementations of collection types of the Swift standard library. But since PATs can be used in generic constraints only, you cannot just have something like LazyMapCollection<SKNode, NodeData> which maps a collection with associated type SKNode to NodeData. This is because their general implementations use protocol abstractions (like CollectionType), which have to be used as type constraints, completely (and not just their element-type). This leads to types like LazyMapCollection<[SKNode], [NodeData]> (after calling flatten() this does not get less complex 😉 ). This works well for exactly no recursion as this would introduce a recursive type. Ok, so we start using a generic type for that again? No way, this didn’t work in the first place, right? Even worse trying to play around with all these possibilities gave me funny compiler errors. Not uncommon was expression too complex, which is really a pain.

So, since you read until here, you might ask whether I really have a solution for this problem as I “promised” in the beginning. Yes I have:

private func recurse() -> AnyForwardCollection<NodeData> {
  let mapping: LazyMapCollection<[SKNode], AnyForwardCollection<NodeData>> = self.children.lazy.map { node in
    if let node = node as? FusionateNode {
      return AnyForwardCollection([node.nodeData])
    }
    return node.recurse()
  }
  return AnyForwardCollection(mapping.flatten())
}

I used Apples wrapper collection types which erase the (garbage) generic constraints of the underlying collection. I use AnyForwardCollection because it is the only one that can take an Array as an argument (although AnyRandomAccessCollection would be much better from my point of view). If you call this method it will return a lazy collection, which has not been built, yet. Only if iterating over the collection or converting it into an array, the structure will be created.

Since Swift is OpenSource we now get a lot more information on how the language evolves. At least regarding the wrapper problem they seem to seek for a solution.

Exciting 😉