lib.rs

use std::mem;
use wasm_bindgen::prelude::*;
use wgpu::util::DeviceExt;
use winit::{
    dpi::PhysicalSize, event_loop::EventLoop, platform::web::WindowBuilderExtWebSys,
    window::WindowBuilder,
};
use std::borrow::Cow;

macro_rules! log {
    ($($t:tt)*) => {
        web_sys::console::log_1(&format!($($t)*).into());
    }
}

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct Vertex {
    pos: [f32; 3],
    color: [f32; 3],
}

impl Vertex {
    fn layout() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[
                wgpu::VertexAttribute {
                    // pos
                    offset: 0,
                    shader_location: 0,
                    format: wgpu::VertexFormat::Float32x3,
                },
                wgpu::VertexAttribute {
                    // color
                    offset: mem::size_of::<[f32; 3]>() as wgpu::BufferAddress,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x3,
                },
            ],
        }
    }
}

const VERTICES: &[Vertex] = &[
    Vertex {
        pos: [-1.0, 1.0, 0.0],  // Top-left
        color: [1.0, 0.0, 1.0], // Magenta
    },
    Vertex {
        pos: [-1.0, -1.0, 0.0], // Bottom-left
        color: [0.0, 0.0, 1.0], // Blue
    },
    Vertex {
        pos: [1.0, 1.0, 0.0],   // Top-right
        color: [1.0, 1.0, 0.0], // Yellow
    },
    Vertex {
        pos: [1.0, -1.0, 0.0],  // Bottom-right
        color: [0.0, 1.0, 0.0], // Green
    },
];

const INDICES: &[u32] = &[0, 1, 2, 2, 1, 3]; // CCW, quad

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable, Debug, Default)]
struct UniformData {
    mouse_move: [f32; 2],
    mouse_click: [f32; 2],
    resolution: [f32; 2],
    time: f32,
    _padding: f32,
}

#[derive(Debug)]
struct State {
    window: web_sys::Window,
    canvas: web_sys::HtmlCanvasElement,
    winit_window: winit::window::Window,
    surface: wgpu::Surface,
    device: wgpu::Device,
    queue: wgpu::Queue,
    surface_config: wgpu::SurfaceConfiguration,
    vertex_buffer: wgpu::Buffer,
    index_buffer: wgpu::Buffer,
    index_num: u32,
    uniform_data: UniformData,
    uniform_buffer: wgpu::Buffer,
    uniform_bind_group: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
    animation_cb: Closure<dyn FnMut(f32)>,
}

static mut STATE: Option<State> = None;

impl State {
    async fn new() -> Self {
        // window
        let window = get_window();
        // canvas
        let canvas = State::init_canvas(&window).expect_throw("Failed to get the canvas");
        // winit window
        let winit_window = State::create_window(&canvas).expect_throw("Failed to create a window");
        // wgpu instance
        let instance = State::create_instance();
        // wgpu surface
        let surface = State::create_surface(&instance, &winit_window)
            .expect_throw("Failed to create a surface");
        // wgpu adapter
        let adapter = State::create_adapter(&instance, &surface)
            .await
            .expect_throw("Failed to create an adapter");
        // wgpu device and queue
        let (device, queue) = State::create_device_and_queue(&adapter)
            .await
            .expect_throw("Failed to create a device and a queue");
        // wgpu surface configuration
        let surface_config = State::create_surface_configuration(
            &surface,
            &adapter,
            &device,
            canvas.width(),
            canvas.height(),
        );
        // wgpu vertex buffer
        let vertex_buffer = State::create_vertex_buffer(&device, bytemuck::cast_slice(VERTICES));
        // wgpu index buffer
        let index_buffer = State::create_index_buffer(&device, bytemuck::cast_slice(INDICES));
        // wgpu uniform buffer
        let uniform_data = UniformData {
            resolution: [canvas.width() as f32, canvas.height() as f32],
            mouse_move: [std::f32::MIN, std::f32::MIN],
            mouse_click: [std::f32::MIN, std::f32::MIN],
            ..Default::default()
        };
        let (uniform_buffer, uniform_layout, uniform_bind_group) =
            State::create_uniform_buffer(&device, bytemuck::cast_slice(&[uniform_data][..]));
        // wgpu shader module
        let shader_module = State::create_shader_module(&device);
        // wgpu render pipeline
        let render_pipeline = State::create_render_pipeline(
            &device,
            &[&uniform_layout],
            &shader_module,
            &surface_config,
        );
        // animation_loop
        let animation_cb = State::create_animation_loop();

        Self {
            window,
            canvas,
            winit_window,
            surface,
            device,
            queue,
            surface_config,
            vertex_buffer,
            index_buffer,
            index_num: INDICES.len() as u32,
            uniform_data,
            uniform_buffer,
            uniform_bind_group,
            render_pipeline,
            animation_cb,
        }
    }

