luau_vm/libs/

math.rs

use luau_common::flags;
use web_time::{SystemTime, UNIX_EPOCH};

use crate::handle::RawHandle;
use crate::native::{NativeCallContext, NativeCallResult, NativeFunction};
use crate::thread::{LUA_TNONE, Thread};
use crate::types::LUA_TNIL;

const LUAU_PI: f64 = core::f64::consts::PI;
const RADIANS_PER_DEGREE: f64 = LUAU_PI / 180.0;
const LUAU_E: f64 = core::f64::consts::E;
#[allow(clippy::excessive_precision)] // Matches the source constant before f64 rounding.
const LUAU_PHI: f64 = 1.61803398874989484820;
const LUAU_SQRT2: f64 = core::f64::consts::SQRT_2;
const LUAU_TAU: f64 = core::f64::consts::TAU;
const PCG32_INC: u64 = 105;

static MATH_LIB: [NativeFunction; 37] = [
    NativeFunction {
        name: "abs",
        function: math_abs,
    },
    NativeFunction {
        name: "acos",
        function: math_acos,
    },
    NativeFunction {
        name: "asin",
        function: math_asin,
    },
    NativeFunction {
        name: "atan2",
        function: math_atan2,
    },
    NativeFunction {
        name: "atan",
        function: math_atan,
    },
    NativeFunction {
        name: "ceil",
        function: math_ceil,
    },
    NativeFunction {
        name: "cosh",
        function: math_cosh,
    },
    NativeFunction {
        name: "cos",
        function: math_cos,
    },
    NativeFunction {
        name: "deg",
        function: math_deg,
    },
    NativeFunction {
        name: "exp",
        function: math_exp,
    },
    NativeFunction {
        name: "floor",
        function: math_floor,
    },
    NativeFunction {
        name: "fmod",
        function: math_fmod,
    },
    NativeFunction {
        name: "frexp",
        function: math_frexp,
    },
    NativeFunction {
        name: "ldexp",
        function: math_ldexp,
    },
    NativeFunction {
        name: "log10",
        function: math_log10,
    },
    NativeFunction {
        name: "log",
        function: math_log,
    },
    NativeFunction {
        name: "max",
        function: math_max,
    },
    NativeFunction {
        name: "min",
        function: math_min,
    },
    NativeFunction {
        name: "modf",
        function: math_modf,
    },
    NativeFunction {
        name: "pow",
        function: math_pow,
    },
    NativeFunction {
        name: "rad",
        function: math_rad,
    },
    NativeFunction {
        name: "random",
        function: math_random,
    },
    NativeFunction {
        name: "randomseed",
        function: math_randomseed,
    },
    NativeFunction {
        name: "sinh",
        function: math_sinh,
    },
    NativeFunction {
        name: "sin",
        function: math_sin,
    },
    NativeFunction {
        name: "sqrt",
        function: math_sqrt,
    },
    NativeFunction {
        name: "tanh",
        function: math_tanh,
    },
    NativeFunction {
        name: "tan",
        function: math_tan,
    },
    NativeFunction {
        name: "noise",
        function: math_noise,
    },
    NativeFunction {
        name: "clamp",
        function: math_clamp,
    },
    NativeFunction {
        name: "sign",
        function: math_sign,
    },
    NativeFunction {
        name: "round",
        function: math_round,
    },
    NativeFunction {
        name: "map",
        function: math_map,
    },
    NativeFunction {
        name: "lerp",
        function: math_lerp,
    },
    NativeFunction {
        name: "isnan",
        function: math_isnan,
    },
    NativeFunction {
        name: "isinf",
        function: math_isinf,
    },
    NativeFunction {
        name: "isfinite",
        function: math_isfinite,
    },
];

