initial commit

3 files changed, 1185 insertions, 0 deletions

diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..a024ad3
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,408 @@
+use std::sync::Arc;
+use winit::{
+    application::ApplicationHandler,
+    event::WindowEvent,
+    event_loop::EventLoop,
+    keyboard::{KeyCode, PhysicalKey},
+    window::Window,
+};
+
+mod osc;
+mod uniform;
+
+const VERTEX_SHADER: &str = "\
+@group(0) @binding(0) var<uniform> _wgsl_resolution: vec2<f32>;
+
+struct VertexOutput {
+    @builtin(position) clip_position: vec4<f32>,
+    @location(0) uv: vec2<f32>,
+    @location(1) resolution: vec2<f32>,
+};
+
+@vertex
+fn vs_main(@builtin(vertex_index) vi: u32) -> VertexOutput {
+    var uv = vec2<f32>(f32(vi % 2u), f32(vi / 2u));
+    var out: VertexOutput;
+    out.clip_position = vec4<f32>(2.0 * uv - 1.0, 0.0, 1.0);
+    out.uv = vec2<f32>(uv.x, 1.0 - uv.y);
+    out.resolution = _wgsl_resolution;
+    return out;
+}
+";
+
+const DEFAULT_FRAGMENT: &str = "\
+@fragment
+fn fs_main(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
+    let check = floor(uv * 10.0);
+    return vec4<f32>(vec3<f32>(fract((check.x + check.y) / 2.0)), 1.0);
+}
+";
+
+struct GpuState {
+    window: Arc<Window>,
+    surface: wgpu::Surface<'static>,
+    device: wgpu::Device,
+    queue: wgpu::Queue,
+    config: wgpu::SurfaceConfiguration,
+    surface_format: wgpu::TextureFormat,
+    vertex_module: wgpu::ShaderModule,
+    vertex_bind_group_layout: wgpu::BindGroupLayout,
+    vertex_bind_group: wgpu::BindGroup,
+    resolution_buffer: wgpu::Buffer,
+    render_pipeline: wgpu::RenderPipeline,
+    uniform_cache: uniform::UniformCache,
+}
+
+struct App {
+    gpu: Option<GpuState>,
+    osc_server: osc::OscServer,
+}
+
+fn load_shader(
+    device: &wgpu::Device,
+    surface_format: wgpu::TextureFormat,
+    vertex_module: &wgpu::ShaderModule,
+    vertex_bgl: &wgpu::BindGroupLayout,
+    fragment_source: &str,
+    cache: &mut uniform::UniformCache,
+) -> Result<wgpu::RenderPipeline, String> {
+    let module = naga::front::wgsl::parse_str(fragment_source)
+        .map_err(|e| format!("WGSL parse error: {}", e))?;
+
+    naga::valid::Validator::new(
+        naga::valid::ValidationFlags::all(),
+        naga::valid::Capabilities::all(),
+    )
+    .validate(&module)
+    .map_err(|e| format!("WGSL validation error: {}", e))?;
+
+    cache.refresh(module, device);
+
+    let fragment_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+        label: Some("fragment_shader"),
+        source: wgpu::ShaderSource::Wgsl(fragment_source.into()),
+    });
+
+    let mut bind_group_layouts: Vec<Option<&wgpu::BindGroupLayout>> = vec![Some(vertex_bgl)];
+    if let Some(ref bgl) = cache.bind_group_layout {
+        bind_group_layouts.push(Some(bgl));
+    }
+
+    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+        label: Some("pipeline_layout"),
+        bind_group_layouts: &bind_group_layouts,
+        immediate_size: 0,
+    });
+
+    let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+        label: Some("render_pipeline"),
+        layout: Some(&pipeline_layout),
+        vertex: wgpu::VertexState {
+            module: vertex_module,
+            entry_point: Some("vs_main"),
+            buffers: &[],
+            compilation_options: wgpu::PipelineCompilationOptions::default(),
+        },
+        fragment: Some(wgpu::FragmentState {
+            module: &fragment_module,
+            entry_point: Some("fs_main"),
+            targets: &[Some(wgpu::ColorTargetState {
+                format: surface_format,
+                blend: Some(wgpu::BlendState::REPLACE),
+                write_mask: wgpu::ColorWrites::ALL,
+            })],
+            compilation_options: wgpu::PipelineCompilationOptions::default(),
+        }),
+        primitive: wgpu::PrimitiveState {
+            topology: wgpu::PrimitiveTopology::TriangleStrip,
+            strip_index_format: None,
+            front_face: wgpu::FrontFace::Ccw,
+            cull_mode: None,
+            polygon_mode: wgpu::PolygonMode::Fill,
+            unclipped_depth: false,
+            conservative: false,
+        },
+        depth_stencil: None,
+        multisample: wgpu::MultisampleState {
+            count: 1,
+            mask: !0,
+            alpha_to_coverage_enabled: false,
+        },
+        multiview_mask: None,
+        cache: None,
+    });
+
+    Ok(pipeline)
+}
+
+impl App {
+    fn new() -> Self {
+        let osc_server = osc::OscServer::new("0.0.0.0:9000").expect("bind OSC socket");
+        Self {
+            gpu: None,
+            osc_server,
+        }
+    }
+
+    fn init_gpu(&mut self, window: Arc<Window>) {
+        let gpu = pollster::block_on(async {
+            let size = window.inner_size();
+            let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
+                backends: wgpu::Backends::VULKAN,
+                flags: wgpu::InstanceFlags::default(),
+                memory_budget_thresholds: wgpu::MemoryBudgetThresholds::default(),
+                backend_options: wgpu::BackendOptions::default(),
+                display: None,
+            });
+
+            let surface = instance
+                .create_surface(window.clone())
+                .expect("create surface");
+
+            let adapter = instance
+                .request_adapter(&wgpu::RequestAdapterOptions {
+                    power_preference: wgpu::PowerPreference::default(),
+                    compatible_surface: Some(&surface),
+                    force_fallback_adapter: false,
+                })
+                .await
+                .expect("find adapter");
+
+            let (device, queue) = adapter
+                .request_device(&wgpu::DeviceDescriptor::default())
+                .await
+                .expect("create device");
+
+            let surface_caps = surface.get_capabilities(&adapter);
+            let surface_format = surface_caps
+                .formats
+                .iter()
+                .find(|f| f.is_srgb())
+                .copied()
+                .unwrap_or(surface_caps.formats[0]);
+
+            let config = wgpu::SurfaceConfiguration {
+                usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
+                format: surface_format,
+                width: size.width.max(1),
+                height: size.height.max(1),
+                present_mode: wgpu::PresentMode::Fifo,
+                alpha_mode: surface_caps.alpha_modes[0],
+                view_formats: vec![],
+                desired_maximum_frame_latency: 2,
+            };
+            surface.configure(&device, &config);
+
+            let mut cache = uniform::UniformCache::new();
+            let vertex_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+                label: Some("vertex_shader"),
+                source: wgpu::ShaderSource::Wgsl(VERTEX_SHADER.into()),
+            });
+
+            let resolution_buffer = device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some("resolution_buffer"),
+                size: 8,
+                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+                mapped_at_creation: false,
+            });
+
+            let vertex_bind_group_layout =
+                device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                    label: Some("vertex_bind_group_layout"),
+                    entries: &[wgpu::BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: wgpu::ShaderStages::VERTEX,
+                        ty: wgpu::BindingType::Buffer {
+                            ty: wgpu::BufferBindingType::Uniform,
+                            has_dynamic_offset: false,
+                            min_binding_size: None,
+                        },
+                        count: None,
+                    }],
+                });
+
+            let vertex_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+                label: Some("vertex_bind_group"),
+                layout: &vertex_bind_group_layout,
+                entries: &[wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: resolution_buffer.as_entire_binding(),
+                }],
+            });
+
+            let pipeline = load_shader(
+                &device,
+                surface_format,
+                &vertex_module,
+                &vertex_bind_group_layout,
+                DEFAULT_FRAGMENT,
+                &mut cache,
+            )
+            .expect("default shader should compile");
+
+            GpuState {
+                window,
+                surface,
+                device,
+                queue,
+                config,
+                surface_format,
+                vertex_module,
+                vertex_bind_group_layout,
+                vertex_bind_group,
+                resolution_buffer,
+                render_pipeline: pipeline,
+                uniform_cache: cache,
+            }
+        });
+
+        self.gpu = Some(gpu);
+    }
+
+    fn handle_osc(&mut self) {
+        let gpu = self.gpu.as_mut().expect("gpu initialized");
+
+        let commands = self.osc_server.poll(&mut gpu.uniform_cache);
+        for cmd in commands {
+            match cmd {
+                osc::OscCommand::Shader(code) => {
+                    match load_shader(
+                        &gpu.device,
+                        gpu.surface_format,
+                        &gpu.vertex_module,
+                        &gpu.vertex_bind_group_layout,
+                        &code,
+                        &mut gpu.uniform_cache,
+                    ) {
+                        Ok(pipeline) => {
+                            gpu.render_pipeline = pipeline;
+                            log::info!("shader loaded successfully");
+                        }
+                        Err(e) => log::error!("{}", e),
+                    }
+                }
+            }
+        }
+    }
+
+    fn render(&mut self) {
+        let gpu = self.gpu.as_mut().expect("gpu initialized");
+
+        gpu.uniform_cache.flush(&gpu.queue);
+
+        let output = match gpu.surface.get_current_texture() {
+            wgpu::CurrentSurfaceTexture::Success(t)
+            | wgpu::CurrentSurfaceTexture::Suboptimal(t) => t,
+            wgpu::CurrentSurfaceTexture::Outdated | wgpu::CurrentSurfaceTexture::Lost => {
+                gpu.surface.configure(&gpu.device, &gpu.config);
+                return;
+            }
+            other => {
+                log::error!("surface error: {:?}", other);
+                return;
+            }
+        };
+
+        let view = output
+            .texture
+            .create_view(&wgpu::TextureViewDescriptor::default());
+
+        let mut encoder = gpu
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
+                label: Some("render_encoder"),
+            });
+
+        {
+            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                label: Some("render_pass"),
+                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                    view: &view,
+                    depth_slice: None,
+                    resolve_target: None,
+                    ops: wgpu::Operations {
+                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
+                        store: wgpu::StoreOp::Store,
+                    },
+                })],
+                depth_stencil_attachment: None,
+                occlusion_query_set: None,
+                timestamp_writes: None,
+                multiview_mask: None,
+            });
+
+            pass.set_pipeline(&gpu.render_pipeline);
+            let resolution = [gpu.config.width as f32, gpu.config.height as f32];
+            gpu.queue
+                .write_buffer(&gpu.resolution_buffer, 0, bytemuck::cast_slice(&resolution));
+            pass.set_bind_group(0, &gpu.vertex_bind_group, &[]);
+            if let Some(ref bg) = gpu.uniform_cache.bind_group {
+                pass.set_bind_group(1, bg, &[]);
+            }
+            pass.draw(0..4, 0..1);
+        }
+
+        gpu.queue.submit(std::iter::once(encoder.finish()));
+        output.present();
+    }
+}
+
+impl ApplicationHandler for App {
+    fn resumed(&mut self, event_loop: &winit::event_loop::ActiveEventLoop) {
+        if self.gpu.is_none() {
+            let window = Arc::new(
+                event_loop
+                    .create_window(Window::default_attributes().with_title("wgsl-view"))
+                    .expect("create window"),
+            );
+            self.init_gpu(window);
+        }
+    }
+
+    fn window_event(
+        &mut self,
+        event_loop: &winit::event_loop::ActiveEventLoop,
+        _window_id: winit::window::WindowId,
+        event: WindowEvent,
+    ) {
+        match event {
+            WindowEvent::CloseRequested => event_loop.exit(),
+            WindowEvent::KeyboardInput {
+                event:
+                    winit::event::KeyEvent {
+                        physical_key: PhysicalKey::Code(KeyCode::Escape),
+                        state: winit::event::ElementState::Pressed,
+                        ..
+                    },
+                ..
+            } => event_loop.exit(),
+            WindowEvent::Resized(size) => {
+                let gpu = self.gpu.as_mut().expect("gpu initialized");
+                gpu.config.width = size.width.max(1);
+                gpu.config.height = size.height.max(1);
+                gpu.surface.configure(&gpu.device, &gpu.config);
+            }
+            WindowEvent::RedrawRequested => {
+                self.handle_osc();
+                self.render();
+            }
+            _ => {}
+        }
+    }
+
+    fn about_to_wait(&mut self, _event_loop: &winit::event_loop::ActiveEventLoop) {
+        self.gpu
+            .as_ref()
+            .expect("gpu initialized")
+            .window
+            .request_redraw();
+    }
+}
+
+fn main() {
+    env_logger::init();
+    let event_loop = EventLoop::new().expect("create event loop");
+    event_loop.set_control_flow(winit::event_loop::ControlFlow::Poll);
+    let mut app = App::new();
+    event_loop.run_app(&mut app).expect("run event loop");
+}
diff --git a/src/osc.rs b/src/osc.rs
new file mode 100644
index 0000000..73f0f35
--- /dev/null
+++ b/src/osc.rs
@@ -0,0 +1,161 @@
+use std::net::UdpSocket;
+
+use rosc::{OscMessage, OscPacket, OscType};
+
+use crate::uniform::{UniformCache, UniformError, UniformRef};
+
+pub struct OscServer {
+    socket: UdpSocket,
+    buf: Vec<u8>,
+}
+
+pub enum OscCommand {
+    Shader(String),
+}
+
+impl OscServer {
+    pub fn new(addr: &str) -> std::io::Result<Self> {
+        let socket = UdpSocket::bind(addr)?;
+        socket.set_nonblocking(true)?;
+        log::info!("listening for OSC on {}", addr);
+        Ok(Self {
+            socket,
+            buf: vec![0u8; 0x10000],
+        })
+    }
+
+    pub fn poll(&mut self, cache: &mut UniformCache) -> Vec<OscCommand> {
+        let mut commands = Vec::new();
+        loop {
+            match self.socket.recv_from(&mut self.buf) {
+                Ok((size, _addr)) => {
+                    let data = &self.buf[..size];
+                    match rosc::decoder::decode_udp(data) {
+                        Ok((_, packet)) => {
+                            self.handle_packet(packet, cache, &mut commands);
+                        }
+                        Err(e) => log::warn!("OSC decode error: {}", e),
+                    }
+                }
+                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
+                Err(e) => {
+                    log::warn!("OSC recv error: {}", e);
+                    break;
+                }
+            }
+        }
+        commands
+    }
+
+    fn handle_packet(
+        &self,
+        packet: OscPacket,
+        cache: &mut UniformCache,
+        commands: &mut Vec<OscCommand>,
+    ) {
+        match packet {
+            OscPacket::Message(msg) => self.handle_message(msg, cache, commands),
+            OscPacket::Bundle(bundle) => {
+                for p in bundle.content {
+                    self.handle_packet(p, cache, commands);
+                }
+            }
+        }
+    }
+
+    fn handle_message(
+        &self,
+        msg: OscMessage,
+        cache: &mut UniformCache,
+        commands: &mut Vec<OscCommand>,
+    ) {
+        let path = &msg.addr;
+
+        if path == "/shader" {
+            match &msg.args[..] {
+                [OscType::String(code)] => commands.push(OscCommand::Shader(code.clone())),
+                _ => log::warn!("/shader: unexpected arguments"),
+            }
+        } else if let Some(rest) = path.strip_prefix("/uniform/") {
+            if let Err(e) = handle_uniform(rest, &msg.args, cache) {
+                log::warn!("error handling {}: {}", path, e);
+            }
+        } else {
+            log::debug!("unhandled OSC message: {}", path);
+        }
+    }
+}
+
+fn handle_uniform(
+    path: &str,
+    args: &[OscType],
+    cache: &mut UniformCache,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let mut parts = path.split('/');
+    let name = parts.next().ok_or("missing uniform name")?;
+
+    let mut uref = cache.get(name).ok_or("uniform not found")?;
+
+    for component in parts {
+        let i = match component {
+            "x" | "r" | "s" => 0,
+            "y" | "g" | "t" => 1,
+            "z" | "b" | "p" => 2,
+            "w" | "a" | "q" => 3,
+            s => s.parse::<usize>()?,
+        };
+        uref = uref.index(i)?;
+    }
+
+    write_osc_values(uref, args)?;
+    Ok(())
+}
+
+/// Write OSC argument values into a UniformRef, converting OSC types
+/// to match the uniform's expected scalar type.
+fn write_osc_values(uref: UniformRef<'_>, args: &[OscType]) -> Result<(), UniformError> {
+    if args.is_empty() {
+        return Err(UniformError::SizeMismatch);
+    }
+
+    let scalar = uref.leaf_scalar()?;
+    match scalar.kind {
+        naga::ScalarKind::Float => {
+            let values: Result<Vec<f32>, _> = args
+                .iter()
+                .map(|a| match a {
+                    OscType::Float(f) => Ok(*f),
+                    OscType::Double(d) => Ok(*d as f32),
+                    OscType::Int(i) => Ok(*i as f32),
+                    _ => Err(UniformError::TypeMismatch),
+                })
+                .collect();
+            uref.set_f32(&values?)
+        }
+        naga::ScalarKind::Sint => {
+            let values: Result<Vec<i32>, _> = args
+                .iter()
+                .map(|a| match a {
+                    OscType::Int(i) => Ok(*i),
+                    OscType::Float(f) => Ok(*f as i32),
+                    OscType::Double(d) => Ok(*d as i32),
+                    _ => Err(UniformError::TypeMismatch),
+                })
+                .collect();
+            uref.set_i32(&values?)
+        }
+        naga::ScalarKind::Uint => {
+            let values: Result<Vec<u32>, _> = args
+                .iter()
+                .map(|a| match a {
+                    OscType::Int(i) => Ok(*i as u32),
+                    OscType::Float(f) => Ok(*f as u32),
+                    OscType::Double(d) => Ok(*d as u32),
+                    _ => Err(UniformError::TypeMismatch),
+                })
+                .collect();
+            uref.set_u32(&values?)
+        }
+        _ => Err(UniformError::TypeMismatch),
+    }
+}
diff --git a/src/uniform.rs b/src/uniform.rs
new file mode 100644
index 0000000..5a66396
--- /dev/null
+++ b/src/uniform.rs
@@ -0,0 +1,616 @@
+use std::collections::HashMap;
+
+use naga::proc::Layouter;
+
+/// A borrowed, typed view into a region of the uniform buffer.
+/// Supports indexing into compound types and writing scalar values.
+pub struct UniformRef<'a> {
+    pub module: &'a naga::Module,
+    pub layouter: &'a Layouter,
+    pub ty: naga::Handle<naga::Type>,
+    pub data: &'a mut [u8],
+}
+
+#[derive(Debug)]
+pub enum UniformError {
+    NotIndexable,
+    IndexOutOfBounds,
+    TypeMismatch,
+    SizeMismatch,
+}
+
+impl std::fmt::Display for UniformError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::NotIndexable => write!(f, "type is not indexable"),
+            Self::IndexOutOfBounds => write!(f, "index out of bounds"),
+            Self::TypeMismatch => write!(f, "scalar type mismatch"),
+            Self::SizeMismatch => write!(f, "value count mismatch"),
+        }
+    }
+}
+
+impl std::error::Error for UniformError {}
+
+impl<'a> UniformRef<'a> {
+    /// Index into a compound type, returning a narrower UniformRef.
+    ///
+    /// - Vector: index selects a scalar component
+    /// - Matrix: index selects a column (vector)
+    /// - Array: index selects an element
+    /// - Struct: index selects a member by position
+    pub fn index(self, i: usize) -> Result<UniformRef<'a>, UniformError> {
+        let ty_inner = &self.module.types[self.ty].inner;
+        match *ty_inner {
+            naga::TypeInner::Vector { size, scalar } => {
+                let count = size as usize;
+                if i >= count {
+                    return Err(UniformError::IndexOutOfBounds);
+                }
+                let scalar_size = scalar.width as usize;
+                let offset = i * scalar_size;
+                let scalar_ty = find_or_expect_scalar_type(self.module, scalar);
+                Ok(UniformRef {
+                    module: self.module,
+                    layouter: self.layouter,
+                    ty: scalar_ty,
+                    data: &mut self.data[offset..offset + scalar_size],
+                })
+            }
+            naga::TypeInner::Matrix {
+                columns,
+                rows,
+                scalar,
+            } => {
+                let col_count = columns as usize;
+                if i >= col_count {
+                    return Err(UniformError::IndexOutOfBounds);
+                }
+                let col_ty = find_or_expect_vector_type(self.module, rows, scalar);
+                let col_layout = self.layouter[col_ty];
+                let stride = col_layout.to_stride();
+                let offset = i * stride as usize;
+                let size = col_layout.size as usize;
+                Ok(UniformRef {
+                    module: self.module,
+                    layouter: self.layouter,
+                    ty: col_ty,
+                    data: &mut self.data[offset..offset + size],
+                })
+            }
+            naga::TypeInner::Array { base, size, stride } => {
+                let len = match size {
+                    naga::ArraySize::Constant(n) => n.get() as usize,
+                    _ => return Err(UniformError::NotIndexable),
+                };
+                if i >= len {
+                    return Err(UniformError::IndexOutOfBounds);
+                }
+                let elem_size = self.layouter[base].size as usize;
+                let offset = i * stride as usize;
+                Ok(UniformRef {
+                    module: self.module,
+                    layouter: self.layouter,
+                    ty: base,
+                    data: &mut self.data[offset..offset + elem_size],
+                })
+            }
+            naga::TypeInner::Struct { ref members, .. } => {
+                if i >= members.len() {
+                    return Err(UniformError::IndexOutOfBounds);
+                }
+                let member = &members[i];
+                let offset = member.offset as usize;
+                let size = self.layouter[member.ty].size as usize;
+                Ok(UniformRef {
+                    module: self.module,
+                    layouter: self.layouter,
+                    ty: member.ty,
+                    data: &mut self.data[offset..offset + size],
+                })
+            }
+            _ => Err(UniformError::NotIndexable),
+        }
+    }
+
+    /// Get the leaf scalar type of this uniform.
+    pub fn leaf_scalar(&self) -> Result<naga::Scalar, UniformError> {
+        leaf_scalar_of(self.module, self.ty)
+    }
+
+    /// Count the total number of leaf scalars in this type.
+    fn leaf_count(&self) -> Result<usize, UniformError> {
+        let ty_inner = &self.module.types[self.ty].inner;
+        match *ty_inner {
+            naga::TypeInner::Scalar(_) => Ok(1),
+            naga::TypeInner::Vector { size, .. } => Ok(size as usize),
+            naga::TypeInner::Matrix { columns, rows, .. } => Ok(columns as usize * rows as usize),
+            naga::TypeInner::Array { base, size, .. } => {
+                let len = match size {
+                    naga::ArraySize::Constant(n) => n.get() as usize,
+                    _ => return Err(UniformError::NotIndexable),
+                };
+                let inner = &self.module.types[base].inner;
+                let per_elem = match *inner {
+                    naga::TypeInner::Scalar(_) => 1,
+                    naga::TypeInner::Vector { size, .. } => size as usize,
+                    naga::TypeInner::Matrix { columns, rows, .. } => {
+                        columns as usize * rows as usize
+                    }
+                    _ => return Err(UniformError::NotIndexable),
+                };
+                Ok(len * per_elem)
+            }
+            _ => Err(UniformError::NotIndexable),
+        }
+    }
+
+    pub fn set_f32(self, values: &[f32]) -> Result<(), UniformError> {
+        self.set_scalars(values, naga::ScalarKind::Float, 4)
+    }
+
+    pub fn set_i32(self, values: &[i32]) -> Result<(), UniformError> {
+        self.set_scalars(values, naga::ScalarKind::Sint, 4)
+    }
+
+    pub fn set_u32(self, values: &[u32]) -> Result<(), UniformError> {
+        self.set_scalars(values, naga::ScalarKind::Uint, 4)
+    }
+
+    fn set_scalars<T: Copy + bytemuck::NoUninit>(
+        self,
+        values: &[T],
+        expected_kind: naga::ScalarKind,
+        expected_width: u8,
+    ) -> Result<(), UniformError> {
+        let expected = self.leaf_count()?;
+        if values.len() != expected {
+            return Err(UniformError::SizeMismatch);
+        }
+        self.write_recursive(values, expected_kind, expected_width, &mut 0)
+    }
+
+    /// `val_idx` tracks position in the flat values array.
+    fn write_recursive<T: Copy + bytemuck::NoUninit>(
+        self,
+        values: &[T],
+        expected_kind: naga::ScalarKind,
+        expected_width: u8,
+        val_idx: &mut usize,
+    ) -> Result<(), UniformError> {
+        let ty_inner = &self.module.types[self.ty].inner;
+        match *ty_inner {
+            naga::TypeInner::Scalar(scalar) => {
+                if scalar.kind != expected_kind || scalar.width != expected_width {
+                    return Err(UniformError::TypeMismatch);
+                }
+                let bytes = bytemuck::bytes_of(&values[*val_idx]);
+                self.data[..bytes.len()].copy_from_slice(bytes);
+                *val_idx += 1;
+                Ok(())
+            }
+            naga::TypeInner::Vector { size, scalar } => {
+                if scalar.kind != expected_kind || scalar.width != expected_width {
+                    return Err(UniformError::TypeMismatch);
+                }
+                let count = size as usize;
+                let w = scalar.width as usize;
+                for c in 0..count {
+                    let offset = c * w;
+                    let bytes = bytemuck::bytes_of(&values[*val_idx]);
+                    self.data[offset..offset + w].copy_from_slice(bytes);
+                    *val_idx += 1;
+                }
+                Ok(())
+            }
+            naga::TypeInner::Matrix {
+                columns,
+                rows,
+                scalar,
+            } => {
+                if scalar.kind != expected_kind || scalar.width != expected_width {
+                    return Err(UniformError::TypeMismatch);
+                }
+                let col_ty = find_or_expect_vector_type(self.module, rows, scalar);
+                let col_stride = self.layouter[col_ty].to_stride() as usize;
+                let w = scalar.width as usize;
+                let row_count = rows as usize;
+                for col in 0..columns as usize {
+                    let col_offset = col * col_stride;
+                    for row in 0..row_count {
+                        let offset = col_offset + row * w;
+                        let bytes = bytemuck::bytes_of(&values[*val_idx]);
+                        self.data[offset..offset + w].copy_from_slice(bytes);
+                        *val_idx += 1;
+                    }
+                }
+                Ok(())
+            }
+            naga::TypeInner::Array { base, size, stride } => {
+                let len = match size {
+                    naga::ArraySize::Constant(n) => n.get() as usize,
+                    _ => return Err(UniformError::NotIndexable),
+                };
+                let elem_size = self.layouter[base].size as usize;
+                let stride = stride as usize;
+                for i in 0..len {
+                    let offset = i * stride;
+                    let elem_ref = UniformRef {
+                        module: self.module,
+                        layouter: self.layouter,
+                        ty: base,
+                        data: &mut self.data[offset..offset + elem_size],
+                    };
+                    elem_ref.write_recursive(values, expected_kind, expected_width, val_idx)?;
+                }
+                Ok(())
+            }
+            _ => Err(UniformError::TypeMismatch),
+        }
+    }
+}
+
+fn leaf_scalar_of(
+    module: &naga::Module,
+    ty: naga::Handle<naga::Type>,
+) -> Result<naga::Scalar, UniformError> {
+    match module.types[ty].inner {
+        naga::TypeInner::Scalar(scalar) => Ok(scalar),
+        naga::TypeInner::Vector { scalar, .. } => Ok(scalar),
+        naga::TypeInner::Matrix { scalar, .. } => Ok(scalar),
+        naga::TypeInner::Array { base, .. } => leaf_scalar_of(module, base),
+        _ => Err(UniformError::TypeMismatch),
+    }
+}
+
+/// Ensure that all component sub-types (scalars for vectors, vectors for matrix columns)
+/// exist as standalone entries in the module's type arena.
+/// naga doesn't always create these when they're only used implicitly.
+fn ensure_component_types(module: &mut naga::Module) {
+    let mut needed_scalars = Vec::new();
+    let mut needed_vectors = Vec::new();
+
+    for (_handle, ty) in module.types.iter() {
+        match ty.inner {
+            naga::TypeInner::Vector { scalar, .. } => {
+                needed_scalars.push(scalar);
+            }
+            naga::TypeInner::Matrix { rows, scalar, .. } => {
+                needed_scalars.push(scalar);
+                needed_vectors.push((rows, scalar));
+            }
+            _ => {}
+        }
+    }
+
+    for scalar in needed_scalars {
+        let exists = module
+            .types
+            .iter()
+            .any(|(_, ty)| ty.inner == naga::TypeInner::Scalar(scalar));
+        if !exists {
+            module.types.insert(
+                naga::Type {
+                    name: None,
+                    inner: naga::TypeInner::Scalar(scalar),
+                },
+                naga::Span::UNDEFINED,
+            );
+        }
+    }
+
+    for (size, scalar) in needed_vectors {
+        let exists = module
+            .types
+            .iter()
+            .any(|(_, ty)| ty.inner == naga::TypeInner::Vector { size, scalar });
+        if !exists {
+            module.types.insert(
+                naga::Type {
+                    name: None,
+                    inner: naga::TypeInner::Vector { size, scalar },
+                },
+                naga::Span::UNDEFINED,
+            );
+        }
+    }
+}
+
+fn find_or_expect_scalar_type(
+    module: &naga::Module,
+    scalar: naga::Scalar,
+) -> naga::Handle<naga::Type> {
+    for (handle, ty) in module.types.iter() {
+        if ty.inner == naga::TypeInner::Scalar(scalar) {
+            return handle;
+        }
+    }
+    panic!("scalar type {:?} not found in module type arena", scalar);
+}
+
+fn find_or_expect_vector_type(
+    module: &naga::Module,
+    size: naga::VectorSize,
+    scalar: naga::Scalar,
+) -> naga::Handle<naga::Type> {
+    for (handle, ty) in module.types.iter() {
+        if ty.inner == (naga::TypeInner::Vector { size, scalar }) {
+            return handle;
+        }
+    }
+    panic!(
+        "vector type vec{}<{:?}> not found in module type arena",
+        size as u8, scalar
+    );
+}
+
+/// Compare two types from potentially different naga Modules for structural equivalence.
+pub fn types_compatible(
+    a_mod: &naga::Module,
+    a_ty: naga::Handle<naga::Type>,
+    b_mod: &naga::Module,
+    b_ty: naga::Handle<naga::Type>,
+) -> bool {
+    let a = &a_mod.types[a_ty].inner;
+    let b = &b_mod.types[b_ty].inner;
+    match (a, b) {
+        (naga::TypeInner::Scalar(a), naga::TypeInner::Scalar(b)) => a == b,
+        (
+            naga::TypeInner::Vector {
+                size: as_,
+                scalar: asc,
+            },
+            naga::TypeInner::Vector {
+                size: bs,
+                scalar: bsc,
+            },
+        ) => as_ == bs && asc == bsc,
+        (
+            naga::TypeInner::Matrix {
+                columns: ac,
+                rows: ar,
+                scalar: asc,
+            },
+            naga::TypeInner::Matrix {
+                columns: bc,
+                rows: br,
+                scalar: bsc,
+            },
+        ) => ac == bc && ar == br && asc == bsc,
+        (
+            naga::TypeInner::Array {
+                base: ab,
+                size: naga::ArraySize::Constant(an),
+                ..
+            },
+            naga::TypeInner::Array {
+                base: bb,
+                size: naga::ArraySize::Constant(bn),
+                ..
+            },
+        ) => an == bn && types_compatible(a_mod, *ab, b_mod, *bb),
+        (
+            naga::TypeInner::Struct {
+                members: am,
+                span: asp,
+            },
+            naga::TypeInner::Struct {
+                members: bm,
+                span: bsp,
+            },
+        ) => {
+            asp == bsp
+                && am.len() == bm.len()
+                && am.iter().zip(bm.iter()).all(|(a, b)| {
+                    a.offset == b.offset && types_compatible(a_mod, a.ty, b_mod, b.ty)
+                })
+        }
+        _ => false,
+    }
+}
+
+/// Manages uniform buffer data for all `var<uniform>` globals in a shader.
+///
+/// Struct members are exposed as individually addressable named uniforms.
+pub struct UniformCache {
+    pub module: naga::Module,
+    pub layouter: Layouter,
+    uniforms: HashMap<String, UniformMember>,
+    buffers: Vec<BufferState>,
+    pub bind_group_layout: Option<wgpu::BindGroupLayout>,
+    pub bind_group: Option<wgpu::BindGroup>,
+    dirty: bool,
+}
+
+struct UniformMember {
+    buffer_idx: usize,
+    ty: naga::Handle<naga::Type>,
+    offset: usize,
+    size: usize,
+}
+
+struct BufferState {
+    /// @group and @binding
+    group: u32,
+    binding: u32,
+    data: Vec<u8>,
+    gpu_buffer: Option<wgpu::Buffer>,
+}
+
+impl UniformCache {
+    pub fn new() -> Self {
+        Self {
+            module: naga::Module::default(),
+            layouter: Layouter::default(),
+            uniforms: HashMap::new(),
+            buffers: Vec::new(),
+            bind_group_layout: None,
+            bind_group: None,
+            dirty: false,
+        }
+    }
+
+    /// Rebuild from a new shader module, transferring compatible uniform values.
+    pub fn refresh(&mut self, mut new_module: naga::Module, device: &wgpu::Device) {
+        ensure_component_types(&mut new_module);
+        let mut new_layouter = Layouter::default();
+        new_layouter.update(new_module.to_ctx()).unwrap();
+
+        let old_module = std::mem::replace(&mut self.module, naga::Module::default());
+        let old_uniforms = std::mem::take(&mut self.uniforms);
+        let old_buffers = std::mem::take(&mut self.buffers);
+
+        let mut new_uniforms: HashMap<String, UniformMember> = HashMap::new();
+        let mut new_buffers: Vec<BufferState> = Vec::new();
+
+        for (_handle, var) in new_module.global_variables.iter() {
+            if var.space != naga::AddressSpace::Uniform {
+                continue;
+            }
+            let binding = match &var.binding {
+                Some(b) => b,
+                None => continue,
+            };
+
+            let ty_inner = &new_module.types[var.ty].inner;
+            let layout = new_layouter[var.ty];
+            let buffer_idx = new_buffers.len();
+
+            let data = vec![0u8; layout.size as usize];
+            new_buffers.push(BufferState {
+                group: binding.group,
+                binding: binding.binding,
+                data,
+                gpu_buffer: None,
+            });
+
+            if let naga::TypeInner::Struct { ref members, .. } = *ty_inner {
+                for member in members {
+                    if let Some(ref name) = member.name {
+                        let member_size = new_layouter[member.ty].size as usize;
+                        new_uniforms.insert(
+                            name.clone(),
+                            UniformMember {
+                                buffer_idx,
+                                ty: member.ty,
+                                offset: member.offset as usize,
+                                size: member_size,
+                            },
+                        );
+                    }
+                }
+            } else if let Some(ref name) = var.name {
+                new_uniforms.insert(
+                    name.clone(),
+                    UniformMember {
+                        buffer_idx,
+                        ty: var.ty,
+                        offset: 0,
+                        size: layout.size as usize,
+                    },
+                );
+            }
+        }
+
+        for (name, new_member) in &new_uniforms {
+            if let Some(old_member) = old_uniforms.get(name) {
+                if old_member.size == new_member.size
+                    && types_compatible(&old_module, old_member.ty, &new_module, new_member.ty)
+                {
+                    let old_data = &old_buffers[old_member.buffer_idx].data
+                        [old_member.offset..old_member.offset + old_member.size];
+                    new_buffers[new_member.buffer_idx].data
+                        [new_member.offset..new_member.offset + new_member.size]
+                        .copy_from_slice(old_data);
+                }
+            }
+        }
+
+        for buf in &mut new_buffers {
+            let gpu_buffer = device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some(&format!("uniform@{}:{}", buf.group, buf.binding)),
+                size: buf.data.len() as u64,
+                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+                mapped_at_creation: false,
+            });
+            buf.gpu_buffer = Some(gpu_buffer);
+        }
+
+        let bind_group_layout = if !new_buffers.is_empty() {
+            let entries: Vec<wgpu::BindGroupLayoutEntry> = new_buffers
+                .iter()
+                .map(|buf| wgpu::BindGroupLayoutEntry {
+                    binding: buf.binding,
+                    visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
+                    ty: wgpu::BindingType::Buffer {
+                        ty: wgpu::BufferBindingType::Uniform,
+                        has_dynamic_offset: false,
+                        min_binding_size: None,
+                    },
+                    count: None,
+                })
+                .collect();
+            Some(
+                device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                    label: Some("uniform_bind_group_layout"),
+                    entries: &entries,
+                }),
+            )
+        } else {
+            None
+        };
+
+        let bind_group = bind_group_layout.as_ref().map(|layout| {
+            let entries: Vec<wgpu::BindGroupEntry> = new_buffers
+                .iter()
+                .map(|buf| wgpu::BindGroupEntry {
+                    binding: buf.binding,
+                    resource: buf.gpu_buffer.as_ref().unwrap().as_entire_binding(),
+                })
+                .collect();
+            device.create_bind_group(&wgpu::BindGroupDescriptor {
+                label: Some("uniform_bind_group"),
+                layout,
+                entries: &entries,
+            })
+        });
+
+        self.module = new_module;
+        self.layouter = new_layouter;
+        self.uniforms = new_uniforms;
+        self.buffers = new_buffers;
+        self.bind_group_layout = bind_group_layout;
+        self.bind_group = bind_group;
+        self.dirty = true;
+    }
+
+    pub fn get(&mut self, name: &str) -> Option<UniformRef<'_>> {
+        // Split the borrow: get member info first, then borrow buffer data
+        let member_info = self.uniforms.get(name)?;
+        let buffer_idx = member_info.buffer_idx;
+        let ty = member_info.ty;
+        let offset = member_info.offset;
+        let size = member_info.size;
+
+        let buf = &mut self.buffers[buffer_idx];
+        self.dirty = true;
+        Some(UniformRef {
+            module: &self.module,
+            layouter: &self.layouter,
+            ty,
+            data: &mut buf.data[offset..offset + size],
+        })
+    }
+
+    pub fn flush(&mut self, queue: &wgpu::Queue) {
+        if !self.dirty {
+            return;
+        }
+        for buf in &self.buffers {
+            if let Some(ref gpu_buffer) = buf.gpu_buffer {
+                queue.write_buffer(gpu_buffer, 0, &buf.data);
+            }
+        }
+        self.dirty = false;
+    }
+}