尾调用

Lua语言是支持尾调用(tail call)消除的。本节介绍并支持尾调用。

首先介绍尾调用这个概念。当一个函数的最后一个动作是调用另一个函数而没有再进行其他工作时,就形成了尾调用。比如下面的示例代码:

function foo(a, b)
    return bar(a + b)
end

foo()函数的最后一个动作(这个例子里也是唯一的动作)就是调用bar()函数。下面先看看在不引入尾调用的情况下,foo()函数的执行过程,如下图所示:

|       |        |       |        |       |        |       |
+-------+        +-------+        +-------+        +-------+
| foo() |        | foo() |        | foo() |      / |  ret1 |
+-------<<       +-------+        +-------<<   /-+ +-------+
|   a   |        |   a   |        |   a   |    | \ |  ret2 |
+-------+        +-------+        +-------+    |   +-------+
|   b   |        |   b   |        |   b   |    |   |       |
+-------+        +-------+        +-------+    |
| bar() |        | bar() |      / |  ret1 | \  |
+-------+        +-------<<   /-+ +-------+  >-/返回值
|  a+b  |        |  a+b  |    | \ |  ret2 | /
+-------+        +-------+    |   +-------+
|       |        :       :    |   |       |
                 +-------+    |
                 |  ret1 | \  |
                 +-------+  >-/返回值
                 |  ret2 | /
                 +-------+
                 |       |
  • 最左边第1个图,是在foo()函数内部准备好了调用bar()函数之前的栈布局。也就是在调用Call(bar)字节码之前。

  • 第2个图是bar()函数调用刚刚完成后的栈布局。也就是bar()函数的Return字节码执行完毕后,但还没有返回到foo()函数的Call(bar)字节码之前。假设这个函数有两个返回值ret1ret2,目前在栈顶。

  • 第3个图是bar()函数返回后的栈布局。也就是foo()Call(bar)字节码执行完毕。即把两个返回值挪到bar()的入口函数位置。

  • 第4个图是foo()函数返回后的栈布局。也就是更外层的调用者的Call(foo)字节码执行完毕后。即把两个返回值挪到foo()的入口函数位置。

后面的3个图是连续执行的。观察下其中的优化空间:

  • 一个比较明显的优化思路是,最后的两次返回值的复制可以一步完成。但这个很难优化,而且并不会优化多少性能;

  • 另外一个不那么明显的地方是,在最左边第1个图bar()函数准备好调用后,foo()函数的栈空间就再没有用到了。所以,可以在调用bar()函数前,先清理掉foo()函数占用的栈空间。按照这个思路,下面重新绘制调用流程:

|       |        |       |        |       |        |       |
+-------+        +-------+        +-------+        +-------+
| foo() |      / | bar() |        | bar() |      / |  ret1 |
+-------<<   /-+ +-------<<       +-------<<   /-+ +-------+
|   a   |    | \ |  a+b  |        |  a+b  |    | \ |  ret2 |
+-------+    |   +-------+        +-------+    |   +-------+
|   b   |    |   |       |        :       :    |   |       |
+-------+    |                    +-------+    |
| bar() | \  |                    |  ret1 | \  |
+-------+  >-/                    +-------+  >-/
|  a+b  | /                       |  ret2 | /
+-------+                         +-------+
|       |                         |       |
  • 最左边第1个图不变,仍然是bar()函数调用前的状态;

  • 第2个图,在调用bar()前,先清理掉foo()函数的栈空间;

  • 第3个图,对应上面的第2个图,是调用完bar()函数后。

  • 第4个图,对应上面最后一个图。由于刚才已经清理过foo()函数的栈空间,所以跳过了上面的第3个图。

跟上面的普通流程对比,这个新流程的操作步骤虽然有改变,但并没有减少,所以对性能并没有优化。但是,在栈空间的使用上有优化!在bar()函数执行之前就已经释放了foo()的栈空间。2层函数调用,但只占用了1层的空间。这带来的优势在这个例子中并不明显,但是在递归调用中就很明显,因为一般递归调用都会有非常多层。如果递归调用的最后一条满足上述尾调用,那么在应用新流程后,就可以支持无限次的递归调用,而不会导致栈溢出!这里的栈溢出,指的是上图中画的Lua虚拟机的栈,而不是Rust程序的栈溢出。

相比于上面的普通流程,这个新流程还有一个小的不同。上面每个图中栈上的<<代表的是当前self.base的位置。可以看到在上面的普通流程中,self.base发生过变化;而在新的流程中,全程没有变化。

在介绍完尾调用的概念后,下面介绍具体实现。

语法分析

在开始语法分析前,再次明确下尾调用的规则:当一个函数的最后一个动作是调用另一个函数而没有再进行其他工作时,就形成了尾调用。下面举一些《Lua程序设计》一书中的反例:

function f1(x)
    g(x)  -- 在f1()返回前,还要丢掉g(x)的返回值
