feat: extend copy propagation to locals

Author: indietyp

libs/@local/hashql/mir/src/body/operand.rs

@@ -3,7 +3,7 @@
 //! Operands represent the inputs to MIR operations. They can either reference
 //! a storage location (place) or contain an immediate constant value.
 
-use super::{constant::Constant, place::Place};
+use super::{constant::Constant, local::Local, place::Place};
 
 /// An operand in a HashQL MIR operation.
 ///
@@ -59,6 +59,12 @@ impl<'heap> Operand<'heap> {
     }
 }
 
+impl From<Local> for Operand<'_> {
+    fn from(local: Local) -> Self {
+        Operand::Place(Place::local(local))
+    }
+}
+
 impl<'heap> From<Place<'heap>> for Operand<'heap> {
     fn from(place: Place<'heap>) -> Self {
         Operand::Place(place)

libs/@local/hashql/mir/src/body/place.rs

@@ -315,10 +315,11 @@ impl<'heap> Place<'heap> {
     /// This is the simplest form of a place, representing direct access to a local variable
     /// without navigating through any structured data. The resulting place has an empty
     /// projection sequence.
-    pub fn local(local: Local, interner: &Interner<'heap>) -> Self {
+    #[must_use]
+    pub const fn local(local: Local) -> Self {
         Self {
             local,
-            projections: interner.projections.intern_slice(&[]),
+            projections: Interned::empty(),
         }
     }
 

libs/@local/hashql/mir/src/builder/body.rs

@@ -70,7 +70,7 @@ impl<'env, 'heap> BodyBuilder<'env, 'heap> {
         };
         let local = self.local_decls.push(decl);
 
-        Place::local(local, self.interner)
+        Place::local(local)
     }
 
     /// Reserves a new basic block and returns its ID.

libs/@local/hashql/mir/src/pass/analysis/data_dependency/resolve.rs

@@ -299,7 +299,7 @@ fn resolve_params_const<'heap, A: Allocator + Clone>(
     } else {
         // We have finished (we have terminated on a param, which is divergent, therefore the place
         // is still valid, just doesn't have a constant value)
-        ResolutionResult::Resolved(Operand::Place(Place::local(place.local, state.interner)))
+        ResolutionResult::Resolved(Operand::Place(Place::local(place.local)))
     }
 }
 
@@ -382,15 +382,14 @@ pub(crate) fn resolve<'heap, A: Allocator + Clone>(
         [] => {
             // Base case: no more projections to resolve.
             // Check for constant propagation through Load.
-            let operand = if let Some(constant) = state
+            let operand = state
                 .graph
                 .constant_bindings
                 .find_by_kind(place.local, EdgeKind::Load)
-            {
-                Operand::Constant(constant)
-            } else {
-                Operand::Place(Place::local(place.local, state.interner))
-            };
+                .map_or_else(
+                    || Operand::Place(Place::local(place.local)),
+                    Operand::Constant,
+                );
 
             return ResolutionResult::Resolved(operand);
         }

libs/@local/hashql/mir/src/pass/transform/administrative_reduction/visitor.rs

