mtl-rs 0.1.10

Rust bindings for Apple's Metal API
use core::ops::Range;

use objc2::{
    Message,
    encode::{Encode, Encoding, RefEncode},
    extern_class, extern_conformance, extern_methods, extern_protocol, msg_send,
    rc::{Allocated, Retained},
};
use objc2_foundation::{CopyingHelper, NSCopying, NSData, NSInteger, NSObject, NSObjectProtocol, NSRange, NSString};

/// Represents a timestamp data entry in a counter heap of type `MTL4CounterHeapTypeTimestamp`.
///
/// See also [Apple's documentation](https://developer.apple.com/documentation/metal/mtl4timestampheapentry?language=objc)
#[repr(C)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MTL4TimestampHeapEntry {
    pub timestamp: u64,
}

unsafe impl Encode for MTL4TimestampHeapEntry {
    const ENCODING: Encoding = Encoding::Struct("?", &[<u64>::ENCODING]);
}

unsafe impl RefEncode for MTL4TimestampHeapEntry {
    const ENCODING_REF: Encoding = Encoding::Pointer(&Self::ENCODING);
}

/// Defines the type of a ``MTL4CounterHeap`` and the contents of its entries.
///
/// See also [Apple's documentation](https://developer.apple.com/documentation/metal/mtl4counterheaptype?language=objc)
// NS_ENUM
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct MTL4CounterHeapType(pub NSInteger);
impl MTL4CounterHeapType {
    /// Specifies that ``MTL4CounterHeap`` entries contain invalid data.
    #[doc(alias = "MTL4CounterHeapTypeInvalid")]
    pub const INVALID: Self = Self(0);
    /// Specifies that ``MTL4CounterHeap`` entries contain GPU timestamp data.
    #[doc(alias = "MTL4CounterHeapTypeTimestamp")]
    pub const TIMESTAMP: Self = Self(1);
}

unsafe impl Encode for MTL4CounterHeapType {
    const ENCODING: Encoding = NSInteger::ENCODING;
}

unsafe impl RefEncode for MTL4CounterHeapType {
    const ENCODING_REF: Encoding = Encoding::Pointer(&Self::ENCODING);
}

/// Provides a hint to the system about the desired accuracy when writing GPU counter timestamps.
///
/// Pass these values to ``MTL4ComputeCommandEncoder/writeTimestampWithGranularity:intoHeap:atIndex:`` and
/// ``MTL4RenderCommandEncoder/writeTimestampWithGranularity:afterStage:intoHeap:atIndex:`` to control the
/// desired accurracy of the counter sampling operation.
///
/// See also [Apple's documentation](https://developer.apple.com/documentation/metal/mtl4timestampgranularity?language=objc)
// NS_ENUM
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct MTL4TimestampGranularity(pub NSInteger);
impl MTL4TimestampGranularity {
    /// A minimally-invasive timestamp which may be less precise.
    ///
    /// Using this granularity incurs in the lowest overhead, at the cost of precision. For example, it may sample at
    /// command encoder boundaries.
    #[doc(alias = "MTL4TimestampGranularityRelaxed")]
    pub const RELAXED: Self = Self(0);
    /// A timestamp as precise as possible.
    ///
    /// Using this granularity may incur in a performance penalty, for example, it may cause splitting of command encoders.
    #[doc(alias = "MTL4TimestampGranularityPrecise")]
    pub const PRECISE: Self = Self(1);
}

unsafe impl Encode for MTL4TimestampGranularity {
    const ENCODING: Encoding = NSInteger::ENCODING;
}

unsafe impl RefEncode for MTL4TimestampGranularity {
    const ENCODING_REF: Encoding = Encoding::Pointer(&Self::ENCODING);
}

extern_class!(
    /// Groups together parameters for configuring a counter heap object at creation time.
    ///
    /// See also [Apple's documentation](https://developer.apple.com/documentation/metal/mtl4counterheapdescriptor?language=objc)
    #[unsafe(super(NSObject))]
    #[derive(Debug, PartialEq, Eq, Hash)]
    pub struct MTL4CounterHeapDescriptor;
);

extern_conformance!(
    unsafe impl NSCopying for MTL4CounterHeapDescriptor {}
);

unsafe impl CopyingHelper for MTL4CounterHeapDescriptor {
    type Result = Self;
}

extern_conformance!(
    unsafe impl NSObjectProtocol for MTL4CounterHeapDescriptor {}
);

