1 files changed, 181 insertions, 45 deletions

diff --git a/pkg/lang/compiler/scope_chain.go b/pkg/lang/compiler/scope_chain.go
index 6b7e693..ad176da 100644
--- a/pkg/lang/compiler/scope_chain.go
+++ b/pkg/lang/compiler/scope_chain.go
@@ -1,102 +1,238 @@
 package compiler
 
+import (
+	"fmt"
+	"jinx/pkg/lang/vm/code"
+)
+
 type ScopeID int
 
 type ScopeChain struct {
-	scopes []Scope
+	nameToSymbol map[string]SymbolID
+	scopes       []Scope
 }
 
 func NewScopeChain() ScopeChain {
 	scopes := make([]Scope, 1)
-	scopes[0] = Scope{
-		kind:         ScopeKindGlobal,
-		nameToSymbol: make(map[string]int),
-		symbols:      make([]Symbol, 0),
-	}
+	scopes[0] = NewFunctionScope("") // Top-most scope is a function scope, so it can have sub-units
 
 	return ScopeChain{
-		scopes: scopes,
+		nameToSymbol: make(map[string]SymbolID),
+		scopes:       scopes,
 	}
 }
 
+func (sc *ScopeChain) CurrentScopeID() ScopeID {
+	return ScopeID(len(sc.scopes) - 1)
+}
+
 func (sc *ScopeChain) Current() *Scope {
-	return &sc.scopes[len(sc.scopes)-1]
+	return &sc.scopes[sc.CurrentScopeID()]
 }
 
-func (sc *ScopeChain) Enter(kind ScopeKind) {
-	sc.scopes = append(sc.scopes, Scope{
-		kind:         kind,
-		nameToSymbol: make(map[string]int),
-		symbols:      make([]Symbol, 0),
-	})
+func (sc *ScopeChain) CurrentFunction() *Scope {
+	// TODO: Probably should make this lookup constant by making a seperate array of function scopes
+	for i := len(sc.scopes) - 1; i <= 0; i++ {
+		if sc.scopes[i].kind == ScopeKindFunction {
+			return &sc.scopes[i]
+		}
+	}
+
+	panic("top scope should always be a function scope")
+}
+
+func (sc *ScopeChain) Enter() {
+	sc.scopes = append(sc.scopes, NewNormalScope())
+}
+
+func (sc *ScopeChain) EnterFunction(unitName string) {
+	sc.scopes = append(sc.scopes, NewFunctionScope(unitName))
 }
 
 func (sc *ScopeChain) Exit() {
-	sc.scopes[len(sc.scopes)-1] = Scope{}
-	sc.scopes = sc.scopes[:len(sc.scopes)-1]
+	if sc.CurrentScopeID() == 0 {
+		return
+	}
+
+	sc.scopes[sc.CurrentScopeID()] = Scope{}
+	sc.scopes = sc.scopes[:sc.CurrentScopeID()]
 }
 
 func (sc *ScopeChain) Declare(name string) (int, bool) {
 	// Check whether the symbol is already declared in any of the scopes.
-	for _, scope := range sc.scopes {
-		if _, ok := scope.nameToSymbol[name]; ok {
-			return 0, false
-		}
+	if _, ok := sc.nameToSymbol[name]; ok {
+		return 0, false
 	}
 
 	current := sc.Current()
-	index := len(current.symbols)
+	indexInScope := len(current.variableSymbols)
+
+	symbolID := SymbolID{
+		symbolKind:   SymbolKindVariable,
+		scopeID:      sc.CurrentScopeID(),
+		indexInScope: indexInScope,
+	}
 
 	// Declare the symbol in the current scope.
-	current.symbols = append(current.symbols, Symbol{
-		kind:       SymbolKindVariable,
-		name:       name,
-		localIndex: index,
+	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) (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,
+		},
 	})
 
-	current.nameToSymbol[name] = index
+	sc.nameToSymbol[name] = symbolID
 
-	return index, true
+	return unitName, true
 }
 
 func (sc *ScopeChain) DeclareAnonymous() int {
 	current := sc.Current()
-	index := len(current.symbols)
+	index := len(current.variableSymbols)
 
-	// Declare the symbol in the current scope.
-	current.symbols = append(current.symbols, Symbol{
-		kind:       SymbolKindVariable,
-		name:       "",
-		localIndex: index,
+	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().symbols)
+	return len(sc.Current().variableSymbols) // :)
 }
 
-func (sc *ScopeChain) Lookup(name string) (Symbol, bool) {
-	for i := len(sc.scopes) - 1; i >= 0; i-- {
-		if symbol, ok := sc.scopes[i].nameToSymbol[name]; ok {
-			return sc.scopes[i].symbols[symbol], true
-		}
+func (sc *ScopeChain) CreateAnonymousFunctionSubUnit() code.Marker {
+	fnScope := sc.CurrentFunction()
+	data := fnScope.data.(ScopeFunction)
+
+	index := data.subUnitCount
+	data.subUnitCount++
+
+	fnScope.data = data
+
+	return sc.CreateFunctionSubUnit(fmt.Sprintf("anon_%d", index))
+}
+
+func (sc *ScopeChain) CreateFunctionSubUnit(subUnitName string) code.Marker {
+	fnScope := sc.CurrentFunction()
+	data := fnScope.data.(ScopeFunction)
+
+	name := data.unitName
+	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.scopes[id.scopeID].variableSymbols[id.indexInScope]
+}
+
+func (sc *ScopeChain) GetFunction(id SymbolID) Symbol[SymbolVariable] {
+	if id.symbolKind != SymbolKindVariable {
+		panic("incorrect symbol id kind given")
 	}
 
-	return Symbol{}, false
+	return sc.scopes[id.scopeID].variableSymbols[id.indexInScope]
+}
+
+type SymbolID struct {
+	symbolKind   SymbolKind
+	scopeID      ScopeID
+	indexInScope int
 }
 
 type ScopeKind int
 
 const (
-	ScopeKindGlobal ScopeKind = iota
+	ScopeKindNormal ScopeKind = iota
 	ScopeKindFunction
-	ScopeKindBlock
+	ScopeKindLoop
 )
 
 type Scope struct {
-	kind         ScopeKind
-	nameToSymbol map[string]int
-	symbols      []Symbol
+	variableSymbols []Symbol[SymbolVariable]
+	functionSymbols []Symbol[SymbolFunction]
+
+	kind ScopeKind
+	data any
+}
+
+func NewNormalScope() Scope {
+	return Scope{
+		variableSymbols: make([]Symbol[SymbolVariable], 0),
+		functionSymbols: make([]Symbol[SymbolFunction], 0),
+		kind:            ScopeKindNormal,
+		data:            ScopeNormal{},
+	}
+}
+
+func NewFunctionScope(unitName string) Scope {
+	return Scope{
+		variableSymbols: make([]Symbol[SymbolVariable], 0),
+		functionSymbols: make([]Symbol[SymbolFunction], 0),
+		kind:            ScopeKindFunction,
+		data: ScopeFunction{
+			unitName:     code.Marker(unitName),
+			subUnitCount: 0,
+		},
+	}
+}
+
+type ScopeNormal struct{}
+
+type ScopeFunction struct {
+	unitName     code.Marker
+	subUnitCount int
+}
+
+type ScopeLoop struct {
+	breakMarker    code.Marker
+	continueMarker code.Marker
 }