Without actually escaping

When dealing with lazy sequences and generic functions, there are subtle pitfalls you need to be aware of. You could easily negate the laziness and end up drastically slowing down your algorithms.

The other day someone asked in the Swift forums for help on creating a function that return the first transformed non-nil element of an array. The person offered this implementation:

extension Array {
  func firstResult<T>(_ predicate: (Element) -> T?) -> T? {
    for element in self {
      if let result = predicate(element) {
        return result
      }
    }
    return nil
  }
}

This does work as intended. It evaluates each element in turn, and if the transform returns a non-nil result, the value is returned and the loop is terminated early without evaluating more elements.

Then someone suggested it could be implemented using a chain of standard library functions instead:

return self.lazy.compactMap(predicate).first

While this also works, it isn’t as efficient. It does transforms the elements of the array, and it removes the null values before returning the first element. However, it transforms each and every element of the array before returning, even if the very first element passes the test!

— But isn’t the .lazy meant to counter exactly that, I hear you say?
Yes, it is. But in this case, it isn’t.

Why?

Escaping and non-escaping closures.

When creating a closure, it captures it surrounding state needed to run the code within the closure. If the closure is passed on as an argument to a function, and this function stores the closure for later evaluation, it must be marked as @escaping, since the state needs to be stored on the heap. For closures that are executed immediately, this isn’t needed, and the compiler can make optimizations when handling the captured state.

For arrays and collections, the compactMap expects a non-escaping closure, as the collection/sequence is evaluated inline and each transformation is performed before the function returns.

However, for their lazy counterparts, the compact map expects an escaping closure, as the actual mapping is deferred to later, and only lazily evaluated when needed. After all, that is the whole idea behind being lazy!

The many compact maps

The lazy sequence returned from self.lazy is either of type LazyCollection or LazySequence, depending on which one the compiler infers. Each one of these have their own implementation of compactMap, and they also inherit a different compact map each from Collection and Sequence respectively. When you just write self.lazy.compactMap(...) the compiler has to infer which implementation you really want.

In the case at hand, the function is declared as func firstResult<T>(_ predicate: (Element) -> T?) -> T?. Note that the parameter predicate in this case is not declared as @escaping.

That means that the argument passed to self.lazy.compactMap(predicate), restricts the compiler’s choices when it tries to infer which overload to use. Since the parameter is non-escaping, it has to fall back to the one of the non-escaping versions declared in the non-lazy Collection/Sequence.

Not actually escaping

In our case we wanted to use .lazy so we didn’t have to first evaluate every element of the input array, only to just pick the first result and throw away the rest. We could solve this by changing our method signature:

extension Array {
  func firstResult<T>(_ predicate: @escaping (Element) -> T?) -> T?
}

However, our function isn’t really escaping is it? Since we terminate the sequence with .first, we will actually walk the sequence, one-by-one, performing the transform on each element as we go, and stop when we find the first passing result! This means our function will block execution and return without actually escaping the scope.

As luck would have it, there is a method named exactly like this. And it does exactly what we want: Allows a non-escaping closure to temporarily be used as if it were allowed to escape. As the documentation says:

You can use this function to call an API that takes an escaping closure in a way that doesn’t allow the closure to escape in practice.

So here it is, without actually escaping:

extension Sequence {
  func firstResult<T>(_ predicate: (Element) -> T?) -> T? {
    return withoutActuallyEscaping(predicate) { escapablePredicate in
      return self.lazy.compactMap(escapablePredicate).first
    }
  }
}

These kinds of bugs can be hard to spot. Since the code is valid without the escaping/non-escaping dance, the compiler doesn’t complain, and the result is correct. For small array, you may not even notice the slowdown. But without actually noticing, you may have added a useless .lazy.