scope_stack.go

package glisp

import (
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
)

// Scope defines a single layer of lexical bindings, mapping symbol numbers to their Sexp values.
type Scope map[int]Sexp

// ScopeStack owns the layers above shared. The shared layer and its ancestors
// may be held by closures or other environments and are managed by Go's GC.
// Only the exclusive prefix is eligible for pooling.
type ScopeStack struct {
	top, bottom *ScopeLayer
	shared      atomic.Pointer[ScopeLayer]
	callBase    *ScopeLayer
}

// ScopeLayer stores one set of lexical bindings. Its parent link is immutable
// while shared; set! may update bindings under mu.
type ScopeLayer struct {
	Scope  Scope
	layout *scopeLayout
	locals []localValue
	mu     sync.RWMutex
	// overlay layers belong to the same lexical frame as their parent.
	overlay bool
	next    *ScopeLayer
}

// NewScopeStack creates and returns an empty ScopeStack.
func NewScopeStack() *ScopeStack {
	return &ScopeStack{}
}

// newScope retrieves a Scope map from the object pool.
func newScope() Scope { return scopePool.Get().(Scope) }

// newScopeLayer retrieves a ScopeLayer from the object pool and initializes it
// with a new, empty Scope.
func newScopeLayer() *ScopeLayer {
	return newScopeLayerWith(newScope())
}

// newScopeLayerWith retrieves a ScopeLayer from the object pool and initializes it
// with the provided Scope.
func newScopeLayerWith(s Scope) *ScopeLayer {
	if s == nil {
		s = newScope()
	}
	layer := scopeLayerPool.Get().(*ScopeLayer)
	layer.Scope = s
	return layer
}

// IsStackElem is a marker method for the StackElem interface.
func (s *ScopeLayer) IsStackElem() {}

// Clone creates a deep copy of the ScopeLayer and its underlying Scope map.
func (s *ScopeLayer) Clone() *ScopeLayer {
	newScope := newScopeLayer()
	s.mu.RLock()
	defer s.mu.RUnlock()
	for k, v := range s.Scope {
		newScope.Scope[k] = v
	}
	newScope.layout = s.layout
	newScope.locals = append(newScope.locals[:0], s.locals...)
	return newScope
}

// Find searches for a symbol's value within this specific scope layer.
func (s *ScopeLayer) Find(n int) (Sexp, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.layout != nil {
		if i, ok := s.layout.index[n]; ok {
			local := s.locals[i]
			return local.value, local.bound
		}
	}
	if expr, ok := s.Scope[n]; ok {
		return expr, true
	}
	return SexpNull, false
}

// Bind sets the value for a symbol in this specific scope layer.
func (s *ScopeLayer) Bind(n int, e Sexp) {
	s.mu.Lock()
	if s.layout != nil {
		if i, ok := s.layout.index[n]; ok {
			s.locals[i] = localValue{value: e, bound: true}
			s.mu.Unlock()
			return
		}
	}
	s.Scope[n] = e
	s.mu.Unlock()
}

// Fork captures the current lexical environment without traversing its parents.
func (stack *ScopeStack) Fork() *ScopeStack {
	stack.shared.Store(stack.top)
	st := &ScopeStack{top: stack.top, bottom: stack.bottom}
	st.shared.Store(st.top)
	return st
}

// ForkBottom shares the global environment. Global layers are always GC-owned.
func (stack *ScopeStack) ForkBottom() *ScopeStack {
	st := &ScopeStack{top: stack.bottom, bottom: stack.bottom}
	st.shared.Store(st.top)
	return st
}

// Only transient call-stack wrappers are pooled, never captured wrappers.
var callScopePool = sync.Pool{New: func() any { return &ScopeStack{} }}

func (stack *ScopeStack) forkCall(capture bool) *ScopeStack {
	frame := callScopePool.Get().(*ScopeStack)
	frame.bottom = stack.bottom
	frame.top = stack.bottom
	if capture {
		stack.shared.Store(stack.top)
		frame.top = stack.top
	}
	frame.shared.Store(frame.top)
	frame.callBase = frame.top
	return frame
}

// IsStackElem is a marker method for the StackElem interface.
func (stack *ScopeStack) IsStackElem() {}

// Clone creates a full, deep copy of the ScopeStack, including all its layers.
// This is a more expensive operation than Fork and is used when a completely
// independent environment is required.
func (stack *ScopeStack) Clone() *ScopeStack {
	return newCloneContext().scope(stack)
}

// push adds a new layer to the top of the stack.
func (stack *ScopeStack) push(layer *ScopeLayer) {
	if layer == nil {
		return
	}
	if stack.top != nil {
		layer.next = stack.top
	} else {
		stack.bottom = layer
		stack.shared.Store(layer)
	}
	stack.top = layer
}

// Pop releases an exclusive layer or drops a reference to a captured layer.
func (stack *ScopeStack) Pop() error {
	if stack.top == nil {
		return errors.New("pop from empty scope stack")
	}
	cur := stack.top
	stack.top = cur.next
	if stack.top == nil {
		stack.bottom = nil
	}
	if cur == stack.shared.Load() {
		stack.shared.Store(stack.top)
	} else {
		recycleScopeLayer(cur)
	}
	return nil
}

// PushScope creates a new, empty scope and pushes it onto the stack.
// This is used for constructs like `let`.
func (stack *ScopeStack) PushScope() {
	stack.Push()
}

