From fdee78a7b3a406ce4fe92c1e166fb45558585bbd Mon Sep 17 00:00:00 2001
From: Eduard Urbach <eduard@urbach.dev>
Date: Fri, 29 Aug 2025 16:48:13 +0200
Subject: [PATCH] Implemented detection for partial resource consumption

---
 src/codegen/hintDestination.go                |  2 +-
 src/codegen/needsRegister.go                  |  6 ++
 src/codegen/reorderPhis.go                    |  4 +-
 src/core/checkResources.go                    |  6 ++
 src/core/compileAssign.go                     |  8 +++
 src/core/errors.go                            | 18 ++++++
 src/core/errors_test.go                       |  3 +-
 src/core/removeDeadCode.go                    |  7 +++
 ...tifier9.q => PartiallyUnknownIdentifier.q} |  2 +-
 src/core/testdata/ResourcePartiallyConsumed.q | 13 ++++
 src/ssa/Block.go                              | 59 ++++++++++++++-----
 src/ssa/Block_test.go                         | 18 ++++--
 src/ssa/Phi.go                                | 35 +++++++++++
 src/ssa/Undefined.go                          | 18 ++++++
 14 files changed, 173 insertions(+), 26 deletions(-)
 rename src/core/testdata/{UnknownIdentifier9.q => PartiallyUnknownIdentifier.q} (69%)
 create mode 100644 src/core/testdata/ResourcePartiallyConsumed.q
 create mode 100644 src/ssa/Undefined.go

