swiftui_native/

metal.rs

//! Metal shader effects — `colorEffect`, `layerEffect` and `distortionEffect`.
//!
//! SwiftUI runs shaders written in Metal Shading Language. Its own
//! `ShaderLibrary` will only take one already compiled — as a file or as bytes
//! — and the offline compiler that produces those ships with Xcode alone, so
//! the obvious reading is that shaders cannot be built on a device.
//!
//! They can. Metal has always compiled source at run time; the missing piece
//! was a way to hand the result to SwiftUI. A `ShaderLibrary` turns out to be a
//! single word wrapping an `RBShaderLibrary`, the RenderBox object SwiftUI
//! renders through, and that class takes an `MTLLibrary` directly. That is
//! [`ShaderLibrary::from_source`], and it is the only one of these three that
//! works on a device without help from a Mac.
//!
//! ```text
//!   source ──(MTLDevice, at run time)──► MTLLibrary ─┐
//!                                                    ├──► ShaderLibrary
//!   source ──(metal compiler, on a Mac)──► .metallib ┘
//!                                                    from_path / from_data
//!        │  .function("name")
//!        ▼
//!   ShaderFunction ──.call(&[args])──► Shader ──► View::color_effect
//! ```
//!
//! [`from_source`](ShaderLibrary::from_source) reaches into SwiftUI's internals
//! and could break on any release; [`from_path`](ShaderLibrary::from_path) and
//! [`from_data`](ShaderLibrary::from_data) go through public initialisers and
//! will not. [`compile`] wraps the toolchain call for a build script that wants
//! to produce a `.metallib` ahead of time, on macOS or cross-built for a device.
//!
//! Shader entry points must be marked `[[ stitchable ]]`. A color effect takes
//! `(float2 position, half4 color, ...)` and returns `half4`.
//!
//! ```ignore
//! use swiftui_native::metal::{ShaderArg, ShaderLibrary};
//!
//! let lib = ShaderLibrary::from_source(SRC)?;   // no toolchain, works on iOS
//! let shader = lib.function("tintify").call(&[ShaderArg::Float(0.3)]);
//! let view = text("hello").color_effect(&shader);
//! ```

use crate::abi;
use crate::resolve;
use crate::views::ViewHandle;
use crate::Color;
use core::ffi::c_void;
// Only the Mac-side compiler helpers take paths.
#[cfg(target_os = "macos")]
use std::path::Path;
use crate::mangle;

mangle::modifier_types! {
    /// `_ShaderFilterEffect` — the modifier all three shader entry points add.
    /// Composed against the receiver, so the result is
    /// `ModifiedContent<Receiver, _ShaderFilterEffect>`.
    SHADER_EFFECT_T = ty("_ShaderFilterEffect");
}

/// An argument passed through to the shader function, after the arguments
/// SwiftUI supplies itself.
#[derive(Clone, Copy, Debug)]
pub enum ShaderArg {
    Float(f32),
    Float2(f32, f32),
    Color(Color),
}

/// A holder for a Swift value of a resilient type, whose size is only known at
/// runtime. Every type in the shader chain is address-only, so they are all
/// carried as raw bytes rather than as Rust structs.
struct SwiftValue {
    bytes: Vec<u8>,
}

impl SwiftValue {
    unsafe fn empty(meta: *const c_void) -> Self {
        // Stride, not size: a callee may write the trailing padding too.
        let n = unsafe { abi::type_bytes(meta) };
        Self { bytes: vec![0u8; n.max(1)] }
    }

    fn as_ptr(&self) -> *mut u8 {
        self.bytes.as_ptr() as *mut u8
    }
}

/// An `NSError`'s `localizedDescription`, or a stand-in.
unsafe fn describe_error(error: *mut c_void) -> String {
    unsafe {
        if error.is_null() {
            return "the Metal compiler failed without saying why".into();
        }
        let text = crate::objc::send(error, crate::objc::sel(c"localizedDescription"));
        if text.is_null() {
            return "the Metal compiler failed without saying why".into();
        }
        let utf8 = crate::objc::send_ret_ptr(text, crate::objc::sel(c"UTF8String"));
        if utf8.is_null() {
            return "the Metal compiler failed without saying why".into();
        }
        core::ffi::CStr::from_ptr(utf8 as *const i8).to_string_lossy().into_owned()
    }
}

