added support for LUA modules

Author: kelindar

fixtures/demo.lua

@@ -0,0 +1,5 @@
+local demo = require("demo_mod")
+
+function main(input)
+    return demo.Mult(5, 5)
+end

fixtures/module.lua

@@ -0,0 +1,7 @@
+local demo_mod = {} -- The main table
+
+function demo_mod.Mult(a, b)
+  return a * b
+end
+
+return demo_mod

gen/binary_test.go

@@ -11,7 +11,7 @@ import (
 //go:generate genny -in=$GOFILE -out=../z_binary_test.go gen "TIn=String,Number,Bool TOut=String,Number,Bool"
 
 func Test_TInTOut(t *testing.T) {
-	m := &Module{Name: "test"}
+	m := &NativeModule{Name: "test"}
 	m.Register("test1", func(v TIn) (TOut, error) {
 		return newTestValue(TypeTOut).(TOut), nil
 	})

gen/unary_in_test.go

@@ -11,7 +11,7 @@ import (
 //go:generate genny -in=$GOFILE -out=../z_unary_test.go gen "TIn=String,Number,Bool"
 
 func Test_In_TIn(t *testing.T) {
-	m := &Module{Name: "test"}
+	m := &NativeModule{Name: "test"}
 	m.Register("test1", func(v TIn) error {
 		return nil
 	})
@@ -45,7 +45,7 @@ func Test_In_TIn(t *testing.T) {
 }
 
 func Test_Out_TIn(t *testing.T) {
-	m := &Module{Name: "test"}
+	m := &NativeModule{Name: "test"}
 	m.Register("test1", func() (TIn, error) {
 		return newTestValue(TypeTIn).(TIn), nil
 	})

json/json.go

@@ -12,6 +12,12 @@ import (
 	"github.com/yuin/gopher-lua"
 )
 
+var (
+	errNested      = errors.New("cannot encode recursively nested tables to JSON")
+	errSparseArray = errors.New("cannot encode sparse array")
+	errInvalidKeys = errors.New("cannot encode mixed or invalid key types")
+)
+
 // Loader is the module loader function.
 func Loader(L *lua.LState) int {
 	t := L.NewTable()
@@ -51,11 +57,7 @@ func apiEncode(L *lua.LState) int {
 	return 1
 }
 
-var (
-	errNested      = errors.New("cannot encode recursively nested tables to JSON")
-	errSparseArray = errors.New("cannot encode sparse array")
-	errInvalidKeys = errors.New("cannot encode mixed or invalid key types")
-)
+// --------------------------------------------------------------------
 
 type invalidTypeError lua.LValueType
 

module.go

@@ -19,8 +19,43 @@ var (
 
 var builtin = make(map[reflect.Type]func(interface{}) lua.LGFunction, 8)
 
-// Module represents a loadable module
-type Module struct {
+// Module represents a loadable module.
+type Module interface {
+	inject(state *lua.LState) error
+}
+
+// --------------------------------------------------------------------
+
+// ScriptModule represents a loadable module written in LUA itself.
+type ScriptModule struct {
+	Script  *Script // The script that contains the module
+	Name    string  // The name of the module
+	Version string  // The module version string
+}
+
+// Inject loads the module into the state
+func (m *ScriptModule) inject(runtime *lua.LState) error {
+
+	// Inject the prerequisite modules of the module
+	if err := m.Script.loadModules(runtime); err != nil {
+		return err
+	}
+
+	// Push the function to the runtime
+	codeFn := runtime.NewFunctionFromProto(m.Script.code)
+	preload := runtime.GetField(runtime.GetField(runtime.Get(lua.EnvironIndex), "package"), "preload")
+	if _, ok := preload.(*lua.LTable); !ok {
+		return errors.New("package.preload must be a table")
+
+	}
+	runtime.SetField(preload, m.Name, codeFn)
+	return nil
+}
+
+// --------------------------------------------------------------------
+
+// NativeModule represents a loadable native module.
+type NativeModule struct {
 	lock    sync.Mutex
 	funcs   map[string]fngen
 	Name    string // The name of the module
@@ -87,7 +122,7 @@ func (g *fngen) generate() lua.LGFunction {
 }
 
 // Register registers a function into the module.
-func (m *Module) Register(name string, function interface{}) error {
+func (m *NativeModule) Register(name string, function interface{}) error {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 
@@ -106,14 +141,14 @@ func (m *Module) Register(name string, function interface{}) error {
 }
 
 // Unregister unregisters a function from the module.
-func (m *Module) Unregister(name string) {
+func (m *NativeModule) Unregister(name string) {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 	delete(m.funcs, name)
 }
 
 // Inject loads the module into the state
-func (m *Module) inject(state *lua.LState) {
+func (m *NativeModule) inject(state *lua.LState) error {
 	table := make(map[string]lua.LGFunction, len(m.funcs))
 	for name, g := range m.funcs {
 		table[name] = g.generate()
@@ -125,6 +160,7 @@ func (m *Module) inject(state *lua.LState) {
 		state.Push(mod)
 		return 1
 	})
+	return nil
 }
 
 // validate validates the function type

module_test.go

@@ -11,8 +11,8 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func testModule() *Module {
-	m := &Module{
+func testModule() Module {
+	m := &NativeModule{
 		Name:    "test",
 		Version: "1.0.0",
 	}
@@ -58,7 +58,7 @@ func Test_Sum(t *testing.T) {
 }
 
 func Test_NotAFunc(t *testing.T) {
-	m := &Module{
+	m := &NativeModule{
 		Name:    "test",
 		Version: "1.0.0",
 	}
@@ -68,3 +68,25 @@ func Test_NotAFunc(t *testing.T) {
 	m.Unregister("hash")
 	assert.Equal(t, 0, len(m.funcs))
 }
+
+func Test_ScriptModule(t *testing.T) {
+
+	m, err := newScript("fixtures/module.lua")
+	assert.NoError(t, err)
+
+	s, err := newScript("fixtures/demo.lua", &ScriptModule{
+		Script:  m,
+		Name:    "demo_mod",
+		Version: "1.0.0",
+	})
+	assert.NoError(t, err)
+
+	out, err := s.Run(context.Background(), 10, m)
+	assert.NoError(t, err)
+	assert.Equal(t, TypeNumber, out.Type())
+	assert.Equal(t, Number(25), out.(Number))
+	assert.Equal(t, "25", out.String())
+
+	err = s.Close()
+	assert.NoError(t, err)
+}

script.go

@@ -15,7 +15,6 @@ import (
 	"github.com/kelindar/lua/json"
 	lua "github.com/yuin/gopher-lua"
 	"github.com/yuin/gopher-lua/parse"
-	//"layeh.com/gopher-json"
 	"layeh.com/gopher-luar"
 )
 
@@ -26,15 +25,16 @@ var (
 // Script represents a LUA script
 type Script struct {
 	lock sync.Mutex
-	name string         // The name of the script
-	argn int            // The number of arguments
-	exec *lua.LState    // The runtime for the script
-	main *lua.LFunction // The main function
-	mods []*Module      // The injected modules
+	name string             // The name of the script
+	argn int                // The number of arguments
+	exec *lua.LState        // The runtime for the script
+	main *lua.LFunction     // The main function
+	mods []Module           // The injected modules
+	code *lua.FunctionProto // The precompiled code
 }
 
 // FromReader reads a script fron an io.Reader
-func FromReader(name string, r io.Reader, modules ...*Module) (*Script, error) {
+func FromReader(name string, r io.Reader, modules ...Module) (*Script, error) {
 	script := &Script{
 		name: name,
 		mods: modules,
@@ -43,7 +43,7 @@ func FromReader(name string, r io.Reader, modules ...*Module) (*Script, error) {
 }
 
 // FromString reads a script fron a string
-func FromString(name, code string, modules ...*Module) (*Script, error) {
+func FromString(name, code string, modules ...Module) (*Script, error) {
 	return FromReader(name, bytes.NewBufferString(code), modules...)
 }
 
@@ -80,40 +80,38 @@ func (s *Script) Run(ctx context.Context, args ...interface{}) (Value, error) {
 }
 
 // Update updates the content of the script.
-func (s *Script) Update(r io.Reader) error {
+func (s *Script) Update(r io.Reader) (err error) {
 	runtime := newVM()
-	fn, err := s.compile(r)
+	s.code, err = s.compile(r)
 	if err != nil {
 		return err
 	}
-
 	// Push the function to the runtime
-	codeFn := runtime.NewFunctionFromProto(fn)
+	codeFn := runtime.NewFunctionFromProto(s.code)
 	runtime.Push(codeFn)
 
 	// Inject the modules
-	runtime.PreloadModule("json", json.Loader)
-	for _, m := range s.mods {
-		m.inject(runtime)
+	if err := s.loadModules(runtime); err != nil {
+		return err
 	}
 
 	// Initialize by calling the script
 	if err := runtime.PCall(0, lua.MultRet, nil); err != nil {
 		return err
 	}
 
-	// Get the main function
-	mainFn, err := findFunction(runtime, "main")
-	if err != nil {
-		return err
-	}
-
-	// Make sure the most recent code is present in the state
+	// Update the fields
 	s.lock.Lock()
 	defer s.lock.Unlock()
-	s.argn = int(mainFn.Proto.NumParameters)
 	s.exec = runtime
-	s.main = mainFn
+	s.argn = 0
+	s.main = nil
+
+	// If we have a main function, set it
+	if mainFn, err := findFunction(runtime, "main"); err == nil {
+		s.argn = int(mainFn.Proto.NumParameters)
+		s.main = mainFn
+	}
 	return nil
 }
 
@@ -129,6 +127,18 @@ func (s *Script) compile(r io.Reader) (*lua.FunctionProto, error) {
 	return lua.Compile(chunk, s.name)
 }
 
+// LoadModules loads in the prerequisite modules
+func (s *Script) loadModules(runtime *lua.LState) error {
+	runtime.PreloadModule("json", json.Loader)
+	for _, m := range s.mods {
+		if err := m.inject(runtime); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // Close closes the script and cleanly disposes of its resources.
 func (s *Script) Close() error {
 	s.lock.Lock()
@@ -140,15 +150,13 @@ func (s *Script) Close() error {
 
 // newVM creates a new LUA state
 func newVM() *lua.LState {
-	state := lua.NewState(lua.Options{
+	return lua.NewState(lua.Options{
 		RegistrySize:        1024 * 20, // this is the initial size of the registry
 		RegistryMaxSize:     1024 * 80, // this is the maximum size that the registry can grow to. If set to `0` (the default) then the registry will not auto grow
 		RegistryGrowStep:    32,        // this is how much to step up the registry by each time it runs out of space. The default is `32`.
 		CallStackSize:       120,       // this is the maximum callstack size of this LState
 		MinimizeStackMemory: true,      // Defaults to `false` if not specified. If set, the callstack will auto grow and shrink as needed up to a max of `CallStackSize`. If not set, the callstack will be fixed at `CallStackSize`.
 	})
-	//state.PreloadModule("relay", vm.mod.loadModule)
-	return state
 }
 
 // findFunction extracts a global function

script_test.go

@@ -11,9 +11,10 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func newScript(file string) (*Script, error) {
+func newScript(file string, mods ...Module) (*Script, error) {
 	f, _ := os.Open(file)
-	return FromReader("test.lua", f, testModule())
+	mods = append(mods, testModule())
+	return FromReader("test.lua", f, mods...)
 }
 
 type Person struct {
@@ -24,7 +25,6 @@ type Person struct {
 // Benchmark_Serial/fib-8         	 5870025	       203 ns/op	      16 B/op	       2 allocs/op
 // Benchmark_Serial/empty-8       	 8592448	       137 ns/op	       0 B/op	       0 allocs/op
 // Benchmark_Serial/update-8      	 1000000	      1069 ns/op	     224 B/op	      14 allocs/op
-// Benchmark_Serial/module-8      	 1900801	       629 ns/op	     160 B/op	       8 allocs/op
 func Benchmark_Serial(b *testing.B) {
 	b.Run("fib", func(b *testing.B) {
 		s, _ := newScript("fixtures/fib.lua")
@@ -157,23 +157,32 @@ func Test_NoMain(t *testing.T) {
 	}
 
 	{
-		_, err := FromString("", `main = 1`)
+		s, err := FromString("", `main = 1`)
+		assert.NoError(t, err)
+
+		_, err = s.Run(context.Background())
 		assert.Error(t, err)
 	}
 
 	{
-		_, err := FromString("", `
+		s, err := FromString("", `
 		function notmain()
 			local x = 1
 		end`)
+		assert.NoError(t, err)
+
+		_, err = s.Run(context.Background())
 		assert.Error(t, err)
 	}
 
 	{
-		_, err := FromString("", `
+		s, err := FromString("", `
 		function xxx()
 			local x = 1
 		end`)
+		assert.NoError(t, err)
+
+		_, err = s.Run(context.Background())
 		assert.Error(t, err)
 	}
 }