//! 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 storage for bits, organized as 32-bit words.

    bits: *mut [u32],

    /// Number of bits this bitset can hold.

    len: u32,

}

/// Create a new bitset backed by the given zero-initialized storage.

export fn new(bits: *mut [u32]) -> Bitset {

    return Bitset { bits, len: bits.len * 32 };

}

/// Create a new bitset backed by the given storage, zeroing it first.

export 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 fn allocate(arena: *mut alloc::Arena, len: u32) -> Bitset throws (alloc::AllocError) {

    let numWords = wordsFor(len);

    let bits = try alloc::allocSlice(arena, @sizeOf(u32), @alignOf(u32), numWords) as *mut [u32];

    return init(bits);

}

/// Set bit `n` in the bitset.

export 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 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 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 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 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 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 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 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 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.

    bs: *Bitset,

    /// Current word index.

    wordIdx: u32,

    /// Remaining bits in the current word (visited bits cleared).

    remaining: u32,

}

/// Create an iterator over set bits.

export fn iter(bs: *Bitset) -> BitIter {

    let remaining = bs.bits[0] if bs.len > 0 else 0;

    return BitIter { bs, wordIdx: 0, remaining };

}

/// Get the next set bit, or nil if none remain.

export 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;

}