/// A compiled Metal library.
pub struct ShaderLibrary(SwiftValue);

impl ShaderLibrary {
    /// The underlying `SwiftUI.ShaderLibrary` value, for diagnostics.
    #[doc(hidden)]
    pub fn value_ptr(&self) -> *mut c_void {
        self.0.as_ptr() as *mut c_void
    }
}

/// One `[[ stitchable ]]` function inside a library.
pub struct ShaderFunction(SwiftValue);

/// A shader function bound to its arguments, ready to apply to a view.
pub struct Shader(SwiftValue);

impl Shader {
    fn as_ptr(&self) -> *mut u8 {
        self.0.as_ptr()
    }
}

impl ShaderLibrary {
    /// Load a `.metallib` from disk.
    pub fn from_path(path: &str) -> Self {
        unsafe {
            // Foundation.URL is resilient, so it is built into a buffer and
            // handed onward by pointer rather than in registers.
            let url_meta = abi::metadata0(resolve::url_ma());
            let url = SwiftValue::empty(url_meta);
            let ss = abi::swift_string(path);
            let s0 = u64::from_le_bytes(ss[..8].try_into().unwrap());
            let s1 = u64::from_le_bytes(ss[8..].try_into().unwrap());
            core::arch::asm!("blr {f}", f = in(reg) resolve::url_file_path_init(),
                in("x0") s0, in("x1") s1,
                in("x8") url.as_ptr(),
                in("x20") url_meta,
                lateout("x2") _, lateout("x3") _, lateout("x4") _, lateout("x5") _,
                lateout("x6") _, lateout("x7") _, lateout("x9") _, lateout("x10") _,
                lateout("x11") _, lateout("x12") _, lateout("x13") _, lateout("x14") _,
                lateout("x15") _, lateout("x16") _, lateout("x17") _,
                lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );

            let lib_meta = abi::metadata0(resolve::shader_library_ma());
            let lib = SwiftValue::empty(lib_meta);
            // ShaderLibrary.init(url:) — x0 = @in URL, x8 = @out ShaderLibrary
            core::arch::asm!("blr {f}", f = in(reg) resolve::shader_library_url_init(),
                in("x0") url.as_ptr(),
                in("x8") lib.as_ptr(),
                in("x20") lib_meta,
                lateout("x1") _, lateout("x2") _, lateout("x3") _, lateout("x4") _,
                lateout("x5") _, lateout("x6") _, lateout("x7") _, lateout("x9") _,
                lateout("x10") _, lateout("x11") _, lateout("x12") _, lateout("x13") _,
                lateout("x14") _, lateout("x15") _, lateout("x16") _, lateout("x17") _,
                lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );
            Self(lib)
        }
    }