// PopScope removes the top scope from the stack.
func (stack *ScopeStack) PopScope() error {
	for stack.top != nil {
		overlay := stack.top.overlay
		if err := stack.Pop(); err != nil {
			return err
		}
		if !overlay {
			return nil
		}
	}
	return errors.New("pop from empty scope stack")
}

// Push adds one or more pre-existing scopes to the top of the stack.
func (stack *ScopeStack) Push(scopes ...Scope) {
	if len(scopes) == 0 {
		layer := newScopeLayer()
		stack.push(layer)
	} else {
		for i := range scopes {
			stack.push(newScopeLayerWith(scopes[i]))
		}
	}
}

// LookupSymbol searches for a symbol starting from the top of the stack and
// moving down through parent scopes until the symbol is found.
func (stack *ScopeStack) LookupSymbol(sym SexpSymbol) (Sexp, error) {
	for ptr := stack.top; ptr != nil; {
		if expr, ok := ptr.Find(sym.number); ok {
			return expr, nil
		}
		ptr = ptr.next
	}
	return SexpNull, fmt.Errorf("symbol `%v` not found", sym.Name())
}

// GlobalFuntions returns a list of function names found in the global scope (bottom layer).
func (stack *ScopeStack) GlobalFuntions() (ret []string) {
	if stack.IsEmpty() {
		return nil
	}
	stack.bottom.mu.RLock()
	defer stack.bottom.mu.RUnlock()
	for _, v := range stack.bottom.Scope {
		if fn, ok := v.(*SexpFunction); ok {
			ret = append(ret, fn.name)
		}
	}
	return
}

// IsEmpty returns true if the stack has no layers.
func (stack *ScopeStack) IsEmpty() bool {
	return stack.top == nil
}

type BindSymbolOption int

const (
	BIND_DEFAULT BindSymbolOption = iota
	BIND_GLOBAL
)

// BindSymbol adds a binding. A captured lexical layer receives a new overlay
// so that def does not replace the bindings seen by existing closures.
func (stack *ScopeStack) BindSymbol(sym SexpSymbol, expr Sexp, options ...BindSymbolOption) error {
	if stack.IsEmpty() {
		return errors.New("no scope available")
	}
	if len(options) > 0 && options[0] == BIND_GLOBAL {
		stack.bottom.Bind(sym.number, expr)
	} else {
		if stack.top != stack.bottom && stack.top == stack.shared.Load() {
			layer := newScopeLayer()
			layer.overlay = true
			stack.push(layer)
		}
		stack.top.Bind(sym.number, expr)
	}
	return nil
}

// SetSymbol updates an existing binding for a symbol. This is the implementation for `set!`.
// It searches up the scope stack for the first occurrence of the symbol and
// modifies its value in-place.
//
// NOTE: By design, this can mutate a shared parent scope. This behavior is
// intentional to mimic `set!` in many Lisp dialects.
func (stack *ScopeStack) SetSymbol(sym SexpSymbol, expr Sexp) error {
	if stack.IsEmpty() {
		return errors.New("no scope available")
	}
	for ptr := stack.top; ptr != nil; {
		if _, ok := ptr.Find(sym.number); ok {
			ptr.Bind(sym.number, expr)
			return nil
		}
		ptr = ptr.next
	}
	return stack.BindSymbol(sym, expr)
}

// releaseOwned recycles only the exclusive prefix. Captured layers retain
// their bindings and parent links for any surviving closures.
func (stack *ScopeStack) releaseOwned() {
	boundary := stack.shared.Load()
	for stack.top != nil && stack.top != boundary {
		cur := stack.top
		stack.top = cur.next
		recycleScopeLayer(cur)
	}
}

// Clear releases this stack's ownership without traversing captured ancestors.
func (stack *ScopeStack) Clear() {
	stack.releaseOwned()
	stack.top = nil
	stack.bottom = nil
	stack.shared.Store(nil)
	stack.callBase = nil
}

// resetCall starts the next tail-call iteration with a fresh parameter frame.
func (stack *ScopeStack) resetCall(layout *scopeLayout) {
	// An uncaptured parameter frame can be cleared in place for the next call.
	if top := stack.top; top != nil && top != stack.shared.Load() && top.next == stack.callBase && !top.overlay && top.layout == layout {
		clear(top.Scope)
		clear(top.locals)
		return
	}
	stack.releaseOwned()
	stack.top = stack.callBase
	stack.shared.Store(stack.callBase)
	stack.pushLayout(layout)
}

// scopePool holds maps for reuse.
var scopePool = sync.Pool{
	New: func() any { return make(Scope) },
}

// scopeLayerPool holds ScopeLayer structs for reuse.
var scopeLayerPool = sync.Pool{
	New: func() any {
		return &ScopeLayer{}
	},
}

// recycleScope clears a scope map and returns it to the pool.
func recycleScope(s Scope) {
	if s != nil {
		clear(s)
		scopePool.Put(s)
	}
}

// recycleScopeLayer is used only for layers exclusively owned by a stack.
func recycleScopeLayer(layer *ScopeLayer) {
	if layer == nil {
		return
	}
	recycleScope(layer.Scope)
	layer.Scope = nil
	layer.next = nil
	layer.overlay = false
	layer.layout = nil
	clear(layer.locals)
	layer.locals = layer.locals[:0]
	scopeLayerPool.Put(layer)
}