Galactica/crates/render/src/gpustate.rs

use std::{iter, rc::Rc};

use anyhow::Result;
use bytemuck;
use galactica_content::Content;
use galactica_system::data::ShipState;
use galactica_util::to_radians;
use glyphon::{FontSystem, Resolution, SwashCache, TextAtlas, TextRenderer};
use nalgebra::{Point2, Point3};
use wgpu;
use winit;

use crate::{
	globaluniform::{GlobalDataContent, GlobalUniform, ObjectData},
	pipeline::PipelineBuilder,
	shaderprocessor::preprocess_shader,
	starfield::Starfield,
	texturearray::TextureArray,
	ui::UiManager,
	vertexbuffer::{consts::SPRITE_INDICES, types::ObjectInstance},
	RenderInput, RenderState, VertexBuffers,
};

/// A high-level GPU wrapper. Reads game state (via RenderInput), produces pretty pictures.
pub struct GPUState {
	pub(crate) device: wgpu::Device,
	pub(crate) config: wgpu::SurfaceConfiguration,
	pub(crate) surface: wgpu::Surface,
	pub(crate) object_pipeline: wgpu::RenderPipeline,
	pub(crate) starfield_pipeline: wgpu::RenderPipeline,
	pub(crate) ui_pipeline: wgpu::RenderPipeline,
	pub(crate) radialbar_pipeline: wgpu::RenderPipeline,
	pub(crate) starfield: Starfield,
	pub(crate) texture_array: TextureArray,
	pub(crate) state: RenderState,
	pub(crate) ui: UiManager,
}

