luau-vm 0.732.0

use crate::VmResult;
use crate::native::{NativeCallContext, NativeCallResult, NativeFunction};
use crate::thread::Thread;

const ALL_ONES: u32 = !0;
const N_BITS: i32 = u32::BITS as i32;

static BIT_LIB: [NativeFunction; 15] = [
    NativeFunction {
        name: "arshift",
        function: b_arshift,
    },
    NativeFunction {
        name: "band",
        function: b_and,
    },
    NativeFunction {
        name: "bnot",
        function: b_not,
    },
    NativeFunction {
        name: "bor",
        function: b_or,
    },
    NativeFunction {
        name: "bxor",
        function: b_xor,
    },
    NativeFunction {
        name: "btest",
        function: b_test,
    },
    NativeFunction {
        name: "extract",
        function: b_extract,
    },
    NativeFunction {
        name: "lrotate",
        function: b_lrot,
    },
    NativeFunction {
        name: "lshift",
        function: b_lshift,
    },
    NativeFunction {
        name: "replace",
        function: b_replace,
    },
    NativeFunction {
        name: "rrotate",
        function: b_rrot,
    },
    NativeFunction {
        name: "rshift",
        function: b_rshift,
    },
    NativeFunction {
        name: "countlz",
        function: b_countlz,
    },
    NativeFunction {
        name: "countrz",
        function: b_countrz,
    },
    NativeFunction {
        name: "byteswap",
        function: b_swap,
    },
];

/// `trim`
fn trim(value: u32) -> u32 {
    value & ALL_ONES
}

/// `mask`
fn mask(width: i32) -> u32 {
    !((ALL_ONES << 1) << (width - 1))
}

/// `andaux`
fn and_aux(ctx: &NativeCallContext) -> VmResult<u32> {
    let mut result = !0u32;
    for argument in ctx.args() {
        result &= argument.unsigned()?;
    }
    Ok(trim(result))
}

/// `b_and`
fn b_and(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(and_aux(&ctx)?)?;
    Ok(1)
}

/// `b_test`
fn b_test(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_boolean(and_aux(&ctx)? != 0)?;
    Ok(1)
}

/// `b_or`
fn b_or(ctx: NativeCallContext) -> NativeCallResult {
    let mut result = 0u32;
    for argument in ctx.args() {
        result |= argument.unsigned()?;
    }
    ctx.push_unsigned(trim(result))?;
    Ok(1)
}

/// `b_xor`
fn b_xor(ctx: NativeCallContext) -> NativeCallResult {
    let mut result = 0u32;
    for argument in ctx.args() {
        result ^= argument.unsigned()?;
    }
    ctx.push_unsigned(trim(result))?;
    Ok(1)
}

/// `b_not`
fn b_not(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(trim(!ctx.arg(1).unsigned()?))?;
    Ok(1)
}

/// `b_shift`
fn b_shift(mut value: u32, mut shift: i32) -> u32 {
    if shift < 0 {
        shift = -shift;
        value = trim(value);
        value = if shift >= N_BITS { 0 } else { value >> shift };
    } else {
        value = if shift >= N_BITS { 0 } else { value << shift };
        value = trim(value);
    }

    value
}

/// `b_lshift`
fn b_lshift(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(b_shift(ctx.arg(1).unsigned()?, ctx.arg(2).integer()?))?;
    Ok(1)
}

/// `b_rshift`
fn b_rshift(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(b_shift(ctx.arg(1).unsigned()?, -ctx.arg(2).integer()?))?;
    Ok(1)
}

/// `b_arshift`
fn b_arshift(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).unsigned()?;
    let shift = ctx.arg(2).integer()?;

    if shift < 0 || (value & (1u32 << (N_BITS - 1))) == 0 {
        ctx.push_unsigned(b_shift(value, -shift))?;
        return Ok(1);
    }

    let result = if shift >= N_BITS {
        ALL_ONES
    } else {
        trim((value >> shift) | !(u32::MAX >> shift))
    };
    ctx.push_unsigned(result)?;
    Ok(1)
}

/// `b_rot`
fn b_rot(mut value: u32, shift: i32) -> u32 {
    let shift = shift & (N_BITS - 1);
    value = trim(value);
    if shift != 0 {
        value = (value << shift) | (value >> (N_BITS - shift));
    }
    trim(value)
}

/// `b_lrot`
fn b_lrot(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(b_rot(ctx.arg(1).unsigned()?, ctx.arg(2).integer()?))?;
    Ok(1)
}

/// `b_rrot`
fn b_rrot(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(b_rot(ctx.arg(1).unsigned()?, -ctx.arg(2).integer()?))?;
    Ok(1)
}

/// `fieldargs`
fn field_args(ctx: &NativeCallContext, field_arg: i32, width_arg: i32) -> VmResult<(i32, i32)> {
    let field = ctx.arg(field_arg).integer()?;
    let width = ctx.arg(width_arg).integer_or(1)?;

    if field < 0 {
        return ctx
            .arg(field_arg)
            .error("field cannot be negative")
            .map_err(Into::into);
    }
    if width <= 0 {
        return ctx
            .arg(width_arg)
            .error("width must be positive")
            .map_err(Into::into);
    }
    if field + width > N_BITS {
        return ctx
            .error("trying to access non-existent bits", [])
            .map_err(Into::into);
    }

    Ok((field, width))
}

/// `b_extract`
fn b_extract(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).unsigned()?;
    let (field, width) = field_args(&ctx, 2, 3)?;
    ctx.push_unsigned((value >> field) & mask(width))?;
    Ok(1)
}

/// `b_replace`
fn b_replace(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).unsigned()?;
    let mut replacement = ctx.arg(2).unsigned()?;
    let (field, width) = field_args(&ctx, 3, 4)?;
    let mask = mask(width);
    replacement &= mask;
    ctx.push_unsigned((value & !(mask << field)) | (replacement << field))?;
    Ok(1)
}

/// `b_countlz`
fn b_countlz(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).unsigned()?;
    let mut result = N_BITS as u32;
    for index in 0..N_BITS {
        if (value & (1u32 << (N_BITS - 1 - index))) != 0 {
            result = index as u32;
            break;
        }
    }
    ctx.push_unsigned(result)?;
    Ok(1)
}

/// `b_countrz`
fn b_countrz(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).unsigned()?;
    let mut result = N_BITS as u32;
    for index in 0..N_BITS {
        if (value & (1u32 << index)) != 0 {
            result = index as u32;
            break;
        }
    }
    ctx.push_unsigned(result)?;
    Ok(1)
}

/// `b_swap`
fn b_swap(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_unsigned(ctx.arg(1).unsigned()?.swap_bytes())?;
    Ok(1)
}

impl Thread {
    /// `luaopen_bit32`
    pub unsafe fn open_bit32(&self) -> NativeCallResult {
        unsafe { self.register(Some(super::LUA_BITLIB_NAME), &BIT_LIB[..])? };
        Ok(1)
    }
}