BE-303: HashQL: Split MIR basic blocks by target support (#8354)

Author: indietyp

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

@@ -27,6 +27,7 @@
 //! - Implement `MixedBitSet::intersect`.
 //! - Implement `DenseBitSet::negate`
 //! - Implement `DenseBitSet::first_unset`
+//! - Re-implement `FiniteBitSet`
 #![expect(
     clippy::integer_division,
     clippy::integer_division_remainder_used,
@@ -39,14 +40,16 @@ use core::{
     alloc::Allocator,
     fmt, iter,
     marker::PhantomData,
-    ops::{BitAnd, BitAndAssign, BitOrAssign, Bound, Not, Range, RangeBounds, Shl},
+    ops::{Bound, Range, RangeBounds},
     slice,
 };
 
 use smallvec::{SmallVec, smallvec};
 
+pub use self::finite::{FiniteBitIter, FiniteBitSet};
 use super::{Id, IdVec};
 
+mod finite;
 #[cfg(test)]
 mod tests;
 
@@ -1995,112 +1998,3 @@ const fn max_bit(word: Word) -> usize {
 fn count_ones(words: &[Word]) -> usize {
     words.iter().map(|word| word.count_ones() as usize).sum()
 }
-
-/// Integral type used to represent the bit set.
-pub trait FiniteBitSetTy:
-    BitAnd<Output = Self>
-    + BitAndAssign
-    + BitOrAssign
-    + Clone
-    + Copy
-    + Shl
-    + Not<Output = Self>
-    + PartialEq
-    + Sized
-{
-    /// Size of the domain representable by this type, e.g. 64 for `u64`.
-    const DOMAIN_SIZE: u32;
-
-    /// Value which represents the `FiniteBitSet` having every bit set.
-    const FILLED: Self;
-    /// Value which represents the `FiniteBitSet` having no bits set.
-    const EMPTY: Self;
-
-    /// Value for one as the integral type.
-    const ONE: Self;
-    /// Value for zero as the integral type.
-    const ZERO: Self;
-
-    /// Perform a checked left shift on the integral type.
-    fn checked_shl(self, rhs: u32) -> Option<Self>;
-    /// Perform a checked right shift on the integral type.
-    fn checked_shr(self, rhs: u32) -> Option<Self>;
-}
-
-impl FiniteBitSetTy for u32 {
-    const DOMAIN_SIZE: Self = 32;
-    const EMPTY: Self = Self::MIN;
-    const FILLED: Self = Self::MAX;
-    const ONE: Self = 1_u32;
-    const ZERO: Self = 0_u32;
-
-    fn checked_shl(self, rhs: u32) -> Option<Self> {
-        self.checked_shl(rhs)
-    }
-
-    fn checked_shr(self, rhs: u32) -> Option<Self> {
-        self.checked_shr(rhs)
-    }
-}
-
-impl fmt::Debug for FiniteBitSet<u32> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{:032b}", self.0)
-    }
-}
-
-/// A fixed-sized bitset type represented by an integer type. Indices outwith than the range
-/// representable by `T` are considered set.
-#[derive(Copy, Clone, Eq, PartialEq)]
-pub struct FiniteBitSet<T: FiniteBitSetTy>(pub T);
-
-impl<T: FiniteBitSetTy> FiniteBitSet<T> {
-    /// Creates a new, empty bitset.
-    pub const fn new_empty() -> Self {
-        Self(T::EMPTY)
-    }
-
-    /// Sets the `index`th bit.
-    pub fn set(&mut self, index: u32) {
-        self.0 |= T::ONE.checked_shl(index).unwrap_or(T::ZERO);
-    }
-
-    /// Unsets the `index`th bit.
-    pub fn clear(&mut self, index: u32) {
-        self.0 &= !T::ONE.checked_shl(index).unwrap_or(T::ZERO);
-    }
-
-    /// Sets the `i`th to `j`th bits.
-    pub fn set_range(&mut self, range: Range<u32>) {
-        let bits = T::FILLED
-            .checked_shl(range.end - range.start)
-            .unwrap_or(T::ZERO)
-            .not()
-            .checked_shl(range.start)
-            .unwrap_or(T::ZERO);
-        self.0 |= bits;
-    }
-
-    /// Is the set empty?
-    pub fn is_empty(&self) -> bool {
-        self.0 == T::EMPTY
-    }
-
-    /// Returns the domain size of the bitset.
-    #[must_use]
-    pub const fn within_domain(index: u32) -> bool {
-        index < T::DOMAIN_SIZE
-    }
-
-    /// Returns if the `index`th bit is set.
-    pub fn contains(&self, index: u32) -> Option<bool> {
-        Self::within_domain(index)
-            .then(|| ((self.0.checked_shr(index).unwrap_or(T::ONE)) & T::ONE) == T::ONE)
-    }
-}
-
-impl<T: FiniteBitSetTy> Default for FiniteBitSet<T> {
-    fn default() -> Self {
-        Self::new_empty()
-    }
-}

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

@@ -1,3 +1,4 @@
+mod array;
 pub mod bit_vec;
 mod index;
 mod slice;
