luau_vm/libs/

vector.rs

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

static VECTOR_LIB: [NativeFunction; 14] = [
    NativeFunction {
        name: "create",
        function: vector_create,
    },
    NativeFunction {
        name: "magnitude",
        function: vector_magnitude,
    },
    NativeFunction {
        name: "normalize",
        function: vector_normalize,
    },
    NativeFunction {
        name: "cross",
        function: vector_cross,
    },
    NativeFunction {
        name: "dot",
        function: vector_dot,
    },
    NativeFunction {
        name: "angle",
        function: vector_angle,
    },
    NativeFunction {
        name: "floor",
        function: vector_floor,
    },
    NativeFunction {
        name: "ceil",
        function: vector_ceil,
    },
    NativeFunction {
        name: "abs",
        function: vector_abs,
    },
    NativeFunction {
        name: "sign",
        function: vector_sign,
    },
    NativeFunction {
        name: "clamp",
        function: vector_clamp,
    },
    NativeFunction {
        name: "max",
        function: vector_max,
    },
    NativeFunction {
        name: "min",
        function: vector_min,
    },
    NativeFunction {
        name: "lerp",
        function: vector_lerp,
    },
];

/// `vector_create`
fn vector_create(ctx: NativeCallContext) -> NativeCallResult {
    let count = ctx.arg_count();
    let x = ctx.arg(1).number()? as f32;
    let y = ctx.arg(2).number()? as f32;
    let z = if count >= 3 {
        ctx.arg(3).number()? as f32
    } else {
        0.0
    };

    #[cfg(feature = "vector4")]
    {
        let w = if count >= 4 {
            ctx.arg(4).number()? as f32
        } else {
            0.0
        };
        ctx.push_vector([x, y, z, w])?;
        Ok(1)
    }

    #[cfg(not(feature = "vector4"))]
    {
        ctx.push_vector([x, y, z])?;
        Ok(1)
    }
}

/// `vector_magnitude`
fn vector_magnitude(ctx: NativeCallContext) -> NativeCallResult {
    let vector = ctx.arg(1).vector()?;
    ctx.push_number(
        vector
            .iter()
            .map(|component| component * component)
            .sum::<f32>()
            .sqrt() as f64,
    )?;
    Ok(1)
}

/// `vector_normalize`
fn vector_normalize(ctx: NativeCallContext) -> NativeCallResult {
    let mut vector = ctx.arg(1).vector()?;
    let inv_sqrt = 1.0f32
        / vector
            .iter()
            .map(|component| component * component)
            .sum::<f32>()
            .sqrt();
    for component in &mut vector {
        *component *= inv_sqrt;
    }
    ctx.push_vector(vector)?;
    Ok(1)
}

/// `vector_cross`
fn vector_cross(ctx: NativeCallContext) -> NativeCallResult {
    let left = ctx.arg(1).vector()?;
    let right = ctx.arg(2).vector()?;

    #[cfg(feature = "vector4")]
    {
        ctx.push_vector([
            left[1] * right[2] - left[2] * right[1],
            left[2] * right[0] - left[0] * right[2],
            left[0] * right[1] - left[1] * right[0],
            0.0,
        ])?;
        Ok(1)
    }

    #[cfg(not(feature = "vector4"))]
    {
        ctx.push_vector([
            left[1] * right[2] - left[2] * right[1],
            left[2] * right[0] - left[0] * right[2],
            left[0] * right[1] - left[1] * right[0],
        ])?;
        Ok(1)
    }
}

/// `vector_dot`
fn vector_dot(ctx: NativeCallContext) -> NativeCallResult {
    let left = ctx.arg(1).vector()?;
    let right = ctx.arg(2).vector()?;
    ctx.push_number(
        left.iter()
            .zip(right.iter())
            .map(|(left, right)| left * right)
            .sum::<f32>() as f64,
    )?;
    Ok(1)
}