    /// Compile Metal Shading Language source at run time.
    ///
    /// No `.metallib`, no Xcode, and it works on a device — which is what
    /// [`compile`] cannot do, since the offline Metal compiler ships with Xcode
    /// alone. Metal itself has always been able to compile at run time; what was
    /// missing was a way to hand the result to SwiftUI, whose `ShaderLibrary`
    /// takes only an already-compiled library's bytes.
    ///
    /// There is one, as it turns out. A `ShaderLibrary` is a single word
    /// wrapping an `RBShaderLibrary` — the RenderBox object SwiftUI actually
    /// renders through — and that class accepts an `MTLLibrary` directly. So
    /// the source goes through `MTLDevice`, which reports errors properly, and
    /// the resulting library is wrapped.
    ///
    /// `RBShaderLibrary` is internal to SwiftUI: it is not API, and a future
    /// release may rename it or change its shape. [`from_data`](Self::from_data)
    /// and [`from_path`](Self::from_path) go through SwiftUI's public
    /// initialisers and will not.
    ///
    /// # Errors
    /// The compiler's own diagnostics, or a note that the runtime pieces are
    /// missing. Worth propagating: a library built from source that does not
    /// compile still *renders*, as a shader that silently does nothing.
    pub fn from_source(source: &str) -> Result<Self, String> {
        unsafe {
            let _pool = crate::objc::AutoreleasePool::new();
            unsafe extern "C" {
                fn MTLCreateSystemDefaultDevice() -> *mut c_void;
            }
            let device = MTLCreateSystemDefaultDevice();
            if device.is_null() {
                return Err("no Metal device".into());
            }

            // -newLibraryWithSource:options:error:
            type NewLibrary = unsafe extern "C" fn(
                *mut c_void, *const c_void, *mut c_void, *mut c_void, *mut *mut c_void,
            ) -> *mut c_void;
            let new_library: NewLibrary =
                core::mem::transmute(crate::objc::msg_send_addr());
            let mut error: *mut c_void = core::ptr::null_mut();
            let library = new_library(
                device,
                crate::objc::sel(c"newLibraryWithSource:options:error:"),
                crate::objc::nsstring(source),
                core::ptr::null_mut(),
                &mut error,
            );
            if library.is_null() {
                return Err(describe_error(error));
            }

            let class = crate::objc::class_named(c"RBShaderLibrary");
            if class.is_null() {
                return Err("SwiftUI's RBShaderLibrary is not available".into());
            }
            let rb = crate::objc::send_id(
                class,
                crate::objc::sel(c"libraryWithMetalLibrary:"),
                library,
            );
            if rb.is_null() {
                return Err("RBShaderLibrary rejected the compiled library".into());
            }
            // The convenience constructor hands back an autoreleased object;
            // the value being built has to own one. The `MTLLibrary` came from
            // a `new...` and is owned here, and `RBShaderLibrary` has taken its
            // own reference by now.
            crate::objc::send(rb, crate::objc::sel(c"retain"));
            crate::objc::send(library, crate::objc::sel(c"release"));

            let meta = abi::metadata0(resolve::shader_library_ma());
            let value = SwiftValue::empty(meta);
            debug_assert_eq!(abi::type_size(meta), 8, "ShaderLibrary is one word");
            *(value.as_ptr() as *mut u64) = rb as u64;
            Ok(Self(value))
        }
    }

    /// Load a `.metallib` from bytes already in memory.
    ///
    /// A `.metallib` has to be built ahead of time — see [`compile_for`] — but
    /// it does not have to be a *file* at run time. `include_bytes!` it into the
    /// binary and this needs no bundle resource and no path, which on iOS is
    /// the difference between shipping a shader and shipping an app layout.
    ///
    /// `Data` is bridged from an `NSData` rather than built directly:
    /// `Data.init(bytes:count:)` is inlinable and so has no symbol to call,
    /// while `init(referencing:)` is a real entry point.
    pub fn from_data(bytes: &[u8]) -> Self {
        unsafe {
            let _pool = crate::objc::AutoreleasePool::new();
            // `alloc`/`init` rather than the autoreleased convenience
            // constructor: `Data` holds onto this reference and there is no
            // pool here whose drain it would survive.
            type InitBytes =
                unsafe extern "C" fn(*mut c_void, *const c_void, *const u8, u64) -> *mut c_void;
            let init: InitBytes = core::mem::transmute(crate::objc::msg_send_addr());
            let nsdata = init(
                crate::objc::send(
                    crate::objc::class(c"OBJC_CLASS_$_NSData"),
                    crate::objc::sel(c"alloc"),
                ),
                crate::objc::sel(c"initWithBytes:length:"),
                bytes.as_ptr(),
                bytes.len() as u64,
            );
            assert!(!nsdata.is_null(), "NSData creation failed");

            // `Data` is sixteen bytes and *loadable*, unlike the resilient
            // `URL` next door: it comes back in x0/x1 rather than through an
            // indirect result, and goes on to `ShaderLibrary` the same way.
            // Reading x8 instead yields an untouched buffer — an empty `Data`,
            // which produces a library that loads and silently shades nothing.
            let data_meta = abi::metadata0(resolve::data_ma());
            let (d0, d1): (u64, u64);
            core::arch::asm!("blr {f}", f = in(reg) resolve::data_referencing_init(),
                in("x0") nsdata,
                in("x20") data_meta,
                lateout("x0") d0, lateout("x1") d1,
                lateout("x2") _, lateout("x3") _, lateout("x4") _, lateout("x5") _,
                lateout("x6") _, lateout("x7") _, lateout("x8") _, lateout("x9") _,
                lateout("x10") _, lateout("x11") _, lateout("x12") _, lateout("x13") _,
                lateout("x14") _, lateout("x15") _, lateout("x16") _, lateout("x17") _,
                lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );

            let lib_meta = abi::metadata0(resolve::shader_library_ma());
            let lib = SwiftValue::empty(lib_meta);
            // ShaderLibrary.init(data:) — x0/x1 = Data, x8 = @out ShaderLibrary
            core::arch::asm!("blr {f}", f = in(reg) resolve::shader_library_data_init(),
                in("x0") d0,
                in("x1") d1,
                in("x8") lib.as_ptr(),
                in("x20") lib_meta,
                lateout("x2") _, lateout("x3") _, lateout("x4") _, lateout("x5") _,
                lateout("x6") _, lateout("x7") _, lateout("x9") _,
                lateout("x10") _, lateout("x11") _, lateout("x12") _, lateout("x13") _,
                lateout("x14") _, lateout("x15") _, lateout("x16") _, lateout("x17") _,
                lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );
            Self(lib)
        }
    }

