[Swift] map, flatMap, compactMap 정복하기!

🤔 글을 시작하며

Swift를 쓴다면 map, flatMap, compactMap정도는 알 것이다. 이 글에서는 이 세 개의 차이를 말하기보단 세부 구현에 대해서 알아보고자 한다. 차이만 알고 싶다면... 아래의 예시만 참고하면 될 것 같다.

 

map은 배열을 1:1로 다른 타입으로 매핑한다.

let numbers = [1, 2, 3]
let strings = numbers.map { "\($0)" }
// ["1", "2", "3"]

compactMap은 배열을 1:1로 다른 타입으로 매핑하면서 nil값은 없애준다.

let strings = ["1", "2", "삼", "4"]

let numbers = strings.compactMap { Int($0) }
// [1, 2, 4]

flatMap은 배열을 펼쳐준다.

let numbers = [[1, 2, 3], [4, 5], [6]]

let flat = numbers.flatMap { $0 }
// [1, 2, 3, 4, 5, 6]

특히 나는 compactMap을 애용한 것 같다. 나는 여기서 세부 구현이 궁금해졌다. 아래부터는 세부구현을 살펴볼 것이다. Swift는 오픈소스이기 때문에 세부구현을 확인할 수 있다.

1️⃣ map

정확한 구현은 swift/stdlib/public/core/Sequence.swift 에서 찾을 수 있었다.

swift/stdlib/public/core/Sequence.swift at 376651df4db8d7633874892b685115e2168a7c41 · swiftlang/swift

그냥 주석도 일단 모조리 긁어왔다. 당당히 extension Sequence의 첫번째 자리를 꿰차고 있는 녀석.

//===----------------------------------------------------------------------===//
// Default implementations for Sequence
//===----------------------------------------------------------------------===//

extension Sequence {
  /// Returns an array containing the results of mapping the given closure
  /// over the sequence's elements.
  ///
  /// In this example, `map` is used first to convert the names in the array
  /// to lowercase strings and then to count their characters.
  ///
  ///     let cast = ["Vivien", "Marlon", "Kim", "Karl"]
  ///     let lowercaseNames = cast.map { $0.lowercased() }
  ///     // 'lowercaseNames' == ["vivien", "marlon", "kim", "karl"]
  ///     let letterCounts = cast.map { $0.count }
  ///     // 'letterCounts' == [6, 6, 3, 4]
  ///
  /// - Parameter transform: A mapping closure. `transform` accepts an
  ///   element of this sequence as its parameter and returns a transformed
  ///   value of the same or of a different type.
  /// - Returns: An array containing the transformed elements of this
  ///   sequence.
  ///
  /// - Complexity: O(*n*), where *n* is the length of the sequence.
  @inlinable
  @_alwaysEmitIntoClient
  public func map<T, E>(
    _ transform: (Element) throws(E) -> T
  ) throws(E) -> [T] {
    let initialCapacity = underestimatedCount
    var result = ContiguousArray<T>()
    result.reserveCapacity(initialCapacity)

    var iterator = self.makeIterator()

    // Add elements up to the initial capacity without checking for regrowth.
    for _ in 0..<initialCapacity {
      result.append(try transform(iterator.next()!))
    }
    // Add remaining elements, if any.
    while let element = iterator.next() {
      result.append(try transform(element))
    }
    return Array(result)
  }
  
  ...

보니까 살짝 막막해졌다. 음... 어디서부터 시작해야 하지? 우선 underestimatedCount를 통해 배열의 크기를 정해준다. 아래는 어떤식으로 count가 되고, underestimatedCount가 결정되는지 나와았다.

  @inlinable
  public var underestimatedCount: Int {
    // TODO: swift-3-indexing-model - review the following
    return count
  }
  
  @inlinable
  public var count: Int {
    return distance(from: startIndex, to: endIndex)
  }

그 뒤로 makeIterator()로 현재 시퀀스를 순회할 수 있는 Iterator를 만들어주고 위에서 구한 underestimatedCount만큼 result에 넣어준다. 여기서 initialCapacity는 주석과 같이 대략적인 원소 개수라고 한다. 이렇게 initialCapacity만큼은 무조건 있다고 생각하기 때문에 강제 언래핑을 하면서 loop를 돈다.