/// `vector_angle`
fn vector_angle(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let angle = unsafe {
        let left = ctx.arg(1).vector()?;
        let right = ctx.arg(2).vector()?;
        let axis = thread.opt_vector(3)?;

        let cross = [
            left[1] * right[2] - left[2] * right[1],
            left[2] * right[0] - left[0] * right[2],
            left[0] * right[1] - left[1] * right[0],
        ];
        let sin = (cross[0] * cross[0] + cross[1] * cross[1] + cross[2] * cross[2]).sqrt();
        let cos = left[0] * right[0] + left[1] * right[1] + left[2] * right[2];
        let mut angle = (sin as f64).atan2(cos as f64);

        if let Some(axis) = axis
            && cross[0] * axis[0] + cross[1] * axis[1] + cross[2] * axis[2] < 0.0
        {
            angle = -angle;
        }

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

/// `vector_floor`
fn vector_floor(ctx: NativeCallContext) -> NativeCallResult {
    let mut vector = ctx.arg(1).vector()?;
    for component in &mut vector {
        *component = component.floor();
    }
    ctx.push_vector(vector)?;
    Ok(1)
}

/// `vector_ceil`
fn vector_ceil(ctx: NativeCallContext) -> NativeCallResult {
    let mut vector = ctx.arg(1).vector()?;
    for component in &mut vector {
        *component = component.ceil();
    }
    ctx.push_vector(vector)?;
    Ok(1)
}

/// `vector_abs`
fn vector_abs(ctx: NativeCallContext) -> NativeCallResult {
    let mut vector = ctx.arg(1).vector()?;
    for component in &mut vector {
        *component = component.abs();
    }
    ctx.push_vector(vector)?;
    Ok(1)
}

/// `vector_sign`
fn vector_sign(ctx: NativeCallContext) -> NativeCallResult {
    let mut vector = ctx.arg(1).vector()?;
    for component in &mut vector {
        *component = crate::number::sign_f(*component);
    }
    ctx.push_vector(vector)?;
    Ok(1)
}

/// `vector_clamp`
fn vector_clamp(ctx: NativeCallContext) -> NativeCallResult {
    let vector = ctx.arg(1).vector()?;
    let min = ctx.arg(2).vector()?;
    let max = ctx.arg(3).vector()?;

    if min[0] > max[0] {
        return ctx
            .arg(3)
            .error("max.x must be greater than or equal to min.x")
            .map_err(Into::into);
    }
    if min[1] > max[1] {
        return ctx
            .arg(3)
            .error("max.y must be greater than or equal to min.y")
            .map_err(Into::into);
    }
    if min[2] > max[2] {
        return ctx
            .arg(3)
            .error("max.z must be greater than or equal to min.z")
            .map_err(Into::into);
    }

    let mut out = vector;
    for index in 0..crate::types::LUA_VECTOR_SIZE {
        out[index] = crate::number::clamp_f(vector[index], min[index], max[index]);
    }
    ctx.push_vector(out)?;
    Ok(1)
}

/// `vector_min`
fn vector_min(ctx: NativeCallContext) -> NativeCallResult {
    let mut result = ctx.arg(1).vector()?;
    for argument in ctx.args().skip(1) {
        let vector = argument.vector()?;
        for component in 0..crate::types::LUA_VECTOR_SIZE {
            if vector[component] < result[component] {
                result[component] = vector[component];
            }
        }
    }
    ctx.push_vector(result)?;
    Ok(1)
}

/// `vector_max`
fn vector_max(ctx: NativeCallContext) -> NativeCallResult {
    let mut result = ctx.arg(1).vector()?;
    for argument in ctx.args().skip(1) {
        let vector = argument.vector()?;
        for component in 0..crate::types::LUA_VECTOR_SIZE {
            if vector[component] > result[component] {
                result[component] = vector[component];
            }
        }
    }
    ctx.push_vector(result)?;
    Ok(1)
}

/// `vector_index`
fn vector_index(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let vector = ctx.arg(1).vector()?;
        let name = ctx.arg(2).string()?;

        if name.len() == 1 {
            let component_index = ((name[0] | b' ') as i32) - ('x' as i32);

            #[cfg(feature = "vector4")]
            let component_index = if component_index == -1 {
                3
            } else {
                component_index
            };

            if (0..crate::types::LUA_VECTOR_SIZE as i32).contains(&component_index) {
                ctx.push_number(vector[component_index as usize] as f64)?;
                return Ok(1);
            }
        }

        crate::error!(thread, "attempt to index vector with '%s'", name).map_err(Into::into)
    }
}

/// `vector_lerp`
fn vector_lerp(ctx: NativeCallContext) -> NativeCallResult {
    let left = ctx.arg(1).vector()?;
    let right = ctx.arg(2).vector()?;
    let t = ctx.arg(3).number()? as f32;

    let mut out = left;
    for index in 0..crate::types::LUA_VECTOR_SIZE {
        out[index] = crate::number::lerp_f(left[index], right[index], t);
    }
    ctx.push_vector(out)?;
    Ok(1)
}

/// `createmetatable`
fn create_metatable(thread: &Thread) -> NativeCallResult {
    unsafe { thread.create_table(0, 1)? };

    #[cfg(feature = "vector4")]
    unsafe {
        thread.push_vector([0.0, 0.0, 0.0, 0.0])?;
    }
    #[cfg(not(feature = "vector4"))]
    unsafe {
        thread.push_vector([0.0, 0.0, 0.0])?;
    }

    unsafe {
        thread.push_value(-2)?;
        thread.set_metatable(-2)?;
        thread.pop(1);
        thread.push_native_closure_k(vector_index, None, 0, None)?;
        thread.raw_set_field(-2, "__index")?;
        thread.set_readonly(-1, 1);
        thread.pop(1);
    }
    Ok(0)
}

impl Thread {
    /// `luaopen_vector`
    pub unsafe fn open_vector(&self) -> NativeCallResult {
        unsafe {
            self.register(Some(super::LUA_VECLIB_NAME), &VECTOR_LIB[..])?;

            #[cfg(feature = "vector4")]
            {
                self.push_vector([0.0, 0.0, 0.0, 0.0])?;
                self.raw_set_field(-2, "zero")?;
                self.push_vector([1.0, 1.0, 1.0, 1.0])?;
                self.raw_set_field(-2, "one")?;
            }
            #[cfg(not(feature = "vector4"))]
            {
                self.push_vector([0.0, 0.0, 0.0])?;
                self.raw_set_field(-2, "zero")?;
                self.push_vector([1.0, 1.0, 1.0])?;
                self.raw_set_field(-2, "one")?;
            }

            create_metatable(self)?;
            Ok(1)
        }
    }
}