const PERLIN_HASH: [u8; 257] = [
    151, 160, 137, 91, 90, 15, 131, 13, 201, 95, 96, 53, 194, 233, 7, 225, 140, 36, 103, 30, 69,
    142, 8, 99, 37, 240, 21, 10, 23, 190, 6, 148, 247, 120, 234, 75, 0, 26, 197, 62, 94, 252, 219,
    203, 117, 35, 11, 32, 57, 177, 33, 88, 237, 149, 56, 87, 174, 20, 125, 136, 171, 168, 68, 175,
    74, 165, 71, 134, 139, 48, 27, 166, 77, 146, 158, 231, 83, 111, 229, 122, 60, 211, 133, 230,
    220, 105, 92, 41, 55, 46, 245, 40, 244, 102, 143, 54, 65, 25, 63, 161, 1, 216, 80, 73, 209, 76,
    132, 187, 208, 89, 18, 169, 200, 196, 135, 130, 116, 188, 159, 86, 164, 100, 109, 198, 173,
    186, 3, 64, 52, 217, 226, 250, 124, 123, 5, 202, 38, 147, 118, 126, 255, 82, 85, 212, 207, 206,
    59, 227, 47, 16, 58, 17, 182, 189, 28, 42, 223, 183, 170, 213, 119, 248, 152, 2, 44, 154, 163,
    70, 221, 153, 101, 155, 167, 43, 172, 9, 129, 22, 39, 253, 19, 98, 108, 110, 79, 113, 224, 232,
    178, 185, 112, 104, 218, 246, 97, 228, 251, 34, 242, 193, 238, 210, 144, 12, 191, 179, 162,
    241, 81, 51, 145, 235, 249, 14, 239, 107, 49, 192, 214, 31, 181, 199, 106, 157, 184, 84, 204,
    176, 115, 121, 50, 45, 127, 4, 150, 254, 138, 236, 205, 93, 222, 114, 67, 29, 24, 72, 243, 141,
    128, 195, 78, 66, 215, 61, 156, 180, 151,
];

const PERLIN_GRAD: [[f32; 3]; 16] = [
    [1.0, 1.0, 0.0],
    [-1.0, 1.0, 0.0],
    [1.0, -1.0, 0.0],
    [-1.0, -1.0, 0.0],
    [1.0, 0.0, 1.0],
    [-1.0, 0.0, 1.0],
    [1.0, 0.0, -1.0],
    [-1.0, 0.0, -1.0],
    [0.0, 1.0, 1.0],
    [0.0, -1.0, 1.0],
    [0.0, 1.0, -1.0],
    [0.0, -1.0, -1.0],
    [1.0, 1.0, 0.0],
    [0.0, -1.0, 1.0],
    [-1.0, 1.0, 0.0],
    [0.0, -1.0, -1.0],
];

/// `pcg32_random`
fn pcg32_random(state: &mut u64) -> u32 {
    let old_state = *state;
    *state = old_state
        .wrapping_mul(6364136223846793005)
        .wrapping_add(PCG32_INC | 1);
    let xorshifted = (((old_state >> 18) ^ old_state) >> 27) as u32;
    let rot = (old_state >> 59) as u32;
    xorshifted.rotate_right(rot)
}

/// `pcg32_seed`
fn pcg32_seed(state: &mut u64, seed: u64) {
    *state = 0;
    let _ = pcg32_random(state);
    *state = state.wrapping_add(seed);
    let _ = pcg32_random(state);
}

/// `perlin_fade`
fn perlin_fade(t: f32) -> f32 {
    t * t * t * (t * (t * 6.0 - 15.0) + 10.0)
}

/// `perlin_lerp`
fn perlin_lerp(t: f32, a: f32, b: f32) -> f32 {
    a + t * (b - a)
}

/// `perlin_grad`
fn perlin_grad(hash: i32, x: f32, y: f32, z: f32) -> f32 {
    let grad = PERLIN_GRAD[(hash & 15) as usize];
    grad[0] * x + grad[1] * y + grad[2] * z
}

