Path-like chunk name component parser

src/chunk.rs

@@ -5,29 +5,35 @@
 //! - Patch: re_chunk_XXX.pak.patch_XXX.pak
 //! - Sub: re_chunk_XXX.pak.sub_XXX.pak
 //! - Sub Patch: re_chunk_XXX.pak.sub_XXX.pak.patch_XXX.pak
+//! - DLC: re_dlc_stm_3308900.pak (and more)
 
 use color_eyre::eyre;
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum ChunkComponent {
+    /// Base chunk with major ID (re_chunk_XXX.pak)
+    Base(u32),
+    /// DLC chunk with DLC ID (re_dlc_stm_3308900.pak)
+    Dlc(String),
+    /// Patch chunk with patch ID (XXX in .patch_XXX.pak)
+    Patch(u32),
+    /// Sub chunk with sub ID (XXX in .sub_XXX.pak)
+    Sub(u32),
+    /// Sub patch chunk with sub patch ID (YYY in .sub_XXX.pak.patch_YYY.pak)
+    SubPatch(u32),
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct ChunkName {
-    /// Major chunk ID (XXX in re_chunk_XXX.pak)
-    pub major_id: u32,
-    /// Patch number (XXX in .patch_XXX.pak)
-    pub patch_id: Option<u32>,
-    /// Sub chunk ID (XXX in .sub_XXX.pak)
-    pub sub_id: Option<u32>,
-    /// Patch number for sub chunk (YYY in .sub_XXX.pak.patch_YYY.pak)
-    pub sub_patch_id: Option<u32>,
+    /// Chunk components
+    pub components: Vec<ChunkComponent>,
 }
 
 impl ChunkName {
     /// Create a new base chunk name (re_chunk_XXX.pak)
     pub fn new(major_id: u32) -> Self {
         Self {
-            major_id,
-            patch_id: None,
-            sub_id: None,
-            sub_patch_id: None,
+            components: vec![ChunkComponent::Base(major_id)],
         }
     }
 
@@ -42,36 +48,102 @@ impl ChunkName {
         }
 
         // every 2 parts is a component
-        let components = dot_parts
+        let part_pairs = dot_parts
             .chunks_exact(2)
             .map(|c| (c[0], c[1]))
             .collect::<Vec<(&str, &str)>>();
+
         // check if all parts have the correct extension
-        if !components.iter().all(|(_, ext)| *ext == "pak") {
+        if !part_pairs.iter().all(|(_, ext)| *ext == "pak") {
             return Err(eyre::eyre!(
                 "Invalid chunk name with invalid extension: {}",
                 name
             ));
         }
 
-        let mut this = Self::new(0);
+        let mut components = Vec::new();
+        let mut has_sub = false;
 
-        for (name, _) in components.iter() {
-            let component = Self::parse_component(name)?;
+        for (part_name, _) in part_pairs.iter() {
+            let component = Self::parse_component(part_name)?;
             match component {
-                Component::Major(id) => this.major_id = id,
-                Component::Sub(id) => this.sub_id = Some(id),
+                Component::Major(id) => {
+                    components.push(ChunkComponent::Base(id));
+                }
+                Component::Dlc(id) => {
+                    components.push(ChunkComponent::Dlc(id));
+                }
+                Component::Sub(id) => {
+                    components.push(ChunkComponent::Sub(id));
+                    has_sub = true;
+                }
                 Component::Patch(id) => {
-                    if this.sub_id.is_some() {
-                        this.sub_patch_id = Some(id);
+                    if has_sub {
+                        components.push(ChunkComponent::SubPatch(id));
                     } else {
-                        this.patch_id = Some(id);
+                        components.push(ChunkComponent::Patch(id));
                     }
                 }
             }
         }
 
-        Ok(this)
+        Ok(Self { components })
+    }
+
+    /// Get the major ID (base chunk ID)
+    pub fn major_id(&self) -> Option<u32> {
+        self.components.iter().find_map(|c| match c {
+            ChunkComponent::Base(id) => Some(*id),
+            _ => None,
+        })
+    }
+
+    /// Get the patch ID
+    pub fn patch_id(&self) -> Option<u32> {
+        self.components.iter().find_map(|c| match c {
+            ChunkComponent::Patch(id) => Some(*id),
+            _ => None,
+        })
+    }
+
+    /// Get the sub ID
+    pub fn sub_id(&self) -> Option<u32> {
+        self.components.iter().find_map(|c| match c {
+            ChunkComponent::Sub(id) => Some(*id),
+            _ => None,
+        })
+    }
+
+    /// Get the sub patch ID
+    pub fn sub_patch_id(&self) -> Option<u32> {
+        self.components.iter().find_map(|c| match c {
+            ChunkComponent::SubPatch(id) => Some(*id),
+            _ => None,
+        })
+    }
+
+    /// Get the DLC ID
+    pub fn dlc_id(&self) -> Option<&str> {
+        self.components.iter().find_map(|c| match c {
+            ChunkComponent::Dlc(id) => Some(id.as_str()),
+            _ => None,
+        })
+    }
+
+    /// Check if this is a DLC chunk
+    pub fn is_dlc(&self) -> bool {
+        self.components
+            .iter()
+            .any(|c| matches!(c, ChunkComponent::Dlc(_)))
+    }
+
+    /// Add a sub patch component with the given ID
+    pub fn with_sub_patch(&self, patch_id: u32) -> Self {
+        let mut new_components = self.components.clone();
+        new_components.push(ChunkComponent::SubPatch(patch_id));
+        Self {
+            components: new_components,
+        }
     }
 
     fn parse_component(name: &str) -> color_eyre::Result<Component> {
@@ -82,6 +154,9 @@ impl ChunkName {
                 .parse::<u32>()
                 .map_err(|e| eyre::eyre!("Chunk name with invalid major ID: {}", e))?;
             Ok(Component::Major(major_id))
+        } else if name.starts_with("re_dlc_") {
+            let dlc_id = name.strip_prefix("re_dlc_").unwrap().to_string();
+            Ok(Component::Dlc(dlc_id))
         } else if name.starts_with("patch_") {
             let patch_id = name
                 .strip_prefix("patch_")
@@ -107,18 +182,18 @@ impl ChunkName {
 
 impl std::fmt::Display for ChunkName {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "re_chunk_{:03}.pak", self.major_id)?;
-        if let Some(patch_id) = self.patch_id {
-            write!(f, ".patch_{:03}.pak", patch_id)?;
-            return Ok(());
+        for (i, component) in self.components.iter().enumerate() {
+            if i > 0 {
+                write!(f, ".")?;
+            }
+            match component {
+                ChunkComponent::Base(id) => write!(f, "re_chunk_{:03}.pak", id)?,
+                ChunkComponent::Dlc(id) => write!(f, "re_dlc_{}.pak", id)?,
+                ChunkComponent::Patch(id) => write!(f, "patch_{:03}.pak", id)?,
+                ChunkComponent::Sub(id) => write!(f, "sub_{:03}.pak", id)?,
+                ChunkComponent::SubPatch(id) => write!(f, "patch_{:03}.pak", id)?,
+            }
         }
-        if let Some(sub_id) = self.sub_id {
-            write!(f, ".sub_{:03}.pak", sub_id)?;
-        }
-        if let Some(sub_patch_id) = self.sub_patch_id {
-            write!(f, ".patch_{:03}.pak", sub_patch_id)?;
-        }
-
         Ok(())
     }
 }
@@ -131,16 +206,45 @@ impl PartialOrd for ChunkName {
 
 impl Ord for ChunkName {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        self.major_id
-            .cmp(&other.major_id)
-            .then(self.sub_id.cmp(&other.sub_id))
-            .then(self.patch_id.cmp(&other.patch_id))
-            .then(self.sub_patch_id.cmp(&other.sub_patch_id))
+        // compare by component count first
+        self.components
+            .len()
+            .cmp(&other.components.len())
+            .then_with(|| {
+                // compare each component
+                for (a, b) in self.components.iter().zip(other.components.iter()) {
+                    let cmp = match (a, b) {
+                        (ChunkComponent::Base(a), ChunkComponent::Base(b)) => a.cmp(b),
+                        (ChunkComponent::Dlc(a), ChunkComponent::Dlc(b)) => a.cmp(b),
+                        (ChunkComponent::Patch(a), ChunkComponent::Patch(b)) => a.cmp(b),
+                        (ChunkComponent::Sub(a), ChunkComponent::Sub(b)) => a.cmp(b),
+                        (ChunkComponent::SubPatch(a), ChunkComponent::SubPatch(b)) => a.cmp(b),
+                        // compare by component type priority
+                        (ChunkComponent::Base(_), _) => std::cmp::Ordering::Less,
+                        (_, ChunkComponent::Base(_)) => std::cmp::Ordering::Greater,
+                        (ChunkComponent::Dlc(_), _) => std::cmp::Ordering::Less,
+                        (_, ChunkComponent::Dlc(_)) => std::cmp::Ordering::Greater,
+                        (ChunkComponent::Sub(_), _) => std::cmp::Ordering::Less,
+                        (_, ChunkComponent::Sub(_)) => std::cmp::Ordering::Greater,
+                        (ChunkComponent::Patch(_), ChunkComponent::SubPatch(_)) => {
+                            std::cmp::Ordering::Less
+                        }
+                        (ChunkComponent::SubPatch(_), ChunkComponent::Patch(_)) => {
+                            std::cmp::Ordering::Greater
+                        }
+                    };
+                    if cmp != std::cmp::Ordering::Equal {
+                        return cmp;
+                    }
+                }
+                std::cmp::Ordering::Equal
+            })
     }
 }
 
 enum Component {
     Major(u32),
+    Dlc(String),
     Patch(u32),
     Sub(u32),
 }
@@ -170,5 +274,51 @@ mod tests {
             sub_patch.to_string(),
             "re_chunk_000.pak.sub_000.pak.patch_001.pak"
         );
+
+        // Test DLC chunk
+        let dlc = ChunkName::try_from_str("re_dlc_stm_3308900.pak").unwrap();
+        assert_eq!(dlc.to_string(), "re_dlc_stm_3308900.pak");
+    }
+
+    #[test]
+    fn test_chunk_helper_methods() {
+        // Test base chunk helper methods
+        let base = ChunkName::new(123);
+        assert_eq!(base.major_id(), Some(123));
+        assert_eq!(base.patch_id(), None);
+        assert_eq!(base.sub_id(), None);
+        assert_eq!(base.sub_patch_id(), None);
+        assert_eq!(base.dlc_id(), None);
+        assert!(!base.is_dlc());
+
+        // Test complex chunk helper methods
+        let complex =
+            ChunkName::try_from_str("re_chunk_456.pak.sub_789.pak.patch_012.pak").unwrap();
+        assert_eq!(complex.major_id(), Some(456));
+        assert_eq!(complex.patch_id(), None);
+        assert_eq!(complex.sub_id(), Some(789));
+        assert_eq!(complex.sub_patch_id(), Some(12));
+        assert_eq!(complex.dlc_id(), None);
+        assert!(!complex.is_dlc());
+
+        // Test DLC chunk helper methods
+        let dlc = ChunkName::try_from_str("re_dlc_stm_3308900.pak").unwrap();
+        assert_eq!(dlc.major_id(), None);
+        assert_eq!(dlc.dlc_id(), Some("stm_3308900"));
+        assert!(dlc.is_dlc());
+    }
+
+    #[test]
+    fn test_with_sub_patch() {
+        let base = ChunkName::try_from_str("re_chunk_000.pak.sub_001.pak").unwrap();
+        let with_patch = base.with_sub_patch(99);
+
+        assert_eq!(with_patch.major_id(), Some(0));
+        assert_eq!(with_patch.sub_id(), Some(1));
+        assert_eq!(with_patch.sub_patch_id(), Some(99));
+        assert_eq!(
+            with_patch.to_string(),
+            "re_chunk_000.pak.sub_001.pak.patch_099.pak"
+        );
     }
 }