1

Let's assume I have an a collection that holds either type A or type B then I think it's quite a natural thing to want a lazy collection of only the A values and one for only the B values.

For example:

struct ABCollection<A, B> {
    fileprivate let elements: [AB]

    enum AB {
        case a(A)
        case b(B)
    }
    
    typealias LazyACollection = LazyMapCollection<LazyFilterCollection<[AB]>, A>
    var onlyAs: LazyACollection {
        return self.elements.lazy.filter {
            switch $0 {
            case .a:
                return true
            case .b:
                return false
            }
        }.map { (input: AB) -> A in
            switch input {
            case .a(let value):
                return value
            case .b:
                fatalError()
            }
        }
    }
}

extension ABCollection: Collection {
    typealias Element = AB
    typealias Iterator = Array<AB>.Iterator

    var startIndex : Int {
        return self.elements.startIndex
    }
    
    var endIndex : Int {
        return self.elements.endIndex
    }
    
    func index(after i: Int) -> Int {
        return self.elements.index(after: i)
    }
    
    subscript(position : Int) -> AB {
        return self.elements[position]
    }
    
    func makeIterator() -> IndexingIterator<Array<AB>> {
        return self.elements.makeIterator()
    }
}

but the onlyAs is quite ugly for two reasons:

  1. the type: LazyMapCollection<LazyFilterCollection<[AB]>, A>
  2. the implementation contains a fatalError() because there can't be any .b(B) values left (as we filtered them out)

sure, with the new compactMap I could write this much neater:

    var onlyAsBetter: LazyCompactMapCollection<[AB], A> {
        return self.elements.lazy.compactMap {
            switch $0 {
            case .a(let a):
                return a
            case .b:
                return nil
            }
        }
    }

except, that LazyCompactMapCollection doesn't exist :disappointed_relieved:. I think this should be added. Incidentally LazyCompactMapCollection is a more general version of LazyFilterCollection and LazyMapCollection and could replace them both, we might however not want to do that for performance reasons.

Please let me know what you think.

Thanks,
Johannes

2

Could you please change the title of this thread or re-post with a different title. It should be more descriptive than “Introduce.”

3

sorry, totally forgot to change the title. Now done.

1 Like
4

I was thinking we should do this as well, though for a completely different reason:

Motivation

The current lazy system is very good for easily defining transforms on collections, but certain constructs can lead to less-than-optimal codegen.

For example, the code collection.map(map1).filter(filter1).map(map2).filter(filter2) will lead to code like this in the formIndex(after:) method:

do {
    do {
        collection.formIndex(after: &i)
    } while !filter1(map1(collection[i]))
} while !filter2(map2(map1(collection[i])))

while it could be represented with this more efficient single loop:

do {
    collection.formIndex(after: &i)
} while !filter1(map1(collection[i])) && !filter2(map2(map1(collection[i])))

Currently, you can get a single loop by instead using this compactMap:

collection.compactMap {
    let a = map1($0)
    guard filter1(a) else { return nil }
    let b = map2(a)
    guard filter2(b) else { return nil }
    return b
}

but this removes the nice composability of the chained map/filter combination.

The standard library recently got an override on LazyMapCollection and LazyFilterCollection which combines multiple filters and maps in a row, however it does not work with alternating maps and filters.

Solution

Define a LazyCompactMapCollection collection (and sequence) which represents a compactMap. Then, add overrides on LazyMapCollection.filter, LazyFilterCollection.map, Lazy*Collection.compactMap, and LazyCompactMapCollection.{filter, map} to return a LazyCompactMapCollection that combines all the maps and filters.

A demo implementation is available here, use .lazy2 to get the new implementation.
Note: Requires Swift 4.1 due to it being based off of the current stdlib code

As an added bonus, you'll never see a giant chain of LazyMapCollection<LazyFilterCollection<...>, ...> again

Note
As this would change the type returned by some methods, this would be a source-breaking change, though since I think .lazy is usually used as an intermediary step between a collection and some aggregator, most uses would probably be fine.

Other Note
The new implementation does perform slightly worse in Debug builds, however using .lazy in a debug build already performs much worse than an equivalent loop (compared to release where the new implementation performs almost as well as the loop), so I'm not sure if this really matters.

3 Likes
5

cool @TellowKrinkle, very happy to proceed with your proposal.

6

Sorry for taking so long on this, but I made a more official proposal and am working on getting my demo code into the stdlib. Does this seem good?

