Skip to main content

rgpre/import/
registry.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use walkdir::WalkDir;
4
5use crate::import::FileType;
6
7/// Represents a file entry in the registry with its metadata.
8pub struct FileEntry {
9    pub name: String,
10    pub path: PathBuf,
11    pub file_type: FileType,
12    /// Pre-loaded file bytes. When set, parsers use this instead of reading from `path`.
13    pub data: Option<Vec<u8>>,
14}
15
16impl FileEntry {
17    /// Creates a registry entry for a discovered file.
18    #[must_use]
19    #[allow(clippy::missing_const_for_fn)]
20    // Heap-owned String/PathBuf constructor is runtime-only in this API.
21    pub fn new(name: String, path: PathBuf, file_type: FileType) -> Self {
22        Self {
23            name,
24            path,
25            file_type,
26            data: None,
27        }
28    }
29}
30
31/// Registry for managing files by their model names
32///
33/// The registry uses the model name (file stem without extension) as the key,
34/// allowing easy lookup of files regardless of their file extension or path.
35pub struct Registry {
36    pub root_path: PathBuf,
37    pub files: HashMap<String, FileEntry>,
38}
39
40impl Registry {
41    /// Creates an empty registry rooted at `root_path`.
42    #[must_use]
43    pub fn new(root_path: PathBuf) -> Self {
44        Self {
45            root_path,
46            files: HashMap::new(),
47        }
48    }
49
50    /// Creates a registry from pre-loaded file data (no filesystem access).
51    #[must_use]
52    pub fn from_data(entries: HashMap<String, (Vec<u8>, FileType)>) -> Self {
53        let mut files = HashMap::new();
54        for (name, (data, file_type)) in entries {
55            files.insert(
56                name.clone(),
57                FileEntry {
58                    name,
59                    path: PathBuf::new(),
60                    file_type,
61                    data: Some(data),
62                },
63            );
64        }
65        Self {
66            root_path: PathBuf::new(),
67            files,
68        }
69    }
70
71    /// Returns a sort key used to prioritize duplicate file sources.
72    #[must_use]
73    pub fn source_rank_key(&self, path: &Path) -> (usize, usize, String) {
74        let rel = path.strip_prefix(&self.root_path).unwrap_or(path);
75        let depth = rel.components().count();
76        let normalized = rel.to_string_lossy().to_ascii_lowercase();
77        let source_rank = if normalized.contains("fxart/") || normalized.contains("fxart\\") {
78            0usize
79        } else if normalized.contains("maps/") || normalized.contains("maps\\") {
80            1usize
81        } else if normalized.contains("3dart/") || normalized.contains("3dart\\") {
82            2usize
83        } else {
84            3usize
85        };
86        (source_rank, depth, normalized)
87    }
88
89    /// Add a file to the registry using its path and file type
90    ///
91    /// The file will be indexed by its name (file stem without extension).
92    /// If a file with the same name already exists, it will be overwritten.
93    pub fn add_file<P: AsRef<Path>>(&mut self, path: P, file_type: FileType) {
94        let path_ref = path.as_ref();
95        let numeric_bank = matches!(
96            file_type,
97            FileType::Bsi | FileType::Texture | FileType::Shade | FileType::Haze
98        ) && path_ref
99            .extension()
100            .is_some_and(|extension| extension.to_string_lossy().parse::<u16>().is_ok());
101        let name = (if numeric_bank {
102            path_ref.file_name()
103        } else {
104            path_ref.file_stem()
105        })
106        .unwrap_or_default()
107        .to_string_lossy()
108        .into_owned();
109        let file_entry = FileEntry::new(name.clone(), path_ref.to_path_buf(), file_type);
110        if let Some(existing) = self.files.get(&name) {
111            // When types differ (e.g. ROB vs COL) or same-type duplicates from
112            // different paths, keep both using synthetic keys so all ROBs remain
113            // searchable for embedded segment lookups.
114            if existing.path != file_entry.path
115                && (existing.file_type != file_entry.file_type
116                    || existing.file_type == FileType::Rob)
117            {
118                let mut suffix = 1usize;
119                loop {
120                    let candidate = format!("{name}#{suffix}");
121                    if let std::collections::hash_map::Entry::Vacant(entry) =
122                        self.files.entry(candidate)
123                    {
124                        entry.insert(file_entry);
125                        return;
126                    }
127                    suffix += 1;
128                }
129            }
130
131            let existing_rank = self.source_rank_key(&existing.path);
132            let new_rank = self.source_rank_key(path_ref);
133            if new_rank < existing_rank {
134                self.files.insert(name, file_entry);
135            }
136        } else {
137            self.files.insert(name, file_entry);
138        }
139    }
140
141    /// Get a file entry by its model name
142    #[must_use]
143    pub fn get_file_by_name(&self, name: &str) -> Option<&FileEntry> {
144        self.files.get(name)
145    }
146
147    /// Get a file entry by its path (extracts the name from the path)
148    #[must_use]
149    pub fn get_file_by_path<P: AsRef<Path>>(&self, path: P) -> Option<&FileEntry> {
150        let path_ref = path.as_ref();
151        let name = path_ref.file_stem().unwrap_or_default().to_string_lossy();
152        path_ref
153            .file_name()
154            .and_then(|name| self.files.get(name.to_string_lossy().as_ref()))
155            .or_else(|| self.files.get(name.as_ref()))
156    }
157
158    /// Get all model names in the registry
159    #[must_use]
160    pub fn get_all_names(&self) -> Vec<&String> {
161        self.files.keys().collect()
162    }
163
164    /// Check if a model name exists in the registry
165    #[must_use]
166    pub fn has_model(&self, name: &str) -> bool {
167        self.files.contains_key(name)
168    }
169}
170
171/// Chooses software art when TEXTURE banks are present and TEXBSI banks are absent.
172/// Mixed installs default to Glide; use explicit edition selection to choose software.
173#[must_use]
174pub fn detect_software_assets(path: &Path) -> bool {
175    let mut software = false;
176    for entry in WalkDir::new(path).into_iter().filter_map(Result::ok) {
177        if !entry.file_type().is_file() {
178            continue;
179        }
180        let name = entry.file_name().to_string_lossy().to_ascii_uppercase();
181        if name.starts_with("TEXBSI.") {
182            return false;
183        }
184        software |= name.starts_with("TEXTURE.");
185    }
186    software
187}
188
189/// Recursively scans a directory and indexes recognized Redguard files.
190pub fn scan_dir<P: AsRef<Path>>(path: P) -> Result<Registry, std::io::Error> {
191    let software = detect_software_assets(path.as_ref());
192    scan_dir_for_edition(path, software)
193}
194
195/// Scans one renderer edition; excludes the other renderer's art directory.
196pub fn scan_dir_for_edition<P: AsRef<Path>>(
197    path: P,
198    software: bool,
199) -> Result<Registry, std::io::Error> {
200    let path = path.as_ref();
201    let mut registry = Registry::new(path.to_path_buf());
202
203    for entry in WalkDir::new(path)
204        .into_iter()
205        .filter_entry(|entry| {
206            !entry.file_type().is_dir()
207                || !entry
208                    .file_name()
209                    .to_string_lossy()
210                    .eq_ignore_ascii_case(if software { "fxart" } else { "3dart" })
211                || entry.depth() == 0
212        })
213        .filter_map(Result::ok)
214    {
215        if entry.file_type().is_file()
216            && let Some(file_type) = FileType::from_path(entry.path())
217        {
218            registry.add_file(entry.path(), file_type);
219        }
220    }
221
222    Ok(registry)
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    #[test]
230    fn source_rank_prefers_fxart_over_3dart() {
231        let reg = Registry::new(PathBuf::from("/root"));
232        let fxart = reg.source_rank_key(Path::new("/root/fxart/FOO.ROB"));
233        let maps = reg.source_rank_key(Path::new("/root/maps/FOO.ROB"));
234        let art3d = reg.source_rank_key(Path::new("/root/3dart/FOO.ROB"));
235        assert!(fxart < maps, "fxart should rank before maps");
236        assert!(maps < art3d, "maps should rank before 3dart");
237    }
238
239    #[test]
240    fn duplicate_rob_from_different_dirs_both_kept() {
241        let mut reg = Registry::new(PathBuf::from("/root"));
242        reg.add_file("/root/3dart/FOO.ROB", FileType::Rob);
243        reg.add_file("/root/fxart/FOO.ROB", FileType::Rob);
244        let has_primary = reg.files.contains_key("FOO");
245        let has_dup = reg.files.contains_key("FOO#1");
246        assert!(has_primary, "primary ROB key must exist");
247        assert!(has_dup, "duplicate ROB must be kept under synthetic key");
248    }
249
250    #[test]
251    fn different_type_same_stem_both_kept() {
252        let mut reg = Registry::new(PathBuf::from("/root"));
253        reg.add_file("/root/fxart/ISLAND.ROB", FileType::Rob);
254        reg.add_file("/root/fxart/ISLAND.COL", FileType::Col);
255        let has_primary = reg.files.contains_key("ISLAND");
256        let has_dup = reg.files.contains_key("ISLAND#1");
257        assert!(has_primary, "primary entry must exist");
258        assert!(
259            has_dup,
260            "different-type same-stem must be kept under synthetic key"
261        );
262    }
263
264    #[test]
265    fn same_type_same_path_not_duplicated() {
266        let mut reg = Registry::new(PathBuf::from("/root"));
267        reg.add_file("/root/fxart/FOO.ROB", FileType::Rob);
268        reg.add_file("/root/fxart/FOO.ROB", FileType::Rob);
269        assert!(reg.files.contains_key("FOO"));
270        assert!(
271            !reg.files.contains_key("FOO#1"),
272            "identical path should not create duplicate"
273        );
274    }
275}