Galactica/src/render/gpustate.rs

use anyhow::Result;
use bytemuck;
use cgmath::{Deg, EuclideanSpace, Matrix4, Point2, Vector2, Vector3};
use std::{iter, rc::Rc};
use wgpu;
use winit::{self, dpi::PhysicalSize, window::Window};

use super::{
	consts::{OPENGL_TO_WGPU_MATRIX, SPRITE_INSTANCE_LIMIT, STARFIELD_INSTANCE_LIMIT},
	globaldata::{GlobalData, GlobalDataContent},
	pipeline::PipelineBuilder,
	texturearray::TextureArray,
	vertexbuffer::{
		consts::{SPRITE_INDICES, SPRITE_VERTICES},
		types::{SpriteInstance, StarfieldInstance, TexturedVertex},
		VertexBuffer,
	},
	Sprite,
};
use crate::{consts, game::Game};

pub struct GPUState {
	device: wgpu::Device,
	config: wgpu::SurfaceConfiguration,
	surface: wgpu::Surface,
	queue: wgpu::Queue,

	pub window: Window,
	pub window_size: winit::dpi::PhysicalSize<u32>,
	window_aspect: f32,

	sprite_pipeline: wgpu::RenderPipeline,
	starfield_pipeline: wgpu::RenderPipeline,
	starfield_count: u32,

	texture_array: TextureArray,
	global_data: GlobalData,
	vertex_buffers: VertexBuffers,
}

struct VertexBuffers {
	sprite: Rc<VertexBuffer>,
	starfield: Rc<VertexBuffer>,
}

impl GPUState {
	pub async fn new(window: Window) -> 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 {
						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 {
			sprite: Rc::new(VertexBuffer::new::<TexturedVertex, SpriteInstance>(
				"sprite",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				SPRITE_INSTANCE_LIMIT,
			)),

			starfield: Rc::new(VertexBuffer::new::<TexturedVertex, StarfieldInstance>(
				"starfield",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				STARFIELD_INSTANCE_LIMIT,
			)),
		};

		// Load uniforms
		let global_data = GlobalData::new(&device);
		let texture_array = TextureArray::new(&device, &queue)?;

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

		// Create render pipelines
		let sprite_pipeline = PipelineBuilder::new("sprite", &device)
			.set_shader(include_str!(concat!(
				env!("CARGO_MANIFEST_DIR"),
				"/src/render/shaders/",
				"sprite.wgsl"
			)))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(&vertex_buffers.sprite)
			.set_bind_group_layouts(bind_group_layouts)
			.build();

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

		return Ok(Self {
			device,
			config,
			surface,
			queue,

			window,
			window_size,
			window_aspect,

			sprite_pipeline,
			starfield_pipeline,

			texture_array,
			global_data,
			vertex_buffers,
			starfield_count: 0,
		});
	}

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