    /// Look up a `[[ stitchable ]]` function by name.
    ///
    /// In Swift this is `library.myFunction`, spelled through
    /// `@dynamicMemberLookup`; the underlying subscript takes the name as an
    /// ordinary string.
    pub fn function(&self, name: &str) -> ShaderFunction {
        unsafe {
            let fn_meta = abi::metadata0(resolve::shader_function_ma());
            let out = SwiftValue::empty(fn_meta);
            let ss = abi::swift_string(name);
            let s0 = u64::from_le_bytes(ss[..8].try_into().unwrap());
            let s1 = u64::from_le_bytes(ss[8..].try_into().unwrap());
            // subscript(dynamicMember:).getter
            //   x0/x1 = name, x8 = @out ShaderFunction, x20 = @in_guaranteed self
            core::arch::asm!("blr {f}", f = in(reg) resolve::shader_library_subscript(),
                in("x0") s0, in("x1") s1,
                in("x8") out.as_ptr(),
                in("x20") self.0.as_ptr(),
                lateout("x2") _, lateout("x3") _, lateout("x4") _, lateout("x5") _,
                lateout("x6") _, lateout("x7") _, lateout("x9") _, lateout("x10") _,
                lateout("x11") _, lateout("x12") _, lateout("x13") _, lateout("x14") _,
                lateout("x15") _, lateout("x16") _, lateout("x17") _,
                lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );
            ShaderFunction(out)
        }
    }
}

impl ShaderFunction {
    /// Bind arguments, producing a `Shader`.
    pub fn call(&self, args: &[ShaderArg]) -> Shader {
        unsafe {
            let arg_meta = abi::metadata0(resolve::shader_argument_ma());
            let stride = abi::type_stride(arg_meta);

            // Each argument is built by its own static factory, then packed into
            // a Swift Array at the element stride.
            let mut packed = vec![0u8; stride * args.len()];
            for (i, a) in args.iter().enumerate() {
                let value = build_argument(*a, arg_meta);
                packed[i * stride..i * stride + value.bytes.len()]
                    .copy_from_slice(&value.bytes);
            }
            let array = abi::swift_array(arg_meta, args.len(), &packed);

            let shader_meta = abi::metadata0(resolve::shader_ma());
            let out = SwiftValue::empty(shader_meta);
            // dynamicallyCall(withArguments:) — x0 = Array, x8 = @out Shader,
            // x20 = @in_guaranteed self
            core::arch::asm!("blr {f}", f = in(reg) resolve::shader_fn_call(),
                in("x0") array,
                in("x8") out.as_ptr(),
                in("x20") self.0.as_ptr(),
                lateout("x1") _, lateout("x2") _, lateout("x3") _, lateout("x4") _,
                lateout("x5") _, lateout("x6") _, lateout("x7") _, lateout("x9") _,
                lateout("x10") _, lateout("x11") _, lateout("x12") _, lateout("x13") _,
                lateout("x14") _, lateout("x15") _, lateout("x16") _, lateout("x17") _,
                lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );
            Shader(out)
        }
    }
}

