introduce state management

moonshark.go

@@ -5,9 +5,7 @@ import (
 	"os"
 	"path/filepath"
 
-	"Moonshark/modules"
+	"Moonshark/state"
-
-	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
 )
 
 func main() {
@@ -18,63 +16,17 @@ func main() {
 
 	scriptPath := os.Args[1]
 
-	// Check if file exists
+	// Create state configured for the script
-	if _, err := os.Stat(scriptPath); os.IsNotExist(err) {
+	luaState, err := state.NewFromScript(scriptPath)
-		fmt.Fprintf(os.Stderr, "Error: script file '%s' not found\n", scriptPath)
-		os.Exit(1)
-	}
-
-	// Initialize global registry
-	if err := modules.Initialize(); err != nil {
-		fmt.Fprintf(os.Stderr, "Error: failed to initialize registry: %v\n", err)
-		os.Exit(1)
-	}
-
-	// Create new Lua state with standard libraries
-	state := luajit.New()
-	if state == nil {
-		fmt.Fprintf(os.Stderr, "Error: failed to create Lua state\n")
-		os.Exit(1)
-	}
-	defer state.Close()
-
-	// Install module system in main state
-	if err := modules.Global.InstallInState(state); err != nil {
-		fmt.Fprintf(os.Stderr, "Error: failed to install module system: %v\n", err)
-		os.Exit(1)
-	}
-
-	// Get the absolute path to the script directory
-	scriptDir := filepath.Dir(scriptPath)
-	absScriptDir, err := filepath.Abs(scriptDir)
 	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error: failed to get absolute path: %v\n", err)
+		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
 		os.Exit(1)
 	}
+	defer luaState.Close()
 
-	// Add script directory to Lua's package.path
+	// Execute the script
-	packagePath := filepath.Join(absScriptDir, "?.lua")
+	if err := luaState.ExecuteFile(scriptPath); err != nil {
-	if err := state.AddPackagePath(packagePath); err != nil {
+		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
-		fmt.Fprintf(os.Stderr, "Warning: failed to add script directory to package.path: %v\n", err)
-	}
-
-	// Read script file
-	scriptContent, err := os.ReadFile(scriptPath)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error reading script file '%s': %v\n", scriptPath, err)
-		os.Exit(1)
-	}
-
-	// Compile script to bytecode
-	bytecode, err := state.CompileBytecode(string(scriptContent), scriptPath)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error compiling '%s': %v\n", scriptPath, err)
-		os.Exit(1)
-	}
-
-	// Execute the compiled bytecode
-	if err := state.LoadAndRunBytecode(bytecode, scriptPath); err != nil {
-		fmt.Fprintf(os.Stderr, "Error executing '%s': %v\n", scriptPath, err)
 		os.Exit(1)
 	}
 }

state/state.go

