metal.rs

use objc2_metal::{
  MTLAllocation,
  MTLBuffer,
  MTLClearColor,
  MTLCompileOptions,
  MTLCreateSystemDefaultDevice,
  MTLDevice,
  MTLDrawable,
  MTLLibrary,
  MTLPixelFormat,
  MTLPrimitiveType,
  MTLRenderPipelineState,
  MTLRenderStages,
  MTLResidencySet,
  MTLResidencySetDescriptor,
  MTLResourceOptions,
  MTLSharedEvent,
  MTLViewport,
  MTL4ArgumentTable,
  MTL4ArgumentTableDescriptor,
  MTL4CommandAllocator,
  MTL4CommandBuffer,
  MTL4CommandEncoder,
  MTL4CommandQueue,
  MTL4Compiler,
  MTL4CompilerDescriptor,
  MTL4CompilerTaskOptions,
  MTL4LibraryFunctionDescriptor,
  MTL4RenderCommandEncoder,
  MTL4RenderPipelineDescriptor,
};

use crate::color::{ Color };

use objc2::{ DefinedClass, MainThreadOnly, define_class, msg_send };
use objc2::rc::{ Retained };
use objc2::runtime::{ ProtocolObject };

use objc2_app_kit::{ NSView };
use objc2_core_foundation::{ CGRect, CGSize };
use objc2_foundation::{ MainThreadMarker, NSObject, NSObjectProtocol, NSString, ns_string };
use objc2_metal_kit::{ MTKView, MTKViewDelegate };
use objc2_quartz_core::{ CAMetalLayer };

use std::cell::{ Cell, RefCell };
use std::ptr::{ NonNull };

const FLAT_2D_SHADER_SOURCE: &str = include_str!("flat_2d.metal");

#[derive(Debug)]
pub struct MetalRenderer {
  command_buffer: Retained<ProtocolObject<dyn MTL4CommandBuffer>>,
  command_queue: Retained<ProtocolObject<dyn MTL4CommandQueue>>,
  triangle_buffers: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
  argument_table: Retained<ProtocolObject<dyn MTL4ArgumentTable>>,
  command_allocators: Vec<Retained<ProtocolObject<dyn MTL4CommandAllocator>>>,
  viewport_size: CGSize,
  viewport_size_buffer: Retained<ProtocolObject<dyn MTLBuffer>>,
  render_pipeline_state: Retained<ProtocolObject<dyn MTLRenderPipelineState>>,
  shared_event: Retained<ProtocolObject<dyn MTLSharedEvent>>,
  frame_number: Cell<u64>,

  // Important to keep a Retained handle to these fields so they don't poof without warning, even if
  // we never touch them directly
  #[expect(unused)]
  device: Retained<ProtocolObject<dyn MTLDevice>>,
  #[expect(unused)]
  shader_library: Retained<ProtocolObject<dyn MTLLibrary>>,
  #[expect(unused)]
  residency_set: Retained<ProtocolObject<dyn MTLResidencySet>>,
}

impl MetalRenderer {
  const FRAMES_IN_FLIGHT: u64 = 3;

  pub fn new(view: &MTKView) -> MetalRenderer {
    let device = view.device()
      .expect("Should have been able to get the Metal device associated with the MetalKit view");

    let compile_options = MTLCompileOptions::new();
    let shader_library = device.newLibraryWithSource_options_error(ns_string!(FLAT_2D_SHADER_SOURCE), Some(&compile_options))
      .expect("Should have been able to compile the Metal shaders");

    let command_buffer = device
      .newCommandBuffer()
      .expect("Should have been able to create command buffer");

    let command_queue = device
      .newMTL4CommandQueue()
      .expect("Should have been able to create command queue");

    let triangle_buffers = Self::make_triangle_buffers(&device);
    let argument_table = Self::make_argument_table(&device);
    let residency_set = Self::make_residency_set(&device);
    let command_allocators = Self::make_command_allocators(&device);

    let viewport_size_buffer = device.newBufferWithLength_options(std::mem::size_of::<MetalUint2>(), MTLResourceOptions::StorageModeShared)
      .expect("Should have been able to create the viewport size buffer");

    let render_pipeline_state = Self::compile_render_pipeline(&device, &shader_library, view.colorPixelFormat());

    let shared_event = device.newSharedEvent()
      .expect("Should have been able to create a shared event");
    shared_event.setSignaledValue(0);

    let viewport_size_allocation: Retained<ProtocolObject<dyn MTLAllocation>> = ProtocolObject::from_retained(viewport_size_buffer.clone());
    residency_set.addAllocation(&viewport_size_allocation);

    for buffer in &triangle_buffers {
      let allocation: Retained<ProtocolObject<dyn MTLAllocation>> = ProtocolObject::from_retained(buffer.clone());
      residency_set.addAllocation(&allocation);
    }

    residency_set.commit();

    command_queue.addResidencySet(&residency_set);

    let view_layer = view
      .downcast_ref::<NSView>()
      .expect("Should have been able to cast the view into an NSView")
      .layer()
      .expect("Should have been able to get the view's layer");
    let metal_layer = view_layer
      .downcast::<CAMetalLayer>()
      .expect("Should have been able to convert the view layer into a Metal layer");
    command_queue.addResidencySet(&metal_layer.residencySet());

    let viewport_size = view.drawableSize();
    let viewport_size_data = MetalUint2 {
      x: viewport_size.width as u32,
      y: viewport_size.height as u32
    };
    unsafe {
      std::ptr::copy_nonoverlapping(
        &viewport_size_data as *const MetalUint2 as *const u8,
        viewport_size_buffer.contents().cast::<u8>().as_mut() as *mut u8,
        std::mem::size_of::<MetalUint2>());
    }

    MetalRenderer {
      command_buffer,
      command_queue,
      triangle_buffers,
      argument_table,
      command_allocators,
      viewport_size,
      viewport_size_buffer,
      render_pipeline_state,
      shared_event,
      frame_number: 0.into(),
      device,
      shader_library,
      residency_set,
    }
  }