/// Build one `Shader.Argument` through its static factory.
unsafe fn build_argument(arg: ShaderArg, meta: *const c_void) -> SwiftValue {
    unsafe {
        let out = SwiftValue::empty(meta);
        match arg {
            // Swift `Float` is 32-bit and travels in s0/s1.
            ShaderArg::Float(v) => {
                core::arch::asm!("blr {f}", f = in(reg) resolve::shader_arg_float(),
                    in("s0") v,
                    in("x8") out.as_ptr(),
                    in("x20") meta,
                    lateout("x0") _, lateout("x1") _, lateout("x2") _, lateout("x3") _,
                    lateout("x4") _, lateout("x5") _, lateout("x6") _, lateout("x7") _,
                    lateout("x9") _, lateout("x10") _, lateout("x11") _, lateout("x12") _,
                    lateout("x13") _, lateout("x14") _, lateout("x15") _, lateout("x16") _,
                    lateout("x17") _, lateout("lr") _, lateout("x21") _, clobber_abi("C"),
                );
            }
            ShaderArg::Float2(x, y) => {
                core::arch::asm!("blr {f}", f = in(reg) resolve::shader_arg_float2(),
                    in("s0") x, in("s1") y,
                    in("x8") out.as_ptr(),
                    in("x20") meta,
                    lateout("x0") _, lateout("x1") _, lateout("x2") _, lateout("x3") _,
                    lateout("x4") _, lateout("x5") _, lateout("x6") _, lateout("x7") _,
                    lateout("x9") _, lateout("x10") _, lateout("x11") _, lateout("x12") _,
                    lateout("x13") _, lateout("x14") _, lateout("x15") _, lateout("x16") _,
                    lateout("x17") _, lateout("lr") _, lateout("x21") _, clobber_abi("C"),
                );
            }
            ShaderArg::Color(c) => {
                let value = crate::views::color_value(c.r, c.g, c.b, c.a);
                core::arch::asm!("blr {f}", f = in(reg) resolve::shader_arg_color(),
                    in("x0") value,
                    in("x8") out.as_ptr(),
                    in("x20") meta,
                    lateout("x1") _, lateout("x2") _, lateout("x3") _,
                    lateout("x4") _, lateout("x5") _, lateout("x6") _, lateout("x7") _,
                    lateout("x9") _, lateout("x10") _, lateout("x11") _, lateout("x12") _,
                    lateout("x13") _, lateout("x14") _, lateout("x15") _, lateout("x16") _,
                    lateout("x17") _, lateout("lr") _, lateout("x21") _, clobber_abi("C"),
                );
            }
        }
        out
    }
}

/// Which shader modifier to apply.
enum Kind {
    Color,
    Layer,
    Distortion,
}