/// `perlin`
fn perlin(x: f32, y: f32, z: f32) -> f32 {
    let x_floor = x.floor();
    let y_floor = y.floor();
    let z_floor = z.floor();

    let xi = (x_floor as i32) & 255;
    let yi = (y_floor as i32) & 255;
    let zi = (z_floor as i32) & 255;

    let xf = x - x_floor;
    let yf = y - y_floor;
    let zf = z - z_floor;

    let u = perlin_fade(xf);
    let v = perlin_fade(yf);
    let w = perlin_fade(zf);

    let p = &PERLIN_HASH;

    let a = (p[xi as usize] as i32 + yi) & 255;
    let aa = (p[a as usize] as i32 + zi) & 255;
    let ab = (p[(a + 1) as usize] as i32 + zi) & 255;

    let b = (p[(xi + 1) as usize] as i32 + yi) & 255;
    let ba = (p[b as usize] as i32 + zi) & 255;
    let bb = (p[(b + 1) as usize] as i32 + zi) & 255;

    let la = perlin_lerp(
        u,
        perlin_grad(p[aa as usize] as i32, xf, yf, zf),
        perlin_grad(p[ba as usize] as i32, xf - 1.0, yf, zf),
    );
    let lb = perlin_lerp(
        u,
        perlin_grad(p[ab as usize] as i32, xf, yf - 1.0, zf),
        perlin_grad(p[bb as usize] as i32, xf - 1.0, yf - 1.0, zf),
    );
    let la1 = perlin_lerp(
        u,
        perlin_grad(p[(aa + 1) as usize] as i32, xf, yf, zf - 1.0),
        perlin_grad(p[(ba + 1) as usize] as i32, xf - 1.0, yf, zf - 1.0),
    );
    let lb1 = perlin_lerp(
        u,
        perlin_grad(p[(ab + 1) as usize] as i32, xf, yf - 1.0, zf - 1.0),
        perlin_grad(p[(bb + 1) as usize] as i32, xf - 1.0, yf - 1.0, zf - 1.0),
    );

    perlin_lerp(w, perlin_lerp(v, la, lb), perlin_lerp(v, la1, lb1))
}

fn modf_parts(value: f64) -> (f64, f64) {
    let integer = value.trunc();
    let fractional = value - integer;
    (integer, fractional)
}

fn frexp_parts(value: f64) -> (f64, i32) {
    if value == 0.0 || !value.is_finite() {
        return (value, 0);
    }

    let bits = value.to_bits();
    let sign = bits & (1u64 << 63);
    let exponent = ((bits >> 52) & 0x7ff) as i32;
    let mantissa = bits & ((1u64 << 52) - 1);

    if exponent == 0 {
        let (mantissa, exponent) = frexp_parts(value * ((1u64 << 53) as f64));
        return (mantissa, exponent - 53);
    }

    let normalized = f64::from_bits(sign | (1022u64 << 52) | mantissa);
    (normalized, exponent - 1022)
}

fn ldexp_value(value: f64, exponent: i32) -> f64 {
    value * (2.0f64).powi(exponent)
}

/// `math_abs`
fn math_abs(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.abs())?;
    Ok(1)
}

/// `math_sin`
fn math_sin(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.sin())?;
    Ok(1)
}

/// `math_sinh`
fn math_sinh(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.sinh())?;
    Ok(1)
}

/// `math_cos`
fn math_cos(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.cos())?;
    Ok(1)
}

/// `math_cosh`
fn math_cosh(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.cosh())?;
    Ok(1)
}

/// `math_tan`
fn math_tan(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.tan())?;
    Ok(1)
}

/// `math_tanh`
fn math_tanh(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.tanh())?;
    Ok(1)
}

/// `math_asin`
fn math_asin(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.asin())?;
    Ok(1)
}

/// `math_acos`
fn math_acos(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.acos())?;
    Ok(1)
}

/// `math_atan`
fn math_atan(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.atan())?;
    Ok(1)
}

/// `math_atan2`
fn math_atan2(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.atan2(ctx.arg(2).number()?))?;
    Ok(1)
}

/// `math_ceil`
fn math_ceil(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.ceil())?;
    Ok(1)
}

/// `math_floor`
fn math_floor(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.floor())?;
    Ok(1)
}

/// `math_fmod`
fn math_fmod(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()? % ctx.arg(2).number()?)?;
    Ok(1)
}

/// `math_modf`
fn math_modf(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let (integer, fractional) = modf_parts(unsafe { thread.check_number(1)? });
    unsafe {
        thread.push_number(integer)?;
        thread.push_number(fractional)?;
    }
    Ok(2)
}