@@ -13,7 +14,9 @@ use core::{
 
 use ::core::sync::atomic;
 
-pub use self::{index::IntoSliceIndex, slice::IdSlice, union_find::IdUnionFind, vec::IdVec};
+pub use self::{
+    array::IdArray, index::IntoSliceIndex, slice::IdSlice, union_find::IdUnionFind, vec::IdVec,
+};
 
 /// Represents errors that can occur when converting values to an [`Id`].
 ///
@@ -258,8 +261,8 @@ macro_rules! newtype {
         $crate::id::newtype!(@parse_attrs [$($other)* #[$attr]] [$($step)*] [$($display)*] ; $(#[$($rest)*])* ; $($tail)*);
     };
 
-    (@parse_attrs [$($other:tt)*] [$($step:tt)*] [$($display:tt)*]; ; $vis:vis struct $name:ident($type:ident is $min:literal..=$max:expr)) => {
-        $crate::id::newtype!(@impl [$($other)*] [$($step)*] [$($display)*] $vis struct $name($type is $min..=$max));
+    (@parse_attrs [$($other:tt)*] [$($step:tt)*] [$($display:tt)*]; ; $($tail:tt)*) => {
+        $crate::id::newtype!(@impl [$($other)*] [$($step)*] [$($display)*] $($tail)*);
     };
 
     // Implementation

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

@@ -249,6 +249,11 @@ where
     pub fn swap(&mut self, lhs: I, rhs: I) {
         self.raw.swap(lhs.as_usize(), rhs.as_usize());
     }
+
+    #[inline]
+    pub fn windows<const N: usize>(&self) -> impl ExactSizeIterator<Item = &[T; N]> {
+        self.raw.array_windows()
+    }
 }
 
 impl<I, T> IdSlice<I, Option<T>>

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

@@ -22,6 +22,7 @@
     binary_heap_into_iter_sorted,
     clone_from_ref,
     const_cmp,
+    const_ops,
     const_trait_impl,
     debug_closure_helpers,
     extend_one,

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

@@ -31,7 +31,8 @@
     string_from_utf8_lossy_owned,
     try_trait_v2,
     temporary_niche_types,
-    const_convert
+    const_convert,
+    variant_count
 )]
 #![expect(clippy::indexing_slicing)]
 extern crate alloc;

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

@@ -15,7 +15,7 @@ use hashql_core::id::{Id as _, bit_vec::DenseBitSet};
 use crate::{
     body::{
         Body,
-        basic_block::BasicBlockSlice,
+        basic_block::{BasicBlockId, BasicBlockSlice},
         basic_blocks::BasicBlocks,
         local::{Local, LocalVec},
         location::Location,
@@ -140,11 +140,12 @@ pub struct StatementCostVec<A: Allocator = Global> {
 
 impl<A: Allocator> StatementCostVec<A> {
     #[expect(unsafe_code)]
-    fn from_iter(mut iter: impl ExactSizeIterator<Item = u32>, alloc: A) -> Self
-    where
-        A: Clone,
-    {
-        let mut offsets = Box::new_uninit_slice_in(iter.len() + 1, alloc.clone());
+    fn offsets(
+        mut iter: impl ExactSizeIterator<Item = u32>,
+        alloc: A,
+    ) -> (Box<BasicBlockSlice<u32>, A>, usize) {
+        // Try to reuse existing offsets if available and of correct length
+        let mut offsets = Box::new_uninit_slice_in(iter.len() + 1, alloc);
 
         let mut offset = 0_u32;
 
@@ -161,12 +162,20 @@ impl<A: Allocator> StatementCostVec<A> {
         debug_assert!(rest.is_empty());
         debug_assert_eq!(iter.len(), 0);
 
-        let costs = alloc::vec::from_elem_in(None, offset as usize, alloc);
-
         // SAFETY: We have initialized all elements of the slice.
         let offsets = unsafe { offsets.assume_init() };
         let offsets = BasicBlockSlice::from_boxed_slice(offsets);
 
+        (offsets, offset as usize)
+    }
+
+    fn from_iter(iter: impl ExactSizeIterator<Item = u32>, alloc: A) -> Self
+    where
+        A: Clone,
+    {
+        let (offsets, length) = Self::offsets(iter, alloc.clone());
+        let costs = alloc::vec::from_elem_in(None, length, alloc);
+
         Self { offsets, costs }
     }
 
@@ -184,10 +193,34 @@ impl<A: Allocator> StatementCostVec<A> {
         )
     }
 
+    #[expect(clippy::cast_possible_truncation)]
+    pub fn remap(&mut self, blocks: &BasicBlocks)
+    where
+        A: Clone,
+    {
+        let alloc = Box::allocator(&self.offsets).clone();
+
+        let (offsets, _) = Self::offsets(
+            blocks.iter().map(|block| block.statements.len() as u32),
+            alloc,
+        );
+        self.offsets = offsets;
+    }
+
     pub fn all_unassigned(&self) -> bool {
         self.costs.iter().all(Option::is_none)
     }
 
+    pub fn of(&self, block: BasicBlockId) -> &[Option<Cost>] {
+        let range = (self.offsets[block] as usize)..(self.offsets[block.plus(1)] as usize);
+
+        &self.costs[range]
+    }
+
+    pub fn allocator(&self) -> &A {
+        Box::allocator(&self.offsets)
+    }
+
     /// Returns the cost at `location`, or `None` if out of bounds or unassigned.
     pub fn get(&self, location: Location) -> Option<Cost> {
         let range = (self.offsets[location.block] as usize)

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

@@ -5,6 +5,7 @@ macro_rules! cost {
 }
 
 mod cost;
+pub mod splitting;
 pub mod statement_placement;
 pub mod target;