1 files changed, 26 insertions, 11 deletions

diff --git a/pkg/lang/compiler/scope_chain.go b/pkg/lang/compiler/scope_chain.go
index ad176da..3a3b819 100644
--- a/pkg/lang/compiler/scope_chain.go
+++ b/pkg/lang/compiler/scope_chain.go
@@ -14,7 +14,7 @@ type ScopeChain struct {
 
 func NewScopeChain() ScopeChain {
 	scopes := make([]Scope, 1)
-	scopes[0] = NewFunctionScope("") // Top-most scope is a function scope, so it can have sub-units
+	scopes[0] = NewFunctionScope(0, "") // Top-most scope is a function scope, so it can have sub-units
 
 	return ScopeChain{
 		nameToSymbol: make(map[string]SymbolID),
@@ -26,13 +26,17 @@ func (sc *ScopeChain) CurrentScopeID() ScopeID {
 	return ScopeID(len(sc.scopes) - 1)
 }
 
+func (sc *ScopeChain) IsRootScope() bool {
+	return sc.CurrentScopeID() == ScopeID(0)
+}
+
 func (sc *ScopeChain) Current() *Scope {
 	return &sc.scopes[sc.CurrentScopeID()]
 }
 
 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++ {
+	for i := len(sc.scopes) - 1; i >= 0; i++ {
 		if sc.scopes[i].kind == ScopeKindFunction {
 			return &sc.scopes[i]
 		}
@@ -45,8 +49,9 @@ 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) EnterFunction(unit code.Marker) {
+	id := sc.CurrentScopeID() + 1
+	sc.scopes = append(sc.scopes, NewFunctionScope(id, unit))
 }
 
 func (sc *ScopeChain) Exit() {
@@ -148,7 +153,7 @@ func (sc *ScopeChain) CreateFunctionSubUnit(subUnitName string) code.Marker {
 	fnScope := sc.CurrentFunction()
 	data := fnScope.data.(ScopeFunction)
 
-	name := data.unitName
+	name := data.unit
 	if name == "" {
 		name = code.Marker(subUnitName)
 	} else {
@@ -174,12 +179,12 @@ func (sc *ScopeChain) GetVariable(id SymbolID) Symbol[SymbolVariable] {
 	return sc.scopes[id.scopeID].variableSymbols[id.indexInScope]
 }
 
-func (sc *ScopeChain) GetFunction(id SymbolID) Symbol[SymbolVariable] {
-	if id.symbolKind != SymbolKindVariable {
+func (sc *ScopeChain) GetFunction(id SymbolID) Symbol[SymbolFunction] {
+	if id.symbolKind != SymbolKindFunction {
 		panic("incorrect symbol id kind given")
 	}
 
-	return sc.scopes[id.scopeID].variableSymbols[id.indexInScope]
+	return sc.scopes[id.scopeID].functionSymbols[id.indexInScope]
 }
 
 type SymbolID struct {
@@ -213,13 +218,14 @@ func NewNormalScope() Scope {
 	}
 }
 
-func NewFunctionScope(unitName string) Scope {
+func NewFunctionScope(id ScopeID, unit code.Marker) Scope {
 	return Scope{
 		variableSymbols: make([]Symbol[SymbolVariable], 0),
 		functionSymbols: make([]Symbol[SymbolFunction], 0),
 		kind:            ScopeKindFunction,
 		data: ScopeFunction{
-			unitName:     code.Marker(unitName),
+			id:           id,
+			unit:         unit,
 			subUnitCount: 0,
 		},
 	}
@@ -228,10 +234,19 @@ func NewFunctionScope(unitName string) Scope {
 type ScopeNormal struct{}
 
 type ScopeFunction struct {
-	unitName     code.Marker
+	id           ScopeID
+	unit         code.Marker
 	subUnitCount int
 }
 
+func (sf ScopeFunction) ID() ScopeID {
+	return sf.id
+}
+
+func (sf ScopeFunction) IsRootScope() bool {
+	return sf.ID() == ScopeID(0)
+}
+
 type ScopeLoop struct {
 	breakMarker    code.Marker
 	continueMarker code.Marker