Define an RAII helper for generic take-and-replace borrow splitting (#10548)

* Define an RAII helper for generic take-and-replace borrow splitting

Follow up to #10524 (comment)

* &mut T

Author: fitzgen

cranelift/codegen/src/egraph.rs

@@ -16,6 +16,7 @@ use crate::opts::generated_code::SkeletonInstSimplification;
 use crate::opts::IsleContext;
 use crate::scoped_hash_map::{Entry as ScopedEntry, ScopedHashMap};
 use crate::settings::Flags;
+use crate::take_and_replace::TakeAndReplace;
 use crate::trace;
 use alloc::vec::Vec;
 use core::cmp::Ordering;
@@ -298,7 +299,8 @@ where
         // A pure node always has exactly one result.
         let orig_value = self.func.dfg.first_result(inst);
 
-        let mut optimized_values = std::mem::take(&mut self.optimized_values);
+        let mut guard = TakeAndReplace::new(self, |x| &mut x.optimized_values);
+        let (ctx, optimized_values) = guard.get();
 
         // Limit rewrite depth. When we apply optimization rules, they
         // may create new nodes (values) and those are, recursively,
@@ -310,28 +312,28 @@ where
         // infinite or problematic recursion, we bound the rewrite
         // depth to a small constant here.
         const REWRITE_LIMIT: usize = 5;
-        if self.rewrite_depth > REWRITE_LIMIT {
-            self.stats.rewrite_depth_limit += 1;
+        if ctx.rewrite_depth > REWRITE_LIMIT {
+            ctx.stats.rewrite_depth_limit += 1;
             return orig_value;
         }
-        self.rewrite_depth += 1;
-        trace!("Incrementing rewrite depth; now {}", self.rewrite_depth);
+        ctx.rewrite_depth += 1;
+        trace!("Incrementing rewrite depth; now {}", ctx.rewrite_depth);
 
         // Invoke the ISLE toplevel constructor, getting all new
         // values produced as equivalents to this value.
         trace!("Calling into ISLE with original value {}", orig_value);
-        self.stats.rewrite_rule_invoked += 1;
+        ctx.stats.rewrite_rule_invoked += 1;
         debug_assert!(optimized_values.is_empty());
         crate::opts::generated_code::constructor_simplify(
-            &mut IsleContext { ctx: self },
+            &mut IsleContext { ctx },
             orig_value,
-            &mut optimized_values,
+            optimized_values,
         );
 
-        self.stats.rewrite_rule_results += optimized_values.len() as u64;
+        ctx.stats.rewrite_rule_results += optimized_values.len() as u64;
 
         // It's not supposed to matter what order `simplify` returns values in.
-        self.ctrl_plane.shuffle(&mut optimized_values);
+        ctx.ctrl_plane.shuffle(optimized_values);
 
         let num_matches = optimized_values.len();
         if num_matches > MATCHES_LIMIT {
@@ -351,10 +353,10 @@ where
         // all returned values.
         let result_value = if let Some(&subsuming_value) = optimized_values
             .iter()
-            .find(|&value| self.subsume_values.contains(value))
+            .find(|&value| ctx.subsume_values.contains(value))
         {
             optimized_values.clear();
-            self.stats.pure_inst_subsume += 1;
+            ctx.stats.pure_inst_subsume += 1;
             subsuming_value
         } else {
             let mut union_value = orig_value;
@@ -366,29 +368,27 @@ where
                 );
                 if optimized_value == orig_value {
                     trace!(" -> same as orig value; skipping");
-                    self.stats.pure_inst_rewrite_to_self += 1;
+                    ctx.stats.pure_inst_rewrite_to_self += 1;
                     continue;
                 }
                 let old_union_value = union_value;
-                union_value = self.func.dfg.union(old_union_value, optimized_value);
-                self.stats.union += 1;
+                union_value = ctx.func.dfg.union(old_union_value, optimized_value);
+                ctx.stats.union += 1;
                 trace!(" -> union: now {}", union_value);
-                self.func.dfg.merge_facts(old_union_value, optimized_value);
-                self.available_block[union_value] =
-                    self.merge_availability(old_union_value, optimized_value);
+                ctx.func.dfg.merge_facts(old_union_value, optimized_value);
+                ctx.available_block[union_value] =
+                    ctx.merge_availability(old_union_value, optimized_value);
             }
             union_value
         };
 
-        self.rewrite_depth -= 1;
-        trace!("Decrementing rewrite depth; now {}", self.rewrite_depth);
-        if self.rewrite_depth == 0 {
-            self.subsume_values.clear();
+        ctx.rewrite_depth -= 1;
+        trace!("Decrementing rewrite depth; now {}", ctx.rewrite_depth);
+        if ctx.rewrite_depth == 0 {
+            ctx.subsume_values.clear();
         }
 
-        debug_assert!(self.optimized_values.is_empty());
-        self.optimized_values = optimized_values;
-
+        debug_assert!(ctx.optimized_values.is_empty());
         result_value
     }
 
@@ -564,36 +564,8 @@ where
             return None;
         }
 
