Keep InstructionInfo metadata consistent during constant propagation

CLAUDE.md

@@ -65,11 +65,12 @@ Available transformations (in order):
 | Name                               | Effect                              |
 | ---------------------------------- | ----------------------------------- |
 | `static-single-assignment` / `ssa` | Convert to SSA form                 |
+| `constant-propagation` / `cp`      | Propagate and fold constants        |
 | `dead-code-elimination` / `dce`    | Remove unused instructions          |
 | `aarch64` / `arm64`                | Translate USM → AArch64 assembly    |
 | `macho` / `macho-obj`              | Emit Mach-O `.o` object file        |
 
-Typical pipeline: `usm input.usm ssa dce aarch64 macho`
+Typical pipeline: `usm input.usm ssa cp dce aarch64 macho`
 
 ## Architecture: Key Patterns
 
@@ -93,6 +94,9 @@ Instructions declare capabilities via trait interfaces:
 - `NonBranchingInstruction` — falls through to the next instruction
 - `Uses(i int)` — returns the i-th used value (for liveness analysis)
 - `Defines()` — returns the defined value (for SSA/DCE)
+- `PropagateConstants(info) []ConstantDefinition` — for constant propagation:
+  returns (register, immediate) pairs that are known-constant after this
+  instruction executes; embed `opt.PropagatesNoConstants` as the default no-op
 
 ### Result/Error System (`core`)
 
@@ -172,12 +176,35 @@ package.
 ## Adding a New Optimization Pass
 
 1. Create a file in `opt/`.
-2. The pass receives a `*gen.FunctionGenerationContext` (or similar IR).
+2. The pass receives a `*gen.FunctionInfo` (or similar IR).
 3. Implement liveness/dataflow analysis if needed — see `opt/dead_code_elimination.go`
-   as a reference.
+   and `opt/constant_propagation.go` as references.
 4. Passes must run **after SSA construction** if they rely on SSA properties.
 5. Register the pass as a `transform.Transformation`.
 
+### Constant Propagation Pass (`opt/constant_propagation.go`)
+
+Propagates known-constant registers to their use sites and folds constant
+expressions. Uses a DFS over the CFG (`ControlFlowGraph.Dfs(0).Timeline`) with
+a per-register reaching-constants stack (mirrors `opt/ssa.ReachingDefinitionsSet`):
+
+- Only tracks registers with exactly **one reachable definition**. Unreachable
+  definitions (dead code, isolated loops) do not count. Registers with multiple
+  reachable definitions are propagated only when all those definitions are
+  sequential redefinitions in the same basic block (later definitions shadow
+  earlier ones on the DFS stack). Diamond joins and cross-block redefinitions
+  are never propagated.
+- After substituting arguments, calls `PropagateConstants` on the instruction to
+  fold constant expressions (e.g. `add #2 #3 → #5`). Folded results propagate
+  to downstream uses in the same DFS scope.
+- Uses CFG DFS (not dominator-tree DFS) to correctly handle unreachable blocks:
+  Lengauer-Tarjan assigns `PreOrder=0` to unvisited nodes, which can corrupt the
+  dominator tree when unreachable blocks are present.
+
+To support a new instruction in CP: implement `PropagateConstants` or embed
+`opt.PropagatesNoConstants`. Binary arithmetic helpers (`foldBinaryConstants`) live
+in `usm/isa/binary_calculation.go`.
+
 ## CI/CD
 
 GitHub Actions (`.github/workflows/ci.yml`):

go.mod

