Extract scopes and scope chain from compiler

3 files changed, 306 insertions, 0 deletions

diff --git a/pkg/lang/compiler/scope/scope_chain.go b/pkg/lang/compiler/scope/scope_chain.go
new file mode 100644
index 0000000..683d53a
--- /dev/null
+++ b/pkg/lang/compiler/scope/scope_chain.go
@@ -0,0 +1,191 @@
+package scope
+
+import (
+	"fmt"
+	"jinx/pkg/lang/vm/code"
+)
+
+type ScopeChain struct {
+	nameToSymbol map[string]SymbolID
+
+	symbolScopes   []SymbolScope // All other scopes are bound to this by ID.
+	functionScopes []FunctionScope
+}
+
+func NewScopeChain() ScopeChain {
+	symbolScopes := make([]SymbolScope, 1)
+	symbolScopes[0] = NewSymbolScope()
+
+	functionScopes := make([]FunctionScope, 1)
+	functionScopes[0] = NewFunctionScope(0, "") // Root function to house top-scope sub units
+
+	return ScopeChain{
+		nameToSymbol: make(map[string]SymbolID),
+
+		symbolScopes:   symbolScopes,
+		functionScopes: functionScopes,
+	}
+}
+
+func (sc *ScopeChain) CurrentScopeID() ScopeID {
+	return ScopeID(len(sc.symbolScopes) - 1)
+}
+
+func (sc *ScopeChain) IsRootScope() bool {
+	return sc.CurrentScopeID() == ScopeID(0)
+}
+
+func (sc *ScopeChain) Current() *SymbolScope {
+	if len(sc.functionScopes) == 0 {
+		panic("root scope is missing")
+	}
+
+	return &sc.symbolScopes[sc.CurrentScopeID()]
+}
+
+func (sc *ScopeChain) CurrentFunction() *FunctionScope {
+	if len(sc.functionScopes) == 0 {
+		panic("root function scope is missing")
+	}
+
+	return &sc.functionScopes[len(sc.functionScopes)-1]
+}
+
+func (sc *ScopeChain) Enter() {
+	sc.symbolScopes = append(sc.symbolScopes, NewSymbolScope())
+}
+
+func (sc *ScopeChain) EnterFunction(unit code.Marker) {
+	sc.Enter()
+	sc.functionScopes = append(sc.functionScopes, NewFunctionScope(sc.CurrentScopeID(), unit))
+}
+
+func (sc *ScopeChain) Exit() {
+	if sc.CurrentScopeID() == 0 {
+		return
+	}
+
+	sc.symbolScopes[sc.CurrentScopeID()] = SymbolScope{}
+	sc.symbolScopes = sc.symbolScopes[:sc.CurrentScopeID()]
+}
+
+func (sc *ScopeChain) Declare(name string) (int, bool) {
+	// Check whether the symbol is already declared in any of the scopes.
+	if _, ok := sc.nameToSymbol[name]; ok {
+		return 0, false
+	}
+
+	current := sc.Current()
+	indexInScope := len(current.variableSymbols)
+
+	symbolID := SymbolID{
+		symbolKind:   SymbolKindVariable,
+		scopeID:      sc.CurrentScopeID(),
+		indexInScope: indexInScope,
+	}
+
+	// Declare the symbol in the current scope.
+	current.variableSymbols = append(current.variableSymbols, Symbol[SymbolVariable]{
+		name: name,
+		data: SymbolVariable{
+			localIndex: indexInScope,
+		},
+	})
+
+	sc.nameToSymbol[name] = symbolID
+
+	return indexInScope, true
+}
+
+func (sc *ScopeChain) DeclareFunction(name string, args uint) (code.Marker, bool) {
+	if _, ok := sc.nameToSymbol[name]; ok {
+		return "", false
+	}
+
+	current := sc.Current()
+	index := len(current.functionSymbols)
+
+	symbolID := SymbolID{
+		symbolKind:   SymbolKindFunction,
+		scopeID:      sc.CurrentScopeID(),
+		indexInScope: index,
+	}
+
+	unitName := sc.CreateFunctionSubUnit(name)
+
+	current.functionSymbols = append(current.functionSymbols, Symbol[SymbolFunction]{
+		name: name,
+		data: SymbolFunction{
+			marker: unitName,
+			args:   args,
+		},
+	})
+
+	sc.nameToSymbol[name] = symbolID
+
+	return unitName, true
+}
+
+func (sc *ScopeChain) DeclareAnonymous() int {
+	current := sc.Current()
+	index := len(current.variableSymbols)
+
+	current.variableSymbols = append(current.variableSymbols, Symbol[SymbolVariable]{
+		name: "", // An anonymous symbol has no name.
+		data: SymbolVariable{
+			localIndex: index,
+		},
+	})
+
+	return index
+}
+
+func (sc *ScopeChain) DeclareTemporary() int {
+	return len(sc.Current().variableSymbols) // :)
+}
+
+func (sc *ScopeChain) CreateAnonymousFunctionSubUnit() code.Marker {
+	fnScope := sc.CurrentFunction()
+
+	index := fnScope.subUnitCount
+	fnScope.subUnitCount++
+
+	return sc.CreateFunctionSubUnit(fmt.Sprintf("anon_%d", index))
+}
+
+func (sc *ScopeChain) CreateFunctionSubUnit(subUnitName string) code.Marker {
+	fnScope := sc.CurrentFunction()
+
+	name := fnScope.unit
+	if name == "" {
+		name = code.Marker(subUnitName)
+	} else {
+		name = name.SubUnit(subUnitName)
+	}
+
+	return name
+}
+
+func (sc *ScopeChain) Lookup(name string) (SymbolID, bool) {
+	if id, ok := sc.nameToSymbol[name]; ok {
+		return id, true
+	}
+
+	return SymbolID{}, false
+}
+
+func (sc *ScopeChain) GetVariable(id SymbolID) Symbol[SymbolVariable] {
+	if id.symbolKind != SymbolKindVariable {
+		panic("incorrect symbol id kind given")
+	}
+
+	return sc.symbolScopes[id.scopeID].variableSymbols[id.indexInScope]
+}
+
+func (sc *ScopeChain) GetFunction(id SymbolID) Symbol[SymbolFunction] {
+	if id.symbolKind != SymbolKindFunction {
+		panic("incorrect symbol id kind given")
+	}
+
+	return sc.symbolScopes[id.scopeID].functionSymbols[id.indexInScope]
+}
diff --git a/pkg/lang/compiler/scope/scopes.go b/pkg/lang/compiler/scope/scopes.go
new file mode 100644
index 0000000..39e48ef
--- /dev/null
+++ b/pkg/lang/compiler/scope/scopes.go
@@ -0,0 +1,50 @@
+package scope
+
+import "jinx/pkg/lang/vm/code"
+
+type ScopeID int
+
+type ScopeKind int
+
+const (
+	ScopeKindNormal ScopeKind = iota
+	ScopeKindFunction
+)
+
+type SymbolScope struct {
+	variableSymbols []Symbol[SymbolVariable]
+	functionSymbols []Symbol[SymbolFunction]
+}
+
+func NewSymbolScope() SymbolScope {
+	return SymbolScope{
+		variableSymbols: make([]Symbol[SymbolVariable], 0),
+		functionSymbols: make([]Symbol[SymbolFunction], 0),
+	}
+}
+
+type FunctionScope struct {
+	id           ScopeID
+	unit         code.Marker
+	subUnitCount int
+}
+
+func NewFunctionScope(id ScopeID, unit code.Marker) FunctionScope {
+	return FunctionScope{
+		id:           id,
+		unit:         unit,
+		subUnitCount: 0,
+	}
+}
+
+func (sf FunctionScope) ID() ScopeID {
+	return sf.id
+}
+
+func (sf FunctionScope) Unit() code.Marker {
+	return sf.unit
+}
+
+func (sf FunctionScope) IsRootScope() bool {
+	return sf.ID() == ScopeID(0)
+}
diff --git a/pkg/lang/compiler/scope/symbol.go b/pkg/lang/compiler/scope/symbol.go
new file mode 100644
index 0000000..df91899
--- /dev/null
+++ b/pkg/lang/compiler/scope/symbol.go
@@ -0,0 +1,65 @@
+package scope
+
+import "jinx/pkg/lang/vm/code"
+
+type SymbolID struct {
+	symbolKind   SymbolKind
+	scopeID      ScopeID
+	indexInScope int
+}
+
+func (id SymbolID) SymbolKind() SymbolKind {
+	return id.symbolKind
+}
+
+type SymbolKind int
+
+const (
+	SymbolKindVariable SymbolKind = iota
+	SymbolKindFunction
+)
+
+func (s SymbolKind) String() string {
+	switch s {
+	case SymbolKindVariable:
+		return "variable"
+	case SymbolKindFunction:
+		return "function"
+	default:
+		panic("unknown symbol kind")
+	}
+}
+
+type Symbol[D SymbolData] struct {
+	name string
+	data D
+}
+
+func (s Symbol[D]) Data() D {
+	return s.data
+}
+
+type SymbolData interface {
+	SymbolVariable | SymbolFunction
+}
+
+type SymbolVariable struct {
+	localIndex int
+}
+
+func (sv SymbolVariable) LocalIndex() int {
+	return sv.localIndex
+}
+
+type SymbolFunction struct {
+	marker code.Marker
+	args   uint
+}
+
+func (sf SymbolFunction) Marker() code.Marker {
+	return sf.marker
+}
+
+func (sf SymbolFunction) Args() uint {
+	return sf.args
+}