BE-239: HashQL: Implement DSE (Dead Store Elimination) pass for MIR (#8184)

Author: indietyp

libs/@local/hashql/compiletest/src/suite/mir_pass_transform_dse.rs

@@ -0,0 +1,120 @@
+use std::io::Write as _;
+
+use hashql_ast::node::expr::Expr;
+use hashql_core::{heap::Heap, r#type::environment::Environment};
+use hashql_diagnostics::DiagnosticIssues;
+use hashql_mir::{
+    body::Body,
+    context::MirContext,
+    def::{DefId, DefIdSlice, DefIdVec},
+    intern::Interner,
+    pass::{TransformPass as _, transform::DeadStoreElimination},
+};
+
+use super::{
+    RunContext, Suite, SuiteDiagnostic, common::process_issues,
+    mir_pass_transform_sroa::mir_pass_transform_sroa,
+};
+use crate::suite::{
+    common::Header,
+    mir_pass_transform_cfg_simplify::mir_pass_transform_cfg_simplify_default_renderer,
+    mir_reify::{d2_output_enabled, mir_format_d2, mir_format_text, mir_spawn_d2},
+};
+
+pub(crate) fn mir_pass_transform_dse<'heap>(
+    heap: &'heap Heap,
+    expr: Expr<'heap>,
+    interner: &Interner<'heap>,
+    render: impl FnOnce(&'heap Heap, &Environment<'heap>, DefId, &DefIdSlice<Body<'heap>>),
+    environment: &mut Environment<'heap>,
+    diagnostics: &mut Vec<SuiteDiagnostic>,
+) -> Result<(DefId, DefIdVec<Body<'heap>>), SuiteDiagnostic> {
+    let (root, mut bodies) =
+        mir_pass_transform_sroa(heap, expr, interner, render, environment, diagnostics)?;
+
+    let mut context = MirContext {
+        heap,
+        env: environment,
+        interner,
+        diagnostics: DiagnosticIssues::new(),
+    };
+
+    let mut pass = DeadStoreElimination::new();
+    for body in bodies.as_mut_slice() {
+        pass.run(&mut context, body);
+    }
+
+    process_issues(diagnostics, context.diagnostics)?;
+    Ok((root, bodies))
+}
+
+pub(crate) struct MirPassTransformDse;
+
+impl Suite for MirPassTransformDse {
+    fn priority(&self) -> usize {
+        1
+    }
+
+    fn name(&self) -> &'static str {
+        "mir/pass/transform/dse"
+    }
+
+    fn description(&self) -> &'static str {
+        "Dead Store Elimination in the MIR"
+    }
+
+    fn secondary_file_extensions(&self) -> &[&str] {
+        &["svg"]
+    }
+
+    fn run<'heap>(
+        &self,
+        RunContext {
+            heap,
+            diagnostics,
+            suite_directives,
+            reports,
+            secondary_outputs,
+            ..
+        }: RunContext<'_, 'heap>,
+        expr: Expr<'heap>,
+    ) -> Result<String, SuiteDiagnostic> {
+        let mut environment = Environment::new(heap);
+        let interner = Interner::new(heap);
+
+        let mut buffer = Vec::new();
+        let mut d2 = d2_output_enabled(self, suite_directives, reports).then(mir_spawn_d2);
+
+        let (root, bodies) = mir_pass_transform_dse(
+            heap,
+            expr,
+            &interner,
+            mir_pass_transform_cfg_simplify_default_renderer(
+                &mut buffer,
+                d2.as_mut().map(|(writer, _)| writer),
+            ),
+            &mut environment,
+            diagnostics,
+        )?;
+
+        let _ = writeln!(buffer, "\n{}\n", Header::new("MIR after DSE"));
+        mir_format_text(heap, &environment, &mut buffer, root, &bodies);
+
+        if let Some((mut writer, handle)) = d2 {
+            writeln!(writer, "final: 'MIR after DSE' {{")
+                .expect("should be able to write to buffer");
+            mir_format_d2(heap, &environment, &mut writer, root, &bodies);
+            writeln!(writer, "}}").expect("should be able to write to buffer");
+
+            writer.flush().expect("should be able to write to buffer");
+            drop(writer);
+
+            let diagram = handle.join().expect("should be able to join handle");
+            let diagram = String::from_utf8_lossy_owned(diagram);
+
+            secondary_outputs.insert("svg", diagram);
+        }
+
+        Ok(String::from_utf8_lossy_owned(buffer))
+    }
+}