@@ -221,7 +221,7 @@ impl<'heap, A: Allocator> AdministrativeReductionVisitor<'_, '_, 'heap, A> {
             .enumerate()
             .map(|(param, argument)| Statement {
                 kind: StatementKind::Assign(Assign {
-                    lhs: Place::local(Local::new(local_offset + param), self.interner),
+                    lhs: Place::local(Local::new(local_offset + param)),
                     rhs: RValue::Load(argument),
                 }),
                 span,

libs/@local/hashql/mir/src/pass/transform/cfg_simplify/mod.rs

@@ -122,12 +122,7 @@ impl<A: BumpAllocator> CfgSimplify<A> {
     /// 3. Replace `A`'s terminator with `B`'s terminator
     ///
     /// SSA invariants may be temporarily broken; the [`SsaRepair`] runs afterward to fix them.
-    fn simplify_goto<'heap>(
-        context: &MirContext<'_, 'heap>,
-        body: &mut Body<'heap>,
-        id: BasicBlockId,
-        goto: Goto<'heap>,
-    ) -> bool {
+    fn simplify_goto<'heap>(body: &mut Body<'heap>, id: BasicBlockId, goto: Goto<'heap>) -> bool {
         // Self-loops cannot be optimized as there's no simplification possible.
         if goto.target.block == id {
             return false;
@@ -170,7 +165,7 @@ impl<A: BumpAllocator> CfgSimplify<A> {
             block.statements.push(Statement {
                 span: block.terminator.span,
                 kind: StatementKind::Assign(Assign {
-                    lhs: Place::local(param, context.interner),
+                    lhs: Place::local(param),
                     rhs: RValue::Load(arg),
                 }),
             });
@@ -424,7 +419,7 @@ impl<A: BumpAllocator> CfgSimplify<A> {
             .transfer_into(&self.alloc);
 
         let changed = match &body.basic_blocks[id].terminator.kind {
-            &TerminatorKind::Goto(goto) => Self::simplify_goto(context, body, id, goto),
+            &TerminatorKind::Goto(goto) => Self::simplify_goto(body, id, goto),
             TerminatorKind::SwitchInt(_) => Self::simplify_switch_int(context, body, id),
             TerminatorKind::Return(_)
             | TerminatorKind::GraphRead(_)

libs/@local/hashql/mir/src/pass/transform/copy_propagation/mod.rs

@@ -1,27 +1,45 @@
 //! Copy and constant propagation transformation pass.
 //!
-//! This pass propagates constant values through the MIR by tracking which locals hold known
-//! constants and substituting uses of those locals with the constant values directly.
+//! This pass propagates both constants and copies through the MIR by tracking which locals hold
+//! known values (either constants or references to other locals) and substituting uses accordingly.
 //!
 //! Unlike [`ForwardSubstitution`], this pass does not perform full data dependency analysis and
 //! cannot resolve values through projections or chained access paths. It is faster but less
-//! comprehensive, making it suitable for quick constant folding in simpler cases.
+//! comprehensive, making it suitable for quick propagation in simpler cases.
 //!
 //! # Algorithm
 //!
 //! The pass operates in a single forward traversal (reverse postorder):
 //!
-//! 1. For each block, propagates constants through block parameters when all predecessors pass the
-//!    same constant value
-//! 2. For each assignment `_x = <operand>`, if the operand is a constant or a local known to hold a
-//!    constant, records that `_x` holds that constant
-//! 3. For each use of a local known to hold a constant, substitutes the use with the constant
+//! 1. For each block, propagates values through block parameters when all predecessors pass the
+//!    same value (constant or local)
+//! 2. For each assignment `_x = <operand>`, records what `_x` holds:
+//!    - If the operand is a constant, records `_x → constant`
+//!    - If the operand is a local (possibly with a known value), records `_x → known value`
+//! 3. For each use of a local with a known value, substitutes the use with that value
+//!
+//! # Examples
+//!
+//! Constant propagation:
+//! ```text
+//! _1 = const 42; use(_1)  →  _1 = const 42; use(const 42)
+//! ```
+//!
+//! Copy propagation:
+//! ```text
+//! _2 = _1; use(_2)  →  _2 = _1; use(_1)
+//! ```
+//!
+//! Chained propagation:
+//! ```text
+//! _2 = _1; _3 = _2; use(_3)  →  _2 = _1; _3 = _1; use(_1)
+//! ```
 //!
 //! # Limitations
 //!
 //! - Does not handle projections: `_2 = (_1,); use(_2.0)` is not simplified
 //! - Does not perform fix-point iteration for loops
-//! - Only tracks constants, not arbitrary value equivalences
+//! - Assumes SSA-like semantics (locals are assigned at most once)
 //!
 //! For more comprehensive value propagation including projections, see [`ForwardSubstitution`].
 //!
@@ -46,6 +64,7 @@ use crate::{
         local::{Local, LocalVec},
         location::Location,
         operand::Operand,
+        place::Place,
         rvalue::RValue,
         statement::Assign,
     },
@@ -165,6 +184,21 @@ where
     .flatten()
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+enum KnownValue<'heap> {
+    Constant(Constant<'heap>),
+    Local(Local),
+}
+
+impl<'heap> From<KnownValue<'heap>> for Operand<'heap> {
+    fn from(value: KnownValue<'heap>) -> Self {
+        match value {
+            KnownValue::Constant(constant) => Operand::Constant(constant),
+            KnownValue::Local(local) => Operand::Place(Place::local(local)),
+        }
+    }
+}
+
 pub struct CopyPropagation<A: BumpAllocator = Scratch> {
     alloc: A,
 }
@@ -196,7 +230,7 @@ impl<'env, 'heap, A: ResetAllocator> TransformPass<'env, 'heap> for CopyPropagat
 
         let mut visitor = CopyPropagationVisitor {
             interner: context.interner,
-            constants: IdVec::with_capacity_in(body.local_decls.len(), &self.alloc),
+            values: IdVec::with_capacity_in(body.local_decls.len(), &self.alloc),
             changed: false,
         };
 
@@ -208,10 +242,10 @@ impl<'env, 'heap, A: ResetAllocator> TransformPass<'env, 'heap> for CopyPropagat
         let mut args = Vec::new_in(&self.alloc);
 
         for &mut id in reverse_postorder {
-            for (local, constant) in
+            for (local, value) in
                 propagate_block_params(&mut args, body, id, |operand| visitor.try_eval(operand))
             {
-                visitor.constants.insert(local, constant);
+                visitor.values.insert(local, value);
             }
 
             Ok(()) =
@@ -224,7 +258,7 @@ impl<'env, 'heap, A: ResetAllocator> TransformPass<'env, 'heap> for CopyPropagat
 
 struct CopyPropagationVisitor<'env, 'heap, A: Allocator> {
     interner: &'env Interner<'heap>,
-    constants: LocalVec<Option<Constant<'heap>>, A>,
+    values: LocalVec<Option<KnownValue<'heap>>, A>,
     changed: bool,
 }
 
@@ -233,19 +267,21 @@ impl<'heap, A: Allocator> CopyPropagationVisitor<'_, 'heap, A> {
     ///
     /// Returns `Int` if the operand is a constant integer or a local known to hold one,
     /// `Place` if it's a non-constant place, or `Other` for operands that can't be simplified.
-    fn try_eval(&self, operand: Operand<'heap>) -> Option<Constant<'heap>> {
-        if let Operand::Constant(constant) = operand {
-            return Some(constant);
+    fn try_eval(&self, operand: Operand<'heap>) -> Option<KnownValue<'heap>> {
+        let place = match operand {
+            Operand::Place(place) => place,
+            Operand::Constant(constant) => return Some(KnownValue::Constant(constant)),
+        };
+
+        if !place.projections.is_empty() {
+            return None;
         }
 
-        if let Operand::Place(place) = operand
-            && place.projections.is_empty()
-            && let Some(&constant) = self.constants.lookup(place.local)
-        {
-            return Some(constant);
+        if let Some(&known) = self.values.lookup(place.local) {
+            return Some(known);
         }
 
-        None
+        Some(KnownValue::Local(place.local))
     }
 }
 
@@ -262,12 +298,10 @@ impl<'heap, A: Allocator> VisitorMut<'heap> for CopyPropagationVisitor<'_, 'heap
     }
 
     fn visit_operand(&mut self, _: Location, operand: &mut Operand<'heap>) -> Self::Result<()> {
-        if let Operand::Place(place) = operand
-            && place.projections.is_empty()
-            && let Some(&constant) = self.constants.lookup(place.local)
-        {
-            *operand = Operand::Constant(constant);
-            self.changed = true;
+        if let Some(known) = self.try_eval(*operand) {
+            let known: Operand<'heap> = known.into();
+            self.changed |= known != *operand;
+            *operand = known;
         }
 
         Ok(())
@@ -292,16 +326,8 @@ impl<'heap, A: Allocator> VisitorMut<'heap> for CopyPropagationVisitor<'_, 'heap
             return Ok(());
         };
 
-        match load {
-            Operand::Place(place) if place.projections.is_empty() => {
-                if let Some(&constant) = self.constants.lookup(place.local) {
-                    self.constants.insert(lhs.local, constant);
-                }
-            }
-            Operand::Place(_) => {}
-            &mut Operand::Constant(constant) => {
-                self.constants.insert(lhs.local, constant);
-            }
+        if let Some(known) = self.try_eval(*load) {
+            self.values.insert(lhs.local, known);
         }
 
         Ok(())