自己动手实现Lua--实现TAILCALL指令
最近在看《自己动手实现Lua—虚拟机、编译器和标准库》。这是本挺不错的书，通过学习此书能够对Lua语言有比较深刻的理解，此外还可以对如何自己实现一门脚本语言有直观的认识。对于想学习Lua的同学，安利一下这本书。废话不多说，书中留了一个作业，让读者自己实现`TAILCALL`指令，实现尾调用的优化。本文就算是交作业吧。
最近在看《自己动手实现Lua—虚拟机、编译器和标准库》。这是本挺不错的书，通过学习此书能够对Lua语言有比较深刻的理解，此外还可以对如何自己实现一门脚本语言有直观的认识。对于想学习Lua的同学，安利一下这本书。
废话不多说，书中留了一个作业，让读者自己实现TAILCALL指令，实现尾调用的优化。本文就算是交作业吧。
本博客已经迁移至CatBro's Blog，那里是我自己搭建的个人博客，页面效果比这边更好，支持站内搜索，评论回复还支持邮件提醒，欢迎关注。这边只会在有时间的时候不定期搬运一下。
尾调用
尾调用，被调函数可以重用主调函数的调用帧，可以有效缓解调用栈溢出。
不过尾调用的条件非常苛刻，必须是直接返回函数调用。下面的是一个尾调用的例子，TAILCALL指令后面肯定紧跟着RETURN指令，并且RETURN指令的操作数A跟TAILCALL相同，RETURN指令的操作数B肯定是0。
$ luac -l -
return f(a)
^D
main <stdin:0,0> (5 instructions at 0x7fa2b9d00070)
0+ params, 2 slots, 1 upvalue, 0 locals, 2 constants, 0 functions
1 [1] GETTABUP 0 0 -1 ; _ENV "f"
2 [1] GETTABUP 1 0 -2 ; _ENV "a"
3 [1] TAILCALL 0 2 0
4 [1] RETURN 0 0
5 [1] RETURN 0 1
下面几个例子，都不是尾调用。
$ luac -l -
return (f(a))
^D
main <stdin:0,0> (5 instructions at 0x7fb5a34035e0)
0+ params, 2 slots, 1 upvalue, 0 locals, 2 constants, 0 functions
1 [1] GETTABUP 0 0 -1 ; _ENV "f"
2 [1] GETTABUP 1 0 -2 ; _ENV "a"
3 [1] CALL 0 2 2
4 [1] RETURN 0 2
5 [1] RETURN 0 1
$ luac -l -
return f(a)+1
^D
main <stdin:0,0> (6 instructions at 0x7f98c3d00070)
0+ params, 2 slots, 1 upvalue, 0 locals, 3 constants, 0 functions
1 [1] GETTABUP 0 0 -1 ; _ENV "f"
2 [1] GETTABUP 1 0 -2 ; _ENV "a"
3 [1] CALL 0 2 2
4 [1] ADD 0 0 -3 ; - 1
5 [1] RETURN 0 2
6 [1] RETURN 0 1
$ luac -l -
return {f(a)}
^D
main <stdin:0,0> (7 instructions at 0x7fb1b95006f0)
0+ params, 3 slots, 1 upvalue, 0 locals, 2 constants, 0 functions
1 [1] NEWTABLE 0 0 0
2 [1] GETTABUP 1 0 -1 ; _ENV "f"
3 [1] GETTABUP 2 0 -2 ; _ENV "a"
4 [1] CALL 1 2 0
5 [1] SETLIST 0 0 1 ; 1
6 [1] RETURN 0 2
7 [1] RETURN 0 1
$ luac -l -
return 1, f(a)
^D
main <stdin:0,0> (6 instructions at 0x7f8f9c500070)
0+ params, 3 slots, 1 upvalue, 0 locals, 3 constants, 0 functions
1 [1] LOADK 0 -1 ; 1
2 [1] GETTABUP 1 0 -2 ; _ENV "f"
3 [1] GETTABUP 2 0 -3 ; _ENV "a"
4 [1] CALL 1 2 0
5 [1] RETURN 0 0
6 [1] RETURN 0 1
CALL指令
我们先来看看普通的CALL指令的操作流程：
// R(A), ... ,R(A+C-2) := R(A)(R(A+1), ... ,R(A+B-1))
func call(i Instruction, vm LuaVM) {
a, b, c := i.ABC()
a += 1
// println(":::"+ vm.StackToString())
nArgs := _pushFuncAndArgs(a, b, vm)
vm.Call(nArgs, c-1)
_popResults(a, c, vm)
}
首先将位于寄存器中的函数和参数推入栈顶，然后调用Call()方法执行函数，最后将栈上的返回值出栈并放入指定寄存器中。
_pushFuncAndArgs()和_popResults()分别要处理操作数B和操作数C为0的特殊情况。操作数B为0的情况，部分参数已经在栈上了（由前面的CALL指令或VARARG指令留在栈上）。
func _pushFuncAndArgs(a, b int, vm LuaVM) (nArgs int) {
if b >= 1 {
vm.CheckStack(b)
for i := a; i < a+b; i++ {
vm.PushValue(i)
}
return b - 1
} else {
_fixStack(a, vm)
return vm.GetTop() - vm.RegisterCount() - 1
}
}
而操作数C为0的情况，则不需要将返回值从栈上弹出。直接将返回值留在栈上，供后面的指令使用。
func _popResults(a, c int, vm LuaVM) {
if c == 1 {
// no results
} else if c > 1 {
for i := a + c - 2; i >= a; i-- {
vm.Replace(i)
}
} else {
// leave results on stack
vm.CheckStack(1)
vm.PushInteger(int64(a))
}
}
然后我们来看下Call()方法做了什么。它首先根据索引找到要调用的函数的值，检查它是否是闭包类型，如果不是直接报错。然后通过callLuaClosure()调用该函数。
// [-(nargs+1), +nresults, e]
// http://www.lua.org/manual/5.3/manual.html#lua_call
func (self *luaState) Call(nArgs, nResults int) {
val := self.stack.get(-(nArgs + 1))
if c, ok := val.(*closure); ok {
fmt.Printf("call %s<%d,%d>\n", c.proto.Source,
c.proto.LineDefined, c.proto.LastLineDefined)
self.callLuaClosure(nArgs, nResults, c)
} else {
panic("not function!")
}
}
callLuaClosure()稍微有点复杂，来看下代码。首先从闭包的函数原型中获取到各种信息，如寄存器个数、固定参数个数、是否是vararg。接着创建一个新的调用帧，并将闭包与调用帧关联。然后将函数和参数值全部从主调帧栈顶弹出，并将固定参数压入被调帧栈顶，如果是vararg且参数个数大于固定参数个数，还要将vararg参数记录下来。
到这新帧就准备就绪了，将新帧推入调用栈顶，然后调用runLuaClosure()开始执行被调函数的指令。执行结束之后，被调帧的任务结束，将其弹出调用栈顶。
此时，返回值还留在被调帧的栈顶，需要移到主调帧栈顶。
func (self *luaState) callLuaClosure(nArgs, nResults int, c *closure) {
nRegs := int(c.proto.MaxStackSize)
nParams := int(c.proto.NumParams)
isVararg := c.proto.IsVararg == 1
// create new lua stack
newStack := newLuaStack(nRegs + 20)
newStack.closure = c
// pass args, pop func
funcAndArgs := self.stack.popN(nArgs + 1)
newStack.pushN(funcAndArgs[1:], nParams)
newStack.top = nRegs
if nArgs > nParams && isVararg {
newStack.varargs = funcAndArgs[nParams+1:]
}
// run closure
self.pushLuaStack(newStack)
self.runLuaClosure()
self.popLuaStack()
// return results, nResults == c - 1
if nResults != 0 {
results := newStack.popN(newStack.top - nRegs)
self.stack.check(len(results))
self.stack.pushN(results, nResults)
}
}
实现TAILCALL指令
知道了CALL指令的流程，我们就可以着手实现TAILCALL指令了。其操作数A和操作数B的含义跟CALL指令完全一致，操作数C没用，相当于固定为0。所以_pushFuncAndArgs()和_popResults()函数可以重用，主要是修改中间的调用流程。我们首先给LuaVM新增一个接口TailCall()。
在api/lua_vm.go文件中添加如下代码：
type LuaVM interface {
...
TailCall(nArgs int) // add this
}
因为TAILCALL指令后面紧跟着RETURN指令，且RETURN指令的操作数B为0，操作数A跟TAILCALL指令一样。所以_popResults()之后，我们啥都不用干，直接把返回值保留在栈上即可。
在vm/inst_call.go文件中修改tailCall()如下：
// return R(A)(R(A+1), ... ,R(A+B-1))
func tailCall(i Instruction, vm LuaVM) {
a, b, _ := i.ABC()
a += 1
nArgs := _pushFuncAndArgs(a, b, vm)
vm.TailCall(nArgs)
_popResults(a, 0, vm)
// no need to _return() as ‘b’ of the following ‘RETURN’ is 0,
// ‘a’ of ‘RETURN’ is same ‘as’ a of ‘TAILCALL’
}
在state/api_call.go文件中添加TailCall()代码如下
// [-(nargs+1), +nresults, e]
func (self *luaState) TailCall(nArgs int) {
val := self.stack.get(-(nArgs + 1))
if c, ok := val.(*closure); ok {
fmt.Printf("tailcall %s<%d,%d>\n", c.proto.Source,
c.proto.LineDefined, c.proto.LastLineDefined)
self.tailCallLuaClosure(nArgs, c)
} else {
panic("not function!")
}
}
然后在tailCallLuaClosure()中实现主要逻辑。在state/api_call.go文件中添加tailCallLuaClosure()代码如下。
首先同样是从闭包中获取函数原型的信息。因为当前帧在TAILCALL之后肯定跟着RETURN，所以保存参数之后可以直接清理掉当前帧的栈，然后直接给被调函数重用。将被调函数的信息设置到当前帧，先检查栈空间是否够，然后将当前帧关联到新的闭包，然后将固定参数推入栈顶，修改栈顶指针指向最后一个寄存器。如果是vararg且参数个数大于固定参数个数，还要将vararg参数记录下来。
一切就绪之后，就可以调用runLuaClosure()开始执行闭包的指令了。执行完毕后返回值保留在栈顶。
func (self *luaState) tailCallLuaClosure(nArgs int, c *closure) {
nRegs := int(c.proto.MaxStackSize)
nParams := int(c.proto.NumParams)
isVararg := c.proto.IsVararg == 1
// store args
args := self.stack.popN(nArgs)
// clean the stack
self.SetTop(0)
// check if stack space is enough
self.stack.check(nRegs + 20)
// substitue the closure to new one
self.stack.closure = c
// push fixed args
self.stack.pushN(args[1:], nParams)
self.stack.top = nRegs
// store varargs
if nArgs > nParams && isVararg {
self.stack.varargs = args[nParams+1:]
}
// run closure
self.runLuaClosure()
}
测试代码
我们重新编译Go代码
cd $LUAGO/go
export GOPATH=$PWD/ch08
export GOBIN=$PWD/ch08/bin
go install luago
然后来编写Lua脚本，我们编写了一个求和的函数
local function sum(n, s, fun)
if n == 0 then
return s
end
s = s + n
return fun(n-1, s, fun)
end
local function assert(v)
if not v then fail() end
end
local v1 = sum(0, 0, sum)
assert(v1 == 0)
local v2 = sum(1, 0, sum)
assert(v2 == 1)
local v3 = sum(3, 0, sum)
assert(v3 == 6)
local v4 = sum(10, 0, sum)
assert(v4 == 55)
local v5 = sum(10000, 0, sum)
assert(v5 == 50005000)
local v6 = sum(1000000, 0, sum)
assert(v6 == 500000500000)
先来看看sum()函数会被编译成什么指令。
$ luac -l -
local function sum(n, s, fun)
if n == 0 then
return s
end
s = s + n
return fun(n-1, s, fun)
end
^D
main <stdin:0,0> (2 instructions at 0x7fe071d00070)
0+ params, 2 slots, 1 upvalue, 1 local, 0 constants, 1 function
1 [7] CLOSURE 0 0 ; 0x7fe071e00110
2 [7] RETURN 0 1
function <stdin:1,7> (11 instructions at 0x7fe071e00110)
3 params, 7 slots, 0 upvalues, 3 locals, 2 constants, 0 functions
1 [2] EQ 0 0 -1 ; - 0
2 [2] JMP 0 1 ; to 4
3 [3] RETURN 1 2
4 [5] ADD 1 1 0
5 [6] MOVE 3 2
6 [6] SUB 4 0 -2 ; - 1
7 [6] MOVE 5 1
8 [6] MOVE 6 2
9 [6] TAILCALL 3 4 0
10 [6] RETURN 3 0
11 [7] RETURN 0 1
可以看到的确是被编译成了TAILCALL，后面紧跟RETURN指令，且RETURN指令的操作数A与TAILCALL相同，操作数B为0。
我们编译Lua脚本，然后用我们的虚拟机进行执行。
luac ../lua/ch08/tailcall.lua
./ch08/bin/luago luac.out
call @../lua/ch08/tailcall.lua<0,0>
call @../lua/ch08/tailcall.lua<1,7>
call @../lua/ch08/tailcall.lua<9,11>
call @../lua/ch08/tailcall.lua<1,7>
call @../lua/ch08/tailcall.lua<9,11>
panic: GETTABUP
哦哦，执行失败了，残念！
于是加日志进行问题排查，把指令执行时的栈打出来
{% fold 点击展开 %}
call @../lua/ch08/tailcall.lua<0,0>
CLOSURE 0 0 [nil][nil][nil][nil][nil][nil][nil]
CLOSURE 1 1 [function][nil][nil][nil][nil][nil][nil]
MOVE 2 0 [function][function][nil][nil][nil][nil][nil]
LOADK 3 -1 [function][function][function][nil][nil][nil][nil]
LOADK 4 -1 [function][function][function][0][nil][nil][nil]
MOVE 5 0 [function][function][function][0][0][nil][nil]
CALL 2 4 2 [function][function][function][0][0][function][nil]
call @../lua/ch08/tailcall.lua<1,7>
EQ 0 0 -1 [0][0][function][nil][nil][nil][nil]
RETURN 1 2 [0][0][function][nil][nil][nil][nil]
RETURN after [0][0][function][nil][nil][nil][nil][0]
CALL after [function][function][0][0][0][function][nil]
MOVE 3 1 [function][function][0][0][0][function][nil]
EQ 1 2 -1 [function][function][0][function][0][function][nil]
JMP 0 1 [function][function][0][function][0][function][nil]
LOADBOOL 4 1 0 [function][function][0][function][0][function][nil]
CALL 3 2 1 [function][function][0][function][true][function][nil]
call @../lua/ch08/tailcall.lua<9,11>
TEST 0 1 [true][nil]
JMP 0 2 [true][nil]
RETURN 0 1 [true][nil]
RETURN after [true][nil]
CALL after [function][function][0][function][true][function][nil]
MOVE 3 0 [function][function][0][function][true][function][nil]
LOADK 4 -2 [function][function][0][function][true][function][nil]
LOADK 5 -1 [function][function][0][function][1][function][nil]
MOVE 6 0 [function][function][0][function][1][0][nil]
CALL 3 4 2 [function][function][0][function][1][0][function]
call @../lua/ch08/tailcall.lua<1,7>
EQ 0 0 -1 [1][0][function][nil][nil][nil][nil]
JMP 0 1 [1][0][function][nil][nil][nil][nil]
ADD 1 1 0 [1][0][function][nil][nil][nil][nil]
MOVE 3 2 [1][1][function][nil][nil][nil][nil]
SUB 4 0 -2 [1][1][function][function][nil][nil][nil]
MOVE 5 1 [1][1][function][function][0][nil][nil]
MOVE 6 2 [1][1][function][function][0][1][nil]
TAILCALL 3 4 0 [1][1][function][function][0][1][function]
RETURN 3 0 [1][function][nil][nil][nil][nil][nil]
RETURN after [1][function][nil][nil][nil][nil]
TAILCALL after [1][function][nil][nil][nil][nil][4]
RETURN 0 1 [1][function][nil][nil][nil][nil][4]
RETURN after [1][function][nil][nil][nil][nil][4]
CALL after [function][function][0][nil][1][0][function]
MOVE 4 1 [function][function][0][nil][1][0][function]
EQ 1 3 -2 [function][function][0][nil][function][0][function]
LOADBOOL 5 0 1 [function][function][0][nil][function][0][function]
CALL 4 2 1 [function][function][0][nil][function][false][function]
call @../lua/ch08/tailcall.lua<9,11>
TEST 0 1 [false][nil]
GETTABUP 1 0 -1 [false][nil]
panic: GETTABUP
{% endfold %}
我们发现在TAILCALL开始执行之后立马调用了RETURN，问题就是出在这里，我们虽然替换了当前帧的闭包为被调函数的闭包，但是忘了更新pc。于是修改tailCallLuaClosure()初始化pc为0。
func (self *luaState) tailCallLuaClosure(nArgs int, c *closure) {
// substitue the closure to new one
self.stack.closure = c
self.stack.pc = 0 // add this
}
重新编译之后测试
$ go install luago
$ ./ch08/bin/luago luac.out
... # 省略前面的日志
call @../lua/ch08/tailcall.lua<1,7>
EQ 0 0 -1 [1][0][function][nil][nil][nil][nil]
JMP 0 1 [1][0][function][nil][nil][nil][nil]
ADD 1 1 0 [1][0][function][nil][nil][nil][nil]
MOVE 3 2 [1][1][function][nil][nil][nil][nil]
SUB 4 0 -2 [1][1][function][function][nil][nil][nil]
MOVE 5 1 [1][1][function][function][0][nil][nil]
MOVE 6 2 [1][1][function][function][0][1][nil]
TAILCALL 3 4 0 [1][1][function][function][0][1][function]
EQ 0 0 -1 [1][function][nil][nil][nil][nil][nil]
JMP 0 1 [1][function][nil][nil][nil][nil][nil]
ADD 1 1 0 [1][function][nil][nil][nil][nil][nil]
panic: arithmetic error!
哦哦，又执行失败了