use crate::isa::reg::{Reg, RegClass};

/// A bit set to track register availability.

pub(crate) struct RegSet {

    /// Bitset to track general purpose register availability.

    gpr: RegBitSet,

    /// Bitset to track floating-point register availability.

    fpr: RegBitSet,

}

use std::ops::{Index, IndexMut};

impl Index<RegClass> for RegSet {

    type Output = RegBitSet;

    fn index(&self, class: RegClass) -> &Self::Output {

        match class {

            RegClass::Int => &self.gpr,

            RegClass::Float => &self.fpr,

            c => unreachable!("Unexpected register class {:?}", c),

        }

    }

}

impl IndexMut<RegClass> for RegSet {

    fn index_mut(&mut self, class: RegClass) -> &mut Self::Output {

        match class {

            RegClass::Int => &mut self.gpr,

            RegClass::Float => &mut self.fpr,

            c => unreachable!("Unexpected register class {:?}", c),

        }

    }

}

/// Bitset for a particular register class.

pub struct RegBitSet {

    /// The register class.

    class: RegClass,

    /// The set of allocatable

    allocatable: u64,

    /// The set of non-alloctable registers.

    non_allocatable: u64,

    /// The max number of registers.

    /// Invariant:

    /// When allocating or freeing a register the encoding (index) of the

    /// register must be less than the max property.

    max: usize,

}

impl RegBitSet {

    /// Creates an integer register class bitset.

    pub fn int(allocatable: u64, non_allocatable: u64, max: usize) -> Self {

        // Assert that one set is the complement of the other.

        debug_assert!(allocatable & non_allocatable == 0);

        Self {

            class: RegClass::Int,

            allocatable,

            non_allocatable,

            max,

        }

    }

    /// Creates a float register class bitset.

    pub fn float(allocatable: u64, non_allocatable: u64, max: usize) -> Self {

        // Assert that one set is the complement of the other.

        debug_assert!(allocatable & non_allocatable == 0);

        Self {

            class: RegClass::Float,

            allocatable,

            non_allocatable,

            max,

        }

    }

}

impl RegSet {

    /// Create a new register set.

    pub fn new(gpr: RegBitSet, fpr: RegBitSet) -> Self {

        debug_assert!(gpr.class == RegClass::Int);

        debug_assert!(fpr.class == RegClass::Float);

        Self { gpr, fpr }

    }

    /// Allocate the next available register of the given class,

    /// returning `None` if there are no more registers available.

    pub fn reg_for_class(&mut self, class: RegClass) -> Option<Reg> {

        self.available(class).then(|| {

            let bitset = &self[class];

            let index = bitset.allocatable.trailing_zeros();

            self.allocate(class, index.into());

            Reg::from(class, index as usize)

        })

    }

    /// Request a specific register.

    pub fn reg(&mut self, reg: Reg) -> Option<Reg> {

        let index = reg.hw_enc();

        self.named_reg_available(reg).then(|| {

            self.allocate(reg.class(), index.try_into().unwrap());

            reg

        })

    }

    /// Marks the specified register as available, utilizing the

    /// register class to determine the bitset that requires updating.

    pub fn free(&mut self, reg: Reg) {

        let bitset = &self[reg.class()];

        let index = reg.hw_enc();

        assert!(index < bitset.max);

        let index = u64::try_from(index).unwrap();

        if !self.is_non_allocatable(reg.class(), index) {

            self[reg.class()].allocatable |= 1 << index;

        }

    }

    /// Returns true if the specified register is allocatable.

    pub fn named_reg_available(&self, reg: Reg) -> bool {

        let bitset = &self[reg.class()];

        assert!(reg.hw_enc() < bitset.max);

        let index = 1 << reg.hw_enc();

        (!bitset.allocatable & index) == 0

            || self.is_non_allocatable(reg.class(), reg.hw_enc().try_into().unwrap())

    }

    fn available(&self, class: RegClass) -> bool {

        let bitset = &self[class];

        bitset.allocatable != 0

    }

    fn allocate(&mut self, class: RegClass, index: u64) {

        if !self.is_non_allocatable(class, index) {

            self[class].allocatable &= !(1 << index);

        }

    }

    fn is_non_allocatable(&self, class: RegClass, index: u64) -> bool {

        let bitset = &self[class];

        let non_allocatable = bitset.non_allocatable;

        non_allocatable != 0 && !non_allocatable & (1 << index) == 0

    }

}

#[cfg(test)]

mod tests {

    use super::{Reg, RegBitSet, RegClass, RegSet};

    const UNIVERSE: u64 = (1 << 16) - 1;

    const MAX: usize = 16;

    #[test]

    fn test_any_gpr() {

        let bitset = RegBitSet::int(UNIVERSE, !UNIVERSE, MAX);

        let zero = RegBitSet::float(0, 0, MAX);

        let mut set = RegSet::new(bitset, zero);

        for _ in 0..16 {

            let gpr = set.reg_for_class(RegClass::Int);

            assert!(gpr.is_some())

        }

        assert!(!set.available(RegClass::Int));

        assert!(set.reg_for_class(RegClass::Int).is_none())

    }

    #[test]

    fn test_gpr() {

        let non_allocatable: u64 = 1 << 5;

        let all = UNIVERSE & !non_allocatable;

        let non_alloc = Reg::int(5);

        let alloc = Reg::int(2);

        let bitset = RegBitSet::int(all, non_allocatable, MAX);

        let zero = RegBitSet::float(0, 0, MAX);

        let mut set = RegSet::new(bitset, zero);

        // Requesting a non alloctable register returns the register

        // and doesn't allocate it.

        assert!(set.reg(non_alloc).is_some());

        assert!(set.reg(non_alloc).is_some());

        // Requesting an allocatable register twice returns none the

        // second time.

        assert!(set.reg(alloc).is_some());

        assert!(set.reg(alloc).is_none());

    }

    #[test]

    fn test_free_reg() {

        let set = RegBitSet::int(UNIVERSE, !UNIVERSE, MAX);

        let zero = RegBitSet::float(0, 0, MAX);

        let mut set = RegSet::new(set, zero);

        let gpr = set.reg_for_class(RegClass::Int).unwrap();

        set.free(gpr);

        assert!(set.reg(gpr).is_some());

    }

}