@@ -5,7 +5,7 @@ go 1.23.0
 require (
 	alon.kr/x/aarch64codegen v0.0.0-20250509160800-a2358754df5d
 	alon.kr/x/faststringmap v0.0.0-20250503134653-20d6364c2c94
-	alon.kr/x/graph v0.0.0-20250319212444-dd67d0281ab7
+	alon.kr/x/graph v0.0.0-20260316082925-2052fa913ff8
 	alon.kr/x/list v0.0.0-20241203223347-3173d76828c0
 	alon.kr/x/macho v0.0.0-20250426181816-9e10903e5803
 	alon.kr/x/set v0.0.0-20250319203758-c0328d5645c9

go.sum

@@ -2,8 +2,8 @@ alon.kr/x/aarch64codegen v0.0.0-20250509160800-a2358754df5d h1:zaesIgOOypUXQT8za
 alon.kr/x/aarch64codegen v0.0.0-20250509160800-a2358754df5d/go.mod h1:WAdZYqOdp9KwoBjWamrMDAphahR5oXkPSICj3+tLIyQ=
 alon.kr/x/faststringmap v0.0.0-20250503134653-20d6364c2c94 h1:EdG6bQh5IT5MZ76MY85/GK4Ma7zoJ4zcFqd2fpxK6NI=
 alon.kr/x/faststringmap v0.0.0-20250503134653-20d6364c2c94/go.mod h1:dPORoTvAFIehHRaI/lfJV3eT6PL5tY7fpAtxzz7rQVw=
-alon.kr/x/graph v0.0.0-20250319212444-dd67d0281ab7 h1:0d/oLvPQWc1B38ZfWfJ5UQXeHx9JT8gmp3K0AOOG8aQ=
-alon.kr/x/graph v0.0.0-20250319212444-dd67d0281ab7/go.mod h1:Cat+EJocX/ddkDX1quMqwe+itOfD7DILXV76nG0tk8Q=
+alon.kr/x/graph v0.0.0-20260316082925-2052fa913ff8 h1:iMPYCLtbIZSc+csc+NVtGgDjgaFQ4oThgKcyq2Nu2Sk=
+alon.kr/x/graph v0.0.0-20260316082925-2052fa913ff8/go.mod h1:Cat+EJocX/ddkDX1quMqwe+itOfD7DILXV76nG0tk8Q=
 alon.kr/x/list v0.0.0-20241203223347-3173d76828c0 h1:NedGw3ZJMKT9Y4RByuogGfBEXT4st8eX/XO5frezx6o=
 alon.kr/x/list v0.0.0-20241203223347-3173d76828c0/go.mod h1:7syBFMuVmxetKppmlKe6nQlp2BRueprI1r8/d7GB7fs=
 alon.kr/x/macho v0.0.0-20250426181816-9e10903e5803 h1:iVNUBqzMx3kA7hJRKQRDsfPX7TQ/kpnNjQ1NnSZ7vcQ=

opt/constant_propagation.go

@@ -0,0 +1,218 @@
+package opt
+
+import (
+	"alon.kr/x/list"
+	"alon.kr/x/stack"
+	"alon.kr/x/usm/core"
+	"alon.kr/x/usm/gen"
+)
+
+// ConstantDefinition pairs a target register with the constant immediate
+// value that an instruction always assigns to it.
+type ConstantDefinition struct {
+	Register  *gen.RegisterInfo
+	Immediate *gen.ImmediateInfo
+}
+
+// ConstantPropagationSupportedInstruction is implemented by ISA instructions
+// that support the constant propagation optimization pass.
+type ConstantPropagationSupportedInstruction interface {
+	gen.InstructionDefinition
+
+	// PropagateConstants returns (register, constant) pairs for each target
+	// that this instruction provably assigns a constant immediate to.
+	// Returns an empty slice if this instruction defines no constants.
+	// Example: "$64 %x = $64 #5" → [{Register: %x, Immediate: $64 #5}]
+	PropagateConstants(info *gen.InstructionInfo) []ConstantDefinition
+}
+
+// PropagatesNoConstants can be embedded in instruction definitions that never
+// assign a constant immediate to any of their targets.
+type PropagatesNoConstants struct{}
+
+func (PropagatesNoConstants) PropagateConstants(*gen.InstructionInfo) []ConstantDefinition {
+	return nil
+}
+
+func newCPNotSupportedError(instruction *gen.InstructionInfo) core.ResultList {
+	return list.FromSingle(core.Result{{
+		Type:     core.InternalErrorResult,
+		Message:  "Instruction does not support constant propagation",
+		Location: instruction.Declaration,
+	}})
+}
+
+// cpBlockSeparator is the sentinel pushed onto cpReachingConstants.registerPushes
+// to mark the boundary between basic blocks in the DFS.
+// Mirrors the blockSeparator sentinel used in ReachingDefinitionsSet (opt/ssa).
+const cpBlockSeparator = ^uint(0)
+
+// cpReachingConstants tracks, for each register, the constant value of its
+// reaching definition at the current position in a DFS traversal.
+type cpReachingConstants struct {
+	// Per-register stacks of reaching constant values. Indices are assigned
+	// lazily via registersToIndex the first time a register is defined.
+	registerStacks []stack.Stack[*gen.ImmediateInfo]
+
+	// Records which register indices were pushed in each block, with
+	// cpBlockSeparator values marking block boundaries.
+	registerPushes stack.Stack[uint]
+
+	// Maps each *RegisterInfo to its index in registerStacks. Entries are
+	// added on demand in define().
+	registersToIndex map[*gen.RegisterInfo]uint
+
+	// Set of reachable blocks (from entry). Used by define() to ignore
+	// definitions in unreachable blocks when checking eligibility.
+	reachable map[*gen.BasicBlockInfo]bool
+}
+
+func newCPReachingConstants(reachable map[*gen.BasicBlockInfo]bool) cpReachingConstants {
+	return cpReachingConstants{
+		registerStacks:   []stack.Stack[*gen.ImmediateInfo]{},
+		registerPushes:   stack.New[uint](),
+		registersToIndex: make(map[*gen.RegisterInfo]uint),
+		reachable:        reachable,
+	}
+}
+
+// get returns the constant value of the reaching definition for reg at the
+// current DFS position, or nil if no constant reaching definition is known.
+func (c *cpReachingConstants) get(reg *gen.RegisterInfo) *gen.ImmediateInfo {
+	idx, ok := c.registersToIndex[reg]
+	if !ok {
+		return nil
+	}
+	stk := c.registerStacks[idx]
+	if len(stk) == 0 {
+		return nil
+	}
+	return stk[len(stk)-1]
+}
+
+// define records a new constant reaching definition for reg in currentBlock.
+//
+// A constant is only pushed when every reachable definition of reg is in
+// currentBlock. Definitions in unreachable blocks are ignored. This covers
+// two cases:
+//   - Exactly one reachable definition, in currentBlock: it dominates all
+//     uses, so the stack value is the true reaching value.
+//   - Multiple reachable definitions, all in currentBlock: sequential
+//     redefinitions within the same block; later ones shadow earlier ones.
+//
+// Any reachable definition in a different block disqualifies the register,
+// because DFS ancestry does not imply domination at join points.
+func (c *cpReachingConstants) define(
+	reg *gen.RegisterInfo,
+	imm *gen.ImmediateInfo,
+	currentBlock *gen.BasicBlockInfo,
+) {
+	if imm == nil {
+		return
+	}
+
+	for _, def := range reg.Definitions {
+		if c.reachable[def.BasicBlockInfo] && def.BasicBlockInfo != currentBlock {
+			return
+		}
+	}
+
+	idx, ok := c.registersToIndex[reg]
+	if !ok {
+		idx = uint(len(c.registerStacks))
+		c.registersToIndex[reg] = idx
+		c.registerStacks = append(c.registerStacks, stack.New[*gen.ImmediateInfo]())
+	}
+	c.registerPushes.Push(idx)
+	c.registerStacks[idx].Push(imm)
+}
+
+func (c *cpReachingConstants) pushBlock() {
+	c.registerPushes.Push(cpBlockSeparator)
+}
+
+func (c *cpReachingConstants) popBlock() {
+	for c.registerPushes.Top() != cpBlockSeparator {
+		idx := c.registerPushes.Top()
+		c.registerPushes.Pop()
+		c.registerStacks[idx].Pop()
+	}
+	c.registerPushes.Pop() // pop the block separator itself
+}
+
+// cpProcessInstruction performs constant propagation for one instruction:
+// it substitutes any register arguments whose reaching definition is a known
+// constant, then records new reaching constants for targets whose value is
+// now known after the substitution.
+func cpProcessInstruction(
+	instruction *gen.InstructionInfo,
+	reaching *cpReachingConstants,
+) core.ResultList {
+	cpInstruction, ok := instruction.Definition.(ConstantPropagationSupportedInstruction)
+	if !ok {
+		return newCPNotSupportedError(instruction)
+	}
+
+	// Substitute arguments whose reaching definition is a known constant.
+	for i, arg := range instruction.Arguments {
+		regArg, ok := arg.(*gen.RegisterArgumentInfo)
+		if !ok {
+			continue
+		}
+		if imm := reaching.get(regArg.Register); imm != nil {
+			instruction.SubstituteArgument(i, imm)
+		}
+	}
+
+	// After substitution, record the reaching constant for each target.
+	for _, def := range cpInstruction.PropagateConstants(instruction) {
+		reaching.define(def.Register, def.Immediate, instruction.BasicBlockInfo)
+	}
+
+	return core.ResultList{}
+}
+
+// ConstantPropagation replaces all uses of registers that are provably
+// assigned a constant immediate value with that immediate directly.
+//
+// The pass performs a DFS from the entry block (block 0), visiting only
+// reachable blocks. Unreachable blocks are skipped entirely.
+//
+// Within the DFS tree, define() only pushes a constant for a register when
+// all of its definitions are confined to the current block. This keeps the
+// analysis sound: registers defined in multiple blocks are skipped because
+// DFS ancestry does not imply domination at join points.
+//
+// All instructions in the function must implement
+// ConstantPropagationSupportedInstruction.
+func ConstantPropagation(function *gen.FunctionInfo) core.ResultList {
+	cfInfo := gen.NewFunctionControlFlowInfo(function)
+	n := uint(len(cfInfo.BasicBlocks))
+	dfs := cfInfo.ControlFlowGraph.Dfs(0)
+
+	reachable := make(map[*gen.BasicBlockInfo]bool, n)
+	for _, event := range dfs.Timeline {
+		if event < n {
+			reachable[cfInfo.BasicBlocks[event]] = true
+		}
+	}
+
+	reaching := newCPReachingConstants(reachable)
+	results := core.ResultList{}
+
+	for _, event := range dfs.Timeline {
+		if event >= n {
+			reaching.popBlock()
+			continue
+		}
+
+		reaching.pushBlock()
+		block := cfInfo.BasicBlocks[event]
+		for _, instruction := range block.Instructions {
+			curResults := cpProcessInstruction(instruction, &reaching)
+			results.Extend(&curResults)
+		}
+	}
+
+	return results
+}

opt/constant_propagation_test.go

@@ -0,0 +1,11 @@
+package opt_test
+
+import (
+	"testing"
+
+	"alon.kr/x/usm/opt"
+)
+
+func TestConstantPropagation(t *testing.T) {
+	RunOptimizationTests(t, "constant_propagation", opt.ConstantPropagation)
+}

opt/testdata/constant_propagation/arithmetic.usm

@@ -0,0 +1,17 @@
+func @use $64 %reg
+
+func @input $64 %param {
+.entry
+	$64 %x = $64 #3
+	$64 %z = add %x %param
+	call @use %z
+	ret
+}
+
+func @expected $64 %param {
+.entry
+	$64 %x = $64 #3
+	$64 %z = add $64 #3 %param
+	call @use %z
+	ret
+}

opt/testdata/constant_propagation/branch_constant.usm

@@ -0,0 +1,31 @@
+func @use $64 %reg
+
+; %x = #5 (non-zero) is a constant. CP substitutes it into the jnz condition
+; and into both branch arms.
+func @input {
+.entry
+	$64 %x = $64 #5
+	jnz %x .left
+.right
+	call @use %x
+	j .exit
+.left
+	call @use %x
+	j .exit
+.exit
+	ret
+}
+
+func @expected {
+.entry
+	$64 %x = $64 #5
+	jnz $64 #5 .left
+.right
+	call @use $64 #5
+	j .exit
+.left
+	call @use $64 #5
+	j .exit
+.exit
+	ret
+}

opt/testdata/constant_propagation/chain.usm

@@ -0,0 +1,17 @@
+func @use $64 %reg
+
+func @input {
+.entry
+	$64 %x = $64 #5
+	$64 %y = %x
+	call @use %y
+	ret
+}
+
+func @expected {
+.entry
+	$64 %x = $64 #5
+	$64 %y = $64 #5
+	call @use $64 #5
+	ret
+}