libs/@local/hashql/compiletest/src/suite/mod.rs

@@ -22,6 +22,7 @@ mod hir_lower_thunking;
 mod hir_reify;
 mod mir_pass_analysis_data_dependency;
 mod mir_pass_transform_cfg_simplify;
+mod mir_pass_transform_dse;
 mod mir_pass_transform_sroa;
 mod mir_reify;
 mod parse_syntax_dump;
@@ -54,8 +55,8 @@ use self::{
     hir_lower_thunking::HirLowerThunkingSuite, hir_reify::HirReifySuite,
     mir_pass_analysis_data_dependency::MirPassAnalysisDataDependency,
     mir_pass_transform_cfg_simplify::MirPassTransformCfgSimplify,
-    mir_pass_transform_sroa::MirPassTransformSroa, mir_reify::MirReifySuite,
-    parse_syntax_dump::ParseSyntaxDumpSuite,
+    mir_pass_transform_dse::MirPassTransformDse, mir_pass_transform_sroa::MirPassTransformSroa,
+    mir_reify::MirReifySuite, parse_syntax_dump::ParseSyntaxDumpSuite,
 };
 use crate::executor::TrialError;
 
@@ -151,6 +152,7 @@ const SUITES: &[&dyn Suite] = &[
     &HirReifySuite,
     &MirPassAnalysisDataDependency,
     &MirPassTransformCfgSimplify,
+    &MirPassTransformDse,
     &MirPassTransformSroa,
     &MirReifySuite,
     &ParseSyntaxDumpSuite,

libs/@local/hashql/core/src/id/bit_vec/mod.rs

@@ -25,15 +25,17 @@
 //! - Removed used of `static_assert_size` macro in favor of `assert_eq!`.
 //! - Removed mentions of `IntervalSet`.
 //! - Implement `MixedBitSet::intersect`.
+//! - Implement `DenseBitSet::negate`
 #![expect(
     clippy::integer_division,
     clippy::integer_division_remainder_used,
     clippy::cast_possible_truncation,
     clippy::too_long_first_doc_paragraph
 )]
 
-use alloc::rc::Rc;
+use alloc::{alloc::Global, rc::Rc};
 use core::{
+    alloc::Allocator,
     fmt, iter,
     marker::PhantomData,
     ops::{BitAnd, BitAndAssign, BitOrAssign, Bound, Not, Range, RangeBounds, Shl},
@@ -359,6 +361,11 @@ impl<T: Id> DenseBitSet<T> {
         // out-of-domain bits, so we need to clear them.
         self.clear_excess_bits();
     }
+
+    pub fn negate(&mut self) {
+        self.words.iter_mut().for_each(|word| *word = !*word);
+        self.clear_excess_bits();
+    }
 }
 
 // dense REL dense
