branchless_shr.rs

// Algorithm that performs a branchless 4-word right shifting, which shifts
// 256bit numbers in 64bit systems or 128bit in 32bit systems.
// 
// @author Lohann Paterno Coutinho Ferreira <developer@lohann.dev>

#[cfg(target_pointer_width = "64")]
pub type Limb = u64;

#[cfg(target_pointer_width = "32")]
pub type Limb = u32;

#[cfg(target_pointer_width = "16")]
pub type Limb = u16;

/// Total number of bits in a 4-word unsigned integer.
const BITS: u32 = Limb::BITS * 4;

/// Mask to extract the `shift`.
const MASK: Limb = Limb::MAX >> (Limb::BITS - BITS.ilog2());

/// number of bits used to represent the maximum `word shift`
const WORD_SHIFT: u32 = Limb::BITS.ilog2();

/// bitmask used to extract the `bit shift`
const BIT_MASK: Limb = MASK >> (BITS.ilog2() - WORD_SHIFT);

#[unsafe(no_mangle)]
pub const fn branchless_shr4(n: &mut [Limb; 4], shift: &[Limb; 4]) {
    // Returns `Limbs::MAX` when `b` is true, otherwise returns zero.
    #[inline(always)]
    const fn to_mask(b: bool) -> Limb {
        (0 as Limb).wrapping_sub(b as Limb)
    }
  
    // load shift
    let mut bit_shift = shift[0];

    // handles the case where `shift >= BITS`
    let mut mask = (bit_shift >> BITS.ilog2()) | shift[1] | shift[2] | shift[3];
    mask = to_mask(mask == 0);

    // load limbs when `shift < BITS`, otherwise zero.
    let mut a = n[0] & mask;
    let mut b = n[1] & mask;
    let mut c = n[2] & mask;
    let mut d = n[3] & mask;

    // use low bits for bit shifting and high bits for word shifting.
    bit_shift &= MASK;
    let word_shift = bit_shift >> WORD_SHIFT;
    let mut bit_shift = (bit_shift & BIT_MASK) as u32;

    // perform bitwise shifting
    {
        // extract low bits
        let a_lo = a.wrapping_shr(bit_shift);
        let b_lo = b.wrapping_shr(bit_shift);
        let c_lo = c.wrapping_shr(bit_shift);
        let d_lo = d.wrapping_shr(bit_shift);

        // ignore high bits when `bit_shift == 0`
        mask = to_mask(bit_shift > 0);
        bit_shift = 64 - bit_shift;

        // extract high bits when `bit_shift > 0`.
        // obs: `a_hi` is discarded when performing right shift.
        // -- let a_hi = a.wrapping_shl(bit_shift) & mask; --
        let b_hi = b.wrapping_shl(bit_shift) & mask;
        let c_hi = c.wrapping_shl(bit_shift) & mask;
        let d_hi = d.wrapping_shl(bit_shift) & mask;

        // merge low and high bits
        d = d_lo; 
        c = c_lo | d_hi;
        b = b_lo | c_hi;
        a = a_lo | b_hi;
    }

    // prepare shift booleans
    // - only one of `shift0`, `shift1`, `shift2` and `shift3` can be true at time.
    let word_shift1 = to_mask((word_shift & 1) > 0);
    let word_shift2 = to_mask((word_shift & 2) > 0);
    let non_zero = word_shift1 | word_shift2;
    let shift0 = !non_zero;
    let shift1 = non_zero ^ word_shift2;
    let shift2 = non_zero ^ word_shift1;
    let shift3 = word_shift1 & word_shift2;

    // perform branchless bitwise word shifting.
    a = (a & shift0) | (b & shift1) | (c & shift2) | (d & shift3);
    b = (b & shift0) | (c & shift1) | (d & shift2);
    c = (c & shift0) | (d & shift1);
    d &= shift0;

    // write result back
    n[0] = a;
    n[1] = b;
    n[2] = c;
    n[3] = d;
}