impl GPUState {
	/// Make a new GPUState that draws on `window`
	pub async fn new(window: winit::window::Window, ct: Rc<Content>) -> Result<Self> {
		let window_size = window.inner_size();
		let window_aspect = window_size.width as f32 / window_size.height as f32;

		let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
			backends: wgpu::Backends::all(),
			..Default::default()
		});

		let surface = unsafe { instance.create_surface(&window) }.unwrap();

		// Basic setup
		let device;
		let queue;
		let config;

		{
			let adapter = instance
				.request_adapter(&wgpu::RequestAdapterOptions {
					power_preference: wgpu::PowerPreference::default(),
					compatible_surface: Some(&surface),
					force_fallback_adapter: false,
				})
				.await
				.unwrap();

			(device, queue) = adapter
					.request_device(
						&wgpu::DeviceDescriptor {
							// TODO: remove nonuniform sampled textures
							features: wgpu::Features::TEXTURE_BINDING_ARRAY | wgpu::Features::SAMPLED_TEXTURE_AND_STORAGE_BUFFER_ARRAY_NON_UNIFORM_INDEXING,
							// We may need limits if we compile for wasm
							limits: wgpu::Limits::default(),
							label: Some("gpu device"),
						},
						None,
					)
					.await
					.unwrap();

			// Assume sRGB
			let surface_caps = surface.get_capabilities(&adapter);
			let surface_format = surface_caps
				.formats
				.iter()
				.copied()
				.filter(|f| f.is_srgb())
				.filter(|f| f.has_stencil_aspect())
				.next()
				.unwrap_or(surface_caps.formats[0]);

			config = wgpu::SurfaceConfiguration {
				usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
				format: surface_format,
				width: window_size.width,
				height: window_size.height,
				present_mode: surface_caps.present_modes[0],
				alpha_mode: surface_caps.alpha_modes[0],
				view_formats: vec![],
			};

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

		let vertex_buffers = VertexBuffers::new(&device, &ct);

		// Load uniforms
		let global_uniform = GlobalUniform::new(&device);
		let texture_array = TextureArray::new(&device, &queue, &ct)?;

		// Make sure these match the indices in each shader
		let bind_group_layouts = &[
			&texture_array.bind_group_layout,
			&global_uniform.bind_group_layout,
		];

		// Text renderer
		let mut text_atlas = TextAtlas::new(&device, &queue, wgpu::TextureFormat::Bgra8UnormSrgb);
		let mut text_font_system = FontSystem::new_with_locale_and_db(
			"en-US".to_string(),
			glyphon::fontdb::Database::new(),
		);

		let conf = ct.get_config();
		for font in &conf.font_files {
			text_font_system.db_mut().load_font_file(font)?;
		}

		// TODO: nice error if no family with this name is found
		text_font_system
			.db_mut()
			.set_sans_serif_family(conf.font_sans.clone());
		text_font_system
			.db_mut()
			.set_serif_family(conf.font_serif.clone());
		text_font_system
			.db_mut()
			.set_monospace_family(conf.font_mono.clone());
		//text_font_system
		//	.db_mut()
		//	.set_cursive_family(conf.font_cursive.clone());
		//text_font_system
		//	.db_mut()
		//	.set_fantasy_family(conf.font_fantasy.clone());

		let text_cache = SwashCache::new();
		let text_renderer = TextRenderer::new(
			&mut text_atlas,
			&device,
			wgpu::MultisampleState::default(),
			None,
		);

		// Create render pipelines
		let object_pipeline = PipelineBuilder::new("object", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"object.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(vertex_buffers.get_object())
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let starfield_pipeline = PipelineBuilder::new("starfield", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"starfield.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(vertex_buffers.get_starfield())
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let ui_pipeline = PipelineBuilder::new("ui", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/shaders/", "ui.wgsl")),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(vertex_buffers.get_ui())
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let radialbar_pipeline = PipelineBuilder::new("radialbar", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"radialbar.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(vertex_buffers.get_radialbar())
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let mut starfield = Starfield::new();
		starfield.regenerate(&ct);

		let mut state = RenderState {
			queue,
			window,
			window_size,
			window_aspect,
			global_uniform,
			vertex_buffers,
			text_atlas,
			text_cache,
			text_font_system,
			text_renderer,
		};

		return Ok(Self {
			ui: UiManager::new(ct, &mut state),
			device,
			config,
			surface,
			starfield,
			texture_array,
			object_pipeline,
			starfield_pipeline,
			ui_pipeline,
			radialbar_pipeline,
			state,
		});
	}
}

impl GPUState {
	/// Get the window we are attached to
	pub fn window(&self) -> &winit::window::Window {
		&self.state.window
	}

	/// Update window size.
	/// This should be called whenever our window is resized.
	pub fn resize(&mut self, ct: &Content) {
		let new_size = self.state.window.inner_size();
		if new_size.width > 0 && new_size.height > 0 {
			self.state.window_size = new_size;
			self.state.window_aspect = new_size.width as f32 / new_size.height as f32;
			self.config.width = new_size.width;
			self.config.height = new_size.height;
			self.surface.configure(&self.device, &self.config);
		}
		self.starfield.update_buffer(ct, &mut self.state);
	}

	/// Initialize the rendering engine
	pub fn init(&mut self, ct: &Content) {
		// Update global values
		self.state.queue.write_buffer(
			&self.state.global_uniform.atlas_buffer,
			0,
			bytemuck::cast_slice(&[self.texture_array.texture_atlas]),
		);

		self.starfield.update_buffer(ct, &mut self.state);
	}

	/// Main render function. Draws sprites on a window.
	pub fn render(&mut self, input: RenderInput) -> Result<(), wgpu::SurfaceError> {
		let input = Rc::new(input);

		// Update global values
		self.state.queue.write_buffer(
			&self.state.global_uniform.data_buffer,
			0,
			bytemuck::cast_slice(&[GlobalDataContent {
				camera_position_x: input.camera_pos.x,
				camera_position_y: input.camera_pos.y,
				camera_zoom: input.camera_zoom,
				camera_zoom_min: input.ct.get_config().zoom_min,
				camera_zoom_max: input.ct.get_config().zoom_max,
				window_size_w: self.state.window_size.width as f32,
				window_size_h: self.state.window_size.height as f32,
				window_scale: self.state.window.scale_factor() as f32,
				window_aspect: self.state.window_aspect,
				starfield_sprite: input.ct.get_config().starfield_texture.into(),
				starfield_tile_size: input.ct.get_config().starfield_size,
				starfield_size_min: input.ct.get_config().starfield_min_size,
				starfield_size_max: input.ct.get_config().starfield_max_size,
			}]),
		);

		self.state.frame_reset();

		let output = self.surface.get_current_texture()?;
		let view = output.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: &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,
			occlusion_query_set: None,
			timestamp_writes: None,
		});

		if self.ui.get_config().show_phys {
			// Create sprite instances

			// Game coordinates (relative to camera) of ne and sw corners of screen.
			// Used to skip off-screen sprites.
			let clip_ne = Point2::new(-self.state.window_aspect, 1.0) * input.camera_zoom;
			let clip_sw = Point2::new(self.state.window_aspect, -1.0) * input.camera_zoom;

			// Order matters, it determines what is drawn on top.
			// The order inside ships and projectiles doesn't matter,
			// but ships should always be under projectiles.
			self.push_system(&input, (clip_ne, clip_sw));
			self.push_ships(&input, (clip_ne, clip_sw));
			self.push_projectiles(&input, (clip_ne, clip_sw));
			self.push_effects(&input, (clip_ne, clip_sw));
		}

		self.ui.draw(input.clone(), &mut self.state).unwrap();

		// These should match the indices in each shader,
		// and should each have a corresponding bind group layout.
		render_pass.set_bind_group(0, &self.texture_array.bind_group, &[]);
		render_pass.set_bind_group(1, &self.state.global_uniform.bind_group, &[]);

		if self.ui.get_config().show_starfield {
			// Starfield pipeline
			self.state
				.vertex_buffers
				.get_starfield()
				.set_in_pass(&mut render_pass);
			render_pass.set_pipeline(&self.starfield_pipeline);
			render_pass.draw_indexed(
				0..SPRITE_INDICES.len() as u32,
				0,
				0..self.state.get_starfield_counter(),
			);
		}

		if self.ui.get_config().show_phys {
			// Sprite pipeline
			self.state
				.vertex_buffers
				.get_object()
				.set_in_pass(&mut render_pass);
			render_pass.set_pipeline(&self.object_pipeline);
			render_pass.draw_indexed(
				0..SPRITE_INDICES.len() as u32,
				0,
				0..self.state.get_object_counter(),
			);
		}

		// Ui pipeline
		self.state
			.vertex_buffers
			.get_ui()
			.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.ui_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..self.state.get_ui_counter(),
		);

		// Radial progress bars
		// TODO: do we need to do this every time?
		self.state
			.vertex_buffers
			.get_radialbar()
			.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.radialbar_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..self.state.get_radialbar_counter(),
		);

		let textareas = self.ui.get_textareas(&input, &self.state);
		self.state
			.text_renderer
			.prepare(
				&self.device,
				&self.state.queue,
				&mut self.state.text_font_system,
				&mut self.state.text_atlas,
				Resolution {
					width: self.state.window_size.width,
					height: self.state.window_size.height,
				},
				textareas,
				&mut self.state.text_cache,
			)
			.unwrap();

		self.state
			.text_renderer
			.render(&self.state.text_atlas, &mut render_pass)
			.unwrap();

		// begin_render_pass borrows encoder mutably,
		// so we need to drop it before calling finish.
		drop(render_pass);

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

		return Ok(());
	}
}

