xgo/compiler/symbol_table.go

package compiler

type SymbolTable struct {
	parent        *SymbolTable
	block         bool
	store         map[string]Symbol
	numDefinition int
	maxDefinition int
	freeSymbols   []Symbol
}

func NewSymbolTable() *SymbolTable {
	return &SymbolTable{
		store: make(map[string]Symbol),
	}
}

func (t *SymbolTable) Define(name string) Symbol {
	symbol := Symbol{Name: name, Index: t.nextIndex()}
	t.numDefinition++

	if t.Parent(true) == nil {
		symbol.Scope = ScopeGlobal
	} else {
		symbol.Scope = ScopeLocal
	}

	t.store[name] = symbol

	t.updateMaxDefs(symbol.Index + 1)

	return symbol
}

func (t *SymbolTable) DefineBuiltin(index int, name string) Symbol {
	symbol := Symbol{
		Name:  name,
		Index: index,
		Scope: ScopeBuiltin,
	}

	t.store[name] = symbol

	return symbol
}

func (t *SymbolTable) Resolve(name string) (symbol Symbol, depth int, ok bool) {
	symbol, ok = t.store[name]
	if !ok && t.parent != nil {
		symbol, depth, ok = t.parent.Resolve(name)
		if !ok {
			return
		}

		if !t.block {
			depth += 1
		}

		// if symbol is defined in parent table and if it's not global/builtin
		// then it's free variable.
		if depth > 0 && symbol.Scope != ScopeGlobal && symbol.Scope != ScopeBuiltin {
			return t.defineFree(symbol), depth, true
		}

		return
	}

	return
}

func (t *SymbolTable) Fork(block bool) *SymbolTable {
	return &SymbolTable{
		store:  make(map[string]Symbol),
		parent: t,
		block:  block,
	}
}

func (t *SymbolTable) Parent(skipBlock bool) *SymbolTable {
	if skipBlock && t.block {
		return t.parent.Parent(skipBlock)
	}

	return t.parent
}

func (t *SymbolTable) MaxSymbols() int {
	return t.maxDefinition
}

func (t *SymbolTable) FreeSymbols() []Symbol {
	return t.freeSymbols
}

func (t *SymbolTable) Names() []string {
	var names []string
	for name := range t.store {
		names = append(names, name)
	}
	return names
}

func (t *SymbolTable) nextIndex() int {
	if t.block {
		return t.parent.nextIndex() + t.numDefinition
	}

	return t.numDefinition
}

func (t *SymbolTable) updateMaxDefs(numDefs int) {
	if numDefs > t.maxDefinition {
		t.maxDefinition = numDefs
	}

	if t.block {
		t.parent.updateMaxDefs(numDefs)
	}
}

func (t *SymbolTable) defineFree(original Symbol) Symbol {
	// TODO: should we check duplicates?

	t.freeSymbols = append(t.freeSymbols, original)

	symbol := Symbol{
		Name:  original.Name,
		Index: len(t.freeSymbols) - 1,
		Scope: ScopeFree,
	}

	t.store[original.Name] = symbol

	return symbol
}