BTW, as we're the only two people who've commented on the proposal, do we need more people to give feedback or is this just because it's small and that's okay?

2 Likes
7

It seems small and uncontroversial from my viewpoint, so it makes sense that there wasn't much activity. The only questions in my mind are ones about compatibility and implementation that someone from the standard library team will probably weigh in on.

8

I performed some measurements (see the code at the bottom). The idea was to a) understand whether chaining lazy.filter.map.filter.map is as bad as the proposal suggests (no it's not that bad), and b) find out whether the implementation of filter and map in terms of the new LazyCompactMapSequence/LazyCompactMapCollection regresses performance or not, because, obviously, there should be a price to be paid for the convenience of a single type, as opposed to the long chain of LazyMapCollection<LazyFilterCollection<LazyMapCollection<...>>>.

Results

What happened next will surprise you!
Seriously, I was quite surprised myself. Here are the results (yet again, the code is a bit long, so it's presented below).

> $ time ../nestedfiltermap bespoke
true
../nestedfiltermap bespoke  0.07s user 0.01s system 94% cpu 0.080 total
                                                                                                                        
> $ time ../nestedfiltermap lazyFilterMap
true
../nestedfiltermap lazyFilterMap  0.26s user 0.01s system 98% cpu 0.273 total
                                                                                                                        
> $ time ../nestedfiltermap lazyFilterMap2
true
../nestedfiltermap lazyFilterMap2  0.25s user 0.01s system 98% cpu 0.262 total
                                                                                                                        
> $ time ../nestedfiltermap lazyCompactMap
true
../nestedfiltermap lazyCompactMap  7.61s user 0.01s system 99% cpu 7.630 total
                                                                                                                        
> $ time ../nestedfiltermap lazyCompactMap2
true
../nestedfiltermap lazyCompactMap2  0.18s user 0.01s system 97% cpu 0.196 total

Tests with 2 suffix use the proposed types.

Outcomes

  • Surprisingly the lazyCompactMap (that I implemented with an intention to establish a reasonable baseline) turned out to be the slowest contestant, by far.
  • Simple chaining (i.e. lazy.filter {}.map{}.filter{}) is not slow. Perhaps due to artificially simple closures, but I would argue that if you chain these calls in this manner, then the closures would be quite simple.
  • lazyFilterMap2 is actually faster than lazyFilterMap, which means that even if there is some extra overhead added by transforming boolean values into Optionals in order to be compatible with compactMap, this overhead is being effectively eliminated by the optimizer.

Some more results

For comparison, here are the results of using AnySequence instead of a simple array of Ints.

> $ time ../nestedfiltermap bespoke
true
../nestedfiltermap bespoke  6.01s user 0.02s system 99% cpu 6.070 total
                                                                                                                        
> $ time ../nestedfiltermap lazyFilterMap
true
../nestedfiltermap lazyFilterMap  6.05s user 0.01s system 99% cpu 6.069 total
                                                                                                                        
> $ time ../nestedfiltermap lazyFilterMap2
true
../nestedfiltermap lazyFilterMap2  6.13s user 0.01s system 99% cpu 6.147 total
                                                                                                                        
> $ time ../nestedfiltermap lazyCompactMap
true
../nestedfiltermap lazyCompactMap  7.42s user 0.01s system 99% cpu 7.439 total
                                                                                                                        
> $ time ../nestedfiltermap lazyCompactMap2
true
../nestedfiltermap lazyCompactMap2  5.96s user 0.01s system 99% cpu 5.971 total

The distribution of the results looks quite similar, just a larger overhead of using AnySequence that's applied to all the cases.

Code

#if false

public let first10k = AnySequence(Array(0...100_000-1))
typealias MainIterator = AnySequence<Int>.Iterator

#else

public let first10k = Array(0...100_000-1)
typealias MainIterator = Array<Int>.Iterator

#endif

@inline(never)
public func lazyFilterMap(_ N: Int) {
  var result = 0
  for _ in 1...10*N {
    let numbers = first10k.lazy
      .filter { $0 % 3 == 0 }
      .map { $0 * 2 }
      .filter { $0 % 8 == 0 }
      .map { $0 / 2 + 1}
    // print(type(of: numbers))
    result = numbers.reduce(into: 0) { $0 += $1 }
  }
  print(result == 416691666)
}

@inline(never)
public func lazyCompactMap(_ N: Int) {
  var result = 0
  for _ in 1...10*N {
    let numbers = first10k.lazy
      .compactMap { (x: Int) -> Int? in
        if x % 3 != 0 { return nil }
        let y = x * 2
        if y % 8 != 0 { return nil }
        let z = y / 2 + 1
        return z
      }
    // print(type(of: numbers))
    result = numbers.reduce(into: 0) { $0 += $1 }
  }
  print(result == 416691666)
}

struct Numbers: Sequence {
  func makeIterator() -> Iterator {
    return Iterator()
  }

  struct Iterator: IteratorProtocol {
    var it: MainIterator

    init() {
      self.it = first10k.makeIterator()
    }

    mutating func next() -> Int? {
      while true {
        if let x = it.next() {
          if x % 3 != 0 { continue }
          let y = x * 2
          if y % 8 != 0 { continue }
          let z = y / 2 + 1
          return z
        } else {
          return nil
        }
      }
    }
  }
}

@inline(never)
public func bespoke(_ N: Int) {
  var result = 0
  for _ in 1...10*N {
    let numbers = Numbers()
    // print(type(of: numbers))
    result = numbers.reduce(into: 0) { $0 += $1 }
  }
  print(result == 416691666)
}

@inline(never)
public func lazyFilterMap2(_ N: Int) {
  var result = 0
  for _ in 1...10*N {
    let numbers = first10k.lazy
      .filter2 { $0 % 3 == 0 }
      .map2 { $0 * 2 }
      .filter2 { $0 % 8 == 0 }
      .map2 { $0 / 2 + 1}
    // print(type(of: numbers))
    result = numbers.reduce(into: 0) { $0 += $1 }
  }
  print(result == 416691666)
}

@inline(never)
public func lazyCompactMap2(_ N: Int) {
  var result = 0
  for _ in 1...10*N {
    let numbers = first10k.lazy
      .compactMap2 { (x: Int) -> Int? in
        if x % 3 != 0 { return nil }
        let y = x * 2
        if y % 8 != 0 { return nil }
        let z = y / 2 + 1
        return z
      }
    // print(type(of: numbers))
    result = numbers.reduce(into: 0) { $0 += $1 }
  }
  print(result == 416691666)
}

if let which = CommandLine.arguments.dropFirst().first {
  switch which {
  case "lazyFilterMap": lazyFilterMap(1)
  case "lazyFilterMap2": lazyFilterMap2(1)
  case "lazyCompactMap": lazyCompactMap(1)
  case "lazyCompactMap2": lazyCompactMap2(1)
  case "bespoke": bespoke(1)
  default: print("No can do")
  }
}
2 Likes
9

In fact! What's even more important, because I did not use the proposed implementation and instead just copied various pieces from there, I ended up adding filter and map methods returning LazyConcatMapSequence directly to LazySequenceProtocol, which, given the performance numbers that I got, means that we can try and get rid of standalone LazyMapSequence and LazyFilterSequence, if that's possible without breaking source.

2 Likes
10

Very interesting findings. Thank you for this.

11

@TellowKrinkle would you be willing to try removing LazyMap and LazyFilter types in favor of the new type and update the proposal accordingly?

12

Dropping LazyMap would lose its O(1) count (and its nonzero underestimatedCount) which might not be preferable, but dropping LazyFilter would probably make sense. The only thing that LazyFilter does that LazyCompactMap doesn't is the check in _customContainsEquatableElement that checks whether or not the element matches the predicate before actually searching the collection. I don't know much that helps things but as long as you think losing that is fine that I'm happy to try the removal.

13

Also did you remember to compile with whole module optimization? All your tests run in 0.00 seconds (except lazyCompactMap due to some weird optimizer bug that's new in Swift 4.1 and I should report) when I test with wmo on.
Edit: In the test I used (the one I linked in my original post) I got a pretty big improvement

14

Good catch!

FWIW, I only added this optimization fairly recently for no other reason than "because I can" ([stdlib] Implement _customContainsEquatableElement for LazySequence a… · apple/swift@dd5ba51 · GitHub).

15

Let me try that.

16

@TellowKrinkle WMO did not change anything. :man_shrugging: Please try the AnySequence version of the code, maybe? It should be significantly slower.

17

Wait a second... Can't we just call the compactMap's transform on the element in question and if it returns nil then it essentially would mean that it's not found?

18

The element passed to contains(_:) would be the transformed element, not the original, so this wouldn't be what we need.

1 Like
19

:man_facepalming:

1 Like
20

I was going to try that, but then it didn't compile and I realized I'd need the inverse of the function, so that won't actually work.