five/hbrt/vm.go

// Copyright (c) 2026 Charles KWON OhJun (charleskwonohjun@gmail.com)
// All rights reserved.

package hbrt

import (
	"fmt"
	"os"
	"runtime/pprof"
	"sync"
)

// VM is the shared state across all threads.
type VM struct {
	mu        sync.RWMutex
	modules   []*Module
	symbols   map[string]*Symbol
	statics   map[string][]Value
	threads   []*Thread           // all threads created (for shutdown cleanup + debugger listing)
	nextTID   uint32              // monotonic thread id
	waFactory func() interface{} // creates WorkAreaManager for new threads
	onExit    func()             // called when Run() finishes (restore terminal etc.)
	Debugger  *Debugger          // nil = no debugging; set by five debug command
}

// SetWAFactory sets the factory for creating WorkAreaManagers.
func (vm *VM) SetWAFactory(f func() interface{}) {
	vm.waFactory = f
}

// SetOnExit sets a callback for when Run() finishes.
func (vm *VM) SetOnExit(f func()) {
	vm.onExit = f
}

// Library modules registered via init() — protected by mutex for FRB concurrent loading.
var (
	libModules     []*Module
	dynamicFuncs   []Symbol // from HB_FUNC() in #pragma BEGINDUMP
	libRegistryMu  sync.Mutex
)

// RegisterLibModule registers a module from a library PRG file.
// Called by init() in generated library code.
func RegisterLibModule(m *Module) {
	libRegistryMu.Lock()
	libModules = append(libModules, m)
	libRegistryMu.Unlock()
}

// RegisterDynamicFunc registers a Go function callable from PRG.
// Called from init() in #pragma BEGINDUMP code via HB_FUNC().
func RegisterDynamicFunc(name string, fn func(*Thread)) {
	libRegistryMu.Lock()
	dynamicFuncs = append(dynamicFuncs, Symbol{
		Name:  name,
		Scope: FsPublic | FsLocal,
		Func:  fn,
	})
	libRegistryMu.Unlock()
}

// RegisterLibModules registers any pending lib modules and dynamic functions.
func (vm *VM) RegisterLibModules() {
	libRegistryMu.Lock()
	mods := libModules
	libModules = nil
	dyns := dynamicFuncs
	dynamicFuncs = nil
	libRegistryMu.Unlock()

	for _, m := range mods {
		vm.RegisterModule(m)
	}
	for i := range dyns {
		sym := &dyns[i]
		vm.RegisterSymbol(sym)
	}
	dynamicFuncs = nil
}

// NewVM creates a new VM instance.
func NewVM() *VM {
	return &VM{
		modules: make([]*Module, 0),
		symbols: make(map[string]*Symbol),
		statics: make(map[string][]Value),
	}
}

// RegisterModule registers a module's symbols with the VM.
func (vm *VM) RegisterModule(m *Module) {
	vm.mu.Lock()
	defer vm.mu.Unlock()
	vm.modules = append(vm.modules, m)
	for i := range m.Symbols {
		sym := &m.Symbols[i]
		vm.symbols[sym.Name] = sym
	}
}

// RegisterSymbol registers a single symbol.
func (vm *VM) RegisterSymbol(sym *Symbol) {
	vm.mu.Lock()
	defer vm.mu.Unlock()
	vm.symbols[sym.Name] = sym
}

// UnregisterSymbol removes a symbol by name. Returns the old symbol if any.
func (vm *VM) UnregisterSymbol(name string) *Symbol {
	vm.mu.Lock()
	defer vm.mu.Unlock()
	old := vm.symbols[name]
	delete(vm.symbols, name)
	return old
}

// SymbolNames returns all registered symbol names.
func (vm *VM) SymbolNames() []string {
	vm.mu.RLock()
	defer vm.mu.RUnlock()
	names := make([]string, 0, len(vm.symbols))
	for n := range vm.symbols {
		names = append(names, n)
	}
	return names
}

// FindSymbol looks up a symbol by name.
func (vm *VM) FindSymbol(name string) *Symbol {
	vm.mu.RLock()
	defer vm.mu.RUnlock()
	return vm.symbols[name]
}

