SolarEngine/solar_engine/src/state.rs

use std::cmp::max;
use std::collections::HashMap;
use std::sync::{Arc};
use cgmath::{perspective, Deg, Matrix4, Point3, Vector3};
use log::info;
use pollster::FutureExt;
use wgpu::util::DeviceExt;
use wgpu::{Adapter, Device, Instance, PresentMode, Queue, Surface, SurfaceCapabilities, SurfaceConfiguration};
use winit::dpi::PhysicalSize;
use winit::window::{Window};
use crate::camera::Camera;
use crate::render::{create_circle_vertices, Geometry, Globals, InstanceRaw, RenderInstance, SampleCount, Shape, Vertex};

pub struct State<'a> {
    surface: Surface<'a>,
    device: Device,
    queue: Queue,
    config: SurfaceConfiguration,
    sample_count: SampleCount,

    size: PhysicalSize<u32>,
    window: Arc<Window>,

    render_pipeline: wgpu::RenderPipeline,

    instances: Vec<RenderInstance>,
    instance_buffer: wgpu::Buffer,

    geometries: HashMap<Shape, Geometry>,
    global_bind_group: wgpu::BindGroup,
    global_buffer: wgpu::Buffer,

    camera: Camera,
}

impl<'a> State<'a> {
    pub(crate) fn new(window: Window) -> Self {
        let window_arc = Arc::new(window);
        let size = window_arc.inner_size();
        let instance = Self::create_gpu_instance();
        let surface = instance.create_surface(window_arc.clone()).unwrap();
        let adapter = Self::create_adapter(instance, &surface);
        let (device, queue) = Self::create_device(&adapter);
        let surface_caps = surface.get_capabilities(&adapter);
        let config = Self::create_surface_config(size, surface_caps);
        surface.configure(&device, &config);

        let sample_count = SampleCount(Self::probe_msaa_support(&device, &config));
        info!("MSAA sample count: {}", sample_count.0);

        let aspect = config.width as f32 / config.height as f32;
        let camera = Camera::new(aspect);
        let view_proj = camera.build_view_projection_matrix();

        let globals = Globals {
            view_proj: view_proj.into(),
        };
        let global_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Global Uniform Buffer"),
            contents: bytemuck::cast_slice(&[globals]),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });
        let global_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Global 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 global_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &global_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: global_buffer.as_entire_binding(),
            }],
            label: Some("Global Bind Group"),
        });

        let render_pipeline = Self::create_render_pipeline(&device, &config, sample_count.0, &global_bind_group_layout);
        let geometries = Self::create_geometries(&device);

        let instances = vec![];
        let instance_data: Vec<InstanceRaw> = instances.iter().map(RenderInstance::to_raw).collect();
        let instance_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Instance Buffer"),
            contents: bytemuck::cast_slice(&instance_data),
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        });

        Self {
            surface,
            device,
            queue,
            config,
            sample_count,
            size,
            window: window_arc,
            render_pipeline,
            geometries,
            global_bind_group,
            global_buffer,
            instances,
            instance_buffer,
            camera,
        }
    }

    fn create_geometries(device: &Device) -> HashMap<Shape, Geometry> {
        let mut geometries = HashMap::new();

        let polygon_vertices = vec![
            Vertex { position: [-0.0868241, 0.49240386, 0.0], color: [0.5, 0.0, 0.5] },
            Vertex { position: [-0.49513406, 0.06958647, 0.0], color: [0.5, 0.0, 0.5] },
            Vertex { position: [-0.21918549, -0.44939706, 0.0], color: [0.5, 0.0, 0.5] },
            Vertex { position: [0.35966998, -0.3473291, 0.0], color: [0.5, 0.0, 0.5] },
            Vertex { position: [0.44147372, 0.2347359, 0.0], color: [0.5, 0.0, 0.5] },
        ];
        let polygon_indices = vec![0, 1, 4, 1, 2, 4, 2, 3, 4];

        let polygon_vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Polygon Vertex Buffer"),
            contents: bytemuck::cast_slice(&polygon_vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });
        let polygon_index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Polygon Index Buffer"),
            contents: bytemuck::cast_slice(&polygon_indices),
            usage: wgpu::BufferUsages::INDEX,
        });
        geometries.insert(Shape::Polygon, Geometry {
            vertex_buffer: polygon_vertex_buffer,
            index_buffer: polygon_index_buffer,
            index_count: polygon_indices.len() as u32,
        });

        let (circle_vertices, circle_indices) = create_circle_vertices(512, 0.5, [0.5, 0.5, 0.5]);
        let circle_vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Circle Vertex Buffer"),
            contents: bytemuck::cast_slice(&circle_vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });
        let circle_index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Circle Index Buffer"),
            contents: bytemuck::cast_slice(&circle_indices),
            usage: wgpu::BufferUsages::INDEX,
        });
        geometries.insert(Shape::Circle, Geometry {
            vertex_buffer: circle_vertex_buffer,
            index_buffer: circle_index_buffer,
            index_count: circle_indices.len() as u32,
        });

        geometries
    }

    fn probe_msaa_support(device: &Device, config: &SurfaceConfiguration) -> u32 {
        pollster::block_on(async {
            for &count in &[16, 8, 4, 2] {
                device.push_error_scope(wgpu::ErrorFilter::Validation);

                let _ = device.create_texture(&wgpu::TextureDescriptor {
                    label: Some("MSAA Probe"),
                    size: wgpu::Extent3d {
                        width: 4,
                        height: 4,
                        depth_or_array_layers: 1,
                    },
                    mip_level_count: 1,
                    sample_count: count,
                    dimension: wgpu::TextureDimension::D2,
                    format: config.format,
                    usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
                    view_formats: &[],
                });

                if device.pop_error_scope().await.is_none() {
                    return count;
                }
            }

            1 // fallback
        })
    }
    
    fn create_render_pipeline(device: &Device, config: &SurfaceConfiguration, sample_count: u32, global_bind_group_layout: &wgpu::BindGroupLayout) -> wgpu::RenderPipeline {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
        });

        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Render Pipeline Layout"),
                bind_group_layouts: &[&global_bind_group_layout],
                push_constant_ranges: &[],
            });

        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[Vertex::desc(), InstanceRaw::desc()],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: config.format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                // Setting this to anything other than Fill requires Features::NON_FILL_POLYGON_MODE
                polygon_mode: wgpu::PolygonMode::Fill,
                // Requires Features::DEPTH_CLIP_CONTROL
                unclipped_depth: false,
                // Requires Features::CONSERVATIVE_RASTERIZATION
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: sample_count,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
            cache: None,
        })
    }

    fn create_surface_config(size: PhysicalSize<u32>, capabilities: SurfaceCapabilities) -> wgpu::SurfaceConfiguration {
        let surface_format = capabilities.formats.iter()
            .find(|f| f.is_srgb())
            .copied()
            .unwrap_or(capabilities.formats[0]);

        SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface_format,
            width: size.width,
            height: size.height,
            present_mode: PresentMode::AutoVsync,
            alpha_mode: capabilities.alpha_modes[0],
            view_formats: vec![],
            desired_maximum_frame_latency: 2,
        }
    }

    fn create_device(adapter: &Adapter) -> (Device, Queue) {
        adapter.request_device(
            &wgpu::DeviceDescriptor {
                required_features: wgpu::Features::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES,
                required_limits: wgpu::Limits::default(),
                memory_hints: Default::default(),
                label: None,
                trace: Default::default(),
            }).block_on().unwrap()
    }

    fn create_adapter(instance: Instance, surface: &Surface) -> Adapter {
        instance.request_adapter(
            &wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::default(),
                compatible_surface: Some(&surface),
                force_fallback_adapter: false,
            }
        ).block_on().unwrap()
    }

    fn create_gpu_instance() -> Instance {
        Instance::new(&wgpu::InstanceDescriptor {
            backends: wgpu::Backends::PRIMARY,
            ..Default::default()
        })
    }

    pub(crate) fn resize(&mut self, new_size: PhysicalSize<u32>) {
        self.size = new_size;

        self.config.width = max(new_size.width, 1);
        self.config.height = max(new_size.height, 1);

        self.surface.configure(&self.device, &self.config);

        self.camera.set_aspect(self.config.width as f32 / self.config.height as f32);
        let view_proj = self.camera.build_view_projection_matrix();

        let new_globals = Globals {
            view_proj: view_proj.into(),
        };
        self.queue.write_buffer(&self.global_buffer, 0, bytemuck::cast_slice(&[new_globals]));

        println!("Resized to {:?} from state!", new_size);
    }

    pub(crate) fn render(&mut self) -> Result<(), wgpu::SurfaceError> {
        let output = self.surface.get_current_texture()?;

        let multisampled_texture = self.device.create_texture(&wgpu::TextureDescriptor {
            label: Some("Multisampled Render Target"),
            size: wgpu::Extent3d {
                width: self.config.width,
                height: self.config.height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: self.sample_count.get(),
            dimension: wgpu::TextureDimension::D2,
            format: self.config.format,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            view_formats: &[],
        });
        let multisampled_view = multisampled_texture.create_view(&Default::default());

        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("Render Encoder"),
        });

        {
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Render Pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &multisampled_view,
                    resolve_target: Some(&output.texture.create_view(&Default::default())),
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color {
                            r: 1.0,
                            g: 0.2,
                            b: 0.3,
                            a: 1.0,
                        }),
                        store: wgpu::StoreOp::Store,
                    }
                })],
                depth_stencil_attachment: None,
                occlusion_query_set: None,
                timestamp_writes: None,
            });

            render_pass.set_pipeline(&self.render_pipeline);
            render_pass.set_bind_group(0, &self.global_bind_group, &[]);

            for shape in [Shape::Polygon, Shape::Circle] {
                let geometry = &self.geometries[&shape];

                let relevant_instances: Vec<_> = self.instances
                    .iter()
                    .enumerate()
                    .filter(|(_, inst)| inst.shape == shape)
                    .map(|(i, _)| i as u32)
                    .collect();

                if relevant_instances.is_empty() {
                    continue;
                }

                render_pass.set_vertex_buffer(0, geometry.vertex_buffer.slice(..));
                render_pass.set_vertex_buffer(1, self.instance_buffer.slice(..));
                render_pass.set_index_buffer(geometry.index_buffer.slice(..), wgpu::IndexFormat::Uint16);
                render_pass.draw_indexed(0..geometry.index_count, 0, 0..relevant_instances.len() as u32);
            }
        }

        self.queue.submit(std::iter::once(encoder.finish()));
        output.present();

        Ok(())
    }

    pub fn set_instances(&mut self, instances: Vec<RenderInstance>) {
        let raw_data: Vec<InstanceRaw> = instances.iter().map(RenderInstance::to_raw).collect();
        let byte_len = (raw_data.len() * size_of::<InstanceRaw>()) as wgpu::BufferAddress;

        // Resize the buffer if necessary
        if byte_len > self.instance_buffer.size() {
            self.instance_buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("Instance Buffer (resized)"),
                contents: bytemuck::cast_slice(&raw_data),
                usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            });
        } else {
            self.queue.write_buffer(&self.instance_buffer, 0, bytemuck::cast_slice(&raw_data));
        }

        self.instances = instances;
    }

    pub fn window(&self) -> &Window {
        &self.window
    }
}