impl GPUState {
	fn push_ships(
		&mut self,
		input: &RenderInput,
		// NE and SW corners of screen
		screen_clip: (Point2<f32>, Point2<f32>),
	) {
		for s in input.phys_img.iter_ships() {
			let ship_pos;
			let ship_ang;
			let ship_cnt;
			match s.ship.get_data().get_state() {
				ShipState::Dead | ShipState::Landed { .. } => continue,

				ShipState::Collapsing { .. } | ShipState::Flying { .. } => {
					let r = &s.rigidbody;
					let pos = *r.translation();
					ship_pos = Point3::new(pos.x, pos.y, 1.0);
					let ship_rot = r.rotation();
					ship_ang = ship_rot.angle();
					ship_cnt = input.ct.get_ship(s.ship.get_data().get_content());
				}

				ShipState::UnLanding { current_z, .. } | ShipState::Landing { current_z, .. } => {
					let r = &s.rigidbody;
					let pos = *r.translation();
					ship_pos = Point3::new(pos.x, pos.y, *current_z);
					let ship_rot = r.rotation();
					ship_ang = ship_rot.angle();
					ship_cnt = input.ct.get_ship(s.ship.get_data().get_content());
				}
			}

			// Position adjusted for parallax
			// TODO: adjust parallax for zoom?
			// 1.0 is z-coordinate, which is constant for ships
			let pos: Point2<f32> =
				(Point2::new(ship_pos.x, ship_pos.y) - input.camera_pos) / ship_pos.z;

			// Game dimensions of this sprite post-scale.
			// Post-scale width or height, whichever is larger.
			// This is in game units.
			//
			// We take the maximum to account for rotated sprites.
			let m =
				(ship_cnt.size / ship_pos.z) * input.ct.get_sprite(ship_cnt.sprite).aspect.max(1.0);

			// Don't draw sprites that are off the screen
			if pos.x < screen_clip.0.x - m
				|| pos.y > screen_clip.0.y + m
				|| pos.x > screen_clip.1.x + m
				|| pos.y < screen_clip.1.y - m
			{
				continue;
			}

			let idx = self.state.get_object_counter();
			// Write this object's location data
			self.state.queue.write_buffer(
				&self.state.global_uniform.object_buffer,
				ObjectData::SIZE * idx as u64,
				bytemuck::cast_slice(&[ObjectData {
					xpos: ship_pos.x,
					ypos: ship_pos.y,
					zpos: ship_pos.z,
					angle: ship_ang,
					size: ship_cnt.size,
					parent: 0,
					is_child: 0,
					_padding: Default::default(),
				}]),
			);

			// Push this object's instance
			let anim_state = s.ship.get_anim_state();
			self.state.push_object_buffer(ObjectInstance {
				texture_index: anim_state.texture_index(),
				texture_fade: anim_state.fade,
				object_index: idx as u32,
				color: [1.0, 1.0, 1.0, 1.0],
			});

			if {
				let is_flying = match s.ship.get_data().get_state() {
					ShipState::Flying { .. }
					| ShipState::UnLanding { .. }
					| ShipState::Landing { .. } => true,
					_ => false,
				};
				is_flying
			} {
				for (engine_point, anim) in s.ship.iter_engine_anim() {
					self.state.queue.write_buffer(
						&self.state.global_uniform.object_buffer,
						ObjectData::SIZE * self.state.get_object_counter() as u64,
						bytemuck::cast_slice(&[ObjectData {
							// Note that we adjust the y-coordinate for half-height,
							// not the x-coordinate, even though our ships point east
							// at 0 degrees. This is because this is placed pre-rotation,
							// and the parent rotation adjustment in our object shader
							// automatically accounts for this.
							xpos: engine_point.pos.x,
							ypos: engine_point.pos.y - engine_point.size / 2.0,
							zpos: 1.0,
							// We still need an adjustment here, though,
							// since engine sprites point north (with exhaust towards the south)
							angle: to_radians(90.0),
							size: engine_point.size,
							parent: idx as u32,
							is_child: 1,
							_padding: Default::default(),
						}]),
					);

					let anim_state = anim.get_texture_idx();
					self.state.push_object_buffer(ObjectInstance {
						texture_index: anim_state.texture_index(),
						texture_fade: anim_state.fade,
						object_index: self.state.get_object_counter() as u32,
						color: [1.0, 1.0, 1.0, 1.0],
					});
				}
			}
		}
	}