	pub fn resize(&mut self, game: &Game, new_size: PhysicalSize<u32>) {
		if new_size.width > 0 && new_size.height > 0 {
			self.window_size = new_size;
			self.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.update_starfield_buffer(game)
	}

	/// Create a SpriteInstance for s and add it to instances.
	/// Also handles child sprites.
	fn push_sprite(
		&self,
		game: &Game,
		instances: &mut Vec<SpriteInstance>,
		clip_ne: Point2<f32>,
		clip_sw: Point2<f32>,
		s: Sprite,
	) {
		// Position adjusted for parallax
		// TODO: adjust parallax for zoom?
		let pos: Point2<f32> = {
			(Point2 {
				x: s.pos.x,
				y: s.pos.y,
			} - game.camera.pos.to_vec())
				/ s.pos.z
		};
		let texture = self.texture_array.get_sprite_texture(s.texture);

		// Game dimensions of this sprite post-scale.
		// Don't divide by 2, we use this later.
		let height = s.size / s.pos.z;
		let width = height * texture.aspect;

		// Don't draw (or compute matrices for)
		// sprites that are off the screen
		if pos.x < clip_ne.x - width
			|| pos.y > clip_ne.y + height
			|| pos.x > clip_sw.x + width
			|| pos.y < clip_sw.y - height
		{
			return;
		}

		// TODO: clean up
		let scale = height / game.camera.zoom;

		// Note that our mesh starts centered at (0, 0).
		// This is essential---we do not want scale and rotation
		// changing our sprite's position!

		// Apply sprite aspect ratio, preserving height.
		// This must be done *before* rotation.
		//
		// We apply the provided scale here as well as a minor optimization
		let sprite_aspect_and_scale =
			Matrix4::from_nonuniform_scale(texture.aspect * scale, scale, 1.0);

		// Apply rotation
		let rotate = Matrix4::from_angle_z(s.angle);

		// Apply screen aspect ratio, again preserving height.
		// This must be done AFTER rotation... think about it!
		let screen_aspect = Matrix4::from_nonuniform_scale(1.0 / self.window_aspect, 1.0, 1.0);

		// After finishing all ops, translate.
		// This must be done last, all other operations
		// require us to be at (0, 0).
		let translate = Matrix4::from_translation(Vector3 {
			x: pos.x / game.camera.zoom / self.window_aspect,
			y: pos.y / game.camera.zoom,
			z: 0.0,
		});

		// Order matters!
		// The rightmost matrix is applied first.
		let t =
			OPENGL_TO_WGPU_MATRIX * translate * screen_aspect * rotate * sprite_aspect_and_scale;

		instances.push(SpriteInstance {
			transform: t.into(),
			texture_index: texture.index,
		});

		// Add children
		if let Some(children) = s.children {
			for c in children {
				self.push_subsprite(game, instances, c, pos, s.angle);
			}
		}
	}

	/// Add a sprite's subsprite to instance.
	/// Only called by push_sprite.
	fn push_subsprite(
		&self,
		game: &Game,
		instances: &mut Vec<SpriteInstance>,
		s: Sprite,
		parent_pos: Point2<f32>,
		parent_angle: Deg<f32>,
	) {
		// TODO: clean up
		if s.children.is_some() {
			panic!("Child sprites must not have child sprites!")
		}

		let texture = self.texture_array.get_sprite_texture(s.texture);
		let scale = s.size / (s.pos.z * game.camera.zoom);
		let sprite_aspect_and_scale =
			Matrix4::from_nonuniform_scale(texture.aspect * scale, scale, 1.0);
		let rotate = Matrix4::from_angle_z(s.angle);
		let screen_aspect = Matrix4::from_nonuniform_scale(1.0 / self.window_aspect, 1.0, 1.0);

		let ptranslate = Matrix4::from_translation(Vector3 {
			x: parent_pos.x / game.camera.zoom / self.window_aspect,
			y: parent_pos.y / game.camera.zoom,
			z: 0.0,
		});
		let protate = Matrix4::from_angle_z(parent_angle);

		let translate = Matrix4::from_translation(Vector3 {
			x: s.pos.x / game.camera.zoom / self.window_aspect,
			y: s.pos.y / game.camera.zoom,
			z: 0.0,
		});

		// Order matters!
		// The rightmost matrix is applied first.
		let t = OPENGL_TO_WGPU_MATRIX
			* ptranslate * screen_aspect
			* protate * translate
			* rotate * sprite_aspect_and_scale;

		instances.push(SpriteInstance {
			transform: t.into(),
			texture_index: texture.index,
		});
	}

	/// Make a SpriteInstance for each of the game's visible sprites.
	/// Will panic if SPRITE_INSTANCE_LIMIT is exceeded.
	///
	/// This is only called inside self.render()
	fn make_sprite_instances(&self, game: &Game) -> Vec<SpriteInstance> {
		let mut instances: Vec<SpriteInstance> = Vec::new();

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

		for s in game.get_sprites() {
			self.push_sprite(game, &mut instances, clip_ne, clip_sw, s);
		}

		// Enforce sprite limit
		if instances.len() as u64 > SPRITE_INSTANCE_LIMIT {
			// TODO: no panic, handle this better.
			panic!("Sprite limit exceeded!")
		}

		return instances;
	}

	/// Make a StarfieldInstance for each star that needs to be drawn.
	/// Will panic if STARFIELD_INSTANCE_LIMIT is exceeded.
	///
	/// Starfield data rarely changes, so this is called only when it's needed.
	pub fn update_starfield_buffer(&mut self, game: &Game) {
		let sz = consts::STARFIELD_SIZE as f32;

		// Compute window size in starfield tiles
		let mut nw_tile: Point2<i32> = {
			// Game coordinates (relative to camera) of nw corner of screen.
			let clip_nw = Point2::from((self.window_aspect, 1.0)) * consts::ZOOM_MAX;

			// Parallax correction.
			// Also, adjust v for mod to work properly
			// (v is centered at 0)
			let v: Point2<f32> = clip_nw * consts::STARFIELD_Z_MIN;
			let v_adj: Point2<f32> = (v.x + (sz / 2.0), v.y + (sz / 2.0)).into();

			#[rustfmt::skip]
			// Compute m = fmod(x, sz)
			let m: Vector2<f32> = (
				(v_adj.x - (v_adj.x / sz).floor() * sz) - (sz / 2.0),
				(v_adj.y - (v_adj.y / sz).floor() * sz) - (sz / 2.0)
			).into();

			// Now, remainder and convert to "relative tile" coordinates
			// ( where (0,0) is center tile, (0, 1) is north, etc)
			let rel = (v - m) / sz;

			// relative coordinates of north-east tile
			(rel.x.round() as i32, rel.y.round() as i32).into()
		};

		// We need to cover the window with stars,
		// but we also need a one-wide buffer to account for motion.
		nw_tile += Vector2::from((1, 1));

		// Truncate tile grid to buffer size
		// (The window won't be full of stars if our instance limit is too small)
		while ((nw_tile.x * 2 + 1) * (nw_tile.y * 2 + 1) * consts::STARFIELD_COUNT as i32)
			> STARFIELD_INSTANCE_LIMIT as i32
		{
			nw_tile -= Vector2::from((1, 1));
		}

		// Add all tiles to buffer
		let mut instances = Vec::new();
		for x in (-nw_tile.x)..=nw_tile.x {
			for y in (-nw_tile.y)..=nw_tile.y {
				let offset = Vector3 {
					x: sz * x as f32,
					y: sz * y as f32,
					z: 0.0,
				};
				for s in &game.system.starfield {
					instances.push(StarfieldInstance {
						position: (s.pos + offset).into(),
						size: s.size,
						tint: s.tint.into(),
					})
				}
			}
		}

		// Enforce starfield limit
		if instances.len() as u64 > STARFIELD_INSTANCE_LIMIT {
			unreachable!("Starfield limit exceeded!")
		}

		self.starfield_count = instances.len() as u32;
		self.queue.write_buffer(
			&self.vertex_buffers.starfield.instances,
			0,
			bytemuck::cast_slice(&instances),
		);
	}

	pub fn render(&mut self, game: &Game) -> Result<(), wgpu::SurfaceError> {
		let output = self.surface.get_current_texture()?;
		let view = output
			.texture
			.create_view(&wgpu::TextureViewDescriptor::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,
		});

		// Update global values
		self.queue.write_buffer(
			&self.global_data.buffer,
			0,
			bytemuck::cast_slice(&[GlobalDataContent {
				camera_position: game.camera.pos.into(),
				camera_zoom: [game.camera.zoom, 0.0],
				camera_zoom_limits: [consts::ZOOM_MIN, consts::ZOOM_MAX],
				window_size: [
					self.window_size.width as f32,
					self.window_size.height as f32,
				],
				window_aspect: [self.window_aspect, 0.0],
				starfield_texture: [self.texture_array.get_starfield_texture().index, 0],
				starfield_tile_size: [consts::STARFIELD_SIZE as f32, 0.0],
				starfield_size_limits: [consts::STARFIELD_SIZE_MIN, consts::STARFIELD_SIZE_MAX],
			}]),
		);

		// Create sprite instances
		let sprite_instances = self.make_sprite_instances(game);
		self.queue.write_buffer(
			&self.vertex_buffers.sprite.instances,
			0,
			bytemuck::cast_slice(&sprite_instances),
		);

		// 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.global_data.bind_group, &[]);

		// Starfield pipeline
		self.vertex_buffers.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.starfield_count);

		// Sprite pipeline
		self.vertex_buffers.sprite.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.sprite_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..sprite_instances.len() as _,
		);

		// begin_render_pass borrows encoder mutably, so we can't call finish()
		// without dropping this variable.
		drop(render_pass);

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

		Ok(())
	}
}