/// `math_sqrt`
fn math_sqrt(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.sqrt())?;
    Ok(1)
}

/// `math_pow`
fn math_pow(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.powf(ctx.arg(2).number()?))?;
    Ok(1)
}

/// `math_log`
fn math_log(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let result = unsafe {
        let x = thread.check_number(1)?;
        if matches!(thread.type_of(2), LUA_TNONE | LUA_TNIL) {
            x.ln()
        } else {
            let base = thread.check_number(2)?;
            if base == 2.0 {
                x.log2()
            } else if base == 10.0 {
                x.log10()
            } else {
                x.ln() / base.ln()
            }
        }
    };

    ctx.push_number(result)?;
    Ok(1)
}

/// `math_log10`
fn math_log10(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.log10())?;
    Ok(1)
}

/// `math_exp`
fn math_exp(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.exp())?;
    Ok(1)
}

/// `math_deg`
fn math_deg(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()? / RADIANS_PER_DEGREE)?;
    Ok(1)
}

/// `math_rad`
fn math_rad(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()? * RADIANS_PER_DEGREE)?;
    Ok(1)
}

/// `math_frexp`
fn math_frexp(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let (mantissa, exponent) = frexp_parts(unsafe { thread.check_number(1)? });
    unsafe {
        thread.push_number(mantissa)?;
        thread.push_integer(exponent)?;
    }
    Ok(2)
}

/// `math_ldexp`
fn math_ldexp(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ldexp_value(ctx.arg(1).number()?, ctx.arg(2).integer()?))?;
    Ok(1)
}

/// `math_min`
fn math_min(ctx: NativeCallContext) -> NativeCallResult {
    let mut result = ctx.arg(1).number()?;
    for argument in ctx.args().skip(1) {
        let value = argument.number()?;
        if value < result {
            result = value;
        }
    }
    ctx.push_number(result)?;
    Ok(1)
}

/// `math_max`
fn math_max(ctx: NativeCallContext) -> NativeCallResult {
    let mut result = ctx.arg(1).number()?;
    for argument in ctx.args().skip(1) {
        let value = argument.number()?;
        if value > result {
            result = value;
        }
    }
    ctx.push_number(result)?;
    Ok(1)
}

/// `math_random`
fn math_random(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let global = unsafe { thread.global() };
    let state = unsafe { &mut global.as_ptr().as_mut().unwrap_unchecked().rng_state };
    unsafe {
        match thread.get_top() {
            0 => {
                let low = pcg32_random(state);
                let high = pcg32_random(state);
                let result = ldexp_value((low as u64 | ((high as u64) << 32)) as f64, -64);
                thread.push_number(result)?;
            }
            1 => {
                let upper = thread.check_integer(1)?;
                if upper < 1 {
                    return thread
                        .lua_arg_error(1, "interval is empty")
                        .map_err(Into::into);
                }

                let x = upper as u64 * pcg32_random(state) as u64;
                let result = 1 + (x >> 32) as i32;
                thread.push_integer(result)?;
            }
            2 => {
                let lower = thread.check_integer(1)?;
                let upper = thread.check_integer(2)?;
                if lower > upper {
                    return thread
                        .lua_arg_error(2, "interval is empty")
                        .map_err(Into::into);
                }

                let interval = (upper as u32).wrapping_sub(lower as u32);
                if interval == u32::MAX {
                    return thread
                        .lua_arg_error(2, "interval is too large")
                        .map_err(Into::into);
                }

                let x = (interval as u64 + 1) * pcg32_random(state) as u64;
                let result = lower.wrapping_add((x >> 32) as i32);
                thread.push_integer(result)?;
            }
            _ => {
                return crate::error!(thread, "wrong number of arguments").map_err(Into::into);
            }
        }
    }
    Ok(1)
}

/// `math_randomseed`
fn math_randomseed(ctx: NativeCallContext) -> NativeCallResult {
    let seed = ctx.arg(1).integer()? as u64;
    let global = unsafe { ctx.raw_thread().global() };
    pcg32_seed(
        unsafe { &mut global.as_ptr().as_mut().unwrap_unchecked().rng_state },
        seed,
    );
    Ok(0)
}

