Add conditional type narrowing

ast/ast.go

@@ -192,6 +192,7 @@ func UnMask(mask Flags) common.Set[model.Flag] {
 
 type BNodeWithSymbol interface {
 	model.NodeWithSymbol
+	BLangNode
 	SetSymbol(symbolRef model.SymbolRef)
 }
 

ast/statements.go

@@ -202,6 +202,14 @@
 	return model.NodeKind_BREAK
 }
 
+func (this *BLangCompoundAssignment) IsDeclaredWithVar() bool {
+	return false
+}
+
+func (this *BLangCompoundAssignment) SetDeclaredWithVar(_ bool) {
+	panic("compound assignemnt can't be declared with var")
+}
+
 func (this *BLangCompoundAssignment) GetOperatorKind() model.OperatorKind {
 	// migrated from BLangCompoundAssignment.java:59:5
 	return this.OpKind

bir/bir_gen.go

@@ -811,7 +811,7 @@ func binaryExpression(ctx *stmtContext, curBB *BIRBasicBlock, expr *ast.BLangBin
 
 func simpleVariableReference(ctx *stmtContext, curBB *BIRBasicBlock, expr *ast.BLangSimpleVarRef) expressionEffect {
 	varName := expr.VariableName.GetValue()
-	symRef := expr.Symbol()
+	symRef := ctx.birCx.CompilerContext.UnnarrowedSymbol(expr.Symbol())
 
 	// Try local variable lookup first
 	if operand, ok := ctx.varMap[symRef]; ok {

bir/corpus_bir_test.go

@@ -119,7 +119,10 @@ func testBIRGeneration(t *testing.T, testPair test_util.TestCase) {
 	// Step 6: Generate control flow graph
 	cfg := semantics.CreateControlFlowGraph(cx, pkg)
 
-	// Step 7: Run semantic analysis
+	// Step 7: Run type narrowing
+	semantics.NarrowTypes(cx, pkg)
+
+	// Step 8: Run semantic analysis
 	semanticAnalyzer := semantics.NewSemanticAnalyzer(cx)
 	semanticAnalyzer.Analyze(pkg)
 

context/context.go

@@ -22,13 +22,15 @@ import (
 	"ballerina-lang-go/tools/diagnostics"
 	"fmt"
 	"strconv"
+	"sync"
 )
 
 type CompilerContext struct {
-	anonTypeCount   map[*model.PackageID]int
-	packageInterner *model.PackageIDInterner
-	symbolSpaces    []*model.SymbolSpace
-	typeEnv         semtypes.Env
+	anonTypeCount    map[*model.PackageID]int
+	packageInterner  *model.PackageIDInterner
+	symbolSpaces     []*model.SymbolSpace
+	typeEnv          semtypes.Env
+	underlyingSymbol sync.Map
 }
 
 func (this *CompilerContext) NewSymbolSpace(packageId model.PackageID) *model.SymbolSpace {
@@ -56,8 +58,27 @@ func (this *CompilerContext) NewBlockScope(parent model.Scope, pkg model.Package
 }
 
 func (this *CompilerContext) GetSymbol(symbol model.SymbolRef) model.Symbol {
-	symbolSpace := this.symbolSpaces[symbol.SpaceIndex]
-	return symbolSpace.Symbols[symbol.Index]
+	return this.symbolSpaces[symbol.SpaceIndex].SymbolAt(symbol.Index)
+}
+
+func (this *CompilerContext) CreateNarrowedSymbol(baseRef model.SymbolRef) model.SymbolRef {
+	symbolSpace := this.symbolSpaces[baseRef.SpaceIndex]
+	underlyingSymbolCopy := *this.GetSymbol(baseRef).(*model.ValueSymbol)
+	symbolIndex := symbolSpace.AppendSymbol(&underlyingSymbolCopy)
+	narrowedSymbol := model.SymbolRef{
+		Package:    baseRef.Package,
+		SpaceIndex: baseRef.SpaceIndex,
+		Index:      symbolIndex,
+	}
+	this.underlyingSymbol.Store(narrowedSymbol, baseRef)
+	return narrowedSymbol
+}
+
+func (this *CompilerContext) UnnarrowedSymbol(symbol model.SymbolRef) model.SymbolRef {
+	if underlying, ok := this.underlyingSymbol.Load(symbol); ok {
+		return underlying.(model.SymbolRef)
+	}
+	return symbol
 }
 
 func (this *CompilerContext) SymbolName(symbol model.SymbolRef) string {