Refactoring lazy functional code in Swift 5

Question

I'm wanting to refactor some lazy functional swift. To explain the situation I'll explain the equivalent eager situation first:

let numbers = 1...10

do {
  print("==================== EAGER INLINE =============================")
  /// We start with a series of transformation on an array:
  let result
    = numbers
      .filter { $0 >= 5 }      /// Drop numbers less than 5
      .map { $0 * $0 }         /// Square the numbers
      .flatMap { [$0, $0+1] }  /// Insert number+1 after each number
      .filter { $0 % 3 != 0 }  /// Drop multiples of 3
  print(result)
  /// [25, 26, 37, 49, 50, 64, 65, 82, 100, 101]
  /// which is [5^2, 5^2+1, 6^2+1, 7^2, 7^2+1, 8^2, 8^2+1, 9^2+1, 10^2, 10^2+1]
  /// (Note 6^2 and 9^2 missing because they are divisible by 3)
}

We can refactor the map and the flatMap into a seperate function:

extension Array where Element == Int {
  func squareAndInsert() -> [Int] {
    self
      .map { $0 * $0 }
      .flatMap { [$0, $0+1] }
  }
}

do {
  print("==================== EAGER REFACTOR =============================")

  let result
    = numbers
      .filter { $0 >= 5 }
      .squareAndInsert()
      .filter { $0 % 3 != 0 }
  print(result)
  /// Gives exactly the same result: [25, 26, 37, 49, 50, 64, 65, 82, 100, 101]
}

So now we'll repeat the process but lazily. First inline:

do {
  print("==================== LAZY INLINE =============================")
  let result: some LazySequenceProtocol /// ": some LazySequenceprotocol" not strictly
  /// required but without it my compiler grumbled about complexity so this is to give the
  /// compiler a nudge in the right direction.
    = numbers
      .lazy  /// Note the ".lazy" added here to make the array lazy.
      .filter { $0 >= 5 }
      .map { $0 * $0 }
      .flatMap { [$0, $0+1] }
      .filter { $0 % 3 != 0 }
  print(result)

}

Which prints: LazyFilterSequence<FlattenSequence<LazyMapSequence<LazyMapSequence<LazyFilterSequence<ClosedRange<Int>>, Int>, Array<Int>>>>(_base: Swift.FlattenSequence<Swift.LazyMapSequence<Swift.LazyMapSequence<Swift.LazyFilterSequence<Swift.ClosedRange<Swift.Int>>, Swift.Int>, Swift.Array<Swift.Int>>>(_base: Swift.LazyMapSequence<Swift.LazyMapSequence<Swift.LazyFilterSequence<Swift.ClosedRange<Swift.Int>>, Swift.Int>, Swift.Array<Swift.Int>>(_base: Swift.LazyMapSequence<Swift.LazyFilterSequence<Swift.ClosedRange<Swift.Int>>, Swift.Int>(_base: Swift.LazyFilterSequence<Swift.ClosedRange<Swift.Int>>(_base: ClosedRange(1...10), _predicate: (Function)), _transform: (Function)), _transform: (Function))), _predicate: (Function))

Yikes!

Looks rather alarming at first sight but this is correct because unlike the eager result which is an array of Ints, the lazy result is an iterator which will provide us with the next number when we ask it to and this needs to know how to work back through all the function calls right back to the initial sequence. That's what this type is describing. Very nice now that we have the "some" keyword as in the past, if we wanted to put in an explicit type we would have to type all the above which is a bit of a mouthful !!

To see the list of numbers we need to force them to be calculated which we can do by putting the lazy sequence into an array: print(Array(result))

And this gives exactly the same result as before: [25, 26, 37, 49, 50, 64, 65, 82, 100, 101]

So now the challenge.

I want to refactor the lazy code in the same way that I did the eager code.

squareAndInsert needs to turn a LazySequenceProtocol<Int> into some LazySequenceProtocol so I try the code below but get various compile errors:

extension LazySequenceProtocol where Element == Int {
  func squareAndInsertLazy() -> some LazySequenceProtocol {
    self
      .map { $0 * $0 }
      .flatMap { [$0, $0+1] }
  }
}

do {
  print("==================== LAZY REFACTOR =============================")

  let result: some LazySequenceProtocol  // Error 1: Property declares an opaque return type, but cannot infer the underlying type from its initializer expression
    = numbers
      .lazy
      .filter { $0 >= 5 }
      .squareAndInsertLazy()  // Error 2: Value of type '[Int]' has no member 'squareAndInsertLazy'
      .filter { $0 % 3 != 0 } // Error 3: Protocol type 'Any' cannot conform to 'LazySequenceProtocol' because only concrete types can conform to protocols
                              // Error 4: Value of type 'Any' has no member 'filter'

  print(result)
}

I think Error 1 would probably go away if I fix the others. I wonder if Error 2 means that trying to pass the lazy sequence into squareAndInsertLazy forces eagerness and this means [Int] is being presented to squareAndInsertLazy. I can't work out how to move forward.

Any help appreciated.

Answer 1

The issue here is that LazySequenceProtocol is a PAT (protocol with associatedtype). So when you call squareAndInsertLazy() it returns some LazySequenceProtocol and it has no idea what the elements are anymore.

You can see this is the issue by commenting out your .filter { $0 % 3 != 0 } and replacing it with .filter { _ in true }. It will be perfectly happy and not complain because it doesn't care what the type of elements is in the sequence.

You can also see this using:

.filter { value in
    let copy = value
    return true
}

If you then Option click on copy it will show you the type is: (some LazySequenceProtocol).Element which cannot be used directly and must be inferred by the compiler. You can't do let copy: (some LazySequenceProtool).Element = value it won't compile.

So now that we have figured out what the problem is what are your possible solutions?

1) Don't return some PAT in this case some LazySequenceProtocol and return the concrete type which would be LazySequence<FlattenSequence<LazyMapSequence<LazyMapSequence<Self.Elements, Int>, [Int]>>>.

2) Go Back to inline.

3) Create a protocol that implements LazySequenceProtocol and refines Element to Int like this:

protocol LazySequenceOfInt: LazySequenceProtocol where Element == Int {}
extension LazySequence: LazySequenceOfInt where Element == Int {}

You will then use some LazySequenceOfInt. If you do this then you will potentially also want to extend the other Lazy types to conform to LazySequenceOfInt so they can also be used. In this particular case LazySequence is the only one you need though.