regalloc2 0.15.1

//! Fuzz the parallel-move resolver.

use crate::moves::{MoveAndScratchResolver, ParallelMoves};
use crate::{Allocation, PReg, RegClass, SpillSlot};
use arbitrary::{Arbitrary, Result, Unstructured};
use std::collections::{HashMap, HashSet};
use std::{vec, vec::Vec};

fn is_stack_alloc(alloc: Allocation) -> bool {
    // Treat registers 20..=29 as fixed stack slots.
    if let Some(reg) = alloc.as_reg() {
        reg.index() > 20
    } else {
        alloc.is_stack()
    }
}

///
#[derive(Clone, Debug)]
pub struct TestCase {
    moves: Vec<(Allocation, Allocation)>,
    available_pregs: Vec<Allocation>,
}

impl Arbitrary<'_> for TestCase {
    fn arbitrary(u: &mut Unstructured) -> Result<Self> {
        let mut ret = TestCase {
            moves: vec![],
            available_pregs: vec![],
        };
        let mut written = HashSet::new();
        // An arbitrary sequence of moves between registers 0 to 29
        // inclusive.
        while bool::arbitrary(u)? {
            let src = if bool::arbitrary(u)? {
                let reg = u.int_in_range(0..=29)?;
                Allocation::reg(PReg::new(reg, RegClass::Int))
            } else {
                let slot = u.int_in_range(0..=31)?;
                Allocation::stack(SpillSlot::new(slot))
            };
            let dst = if bool::arbitrary(u)? {
                let reg = u.int_in_range(0..=29)?;
                Allocation::reg(PReg::new(reg, RegClass::Int))
            } else {
                let slot = u.int_in_range(0..=31)?;
                Allocation::stack(SpillSlot::new(slot))
            };

            // Stop if we are going to write a reg more than once:
            // that creates an invalid parallel move set.
            if written.contains(&dst) {
                break;
            }
            written.insert(dst);

            ret.moves.push((src, dst));
        }

        // We might have some unallocated registers free for scratch
        // space...
        for i in 0..u.int_in_range(0..=2)? {
            let reg = PReg::new(30 + i, RegClass::Int);
            ret.available_pregs.push(Allocation::reg(reg));
        }
        Ok(ret)
    }
}

pub fn check(t: TestCase) {
    let mut par = ParallelMoves::new();
    for &(src, dst) in &t.moves {
        par.add(src, dst, ());
    }

    let moves = par.resolve();
    log::trace!("raw resolved moves: {:?}", moves);

    // Resolve uses of scratch reg and stack-to-stack moves with the scratch
    // resolver.
    let mut avail = t.available_pregs.clone();
    let find_free_reg = || avail.pop();
    let mut next_slot = 32;
    let get_stackslot = || {
        let slot = next_slot;
        next_slot += 1;
        Allocation::stack(SpillSlot::new(slot))
    };
    let preferred_victim = PReg::new(0, RegClass::Int);
    let scratch_resolver = MoveAndScratchResolver {
        find_free_reg,
        get_stackslot,
        is_stack_alloc,
        borrowed_scratch_reg: preferred_victim,
    };
    let moves = scratch_resolver.compute(moves);
    log::trace!("resolved moves: {:?}", moves);

    // Compute the final source reg for each dest reg in the original
    // parallel-move set.
    let mut final_src_per_dest: HashMap<Allocation, Allocation> = HashMap::new();
    for &(src, dst) in &t.moves {
        final_src_per_dest.insert(dst, src);
    }
    log::trace!("expected final state: {:?}", final_src_per_dest);

    // Simulate the sequence of moves.
    let mut locations: HashMap<Allocation, Allocation> = HashMap::new();
    for (src, dst, _) in moves {
        let data = locations.get(&src).cloned().unwrap_or(src);
        locations.insert(dst, data);
    }
    log::trace!("simulated final state: {:?}", locations);

    // Assert that the expected register-moves occurred.
    for (reg, data) in locations {
        if let Some(&expected_data) = final_src_per_dest.get(&reg) {
            assert_eq!(expected_data, data);
        } else {
            if data != reg {
                // If not just the original value, then this location has been
                // modified, but it was not part of the original parallel move.
                // It must have been an available preg or a scratch stackslot.
                assert!(
                    t.available_pregs.contains(&reg)
                        || (reg.is_stack() && reg.as_stack().unwrap().index() >= 32)
                );
            }
        }
    }
}

#[test]
fn smoke() {
    arbtest::arbtest(|u| {
        let test_case = TestCase::arbitrary(u)?;
        check(test_case);
        Ok(())
    })
    .budget_ms(1_000);
}