Improved dead store optimization

internal/optimizer/optimizer_test.go

@@ -572,6 +572,7 @@ func TestPassRegDead_KeptStoreWithLoad(t *testing.T) {
 }
 
 func TestPassRegDead_KeptStoreAtLabel(t *testing.T) {
+	// sta regVar before a label — no lda regVar anywhere → globally dead
 	input := []string{
 		"\tlda #$05",
 		"\tsta myFunc_temp",
@@ -586,11 +587,15 @@ func TestPassRegDead_KeptStoreAtLabel(t *testing.T) {
 		},
 	}
 
-	output, _ := Optimize(input, cfg)
+	output, dissolved := Optimize(input, cfg)
+
+	if !dissolved["myFunc_temp"] {
+		t.Error("expected myFunc_temp to be dissolved — globally dead")
+	}
 
 	joined := strings.Join(output, "\n")
-	if !strings.Contains(joined, "sta myFunc_temp") {
+	if strings.Contains(joined, "sta myFunc_temp") {
-		t.Errorf("expected sta myFunc_temp to be kept (label), got:\n%s", joined)
+		t.Errorf("expected sta myFunc_temp to be removed (globally dead), got:\n%s", joined)
 	}
 }
 
@@ -703,10 +708,9 @@ func TestPassRegDead_DeadStoreNeverReloaded(t *testing.T) {
 	}
 }
 
