lua/vm.go

// Package lua ...
//
// Description : lua ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2024-11-14 17:46
package lua

import (
	"fmt"
	libs "github.com/vadv/gopher-lua-libs"
	luaCompile "github.com/yuin/gopher-lua"
	"github.com/yuin/gopher-lua/ast"
	"github.com/yuin/gopher-lua/parse"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

var (
	VMInstance *VM
)

// InitVM 初始化lua虚拟机
func InitVM(poolSize int, libPath []string, preComplierScript map[string]string) error {
	if len(libPath) == 0 {
		libPath = []string{"./lib"} // 默认使用内置的lib
	}
	pathStr := ""
	for _, path := range libPath {
		// 绝对路径
		libDir, _ := filepath.Abs(path)
		libDir = filepath.Join(libDir, "?.lua")
		pathStr += ";" + libDir
	}
	pathStr += ";;"
	os.Setenv("LUA_PATH", pathStr)
	if poolSize <= 0 {
		poolSize = 32
	}
	VMInstance = &VM{
		l:                 &sync.RWMutex{},
		scriptLock:        &sync.RWMutex{},
		pool:              make([]*luaCompile.LState, 0),
		preComplierScript: make(map[string]*luaCompile.FunctionProto),
	}
	for scriptID, scriptContent := range preComplierScript {
		if err := VMInstance.AddScript(scriptID, scriptContent); nil != err {
			return err
		}
	}
	return nil
}

type VM struct {
	l                 *sync.RWMutex
	scriptLock        *sync.RWMutex
	pool              []*luaCompile.LState                 // lua虚拟机实例池
	preComplierScript map[string]*luaCompile.FunctionProto // 预编译的脚本内容
}

// AddScript 脚本编译
func (v *VM) AddScript(scriptID string, script string) error {
	v.scriptLock.Lock()
	defer v.scriptLock.Unlock()
	var (
		err   error
		chunk []ast.Stmt
		proto *luaCompile.FunctionProto
	)

	if chunk, err = parse.Parse(strings.NewReader(script), scriptID); nil != err {
		return err
	}
	if proto, err = luaCompile.Compile(chunk, scriptID); nil != err {
		return err
	}
	v.preComplierScript[scriptID] = proto
	return nil
}

// RemoveScript 删除一个已编译的脚本
func (v *VM) RemoveScript(scriptID string) {
	v.scriptLock.Lock()
	defer v.scriptLock.Unlock()
	delete(v.preComplierScript, scriptID)
}

// GetVm 获取一个VM实例
func (v *VM) GetVm() *luaCompile.LState {
	v.l.Lock()
	defer v.l.Unlock()
	if len(v.pool) == 0 {
		return luaCompile.NewState()
	}
	l := v.pool[len(v.pool)-1]
	// 基于libs预导入lua的常用模块
	libs.Preload(l)
	v.pool = v.pool[:len(v.pool)-1]
	return l
}

// Close 关闭指定虚拟机
func (v *VM) Close(l *luaCompile.LState) {
	l.Close()
	v.l.Lock()
	defer v.l.Unlock()
	// 实例放回实例实例池
	v.pool = append(v.pool, l)
}

// Shutdown 关闭全部虚拟机
func (v *VM) Shutdown() {
	for _, item := range v.pool {
		item.Close()
	}
}

// RunPreCompileScript 执行预编译过的脚本
func (v *VM) RunPreCompileScript(scriptID string) (*luaCompile.LState, error) {
	var (
		err    error
		proto  *luaCompile.FunctionProto
		exists bool
	)
	v.scriptLock.Lock()
	if proto, exists = v.preComplierScript[scriptID]; !exists {
		v.scriptLock.Unlock()
		return nil, fmt.Errorf("preComplier script %s not found", scriptID)
	}
	v.scriptLock.Unlock()
	l := v.GetVm()
	lFunc := l.NewFunctionFromProto(proto)
	l.Push(lFunc)
	if err = l.PCall(0, luaCompile.MultRet, nil); nil != err {
		v.Close(l)
		return nil, err
	}
	return l, nil
}

// Run 运行脚本
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 18:08 2024/11/14
func (v *VM) Run(script string) (*luaCompile.LState, error) {
	l := v.GetVm()
	if err := l.DoString(script); err != nil {
		// 直接归还实例
		l.Close()
		return nil, err
	}
	return l, nil
}

// GetResultAndRemove 获取最后一个返回值, 并从结果中移除
func (v *VM) GetResultAndRemove(l *luaCompile.LState) luaCompile.LValue {
	val := l.Get(-1)
	l.Pop(1)
	return val
}

// GetResultByIndex 根据索引获取返回值
func (v *VM) GetResultByIndex(l *luaCompile.LState, index int) luaCompile.LValue {
	val := l.Get(index)
	return val
}