main.rs

use clap::Parser;

mod examples;
use examples::*;

#[derive(Parser)]
#[command(
    name = "example",
    about = "WGSL module examples and hello_triangle shader demo"
)]
enum Cli {
    /// Show available example shaders.
    Show,
    /// Validate and print the source.
    ///
    /// Prints raw source and source generated by `naga`.
    Source { module: Option<String> },
    // Linkage,
    /// Run the hello_triangle shader
    // TODO(schell): expand this to run any shader.
    // This means making legitimate shaders that actually do something.
    Run,
}

fn print_available_modules() {
    eprintln!("Available modules:");
    for module in EXAMPLE_MODULES {
        eprintln!("  {}", module.name);
    }
}

fn validate_and_print_source(module: &wgsl_rs::Module) {
    println!("## {}", module.name);

    let source = module.wgsl_source();
    println!("raw source:\n\n{source}\n\n");

    // Template modules can't be parsed/validated standalone — their IR
    // contains unresolved `Type::TypeParam` nodes (rendered as
    // `__TP{name}__` placeholders, which aren't valid WGSL identifiers).
    // Just print the raw template source.
    if module.is_template() {
        println!(
            "(this is a template module with type parameters {:?}; call \
             `instantiate(&[ir::Type::...])` to produce a concrete shader)",
            module.module_type_params
        );
        return;
    }

    // Parse the source into a Module.
    let module: naga::Module = naga::front::wgsl::parse_str(&source).unwrap();

    // Validate the module.
    // Validation can be made less restrictive by changing the ValidationFlags.
    let result = naga::valid::Validator::new(
        naga::valid::ValidationFlags::all(),
        naga::valid::Capabilities::all(),
    )
    .subgroup_stages(naga::valid::ShaderStages::all())
    .subgroup_operations(naga::valid::SubgroupOperationSet::all())
    .validate(&module);

    let info = match result {
        Err(e) => {
            panic!("{}", e.emit_to_string(&source));
        }
        Ok(i) => i,
    };

    let wgsl =
        naga::back::wgsl::write_string(&module, &info, naga::back::wgsl::WriterFlags::empty())
            .unwrap();
    println!("naga source:\n\n{wgsl}");
}

// /// Print the linkage within a module.
// fn print_linkage(module: &wgsl_rs::Module) {
//     let name = module.name;
//     // Test hello_triangle linkage
//     println!(
//         "{name}::linkage::SHADER_SOURCE length: {}",
//         hello_triangle::linkage::SHADER_SOURCE.len()
//     );
//     println!(
//         "hello_triangle bind_group_0 layout entries: {}",
//         hello_triangle::linkage::bind_group_0::LAYOUT_ENTRIES.len()
//     );
//     println!(
//         "hello_triangle vtx_main entry point: {}",
//         hello_triangle::linkage::vtx_main::ENTRY_POINT
//     );
//     println!(
//         "hello_triangle frag_main entry point: {}",
//         hello_triangle::linkage::frag_main::ENTRY_POINT
//     );

//     // Test compute_shader linkage
//     println!(
//         "\ncompute_shader::linkage::SHADER_SOURCE length: {}",
//         compute_shader::linkage::SHADER_SOURCE.len()
//     );
//     println!(
//         "compute_shader bind_group_0 layout entries: {}",
//         compute_shader::linkage::bind_group_0::LAYOUT_ENTRIES.len()
//     );
//     println!(
//         "compute_shader main entry point: {}",
//         compute_shader::linkage::main::ENTRY_POINT
//     );
//     println!(
//         "compute_shader main workgroup size: {:?}",
//         compute_shader::linkage::main::WORKGROUP_SIZE
//     );

//     // Test structs linkage
//     println!(
//         "\nstructs::linkage::SHADER_SOURCE length: {}",
//         structs::linkage::SHADER_SOURCE.len()
//     );
//     println!(
//         "structs frag_shader entry point: {}",
//         structs::linkage::frag_shader::ENTRY_POINT
//     );

//     println!("\n=== Linkage API tests passed ===");
// }