end
function f2(x)
    return g(x) + 1  -- 还要执行+1
end
function f3(x)
    return x or g(x)  -- 还要把g(x)的返回值限制为1个
end
function f4(x)
    return (g(x))  -- 还要把g(x)的返回值限制为1个
end

在Lua语言中,只有形如return func(args)的调用才是尾调用。当然这里的funcargs可以很复杂,比如return t.k(a+b.f())也是尾调用。

在明确规则后,语法分析时判断尾调用就比较简单。在解析return语句时,增加对尾调用的判断:

    let iret = self.sp;
    let (nexp, last_exp) = self.explist();

    if let (0, &ExpDesc::Local(i)) = (nexp, &last_exp) {
        // 只有1个返回值并且是局部变量
        ByteCode::Return(i as u8, 1)

    } else if let (0, &ExpDesc::Call(func, narg_plus)) = (nexp, &last_exp) {
        // 新增尾调用:只有1个返回值,并且是函数调用
        ByteCode::TailCall(func as u8, narg_plus as u8)

    } else if self.discharge_expand(last_exp) {
        // 最后一个返回值是可变类型,比如可变参数或者函数调用,
        // 则在语法分析阶段无法得知返回值个数
        ByteCode::Return(iret as u8, 0)

    } else {
        // 最后一个返回值不是可变类型
        ByteCode::Return(iret as u8, nexp as u8 + 1)
    }

上述代码中一共有4个情况。其中第2个情况是新增的尾调用,另外3种情况都是本章之前小节中已经支持的,这里不再介绍。

新增的字节码TailCall类似于函数调用字节码Call,但是由于尾调用的返回值肯定是函数调用,返回值个数肯定未知,所以就省略第3个关联参数。至此,函数调用相关的字节码就有3个了:

pub enum ByteCode {
    Call(u8, u8, u8),
    CallSet(u8, u8, u8),
    TailCall(u8, u8),  // 新增尾调用

虚拟机执行

接下来看尾调用的虚拟机执行部分。由本节开头介绍的尾调用的流程,可以得出相对于普通的函数调用,尾调用的执行有3点不同:

  • 在调用内层函数前,先提前清理外层函数的栈空间,这也是尾调用的意义所在;
  • 在内层函数返回后,由于外层函数已经被清理,所以没必要返回给外层函数,而是直接返回给更外层的调用函数。
  • 全程不需要调整self.base

由此,可以实现TailCall字节码的执行流程如下:

    ByteCode::TailCall(func, narg_plus) => {
        self.stack.drain(self.base-1 .. self.base+func as usize);
        return self.do_call_function(narg_plus);
    }

非常简单,只有两行代码:

第1行,通过self.stack.drain()来清理外层函数的栈空间。

第2行,通过return语句直接从当前execute()中返回,也就是说当内层函数执行完毕后,无需返回给当前函数,而是直接返回给更外层的调用者。另外,根据上面列出的尾调用的规则,这一行Rust代码本身也属于尾调用。所以只要Rust语言也支持尾调用消除,那么我们的Lua解释器在执行过程中,其本身的栈也不会有增加。

另外,第2行中新增的do_call_function()方法执行具体的函数调用,其代码是从之前小节中CallCallSet字节码调用的call_function()方法中提取出来的,只不过去掉了对self.base的更新。而call_function()方法就修改为对这个新方法的包装:

    fn call_function(&mut self, func: u8, narg_plus: u8) -> usize {
        self.base += func as usize + 1; // get into new world
        let nret = self.do_call_function(narg_plus);
        self.base -= func as usize + 1; // come back
        nret
    }

测试

至此,我们完成了尾调用。用下面的Lua代码验证下:

function f(n)
    if n > 10000 then return n end
    return f(n+1)
end
print(f(0))

但是执行时出现栈溢出错误:

$ cargo r -- test_lua/tailcall.lua

thread 'main' has overflowed its stack
fatal runtime error: stack overflow
[1]    85084 abort      cargo r -- test_lua/tailcall.lua

起初我以为是Rust的debug版本没有进行尾调用优化,但是后来加上--release后,也只是可以支持更大的递归深度,推迟了栈溢出,但最终还是会栈溢出。这就要回到刚才说的:“所以只要Rust语言也支持尾调用消除,那么。。。”,这句话前面的假设可能不成立,即Rust语言可能不支持尾调用消除。这里有篇文章介绍了Rust语言对于尾调用的讨论,结论大概就是由于实现太复杂(可能是要涉及资源drop),并且收益有限(如果真有必要那么程序员可以手动改递归为循环),所以最终Rust语言并不支持尾调用消除。这么说来,为了让本节完成的Lua的尾调用消除有意义,只能把对execute()函数的递归调用改成循环。这个改动本身并不难,但是后续还有两处要修改函数调用流程的地方,一是整个程序的入口函数的调用方式,二是支持协程后函数的状态保存。所以计划在完成最终的函数调用流程后,再做这个改动。