impl MTL4CounterHeapDescriptor {
    extern_methods!(
        /// Assigns the type of data that the heap contains.
        #[unsafe(method(type))]
        #[unsafe(method_family = none)]
        pub fn r#type(&self) -> MTL4CounterHeapType;

        /// Setter for [`type`][Self::type].
        #[unsafe(method(setType:))]
        #[unsafe(method_family = none)]
        pub fn set_type(
            &self,
            r#type: MTL4CounterHeapType,
        );

        /// Assigns the number of entries in the heap.
        ///
        /// Each entry represents one item in the heap. The size of the individual entries depends on the heap type.
        #[unsafe(method(count))]
        #[unsafe(method_family = none)]
        pub fn count(&self) -> usize;

        /// Setter for [`count`][Self::count].
        #[unsafe(method(setCount:))]
        #[unsafe(method_family = none)]
        pub fn set_count(
            &self,
            count: usize,
        );
    );
}

/// Methods declared on superclass `NSObject`.
impl MTL4CounterHeapDescriptor {
    extern_methods!(
        #[unsafe(method(init))]
        #[unsafe(method_family = init)]
        pub fn init(this: Allocated<Self>) -> Retained<Self>;

        #[unsafe(method(new))]
        #[unsafe(method_family = new)]
        pub fn new() -> Retained<Self>;
    );
}

extern_protocol!(
    /// Represents an opaque, driver-controlled section of memory that can store GPU counter data.
    ///
    /// The data instances that this type stores correspond to the ``MTL4CounterHeapType`` heap type that you assign at creation time.
    ///
    /// See also [Apple's documentation](https://developer.apple.com/documentation/metal/mtl4counterheap?language=objc)
    pub unsafe trait MTL4CounterHeap: NSObjectProtocol {
        /// Queries the number of entries in the heap.
        #[unsafe(method(count))]
        #[unsafe(method_family = none)]
        fn count(&self) -> usize;

        /// Queries the type of the heap.
        #[unsafe(method(type))]
        #[unsafe(method_family = none)]
        fn r#type(&self) -> MTL4CounterHeapType;
    }
);

pub trait MTL4CounterHeapExt: MTL4CounterHeap + Message {
    /// Assigns a label for later inspection or visualization.
    fn label(&self) -> Option<String> {
        let s: Option<Retained<NSString>> = unsafe { msg_send![self, label] };
        s.map(|v| v.to_string())
    }

    /// Setter for [`label`][Self::label].
    fn set_label(
        &self,
        label: Option<&str>,
    ) {
        unsafe {
            let _: () = msg_send![self, setLabel: label.map(NSString::from_str).as_deref()];
        }
    }

    /// Resolves heap data on the CPU timeline.
    ///
    /// This method resolves heap data in the CPU timeline. Your app needs to ensure the GPU work has completed in order to
    /// retrieve the data correctly. You can alternatively resolve the heap data in the GPU timeline by calling
    /// ``MTL4CommandBuffer/resolveCounterHeap:withRange:intoBuffer:waitFence:updateFence:``.
    ///
    /// - Note: When resolving counters in the CPU timeline, signaling an instance of ``MTLSharedEvent`` after any workloads
    /// write counters (and waiting on that signal on the CPU) is sufficient to ensure synchronization.
    ///
    /// - Parameter range: The range in the heap to resolve.
    /// - Returns a newly allocated autoreleased NSData containing tightly packed resolved heap counter values.
    fn resolve_counter_range(
        &self,
        range: Range<usize>,
    ) -> Option<Retained<NSData>> {
        let ns_range = NSRange {
            location: range.start,
            length: range.end.saturating_sub(range.start),
        };
        unsafe { msg_send![self, resolveCounterRange: ns_range] }
    }

    /// Invalidates a range of entries in this counter heap.
    ///
    /// The effect of this call is immediate on the CPU timeline. You are responsible for ensuring that this counter heap is not currently in use on the GPU.
    ///
    /// - Note: Invalidated entries produce 0 when resolved.
    ///
    /// - Parameters:
    /// - range: A heap index range to invalidate.
    fn invalidate_counter_range(
        &self,
        range: Range<usize>,
    ) {
        let ns_range = NSRange {
            location: range.start,
            length: range.end.saturating_sub(range.start),
        };
        unsafe { msg_send![self, invalidateCounterRange: ns_range] }
    }
}

impl<T: MTL4CounterHeap + Message> MTL4CounterHeapExt for T {}