/// `math_noise`
fn math_noise(ctx: NativeCallContext) -> NativeCallResult {
    let mut x = ctx.arg(1).number()?;
    let mut y = ctx.arg(2).number_or(0.0)?;
    let mut z = ctx.arg(3).number_or(0.0)?;

    if flags::FixMathNoisePrecision.get() {
        x %= 256.0;
        y %= 256.0;
        z %= 256.0;
    }

    ctx.push_number(perlin(x as f32, y as f32, z as f32) as f64)?;
    Ok(1)
}

/// `math_clamp`
fn math_clamp(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).number()?;
    let min = ctx.arg(2).number()?;
    let max = ctx.arg(3).number()?;
    if min > max {
        return ctx
            .arg(3)
            .error("max must be greater than or equal to min")
            .map_err(Into::into);
    }

    ctx.push_number(value.clamp(min, max))?;
    Ok(1)
}

/// `math_sign`
fn math_sign(ctx: NativeCallContext) -> NativeCallResult {
    let value = ctx.arg(1).number()?;
    ctx.push_number(if value > 0.0 {
        1.0
    } else if value < 0.0 {
        -1.0
    } else {
        0.0
    })?;
    Ok(1)
}

/// `math_round`
fn math_round(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_number(ctx.arg(1).number()?.round())?;
    Ok(1)
}

/// `math_map`
fn math_map(ctx: NativeCallContext) -> NativeCallResult {
    let x = ctx.arg(1).number()?;
    let in_min = ctx.arg(2).number()?;
    let in_max = ctx.arg(3).number()?;
    let out_min = ctx.arg(4).number()?;
    let out_max = ctx.arg(5).number()?;

    ctx.push_number(out_min + (x - in_min) * (out_max - out_min) / (in_max - in_min))?;
    Ok(1)
}

/// `math_lerp`
fn math_lerp(ctx: NativeCallContext) -> NativeCallResult {
    let a = ctx.arg(1).number()?;
    let b = ctx.arg(2).number()?;
    let t = ctx.arg(3).number()?;
    ctx.push_number(if t == 1.0 { b } else { a + (b - a) * t })?;
    Ok(1)
}

/// `math_isnan`
fn math_isnan(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_boolean(ctx.arg(1).number()?.is_nan())?;
    Ok(1)
}

/// `math_isinf`
fn math_isinf(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_boolean(ctx.arg(1).number()?.is_infinite())?;
    Ok(1)
}

/// `math_isfinite`
fn math_isfinite(ctx: NativeCallContext) -> NativeCallResult {
    ctx.push_boolean(ctx.arg(1).number()?.is_finite())?;
    Ok(1)
}

impl Thread {
    /// `luaopen_math`
    pub unsafe fn open_math(&self) -> NativeCallResult {
        unsafe {
            let mut seed = self.encode_pointer(self.as_ptr() as usize) as u64;
            seed ^= SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .map_or(0, |duration| duration.as_secs());
            seed ^= crate::clock().to_bits();

            pcg32_seed(
                &mut self.global().as_ptr().as_mut().unwrap_unchecked().rng_state,
                seed,
            );

            self.register(Some(super::LUA_MATHLIB_NAME), &MATH_LIB[..])?;
            self.push_number(LUAU_PI)?;
            self.raw_set_field(-2, "pi")?;
            self.push_number(f64::INFINITY)?;
            self.raw_set_field(-2, "huge")?;
            self.push_number(f64::NAN)?;
            self.raw_set_field(-2, "nan")?;
            self.push_number(LUAU_E)?;
            self.raw_set_field(-2, "e")?;
            self.push_number(LUAU_PHI)?;
            self.raw_set_field(-2, "phi")?;
            self.push_number(LUAU_SQRT2)?;
            self.raw_set_field(-2, "sqrt2")?;
            self.push_number(LUAU_TAU)?;
            self.raw_set_field(-2, "tau")?;
            Ok(1)
        }
    }
}