initial commit

1 files changed, 408 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");
+}