    fn init_canvas(window: &web_sys::Window) -> Option<web_sys::HtmlCanvasElement> {
        let element = window.document()?.get_element_by_id("canvas0")?;
        let canvas = element.dyn_into::<web_sys::HtmlCanvasElement>().ok()?;
        let (width, height) = (canvas.client_width() as u32, canvas.client_height() as u32);
        canvas.set_width(width);
        canvas.set_height(height);
        log!("Canvas size: ({}, {})", canvas.width(), canvas.height());
        Some(canvas)
    }

    fn create_window(canvas: &web_sys::HtmlCanvasElement) -> Option<winit::window::Window> {
        // winit's window writes its physical and logical sizes in the HTML element
        // as attributes and inner style respectively
        // So, we should put "!important" in CSS to overwrite the inner style
        let window = WindowBuilder::new()
            .with_inner_size(PhysicalSize::new(
                canvas.client_width(),
                canvas.client_height(),
            ))
            .with_canvas(Some(canvas.clone()))
            .build(&EventLoop::new())
            .ok()?;
        log!(
            "Window size: ({}, {})",
            window.inner_size().width,
            window.inner_size().height
        );
        Some(window)
    }

    fn create_instance() -> wgpu::Instance {
        wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::BROWSER_WEBGPU,
            ..Default::default()
        })
    }

    fn create_surface(
        instance: &wgpu::Instance,
        winit_window: &winit::window::Window,
    ) -> Option<wgpu::Surface> {
        unsafe { instance.create_surface(winit_window).ok() }
    }

    async fn create_adapter(
        instance: &wgpu::Instance,
        surface: &wgpu::Surface,
    ) -> Option<wgpu::Adapter> {
        instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                compatible_surface: Some(surface),
                force_fallback_adapter: false,
                ..Default::default()
            })
            .await
    }

    async fn create_device_and_queue(
        adapter: &wgpu::Adapter,
    ) -> Option<(wgpu::Device, wgpu::Queue)> {
        adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    features: wgpu::Features::empty(),
                    limits: wgpu::Limits::downlevel_webgl2_defaults(),
                    label: None,
                },
                None,
            )
            .await
            .ok()
    }

    fn create_surface_configuration(
        surface: &wgpu::Surface,
        adapter: &wgpu::Adapter,
        device: &wgpu::Device,
        width: u32,
        height: u32,
    ) -> wgpu::SurfaceConfiguration {
        let surface_caps = surface.get_capabilities(adapter);
        let surface_config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface_caps.formats[0],
            width,
            height,
            present_mode: surface_caps.present_modes[0],
            alpha_mode: surface_caps.alpha_modes[0],
            view_formats: vec![],
        };
        surface.configure(device, &surface_config);
        surface_config
    }

    fn create_vertex_buffer(device: &wgpu::Device, contents: &[u8]) -> wgpu::Buffer {
        device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Vertex buffer"),
            contents,
            usage: wgpu::BufferUsages::VERTEX,
        })
    }

    fn create_index_buffer(device: &wgpu::Device, contents: &[u8]) -> wgpu::Buffer {
        device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Index buffer"),
            contents,
            usage: wgpu::BufferUsages::INDEX,
        })
    }

    fn create_uniform_buffer(
        device: &wgpu::Device,
        contents: &[u8],
    ) -> (wgpu::Buffer, wgpu::BindGroupLayout, wgpu::BindGroup) {
        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Uniform buffer"),
            contents,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });
        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Uniform bind group layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Uniform bind group"),
            layout: &bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: buffer.as_entire_binding(),
            }],
        });
        (buffer, bind_group_layout, bind_group)
    }

    fn create_shader_module(device: &wgpu::Device) -> wgpu::ShaderModule {
        // Creates shader module from a single file.
        let by_single_file = || -> wgpu::ShaderModule {
            device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("Shader module"),
                source: wgpu::ShaderSource::Wgsl(include_str!("monolithic.wgsl").into()),
            })
        };

        // Composes shader module using naga_oil composition.
        let by_naga_oil = || -> wgpu::ShaderModule {
            use naga_oil::compose::{Composer, ComposableModuleDescriptor, NagaModuleDescriptor};
            let mut composer = Composer::default();
            composer.add_composable_module(ComposableModuleDescriptor {
                source: include_str!("uniform.wgsl"), // Path to the file relative to the current file
                file_path: "uniform.wgsl", // Path to the file relative to the Cargo.toml
                ..Default::default()
            }).unwrap();
            let naga_module = composer.make_naga_module(NagaModuleDescriptor {
                source: include_str!("top.wgsl"), // Path to the file relative to the current file
                file_path: "example.wgsl", // Path to the file relative to the Cargo.toml
                shader_defs: [("UNIFORM".to_owned(), Default::default())].into(),
                ..Default::default()
            }).unwrap();
            device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("Shader module"),
                source: wgpu::ShaderSource::Naga(Cow::Owned(naga_module)),
            })
        };

        let by_my_wgsl = || -> wgpu::ShaderModule {
            use my_wgsl::*;

            #[wgsl_decl_struct]
            struct UniformData {
                mouse_move: vec2<f32>,
                mouse_click: vec2<f32>,
                resolution: vec2<f32>,
                time: f32,
            }

            #[wgsl_decl_struct]
            struct VertexInput {
                #[location(0)] pos: vec3<f32>,
                #[location(1)] color: vec3<f32>
            }
            
            #[wgsl_decl_struct]
            struct VertexOutput {
                #[builtin(position)] pos: vec4<f32>,
                #[location(1)] color: vec3<f32>
            }

            let mut builder = Builder::new();

            wgsl_structs!(builder, UniformData, VertexInput, VertexOutput);

            wgsl_bind!(builder, group(0) binding(0) var<uniform> uni: UniformData);

            wgsl_fn!(builder,
                #[vertex]
                fn v_main(input: VertexInput) -> VertexOutput {
                    var output: VertexOutput;
                    output.pos = vec4<f32>(input.pos, 1.0);
                    output.color = input.color;
                    return output;
                }
            );

            wgsl_fn!(builder,
                #[fragment]
                fn f_main(input: VertexOutput) -> #[location(0)] vec4<f32> {
                    let x = select(0.0, 0.3, distance(input.pos.xy, uni.mouse_move) < 25.0);
                    let y = select(0.0, 0.3, distance(input.pos.xy, uni.mouse_click) < 25.0);
                    return vec4f(input.color + x - y, 1.0);
                }
            );

            let wgsl = builder.build();

            device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("Shader module"),
                source: wgpu::ShaderSource::Wgsl(wgsl.into()),
            })
        };

        // With measurement.
        fn with_measure (title: &str, id: &str, f: impl FnOnce() -> wgpu::ShaderModule) -> wgpu::ShaderModule {
            let s = now();
            let res = f();
            let e = now();
            let msg = format!("[{title}] took {} ms", (e - s).floor());
            set_inner_html(id, &msg);
            res
        }

        // Let's try all of them.
        with_measure("Creating shader module from a single file", "measure_single_file", by_single_file);
        with_measure("Compositing shader module using naga_oil", "measure_naga_oil", by_naga_oil);
        with_measure("Compositing shader module using my_wgsl", "measure_my_wgsl", by_my_wgsl)
    }

    fn create_render_pipeline(
        device: &wgpu::Device,
        bind_group_layouts: &[&wgpu::BindGroupLayout],
        shader_module: &wgpu::ShaderModule,
        surface_config: &wgpu::SurfaceConfiguration,
    ) -> wgpu::RenderPipeline {
        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Render pipeline layout"),
                bind_group_layouts,
                push_constant_ranges: &[],
            });
        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: shader_module,
                entry_point: "v_main",
                buffers: &[Vertex::layout()],
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                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,
            },
            fragment: Some(wgpu::FragmentState {
                module: shader_module,
                entry_point: "f_main",
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_config.format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            multiview: None,
        })
    }

    fn create_animation_loop() -> Closure<dyn FnMut(f32)> {
        Closure::<dyn FnMut(f32)>::new(|time: f32| unsafe {
            let state = STATE.as_mut().unwrap_unchecked();
            state.render(time);
            state.request_animation_frame();
        })
    }

    #[inline(always)]
    fn request_animation_frame(&self) {
        self.window
            .request_animation_frame(self.animation_cb.as_ref().unchecked_ref())
            .expect_throw("Failed to request an animation frame");
    }

    fn resize(&mut self) {
        let new_width = self.canvas.client_width() as u32;
        let new_height = self.canvas.client_height() as u32;
        if new_width != self.canvas.width() || new_height != self.canvas.height() {
            // Synchronize manually
            self.canvas.set_width(new_width);
            self.canvas.set_height(new_height);
            self.winit_window
                .set_inner_size(PhysicalSize::new(new_width, new_height));
            self.surface_config.width = new_width;
            self.surface_config.height = new_height;
            self.surface.configure(&self.device, &self.surface_config);

            // Update uniform data
            self.uniform_data.resolution = [new_width as f32, new_height as f32];

            log!("Resized: ({}, {})", new_width, new_height);
        }
    }

    fn mousemove(&mut self, event: web_sys::MouseEvent) {
        // Update uniform data
        self.uniform_data.mouse_move = [event.offset_x() as f32, event.offset_y() as f32];
    }

    fn click(&mut self, event: web_sys::MouseEvent) {
        // Update uniform data
        self.uniform_data.mouse_click = [event.offset_x() as f32, event.offset_y() as f32];
    }

    fn render(&mut self, time: f32) {
        // Write uniform data to its buffer
        self.uniform_data.time = time * 0.001;
        self.queue.write_buffer(
            &self.uniform_buffer,
            0,
            bytemuck::cast_slice(&[self.uniform_data][..]),
        );

        let surface_texture = self.surface.get_current_texture().unwrap();
        let texture_view = surface_texture.texture.create_view(&Default::default());
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Render command encoder"),
            });

        {
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Render pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &texture_view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color {
                            r: 0.0,
                            g: 0.0,
                            b: 0.0,
                            a: 1.0,
                        }),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            render_pass.set_pipeline(&self.render_pipeline);
            render_pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));
            render_pass.set_index_buffer(self.index_buffer.slice(..), wgpu::IndexFormat::Uint32);
            render_pass.set_bind_group(0, &self.uniform_bind_group, &[]);
            render_pass.draw_indexed(0..self.index_num, 0, 0..1);
        }
        self.queue.submit(std::iter::once(encoder.finish()));
        surface_texture.present();
    }
}

