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
        }
    }
}
复制代码