Alvis

Defunctionalize the Continuation

发表于更新于阅读时长 3 分钟

如何把任意递归函数转化成循环

标题可能看起来有点抽象，continuation 我相信读者已经知道是什么了，defunctionalize 指的是一种把用到了高阶函数的代码变成无需高阶函数的代码的技术。可能这么说还是有点抽象，不如先讨论一个实际问题：中序遍历二叉树。我相信读者可以轻松写出它的递归版本

fsharp
type Tree =
    { Value: int
      Left: Option<Tree>
      Right: Option<Tree> }
let traverse (tree: Option<Tree>) =
    let rec traverse (tree: Option<Tree>) =
        match tree with
        | None -> ()
        | Some tree ->
            traverse tree.Left
            printfn "%i" tree.Value
            traverse tree.Right
    traverse tree

非常简单，只有十几行，但是要把变成循环比较麻烦。我们可以用 defunctionalize 机械化地达成这个要求。可是traverse里并没有用到高阶函数，我们要怎么对它 defunctionalize？这时，程序的数据都在栈上，互相之间由栈帧指针连接，程序的控制流被隐藏了起来。那自然可以想到把控制流暴露出来，也就是使用 CPS 转换

fsharp
let traverseCps (tree: Option<Tree>) =
    let rec traverse (tree: Option<Tree>) (kont: unit -> unit) =
        match tree with
        | None -> kont ()
        | Some tree ->
            let kont () =
                printfn "%i" tree.Value
                traverse tree.Right kont
            traverse tree.Left kont
    traverse tree (fun () -> ())

这时，数据都在闭包里，它们就像洋葱一样从里向外捕获前一次递归中的数据。但重点是控制流终于以高阶函数的形式暴露了出来，我们终于可以 defunctionalize 这些 continutation 了。defunctionalize 指的是这样一种过程：

  1. 分析函数的所有调用 -- traverse一共被调用了两次
  2. 并给每次调用传入的函数分配一个 sum type（F# 叫 discriminated union）的 variant，这里分别叫它们 DoNothing 和 PrintAndKont
  3. 每个 variant 上携带的值就是对应的闭包所捕获的变量 -- 第 6 行捕获了 tree 和前一个 kont，第 12 行没捕获任何东西
fsharp
type Kont =
    | DoNothing
    | PrintAndKont of Tree * Kont
  1. 再把每个函数实际运行的逻辑写到一个 match 里
fsharp
let traverseCpsDefun (tree: Option<Tree>) =
    let rec traverse (tree: Option<Tree>) (kont: Kont) =
        match tree with
        | None -> run kont
        | Some tree -> traverse tree.Left (PrintAndKont(tree, kont))
    and run kont =
        match kont with
        | DoNothing -> ()
        | PrintAndKont(tree, next) ->
            printfn "%i" tree.Value
            traverse tree.Right next
    traverse tree DoNothing

现在，我们的代码里已经不存在高阶函数了，只有两个互相递归的函数。数据从闭包里转移到了我们自己定义的 sum type 里。因为run只调用了一次，让我们把它 inline 掉

fsharp
let traverseCpsDefunInline (tree: Option<Tree>) =
    let rec traverse (tree: Option<Tree>) (kont: Kont) =
        match tree with
        | None ->
            match kont with
            | DoNothing -> ()
            | PrintAndKont(tree, next) ->
                printfn "%i" tree.Value
                traverse tree.Right next
        | Some tree -> traverse tree.Left (PrintAndKont(tree, kont))
    traverse tree DoNothing

我们惊讶地发现，这个函数的性质非常好：不仅没有了麻烦的高阶函数，它还是个尾递归函数。把尾递归函数转换成循环是很容易的

fsharp
let traverseLoop (tree: Option<Tree>) =
    let mutable tree = tree
    let mutable kont = DoNothing
    let mutable shouldEnd = false
    while not shouldEnd do
        match tree with
        | None ->
            match kont with
            | DoNothing -> shouldEnd <- true
            | PrintAndKont(t, next) ->
                printfn "%i" t.Value
                tree <- t.Right
                kont <- next
        | Some t ->
            tree <- t.Left
            kont <- PrintAndKont(t, kont)

眼光敏锐的读者可能已经发现了，我们定义的Kont，它和链表是同构的，那同时和栈也是同构的，第 15 行是在弹出，第 18 行是在压栈，那只要把Kont替换成栈

fsharp
let traverseIterative (tree: Option<Tree>) =
    let mutable tree = tree
    let stack = new Stack<Tree>()
    let mutable shouldEnd = false
    while not shouldEnd do
        match tree with
        | None ->
            if stack.Count = 0 then
                shouldEnd <- true
            else
                let t = stack.Pop()
                printfn "%i" t.Value
                tree <- t.Right
        | Some t ->
            tree <- t.Left
            stack.Push t

就能获得传统意义上，手工用栈模拟的循环式中序遍历二叉树。现在所有数据都在这个手动引入的栈上了。

读到现在，可能会有人有疑问，这有什么用呢？的确，defunctionalize 的理论意义远大于它的使用意义，但是上面的这些转换揭示了一些有趣的事实。首先，众所周知的一点：CPS 变换把隐藏的控制流暴露出来，让我们好对它做出各种优化。而对 continuation 的 defunctionalize 进一步地把无法优化的匿名函数变成了已知的 sum type，这又可以开展进一步的优化。而且，作为一种通用的递归转循环的手段，尽管最后一步转变成栈倒是需要一些领域知识，但这件事本身也已经很有趣了。

© 2016 - 2024Austaras Devas