SolarEngine/simulator/src/main.rs

use std::cmp::max;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::thread;
use std::time::{Duration, Instant};
use cgmath::num_traits::{pow, ToPrimitive};
use cgmath::Rotation3;
use pollster::FutureExt;
use wgpu::util::DeviceExt;
use wgpu::{Adapter, Device, Instance, PresentMode, Queue, Surface, SurfaceCapabilities, SurfaceConfiguration};
use winit::application::ApplicationHandler;
use winit::dpi::PhysicalSize;
use winit::event::{ElementState, WindowEvent};
use winit::event::WindowEvent::KeyboardInput;
use winit::event_loop::{ActiveEventLoop, EventLoop};
use winit::keyboard::Key;
use winit::platform::modifier_supplement::KeyEventExtModifierSupplement;
use winit::window::{Window, WindowId};
use solar_engine::{Body, Simulator};

pub async fn run() {
    let event_loop = EventLoop::new().unwrap();

    let mut window_state = StateApplication::new();
    let _ = event_loop.run_app(&mut window_state);
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Globals {
    aspect_ratio: f32,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
enum Shape {
    Polygon,
    Circle,
}

struct Geometry {
    vertex_buffer: wgpu::Buffer,
    index_buffer: wgpu::Buffer,
    index_count: u32,
}

struct RenderInstance {
    position: cgmath::Vector3<f32>,
    rotation: cgmath::Quaternion<f32>,
    color: [f32; 3],
    scale: f32,
    shape: Shape,
}

impl RenderInstance {
    fn to_raw(&self) -> InstanceRaw {
        let model = cgmath::Matrix4::from_translation(self.position)
            * cgmath::Matrix4::from(self.rotation)
            * cgmath::Matrix4::from_scale(self.scale);
        InstanceRaw {
            model: model.into(),
            color: self.color,
        }
    }
}

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

impl InstanceRaw {
    fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: size_of::<InstanceRaw>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Instance,
            attributes: &[
                wgpu::VertexAttribute { offset: 0, shader_location: 5, format: wgpu::VertexFormat::Float32x4 },
                wgpu::VertexAttribute { offset: 16, shader_location: 6, format: wgpu::VertexFormat::Float32x4 },
                wgpu::VertexAttribute { offset: 32, shader_location: 7, format: wgpu::VertexFormat::Float32x4 },
                wgpu::VertexAttribute { offset: 48, shader_location: 8, format: wgpu::VertexFormat::Float32x4 },
                wgpu::VertexAttribute { offset: 64, shader_location: 9, format: wgpu::VertexFormat::Float32x3 },
            ],
        }
    }
}

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

impl Vertex {
    const ATTRIBS: [wgpu::VertexAttribute; 2] =
        wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3];

    fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: size_of::<Self>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &Self::ATTRIBS,
        }
    }
}

struct StateApplication<'a> {
    state: Option<State<'a>>,
}

impl<'a> StateApplication<'a> {
    pub fn new() -> Self {
        Self {
            state: None,
        }
    }
}

impl<'a> ApplicationHandler for StateApplication<'a>{
    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
        let window = event_loop.create_window(Window::default_attributes().with_title("Hello!")).unwrap();
        self.state = Some(State::new(window));
    }

    fn window_event(&mut self, event_loop: &ActiveEventLoop, window_id: WindowId, event: WindowEvent) {
        let window = self.state.as_ref().unwrap().window();

        if window.id() == window_id {
            match event {
                WindowEvent::CloseRequested => {
                    event_loop.exit();
                }
                WindowEvent::Resized(physical_size) => {
                    self.state.as_mut().unwrap().resize(physical_size);
                }
                WindowEvent::RedrawRequested => {
                    self.state.as_mut().unwrap().update();
                    self.state.as_mut().unwrap().render().unwrap();
                }
                WindowEvent::KeyboardInput { .. } => {
                    if let Some(state) = self.state.as_mut() {
                        if state.input(&event) {
                            return;
                        }
                    }
                }
                _ => {}
            }
        }
    }

    fn about_to_wait(&mut self, _event_loop: &ActiveEventLoop) {
        let window = self.state.as_ref().unwrap().window();
        window.request_redraw();
    }
}

fn create_circle_vertices(segment_count: usize, radius: f32, color: [f32; 3]) -> (Vec<Vertex>, Vec<u16>) {
    let mut vertices = vec![Vertex { position: [0.0, 0.0, 0.0], color }];
    let mut indices = vec![];

    for i in 0..=segment_count {
        let theta = (i as f32) / (segment_count as f32) * std::f32::consts::TAU;
        let x = radius * theta.cos();
        let y = radius * theta.sin();
        vertices.push(Vertex { position: [x, y, 0.0], color });
    }

    for i in 1..=segment_count {
        indices.push(0);
        indices.push(i as u16);
        indices.push((i % segment_count + 1) as u16);
    }

    (vertices, indices)
}

struct State<'a> {
    surface: Surface<'a>,
    device: Device,
    queue: Queue,
    config: wgpu::SurfaceConfiguration,
    // sample_count: u32,

    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,

    simulator: Arc<RwLock<Simulator>>,
}

