Add EntryRegDiversions to record diversions for each block entry.

cranelift/codegen/src/binemit/mod.rs

@@ -176,7 +176,7 @@ where
 {
     let mut divert = RegDiversions::new();
     for ebb in func.layout.ebbs() {
-        divert.clear();
+        divert.at_ebb(&func.entry_diversions, ebb);
         debug_assert_eq!(func.offsets[ebb], sink.offset());
         for inst in func.layout.ebb_insts(ebb) {
             emit_inst(func, inst, &mut divert, sink, isa);

cranelift/codegen/src/binemit/relaxation.rs

@@ -74,7 +74,7 @@ pub fn relax_branches(
     {
         let mut cur = FuncCursor::new(func);
         while let Some(ebb) = cur.next_ebb() {
-            divert.clear();
+            divert.at_ebb(&cur.func.entry_diversions, ebb);
             cur.func.offsets[ebb] = offset;
             while let Some(inst) = cur.next_inst() {
                 divert.apply(&cur.func.dfg[inst]);
@@ -93,7 +93,8 @@ pub fn relax_branches(
         // Visit all instructions in layout order.
         let mut cur = FuncCursor::new(func);
         while let Some(ebb) = cur.next_ebb() {
-            divert.clear();
+            divert.at_ebb(&cur.func.entry_diversions, ebb);
+
             // Record the offset for `ebb` and make sure we iterate until offsets are stable.
             if cur.func.offsets[ebb] != offset {
                 cur.func.offsets[ebb] = offset;

cranelift/codegen/src/binemit/shrink.rs

@@ -20,7 +20,9 @@ pub fn shrink_instructions(func: &mut Function, isa: &dyn TargetIsa) {
     let mut divert = RegDiversions::new();
 
     for ebb in func.layout.ebbs() {
-        divert.clear();
+        // Load diversions from predecessors.
+        divert.at_ebb(&func.entry_diversions, ebb);
+
         for inst in func.layout.ebb_insts(ebb) {
             let enc = func.encodings[inst];
             if enc.is_legal() {

cranelift/codegen/src/ir/function.rs

@@ -14,7 +14,7 @@ use crate::ir::{
 use crate::ir::{EbbOffsets, InstEncodings, SourceLocs, StackSlots, ValueLocations};
 use crate::ir::{JumpTableOffsets, JumpTables};
 use crate::isa::{CallConv, EncInfo, Encoding, Legalize, TargetIsa};
-use crate::regalloc::RegDiversions;
+use crate::regalloc::{EntryRegDiversions, RegDiversions};
 use crate::value_label::ValueLabelsRanges;
 use crate::write::write_function;
 use core::fmt;
@@ -62,6 +62,12 @@ pub struct Function {
     /// Location assigned to every value.
     pub locations: ValueLocations,
 
+    /// Non-default locations assigned to value at the entry of basic blocks.
+    ///
+    /// At the entry of each basic block, we might have values which are not in their default
+    /// ValueLocation. This field records these register-to-register moves as Diversions.
+    pub entry_diversions: EntryRegDiversions,
+
     /// Code offsets of the EBB headers.
     ///
     /// This information is only transiently available after the `binemit::relax_branches` function
@@ -94,6 +100,7 @@ impl Function {
             layout: Layout::new(),
             encodings: SecondaryMap::new(),
             locations: SecondaryMap::new(),
+            entry_diversions: EntryRegDiversions::new(),
             offsets: SecondaryMap::new(),
             jt_offsets: SecondaryMap::new(),
             srclocs: SecondaryMap::new(),
@@ -112,6 +119,7 @@ impl Function {
         self.layout.clear();
         self.encodings.clear();
         self.locations.clear();
+        self.entry_diversions.clear();
         self.offsets.clear();
         self.jt_offsets.clear();
         self.srclocs.clear();
@@ -198,10 +206,12 @@ impl Function {
             !self.offsets.is_empty(),
             "Code layout must be computed first"
         );
+        let mut divert = RegDiversions::new();
+        divert.at_ebb(&self.entry_diversions, ebb);
         InstOffsetIter {
             encinfo: encinfo.clone(),
             func: self,
-            divert: RegDiversions::new(),
+            divert,
             encodings: &self.encodings,
             offset: self.offsets[ebb],
             iter: self.layout.ebb_insts(ebb),

cranelift/codegen/src/regalloc/coloring.rs

@@ -199,6 +199,7 @@ impl<'a> Context<'a> {
             self.domtree,
         );
 
+        self.divert.at_ebb(&self.cur.func.entry_diversions, ebb);
         if self.cur.func.layout.entry_block() == Some(ebb) {
             // Parameters on the entry block have ABI constraints.
             self.color_entry_params(tracker.live())

cranelift/codegen/src/regalloc/diversion.rs

@@ -9,10 +9,11 @@
 
 use crate::fx::FxHashMap;
 use crate::hash_map::{Entry, Iter};
+use crate::ir::{Ebb, StackSlot, Value, ValueLoc, ValueLocations};
 use crate::ir::{InstructionData, Opcode};
-use crate::ir::{StackSlot, Value, ValueLoc, ValueLocations};
 use crate::isa::{RegInfo, RegUnit};
 use core::fmt;
+use cranelift_entity::{SparseMap, SparseMapValue};
 
 /// A diversion of a value from its original location to a new register or stack location.
 ///
@@ -38,10 +39,23 @@ impl Diversion {
 }
 
 /// Keep track of diversions in an EBB.
+#[derive(Clone)]
 pub struct RegDiversions {
     current: FxHashMap<Value, Diversion>,
 }
 
+/// Keep track of diversions at the entry of EBB.
+#[derive(Clone)]
+struct EntryRegDiversionsValue {
+    key: Ebb,
+    divert: RegDiversions,
+}
+
+/// Map EBB to their matching RegDiversions at basic blocks entry.
+pub struct EntryRegDiversions {
+    map: SparseMap<Ebb, EntryRegDiversionsValue>,
+}
+
 impl RegDiversions {
     /// Create a new empty diversion tracker.
     pub fn new() -> Self {
@@ -50,7 +64,7 @@ impl RegDiversions {
         }
     }
 
-    /// Clear the tracker, preparing for a new EBB.
+    /// Clear the content of the diversions, to reset the state of the compiler.
     pub fn clear(&mut self) {
         self.current.clear()
     }
@@ -92,7 +106,7 @@ impl RegDiversions {
     /// Record any kind of move.
     ///
     /// The `from` location must match an existing `to` location, if any.
-    pub fn divert(&mut self, value: Value, from: ValueLoc, to: ValueLoc) {
+    fn divert(&mut self, value: Value, from: ValueLoc, to: ValueLoc) {
         debug_assert!(from.is_assigned() && to.is_assigned());
         match self.current.entry(value) {
             Entry::Occupied(mut e) => {
@@ -163,9 +177,92 @@ impl RegDiversions {
         self.current.remove(&value).map(|d| d.to)
     }
 
+    /// Resets the state of the current diversions to the recorded diversions at the entry of the
+    /// given `ebb`. The recoded diversions is available after coloring on `func.entry_diversions`
+    /// field.
+    pub fn at_ebb(&mut self, entry_diversions: &EntryRegDiversions, ebb: Ebb) {
+        self.clear();
+        if let Some(entry_divert) = entry_diversions.map.get(ebb) {
+            let iter = entry_divert.divert.current.iter();
+            self.current.extend(iter);
+        }
+    }
+
+    /// Copy the current state of the diversions, and save it for the entry of the `ebb` given as
+    /// argument.
+    ///
+    /// Note: This function can only be called once on an `ebb` with a given `entry_diversions`
+    /// argument, otherwise it would panic.
+    pub fn save_for_ebb(&mut self, entry_diversions: &mut EntryRegDiversions, target: Ebb) {
+        // No need to save anything if there is no diversions to be recorded.
+        if self.is_empty() {
+            return;
+        }
+        debug_assert!(!entry_diversions.map.contains_key(target));
+        let iter = self.current.iter();
+        let mut entry_divert = RegDiversions::new();
+        entry_divert.current.extend(iter);
+        entry_diversions.map.insert(EntryRegDiversionsValue {
+            key: target,
+            divert: entry_divert,
+        });
+    }
+
+    /// Check that the recorded entry for a given `ebb` matches what is recorded in the
+    /// `entry_diversions`.
+    pub fn check_ebb_entry(&self, entry_diversions: &EntryRegDiversions, target: Ebb) -> bool {
+        let entry_divert = match entry_diversions.map.get(target) {
+            Some(entry_divert) => entry_divert,
+            None => return self.is_empty(),
+        };
+
+        if entry_divert.divert.current.len() != self.current.len() {
+            return false;
+        }
+
+        for (val, _) in entry_divert.divert.current.iter() {
+            if !self.current.contains_key(val) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     /// Return an object that can display the diversions.
     pub fn display<'a, R: Into<Option<&'a RegInfo>>>(&'a self, regs: R) -> DisplayDiversions<'a> {
-        DisplayDiversions(self, regs.into())
+        DisplayDiversions(&self, regs.into())
+    }
+}
+
+impl EntryRegDiversions {
+    /// Create a new empty entry diversion, to associate diversions to each EBB entry.
+    pub fn new() -> Self {
+        EntryRegDiversions {
+            map: SparseMap::new(),
+        }
+    }
+
+    pub fn clear(&mut self) {
+        self.map.clear();
+    }
+}
+
+impl Clone for EntryRegDiversions {
+    /// The Clone trait is required by `ir::Function`.
+    fn clone(&self) -> Self {
+        let mut tmp = Self::new();
+        for v in self.map.values() {
+            tmp.map.insert(v.clone());
+        }
+        tmp
+    }
+}
+
+/// Implement `SparseMapValue`, as required to make use of a `SparseMap` for mapping the entry
+/// diversions for each EBB.
+impl SparseMapValue<Ebb> for EntryRegDiversionsValue {
+    fn key(&self) -> Ebb {
+        self.key
     }
 }
 
@@ -175,7 +272,7 @@ pub struct DisplayDiversions<'a>(&'a RegDiversions, Option<&'a RegInfo>);
 impl<'a> fmt::Display for DisplayDiversions<'a> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "{{")?;
-        for (value, div) in self.0.iter() {
+        for (value, div) in self.0.current.iter() {
             write!(
                 f,
                 " {}: {} -> {}",

cranelift/codegen/src/regalloc/mod.rs

@@ -20,6 +20,6 @@ mod solver;
 mod spilling;
 
 pub use self::context::Context;
-pub use self::diversion::RegDiversions;
+pub use self::diversion::{EntryRegDiversions, RegDiversions};
 pub use self::register_set::RegisterSet;
 pub use self::safepoint::emit_stackmaps;

cranelift/codegen/src/value_label.rs

@@ -118,7 +118,7 @@ where
     let mut tracked_values: Vec<(Value, ValueLabel, u32, ValueLoc)> = Vec::new();
     let mut divert = RegDiversions::new();
     for ebb in ebbs {
-        divert.clear();
+        divert.at_ebb(&func.entry_diversions, ebb);
         let mut last_srcloc: Option<T> = None;
         for (offset, inst, size) in func.inst_offsets(ebb, &encinfo) {
             divert.apply(&func.dfg[inst]);

cranelift/codegen/src/verifier/locations.rs

@@ -51,9 +51,8 @@ impl<'a> LocationVerifier<'a> {
         let mut divert = RegDiversions::new();
 
         for ebb in self.func.layout.ebbs() {
-            // Diversions are reset at the top of each EBB. No diversions can exist across control
+            divert.at_ebb(&self.func.entry_diversions, ebb);
-            // flow edges.
+
-            divert.clear();
             for inst in self.func.layout.ebb_insts(ebb) {
                 let enc = self.func.encodings[inst];