@@ -1751,9 +1758,9 @@ impl<R: Id, C: Id> fmt::Debug for BitMatrix<R, C> {
 /// `R` and `C` are index types used to identify rows and columns respectively;
 /// typically newtyped `usize` wrappers, but they can also just be `usize`.
 #[derive(Clone)]
-pub struct SparseBitMatrix<R, C> {
+pub struct SparseBitMatrix<R, C, A: Allocator = Global> {
     num_columns: usize,
-    rows: IdVec<R, Option<DenseBitSet<C>>>,
+    rows: IdVec<R, Option<DenseBitSet<C>>, A>,
 }
 
 impl<R: Id, C: Id> SparseBitMatrix<R, C> {
@@ -1765,6 +1772,17 @@ impl<R: Id, C: Id> SparseBitMatrix<R, C> {
             rows: IdVec::new(),
         }
     }
+}
+
+impl<R: Id, C: Id, A: Allocator> SparseBitMatrix<R, C, A> {
+    /// Creates a new empty sparse bit matrix with no rows or columns.
+    #[must_use]
+    pub const fn new_in(num_columns: usize, alloc: A) -> Self {
+        Self {
+            num_columns,
+            rows: IdVec::new_in(alloc),
+        }
+    }
 
     fn ensure_row(&mut self, row: R) -> &mut DenseBitSet<C> {
         // Instantiate any missing rows up to and including row `row` with an empty `DenseBitSet`.
@@ -1848,6 +1866,20 @@ impl<R: Id, C: Id> SparseBitMatrix<R, C> {
         self.rows.get(row)?.as_ref()
     }
 
+    pub fn superset_row(&self, row: R, other: &DenseBitSet<C>) -> Option<bool> {
+        match self.rows.get(row) {
+            Some(Some(row)) => Some(row.superset(other)),
+            _ => None,
+        }
+    }
+
+    pub fn subset_row(&self, row: R, other: &DenseBitSet<C>) -> Option<bool> {
+        match self.rows.get(row) {
+            Some(Some(row)) => Some(other.superset(row)),
+            _ => None,
+        }
+    }
+
     /// Intersects `row` with `set`. `set` can be either `DenseBitSet` or
     /// `ChunkedBitSet`. Has no effect if `row` does not exist.
     ///

libs/@local/hashql/core/src/id/vec.rs

@@ -370,7 +370,7 @@ where
 }
 
 // Map-like APIs for IdVec<I, Option<T>>
-impl<I, T> IdVec<I, Option<T>>
+impl<I, T, A: Allocator> IdVec<I, Option<T>, A>
 where
     I: Id,
 {

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

@@ -7,6 +7,7 @@
 use hashql_core::{id, span::SpanId, symbol::Symbol, r#type::TypeId};
 
 id::newtype!(
+    #[steppable]
     /// A unique identifier for a local variable in the HashQL MIR.
     ///
     /// Local variables represent storage locations within a function's execution context.

libs/@local/hashql/mir/src/lib.rs

@@ -21,6 +21,8 @@
     iter_intersperse,
     string_from_utf8_lossy_owned,
     try_trait_v2,
+    step_trait,
+    exact_length_collection
 )]
 #![expect(clippy::indexing_slicing)]
 extern crate alloc;

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

@@ -2,7 +2,7 @@ use alloc::alloc::Global;
 use core::{alloc::Allocator, fmt};
 
 use hashql_core::{
-    graph::{LinkedGraph, NodeId, linked::IncidentEdges},
+    graph::{LinkedGraph, NodeId, Predecessors as _, Successors as _, linked::IncidentEdges},
     id::Id as _,
     intern::Interned,
     symbol::Symbol,
@@ -216,6 +216,14 @@ impl<'heap, A: Allocator> TransientDataDependencyGraph<'heap, A> {
     {
         self.graph.resolve(interner, place)
     }
+
+    pub fn depends_on(&self, local: Local) -> impl Iterator<Item = Local> {
+        self.graph.depends_on(local)
+    }
+
+    pub fn dependent_on(&self, local: Local) -> impl Iterator<Item = Local> {
+        self.graph.dependent_on(local)
+    }
 }
 
 impl<A: Allocator> fmt::Display for TransientDataDependencyGraph<'_, A> {
@@ -236,6 +244,18 @@ pub struct DataDependencyGraph<'heap, A: Allocator = Global> {
 }
 
 impl<'heap, A: Allocator> DataDependencyGraph<'heap, A> {
+    pub fn depends_on(&self, local: Local) -> impl Iterator<Item = Local> {
+        self.graph
+            .successors(NodeId::new(local.as_usize()))
+            .map(|node| Local::new(node.as_usize()))
+    }
+
+    pub fn dependent_on(&self, local: Local) -> impl Iterator<Item = Local> {
+        self.graph
+            .predecessors(NodeId::new(local.as_usize()))
+            .map(|node| Local::new(node.as_usize()))
+    }
+
     /// Creates a transient graph with all edges resolved to their ultimate sources.
     ///
     /// Each edge in the original graph is traced through [`resolve`](Self::resolve), producing