	fn push_projectiles(
		&mut self,
		input: &RenderInput,
		// NE and SW corners of screen
		screen_clip: (Point2<f32>, Point2<f32>),
	) {
		for p in input.phys_img.iter_projectiles() {
			let r = &p.rigidbody;
			let proj_pos = *r.translation();
			let proj_rot = r.rotation();
			let proj_ang = proj_rot.angle();
			let proj_cnt = &p.projectile.content; // TODO: don't clone this?

			// Position adjusted for parallax
			// TODO: adjust parallax for zoom?
			// 1.0 is z-coordinate, which is constant for projectiles
			let pos = (proj_pos - input.camera_pos) / 1.0;

			// Game dimensions of this sprite post-scale.
			// Post-scale width or height, whichever is larger.
			// This is in game units.
			//
			// We take the maximum to account for rotated sprites.
			let m = (proj_cnt.size / 1.0) * input.ct.get_sprite(proj_cnt.sprite).aspect.max(1.0);

			// Don't draw sprites that are off the screen
			if pos.x < screen_clip.0.x - m
				|| pos.y > screen_clip.0.y + m
				|| pos.x > screen_clip.1.x + m
				|| pos.y < screen_clip.1.y - m
			{
				continue;
			}

			let idx = self.state.get_object_counter();
			// Write this object's location data
			self.state.queue.write_buffer(
				&self.state.global_uniform.object_buffer,
				ObjectData::SIZE * idx as u64,
				bytemuck::cast_slice(&[ObjectData {
					xpos: proj_pos.x,
					ypos: proj_pos.y,
					zpos: 1.0,
					angle: proj_ang,
					size: 0f32.max(proj_cnt.size + p.projectile.size_rng),
					parent: 0,
					is_child: 0,
					_padding: Default::default(),
				}]),
			);

			let anim_state = p.projectile.get_anim_state();
			self.state.push_object_buffer(ObjectInstance {
				texture_index: anim_state.texture_index(),
				texture_fade: anim_state.fade,
				object_index: idx as u32,
				color: [1.0, 1.0, 1.0, 1.0],
			});
		}
	}