-        /// A small RAII helper for temporarily taking out our `optimized_insts`
-        /// vec and then replacing it upon drop.
-        struct WithOptimizedInsts<'a, 'opt, 'analysis> {
-            ctx: &'a mut OptimizeCtx<'opt, 'analysis>,
-            optimized_insts: SmallVec<[SkeletonInstSimplification; MATCHES_LIMIT]>,
-        }
-
-        impl Drop for WithOptimizedInsts<'_, '_, '_> {
-            fn drop(&mut self) {
-                self.optimized_insts.clear();
-                self.ctx.optimized_insts = std::mem::take(&mut self.optimized_insts);
-            }
-        }
-
-        impl<'a, 'b, 'c> WithOptimizedInsts<'a, 'b, 'c> {
-            fn new(ctx: &'a mut OptimizeCtx<'b, 'c>) -> Self {
-                let optimized_insts = std::mem::take(&mut ctx.optimized_insts);
-                debug_assert!(optimized_insts.is_empty());
-                WithOptimizedInsts {
-                    ctx,
-                    optimized_insts,
-                }
-            }
-        }
-
-        let mut guard = WithOptimizedInsts::new(self);
-        let WithOptimizedInsts {
-            ctx,
-            optimized_insts,
-        } = &mut guard;
+        let mut guard = TakeAndReplace::new(self, |x| &mut x.optimized_insts);
+        let (ctx, optimized_insts) = guard.get();
 
         crate::opts::generated_code::constructor_simplify_skeleton(
             &mut IsleContext { ctx },

cranelift/codegen/src/lib.rs

@@ -82,13 +82,15 @@ mod ranges;
 mod remove_constant_phis;
 mod result;
 mod scoped_hash_map;
+mod take_and_replace;
 mod unreachable_code;
 mod value_label;
 
 #[cfg(feature = "souper-harvest")]
 mod souper_harvest;
 
 pub use crate::result::{CodegenError, CodegenResult, CompileError};
+pub use crate::take_and_replace::TakeAndReplace;
 
 #[cfg(feature = "incremental-cache")]
 pub mod incremental_cache;

cranelift/codegen/src/take_and_replace.rs

@@ -0,0 +1,91 @@
+//! Helper for temporarily taking values out and then putting them back in.
+
+/// An RAII type to temporarily take a `U` out of a `T` and then put it back
+/// again on drop.
+///
+/// This allows you to split borrows, if necessary, to satisfy the borrow
+/// checker.
+///
+/// The `F` type parameter must project from the container type `T` to its `U`
+/// that we want to temporarily take out of it.
+///
+/// # Example
+///
+/// ```
+/// use cranelift_codegen::TakeAndReplace;
+///
+/// #[derive(Default)]
+/// struct BigContextStruct {
+///     items: Vec<u32>,
+///     count: usize,
+/// }
+///
+/// impl BigContextStruct {
+///     fn handle_item(&mut self, item: u32) {
+///         self.count += 1;
+///         println!("Handled {item}!");
+///     }
+/// }
+///
+/// let mut ctx = BigContextStruct::default();
+/// ctx.items.extend([42, 1337, 1312]);
+///
+/// {
+///     // Temporarily take `self.items` out of `ctx`.
+///     let mut guard = TakeAndReplace::new(&mut ctx, |ctx| &mut ctx.items);
+///     let (ctx, items) = guard.get();
+///
+///     // Now we can both borrow/iterate/mutate `items` and call `&mut self` helper
+///     // methods on `ctx`. This would not otherwise be possible if we didn't split
+///     // the borrows, since Rust's borrow checker doesn't see through methods and
+///     // know that `handle_item` doesn't use `self.items`.
+///     for item in items.drain(..) {
+///         ctx.handle_item(item);
+///     }
+/// }
+///
+/// // When `guard` is dropped, `items` is replaced in `ctx`, allowing us to
+/// // reuse its capacity and avoid future allocations.  ```
+/// assert!(ctx.items.capacity() >= 3);
+/// ```
+pub struct TakeAndReplace<'a, T, U, F>
+where
+    F: Fn(&mut T) -> &mut U,
+    U: Default,
+{
+    container: &'a mut T,
+    value: U,
+    proj: F,
+}
+
+impl<'a, T, U, F> Drop for TakeAndReplace<'a, T, U, F>
+where
+    F: Fn(&mut T) -> &mut U,
+    U: Default,
+{
+    fn drop(&mut self) {
+        *(self.proj)(self.container) = std::mem::take(&mut self.value);
+    }
+}
+
+impl<'a, T, U, F> TakeAndReplace<'a, T, U, F>
+where
+    F: Fn(&mut T) -> &mut U,
+    U: Default,
+{
+    /// Create a new `TakeAndReplace` that temporarily takes out
+    /// `proj(container)`.
+    pub fn new(mut container: &'a mut T, proj: F) -> Self {
+        let value = std::mem::take(proj(&mut container));
+        TakeAndReplace {
+            container,
+            value,
+            proj,
+        }
+    }
+
+    /// Get the underlying container and taken-out value.
+    pub fn get(&mut self) -> (&mut T, &mut U) {
+        (&mut *self.container, &mut self.value)
+    }
+}