lua.go

package modules

import (
	"config"
	"utils"
	"fmt"
	"reflect"
	"sync"

	gopherjson "github.com/alicebob/gopher-json"
	"github.com/sirupsen/logrus"

	"github.com/yuin/gopher-lua"
)

var (
	lstates     *utils.Pool
	fileLoader  func(string) string
	scriptlet   map[string]*lua.LFunction
	scriptletMu sync.Mutex
)

func loadModule(L *lua.LState) int {
	fname := L.CheckString(lua.UpvalueIndex(1))
	content := fileLoader(fname)
	if len(content) > 0 {
		fn, err := L.LoadString(content)
		if err != nil {
			panic(fmt.Errorf("lua.loadModule %s failed: %s", fname, err))
		}
		L.Push(fn)
		L.Call(0, 1)
		return 1
	}
	panic(ErrNotFound)
}

func preloader(L *lua.LState) int {
	L.Push(L.NewClosure(loadModule, L.Get(2)))
	return 1
}

func createLState() *lua.LState {
	logrus.Debug("createLState")
	L := lua.NewState()
	L.PreloadModule("json", gopherjson.Loader)

	// setup loader
	meta := L.NewTable()
	L.SetField(meta, "__index", L.NewFunction(preloader))
	preload := L.GetField(L.GetField(L.Get(lua.EnvironIndex), "package"), "preload")
	L.SetMetatable(preload, meta)
	return L
}

// InitLuaModuleWithLoader init with application
func InitLuaModuleWithLoader(loader func(string) string) {
	fileLoader = loader
	lstates = utils.NewPool()
	scriptlet = make(map[string]*lua.LFunction)
	for i := 0; i < config.MaxLuaStateCount; i++ {
		lstates.Put(createLState())
	}
}

func CloseLuaModule() {
	lstates.Wait()
	lstates.ForEach(func(obj interface{}) {
		obj.(*lua.LState).Close()
	})
	lstates = nil
}

// GetLState get LState
func GetLState() *utils.PoolElem {
	return lstates.Get()
}

// ReleaseLState release LState
func ReleaseLState(elem *utils.PoolElem) {
	lstates.Put(elem)
}

// AddLScriptlet adds named scriptlet
func AddLScriptlet(name, content string) (e error) {
	scriptletMu.Lock()
	defer scriptletMu.Unlock()
	if _, hasValue := scriptlet[name]; hasValue {
		return ErrAlreadyExists
	}

	defer func() {
		if err := recover(); err != nil {
			e = err.(error)
		}
	}()
	lstates.ForEach(func(obj interface{}) {
		L := obj.(*lua.LState)
		loader, err := L.LoadString(content)
		if err != nil {
			panic(err)
		}
		L.CallByParam(lua.P{
			Fn:      loader,
			NRet:    1,
			Protect: false,
		})
		scriptlet[name] = L.ToFunction(-1)
		L.Pop(1)
	})
	return nil
}

func toLValue(L *lua.LState, value interface{}) (lua.LValue, error) {
	switch v := value.(type) {
	case *lua.LFunction:
		return v, nil
	case lua.LGFunction:
		return L.NewFunction(v), nil
	case uint8:
		return lua.LNumber(v), nil
	case int:
		return lua.LNumber(v), nil
	case uint16:
		return lua.LNumber(v), nil
	case uint32:
		return lua.LNumber(v), nil
	case uint64:
		return lua.LNumber(v), nil
	case float32:
		return lua.LNumber(v), nil
	case float64:
		return lua.LNumber(v), nil
	case string:
		return lua.LString(v), nil
	case []byte:
		return lua.LString(v), nil
	case bool:
		return lua.LBool(v), nil
	case lua.LValue:
		return v, nil
	case chan lua.LValue:
		return lua.LChannel(v), nil
	case reflect.Value:
		return toLValue(L, v.Interface())
	default:
		t := reflect.TypeOf(value)
		switch t.Kind() {
		case reflect.Slice, reflect.Array:
			s := reflect.ValueOf(value)
			arr := L.CreateTable(s.Len(), 0)
			for i := 0; i < s.Len(); i++ {
				lv, err := toLValue(L, s.Index(i))
				if err != nil {
					return lua.LNil, err
				}
				arr.RawSetInt(i+1, lv)
			}
			return arr, nil
		case reflect.Map:
			s := reflect.ValueOf(value)
			keys := s.MapKeys()
			m := L.CreateTable(0, len(keys))
			for _, key := range keys {
				lkey, err := toLValue(L, key)
				if err != nil {
					return lua.LNil, err
				}
				lv, err := toLValue(L, s.MapIndex(key))
				if err != nil {
					return lua.LNil, err
				}
				m.RawSet(lkey, lv)
			}
			return m, nil
		default:
			return lua.LNil, fmt.Errorf("unsupported type %s", t)
		}
	}
}

// RunLScriptlet runs a named scriptlet added by AddScriptlet
func RunLScriptlet(L *lua.LState, name string, nRets int, args ...interface{}) (r interface{}, e error) {
	scriptletMu.Lock()
	lfn, hasValue := scriptlet[name]
	scriptletMu.Unlock()
	if hasValue {
		top := L.GetTop()
		L.Push(lfn)
		for _, arg := range args {
			lv, err := toLValue(L, arg)
			if err != nil {
				L.SetTop(top)
				return nil, err
			}
			L.Push(lv)
		}
		err := L.PCall(len(args), nRets, nil)
		if err != nil {
			return nil, fmt.Errorf("run scriptlet %s %s", name, err)
		}
		count := L.GetTop() - top
		defer L.Pop(count)
		if nRets == 0 {
			return nil, nil
		} else if nRets == 1 {
			return L.Get(-1), nil
		}
		rets := make([]lua.LValue, count)
		for i := 0; i < count; i++ {
			rets[i] = L.Get(-(count - i))
		}
		return rets, nil
	}
	return nil, ErrNotFound
}