    var iterator = self.makeIterator()

    // Add elements up to the initial capacity without checking for regrowth.
    for _ in 0..<initialCapacity {
      result.append(try transform(iterator.next()!))
    }

위에서 보여준 underestimatedCount의 경우 Collection.swift파일에 있는 구현내용이고, 시퀀스는 아래와 같이 기본값이 0을 갖는다.

  /// A value less than or equal to the number of elements in the sequence,
  /// calculated nondestructively.
  ///
  /// The default implementation returns 0. If you provide your own
  /// implementation, make sure to compute the value nondestructively.
  ///
  /// - Complexity: O(1), except if the sequence also conforms to `Collection`.
  ///   In this case, see the documentation of `Collection.underestimatedCount`.
  @inlinable
  public var underestimatedCount: Int {
    return 0
  }

 그리고 위에서 정확히 측정할 수 있는 양은 빠르게 loop를 돌고, 아래와 같이 나머지는 다음 구현을 넣어주는 식으로 작동한다.

    // Add remaining elements, if any.
    while let element = iterator.next() {
      result.append(try transform(element))
    }
    return Array(result)

결국 아래와 같이 Element를 T타입으로 바꾸어서 반환하는 함수가 내부적으로 구현되어 있는 것을 알 수 있다.

  public func map<T, E>(
    _ transform: (Element) throws(E) -> T
  ) throws(E) -> [T]

2️⃣ flatMap

swift/stdlib/public/core/Sequence.swift, swift/stdlib/public/core/FlatMap.swift 에서 확인할 수 있었다.

swift/stdlib/public/core/FlatMap.swift at 376651df4db8d7633874892b685115e2168a7c41 · swiftlang/swift

두 가지로 구현되어 있다. 바로 Sequence와 LazySequence. 앞의 구현은 단순하게 result를 만들어주고 loop를 통해 넣어주는 것이다. 그 뒤는 뭔가 <>가 많지만 내부 구현을 보면 아주아주 단순하다. map을 해준 뒤 join()을 해준 것이다. 

  // MARK: - Sequence
  
  @inlinable // protocol-only
  @inline(__always)
  public func _compactMap<ElementOfResult>(
    _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    var result: [ElementOfResult] = []
    for element in self {
      if let newElement = try transform(element) {
        result.append(newElement)
      }
    }
    return result
  }
  
  // MARK: - LazySequenceProtocol
  
  /// Returns the concatenated results of mapping the given transformation over
  /// this sequence.
  ///
  /// Use this method to receive a single-level sequence when your
  /// transformation produces a sequence or collection for each element.
  /// Calling `flatMap(_:)` on a sequence `s` is equivalent to calling
  /// `s.map(transform).joined()`.
  ///
  /// - Complexity: O(1)
  @inlinable // lazy-performance
  public func flatMap<SegmentOfResult>(
    _ transform: @escaping (Elements.Element) -> SegmentOfResult
  ) -> LazySequence<
    FlattenSequence<LazyMapSequence<Elements, SegmentOfResult>>> {
    return self.map(transform).joined()
  }

그러면 flatMap대신에 난 map과 join()을 쓰겠어!라고 생각할 수도 있지만. 실제로 해보면 아래와 같이 나오게 된다. 

let numbers = [[1, 2, 3], [4, 5], [6]]

let flat = numbers.map { $0 }.joined()
// FlattenSequence<Array<Array<Int>>>(_base: [[1, 2, 3], [4, 5], [6]])

LazySequence가 나오기 때문인데 아래와 같이 실제 prefix(2)처럼 결과를 사용하게 된다면 문제없이 나온다.

let numbers = [[1, 2, 3], [4, 5], [6]]

let flat = numbers.flatMap { $0 }.prefix(2)
// [1, 2]

 

참고로 이전 구현은 아래와 같았다. 구현을 보면 flatMap이 compact의 역할도 해줬다. 지금도 사용 가능하지만 warning을 보게 될 것이다.

