BE-229, BE-230: HashQL: Implement SROA (Scalar Replacement of Aggregates) pass for MIR (#8183)

Author: indietyp

libs/@local/hashql/compiletest/src/suite/mir_pass_transform_sroa.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::Sroa},
+};
+
+use super::{
+    RunContext, Suite, SuiteDiagnostic, common::process_issues,
+    mir_pass_transform_cfg_simplify::mir_pass_transform_cfg_simplify,
+};
+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_sroa<'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_cfg_simplify(heap, expr, interner, render, environment, diagnostics)?;
+
+    let mut context = MirContext {
+        heap,
+        env: environment,
+        interner,
+        diagnostics: DiagnosticIssues::new(),
+    };
+
+    let mut pass = Sroa::new();
+    for body in bodies.as_mut_slice() {
+        pass.run(&mut context, body);
+    }
+
+    process_issues(diagnostics, context.diagnostics)?;
+    Ok((root, bodies))
+}
+
+pub(crate) struct MirPassTransformSroa;
+
+impl Suite for MirPassTransformSroa {
+    fn priority(&self) -> usize {
+        1
+    }
+
+    fn name(&self) -> &'static str {
+        "mir/pass/transform/sroa"
+    }
+
+    fn description(&self) -> &'static str {
+        "Scalar Replacement of Aggregates 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_sroa(
+            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 SROA"));
+        mir_format_text(heap, &environment, &mut buffer, root, &bodies);
+
+        if let Some((mut writer, handle)) = d2 {
+            writeln!(writer, "final: 'MIR after SROA' {{")
+                .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_sroa;
 mod mir_reify;
 mod parse_syntax_dump;
 
@@ -52,7 +53,8 @@ use self::{
     hir_lower_specialization::HirLowerSpecializationSuite,
     hir_lower_thunking::HirLowerThunkingSuite, hir_reify::HirReifySuite,
     mir_pass_analysis_data_dependency::MirPassAnalysisDataDependency,
-    mir_pass_transform_cfg_simplify::MirPassTransformCfgSimplify, mir_reify::MirReifySuite,
+    mir_pass_transform_cfg_simplify::MirPassTransformCfgSimplify,
+    mir_pass_transform_sroa::MirPassTransformSroa, mir_reify::MirReifySuite,
     parse_syntax_dump::ParseSyntaxDumpSuite,
 };
 use crate::executor::TrialError;
@@ -149,6 +151,7 @@ const SUITES: &[&dyn Suite] = &[
     &HirReifySuite,
     &MirPassAnalysisDataDependency,
     &MirPassTransformCfgSimplify,
+    &MirPassTransformSroa,
     &MirReifySuite,
     &ParseSyntaxDumpSuite,
 ];

libs/@local/hashql/core/src/heap.rs libs/@local/hashql/core/src/heap/mod.rslibs/@local/hashql/core/src/heap.rs renamed to libs/@local/hashql/core/src/heap/mod.rs

@@ -37,12 +37,15 @@
 //! - `heap::VecDeque<'heap, T>`: A double-ended queue allocated on the heap
 //! - `heap::HashMap<'heap, K, V, S>`: A hash map allocated on the heap
 
+mod scratch;
+
 use core::{alloc::Allocator, ptr};
 use std::sync::Mutex;
 
 use bumpalo::Bump;
 use hashbrown::HashSet;
 
+pub use self::scratch::Scratch;
 use crate::{
     collections::{FastHashSet, fast_hash_set_with_capacity},
     symbol::{Symbol, sym::TABLES},

libs/@local/hashql/core/src/heap/scratch.rs

@@ -0,0 +1,106 @@
+//! Scratch allocator for temporary allocations.
+//!
+//! This module provides [`Scratch`], a resettable bump allocator designed for
+//! temporary allocations that can be bulk-freed by calling [`Scratch::reset`].
+
+use core::alloc::Allocator;
+
+use bumpalo::Bump;
+
+/// A resettable scratch allocator for temporary allocations.
+///
+/// `Scratch` wraps a [`Bump`] allocator, providing a simple interface for
+/// allocating memory that can be efficiently freed in bulk. This is useful
+/// for temporary allocations during query processing where individual
+/// deallocations are unnecessary.
+///
+/// # Usage
+///
+/// The allocator can be used directly with collections via the [`Allocator`] trait:
+///
+/// ```
+/// # #![feature(allocator_api)]
+/// # use hashql_core::heap::Scratch;
+/// let mut scratch = Scratch::new();
+/// let mut vec: Vec<u32, &Scratch> = Vec::new_in(&scratch);
+/// vec.push(42);
+/// # drop(vec);
+/// // When done, reset to free all allocations at once
+/// scratch.reset();
+/// ```
+#[derive(Debug)]
+pub struct Scratch {
+    bump: Bump,
+}
+
+impl Scratch {
+    /// Creates a new scratch allocator with an empty arena.
+    #[must_use]
+    pub fn new() -> Self {
+        Self { bump: Bump::new() }
+    }
+
+    /// Resets the allocator, freeing all allocations at once.
+    pub fn reset(&mut self) {
+        self.bump.reset();
+    }
+}
+
+impl Default for Scratch {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[expect(unsafe_code, reason = "proxy to bump")]
+// SAFETY: this simply delegates to the bump allocator
+unsafe impl Allocator for &Scratch {
+    fn allocate_zeroed(
+        &self,
+        layout: core::alloc::Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
+        (&self.bump).allocate_zeroed(layout)
+    }
+
+    unsafe fn grow(
+        &self,
+        ptr: core::ptr::NonNull<u8>,
+        old_layout: core::alloc::Layout,
+        new_layout: core::alloc::Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
+        // SAFETY: this simply delegates to the bump allocator
+        unsafe { (&self.bump).grow(ptr, old_layout, new_layout) }
+    }
+
+    unsafe fn grow_zeroed(
+        &self,
+        ptr: core::ptr::NonNull<u8>,
+        old_layout: core::alloc::Layout,
+        new_layout: core::alloc::Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
+        // SAFETY: this simply delegates to the bump allocator
+        unsafe { (&self.bump).grow_zeroed(ptr, old_layout, new_layout) }
+    }
+
+    unsafe fn shrink(
+        &self,
+        ptr: core::ptr::NonNull<u8>,
+        old_layout: core::alloc::Layout,
+        new_layout: core::alloc::Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
+        // SAFETY: this simply delegates to the bump allocator
+        unsafe { (&self.bump).shrink(ptr, old_layout, new_layout) }
+    }
+
+    fn allocate(
+        &self,
+        layout: core::alloc::Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
+        (&self.bump).allocate(layout)
+    }
+
+    unsafe fn deallocate(&self, ptr: core::ptr::NonNull<u8>, layout: core::alloc::Layout) {
+        // SAFETY: this simply delegates to the bump allocator
+        unsafe { (&self.bump).deallocate(ptr, layout) }
+    }
+}

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

@@ -24,3 +24,10 @@ pub struct Location {
     /// complete (terminator).
     pub statement_index: usize,
 }
+
+impl Location {
+    pub const PLACEHOLDER: Self = Self {
+        block: BasicBlockId::PLACEHOLDER,
+        statement_index: usize::MAX,
+    };
+}

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

@@ -1,6 +1,9 @@
 mod cfg_simplify;
 mod dbe;
 pub mod error;
+mod sroa;
 mod ssa_repair;
 
-pub use self::{cfg_simplify::CfgSimplify, dbe::DeadBlockElimination, ssa_repair::SsaRepair};
+pub use self::{
+    cfg_simplify::CfgSimplify, dbe::DeadBlockElimination, sroa::Sroa, ssa_repair::SsaRepair,
+};