improve stacktraces on errors

1 files changed, 39 insertions, 6 deletions

diff --git a/luajit.go b/luajit.go
index 9414787..7439023 100644
--- a/luajit.go
+++ b/luajit.go
@@ -39,6 +39,7 @@ type Lua struct {
 	running **Lua
 	yield func()
 	resume func() bool
+	tracebackHandler LuaRef
 }
 
 //export luna_run_go_func
@@ -53,6 +54,10 @@ func (l *Lua) Start() {
 	l.l = C.luaL_newstate()
 	l.running = &l
 	C.luaL_openlibs(l.l)
+
+	// Put debug handler into the registry for later copies to the stack.
+	l.PushGoFunction(TracebackHandler)
+	l.tracebackHandler = l.PopToRef()
 }
 
 // Close closes the Lua context.
@@ -62,12 +67,13 @@ func (l *Lua) Close() {
 
 // Require loads and executes the file pushing results onto the Lua stack.
 func (l *Lua) Require(file string) error {
+	l.PushTracebackHandler()
 	if !l.LoadFile(file) {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
 		return errors.New("could not open the file:\n" + errMsg)
 	}
-	err := l.PCall(0, C.LUA_MULTRET)
+	err := l.PCall(0, C.LUA_MULTRET, -2)
 	if err != nil {
 		return errors.New("could not execute the file:\n" + err.Error())
 	}
@@ -76,9 +82,9 @@ func (l *Lua) Require(file string) error {
 
 // PCall calls a function with the stack index (-nargs-1) expecting nresults
 // number of results.
-func (l *Lua) PCall(nargs int, nresults int) error {
+func (l *Lua) PCall(nargs int, nresults int, errfunc int) error {
 	*l.running = l
-	res := C.lua_pcall(l.l, C.int(nargs), C.int(nresults), 0)
+	res := C.lua_pcall(l.l, C.int(nargs), C.int(nresults), C.int(errfunc))
 	if res != C.LUA_OK {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
@@ -183,6 +189,14 @@ func (l *Lua) SetTableItem(key string) {
 	C.lua_setfield(l.l, -2, C.CString(key))
 }
 
+// GetTableItem gets a value from the table (offset -1) under the specified key
+// and puts it onto the stack.
+func (l *Lua) GetTableItem(key string) {
+	cstr := C.CString(key)
+	defer C.free(unsafe.Pointer(cstr))
+	C.lua_getfield(l.l, -1, C.CString(key))
+}
+
 // PushAny pushes value v onto the stack.
 func (l *Lua) PushAny(v any) error {
 	switch v.(type) {
@@ -256,6 +270,11 @@ func (l *Lua) PushFromRef(ref LuaRef) {
 	C.lua_rawgeti(l.l, C.LUA_REGISTRYINDEX, C.int(ref));
 }
 
+// Unref removes the ref from the registry.
+func (l *Lua) Unref(ref LuaRef) {
+	C.luaL_unref(l.l, C.LUA_REGISTRYINDEX, C.int(ref))
+}
+
 // Type returns type of the value sitting at n index on the stack.
 func (l *Lua) Type(n int) int {
 	return int(C.lua_type(l.l, C.int(n)))
@@ -301,12 +320,13 @@ func (l *Lua) Next() bool {
 func (l *Lua) LoadString(code string) error {
 	cstr := C.CString(code)
 	defer C.free(unsafe.Pointer(cstr))
+	l.PushTracebackHandler()
 	if C.luaL_loadstring(l.l, cstr) != C.LUA_OK {
 		errMsg := l.ToString(-1)
 		l.Pop(1)
 		return errors.New(errMsg)
 	}
-	err := l.PCall(0, 0)
+	err := l.PCall(0, 0, -2)
 	if err != nil {
 		return err
 	}
@@ -461,9 +481,22 @@ func (l *Lua) NewThread(yield func(), resume func() bool) *Lua {
 		running: l.running,
 		resume: resume,
 		yield: yield,
+		tracebackHandler: l.tracebackHandler,
 	}
 }
 
-func (l *Lua) Unref(ref LuaRef) {
-	C.luaL_unref(l.l, C.LUA_REGISTRYINDEX, C.int(ref))
+// PushTracebackHandler puts the traceback handler onto the stack for using
+// as a errfunc argument when calling PCall.
+func (l *Lua) PushTracebackHandler() {
+	l.PushFromRef(l.tracebackHandler)
+}
+
+// TracebackHandler handles stack trace formatting on errors.
+func TracebackHandler(l *Lua) int {
+	msg := l.ToString(-1)
+	l.GetGlobal("debug")
+	l.GetTableItem("traceback")
+	l.PushString(msg)
+	C.lua_call(l.l, 1, 1)
+	return 1
 }