  fn make_triangle_buffers(device: &ProtocolObject<dyn MTLDevice>) -> Vec<Retained<ProtocolObject<dyn MTLBuffer>>> {
    let mut result = Vec::new();

    for _ in 0..Self::FRAMES_IN_FLIGHT {
      let buffer = device.newBufferWithLength_options(std::mem::size_of::<VertexData>() * 3, MTLResourceOptions::StorageModeShared);
      let Some(buffer) = buffer else {
        panic!("Should be able to create a Metal buffer for the triangle data");
      };

      result.push(buffer);
    }

    result
  }

  fn make_argument_table(device: &ProtocolObject<dyn MTLDevice>) -> Retained<ProtocolObject<dyn MTL4ArgumentTable>> {
    let descriptor = MTL4ArgumentTableDescriptor::new();
    descriptor.setMaxBufferBindCount(2);

    device.newArgumentTableWithDescriptor_error(&descriptor)
      .expect("Should have been able to create a Metal argument table")
  }

  fn make_residency_set(device: &ProtocolObject<dyn MTLDevice>) -> Retained<ProtocolObject<dyn MTLResidencySet>> {
    let descriptor = MTLResidencySetDescriptor::new();
    device.newResidencySetWithDescriptor_error(&descriptor)
      .expect("Should have been able to create a Metal residency set")
  }

  fn make_command_allocators(device: &ProtocolObject<dyn MTLDevice>) -> Vec<Retained<ProtocolObject<dyn MTL4CommandAllocator>>> {
    let mut result = Vec::new();

    for _ in 0..Self::FRAMES_IN_FLIGHT {
      let allocator = device.newCommandAllocator();
      let Some(allocator) = allocator else {
        panic!("Should have been able to create a Metal command allocator");
      };

      result.push(allocator);
    }

    result
  }

  fn compile_render_pipeline(device: &ProtocolObject<dyn MTLDevice>, shader_library: &ProtocolObject<dyn MTLLibrary>, pixel_format: MTLPixelFormat) -> Retained<ProtocolObject<dyn MTLRenderPipelineState>> {
    let compiler_descriptor = MTL4CompilerDescriptor::new();
    let compiler = device.newCompilerWithDescriptor_error(&compiler_descriptor)
      .expect("Should have been able to create a Metal compiler");

    let pipeline_descriptor = Self::describe_render_pipeline(shader_library, pixel_format);
    let compiler_task_options = MTL4CompilerTaskOptions::new();

    compiler.newRenderPipelineStateWithDescriptor_compilerTaskOptions_error(&pipeline_descriptor, Some(&compiler_task_options))
      .expect("Should have been able to create a render pipeline state")
  }

  fn describe_render_pipeline(shader_library: &ProtocolObject<dyn MTLLibrary>, pixel_format: MTLPixelFormat) -> Retained<MTL4RenderPipelineDescriptor> {
    let pipeline_descriptor = MTL4RenderPipelineDescriptor::new();
    pipeline_descriptor.setLabel(Some(ns_string!("Crucible Metal 4 Render Pipeline")));

    unsafe {
      pipeline_descriptor
        .colorAttachments()
        .objectAtIndexedSubscript(0)
        .setPixelFormat(pixel_format);
    };

    let vertex_descriptor = Self::describe_vertex_function(shader_library);
    pipeline_descriptor.setVertexFunctionDescriptor(Some(&vertex_descriptor));

    let fragment_descriptor = Self::describe_fragment_function(shader_library);
    pipeline_descriptor.setFragmentFunctionDescriptor(Some(&fragment_descriptor));

    pipeline_descriptor
  }