	fn push_system(
		&mut self,
		input: &RenderInput,
		// NE and SW corners of screen
		screen_clip: (Point2<f32>, Point2<f32>),
	) {
		let system = input.ct.get_system(input.current_system);

		for o in &system.objects {
			// Position adjusted for parallax
			let pos: Point2<f32> = (Point2::new(o.pos.x, o.pos.y) - input.camera_pos) / o.pos.z;

			// Game dimensions of this sprite post-scale.
			// Post-scale width or height, whichever is larger.
			// This is in game units.
			//
			// We take the maximum to account for rotated sprites.
			let m = (o.size / o.pos.z) * input.ct.get_sprite(o.sprite).aspect.max(1.0);

			// Don't draw sprites that are off the screen
			if pos.x < screen_clip.0.x - m
				|| pos.y > screen_clip.0.y + m
				|| pos.x > screen_clip.1.x + m
				|| pos.y < screen_clip.1.y - m
			{
				continue;
			}

			let idx = self.state.get_object_counter();
			// Write this object's location data
			self.state.queue.write_buffer(
				&self.state.global_uniform.object_buffer,
				ObjectData::SIZE * idx as u64,
				bytemuck::cast_slice(&[ObjectData {
					xpos: o.pos.x,
					ypos: o.pos.y,
					zpos: o.pos.z,
					angle: o.angle,
					size: o.size,
					parent: 0,
					is_child: 0,
					_padding: Default::default(),
				}]),
			);

			let sprite = input.ct.get_sprite(o.sprite);
			let texture_a = sprite.get_first_frame(); // ANIMATE

			// Push this object's instance
			self.state.push_object_buffer(ObjectInstance {
				texture_index: [texture_a, texture_a],
				texture_fade: 1.0,
				object_index: idx as u32,
				color: [1.0, 1.0, 1.0, 1.0],
			});
		}
	}

	fn push_effects(
		&mut self,
		input: &RenderInput,
		// NE and SW corners of screen
		screen_clip: (Point2<f32>, Point2<f32>),
	) {
		for p in input.phys_img.iter_effects() {
			let r = &p.rigidbody;
			let pos = *r.translation();
			let rot = r.rotation();
			let ang = rot.angle();

			// Position adjusted for parallax
			// TODO: adjust parallax for zoom?
			// 1.0 is z-coordinate, which is constant for projectiles
			let adjusted_pos = (pos - input.camera_pos) / 1.0;

			// Game dimensions of this sprite post-scale.
			// Post-scale width or height, whichever is larger.
			// This is in game units.
			//
			// We take the maximum to account for rotated sprites.
			let m = (p.effect.size / 1.0)
				* input
					.ct
					.get_sprite(p.effect.anim.get_sprite())
					.aspect
					.max(1.0);

			// Don't draw sprites that are off the screen
			if adjusted_pos.x < screen_clip.0.x - m
				|| adjusted_pos.y > screen_clip.0.y + m
				|| adjusted_pos.x > screen_clip.1.x + m
				|| adjusted_pos.y < screen_clip.1.y - m
			{
				continue;
			}

			let idx = self.state.get_object_counter();
			// Write this object's location data
			self.state.queue.write_buffer(
				&self.state.global_uniform.object_buffer,
				ObjectData::SIZE * idx as u64,
				bytemuck::cast_slice(&[ObjectData {
					xpos: pos.x,
					ypos: pos.y,
					zpos: 1.0,
					angle: ang,
					size: p.effect.size,
					parent: 0,
					is_child: 0,
					_padding: Default::default(),
				}]),
			);

			let anim_state = p.effect.anim.get_texture_idx();
			self.state.push_object_buffer(ObjectInstance {
				texture_index: anim_state.texture_index(),
				texture_fade: anim_state.fade,
				object_index: idx as u32,
				color: [1.0, 1.0, 1.0, p.get_fade()],
			});
		}
	}
}