  @inlinable
  public func flatMap<SegmentOfResult: Sequence>(
    _ transform: (Element) throws -> SegmentOfResult
  ) rethrows -> [SegmentOfResult.Element] {
    var result: [SegmentOfResult.Element] = []
    for element in self {
      result.append(contentsOf: try transform(element))
    }
    return result
  }
  
  @available(swift, deprecated: 4.1, obsoleted: 5.0, renamed: "compactMap(_:)",
    message: "Please use compactMap(_:) for the case where closure returns an optional value")
  public func flatMap(
    _ transform: (Element) throws -> String?
  ) rethrows -> [String] {
    return try _compactMap(transform)
  }
  
  // The implementation of compactMap accepting a closure with an optional result.
  // Factored out into a separate function in order to be used in multiple
  // overloads.
  @inlinable // protocol-only
  @inline(__always)
  public func _compactMap<ElementOfResult>(
    _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    var result: [ElementOfResult] = []
    for element in self {
      if let newElement = try transform(element) {
        result.append(newElement)
      }
    }
    return result
  }

3️⃣ compactMap

이것은 타입별로 조금씩 다른데 FlatMap.swift, SequenceAlgorithms.swift 파일에서 각각 있다. 순서대로 flatMap에서 사용하던 compactMap로직, LazySequence에서의 로직이 있다.

  @inlinable // protocol-only
  public func compactMap<ElementOfResult>(
    _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    return try _compactMap(transform)
  }

  // The implementation of compactMap accepting a closure with an optional result.
  // Factored out into a separate function in order to be used in multiple
  // overloads.
  @inlinable // protocol-only
  @inline(__always)
  public func _compactMap<ElementOfResult>(
    _ transform: (Element) throws -> ElementOfResult?
  ) rethrows -> [ElementOfResult] {
    var result: [ElementOfResult] = []
    for element in self {
      if let newElement = try transform(element) {
        result.append(newElement)
      }
    }
    return result
  }
  
  @inlinable // lazy-performance
  public func compactMap<ElementOfResult>(
    _ transform: @escaping (Elements.Element) -> ElementOfResult?
  ) -> LazyMapSequence<
    LazyFilterSequence<
      LazyMapSequence<Elements, ElementOfResult?>>,
    ElementOfResult
  > {
    return self.map(transform).filter { $0 != nil }.map { $0! }
  }

이것도 차근차근 보면 그리 어렵진 않은 것이 if let으로 nil이 아닌 경우에 append 해주거나 map, filter, map을 활용하여 옵셔널을 벗기는 로직이다. 하지만 이것은 직접 구현해 보면 joined()가 없어서 즉시 Int 배열이 정상적으로 나오는 모습이 보인다.

let numbers = ["1", "2", "3", "사"]

let compact = numbers.map { Int($0) }.filter { $0 != nil }.map { $0! }
// [1, 2, 3]

⏱️ lazySequence 간단 성능 테스트

아래와 같이 0..<1000까지 map, filter, map을 해준다면 두 개는 성능차이가 얼마나 날까?

// MARK: Eager

let start = Date()

let numbers = (0..<10_000_000)
    .map { $0 * 2 }
    .filter { $0 % 3 == 0 }
    .map { "\($0)" }
    .prefix(1)

let end = Date()
print("eager time: \(end.timeIntervalSince(start)) sec")

// MARK: Lazy

let lazyStart = Date()

let lazyNumbers = (0..<10_000_000).lazy
    .map { $0 * 2 }
    .filter { $0 % 3 == 0 }
    .map { "\($0)" }
    .prefix(1)

_ = Array(lazyNumbers) // 평가 강제!

let lazyEnd = Date()
print("lazy time: \(lazyEnd.timeIntervalSince(lazyStart)) sec")

결과는 아래와 같이 1000만 개의 배열에서 약 2만 배 성능이 차이가 난다. 1000만이 아니고 숫자가 더 커질수록 10억이 되어도 아래의 lazy는 속도가 똑같다.

eager time: 2.5064669847488403 sec
lazy time: 0.0001270771026611328 sec

이처럼 LazySequence는 사용하는 것만 사용하여 불필요한 계산을 줄여주기 때문에 성능이 훨씬 좋다!