unsafe fn apply(view: &ViewHandle, shader: &Shader, kind: Kind, max: (f64, f64)) -> ViewHandle {
    unsafe {
        let (func, has_offset) = match kind {
            Kind::Color => (resolve::color_effect_fn(), false),
            Kind::Layer => (resolve::layer_effect_fn(), true),
            Kind::Distortion => (resolve::distortion_effect_fn(), true),
        };
        let (meta, wt, size) = SHADER_EFFECT_T.applied_to(view.meta);
        let mut result = vec![0u8; size];
        let this = view.value_ptr();

        // colorEffect(_:isEnabled:)          x0 = @in Shader, x1 = isEnabled
        // layerEffect(_:maxSampleOffset:isEnabled:) adds a CGSize in d0/d1
        //   then Self.Type and the Self: View witness table, x8 = @out, x20 = self
        if has_offset {
            core::arch::asm!("blr {f}", f = in(reg) func,
                in("x0") shader.as_ptr(),
                in("d0") max.0, in("d1") max.1,
                in("x1") 1u64,
                in("x2") view.meta,
                in("x3") view.wt,
                in("x8") result.as_mut_ptr(),
                in("x20") this,
                lateout("x4") _, lateout("x5") _, lateout("x6") _, lateout("x7") _,
                lateout("x9") _, lateout("x10") _, lateout("x11") _, lateout("x12") _,
                lateout("x13") _, lateout("x14") _, lateout("x15") _, lateout("x16") _,
                lateout("x17") _, lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );
        } else {
            core::arch::asm!("blr {f}", f = in(reg) func,
                in("x0") shader.as_ptr(),
                in("x1") 1u64,
                in("x2") view.meta,
                in("x3") view.wt,
                in("x8") result.as_mut_ptr(),
                in("x20") this,
                lateout("x4") _, lateout("x5") _, lateout("x6") _, lateout("x7") _,
                lateout("x9") _, lateout("x10") _, lateout("x11") _, lateout("x12") _,
                lateout("x13") _, lateout("x14") _, lateout("x15") _, lateout("x16") _,
                lateout("x17") _, lateout("lr") _, lateout("x21") _, clobber_abi("C"),
            );
        }
        crate::views::wrap_owned(result, meta, wt)
    }
}

/// Apply `.colorEffect(_:)` — recolours each pixel.
pub fn color_effect(view: &ViewHandle, shader: &Shader) -> ViewHandle {
    unsafe { apply(view, shader, Kind::Color, (0.0, 0.0)) }
}

/// Apply `.layerEffect(_:maxSampleOffset:)` — can sample neighbouring pixels.
pub fn layer_effect(view: &ViewHandle, shader: &Shader, max_w: f64, max_h: f64) -> ViewHandle {
    unsafe { apply(view, shader, Kind::Layer, (max_w, max_h)) }
}

/// Apply `.distortionEffect(_:maxSampleOffset:)` — warps geometry.
pub fn distortion_effect(view: &ViewHandle, shader: &Shader, max_w: f64, max_h: f64) -> ViewHandle {
    unsafe { apply(view, shader, Kind::Distortion, (max_w, max_h)) }
}

/// Compile Metal Shading Language source into a `.metallib` for the Mac.
///
/// Only on the Mac, and deliberately: the offline Metal compiler ships with
/// Xcode and nowhere else, so there is nothing for this to call on a device.
/// Anything shipping to iOS builds its library ahead of time with
/// [`compile_for`] and loads it with [`ShaderLibrary::from_path`] or
/// [`ShaderLibrary::from_data`] — both of which work on every platform.
///
/// Metal *can* compile source at runtime on iOS, through
/// `MTLDevice.makeLibrary(source:options:)`. It is no help here: that hands
/// back an `MTLLibrary` object, and SwiftUI's `ShaderLibrary` takes a compiled
/// library's *bytes*. There is no public way from one to the other.
#[cfg(target_os = "macos")]
pub fn compile(source: &str, out_path: impl AsRef<Path>) -> std::io::Result<()> {
    compile_for("macosx", source, out_path)
}

/// Compile for a named SDK — `macosx`, `iphonesimulator`, `iphoneos`.
///
/// A `.metallib` is per-platform, so building one for a device or a simulator
/// has to name the SDK explicitly. This runs on the Mac whichever target it is
/// producing for.
#[cfg(target_os = "macos")]
pub fn compile_for(sdk: &str, source: &str, out_path: impl AsRef<Path>) -> std::io::Result<()> {
    use std::process::Command;
    let out = out_path.as_ref();
    let src = out.with_extension("metal");
    std::fs::write(&src, source)?;
    let status = Command::new("xcrun")
        .args(["-sdk", sdk, "metal", "-o"])
        .arg(out)
        .arg(&src)
        .output()?;
    if !status.status.success() {
        return Err(std::io::Error::other(String::from_utf8_lossy(&status.stderr).to_string()));
    }
    Ok(())
}