//! Bitset utilities for register tracking.
//!
//! Provides efficient bit-level set operations for tracking live registers
//! during liveness analysis and register allocation.
@test mod tests;

use std::lang::alloc;

/// Return the minimum of two 32-bit values.
fn min(a: u32, b: u32) -> u32 {
    if a < b {
        return a;
    }
    return b;
}

/// Calculate the number of 32-bit words needed to store `n` bits.
export fn wordsFor(n: u32) -> u32 {
    return (n + 31) / 32;
}

/// A fixed-size bitset backed by an array of 32-bit words.
export record Bitset: Copy {
    /// Backing words. The storage must outlive the bitset and its iterators.
    bits: *unsafe mut [u32],
    /// Number of bits this bitset can hold.
    len: u32,
}

/// Create a new bitset backed by the given zero-initialized storage.
/// The storage must outlive the bitset and its iterators.
export unsafe fn new(bits: &mut [u32]) -> Bitset {
    let len = bits.len * 32;
    return Bitset { bits: bits as *unsafe mut [u32], len };
}

/// Create a new bitset backed by the given storage, zeroing it first.
/// The storage must outlive the bitset and its iterators.
export unsafe fn init(bits: &mut [u32]) -> Bitset {
    for i in 0..bits.len {
        set bits[i] = 0;
    }
    return new(bits);
}

/// Create a bitset from arena allocation.
export unsafe fn allocate(arena: &mut alloc::Arena, len: u32) -> Bitset throws (alloc::AllocError) {
    let numWords = wordsFor(len);
    let bits = try alloc::allocRawSlice(arena, @sizeOf(u32), @alignOf(u32), numWords) as *unsafe mut [u32];

    return init(bits);
}

/// Set bit `n` in the bitset.
export unsafe fn put(bs: &mut Bitset, n: u32) {
    if n >= bs.len {
        return;
    }
    let word = n / 32;
    let b = n % 32;

    set bs.bits[word] |= (1 << b);
}

/// Clear bit `n` in the bitset.
export unsafe fn clear(bs: &mut Bitset, n: u32) {
    if n >= bs.len {
        return;
    }
    let word = n / 32;
    let b = n % 32;

    set bs.bits[word] &= ~(1 << b);
}

/// Check if bit `n` is set.
export unsafe fn contains(bs: &Bitset, n: u32) -> bool {
    if n >= bs.len {
        return false;
    }
    let word = n / 32;
    let b = n % 32;

    return (bs.bits[word] & (1 << b)) <> 0;
}

/// Count the number of set bits.
export unsafe fn count(bs: &Bitset) -> u32 {
    let mut total: u32 = 0;
    let numWords = bs.bits.len;
    for i in 0..numWords {
        let word = bs.bits[i];
        if word <> 0 {
            set total += popCount(word);
        }
    }
    return total;
}

/// Population count for a 32-bit word.
fn popCount(x: u32) -> u32 {
    let mut n = x;
    set n -= ((n >> 1) & 0x55555555);
    set n = (n & 0x33333333) + ((n >> 2) & 0x33333333);
    set n = (n + (n >> 4)) & 0x0F0F0F0F;
    set n += (n >> 8);
    set n += (n >> 16);

    return n & 0x3F;
}

/// Union: `dst = dst | src`.
export unsafe fn union_(dst: &mut Bitset, src: &Bitset) {
    let numWords = dst.bits.len;
    let srcWords = src.bits.len;
    let minWords = min(numWords, srcWords);
    for i in 0..minWords {
        set dst.bits[i] |= src.bits[i];
    }
}

/// Subtract: `dst = dst - src`.
export unsafe fn subtract(dst: &mut Bitset, src: &Bitset) {
    let numWords = dst.bits.len;
    let srcWords = src.bits.len;
    let minWords = min(numWords, srcWords);
    for i in 0..minWords {
        set dst.bits[i] &= ~src.bits[i];
    }
}

/// Check if two bitsets are equal.
export unsafe fn eq(a: &Bitset, b: &Bitset) -> bool {
    let numWordsA = a.bits.len;
    let numWordsB = b.bits.len;
    let minWords = min(numWordsA, numWordsB);

    for i in 0..minWords {
        if a.bits[i] <> b.bits[i] {
            return false;
        }
    }
    for i in minWords..numWordsA {
        if a.bits[i] <> 0 {
            return false;
        }
    }
    for i in minWords..numWordsB {
        if b.bits[i] <> 0 {
            return false;
        }
    }
    return true;
}

/// Copy bits from source to destination.
export unsafe fn copy(dst: &mut Bitset, src: &Bitset) {
    let numWords = dst.bits.len;
    let srcWords = src.bits.len;
    let minWords = min(numWords, srcWords);

    for i in 0..minWords {
        set dst.bits[i] = src.bits[i];
    }
    // Clear remaining words if destination is larger.
    for i in minWords..numWords {
        set dst.bits[i] = 0;
    }
}

/// Clear all bits.
export unsafe fn clearAll(bs: &mut Bitset) {
    let numWords = bs.bits.len;
    for i in 0..numWords {
        set bs.bits[i] = 0;
    }
}

/// Iterator state for iterating set bits.
export record BitIter: Copy {
    /// Bitset being iterated. It must outlive this iterator.
    bs: *unsafe Bitset,
    /// Current word index.
    wordIdx: u32,
    /// Remaining bits in the current word (visited bits cleared).
    remaining: u32,
}

/// Create an iterator over set bits.
/// The bitset and its backing words must outlive the iterator.
export unsafe fn iter(bs: &Bitset) -> BitIter {
    let remaining = bs.bits[0] if bs.len > 0 else 0;
    return BitIter { bs: bs as *unsafe Bitset, wordIdx: 0, remaining };
}

/// Get the next set bit, or nil if none remain.
export unsafe fn iterNext(it: &mut BitIter) -> ?u32 {
    let numWords = it.bs.bits.len;
    // Skip to next non-zero word.
    while it.remaining == 0 {
        set it.wordIdx += 1;
        if it.wordIdx >= numWords {
            return nil;
        }
        set it.remaining = it.bs.bits[it.wordIdx];
    }
    // Find lowest set bit position using de Bruijn sequence.
    let b = ctz(it.remaining);
    let n = it.wordIdx * 32 + b;
    if n >= it.bs.len {
        return nil;
    }
    // Clear the lowest set bit.
    set it.remaining &= it.remaining - 1;

    return n;
}

/// Count trailing zeros in a 32-bit value.
/// Returns the bit position of the lowest set bit (0-31).
/// Behavior is undefined if `x` is 0.
fn ctz(x: u32) -> u32 {
    // De Bruijn sequence for 32-bit CTZ.
    constant DEBRUIJN: u32 = 0x077CB531;
    constant TABLE: [u8; 32] = [
         0,  1, 28,  2, 29, 14, 24,  3, 30, 22, 20, 15, 25, 17,  4,  8,
        31, 27, 13, 23, 21, 19, 16,  7, 26, 12, 18,  6, 11,  5, 10,  9,
    ];
    // Isolate lowest set bit, multiply by de Bruijn constant, look up.
    let isolated = x & (~x + 1);

    return TABLE[(isolated * DEBRUIJN) >> 27] as u32;
}