impl<'a> State<'a> {
    pub 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 = Self::probe_msaa_support(&device, &config);
        println!("Probe: {:?}", sample_count);

        let globals = Globals {
            aspect_ratio: config.width as f32 / config.height as f32,
        };
        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, &global_bind_group_layout);

        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,
        });

        let mut sim = Simulator::new();
        sim.add_body(Body {
            name: "Sun".to_string(),
            position: [0.0, 0.0],
            velocity: [0.0, 0.0],
            mass: 1.989e30,
        });
        sim.add_body(Body {
            name: "Earth".to_string(),
            position: [1.496e11, 0.0],
            velocity: [0.0, 29780.0],
            mass: 5.972e24,
        });

        let instances = {
            let sun_pos = sim.bodies[0].position;
            sim.bodies.iter().enumerate().map(|(i, b)| RenderInstance {
                position: cgmath::Vector3::new(
                    ((b.position[0] - sun_pos[0]) / 1.496e11) as f32,
                    ((b.position[1] - sun_pos[1]) / 1.496e11) as f32,
                    0.0,
                ),
                rotation: cgmath::Quaternion::from_angle_z(cgmath::Deg(0.0)),
                color: match i {
                    0 => [1.0, 1.0, 0.0],
                    1 => [0.0, 0.0, 1.0],
                    _ => [0.5, 0.5, 0.5],
                },
                scale: 0.1,
                shape: Shape::Circle
            }).collect::<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,
        });

        let simulator = Arc::new(RwLock::new(sim));
        let simulator_clone = simulator.clone();
        thread::spawn(move || {
            let mut last = Instant::now();
            loop {
                let now = Instant::now();
                let dt = now.duration_since(last).as_secs_f64();
                last = now;

                {
                    let mut sim = simulator_clone.write().unwrap();
                    let timewarp = sim.get_timewarp();
                    sim.step(dt * timewarp as f64);
                }

                thread::sleep(Duration::from_millis(1));
            }
        });

        Self {
            surface,
            device,
            queue,
            config,
            size,
            window: window_arc,
            
            global_bind_group,
            global_buffer,

            render_pipeline,
            geometries,

            instances,
            instance_buffer,

            simulator,
        }
    }

    fn probe_msaa_support(device: &wgpu::Device, config: &wgpu::SurfaceConfiguration) -> u32 {
        for &count in &[16, 8, 4, 2] {
            let texture_desc = 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: &[],
            };

            // This will panic if not supported, so run this ONLY if you are confident (not ideal)
            if let Ok(_) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                device.create_texture(&texture_desc);
            })) {
                return count;
            }
        }

        1 // fallback
    }

    pub fn input(&mut self, event: &WindowEvent) -> bool {
        if let KeyboardInput { event, .. } = event {
            if event.state == ElementState::Pressed {
                return match event.key_without_modifiers().as_ref() {
                    Key::Character(".") => {
                        let mut sim = self.simulator.write().unwrap();
                        sim.increase_timewarp();
                        println!("Timewarp: {}", sim.get_timewarp());
                        true
                    }
                    Key::Character(",") => {
                        let mut sim = self.simulator.write().unwrap();
                        sim.decrease_timewarp();
                        println!("Timewarp: {}", sim.get_timewarp());
                        true
                    }
                    Key::Character("-") => {
                        let mut sim = self.simulator.write().unwrap();
                        sim.reset_timewarp();
                        println!("Timewarp: {}", sim.get_timewarp());
                        true
                    }
                    _ => {
                        false
                    }
                }
            }
        }
        
        false
    }

    fn update(&mut self) {
        let sim = self.simulator.read().unwrap();

        let updated_instances: Vec<RenderInstance> = {
            sim.bodies.iter().enumerate().map(|(i, b)| RenderInstance {
                position: cgmath::Vector3::new((b.position[0] / 1.496e11) as f32, (b.position[1] / 1.496e11) as f32, 0.0),
                rotation: cgmath::Quaternion::from_angle_z(cgmath::Deg(0.0)),
                color: match i {
                    0 => [1.0, 1.0, 0.0],
                    1 => [0.0, 0.0, 1.0],
                    _ => [0.5, 0.5, 0.5],
                },
                scale: 0.5,
                shape: Shape::Circle
            }).collect()
        };

        let instance_data: Vec<InstanceRaw> = updated_instances.iter().map(RenderInstance::to_raw).collect();
        self.queue.write_buffer(&self.instance_buffer, 0, bytemuck::cast_slice(&instance_data));
        self.instances = updated_instances;
    }
    
    fn create_render_pipeline(device: &Device, config: &wgpu::SurfaceConfiguration, 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, // 1.
                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: 8,
                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]);

        wgpu::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 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);

        let new_globals = Globals {
            aspect_ratio: self.config.width as f32 / self.config.height as f32,
        };
        self.queue.write_buffer(&self.global_buffer, 0, bytemuck::cast_slice(&[new_globals]));

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

    pub 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: 8,
            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 window(&self) -> &Window {
        &self.window
    }
}

fn main() {
    pollster::block_on(run());
}