mtl-rs 0.2.2

use core::ops::Range;

use objc2::{Message, extern_protocol, msg_send, runtime::ProtocolObject};
use objc2_foundation::NSRange;

use super::MTLIntersectionFunctionSignature;
use crate::{
    MTLBuffer, MTLFunctionHandle, MTLResource, MTLVisibleFunctionTable, types::MTLResourceID,
    util::opt_ref_slice_as_ptr,
};

extern_protocol!(
    /// Intersection function table
    ///
    /// Availability: macOS 11.0+, iOS 14.0+, tvOS 16.0+
    ///
    /// # Safety
    ///
    /// Implementors must be Objective-C objects that conform to the
    /// `MTLIntersectionFunctionTable` protocol.
    pub unsafe trait MTLIntersectionFunctionTable: MTLResource {
        #[unsafe(method(setBuffer:offset:atIndex:))]
        #[unsafe(method_family = none)]
        fn set_buffer(
            &self,
            buffer: Option<&ProtocolObject<dyn MTLBuffer>>,
            offset: usize,
            index: usize,
        );

        /// Handle of the GPU resource suitable for storing in an Argument Buffer
        ///
        /// Availability: macOS 13.0+, iOS 16.0+
        #[unsafe(method(gpuResourceID))]
        #[unsafe(method_family = none)]
        fn gpu_resource_id(&self) -> MTLResourceID;

        #[unsafe(method(setFunction:atIndex:))]
        #[unsafe(method_family = none)]
        fn set_function_at_index(
            &self,
            function: Option<&ProtocolObject<dyn MTLFunctionHandle>>,
            index: usize,
        );

        #[unsafe(method(setOpaqueTriangleIntersectionFunctionWithSignature:atIndex:))]
        #[unsafe(method_family = none)]
        fn set_opaque_triangle_intersection_function_with_signature_at_index(
            &self,
            signature: MTLIntersectionFunctionSignature,
            index: usize,
        );

        #[unsafe(method(setOpaqueCurveIntersectionFunctionWithSignature:atIndex:))]
        #[unsafe(method_family = none)]
        fn set_opaque_curve_intersection_function_with_signature_at_index(
            &self,
            signature: MTLIntersectionFunctionSignature,
            index: usize,
        );

        #[unsafe(method(setVisibleFunctionTable:atBufferIndex:))]
        #[unsafe(method_family = none)]
        fn set_visible_function_table_at_buffer_index(
            &self,
            function_table: Option<&ProtocolObject<dyn MTLVisibleFunctionTable>>,
            buffer_index: usize,
        );
    }
);

pub trait MTLIntersectionFunctionTableExt: MTLIntersectionFunctionTable + Message {
    /// Set an array of buffers at the given bind point index range.
    fn set_buffers(
        &self,
        buffers: &[Option<&ProtocolObject<dyn MTLBuffer>>],
        offsets: &[usize],
        range: Range<usize>,
    ) where
        Self: Sized,
    {
        assert_eq!(buffers.len(), offsets.len(), "buffer and offset slices must have equal lengths");
        let range = checked_range(range, buffers.len());
        let ptr = opt_ref_slice_as_ptr(buffers);
        unsafe { msg_send![self, setBuffers: ptr, offsets: offsets.as_ptr(), withRange: range] }
    }

    /// Set an array of functions at the given index range.
    fn set_functions(
        &self,
        functions: &[Option<&ProtocolObject<dyn MTLFunctionHandle>>],
        range: Range<usize>,
    ) where
        Self: Sized,
    {
        let range = checked_range(range, functions.len());
        let ptr = opt_ref_slice_as_ptr(functions);
        unsafe { msg_send![self, setFunctions: ptr, withRange: range] }
    }

    /// Set an array of visible function tables at the given buffer index range.
    fn set_visible_function_tables(
        &self,
        tables: &[Option<&ProtocolObject<dyn MTLVisibleFunctionTable>>],
        range: Range<usize>,
    ) where
        Self: Sized,
    {
        let range = checked_range(range, tables.len());
        let ptr = opt_ref_slice_as_ptr(tables);
        unsafe { msg_send![self, setVisibleFunctionTables: ptr, withBufferRange: range] }
    }

    /// Set the opaque triangle intersection function for an index range.
    fn set_opaque_triangle_intersection_function_with_signature(
        &self,
        signature: MTLIntersectionFunctionSignature,
        range: Range<usize>,
    ) where
        Self: Sized,
    {
        let range = NSRange::from(range);
        unsafe {
            msg_send![
                self,
                setOpaqueTriangleIntersectionFunctionWithSignature: signature,
                withRange: range
            ]
        }
    }

    /// Set the opaque curve intersection function for an index range.
    fn set_opaque_curve_intersection_function_with_signature(
        &self,
        signature: MTLIntersectionFunctionSignature,
        range: Range<usize>,
    ) where
        Self: Sized,
    {
        let range = NSRange::from(range);
        unsafe {
            msg_send![
                self,
                setOpaqueCurveIntersectionFunctionWithSignature: signature,
                withRange: range
            ]
        }
    }
}

impl<T: MTLIntersectionFunctionTable + Message> MTLIntersectionFunctionTableExt for T {}

fn checked_range(
    range: Range<usize>,
    expected_len: usize,
) -> NSRange {
    let range = NSRange::from(range);
    assert_eq!(range.length, expected_len, "range length must match slice length");
    range
}

#[cfg(test)]
mod tests {
    use super::checked_range;

    #[test]
    fn checked_range_accepts_a_matching_slice_length() {
        assert_eq!(checked_range(4..7, 3), objc2_foundation::NSRange::new(4, 3));
    }

    #[test]
    #[should_panic(expected = "range length must match slice length")]
    fn checked_range_rejects_a_mismatched_slice_length() {
        let _ = checked_range(4..7, 2);
    }
}