rpkg_rs/resource/
partition_manager.rs

1use std::path::{Path, PathBuf};
2
3use itertools::Itertools;
4use thiserror::Error;
5use crate::resource::partition_manager::PartitionManagerError::PartitionNotFound;
6
7use crate::resource::pdefs::{
8    GameDiscoveryError, GamePaths, PackageDefinitionError, PackageDefinitionSource, PartitionId,
9    PartitionInfo,
10};
11use crate::resource::resource_info::ResourceInfo;
12use crate::resource::runtime_resource_id::RuntimeResourceID;
13use crate::WoaVersion;
14
15use super::resource_partition::{PatchId, ResourcePartition, ResourcePartitionError};
16
17#[derive(Debug, Error)]
18pub enum PartitionManagerError {
19    #[error("Cannot use packagedefinition config: {0}")]
20    PackageDefinitionError(#[from] PackageDefinitionError),
21
22    #[error("partition '{0}' error: {1}")]
23    PartitionError(PartitionId, ResourcePartitionError),
24
25    #[error("partition {0} could not be found")]
26    PartitionNotFound(String),
27    
28    #[error("resource {0} could not be found")]
29    ResourceNotFound(String),
30
31    #[error("Could not discover game paths: {0}")]
32    GameDiscoveryError(#[from] GameDiscoveryError),
33    
34    #[error("Could not find a root partition")]
35    NoRootPartition(),
36}
37
38#[allow(dead_code)]
39#[derive(Clone, Debug, Copy)]
40pub struct PartitionState {
41    pub installing: bool,
42    pub mounted: bool,
43    pub install_progress: f32,
44}
45
46pub struct PartitionManager {
47    runtime_directory: PathBuf,
48    partition_infos: Vec<PartitionInfo>, //All potential partitions which could be mounted with this manager
49    pub partitions: Vec<ResourcePartition>, //All mounted partitions
50}
51
52impl PartitionManager {
53    /// Create a new PartitionManager for the game at the given path, and a custom package definition.
54    ///
55    /// # Arguments
56    /// - `runtime_directory` - The path to the game's runtime directory.
57    /// - `package_definition` - The package definition to use.
58    pub fn new(
59        runtime_directory: PathBuf,
60        package_definition: &PackageDefinitionSource,
61    ) -> Result<Self, PartitionManagerError> {
62        let partition_infos = package_definition
63            .read()
64            .map_err(PartitionManagerError::PackageDefinitionError)?;
65
66        Ok(Self {
67            runtime_directory,
68            partition_infos,
69            partitions: vec![],
70        })
71    }
72
73    /// Create a new PartitionManager by mounting the game at the given path.
74    ///
75    /// # Arguments
76    /// - `retail_path` - The path to the game's retail directory.
77    /// - `game_version` - The version of the game.
78    /// - `mount` - Indicates whether to automatically mount the partitions, can eliminate the need to call `mount_partitions` separately
79    pub fn from_game(
80        retail_directory: PathBuf,
81        game_version: WoaVersion,
82        mount: bool,
83    ) -> Result<Self, PartitionManagerError> {
84        Self::from_game_with_callback(retail_directory, game_version, mount, |_, _| {})
85    }
86
87    /// Create a new PartitionManager by mounting the game at the given path.
88    ///
89    /// # Arguments
90    /// - `retail_path` - The path to the game's retail directory.
91    /// - `game_version` - The version of the game.
92    /// - `mount` - Indicates whether to automatically mount the partitions, can eliminate the need to call `mount_partitions` separately
93    /// - `progress_callback` - A callback function that will be called with the current mounting progress.
94    pub fn from_game_with_callback<F>(
95        retail_directory: PathBuf,
96        game_version: WoaVersion,
97        mount: bool,
98        progress_callback: F,
99    ) -> Result<Self, PartitionManagerError>
100    where
101        F: FnMut(usize, &PartitionState),
102    {
103        let game_paths = GamePaths::from_retail_directory(retail_directory)?;
104        let package_definition =
105            PackageDefinitionSource::from_file(game_paths.package_definition_path, game_version)?;
106
107        // And read all the partition infos.
108        let partition_infos = package_definition
109            .read()
110            .map_err(PartitionManagerError::PackageDefinitionError)?;
111
112        let mut package_manager = Self {
113            runtime_directory: game_paths.runtime_path,
114            partition_infos,
115            partitions: vec![],
116        };
117
118        // If the user requested auto mounting, do it.
119        if mount {
120            package_manager.mount_partitions(progress_callback)?;
121        }
122
123        Ok(package_manager)
124    }
125
126    fn try_read_partition<F>(
127        runtime_directory: &Path,
128        partition_info: PartitionInfo,
129        mut progress_callback: F,
130    ) -> Result<Option<ResourcePartition>, PartitionManagerError>
131    where
132        F: FnMut(&PartitionState),
133    {
134        let mut partition = ResourcePartition::new(partition_info.clone());
135        let mut state_result: PartitionState = PartitionState {
136            installing: false,
137            mounted: false,
138            install_progress: 0.0,
139        };
140
141        let callback = |state: &_| {
142            progress_callback(state);
143            state_result = *state;
144        };
145
146        partition
147            .mount_resource_packages_in_partition_with_callback(runtime_directory, callback)
148            .map_err(|e| PartitionManagerError::PartitionError(partition_info.id, e))?;
149
150        if state_result.mounted {
151            Ok(Some(partition))
152        } else {
153            Ok(None)
154        }
155    }
156
157    /// Mount all the partitions in the game.
158    ///
159    /// # Arguments
160    /// - `progress_callback` - A callback function that will be called with the current mounting progress.
161    pub fn mount_partitions<F>(
162        &mut self,
163        mut progress_callback: F,
164    ) -> Result<(), PartitionManagerError>
165    where
166        F: FnMut(usize, &PartitionState),
167    {
168        let partitions = self
169            .partition_infos
170            .iter()
171            .enumerate()
172            .map(|(index, partition_info)| {
173                let callback = |state: &_| {
174                    progress_callback(index + 1, state);
175                };
176
177                Self::try_read_partition(&self.runtime_directory, partition_info.clone(), callback)
178            })
179            .collect::<Result<Vec<Option<ResourcePartition>>, PartitionManagerError>>()?
180            .into_iter()
181            .flatten()
182            .collect::<Vec<ResourcePartition>>();
183
184        for partition in partitions {
185            self.partitions.push(partition);
186        }
187
188        Ok(())
189    }
190
191    /// Mount a single partition in the game.
192    ///
193    /// # Arguments
194    /// - `partition_info` - The partition info to mount.
195    /// - `progress_callback` - A callback function that will be called with the current mounting progress.
196    pub fn mount_partition<F>(
197        &mut self,
198        partition_info: PartitionInfo,
199        progress_callback: F,
200    ) -> Result<(), PartitionManagerError>
201    where
202        F: FnMut(&PartitionState),
203    {
204        if let Some(partition) =
205            Self::try_read_partition(&self.runtime_directory, partition_info, progress_callback)?
206        {
207            self.partitions.push(partition)
208        }
209
210        Ok(())
211    }
212
213    pub fn read_resource_from(
214        &self,
215        partition_id: PartitionId,
216        rrid: RuntimeResourceID,
217    ) -> Result<Vec<u8>, PartitionManagerError> {
218        let partition = self
219            .partitions
220            .iter()
221            .find(|partition| partition.partition_info().id == partition_id);
222
223        if let Some(partition) = partition {
224            match partition.read_resource(&rrid) {
225                Ok(data) => Ok(data),
226                Err(e) => Err(PartitionManagerError::PartitionError(partition_id, e)),
227            }
228        } else {
229            Err(PartitionManagerError::PartitionNotFound(
230                partition_id.to_string(),
231            ))
232        }
233    }
234
235    pub fn find_partition(&self, partition_id: PartitionId) -> Option<&ResourcePartition> {
236        self.partitions
237            .iter()
238            .find(|partition| partition.partition_info().id == partition_id)
239    }
240    
241    pub fn root_partition(
242        &self
243    ) -> Result<PartitionId, PartitionManagerError> {
244        if let Some(mut partition) = self.partition_infos.first(){
245            loop {
246                match &partition.parent{
247                    Some(parent) => {
248                        match self.find_partition(parent.clone()){
249                            Some(part) => {partition = part.partition_info()}
250                            None => {return Err(PartitionNotFound(parent.to_string()))}
251                        };
252                    },
253                    None => return Ok(partition.id.clone()),
254                }
255            }
256        }
257        Err(PartitionManagerError::NoRootPartition())
258    }
259    
260    pub fn partitions_with_resource(&self, rrid: &RuntimeResourceID) -> Vec<PartitionId> {
261        self.partitions
262            .iter()
263            .filter_map(|partition| {
264                if partition.contains(rrid) {
265                    Some(partition.partition_info().id.clone())
266                } else {
267                    None
268                }
269            })
270            .collect()
271    }
272
273    /// Iterates over all `RuntimeResourceID`s across all partitions.
274    ///
275    /// # Returns
276    /// An iterator yielding references to `RuntimeResourceID`.
277    pub fn iter_all_runtime_resource_ids(&self) -> impl Iterator<Item = &RuntimeResourceID> + '_ {
278        self.partitions.iter().flat_map(|partition| {
279            partition.resources.keys()
280        })
281    }
282    
283    pub fn resource_mounted(&self, rrid: &RuntimeResourceID) -> bool{
284        self.iter_all_runtime_resource_ids().contains(rrid)
285    }
286    
287    pub fn resource_infos(&self, rrid: &RuntimeResourceID) -> Vec<(PartitionId, &ResourceInfo)> {
288        self.partitions_with_resource(rrid)
289            .into_iter()
290            .filter_map(|p_id| {
291                if let Ok(info) = self.resource_info_from(&p_id, rrid) {
292                    Some((p_id, info))
293                } else {
294                    None
295                }
296            })
297            .collect()
298    }
299
300    pub fn resource_info_from(
301        &self,
302        partition_id: &PartitionId,
303        rrid: &RuntimeResourceID,
304    ) -> Result<&ResourceInfo, PartitionManagerError> {
305        let partition = self
306            .partitions
307            .iter()
308            .find(|partition| partition.partition_info().id == *partition_id);
309
310        if let Some(partition) = partition {
311            match partition.get_resource_info(rrid) {
312                Ok(info) => Ok(info),
313                Err(e) => Err(PartitionManagerError::PartitionError(partition_id.clone(), e)),
314            }
315        } else {
316            Err(PartitionManagerError::PartitionNotFound(
317                partition_id.to_string(),
318            ))
319        }
320    }
321
322    /// Resolves a resource by traversing the partition hierarchy.
323    ///
324    /// # Arguments
325    /// - `partition_id` - The ID of the partition to start resolving from.
326    /// - `resource_id` - The ID of the resource to resolve.
327    pub fn resolve_resource_from(
328        &self,
329        partition_id: PartitionId,
330        resource_id: &RuntimeResourceID,
331    ) -> Result<(&ResourceInfo, PartitionId), PartitionManagerError> {
332        match self.find_partition(partition_id.clone()) {
333            Some(partition) => {
334                if partition.contains(resource_id) {
335                    match partition.get_resource_info(resource_id) {
336                        Ok(info) => Ok((info, partition_id.clone())),
337                        Err(_) => Err(PartitionManagerError::ResourceNotFound(resource_id.to_string())),
338                    }
339                } else {
340                    match &partition.partition_info().parent {
341                        
342                        Some(parent_id) => {
343                            self.resolve_resource_from(parent_id.clone(), resource_id)
344                        },
345                        None => {
346                            Err(PartitionManagerError::ResourceNotFound(resource_id.to_string()))
347                        }
348                    }
349                }
350            },
351            None => {
352                Err(PartitionNotFound(partition_id.to_string()))
353            }
354        }
355    }
356    #[deprecated(
357        since = "1.0.0",
358        note = "prefer direct access through the partitions field"
359    )]
360    pub fn partitions(&self) -> &Vec<ResourcePartition> {
361        &self.partitions
362    }
363
364    #[deprecated(
365        since = "1.1.0",
366        note = "please implement this yourself, it is out of scope for this struct"
367    )]
368    pub fn print_resource_changelog(&self, rrid: &RuntimeResourceID) {
369        println!("Resource: {rrid}");
370
371        for partition in &self.partitions {
372            let mut last_occurence: Option<&ResourceInfo> = None;
373
374            let size =
375                |info: &ResourceInfo| info.compressed_size().unwrap_or(info.header.data_size);
376
377            let changes = partition.resource_patch_indices(rrid);
378            let deletions = partition.resource_removal_indices(rrid);
379            let occurrences = changes
380                .clone()
381                .into_iter()
382                .chain(deletions.clone().into_iter())
383                .collect::<Vec<PatchId>>();
384
385            for occurence in occurrences.iter().sorted() {
386                println!(
387                    "{}: {}",
388                    match occurence {
389                        PatchId::Base => {
390                            "Base"
391                        }
392                        PatchId::Patch(_) => {
393                            "Patch"
394                        }
395                    },
396                    partition.partition_info().filename(*occurence)
397                );
398
399                if deletions.contains(occurence) {
400                    println!("\t- Removal: resource deleted");
401                    last_occurence = None;
402                }
403
404                if changes.contains(occurence) {
405                    if let Ok(info) = partition.resource_info_from(rrid, *occurence) {
406                        if let Some(last_info) = last_occurence {
407                            println!(
408                                "\t- Modification: Size changed from {} to {}",
409                                size(last_info),
410                                size(info)
411                            );
412                        } else {
413                            println!("\t- Addition: New occurrence, Size {} bytes", size(info))
414                        }
415                        last_occurence = Some(info);
416                    }
417                }
418            }
419        }
420    }
421}