blonkus/fw/data/model.rs

use crate::{
	data::{read, vertex::Vertex},
	ffi,
	math::*,
	types::{Cast, StkStr},
};
use std::{
	collections::HashMap,
	io::{self, Read, Seek, SeekFrom},
	ops::Range,
	str::Utf8Error,
};

#[derive(thiserror::Error, Debug)]
#[non_exhaustive]
pub enum Err {
	#[error(transparent)]
	Io(#[from] io::Error),
	#[error("Bad magic number")]
	Magic,
	#[error("Incorrect number of poses")]
	NumPoses,
	#[error("Bad text offset")]
	TextOffset,
	#[error(transparent)]
	Utf8(#[from] Utf8Error),
	#[error("Bad parent joint")]
	Parent,
	#[error("Bad frame count")]
	FrameCount,
	#[error("Unknown vertex array type")]
	VaType,
	#[error("Unsupported vertex array format")]
	VaFormat,
}

struct Pose {
	par: usize,
	msk: u32,
	ofs: [f32; 10],
	scl: [f32; 10],
}

struct Joint {
	#[allow(dead_code)]
	nam: StkStr,
	mat: Mat4,
	inv: Mat4,
}

pub struct Mesh {
	mat: StkStr,
	vtx: Range<usize>,
	idx: Range<usize>,
}

/// Data to be uploaded to the GPU. Separate from Model so that it may
/// be dropped after upload.
pub struct ModelData {
	pub vtx: Vec<Vertex>,
	pub idx: Vec<u32>,
}

/// Run-time model data. Anything necessary for passing commands to
/// the GPU.
pub struct Model {
	meshes: HashMap<StkStr, Mesh>,
	frames: Vec<Vec<Mat4>>,
	num_frms: usize,
}

const FMT_I8: u32 = 0;
const FMT_U8: u32 = 1;
const FMT_I16: u32 = 2;
const FMT_U16: u32 = 3;
const FMT_I32: u32 = 4;
const FMT_U32: u32 = 5;
const FMT_F16: u32 = 6;
const FMT_F32: u32 = 7;
const FMT_F64: u32 = 8;

const fn fmt_check(fm: u32) -> Result<(), Err> {
	match fm {
		| FMT_I8 | FMT_U8 | FMT_I16 | FMT_U16 | FMT_I32 | FMT_U32 | FMT_F16
		| FMT_F32 | FMT_F64 => Ok(()),
		| _ => Err(Err::VaFormat),
	}
}

fn fmt_size(fm: u32) -> usize {
	match fm {
		| FMT_I8 | FMT_U8 => 1,
		| FMT_I16 | FMT_U16 | FMT_F16 => 2,
		| FMT_I32 | FMT_U32 | FMT_F32 => 4,
		| FMT_F64 => 8,
		| _ => unsafe { std::hint::unreachable_unchecked() },
	}
}

fn fmt_read_abs(fm: u32, vec: &[u8], pos: usize) -> f32 {
	match fm {
		| FMT_I8 => read::i8x(vec, pos).into(),
		| FMT_U8 => read::u8x(vec, pos).into(),
		| FMT_I16 => read::i16le(vec, pos).into(),
		| FMT_U16 => read::u16le(vec, pos).into(),
		| FMT_I32 => read::i32le(vec, pos).cast(),
		| FMT_U32 => read::u32le(vec, pos).cast(),
		| FMT_F16 => read::f16le(vec, pos).to_f32(),
		| FMT_F32 => read::f32le(vec, pos),
		| FMT_F64 => read::f64le(vec, pos) as f32,
		| _ => unsafe { std::hint::unreachable_unchecked() },
	}
}

fn fmt_read_nrm(fm: u32, vec: &[u8], pos: usize) -> f32 {
	match fm {
		| FMT_I8 => f32::from(read::i8x(vec, pos)).abs() / 128.0,
		| FMT_U8 => f32::from(read::u8x(vec, pos)) / 255.0,
		| FMT_I16 => f32::from(read::i16le(vec, pos)).abs() / 32768.0,
		| FMT_U16 => f32::from(read::u16le(vec, pos)) / 65535.0,
		| FMT_I32 => (f64::from(read::i32le(vec, pos)) / 2147483648.0) as f32,
		| FMT_U32 => (f64::from(read::u32le(vec, pos)) / 4294967295.0) as f32,
		| FMT_F16 => f32::clamp(read::f16le(vec, pos).to_f32(), 0.0, 1.0),
		| FMT_F32 => f32::clamp(read::f32le(vec, pos), 0.0, 1.0),
		| FMT_F64 => f32::clamp(read::f64le(vec, pos) as f32, 0.0, 1.0),
		| _ => unsafe { std::hint::unreachable_unchecked() },
	}
}

impl Mesh {
	pub fn get_mat(&self) -> &str {
		&self.mat
	}

	pub const fn vtx_beg(&self) -> usize {
		self.vtx.start
	}

	pub const fn vtx_end(&self) -> usize {
		self.vtx.end
	}

	pub const fn vtx_len(&self) -> usize {
		self.vtx.end - self.vtx.start
	}

	pub const fn idx_beg(&self) -> usize {
		self.idx.start
	}

	pub const fn idx_end(&self) -> usize {
		self.idx.end
	}

	pub const fn idx_len(&self) -> usize {
		self.idx.end - self.idx.start
	}
}

impl Model {
	pub fn read<R>(rd: &mut R) -> Result<(Self, ModelData), Err>
	where
		R: Read + Seek,
	{
		fn from_stab(stab: &[u8], strn: usize) -> Result<&str, Err> {
			let strn = stab.get(strn..).ok_or(Err::TextOffset)?;
			ffi::cstr_from_slice(strn).ok_or(Err::TextOffset)
		}

		let mut head = [0; 124];
		rd.read_exact(&mut head)?;

		if &head[0..20] != b"INTERQUAKEMODEL\0\x02\0\0\0" {
			return Err(Err::Magic);
		}

		let num_text = read::u32le(&head, 28) as usize;
		let ofs_text = read::u32le(&head, 32).into();
		let num_mesh = read::u32le(&head, 36) as usize;
		let ofs_mesh = read::u32le(&head, 40).into();
		let num_vinf = read::u32le(&head, 44) as usize;
		let num_vert = read::u32le(&head, 48) as usize;
		let ofs_vinf = read::u32le(&head, 52).into();
		let num_tris = read::u32le(&head, 56) as usize;
		let ofs_tris = read::u32le(&head, 60).into();
		let num_join = read::u32le(&head, 68) as usize;
		let ofs_join = read::u32le(&head, 72).into();
		let num_pose = read::u32le(&head, 76) as usize;
		let ofs_pose = read::u32le(&head, 80).into();
		//let num_anim = read::u32le(&head, 84) as usize;
		//let ofs_anim = read::u32le(&head, 88).into();
		let num_frms = read::u32le(&head, 92) as usize;
		let num_frmc = read::u32le(&head, 96) as usize;
		let ofs_frms = read::u32le(&head, 100).into();
		//let ofs_bnds = read::u32le(&head, 104).into();

		if num_join != 0 && num_pose != 0 && num_pose != num_join {
			return Err(Err::NumPoses);
		}

		// collect vertex info
		let mut pos_v = Vec::with_capacity(num_vert);
		let mut tex_v = Vec::with_capacity(num_vert);
		let mut nrm_v = Vec::with_capacity(num_vert);
		let mut tan_v = Vec::with_capacity(num_vert);
		let mut idx_v = Vec::with_capacity(num_vert);
		let mut wgt_v = Vec::with_capacity(num_vert);
		let mut clr_v = Vec::with_capacity(num_vert);

		rd.seek(SeekFrom::Start(ofs_vinf))?;

		for array_def in read::hunk(rd, num_vinf * 20)?.chunks(20) {
			let ty = read::u32le(array_def, 0);
			//let fl = read::u32le(array_def, 4);
			let fm = read::u32le(array_def, 8);
			let sz = read::u32le(array_def, 12) as usize;
			let of = read::u32le(array_def, 16).into();

			fmt_check(fm)?;

			rd.seek(SeekFrom::Start(of))?;

			let obj_size = fmt_size(fm);
			let obj_blks = obj_size * sz;

			let hunk = read::hunk(rd, obj_blks * num_vert)?;

			let v2 = |v: &mut Vec<Vec3>, rd: fn(u32, &[u8], usize) -> f32| {
				if sz != 2 {
					Err(Err::VaFormat)
				} else {
					for vec in hunk.chunks(obj_blks) {
						let x = rd(fm, vec, 0);
						let y = rd(fm, vec, obj_size);
						v.push(Vec3::new(x, y, 0.0));
					}
					Ok(())
				}
			};

			let v3 = |v: &mut Vec<Vec3>, rd: fn(u32, &[u8], usize) -> f32| {
				if sz != 3 {
					Err(Err::VaFormat)
				} else {
					for vec in hunk.chunks(obj_blks) {
						let x = rd(fm, vec, 0);
						let y = rd(fm, vec, obj_size);
						let z = rd(fm, vec, obj_size * 2);
						v.push(Vec3::new(x, y, z));
					}
					Ok(())
				}
			};

			let v4 = |v: &mut Vec<Vec4>, rd: fn(u32, &[u8], usize) -> f32| {
				if sz != 4 {
					Err(Err::VaFormat)
				} else {
					for vec in hunk.chunks(obj_blks) {
						let x = rd(fm, vec, 0);
						let y = rd(fm, vec, obj_size);
						let z = rd(fm, vec, obj_size * 2);
						let w = rd(fm, vec, obj_size * 3);
						v.push(Vec4::new(x, y, z, w));
					}
					Ok(())
				}
			};

			match ty {
				| 0 => v3(&mut pos_v, fmt_read_abs)?,
				| 1 => v2(&mut tex_v, fmt_read_abs)?,
				| 2 => v3(&mut nrm_v, fmt_read_abs)?,
				| 3 => v4(&mut tan_v, fmt_read_abs)?,
				| 4 => v4(&mut idx_v, fmt_read_abs)?,
				| 5 => v4(&mut wgt_v, fmt_read_nrm)?,
				| 6 => v4(&mut clr_v, fmt_read_nrm)?,
				| _ => return Err(Err::VaType),
			}
		}

		// zip the vertex info
		let mut vert_dat = Vec::with_capacity(num_vert);

		for (pos, tex, nrm, tan, idx, wgt, clr) in crate::iter::OpenZip((
			pos_v.into_iter(),
			tex_v.into_iter(),
			nrm_v.into_iter(),
			tan_v.into_iter(),
			idx_v.into_iter(),
			wgt_v.into_iter(),
			clr_v.into_iter(),
		)) {
			vert_dat.push(Vertex {
				pos: pos.unwrap_or(Vec3::ZERO),
				tex: tex.unwrap_or(Vec3::ZERO),
				nrm: nrm.unwrap_or(Vec3::ZERO),
				tan: tan.unwrap_or(Vec4::ZERO),
				idx: idx.unwrap_or(Vec4::ZERO),
				wgt: wgt.unwrap_or(Vec4::ZERO),
				clr: clr.unwrap_or(Vec4::ONE),
			});
		}

		// collect index info
		let mut indx_dat = Vec::with_capacity(num_tris * 3);

		rd.seek(SeekFrom::Start(ofs_tris))?;

		for tri in read::hunk(rd, num_tris * 4 * 3)?.chunks(4 * 3) {
			indx_dat.push(read::u32le(tri, 0));
			indx_dat.push(read::u32le(tri, 4));
			indx_dat.push(read::u32le(tri, 8));
		}

		// collect text
		rd.seek(SeekFrom::Start(ofs_text))?;

		let stab = read::hunk(rd, num_text)?;

		// collect mesh info
		let mut meshes = HashMap::with_capacity(num_mesh);

		rd.seek(SeekFrom::Start(ofs_mesh))?;

		for mesh in read::hunk(rd, num_mesh * 4 * 6)?.chunks(4 * 6) {
			let nam = read::u32le(mesh, 0) as usize;
			let mat = read::u32le(mesh, 4) as usize;
			let vtx = read::u32le(mesh, 8) as usize;
			let vnu = read::u32le(mesh, 12) as usize;
			let idx = read::u32le(mesh, 16) as usize * 3;
			let inu = read::u32le(mesh, 20) as usize * 3;

			let nam = StkStr::new(from_stab(&stab, nam)?);
			let mat = StkStr::new(from_stab(&stab, mat)?);

			let vtx = vtx..vtx + vnu;
			let idx = idx..idx + inu;

			meshes.insert(nam, Mesh { mat, vtx, idx });
		}

		// collect joint info
		let mut joints: Vec<Joint> = Vec::with_capacity(num_join);

		rd.seek(SeekFrom::Start(ofs_join))?;

		for join in read::hunk(rd, num_join * 4 * 12)?.chunks(4 * 12) {
			let nam = read::u32le(join, 0) as usize;
			let par = read::u32le(join, 4) as usize;
			let trs = read::array(read::f32le, join, [0.0; 10], 4, 8);

			let nam = StkStr::new(from_stab(&stab, nam)?);

			let xlt = Vec3::new(trs[0], trs[1], trs[2]);
			let rot = Quat::from_xyzw(trs[3], trs[4], trs[5], trs[6]).normalize();
			let scl = Vec3::new(trs[7], trs[8], trs[9]);

			let mut mat = Mat4::from_scale_rotation_translation(scl, rot, xlt);
			let mut inv = mat.inverse();

			if par & 0x8000_0000 == 0 {
				let parent = joints.get(par).ok_or(Err::Parent)?;
				mat = parent.mat * mat;
				inv = inv * parent.inv;
			}

			joints.push(Joint { nam, mat, inv });
		}

		// collect frames
		let mut frames = Vec::with_capacity(num_frms * num_frmc);

		rd.seek(SeekFrom::Start(ofs_frms))?;

		for samp in read::hunk(rd, num_frms * num_frmc * 2)?.chunks(2) {
			frames.push(read::u16le(samp, 0));
		}

		// collect pose info
		let mut poses = Vec::with_capacity(num_pose);

		rd.seek(SeekFrom::Start(ofs_pose))?;

		for pose in read::hunk(rd, num_pose * 4 * 22)?.chunks(4 * 22) {
			let par = read::u32le(pose, 0) as usize;
			let msk = read::u32le(pose, 4);
			let ofs = read::array(read::f32le, pose, [0.0; 10], 4, 8);
			let scl = read::array(read::f32le, pose, [0.0; 10], 4, 48);

			poses.push(Pose { par, msk, ofs, scl });
		}

		// calculate frame matrices
		let mut frame = frames.iter();
		let mut frames = Vec::with_capacity(num_frms);

		for _ in 0..num_frms {
			let mut frame_mats = Vec::with_capacity(num_pose);

			for (pose, joint) in poses.iter().zip(joints.iter()) {
				let mut ofs = pose.ofs;

				for (i, off) in ofs.iter_mut().enumerate() {
					if pose.msk & 1 << i != 0 {
						let frm = f32::from(*frame.next().ok_or(Err::FrameCount)?);
						*off += frm * pose.scl[i];
					}
				}

				let xlt = Vec3::new(ofs[0], ofs[1], ofs[2]);
				let rot =
					Quat::from_xyzw(ofs[3], ofs[4], ofs[5], ofs[6]).normalize();
				let scl = Vec3::new(ofs[7], ofs[8], ofs[9]);

				let mut mat = Mat4::from_scale_rotation_translation(scl, rot, xlt);

				if pose.par & 0x8000_0000 == 0 {
					let parent = joints.get(pose.par).ok_or(Err::Parent)?;
					mat = parent.mat * mat * joint.inv;

					let parent = frame_mats.get(pose.par).ok_or(Err::Parent)?;
					mat = *parent * mat;
				} else {
					mat = mat * joint.inv;
				}

				frame_mats.push(mat);
			}

			frames.push(frame_mats);
		}

		Ok((
			Self { meshes, frames, num_frms },
			ModelData { vtx: vert_dat, idx: indx_dat },
		))
	}

	pub fn frame(&self, joint: usize, frame: usize) -> Option<Mat4> {
		Some(*self.frames.get(frame)?.get(joint)?)
	}

	pub const fn num_frames(&self) -> usize {
		self.num_frms
	}
}

impl std::ops::Deref for Model {
	type Target = HashMap<StkStr, Mesh>;

	fn deref(&self) -> &Self::Target {
		&self.meshes
	}
}

// EOF