  fn describe_vertex_function(shader_library: &ProtocolObject<dyn MTLLibrary>) -> Retained<MTL4LibraryFunctionDescriptor> {
    let descriptor = MTL4LibraryFunctionDescriptor::new();
    descriptor.setLibrary(Some(shader_library));
    descriptor.setName(Some(ns_string!("vertex_shader")));
    descriptor
  }

  fn describe_fragment_function(shader_library: &ProtocolObject<dyn MTLLibrary>) -> Retained<MTL4LibraryFunctionDescriptor> {
    let descriptor = MTL4LibraryFunctionDescriptor::new();
    descriptor.setLibrary(Some(shader_library));
    descriptor.setName(Some(ns_string!("fragment_shader")));
    descriptor
  }

  fn render_frame(&self, view: &MTKView) {
    let frame_number = self.frame_number.get();
    self.frame_number.set(frame_number + 1);

    let frame_index = (frame_number % Self::FRAMES_IN_FLIGHT) as usize;

    tracing::trace!("Rendering frame {}", frame_number);

    if frame_number >= Self::FRAMES_IN_FLIGHT {
      self.wait_for_free_frame();
    }

    let frame_allocator = &self.command_allocators[frame_index];
    frame_allocator.reset();

    let frame_label = NSString::from_str(&format!("Frame: {}", frame_number));

    self.command_buffer.beginCommandBufferWithAllocator(frame_allocator);
    self.command_buffer.setLabel(Some(&frame_label));

    let render_pass_descriptor = view.currentMTL4RenderPassDescriptor()
      .expect("Should have been able to get the current render pass descriptor from the MetalKit view");
    let render_pass_encoder = self.command_buffer.renderCommandEncoderWithDescriptor(&render_pass_descriptor)
      .expect("Should have been able to create a render command encoder");
    render_pass_encoder.setLabel(Some(&frame_label));

    render_pass_encoder.setRenderPipelineState(&self.render_pipeline_state);

    let viewport = MTLViewport {
      originX: 0.0,
      originY: 0.0,
      width: self.viewport_size.width,
      height: self.viewport_size.height,
      znear: 0.0,
      zfar: 1.0,
    };
    render_pass_encoder.setViewport(viewport);

    let vertex_buffer = &self.triangle_buffers[frame_index];
    self.make_triangle_in_buffer(vertex_buffer);
    unsafe {
      self.argument_table.setAddress_atIndex(vertex_buffer.gpuAddress(), 0);
      self.argument_table.setAddress_atIndex(self.viewport_size_buffer.gpuAddress(), 1);
    }
    render_pass_encoder.setArgumentTable_atStages(&self.argument_table, MTLRenderStages::Vertex);

    unsafe {
      render_pass_encoder.drawPrimitives_vertexStart_vertexCount(MTLPrimitiveType::Triangle, 0, 3);
    }

    render_pass_encoder.endEncoding();
    self.command_buffer.endCommandBuffer();

    let current_drawable = view.currentDrawable()
      .expect("View should have a current drawable");
    let current_drawable = ProtocolObject::from_retained(current_drawable.clone());
    self.command_queue.waitForDrawable(&current_drawable);
    let command_buffers_to_commit = NonNull::new([NonNull::from_ref(&*self.command_buffer)].as_mut_ptr())
      .expect("Nonempty slice should always be non-null");

    unsafe {
      self.command_queue.commit_count(command_buffers_to_commit, 1);
    }

    self.command_queue.signalDrawable(&current_drawable);
    current_drawable.present();

    let event = ProtocolObject::from_retained(self.shared_event.clone());
    self.command_queue.signalEvent_value(&event, frame_number);
  }

  fn wait_for_free_frame(&self) {
    loop {
      let was_signaled = self.shared_event.waitUntilSignaledValue_timeoutMS(self.frame_number.get() - Self::FRAMES_IN_FLIGHT, 10);
      if was_signaled {
        break;
      } else {
        tracing::debug!("Renderer waiting on shared event {}", self.frame_number.get() - Self::FRAMES_IN_FLIGHT);
      }
    }
  }

