wasmtime/lib/cretonne/src/regalloc/virtregs.rs

//! Virtual registers.
//!
//! A virtual register is a set of related SSA values whose live ranges don't interfere. If all the
//! values in a virtual register are assigned to the same location, fewer copies will result in the
//! output.
//!
//! A virtual register is typically built by merging together SSA values that are "phi-related" -
//! that is, one value is passed as an EBB argument to a branch and the other is the EBB parameter
//! value itself.
//!
//! If any values in a virtual register are spilled, they will use the same stack slot. This avoids
//! memory-to-memory copies when a spilled value is passed as an EBB argument.

use entity_list::{EntityList, ListPool};
use entity_map::{EntityMap, PrimaryEntityData};
use ir::Value;
use packed_option::PackedOption;
use ref_slice::ref_slice;

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
pub struct VirtReg(u32);
entity_impl!(VirtReg, "vreg");

type ValueList = EntityList<Value>;
impl PrimaryEntityData for ValueList {}

/// Collection of virtual registers.
///
/// Each virtual register is a list of values. Also maintain a map from values to their unique
/// virtual register, if any.
pub struct VirtRegs {
    /// Memory pool for the value lists.
    pool: ListPool<Value>,

    /// The primary table of virtual registers.
    ///
    /// The list of values ion a virtual register is kept sorted according to the dominator tree's
    /// RPO of the value defs.
    vregs: EntityMap<VirtReg, ValueList>,

    /// Each value belongs to at most one virtual register.
    value_vregs: EntityMap<Value, PackedOption<VirtReg>>,
}

#[allow(dead_code)]
impl VirtRegs {
    /// Create a new virtual register collection.
    pub fn new() -> VirtRegs {
        VirtRegs {
            pool: ListPool::new(),
            vregs: EntityMap::new(),
            value_vregs: EntityMap::new(),
        }
    }

    /// Clear all virtual registers.
    pub fn clear(&mut self) {
        self.vregs.clear();
        self.value_vregs.clear();
        self.pool.clear();
    }

    /// Get the virtual register containing `value`, if any.
    pub fn get(&self, value: Value) -> Option<VirtReg> {
        self.value_vregs.get_or_default(value).into()
    }

    /// Get the list of values in `vreg`. The values are ordered according to `DomTree::rpo_cmp` of
    /// their definition points.
    pub fn values(&self, vreg: VirtReg) -> &[Value] {
        self.vregs[vreg].as_slice(&self.pool)
    }

    /// Get the congruence class of `value`.
    ///
    /// If `value` belongs to a virtual register, the congruence class is the values of the virtual
    /// register. Otherwise it is just the value itself.
    pub fn congruence_class<'a, 'b>(&'a self, value: &'b Value) -> &'b [Value]
        where 'a: 'b
    {
        self.get(*value)
            .map(|vr| self.values(vr))
            .unwrap_or(ref_slice(value))
    }

    /// Check if `a` and `b` belong to the same congruence class.
    pub fn same_class(&self, a: Value, b: Value) -> bool {
        match (self.get(a), self.get(b)) {
            (Some(va), Some(vb)) => va == vb,
            _ => a == b,
        }
    }

    /// Unify `values` into a single virtual register.
    ///
    /// The values in the slice can be singletons or they can belong to a virtual register already.
    /// If a value belongs to a virtual register, all of the values in that register must be
    /// present.
    ///
    /// The values are assumed to already be in RPO order.
    pub fn unify(&mut self, values: &[Value]) -> VirtReg {
        // Start by clearing all virtual registers involved.
        // Pick a virtual register to reuse (the smallest number) or allocate a new one.
        let mut singletons = 0;
        let mut cleared = 0;
        let vreg = values
            .iter()
            .filter_map(|&v| {
                let vr = self.get(v);
                match vr {
                    None => singletons += 1,
                    Some(vr) => {
                        if !self.vregs[vr].is_empty() {
                            cleared += self.vregs[vr].len(&self.pool);
                            self.vregs[vr].clear(&mut self.pool);
                        }
                    }
                }
                vr
            })
            .min()
            .unwrap_or_else(|| self.vregs.push(Default::default()));

        assert_eq!(values.len(),
                   singletons + cleared,
                   "Can't unify partial virtual registers");

        self.vregs[vreg].extend(values.iter().cloned(), &mut self.pool);
        for &v in values {
            *self.value_vregs.ensure(v) = vreg.into();
        }

        vreg
    }
}