// GetSym returns the cached Symbol, performing a one-time FindSymbol
// lookup on first access and stashing the pointer in *cache for all
// subsequent calls. Generated code (gengo) declares a package-level
// `var _sym_NAME *Symbol` per unique call target and routes every
// PushSymbol through this helper so the hot path becomes a single
// non-nil check instead of vm.symbols map + RWMutex per invocation.
func (t *Thread) GetSym(cache **Symbol, name string) *Symbol {
	if s := *cache; s != nil {
		return s
	}
	s := t.vm.FindSymbol(name)
	if s != nil {
		// Only cache successful resolutions — nil might be due to
		// init-order (another module's registrations pending);
		// retry on next call once those complete.
		*cache = s
	}
	return s
}

// NewThread creates a new Thread attached to this VM.
func (vm *VM) NewThread() *Thread {
	t := NewThread(vm)
	vm.mu.Lock()
	vm.nextTID++
	t.tid = vm.nextTID
	vm.threads = append(vm.threads, t)
	vm.mu.Unlock()
	return t
}

// Threads returns a snapshot of all threads currently tracked by the
// VM. Used by the debugger's `threads` command. Returned slice is a
// copy — callers can iterate without holding any lock.
func (vm *VM) Threads() []*Thread {
	vm.mu.RLock()
	defer vm.mu.RUnlock()
	out := make([]*Thread, len(vm.threads))
	copy(out, vm.threads)
	return out
}

// TID returns this thread's VM-unique id. Main thread gets 1.
func (t *Thread) TID() uint32 { return t.tid }

// Run starts execution from the named function.
func (vm *VM) Run(funcName string) Value {
	// Register any library modules from init()
	for _, m := range libModules {
		vm.RegisterModule(m)
	}
	libModules = nil

	sym := vm.FindSymbol(funcName)
	if sym == nil {
		panic("function not found: " + funcName)
	}
	if sym.Func == nil {
		panic("function has no implementation: " + funcName)
	}

	t := vm.NewThread()

	// Auto-initialize WorkAreaManager if not set
	if t.WA == nil && vm.waFactory != nil {
		t.WA = vm.waFactory()
	}

	// Copy statics to thread
	vm.mu.RLock()
	for k, v := range vm.statics {
		t.statics[k] = v
	}
	vm.mu.RUnlock()

	// Install signal handlers for clean shutdown
	vm.InstallSignalHandlers()

	// Optional CPU profiling — FIVE_CPUPROFILE=<path> writes a pprof
	// file covering the whole program run. Used to collect default.pgo
	// input for profile-guided compilation of Five-runtime code.
	if path := os.Getenv("FIVE_CPUPROFILE"); path != "" {
		if f, err := os.Create(path); err == nil {
			if werr := pprof.StartCPUProfile(f); werr == nil {
				defer f.Close()
				defer pprof.StopCPUProfile()
				defer fmt.Fprintf(os.Stderr, "CPU profile written to %s\n", path)
			} else {
				fmt.Fprintf(os.Stderr, "FIVE_CPUPROFILE: StartCPUProfile: %v\n", werr)
				f.Close()
			}
		} else {
			fmt.Fprintf(os.Stderr, "FIVE_CPUPROFILE: cannot create %s: %v\n", path, err)
		}
	}

	// Call the function, ensure full shutdown on exit.
	// On unhandled *HbError, route through DefaultErrorHook (writes
	// error.log) before letting the panic surface. The Go panic still
	// propagates — we only add the diagnostic side effect.
	defer vm.Shutdown()
	defer func() {
		if r := recover(); r != nil {
			if DefaultErrorHook != nil {
				DefaultErrorHook(t, r)
			}
			panic(r)
		}
	}()
	// Attach the symbol so the entry frame shows its name in stack traces.
	// Normal calls go through Function() which sets pendingCallSym; direct
	// VM.Run needs to do it manually.
	t.pendingCallSym = sym
	sym.Func(t)

	return t.retVal
}

// DefaultErrorHook runs when an unhandled panic escapes Main. hbrtl sets
// this at package init to dump error.log. Nil by default — set once at
// startup, not swapped at runtime, so no synchronization.
var DefaultErrorHook func(t *Thread, panicValue interface{})

// DebugDiagnosticHook renders the error.log-style state dump (workareas,
// SET flags, runtime memory) for the debugger's `diag` command. hbrtl
// sets this at init time — keeping the renderers in hbrtl avoids a
// circular import (hbrdd → hbrt ← hbrt needs hbrdd types).
//
// section values: "" (everything), "wa", "set", "mem". Unknown sections
// fall back to everything. The hook writes one line per call to `emit`.
var DebugDiagnosticHook func(t *Thread, section string, emit func(string))