/// Build the linkage into a working `winit` + `wgpu` app, as a
/// dogfooding test.
fn build_linkage() {
    use std::sync::Arc;

    use futures::executor::block_on;
    use winit::{
        application::ApplicationHandler,
        event::WindowEvent,
        event_loop::{ControlFlow, EventLoop},
        window::Window,
    };

    let event_loop = EventLoop::new().unwrap();

    struct WgpuStuff {
        _instance: wgpu::Instance,
        surface: wgpu::Surface<'static>,
        surface_config: wgpu::SurfaceConfiguration,
        _adapter: wgpu::Adapter,
        device: wgpu::Device,
        queue: wgpu::Queue,
    }

    impl WgpuStuff {
        fn new(window: Arc<Window>, display_handle: winit::event_loop::OwnedDisplayHandle) -> Self {
            let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_with_display_handle(
                Box::new(display_handle),
            ));
            let surface = instance.create_surface(window).unwrap();
            let adapter = block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::default(),
                compatible_surface: Some(&surface),
                force_fallback_adapter: false,
            }))
            .expect("Failed to find an appropriate adapter");

            let (device, queue) =
                block_on(adapter.request_device(&wgpu::DeviceDescriptor::default()))
                    .expect("Failed to create device");
            let surface_config = surface
                .get_default_config(&adapter, 800, 600)
                .expect("no default surface config");
            surface.configure(&device, &surface_config);

            Self {
                _instance: instance,
                surface,
                surface_config,
                _adapter: adapter,
                device,
                queue,
            }
        }
    }

    struct HelloTriangle {
        frame: u32,
        frame_uniform_buffer: wgpu::Buffer,
        bindgroup: wgpu::BindGroup,
        render_pipeline: wgpu::RenderPipeline,
    }

    impl HelloTriangle {
        fn new(wgpu_stuff: &WgpuStuff) -> Self {
            let device = &wgpu_stuff.device;
            let queue = &wgpu_stuff.queue;
            let frame = 0u32;
            let frame_uniform_buffer = hello_triangle::create_frame_buffer(device);
            queue.write_buffer(&frame_uniform_buffer, 0, &frame.to_ne_bytes());

            let bindgroup_layout = hello_triangle::linkage::bind_group_0::layout(device);
            let bindgroup = hello_triangle::linkage::bind_group_0::create(
                device,
                &bindgroup_layout,
                frame_uniform_buffer.as_entire_binding(),
            );

            let module = hello_triangle::linkage::shader_module(device);
            let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("hello_triangle"),
                bind_group_layouts: &[Some(&bindgroup_layout)],
                immediate_size: 0,
            });
            let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                label: Some("hello_triangle"),
                layout: Some(&pipeline_layout),
                vertex: hello_triangle::linkage::vtx_main::vertex_state(&module),
                primitive: wgpu::PrimitiveState {
                    topology: wgpu::PrimitiveTopology::TriangleList,
                    strip_index_format: None,
                    front_face: wgpu::FrontFace::Ccw,
                    cull_mode: None,
                    unclipped_depth: false,
                    polygon_mode: wgpu::PolygonMode::Fill,
                    conservative: false,
                },
                depth_stencil: None,
                multisample: wgpu::MultisampleState::default(),
                fragment: Some(hello_triangle::linkage::frag_main::fragment_state(
                    &module,
                    &[Some(wgpu::ColorTargetState {
                        format: wgpu_stuff
                            .surface
                            .get_configuration()
                            .expect("missing surface configuration")
                            .format,
                        blend: None,
                        write_mask: wgpu::ColorWrites::all(),
                    })],
                )),
                multiview_mask: None,
                cache: None,
            });

            Self {
                frame,
                frame_uniform_buffer,
                bindgroup,
                render_pipeline,
            }
        }
    }

    struct AppInner {
        window: Arc<Window>,
        wgpu_stuff: WgpuStuff,
        hello_triangle: HelloTriangle,
    }

    impl AppInner {
        pub fn new(event_loop: &winit::event_loop::ActiveEventLoop) -> Self {
            let window = Arc::new(
                event_loop
                    .create_window(Window::default_attributes())
                    .unwrap(),
            );
            let wgpu_stuff = WgpuStuff::new(window.clone(), event_loop.owned_display_handle());
            let hello_triangle = HelloTriangle::new(&wgpu_stuff);
            Self {
                window,
                wgpu_stuff,
                hello_triangle,
            }
        }
    }

    #[derive(Default)]
    struct App {
        inner: Option<AppInner>,
    }

    impl ApplicationHandler for App {
        fn resumed(&mut self, event_loop: &winit::event_loop::ActiveEventLoop) {
            self.inner = Some(AppInner::new(event_loop));
        }

        fn window_event(
            &mut self,
            event_loop: &winit::event_loop::ActiveEventLoop,
            _window_id: winit::window::WindowId,
            event: winit::event::WindowEvent,
        ) {
            match event {
                WindowEvent::CloseRequested => {
                    println!("Closing");
                    event_loop.exit();
                }
                WindowEvent::RedrawRequested => {
                    if let Some(AppInner {
                        window,
                        wgpu_stuff,
                        hello_triangle,
                    }) = self.inner.as_mut()
                    {
                        let device = &wgpu_stuff.device;
                        let queue = &wgpu_stuff.queue;

                        hello_triangle.frame += 1;
                        queue.write_buffer(
                            &hello_triangle.frame_uniform_buffer,
                            0,
                            &hello_triangle.frame.to_ne_bytes(),
                        );

                        let mut encoder =
                            device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                                label: Some("pass"),
                            });
                        let texture = match wgpu_stuff.surface.get_current_texture() {
                            wgpu::CurrentSurfaceTexture::Success(t)
                            | wgpu::CurrentSurfaceTexture::Suboptimal(t) => t,
                            wgpu::CurrentSurfaceTexture::Outdated => {
                                log::warn!(
                                    "surface acquire returned Outdated; reconfiguring and \
                                     skipping frame"
                                );
                                wgpu_stuff
                                    .surface
                                    .configure(&wgpu_stuff.device, &wgpu_stuff.surface_config);
                                window.request_redraw();
                                return;
                            }
                            wgpu::CurrentSurfaceTexture::Timeout
                            | wgpu::CurrentSurfaceTexture::Occluded => {
                                log::debug!(
                                    "surface acquire returned Timeout/Occluded; skipping frame"
                                );
                                window.request_redraw();
                                return;
                            }
                            wgpu::CurrentSurfaceTexture::Lost => {
                                log::warn!(
                                    "surface acquire returned Lost; reconfiguring and skipping \
                                     frame"
                                );
                                let size = window.inner_size();
                                wgpu_stuff.surface_config.width = size.width;
                                wgpu_stuff.surface_config.height = size.height;
                                wgpu_stuff
                                    .surface
                                    .configure(&wgpu_stuff.device, &wgpu_stuff.surface_config);
                                window.request_redraw();
                                return;
                            }
                            wgpu::CurrentSurfaceTexture::Validation => {
                                log::warn!("surface acquire returned Validation; skipping frame");
                                window.request_redraw();
                                return;
                            }
                        };
                        {
                            let mut render_pass =
                                encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                                    label: Some("pass"),
                                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                                        view: &texture
                                            .texture
                                            .create_view(&wgpu::TextureViewDescriptor::default()),
                                        depth_slice: None,
                                        resolve_target: None,
                                        ops: wgpu::Operations {
                                            load: wgpu::LoadOp::Clear(wgpu::Color {
                                                r: 0.0,
                                                g: 0.0,
                                                b: 0.0,
                                                a: 0.0,
                                            }),
                                            store: wgpu::StoreOp::Store,
                                        },
                                    })],
                                    ..Default::default()
                                });
                            render_pass.set_pipeline(&hello_triangle.render_pipeline);
                            render_pass.set_bind_group(0, &hello_triangle.bindgroup, &[]);
                            render_pass.draw(0..3, 0..1);
                        }

                        let index = queue.submit(Some(encoder.finish()));
                        texture.present();

                        // TODO: maybe do this elsewhere?
                        device
                            .poll(wgpu::PollType::Wait {
                                submission_index: Some(index),
                                timeout: None,
                            })
                            .unwrap();

                        window.request_redraw();
                    }
                }
                _ => {}
            }
        }
    }

    // ControlFlow::Poll continuously runs the event loop, even if the OS hasn't
    // dispatched any events. This is ideal for games and similar applications.
    event_loop.set_control_flow(ControlFlow::Poll);

    let mut app = App::default();
    event_loop.run_app(&mut app).unwrap();
}

fn print_all_modules() {
    for module in EXAMPLE_MODULES {
        validate_and_print_source(module);
    }
}

pub fn main() {
    let cli = Cli::parse();

    match cli {
        Cli::Show => print_available_modules(),
        Cli::Source { module } => {
            if let Some(name) = module {
                if get_module_by_name(&name)
                    .map(validate_and_print_source)
                    .is_none()
                {
                    eprintln!("Unknown module: {name}");
                    print_available_modules();
                    std::process::exit(1);
                }
            } else {
                print_all_modules();
            }
        }
        Cli::Run => build_linkage(),
    }
}