diff --git a/src/codegen/hintDestination.go b/src/codegen/hintDestination.go
index 9f76cf5..f05d041 100644
--- a/src/codegen/hintDestination.go
+++ b/src/codegen/hintDestination.go
@@ -20,7 +20,7 @@ func (f *Function) hintDestination(step *Step) {
 		// as the register for the left operand.
 		f.ValueToStep[instr.Left].hint(step.Register)
 	case *ssa.Phi:
-		for _, variant := range instr.Arguments {
+		for variant := range instr.DefinedArguments {
 			variant := f.ValueToStep[variant]
 			variant.Phis.Add(step)
 			variant.hint(step.Register)
diff --git a/src/codegen/needsRegister.go b/src/codegen/needsRegister.go
index 1f787f0..e1f66c8 100644
--- a/src/codegen/needsRegister.go
+++ b/src/codegen/needsRegister.go
@@ -13,6 +13,12 @@ func (f *Function) needsRegister(s *Step) bool {
 		return false
 	}
 
+	_, isPhi := s.Value.(*ssa.Phi)
+
+	if isPhi {
+		return true
+	}
+
 	_, isStruct := typ.(*types.Struct)
 
 	if isStruct {
diff --git a/src/codegen/reorderPhis.go b/src/codegen/reorderPhis.go
index bfdcd9a..6b862f8 100644
--- a/src/codegen/reorderPhis.go
+++ b/src/codegen/reorderPhis.go
@@ -21,8 +21,8 @@ func (f *Function) reorderPhis() {
 			phis := f.Steps[start:end]
 
 			slices.SortStableFunc(phis, func(a *Step, b *Step) int {
-				aIndex := f.ValueToStep[a.Value.(*ssa.Phi).Arguments[0]].Index
-				bIndex := f.ValueToStep[b.Value.(*ssa.Phi).Arguments[0]].Index
+				aIndex := f.ValueToStep[a.Value.(*ssa.Phi).FirstDefined()].Index
+				bIndex := f.ValueToStep[b.Value.(*ssa.Phi).FirstDefined()].Index
 				return aIndex - bIndex
 			})
 
diff --git a/src/core/checkResources.go b/src/core/checkResources.go
index d2a13c2..e0d53e1 100644
--- a/src/core/checkResources.go
+++ b/src/core/checkResources.go
@@ -26,6 +26,12 @@ func (f *Function) checkResources() error {
 				continue
 			}
 
+			phi, isPhi := value.(*ssa.Phi)
+
+			if isPhi && phi.IsPartiallyUndefined() {
+				return errors.New(&ResourcePartiallyConsumed{TypeName: resource.Name()}, f.File, phi.FirstDefined().(ssa.HasSource).Start())
+			}
+
 			return errors.New(&ResourceNotConsumed{TypeName: resource.Name()}, f.File, value.(ssa.HasSource).Start())
 		}
 	}
diff --git a/src/core/compileAssign.go b/src/core/compileAssign.go
index 14fed2a..d214ce8 100644
--- a/src/core/compileAssign.go
+++ b/src/core/compileAssign.go
@@ -23,6 +23,14 @@ func (f *Function) compileAssign(node *ast.Assign) error {
 	name := left.String(f.File.Bytes)
 	leftValue, exists := f.Block().FindIdentifier(name)
 
+	if exists {
+		phi, isPhi := leftValue.(*ssa.Phi)
+
+		if isPhi && phi.IsPartiallyUndefined() {
+			return errors.New(&PartiallyUnknownIdentifier{Name: name}, f.File, left.Source().StartPos)
+		}
+	}
+
 	if !exists {
 		return errors.New(&UnknownIdentifier{Name: name}, f.File, left.Source().StartPos)
 	}
diff --git a/src/core/errors.go b/src/core/errors.go
index 0c899eb..eef24df 100644
--- a/src/core/errors.go
+++ b/src/core/errors.go
@@ -76,6 +76,15 @@ func (err *ResourceNotConsumed) Error() string {
 	return fmt.Sprintf("Resource of type '%s' not consumed", err.TypeName)
 }
 
+// ResourcePartiallyConsumed error is created when a resource has only partially been consumed in an exit block.
+type ResourcePartiallyConsumed struct {
+	TypeName string
+}
+
+func (err *ResourcePartiallyConsumed) Error() string {
+	return fmt.Sprintf("Resource of type '%s' not consumed in all branches", err.TypeName)
+}
+
 // ReturnCountMismatch error is created when the number of returned values doesn't match the return type.
 type ReturnCountMismatch struct {
 	Count         int
@@ -135,6 +144,15 @@ func (err *UnknownIdentifier) Error() string {
 	return fmt.Sprintf("Unknown identifier '%s'", err.Name)
 }
 
+// PartiallyUnknownIdentifier represents identifiers that were only defined in one branch.
+type PartiallyUnknownIdentifier struct {
+	Name string
+}
+
+func (err *PartiallyUnknownIdentifier) Error() string {
+	return fmt.Sprintf("Identifier '%s' is not defined in every branch", err.Name)
+}
+
 // UndefinedStructField is created when an undefined struct field is accessed.
 type UndefinedStructField struct {
 	Identifier string
diff --git a/src/core/errors_test.go b/src/core/errors_test.go
index 9fe5037..0cbf745 100644
--- a/src/core/errors_test.go
+++ b/src/core/errors_test.go
@@ -35,11 +35,13 @@ var errs = []struct {
 	{"NotDataStruct.q", &core.NotDataStruct{TypeName: "int"}},
 	{"ParameterCountMismatch.q", &core.ParameterCountMismatch{Function: "main.f", Count: 0, ExpectedCount: 1}},
 	{"ParameterCountMismatch2.q", &core.ParameterCountMismatch{Function: "main.f", Count: 2, ExpectedCount: 1}},
+	{"PartiallyUnknownIdentifier.q", &core.PartiallyUnknownIdentifier{Name: "x"}},
 	{"ResourceAlreadyConsumed.q", &core.UnknownIdentifier{Name: "x"}},
 	{"ResourceAlreadyConsumed2.q", &core.UnknownIdentifier{Name: "x"}},
 	{"ResourceAlreadyConsumed3.q", &core.UnknownIdentifier{Name: "x"}},
 	{"ResourceAlreadyConsumed4.q", &core.UnknownIdentifier{Name: "x"}},
 	{"ResourceNotConsumed.q", &core.ResourceNotConsumed{TypeName: "!int64"}},
+	{"ResourcePartiallyConsumed.q", &core.ResourcePartiallyConsumed{TypeName: "!int64"}},
 	{"ResourceTypeMismatch.q", &core.TypeMismatch{Encountered: "int64", Expected: "!int64", ParameterName: "_", IsReturn: false}},
 	{"ReturnCountMismatch.q", &core.ReturnCountMismatch{Count: 1, ExpectedCount: 0}},
 	{"ReturnCountMismatch2.q", &core.ReturnCountMismatch{Count: 1, ExpectedCount: 2}},
@@ -54,7 +56,6 @@ var errs = []struct {
 	{"UnknownIdentifier6.q", &core.UnknownIdentifier{Name: "os"}},
 	{"UnknownIdentifier7.q", &core.UnknownIdentifier{Name: "os.unknown"}},
 	{"UnknownIdentifier8.q", &core.UnknownIdentifier{Name: "x"}},
-	{"UnknownIdentifier9.q", &core.UnknownIdentifier{Name: "x"}},
 	{"UnknownStructField.q", &core.UnknownStructField{StructName: "string", FieldName: "unknown"}},
 	{"UnusedValue.q", &core.UnusedValue{Value: "42"}},
 	{"UnusedValue2.q", &core.UnusedValue{Value: "2 + 3"}},
diff --git a/src/core/removeDeadCode.go b/src/core/removeDeadCode.go
index 7470647..ee2d007 100644
--- a/src/core/removeDeadCode.go
+++ b/src/core/removeDeadCode.go
@@ -17,6 +17,13 @@ func (f *Function) removeDeadCode() error {
 				continue
 			}
 
+			_, isPhi := value.(*ssa.Phi)
+
+			if isPhi {
+				block.RemoveAt(i)
+				continue
+			}
+
 			structField, isFieldOfStruct := value.(ssa.StructField)
 
 			if isFieldOfStruct && structField.Struct() != nil {
diff --git a/src/core/testdata/UnknownIdentifier9.q b/src/core/testdata/PartiallyUnknownIdentifier.q
similarity index 69%
rename from src/core/testdata/UnknownIdentifier9.q
rename to src/core/testdata/PartiallyUnknownIdentifier.q
index 56829e9..e2c3d51 100644
--- a/src/core/testdata/UnknownIdentifier9.q
+++ b/src/core/testdata/PartiallyUnknownIdentifier.q
@@ -1,5 +1,5 @@
 main() {
-	if 0 == 0 {
+	if true {
 		x := 1
 	}
 
diff --git a/src/core/testdata/ResourcePartiallyConsumed.q b/src/core/testdata/ResourcePartiallyConsumed.q
new file mode 100644
index 0000000..e8b34c5
--- /dev/null
+++ b/src/core/testdata/ResourcePartiallyConsumed.q
@@ -0,0 +1,13 @@
+main() {
+	x := acquire()
+
+	if true {
+		use(x)
+	} else {
+		free(x)
+	}
+}
+
+acquire() -> !int { return 1 }
+use(_ int) {}
+free(_ !int) {}
\ No newline at end of file
diff --git a/src/ssa/Block.go b/src/ssa/Block.go
index 0e84849..a7d34d2 100644
--- a/src/ssa/Block.go
+++ b/src/ssa/Block.go
@@ -40,35 +40,62 @@ func (b *Block) AddSuccessor(successor *Block) {
 		}
 	}
 
-	for name, oldValue := range successor.Identifiers {
-		newValue, exists := b.Identifiers[name]
+	keys := make(map[string]struct{}, max(len(b.Identifiers), len(successor.Identifiers)))
 
-		if !exists {
-			delete(successor.Identifiers, name)
-			continue
-		}
+	for name := range successor.Identifiers {
+		keys[name] = struct{}{}
+	}
 
-		if oldValue == newValue {
-			continue
-		}
+	for name := range b.Identifiers {
+		keys[name] = struct{}{}
+	}
+
+	for name := range keys {
+		oldValue, oldExists := successor.Identifiers[name]
+		newValue, newExists := b.Identifiers[name]
+
+		switch {
+		case oldExists:
+			if oldValue == newValue {
+				continue
+			}
+
+			phi, isPhi := oldValue.(*Phi)
 
-		phi, isPhi := oldValue.(*Phi)
+			if !isPhi || successor.Index(phi) == -1 {
+				phi = &Phi{
+					Arguments: make([]Value, len(successor.Predecessors)-1, len(successor.Predecessors)),
+					Typ:       oldValue.Type(),
+				}
+
+				for i := range phi.Arguments {
+					phi.Arguments[i] = oldValue
+				}
+
+				successor.InsertAt(phi, 0)
+				successor.Identifiers[name] = phi
+			}
+
+			if newExists {
+				phi.Arguments = append(phi.Arguments, newValue)
+			} else {
+				phi.Arguments = append(phi.Arguments, Undefined)
+			}
 
-		if !isPhi || successor.Index(phi) == -1 {
-			phi = &Phi{
+		case newExists:
+			phi := &Phi{
 				Arguments: make([]Value, len(successor.Predecessors)-1, len(successor.Predecessors)),
-				Typ:       oldValue.Type(),
+				Typ:       newValue.Type(),
 			}
 
 			for i := range phi.Arguments {
-				phi.Arguments[i] = oldValue
+				phi.Arguments[i] = Undefined
 			}
 
 			successor.InsertAt(phi, 0)
 			successor.Identifiers[name] = phi
+			phi.Arguments = append(phi.Arguments, newValue)
 		}
-
-		phi.Arguments = append(phi.Arguments, newValue)
 	}
 }
 
diff --git a/src/ssa/Block_test.go b/src/ssa/Block_test.go
index 3dd8a99..e5143db 100644
--- a/src/ssa/Block_test.go
+++ b/src/ssa/Block_test.go
@@ -150,8 +150,12 @@ func TestBlockFindIdentifier(t *testing.T) {
 	_, exists = elseBlock.FindIdentifier("then")
 	assert.False(t, exists)
 
-	_, exists = mergeBlock.FindIdentifier("then")
-	assert.False(t, exists)
+	partial, exists := mergeBlock.FindIdentifier("then")
+	assert.True(t, exists)
+	phi, isPhi := partial.(*ssa.Phi)
+	assert.True(t, isPhi)
+	assert.True(t, phi.IsPartiallyUndefined())
+	assert.Equal(t, phi.FirstDefined(), value)
 
 	// Else
 	_, exists = branch.FindIdentifier("else")
@@ -164,8 +168,12 @@ func TestBlockFindIdentifier(t *testing.T) {
 	assert.True(t, exists)
 	assert.Equal(t, value, elseValue)
 
-	_, exists = mergeBlock.FindIdentifier("else")
-	assert.False(t, exists)
+	partial, exists = mergeBlock.FindIdentifier("else")
+	assert.True(t, exists)
+	phi, isPhi = partial.(*ssa.Phi)
+	assert.True(t, isPhi)
+	assert.True(t, phi.IsPartiallyUndefined())
+	assert.Equal(t, phi.FirstDefined(), value)
 
 	// Merge
 	_, exists = branch.FindIdentifier("merge")
@@ -184,7 +192,7 @@ func TestBlockFindIdentifier(t *testing.T) {
 	// Phi
 	value, exists = mergeBlock.FindIdentifier("x")
 	assert.True(t, exists)
-	phi, isPhi := value.(*ssa.Phi)
+	phi, isPhi = value.(*ssa.Phi)
 	assert.True(t, isPhi)
 	assert.Equal(t, phi.Arguments[0], thenValue)
 	assert.Equal(t, phi.Arguments[1], elseValue)
diff --git a/src/ssa/Phi.go b/src/ssa/Phi.go
index 321f941..cab490c 100644
--- a/src/ssa/Phi.go
+++ b/src/ssa/Phi.go
@@ -13,6 +13,19 @@ type Phi struct {
 	Liveness
 }
 
+// DefinedArguments is an iterator over arguments that are defined.
+func (p *Phi) DefinedArguments(yield func(Value) bool) {
+	for _, arg := range p.Arguments {
+		if arg == Undefined {
+			continue
+		}
+
+		if !yield(arg) {
+			return
+		}
+	}
+}
+
 // Equals returns true if the phi nodes are equal.
 func (a *Phi) Equals(v Value) bool {
 	b, sameType := v.(*Phi)
@@ -24,11 +37,33 @@ func (a *Phi) Equals(v Value) bool {
 	return a.Arguments.Equals(b.Arguments)
 }
 
+// FirstDefined returns the first input value that is not undefined.
+func (p *Phi) FirstDefined() Value {
+	for _, arg := range p.Arguments {
+		if arg != Undefined {
+			return arg
+		}
+	}
+
+	panic("phi composed of only undefined values must not exist")
+}
+
 // IsConst returns true because two equal phi nodes produce the same value.
 func (p *Phi) IsConst() bool {
 	return true
 }
 
+// IsPartiallyUndefined returns true if one of the input values is undefined.
+func (p *Phi) IsPartiallyUndefined() bool {
+	for _, arg := range p.Arguments {
+		if arg == Undefined {
+			return true
+		}
+	}
+
+	return false
+}
+
 // String returns a human-readable representation of the phi node.
 func (p *Phi) String() string {
 	return fmt.Sprintf("phi(%s)", p.Arguments.String())
diff --git a/src/ssa/Undefined.go b/src/ssa/Undefined.go
new file mode 100644
index 0000000..908a0d1
--- /dev/null
+++ b/src/ssa/Undefined.go
@@ -0,0 +1,18 @@
+package ssa
+
+// Undefined is used in phi values to represent an undefined value.
+var Undefined = &undefined{}
+
+type undefined struct{ Void }
+
+// Equals always returns false.
+func (v *undefined) Equals(Value) bool { return false }
+
+// Inputs always returns nil.
+func (v *undefined) Inputs() []Value { return nil }
+
+// Replace does nothing.
+func (v *undefined) Replace(Value, Value) {}
+
+// String returns a human-readable representation of the undefined value.
+func (v *undefined) String() string { return "undefined" }
\ No newline at end of file
-- 
2.51.2