  fn make_triangle_in_buffer(&self, vertex_buffer: &ProtocolObject<dyn MTLBuffer>) {
    let radius = 350.0;
    let rotation = (self.frame_number.get() % 360) as f32;

    let angle_0 = rotation * std::f32::consts::PI / 180.0;
    let angle_1 = angle_0 + 2.0 * std::f32::consts::PI / 3.0;
    let angle_2 = angle_1 + 2.0 * std::f32::consts::PI / 3.0;

    let mesh = Mesh {
      data: vec![
        VertexData {
          position: MetalFloat2 {
            x: radius * angle_0.cos(),
            y: radius * angle_0.sin(),
          },
          color: MetalFloat4 {
            x: 1.0,
            y: 0.0,
            z: 0.0,
            w: 1.0,
          },
        },
        VertexData {
          position: MetalFloat2 {
            x: radius * angle_1.cos(),
            y: radius * angle_1.sin(),
          },
          color: MetalFloat4 {
            x: 0.0,
            y: 1.0,
            z: 0.0,
            w: 1.0,
          },
        },
        VertexData {
          position: MetalFloat2 {
            x: radius * angle_2.cos(),
            y: radius * angle_2.sin(),
          },
          color: MetalFloat4 {
            x: 0.0,
            y: 0.0,
            z: 1.0,
            w: 1.0,
          },
        },
      ]
    };

    unsafe {
      std::ptr::copy_nonoverlapping(
        mesh.data.as_ptr() as *const u8,
        vertex_buffer.contents().cast::<u8>().as_mut() as *mut u8,
        std::mem::size_of::<VertexData>() * mesh.data.len());
    }
  }
}

#[derive(Debug)]
pub struct MetalCanvas {
  mtk_view: Retained<MTKView>,

  // Important to keep a Retained handle to this field so it doesn't poof without warning, even
  // though we don't touch it directly
  #[expect(unused)]
  delegate: Retained<MetalCanvasDelegate>,
}

impl MetalCanvas {
  pub fn new(mtm: MainThreadMarker, frame: CGRect, clear_color: Color) -> MetalCanvas {
    let device = MTLCreateSystemDefaultDevice()
      .expect("Should have been able to create Metal device");

    let mtk_view = MTKView::initWithFrame_device(MTKView::alloc(mtm), frame, Some(&device));
    mtk_view.setClearColor(clear_color.into());

    let delegate = MetalCanvasDelegate::new(mtm, &mtk_view);
    mtk_view.setDelegate(Some(ProtocolObject::from_ref(&*delegate)));

    MetalCanvas {
      mtk_view,
      delegate,
    }
  }

  pub fn view(&self) -> &MTKView {
    &self.mtk_view
  }
}

define_class!{
  #[derive(Debug)]
  #[unsafe(super = NSObject)]
  #[thread_kind = MainThreadOnly]
  #[ivars = MetalCanvasDelegateIvars]
  struct MetalCanvasDelegate;

  unsafe impl NSObjectProtocol for MetalCanvasDelegate {}

  unsafe impl MTKViewDelegate for MetalCanvasDelegate {
    #[unsafe(method(mtkView:drawableSizeWillChange:))]
    fn mtk_view_drawable_size_will_change(&self, _view: &MTKView, _size: CGSize) {
    }

    #[unsafe(method(drawInMTKView:))]
    fn draw_in_mtk_view(&self, view: &MTKView) {
      self
        .ivars()
        .renderer
        .borrow()
        .render_frame(view);
    }
  }
}

impl MetalCanvasDelegate {
  fn new(mtm: MainThreadMarker, view: &MTKView) -> Retained<MetalCanvasDelegate> {
    let this = MetalCanvasDelegate::alloc(mtm)
      .set_ivars(MetalCanvasDelegateIvars::new(MetalRenderer::new(view)));

    unsafe {
      msg_send![super(this), init]
    }
  }
}

#[derive(Debug)]
struct MetalCanvasDelegateIvars {
  renderer: RefCell<MetalRenderer>,
}

impl MetalCanvasDelegateIvars {
  fn new(renderer: MetalRenderer) -> MetalCanvasDelegateIvars {
    MetalCanvasDelegateIvars {
      renderer: renderer.into(),
    }
  }
}

impl From<Color> for MTLClearColor {
  fn from(color: Color) -> MTLClearColor {
    MTLClearColor {
      red: color.r() as f64,
      green: color.g() as f64,
      blue: color.b() as f64,
      alpha: color.a() as f64,
    }
  }
}

#[derive(Clone, Debug)]
struct Mesh {
  data: Vec<VertexData>,
}

#[derive(Clone, Copy, Debug)]
#[repr(C)]
struct VertexData {
  position: MetalFloat2,
  color: MetalFloat4,
}

#[derive(Clone, Copy, Debug)]
#[repr(C, align(8))]
struct MetalUint2 {
  x: u32,
  y: u32,
}

#[derive(Clone, Copy, Debug)]
#[repr(C, align(8))]
struct MetalFloat2 {
  x: f32,
  y: f32,
}

#[derive(Clone, Copy, Debug)]
#[repr(C, align(16))]
struct MetalFloat4 {
  x: f32,
  y: f32,
  z: f32,
  w: f32,
}