-func TestPassRegDead_CallSiteKeepsStore(t *testing.T) {
+func TestPassRegDead_GloballyDeadBeforeJsr(t *testing.T) {
-	// sta regVar; jsr foo; ... no lda regVar ...; rts → store kept (conservative)
+	// sta regVar; jsr foo; no lda regVar anywhere → globally dead
-	// The callee may modify A, so the store survives even though nothing
+	// jsr does not protect a store with zero readers.
-	// reloads regVar after the call.
 	input := []string{
 		"\tlda 53281",
 		"\tsta call_val",
@@ -723,13 +727,45 @@ func TestPassRegDead_CallSiteKeepsStore(t *testing.T) {
 
 	output, dissolved := Optimize(input, cfg)
 
+	if !dissolved["call_val"] {
+		t.Error("call_val should be dissolved — globally dead")
+	}
+
+	joined := strings.Join(output, "\n")
+	if strings.Contains(joined, "sta call_val") {
+		t.Errorf("expected sta call_val to be removed (globally dead), got:\n%s", joined)
+	}
+}
+
+func TestPassRegDead_JsrProtectsStoreWhenReloaded(t *testing.T) {
+	// sta regVar; jsr foo; lda regVar → store kept
+	// lda exists (pre-scan passes), but jsr blocks local scan before reaching it.
+	// The callee may modify A, so the reload after jsr is genuine.
+	input := []string{
+		"\tlda 53281",
+		"\tsta call_val",
+		"\tjsr helper",
+		"\tlda call_val",
+		"\tsta 53280",
+		"\trts",
+	}
+
+	cfg := &Config{
+		EnableRegisterVars: true,
+		RegisterVars: map[string]bool{
+			"call_val": true,
+		},
+	}
+
+	output, dissolved := Optimize(input, cfg)
+
 	if dissolved["call_val"] {
-		t.Error("call_val should NOT be dissolved — jsr is conservative")
+		t.Error("call_val should NOT be dissolved — lda exists, jsr protects store")
 	}
 
 	joined := strings.Join(output, "\n")
 	if !strings.Contains(joined, "sta call_val") {
-		t.Errorf("expected sta call_val to be kept (before jsr), got:\n%s", joined)
+		t.Errorf("expected sta call_val to be kept (jsr before reload), got:\n%s", joined)
 	}
 }
 
@@ -765,3 +801,68 @@ func TestPassRegDead_InitValueFlowsThrough(t *testing.T) {
 		t.Error("init value lda #$2a should survive")
 	}
 }
+
+func TestPassRegDead_GloballyDeadBeforeLabel(t *testing.T) {
+	// sta regVar before a label, no lda regVar anywhere → globally dead
+	// Typical copy-loop pattern: the value flows through A, never reloaded.
+	input := []string{
+		"_LOOPSTART",
+		"\tlda (src),y",
+		"\tsta loop_val",
+		"\tsta (dst),y",
+		"\tbne _SKIP",
+		"_SKIP:",
+		"\tjmp _LOOPSTART",
+		"\trts",
+	}
+
+	cfg := &Config{
+		EnableRegisterVars: true,
+		RegisterVars: map[string]bool{
+			"loop_val": true,
+		},
+	}
+
+	output, dissolved := Optimize(input, cfg)
+
+	if !dissolved["loop_val"] {
+		t.Error("loop_val should be dissolved — globally dead (no lda anywhere)")
+	}
+
+	joined := strings.Join(output, "\n")
+	if strings.Contains(joined, "sta loop_val") {
+		t.Errorf("expected sta loop_val to be removed (globally dead), got:\n%s", joined)
+	}
+}
+
+func TestPassRegDead_NotGloballyDeadWithLda(t *testing.T) {
+	// sta regVar before a label, but lda regVar exists elsewhere → not globally dead
+	// Falls through to local scan. With jsr before the lda, store is kept.
+	input := []string{
+		"\tsta spill_me_bg",
+		"\tlda 53282",
+		"_SKIP:",
+		"\tjsr helper",
+		"\tlda spill_me_bg",
+		"\tsta 53282",
+		"\trts",
+	}
+
+	cfg := &Config{
+		EnableRegisterVars: true,
+		RegisterVars: map[string]bool{
+			"spill_me_bg": true,
+		},
+	}
+
+	output, dissolved := Optimize(input, cfg)
+
+	if dissolved["spill_me_bg"] {
+		t.Error("spill_me_bg should NOT be dissolved — has a real lda")
+	}
+
+	joined := strings.Join(output, "\n")
+	if !strings.Contains(joined, "sta spill_me_bg") {
+		t.Errorf("expected sta spill_me_bg to be kept, got:\n%s", joined)
+	}
+}

internal/optimizer/pass_regdead.go

@@ -1,16 +1,33 @@
 package optimizer
 
 // passRegDead eliminates dead stores to REGISTER variables.
-// A store is dead if no matching load (lda regVar) follows before A is
+// Two-phase approach:
-// clobbered or a label is reached. This is safe because REGISTER variables
+//  1. Global pre-scan: if no lda regVar exists anywhere in the input,
-// have a contract: nothing external can observe their memory location.
+//     all stores to that variable are provably dead (labels are irrelevant).
+//  2. Per-store local scan: for variables that DO have at least one lda,
+//     scan forward through the current basic block.
+//
+// This is safe because REGISTER variables have a contract: nothing external
+// can observe their memory location (ASM/Script/Macro references are compile errors).
 func passRegDead(lines []asmLine, registerVars map[string]bool) []asmLine {
-	var result []asmLine
+	// Phase 1 — global pre-scan: which REGISTER vars have at least one lda?
+	hasLda := make(map[string]bool, len(registerVars))
+	for _, l := range lines {
+		if l.isCode && l.opcode == "lda" && l.operand != "" && registerVars[l.operand] {
+			hasLda[l.operand] = true
+		}
+	}
 
+	// Phase 2 — per-store decision
+	var result []asmLine
 	for i := 0; i < len(lines); i++ {
 		line := lines[i]
 
 		if line.isCode && line.opcode == "sta" && line.operand != "" && registerVars[line.operand] {
+			if !hasLda[line.operand] {
+				// Globally dead — no lda anywhere in the program
+				continue
+			}
 			if isRegStoreDead(lines, i+1, line.operand) {
 				continue
 			}