diff options
Diffstat (limited to 'crates/fparkan-assets/src/research.rs')
| -rw-r--r-- | crates/fparkan-assets/src/research.rs | 1151 |
1 files changed, 1151 insertions, 0 deletions
diff --git a/crates/fparkan-assets/src/research.rs b/crates/fparkan-assets/src/research.rs new file mode 100644 index 0000000..b261928 --- /dev/null +++ b/crates/fparkan-assets/src/research.rs @@ -0,0 +1,1151 @@ +use crate::{decode_nres_payload, NresDocument}; +use fparkan_binary::{checked_allocation_len, checked_count_bytes, Cursor, DecodeError, Limits}; +use fparkan_nres::{EntryId, NresEntry}; +use fparkan_path::NormalizedPath; +use fparkan_prototype::{ + decode_unit_dat, unit_component_tree, PrototypeError, UnitComponentRecord, + UnitComponentTreeNode, UnitDat, +}; +use fparkan_vfs::Vfs; +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +const TRF0: u32 = u32::from_le_bytes(*b"TRF0"); +const TRF1: u32 = u32::from_le_bytes(*b"TRF1"); +const TRF2: u32 = u32::from_le_bytes(*b"TRF2"); +const TRF3: u32 = u32::from_le_bytes(*b"TRF3"); +const TRF4: u32 = u32::from_le_bytes(*b"TRF4"); +const TRF5: u32 = u32::from_le_bytes(*b"TRF5"); +const TRF6: u32 = u32::from_le_bytes(*b"TRF6"); +const TRF7: u32 = u32::from_le_bytes(*b"TRF7"); +const TRF8: u32 = u32::from_le_bytes(*b"TRF8"); +const TRF9: u32 = u32::from_le_bytes(*b"TRF9"); +const TRFA: u32 = u32::from_le_bytes(*b"TRFA"); +const TRFB: u32 = u32::from_le_bytes(*b"TRFB"); +const UNIT_DAT_MAGIC: u32 = 0x0000_F0F1; +const UNIT_DAT_HEADER_BYTES: usize = 8; +const UNIT_DAT_RECORD_BYTES: usize = 112; +const INITIAL_RESEARCH_COST_BOUND: f32 = f32::from_bits(0x4b18_967f); + +const RESEARCH_CHUNKS: [u32; 12] = [ + TRF0, TRF1, TRF2, TRF3, TRF4, TRF5, TRF6, TRF7, TRF8, TRF9, TRFA, TRFB, +]; + +/// Ore and power outputs from one native CalcFullResearchCost call. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct ResearchCost { + /// First native output, labeled ore by the Calc debug string. + pub ore: f32, + /// Second native output, labeled power by the Calc debug string. + pub power: f32, +} + +/// One decoded MisLoad research table and its original NRes archive. +#[derive(Debug)] +pub struct ResearchTrfDocument { + archive: NresDocument, + records: Vec<[u32; 10]>, + flags: Vec<u8>, + aliases: Vec<ResearchAlias>, + name_pool: Vec<u8>, + relation_2_3: ResearchRelations, + relation_4_5: ResearchRelations, +} + +/// One ordered research candidate and its native per-entry status byte. +#[derive(Clone, Debug)] +pub struct ResearchCandidate { + /// Unit DAT path read through the caller's VFS. + pub path: NormalizedPath, + /// Native status: zero entries are considered by the next refresh. + pub status: u8, +} + +/// Mutable candidate catalog and selected-index state for the research chooser. +#[derive(Clone, Debug, Default)] +pub struct ResearchCandidateSet { + /// Candidate entries in native supplied order. + pub entries: Vec<ResearchCandidate>, + /// Selected candidate index, preserved when no lower-cost winner exists. + pub selected_index: u32, + /// Whether the candidate catalog has been loaded. + pub loaded: bool, + /// Native candidate-set scan status byte. + pub status: u8, +} + +impl ResearchCandidateSet { + /// Refreshes this set's native selector state from the supplied research table. + /// + /// # Errors + /// + /// Returns explicit errors for unsupported lookup semantics or malformed + /// research-table fields reached by the predicates. + /// Missing paths, invalid candidate DAT framing and an empty-table cost + /// failure skip only that entry. + pub fn refresh_selection( + &mut self, + vfs: &dyn Vfs, + research: Option<&ResearchTrfDocument>, + ) -> Result<(), ResearchError> { + let Some(research) = research else { + return Ok(()); + }; + if !self.loaded { + return Ok(()); + } + + self.status = 1; + let mut running_bound = INITIAL_RESEARCH_COST_BOUND; + for index in 0..self.entries.len() { + if self.entries[index].status != 0 { + continue; + } + + self.status = 0; + let Ok(bytes) = vfs.read(&self.entries[index].path) else { + continue; + }; + let Ok((candidate, tree)) = decode_candidate_first_tree(&bytes) else { + continue; + }; + + let first = first_research_predicate(&candidate, &tree, research)?; + self.entries[index].status = u8::from(first); + if !first { + continue; + } + + let second = second_research_predicate(&candidate, &tree, research)?; + self.entries[index].status = u8::from(second); + let cost = match calc_full_research_cost(&candidate, research) { + Ok(cost) => cost, + Err(ResearchError::EmptyResearchTable) => continue, + Err(error) => return Err(error), + }; + if !second && cost.ore < running_bound { + let selected_index = u32::try_from(index) + .map_err(|_| ResearchError::Binary(DecodeError::IntegerOverflow))?; + running_bound = cost.ore; + self.selected_index = selected_index; + } + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +struct ResearchAlias { + name_range: std::ops::Range<usize>, + row_index: u16, +} + +#[derive(Debug)] +struct ResearchRelations { + table_id: u32, + group_starts: Vec<usize>, + indices: Vec<u32>, +} + +#[derive(Debug)] +struct Chunk<'a> { + entry: &'a NresEntry, + payload: &'a [u8], +} + +/// TRF decoding and cost-calculation error. +#[derive(Debug)] +pub enum ResearchError { + /// The outer NRes archive is malformed. + Nres(fparkan_nres::NresError), + /// A bounded TRF field is truncated or overflows. + Binary(DecodeError), + /// A native-required TRF chunk is absent. + MissingChunk { type_id: u32 }, + /// A recognized TRF chunk occurs more than once. + DuplicateChunk { type_id: u32 }, + /// The native loader requires TRF0 attribute 2 to equal three. + InvalidTrf0Attribute { got: u32 }, + /// A TRF1 payload is too large for the native zero-initialized flag buffer. + FlagsTooLong { row_count: usize, actual: usize }, + /// A TRFB pair points outside the TRF6 name pool. + InvalidNameOffset { pair_index: usize, offset: u16 }, + /// A TRF6 string does not end in NUL before the pool ends. + UnterminatedName { pair_index: usize, offset: u16 }, + /// A TRF0 row's low 16-bit pair link is outside TRFB. + InvalidRecordPair { + record_index: usize, + pair_index: u16, + pair_count: usize, + }, + /// A key or relation resolves to a research row outside TRF0. + InvalidResearchIndex { + type_id: u32, + index: u32, + row_count: usize, + }, + /// A relation table has no group at this research row. + MissingRelationGroup { + type_id: u32, + index: usize, + group_count: usize, + }, + /// A required non-empty relation stream is absent. + MissingRelationStream { type_id: u32, required_ids: u64 }, + /// A candidate Unit DAT resource field is not NUL-terminated. + MissingCandidateKeyTerminator { component_index: usize }, + /// No TRFB key matches an enabled candidate component. + UnknownResearchKey { + component_index: usize, + key: Vec<u8>, + }, + /// A non-ASCII lookup would require an unverified locale-specific fold. + UnsupportedNonAsciiLookup { + query: Vec<u8>, + alias: Option<Vec<u8>>, + }, + /// The research table has no rows. + EmptyResearchTable, + /// The candidate Unit DAT has no root component. + EmptyCandidateTree, + /// The existing Unit DAT tree builder rejected the child-count framing. + CandidateTree(PrototypeError), +} + +impl fmt::Display for ResearchError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Nres(source) => write!(formatter, "{source}"), + Self::Binary(source) => write!(formatter, "{source}"), + Self::MissingChunk { type_id } => { + write!(formatter, "research NRes is missing chunk {type_id:#010x}") + } + Self::DuplicateChunk { type_id } => { + write!(formatter, "research NRes has duplicate chunk {type_id:#010x}") + } + Self::InvalidTrf0Attribute { got } => { + write!(formatter, "TRF0 attribute 2 is {got}, expected 3") + } + Self::FlagsTooLong { row_count, actual } => write!( + formatter, + "TRF1 has {actual} flags, exceeding the {row_count}-byte native buffer" + ), + Self::InvalidNameOffset { pair_index, offset } => write!( + formatter, + "TRFB pair {pair_index} name offset {offset} is outside TRF6" + ), + Self::UnterminatedName { pair_index, offset } => write!( + formatter, + "TRFB pair {pair_index} name at TRF6 offset {offset} is not NUL-terminated" + ), + Self::InvalidRecordPair { + record_index, + pair_index, + pair_count, + } => write!( + formatter, + "TRF0 row {record_index} links to TRFB pair {pair_index}, but there are {pair_count} pairs" + ), + Self::InvalidResearchIndex { + type_id, + index, + row_count, + } => write!( + formatter, + "research index {index} from {type_id:#010x} exceeds TRF0 row count {row_count}" + ), + Self::MissingRelationGroup { + type_id, + index, + group_count, + } => write!( + formatter, + "{type_id:#010x} has no relation group for row {index} (group count {group_count})" + ), + Self::MissingRelationStream { + type_id, + required_ids, + } => write!( + formatter, + "{type_id:#010x} is absent but its count table requires {required_ids} indices" + ), + Self::MissingCandidateKeyTerminator { component_index } => write!( + formatter, + "candidate component {component_index} has no NUL in its 32-byte resource key" + ), + Self::UnknownResearchKey { + component_index, + key, + } => write!( + formatter, + "candidate component {component_index} key {:?} is absent from TRFB", + String::from_utf8_lossy(key) + ), + Self::UnsupportedNonAsciiLookup { query, alias } => { + if let Some(alias) = alias { + write!( + formatter, + "research key comparison between {:?} and {:?} needs unverified non-ASCII case folding", + String::from_utf8_lossy(query), + String::from_utf8_lossy(alias) + ) + } else { + write!( + formatter, + "research key {:?} needs unverified non-ASCII case folding", + String::from_utf8_lossy(query) + ) + } + } + Self::EmptyResearchTable => write!(formatter, "research table has no rows"), + Self::EmptyCandidateTree => write!(formatter, "candidate Unit DAT has no root component"), + Self::CandidateTree(source) => write!(formatter, "{source}"), + } + } +} + +impl std::error::Error for ResearchError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Nres(source) => Some(source), + Self::Binary(source) => Some(source), + Self::CandidateTree(source) => Some(source), + Self::MissingChunk { .. } + | Self::DuplicateChunk { .. } + | Self::InvalidTrf0Attribute { .. } + | Self::FlagsTooLong { .. } + | Self::InvalidNameOffset { .. } + | Self::UnterminatedName { .. } + | Self::InvalidRecordPair { .. } + | Self::InvalidResearchIndex { .. } + | Self::MissingRelationGroup { .. } + | Self::MissingRelationStream { .. } + | Self::MissingCandidateKeyTerminator { .. } + | Self::UnknownResearchKey { .. } + | Self::UnsupportedNonAsciiLookup { .. } + | Self::EmptyResearchTable + | Self::EmptyCandidateTree => None, + } + } +} + +impl From<fparkan_nres::NresError> for ResearchError { + fn from(value: fparkan_nres::NresError) -> Self { + Self::Nres(value) + } +} + +impl From<DecodeError> for ResearchError { + fn from(value: DecodeError) -> Self { + Self::Binary(value) + } +} + +impl From<PrototypeError> for ResearchError { + fn from(value: PrototypeError) -> Self { + Self::CandidateTree(value) + } +} + +impl ResearchTrfDocument { + /// Decodes the native research table chunks from an NRes archive. + /// + /// NRes validates the archive framing and payload ranges. This reader then + /// validates only the TRF bytes needed by its row, key, and relation views. + /// + /// # Errors + /// + /// Returns an error for malformed NRes framing, a missing required chunk, + /// an invalid TRF0 guard, or a referenced byte range outside its chunk. + pub fn decode(bytes: Arc<[u8]>) -> Result<Self, ResearchError> { + let archive = decode_nres_payload(bytes)?; + let mut chunk_ids = HashMap::new(); + for entry in archive.entries() { + let type_id = entry.meta().type_id; + if RESEARCH_CHUNKS.contains(&type_id) && chunk_ids.insert(type_id, entry.id()).is_some() + { + return Err(ResearchError::DuplicateChunk { type_id }); + } + } + + let trf0 = required_chunk(&archive, &chunk_ids, TRF0)?; + if trf0.entry.meta().attr2 != 3 { + return Err(ResearchError::InvalidTrf0Attribute { + got: trf0.entry.meta().attr2, + }); + } + let row_count = trf0.entry.meta().attr1; + let row_bytes = counted_prefix(trf0.payload, u64::from(row_count), 40)?; + let records = read_research_rows(row_bytes, row_count)?; + + let trf1 = required_chunk(&archive, &chunk_ids, TRF1)?; + let flag_count = usize::try_from(row_count).map_err(|_| DecodeError::IntegerOverflow)?; + if trf1.payload.len() > flag_count { + return Err(ResearchError::FlagsTooLong { + row_count: flag_count, + actual: trf1.payload.len(), + }); + } + let mut flags = vec![0; flag_count]; + flags[..trf1.payload.len()].copy_from_slice(trf1.payload); + + let trfb = required_chunk(&archive, &chunk_ids, TRFB)?; + let alias_count = trfb.entry.meta().attr1; + let pair_bytes = counted_prefix(trfb.payload, u64::from(alias_count), 4)?; + let trf6 = required_chunk(&archive, &chunk_ids, TRF6)?; + let name_pool = trf6.payload.to_vec(); + let aliases = read_aliases(pair_bytes, &name_pool, alias_count)?; + + for (record_index, record) in records.iter().enumerate() { + // The native reader casts the packed word to u16 before indexing TRFB. + let pair_index = record[8] as u16; + if usize::from(pair_index) >= aliases.len() { + return Err(ResearchError::InvalidRecordPair { + record_index, + pair_index, + pair_count: aliases.len(), + }); + } + } + + for type_id in [TRF7, TRF8, TRF9, TRFA] { + required_chunk(&archive, &chunk_ids, type_id)?; + } + + let relation_2_3 = read_relations(&archive, &chunk_ids, TRF2, TRF3)?; + let relation_4_5 = read_relations(&archive, &chunk_ids, TRF4, TRF5)?; + + Ok(Self { + archive, + records, + flags, + aliases, + name_pool, + relation_2_3, + relation_4_5, + }) + } + + /// Returns the retained NRes document, including opaque chunks and unused payload tails. + #[must_use] + pub fn archive(&self) -> &NresDocument { + &self.archive + } + + /// Number of TRF0 rows, taken from that chunk's own attr1. + #[must_use] + pub fn row_count(&self) -> usize { + self.records.len() + } + + /// Looks up the first ASCII-case-insensitive TRFB alias in source order. + /// + /// The shipped research keys are ASCII and use ASCII case-insensitive matching. + /// Comparisons involving non-ASCII bytes fail explicitly because native + /// locale-specific folding for those bytes is unknown. + /// + /// # Errors + /// + /// Returns an error if matching requires unverified non-ASCII folding or + /// the selected alias contains an out-of-range row id. + pub fn find_key(&self, key: &[u8]) -> Result<Option<usize>, ResearchError> { + let key = key.split(|byte| *byte == 0).next().unwrap_or_default(); + if !key.is_ascii() { + return Err(ResearchError::UnsupportedNonAsciiLookup { + query: key.to_vec(), + alias: None, + }); + } + for alias in &self.aliases { + let name = &self.name_pool[alias.name_range.clone()]; + if !name.is_ascii() { + return Err(ResearchError::UnsupportedNonAsciiLookup { + query: key.to_vec(), + alias: Some(name.to_vec()), + }); + } + if name.eq_ignore_ascii_case(key) { + let row_index = u32::from(alias.row_index); + if usize::from(alias.row_index) >= self.records.len() { + return Err(ResearchError::InvalidResearchIndex { + type_id: TRFB, + index: row_index, + row_count: self.records.len(), + }); + } + return Ok(Some(usize::from(alias.row_index))); + } + } + Ok(None) + } + + /// Returns the preserved ten raw TRF0 words for one row. + #[must_use] + pub fn row_words(&self, index: usize) -> Option<&[u32; 10]> { + self.records.get(index) + } + + /// Returns the original TRF1 byte for one row. + #[must_use] + pub fn flag_byte(&self, index: usize) -> Option<u8> { + self.flags.get(index).copied() + } + + /// Returns the raw name bytes selected by a row's low-16-bit TRFB link. + /// + /// # Errors + /// + /// Returns an error when the row is outside TRF0. + pub fn row_name(&self, index: usize) -> Result<&[u8], ResearchError> { + let row = self + .records + .get(index) + .ok_or(ResearchError::InvalidResearchIndex { + type_id: TRF0, + index: u32::try_from(index).unwrap_or(u32::MAX), + row_count: self.records.len(), + })?; + let pair_index = usize::from(row[8] as u16); + let alias = self + .aliases + .get(pair_index) + .ok_or(ResearchError::InvalidRecordPair { + record_index: index, + pair_index: row[8] as u16, + pair_count: self.aliases.len(), + })?; + Ok(&self.name_pool[alias.name_range.clone()]) + } + + /// Returns the ordered TRF2/TRF3 indices used by CalcFullResearchCost. + /// + /// # Errors + /// + /// Returns an error if this row has no TRF2 group or its TRF3 stream + /// contains a row index outside TRF0. + pub fn relation_2_3(&self, index: usize) -> Result<&[u32], ResearchError> { + let related = self.relation_2_3.get(index)?; + validate_relation_indices(related, TRF3, self.records.len())?; + Ok(related) + } + + /// Returns the separate ordered TRF4/TRF5 index group. + /// + /// # Errors + /// + /// Returns an error if this row has no TRF4 group or its TRF5 stream + /// contains a row index outside TRF0. + pub fn relation_4_5(&self, index: usize) -> Result<&[u32], ResearchError> { + let related = self.relation_4_5.get(index)?; + validate_relation_indices(related, TRF5, self.records.len())?; + Ok(related) + } + + fn gates(&self, index: usize) -> Result<(u32, u32, u32), ResearchError> { + let flag = self + .flags + .get(index) + .copied() + .ok_or(ResearchError::InvalidResearchIndex { + type_id: TRF1, + index: u32::try_from(index).unwrap_or(u32::MAX), + row_count: self.records.len(), + })?; + Ok(( + u32::from((flag >> 1) & 1), + u32::from(flag & 1), + u32::from((flag >> 2) & 1), + )) + } +} + +pub(super) fn decode_candidate_first_tree( + payload: &[u8], +) -> Result<(UnitDat, Vec<UnitComponentTreeNode>), PrototypeError> { + let limits = Limits::default(); + let payload_len = u64::try_from(payload.len()).map_err(|_| DecodeError::IntegerOverflow)?; + if payload_len > limits.max_file_bytes { + return Err(DecodeError::LimitExceeded { + count: payload_len, + limit: limits.max_file_bytes, + } + .into()); + } + if payload.len() < UNIT_DAT_HEADER_BYTES { + return Err(PrototypeError::InvalidSize); + } + let magic = u32::from_le_bytes( + payload[..4] + .try_into() + .map_err(|_| PrototypeError::InvalidSize)?, + ); + if magic != UNIT_DAT_MAGIC { + return Err(PrototypeError::InvalidUnitDatMagic(magic)); + } + + let available_body = payload.len() - UNIT_DAT_HEADER_BYTES; + let complete_records = available_body / UNIT_DAT_RECORD_BYTES; + if complete_records == 0 { + return Err(PrototypeError::InvalidSize); + } + let maximum_records = + usize::try_from(limits.max_entries).map_err(|_| DecodeError::IntegerOverflow)?; + let scan_records = complete_records.min(maximum_records); + let mut pending_nodes = 1_u64; + let mut first_tree_count = None; + + for record_index in 0..scan_records { + let record_offset = record_index + .checked_mul(UNIT_DAT_RECORD_BYTES) + .and_then(|offset| UNIT_DAT_HEADER_BYTES.checked_add(offset)) + .ok_or(DecodeError::IntegerOverflow)?; + let children_offset = record_offset + .checked_add(108) + .ok_or(DecodeError::IntegerOverflow)?; + let children_end = children_offset + .checked_add(4) + .ok_or(DecodeError::IntegerOverflow)?; + let child_count = u32::from_le_bytes( + payload + .get(children_offset..children_end) + .ok_or(PrototypeError::InvalidSize)? + .try_into() + .map_err(|_| PrototypeError::InvalidSize)?, + ); + pending_nodes = pending_nodes + .checked_sub(1) + .and_then(|pending| pending.checked_add(u64::from(child_count))) + .ok_or(DecodeError::IntegerOverflow)?; + if pending_nodes > u64::from(limits.max_entries) { + return Err(DecodeError::LimitExceeded { + count: pending_nodes, + limit: u64::from(limits.max_entries), + } + .into()); + } + if pending_nodes == 0 { + first_tree_count = Some( + record_index + .checked_add(1) + .ok_or(DecodeError::IntegerOverflow)?, + ); + break; + } + } + + let Some(record_count) = first_tree_count else { + if scan_records == maximum_records { + return Err(DecodeError::LimitExceeded { + count: u64::from(limits.max_entries) + 1, + limit: u64::from(limits.max_entries), + } + .into()); + } + return Err(PrototypeError::InvalidSize); + }; + let record_count = u64::try_from(record_count).map_err(|_| DecodeError::IntegerOverflow)?; + let _record_capacity = checked_allocation_len(record_count, u64::from(limits.max_entries))?; + let record_bytes = checked_count_bytes( + record_count, + UNIT_DAT_RECORD_BYTES as u64, + u64::try_from(available_body).map_err(|_| DecodeError::IntegerOverflow)?, + )?; + let record_size = u64::try_from(std::mem::size_of::<UnitComponentRecord>()) + .map_err(|_| DecodeError::IntegerOverflow)?; + let tree_record_size = std::mem::size_of::<UnitComponentTreeNode>() + .checked_add(2 * std::mem::size_of::<usize>()) + .ok_or(DecodeError::IntegerOverflow)?; + let tree_record_size = + u64::try_from(tree_record_size).map_err(|_| DecodeError::IntegerOverflow)?; + let decoded_bytes = record_size + .checked_add(tree_record_size) + .and_then(|size| size.checked_mul(record_count)) + .ok_or(DecodeError::IntegerOverflow)?; + let _decoded_capacity = checked_allocation_len(decoded_bytes, limits.max_decoded_bytes)?; + let prefix_end = UNIT_DAT_HEADER_BYTES + .checked_add(record_bytes) + .ok_or(DecodeError::IntegerOverflow)?; + + let candidate = decode_unit_dat(&payload[..prefix_end])?; + let tree = unit_component_tree(&candidate.records)?; + Ok((candidate, tree)) +} + +fn candidate_research_predicate( + candidate: &UnitDat, + tree: &[UnitComponentTreeNode], + research: &ResearchTrfDocument, + first_info_word: usize, + second_info_word: usize, +) -> Result<bool, ResearchError> { + let mut stack = vec![0_usize]; + while let Some(node_index) = stack.pop() { + let Some(node) = tree.get(node_index) else { + return Err(ResearchError::CandidateTree( + PrototypeError::InvalidUnitHierarchy(format!( + "component index {node_index} is outside the reconstructed tree" + )), + )); + }; + let Some(component) = candidate.records.get(node.record_index) else { + return Err(ResearchError::CandidateTree( + PrototypeError::InvalidUnitHierarchy(format!( + "tree references missing component {}", + node.record_index + )), + )); + }; + if component.kind == 0 { + continue; + } + + let Some(key_end) = component.resource_raw.iter().position(|byte| *byte == 0) else { + return Err(ResearchError::MissingCandidateKeyTerminator { + component_index: node.record_index, + }); + }; + let Some(research_index) = research.find_key(&component.resource_raw[..key_end])? else { + return Ok(false); + }; + let (word0, word1, word2) = research.gates(research_index)?; + let info_words = [word0, word1, word2]; + if info_words[first_info_word] == 0 || info_words[second_info_word] == 0 { + return Ok(false); + } + for child in node.children.iter().rev() { + stack.push(*child); + } + } + Ok(true) +} + +fn first_research_predicate( + candidate: &UnitDat, + tree: &[UnitComponentTreeNode], + research: &ResearchTrfDocument, +) -> Result<bool, ResearchError> { + candidate_research_predicate(candidate, tree, research, 1, 2) +} + +fn second_research_predicate( + candidate: &UnitDat, + tree: &[UnitComponentTreeNode], + research: &ResearchTrfDocument, +) -> Result<bool, ResearchError> { + candidate_research_predicate(candidate, tree, research, 0, 2) +} + +/// Calculates the ordered ore and power cost for one candidate Unit DAT tree. +/// +/// # Errors +/// +/// Returns an error for an empty research table, malformed candidate tree, +/// missing or unsupported key comparison, invalid relation group, or invalid +/// row reference. +pub fn calc_full_research_cost( + candidate: &UnitDat, + research: &ResearchTrfDocument, +) -> Result<ResearchCost, ResearchError> { + if research.records.is_empty() { + return Err(ResearchError::EmptyResearchTable); + } + let tree = unit_component_tree(&candidate.records)?; + if tree.is_empty() { + return Err(ResearchError::EmptyCandidateTree); + } + + let mut visited = vec![false; research.records.len()]; + let direct = collect_direct_indices(candidate, &tree, research, &mut visited)?; + let mut closure = Vec::new(); + for &index in &direct { + closure.push(index); + let (_, recurse, _) = research.gates(index)?; + if recurse == 0 { + append_relation_closure(index, research, &mut visited, &mut closure)?; + } + } + + let mut ore = 0.0_f32; + let mut power = 0.0_f32; + for index in closure { + let (included, _, _) = research.gates(index)?; + if included == 0 { + let row = research + .records + .get(index) + .ok_or(ResearchError::InvalidResearchIndex { + type_id: TRF0, + index: u32::try_from(index).unwrap_or(u32::MAX), + row_count: research.records.len(), + })?; + ore = ore + f32::from_bits(row[1]); + power = power + f32::from_bits(row[0]); + } + } + + Ok(ResearchCost { ore, power }) +} + +fn collect_direct_indices( + candidate: &UnitDat, + tree: &[UnitComponentTreeNode], + research: &ResearchTrfDocument, + visited: &mut [bool], +) -> Result<Vec<usize>, ResearchError> { + let mut direct = Vec::new(); + let mut stack = vec![0_usize]; + let mut first_error = None; + + while let Some(node_index) = stack.pop() { + let Some(node) = tree.get(node_index) else { + return Err(ResearchError::CandidateTree( + PrototypeError::InvalidUnitHierarchy(format!( + "component index {node_index} is outside the reconstructed tree" + )), + )); + }; + let Some(component) = candidate.records.get(node.record_index) else { + return Err(ResearchError::CandidateTree( + PrototypeError::InvalidUnitHierarchy(format!( + "tree references missing component {}", + node.record_index + )), + )); + }; + + if component.kind == 0 { + continue; + } + + let key_end = component.resource_raw.iter().position(|byte| *byte == 0); + let key = match key_end { + Some(end) => &component.resource_raw[..end], + None => { + if first_error.is_none() { + first_error = Some(ResearchError::MissingCandidateKeyTerminator { + component_index: node.record_index, + }); + } + for child in node.children.iter().rev() { + stack.push(*child); + } + continue; + } + }; + + match research.find_key(key) { + Ok(Some(index)) => match research.gates(index) { + Ok((included, _, _)) => { + if included == 0 && !visited[index] { + visited[index] = true; + direct.push(index); + } + } + Err(error) => { + if first_error.is_none() { + first_error = Some(error); + } + } + }, + Ok(None) => { + if first_error.is_none() { + first_error = Some(ResearchError::UnknownResearchKey { + component_index: node.record_index, + key: key.to_vec(), + }); + } + } + Err(error) => { + if first_error.is_none() { + first_error = Some(error); + } + } + } + + for child in node.children.iter().rev() { + stack.push(*child); + } + } + + if let Some(error) = first_error { + return Err(error); + } + Ok(direct) +} + +#[derive(Clone, Copy)] +struct RelationFrame { + row_index: usize, + next_child: usize, +} + +fn append_relation_closure( + root: usize, + research: &ResearchTrfDocument, + visited: &mut [bool], + closure: &mut Vec<usize>, +) -> Result<(), ResearchError> { + let mut stack = vec![RelationFrame { + row_index: root, + next_child: 0, + }]; + + while let Some(frame) = stack.last().copied() { + let related = research.relation_2_3.get(frame.row_index)?; + if frame.next_child >= related.len() { + stack.pop(); + continue; + } + let related_index = related[frame.next_child]; + stack + .last_mut() + .expect("non-empty relation stack") + .next_child += 1; + + let index = usize::try_from(related_index).map_err(|_| DecodeError::IntegerOverflow)?; + if index >= research.records.len() { + return Err(ResearchError::InvalidResearchIndex { + type_id: TRF3, + index: related_index, + row_count: research.records.len(), + }); + } + let (included, recurse, _) = research.gates(index)?; + if included != 0 || visited[index] { + continue; + } + + visited[index] = true; + closure.push(index); + if recurse == 0 { + stack.push(RelationFrame { + row_index: index, + next_child: 0, + }); + } + } + + Ok(()) +} + +impl ResearchRelations { + fn get(&self, index: usize) -> Result<&[u32], ResearchError> { + let group_count = self.group_starts.len().saturating_sub(1); + let missing_group = || ResearchError::MissingRelationGroup { + type_id: self.table_id, + index, + group_count, + }; + let Some(end_index) = index.checked_add(1) else { + return Err(missing_group()); + }; + let Some(&start) = self.group_starts.get(index) else { + return Err(missing_group()); + }; + let Some(&end) = self.group_starts.get(end_index) else { + return Err(missing_group()); + }; + Ok(&self.indices[start..end]) + } +} + +fn validate_relation_indices( + related: &[u32], + type_id: u32, + row_count: usize, +) -> Result<(), ResearchError> { + for &index in related { + if usize::try_from(index).map_or(true, |value| value >= row_count) { + return Err(ResearchError::InvalidResearchIndex { + type_id, + index, + row_count, + }); + } + } + Ok(()) +} + +fn required_chunk<'a>( + archive: &'a NresDocument, + chunk_ids: &HashMap<u32, EntryId>, + type_id: u32, +) -> Result<Chunk<'a>, ResearchError> { + let id = chunk_ids + .get(&type_id) + .copied() + .ok_or(ResearchError::MissingChunk { type_id })?; + let entry = archive + .entry(id) + .ok_or(fparkan_nres::NresError::EntryIdOutOfRange { + id: id.0, + entry_count: archive.header().entry_count, + })?; + let payload = archive.payload(id)?; + Ok(Chunk { entry, payload }) +} + +fn optional_chunk<'a>( + archive: &'a NresDocument, + chunk_ids: &HashMap<u32, EntryId>, + type_id: u32, +) -> Result<Option<Chunk<'a>>, ResearchError> { + let Some(id) = chunk_ids.get(&type_id).copied() else { + return Ok(None); + }; + let entry = archive + .entry(id) + .ok_or(fparkan_nres::NresError::EntryIdOutOfRange { + id: id.0, + entry_count: archive.header().entry_count, + })?; + let payload = archive.payload(id)?; + Ok(Some(Chunk { entry, payload })) +} + +fn counted_prefix(payload: &[u8], count: u64, stride: u64) -> Result<&[u8], ResearchError> { + let remaining = u64::try_from(payload.len()).map_err(|_| DecodeError::IntegerOverflow)?; + let size = checked_count_bytes(count, stride, remaining)?; + Ok(&payload[..size]) +} + +fn read_research_rows(bytes: &[u8], row_count: u32) -> Result<Vec<[u32; 10]>, ResearchError> { + let capacity = checked_allocation_len( + u64::from(row_count), + u64::from(Limits::default().max_array_items), + )?; + let mut cursor = Cursor::new(bytes); + let mut records = Vec::with_capacity(capacity); + for _ in 0..row_count { + let mut row = [0_u32; 10]; + for word in &mut row { + *word = cursor.read_u32_le()?; + } + records.push(row); + } + cursor.require_eof()?; + Ok(records) +} + +fn read_aliases( + bytes: &[u8], + name_pool: &[u8], + alias_count: u32, +) -> Result<Vec<ResearchAlias>, ResearchError> { + let capacity = checked_allocation_len( + u64::from(alias_count), + u64::from(Limits::default().max_array_items), + )?; + let mut cursor = Cursor::new(bytes); + let mut aliases = Vec::with_capacity(capacity); + for pair_index in 0..capacity { + let name_offset = cursor.read_u16_le()?; + let row_index = cursor.read_u16_le()?; + let offset = usize::from(name_offset); + let suffix = name_pool + .get(offset..) + .ok_or(ResearchError::InvalidNameOffset { + pair_index, + offset: name_offset, + })?; + let length = + suffix + .iter() + .position(|byte| *byte == 0) + .ok_or(ResearchError::UnterminatedName { + pair_index, + offset: name_offset, + })?; + aliases.push(ResearchAlias { + name_range: offset..offset + length, + row_index, + }); + } + cursor.require_eof()?; + Ok(aliases) +} + +fn read_relations( + archive: &NresDocument, + chunk_ids: &HashMap<u32, EntryId>, + count_type: u32, + index_type: u32, +) -> Result<ResearchRelations, ResearchError> { + let count_chunk = required_chunk(archive, chunk_ids, count_type)?; + let count_count = count_chunk.entry.meta().attr1; + let count_bytes = counted_prefix(count_chunk.payload, u64::from(count_count), 4)?; + let mut cursor = Cursor::new(count_bytes); + let max_items = u64::from(Limits::default().max_array_items); + let group_count = u64::from(count_count); + let offset_count = group_count + .checked_add(1) + .ok_or(DecodeError::IntegerOverflow)?; + let offset_capacity = checked_allocation_len( + offset_count, + max_items + .checked_add(1) + .ok_or(DecodeError::IntegerOverflow)?, + )?; + let mut group_starts = Vec::with_capacity(offset_capacity); + group_starts.push(0_usize); + let mut total_ids = 0_u64; + for _ in 0..count_count { + let count = cursor.read_u32_le()?; + total_ids = total_ids + .checked_add(u64::from(count)) + .ok_or(DecodeError::IntegerOverflow)?; + let next = checked_allocation_len(total_ids, max_items)?; + group_starts.push(next); + } + + let index_chunk = optional_chunk(archive, chunk_ids, index_type)?; + let indices = match index_chunk { + None if total_ids == 0 => Vec::new(), + None => { + return Err(ResearchError::MissingRelationStream { + type_id: index_type, + required_ids: total_ids, + }) + } + Some(chunk) => { + let index_bytes = counted_prefix(chunk.payload, total_ids, 4)?; + let mut cursor = Cursor::new(index_bytes); + let capacity = checked_allocation_len(total_ids, max_items)?; + let mut indices = Vec::with_capacity(capacity); + for _ in 0..total_ids { + indices.push(cursor.read_u32_le()?); + } + cursor.require_eof()?; + indices + } + }; + + Ok(ResearchRelations { + table_id: count_type, + group_starts, + indices, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decoded_array_limits_fail_before_large_allocations() { + let limit = Limits::default().max_array_items; + let too_many = limit.checked_add(1).expect("test count"); + + assert!(matches!( + read_research_rows(&[], too_many), + Err(ResearchError::Binary(DecodeError::LimitExceeded { count, limit: actual_limit })) + if count == u64::from(too_many) && actual_limit == u64::from(limit) + )); + assert!(matches!( + read_aliases(&[], &[], too_many), + Err(ResearchError::Binary(DecodeError::LimitExceeded { count, limit: actual_limit })) + if count == u64::from(too_many) && actual_limit == u64::from(limit) + )); + } +} |