Optional 中的 map 和 flatMap

在面试的时候，面试官问了我一个问题。你使用过 Optional 吗？Optional 的原理是什么？我一听，这还不简单。Optional是一个枚举。

public enum Optional<Wrapped> : ExpressibleByNilLiteral {
    case none
    case some(Wrapped)
}

然后面试官又问，那你使用过Optional的map和flatMap吗。我一般是在Sequence里使用map和flatMap的。在Optional里还真没用过。所以在面试之后，赶紧来补习一波。

首先，我试用了一下的代码来测试

let h: Int? = 10
let h1 = h.map { $0 + 2}
let h2 = h.flatMap { $0 + 2}
print(h1) // Optional(10)
print(h2) // Optional(10)

咦，好像没区别？好吧，换种测试方法，在闭包里再转成一个Optional看一下

let g: String? = "10"
let g1 = g.map {Int($0)}
let g2 = g.flatMap {Int($0)}

print(g1) // Optional(Optional(10))
print(g2) // Optional(10)

嗯，这次的返回的结果就不一样了。flatMap对结果做了一次展开。那它是怎么做到的呢？源代码在github的swift/stdlib/public/core/Optional.swift文件里。来看一下这2个方法的源代码有什么区别吧。

  public func map<U>(
    _ transform: (Wrapped) throws -> U
  ) rethrows -> U? {
    switch self {
    case .some(let y):
      return .some(try transform(y))
    case .none:
      return .none
    }
  }

  public func flatMap<U>(
    _ transform: (Wrapped) throws -> U?
  ) rethrows -> U? {
    switch self {
    case .some(let y):
      return try transform(y)
    case .none:
      return .none
    }

一眼看上去，2个方法的实现只有2处不一致。一处是方法签名

func map<U>(_ transform: (Wrapped) throws -> U) rethrows -> U?      
func flatMap<U>(_ transform: (Wrapped) throws -> U?) rethrows -> U?

另一处是当有值的时候的返回值

return .some(try transform(y)) // map
return try transform(y) // flatMap

但我们一般申明一个Optional的返回值时，如果是一个确定的值，也可以直接返回。编译器会帮你包装成.some()的形式。所以这里map方法里的代码也可以写成和下面的flatMap方法一致。

那么这样看来，就只有方法签名不一致了，那它到底是怎么实现的呢？这里我们就以上面的例子，把泛型转成具体的类型来看一下吧。

在上述方法里，map方法transform闭包的返回值是Int($0)，是一个Int?的类型，那么U就是一个Int?,U?就是Int??。所以转成具体类型之后，方法签名变成了

func map<U>(_ transform: (Wrapped) throws -> U) rethrows -> U?      
func myMap(_ transform: (Wrapped) throws -> Int?) rethrows -> Int??

而flatMap方法tranform的返回值签名是U?，对应Int?，那么这里U就是Int，所以flatMap方法的方法签名就变成了

func flatMap<U>(_ transform: (Wrapped) throws -> U?) rethrows -> U?
func myFlatMap(_ transform: (Wrapped) throws -> Int?) rethrows -> Int?

这样一看是不是就一目了然了。只能说swift的泛型机制真是太强大太巧妙了。我修改签名的源代码放在下面，有兴趣的可以跑起来看看

extension Optional {
    func myMap(_ transform: (Wrapped) throws -> Int?) rethrows -> Int?? {
        switch self {
        case .some(let y):
            return try transform(y)
        case .none:
            return .none
        }
    }
    
    func myFlatMap(_ transform: (Wrapped) throws -> Int?) rethrows -> Int? {
        switch self {
        case .some(let y):
            return try transform(y)
        case .none:
            return .none
        }
    }
}