#[wasm_bindgen(start)]
pub fn main() {
    std::panic::set_hook(Box::new(console_error_panic_hook::hook));
    console_log::init_with_level(log::Level::Warn).expect("Couldn't initialize logger");
}

#[wasm_bindgen]
pub async fn run() {
    unsafe {
        STATE = Some(State::new().await);
        let state = STATE.as_mut().unwrap_unchecked();

        // Sets resize event listener on window.
        add_event_listener("", "resize", || STATE.as_mut().unwrap_unchecked().resize());

        // Sets mousemove/click event listener on canvas
        add_event_listener_with_mouseevent("canvas0", "mousemove", |event: web_sys::MouseEvent| {
            STATE.as_mut().unwrap_unchecked().mousemove(event)
        });
        add_event_listener_with_mouseevent("canvas0", "click", |event: web_sys::MouseEvent| {
            STATE.as_mut().unwrap_unchecked().click(event)
        });
        state.request_animation_frame();
    };
}

fn get_window() -> web_sys::Window {
    web_sys::window().expect_throw("Failed to get window")
}

fn get_document() -> web_sys::Document {
    get_window().document().expect_throw("Failed to get document")
}

fn get_element_by_id(id: &str) -> web_sys::Element {
    get_document()
        .get_element_by_id(id)
        .expect_throw("Failed to get element")
}