@@ -0,0 +1,266 @@
+package state
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"Moonshark/modules"
+
+	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
+)
+
+// State wraps luajit.State with enhanced functionality
+type State struct {
+	*luajit.State
+	initialized bool
+	isWorker    bool
+}
+
+// Config holds state initialization options
+type Config struct {
+	OpenLibs       bool
+	InstallModules bool
+	ScriptDir      string
+	IsWorker       bool
+}
+
+// DefaultConfig returns default configuration
+func DefaultConfig() Config {
+	return Config{
+		OpenLibs:       true,
+		InstallModules: true,
+		IsWorker:       false,
+	}
+}
+
+// New creates a new enhanced Lua state
+func New(config ...Config) (*State, error) {
+	cfg := DefaultConfig()
+	if len(config) > 0 {
+		cfg = config[0]
+	}
+
+	// Create base state
+	baseState := luajit.New(cfg.OpenLibs)
+	if baseState == nil {
+		return nil, fmt.Errorf("failed to create Lua state")
+	}
+
+	state := &State{
+		State:    baseState,
+		isWorker: cfg.IsWorker,
+	}
+
+	// Set worker global flag if this is a worker state
+	if cfg.IsWorker {
+		state.PushBoolean(true)
+		state.SetGlobal("__IS_WORKER")
+	}
+
+	// Install module system if requested
+	if cfg.InstallModules {
+		if err := state.initializeModules(); err != nil {
+			state.Close()
+			return nil, fmt.Errorf("failed to install modules: %w", err)
+		}
+	}
+
+	// Set script directory if provided
+	if cfg.ScriptDir != "" {
+		if err := state.SetScriptDirectory(cfg.ScriptDir); err != nil {
+			state.Close()
+			return nil, fmt.Errorf("failed to set script directory: %w", err)
+		}
+	}
+
+	state.initialized = true
+	return state, nil
+}
+
+// NewFromScript creates a state configured for a specific script file
+func NewFromScript(scriptPath string, config ...Config) (*State, error) {
+	cfg := DefaultConfig()
+	if len(config) > 0 {
+		cfg = config[0]
+	}
+
+	// Set script directory from file path
+	scriptDir := filepath.Dir(scriptPath)
+	absScriptDir, err := filepath.Abs(scriptDir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get absolute path: %w", err)
+	}
+	cfg.ScriptDir = absScriptDir
+
+	return New(cfg)
+}
+
+// NewWorker creates a new worker state with __IS_WORKER flag set
+func NewWorker(config ...Config) (*State, error) {
+	cfg := DefaultConfig()
+	if len(config) > 0 {
+		cfg = config[0]
+	}
+	cfg.IsWorker = true
+
+	return New(cfg)
+}
+
+// NewWorkerFromScript creates a worker state configured for a specific script
+func NewWorkerFromScript(scriptPath string, config ...Config) (*State, error) {
+	cfg := DefaultConfig()
+	if len(config) > 0 {
+		cfg = config[0]
+	}
+	cfg.IsWorker = true
+
+	scriptDir := filepath.Dir(scriptPath)
+	absScriptDir, err := filepath.Abs(scriptDir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get absolute path: %w", err)
+	}
+	cfg.ScriptDir = absScriptDir
+
+	return New(cfg)
+}
+
+// initializeModules sets up the module system
+func (s *State) initializeModules() error {
+	// Initialize global registry if needed
+	if modules.Global == nil {
+		if err := modules.Initialize(); err != nil {
+			return err
+		}
+	}
+
+	return modules.Global.InstallInState(s.State)
+}
+
+// SetScriptDirectory adds a directory to Lua's package.path
+func (s *State) SetScriptDirectory(dir string) error {
+	packagePath := filepath.Join(dir, "?.lua")
+	return s.AddPackagePath(packagePath)
+}
+
+// ExecuteFile compiles and runs a Lua script file
+func (s *State) ExecuteFile(scriptPath string) error {
+	// Check if file exists
+	if _, err := os.Stat(scriptPath); os.IsNotExist(err) {
+		return fmt.Errorf("script file '%s' not found", scriptPath)
+	}
+
+	// Read script content
+	scriptContent, err := os.ReadFile(scriptPath)
+	if err != nil {
+		return fmt.Errorf("failed to read script file '%s': %w", scriptPath, err)
+	}
+
+	return s.ExecuteString(string(scriptContent), scriptPath)
+}
+
+// ExecuteString compiles and runs Lua code
+func (s *State) ExecuteString(code, name string) error {
+	// Use bytecode compilation for better performance
+	bytecode, err := s.CompileBytecode(code, name)
+	if err != nil {
+		return fmt.Errorf("compilation error in '%s': %w", name, err)
+	}
+
+	if err := s.LoadAndRunBytecode(bytecode, name); err != nil {
+		return fmt.Errorf("execution error in '%s': %w", name, err)
+	}
+
+	return nil
+}
+
+// ExecuteFileWithResults executes a script and returns results
+func (s *State) ExecuteFileWithResults(scriptPath string) ([]any, error) {
+	if _, err := os.Stat(scriptPath); os.IsNotExist(err) {
+		return nil, fmt.Errorf("script file '%s' not found", scriptPath)
+	}
+
+	scriptContent, err := os.ReadFile(scriptPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read script file '%s': %w", scriptPath, err)
+	}
+
+	return s.ExecuteStringWithResults(string(scriptContent), scriptPath)
+}
+
+// ExecuteStringWithResults executes code and returns all results
+func (s *State) ExecuteStringWithResults(code, name string) ([]any, error) {
+	baseTop := s.GetTop()
+	defer s.SetTop(baseTop)
+
+	nresults, err := s.Execute(code)
+	if err != nil {
+		return nil, fmt.Errorf("execution error in '%s': %w", name, err)
+	}
+
+	if nresults == 0 {
+		return nil, nil
+	}
+
+	results := make([]any, nresults)
+	for i := 0; i < nresults; i++ {
+		val, err := s.ToValue(baseTop + i + 1)
+		if err != nil {
+			results[i] = nil
+		} else {
+			results[i] = val
+		}
+	}
+
+	return results, nil
+}
+
+// IsInitialized returns whether the state has been properly initialized
+func (s *State) IsInitialized() bool {
+	return s.initialized
+}
+
+// IsWorker returns whether this state is a worker state
+func (s *State) IsWorker() bool {
+	return s.isWorker
+}
+
+// Close cleans up the state and releases resources
+func (s *State) Close() {
+	if s.State != nil {
+		s.Cleanup()
+		s.State.Close()
+		s.State = nil
+	}
+	s.initialized = false
+}
+
+// Quick creates a minimal state for one-off script execution
+func Quick() (*State, error) {
+	return New(Config{
+		OpenLibs:       true,
+		InstallModules: true,
+	})
+}
+
+// QuickExecute creates a state, executes code, and cleans up
+func QuickExecute(code, name string) error {
+	state, err := Quick()
+	if err != nil {
+		return err
+	}
+	defer state.Close()
+
+	return state.ExecuteString(code, name)
+}
+
+// QuickExecuteWithResults executes code and returns results
+func QuickExecuteWithResults(code, name string) ([]any, error) {
+	state, err := Quick()
+	if err != nil {
+		return nil, err
+	}
+	defer state.Close()
+
+	return state.ExecuteStringWithResults(code, name)
+}