fn set_inner_html(id: &str, value: &str) {
    get_element_by_id(id).set_inner_html(value)
}

fn add_event_listener(id: &str, type_: &str, f: impl Fn() + 'static) {
    let listener = Closure::<dyn Fn()>::new(f);
    if id.is_empty() {
        get_window()
            .add_event_listener_with_callback(type_, listener.as_ref().unchecked_ref())
            .expect("Failed to add an event listener");
    } else {
        get_element_by_id(id)
            .add_event_listener_with_callback(type_, listener.as_ref().unchecked_ref())
            .expect("Failed to add an event listener");
    }
    listener.forget(); // Leak, but it occurs just once
}

fn add_event_listener_with_mouseevent(
    id: &str,
    type_: &str,
    f: impl Fn(web_sys::MouseEvent) + 'static,
) {
    let listener = Closure::<dyn Fn(_)>::new(f);
    if id.is_empty() {
        get_window()
            .add_event_listener_with_callback(type_, listener.as_ref().unchecked_ref())
            .expect("Failed to add an event listener");
    } else {
        get_element_by_id(id)
            .add_event_listener_with_callback(type_, listener.as_ref().unchecked_ref())
            .expect("Failed to add an event listener");
    }
    listener.forget(); // Leak, but it occurs just once
}

fn get_performance() -> web_sys::Performance {
    get_window().performance().expect_throw("Failed to get performance")
} 

fn now() -> f64 {
    get_performance().now()
}