chunking: Fix specific component file allocation

crates/ostree-ext/src/chunking.rs

@@ -233,6 +233,17 @@ impl Chunk {
 }
 
 impl Chunking {
+    /// Creates a reverse map from content IDs to checksums
+    fn create_content_id_map(
+        map: &IndexMap<String, ContentID>,
+    ) -> IndexMap<ContentID, Vec<&String>> {
+        let mut rmap = IndexMap::<ContentID, Vec<&String>>::new();
+        for (checksum, contentid) in map.iter() {
+            rmap.entry(Rc::clone(contentid)).or_default().push(checksum);
+        }
+        rmap
+    }
+
     /// Generate an initial single chunk.
     pub fn new(repo: &ostree::Repo, rev: &str) -> Result<Self> {
         // Find the target commit
@@ -310,21 +321,17 @@ impl Chunking {
             return Ok(());
         }
 
-        // Reverses `contentmeta.map` i.e. contentid -> Vec<checksum>
-        let mut rmap = IndexMap::<ContentID, Vec<&String>>::new();
-        for (checksum, contentid) in meta.map.iter() {
-            rmap.entry(Rc::clone(contentid)).or_default().push(checksum);
-        }
-
         // Create exclusive chunks first if specified
         let mut processed_specific_components = BTreeSet::new();
         if let Some(specific_meta) = specific_contentmeta {
+            let specific_map = Self::create_content_id_map(&specific_meta.map);
+
             for component in &specific_meta.sizes {
                 let mut chunk = Chunk::new(&component.meta.name);
                 chunk.packages = vec![component.meta.name.to_string()];
 
                 // Move all objects belonging to this exclusive component
-                if let Some(objects) = rmap.get(&component.meta.identifier) {
+                if let Some(objects) = specific_map.get(&component.meta.identifier) {
                     for &obj in objects {
                         self.remainder.move_obj(&mut chunk, obj);
                     }
@@ -353,6 +360,8 @@ impl Chunking {
             .cloned()
             .collect();
 
+        let rmap = Self::create_content_id_map(&meta.map);
+
         // Process regular components with bin packing if we have remaining layers
         if let Some(remaining) = NonZeroU32::new(self.remaining()) {
             let start = Instant::now();
@@ -1221,4 +1230,183 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_process_mapping_specific_components_contain_correct_objects() -> Result<()> {
+        // This test validates that specific components get their own dedicated layers
+        // and that their objects are properly isolated from regular package layers
+
+        // Setup: Create 5 packages
+        // - pkg1, pkg2: Will be marked as "specific components" (get their own layers)
+        // - pkg3, pkg4, pkg5: Regular packages (will be bin-packed together)
+        let packages = [
+            (1, 100, 50000), // pkg1 - SPECIFIC COMPONENT
+            (2, 200, 40000), // pkg2 - SPECIFIC COMPONENT
+            (3, 300, 30000), // pkg3 - regular package
+            (4, 400, 20000), // pkg4 - regular package
+            (5, 500, 10000), // pkg5 - regular package
+        ];
+
+        let (contentmeta, mut system_metadata, mut specific_components_meta, mut chunking) =
+            setup_exclusive_test(&packages, 8, None)?;
+
+        // Create object mappings
+        // - system_objects_map: Contains ALL objects in the system (both specific and regular)
+        // - specific_components_objects_map: Contains ONLY objects from specific components
+        let mut system_objects_map = IndexMap::new();
+        let mut specific_components_objects_map = IndexMap::new();
+
+        // SPECIFIC COMPONENT 1 (pkg1): owns 3 objects
+        let pkg1_objects = ["checksum_1_a", "checksum_1_b", "checksum_1_c"];
+        for obj in &pkg1_objects {
+            system_objects_map.insert(obj.to_string(), contentmeta[0].meta.identifier.clone());
+            specific_components_objects_map
+                .insert(obj.to_string(), contentmeta[0].meta.identifier.clone());
+        }
+
+        // SPECIFIC COMPONENT 2 (pkg2): owns 2 objects
+        let pkg2_objects = ["checksum_2_a", "checksum_2_b"];
+        for obj in &pkg2_objects {
+            system_objects_map.insert(obj.to_string(), contentmeta[1].meta.identifier.clone());
+            specific_components_objects_map
+                .insert(obj.to_string(), contentmeta[1].meta.identifier.clone());
+        }
+
+        // REGULAR PACKAGE 1 (pkg3): owns 2 objects
+        let pkg3_objects = ["checksum_3_a", "checksum_3_b"];
+        for obj in &pkg3_objects {
+            system_objects_map.insert(obj.to_string(), contentmeta[2].meta.identifier.clone());
+        }
+
+        // REGULAR PACKAGE 2 (pkg4): owns 1 object
+        let pkg4_objects = ["checksum_4_a"];
+        for obj in &pkg4_objects {
+            system_objects_map.insert(obj.to_string(), contentmeta[3].meta.identifier.clone());
+        }
+
+        // REGULAR PACKAGE 3 (pkg5): owns 3 objects
+        let pkg5_objects = ["checksum_5_a", "checksum_5_b", "checksum_5_c"];
+        for obj in &pkg5_objects {
+            system_objects_map.insert(obj.to_string(), contentmeta[4].meta.identifier.clone());
+        }
+
+        // Set up metadata
+        system_metadata.map = system_objects_map;
+        specific_components_meta.map = specific_components_objects_map;
+        specific_components_meta.sizes = vec![contentmeta[0].clone(), contentmeta[1].clone()];
+
+        // Initialize: Add ALL objects to the remainder chunk before processing
+        // This includes both specific component objects and regular package objects
+        // because process_mapping needs to move them from remainder to their final layers
+        for (checksum, _) in &system_metadata.map {
+            chunking.remainder.content.insert(
+                RcStr::from(checksum.as_str()),
+                (
+                    1000,
+                    vec![Utf8PathBuf::from(format!("/path/to/{}", checksum))],
+                ),
+            );
+            chunking.remainder.size += 1000;
+        }
+
+        // Process the mapping
+        // - system_metadata contains ALL objects in the system
+        // - specific_components_meta tells process_mapping which objects belong to specific components
+        chunking.process_mapping(
+            &system_metadata,
+            &Some(NonZeroU32::new(8).unwrap()),
+            None,
+            Some(&specific_components_meta),
+        )?;
+
+        // VALIDATION PART 1: Specific components get their own dedicated chunks
+        assert!(
+            chunking.chunks.len() >= 2,
+            "Should have at least 2 chunks for specific components"
+        );
+
+        // Specific Component Layer 1: pkg1 only
+        let specific_component_1_layer = &chunking.chunks[0];
+        assert_eq!(specific_component_1_layer.name, "pkg1");
+        assert_eq!(specific_component_1_layer.packages, vec!["pkg1"]);
+        assert_eq!(specific_component_1_layer.content.len(), 3);
+        for obj in &pkg1_objects {
+            assert!(
+                specific_component_1_layer.content.contains_key(*obj),
+                "Specific component 1 layer should contain {}",
+                obj
+            );
+        }
+
+        // Specific Component Layer 2: pkg2 only
+        let specific_component_2_layer = &chunking.chunks[1];
+        assert_eq!(specific_component_2_layer.name, "pkg2");
+        assert_eq!(specific_component_2_layer.packages, vec!["pkg2"]);
+        assert_eq!(specific_component_2_layer.content.len(), 2);
+        for obj in &pkg2_objects {
+            assert!(
+                specific_component_2_layer.content.contains_key(*obj),
+                "Specific component 2 layer should contain {}",
+                obj
+            );
+        }
+
+        // VALIDATION PART 2: Specific component layers contain NO regular package objects
+        for specific_layer in &chunking.chunks[0..2] {
+            for obj in pkg3_objects
+                .iter()
+                .chain(&pkg4_objects)
+                .chain(&pkg5_objects)
+            {
+                assert!(
+                    !specific_layer.content.contains_key(*obj),
+                    "Specific component layer '{}' should NOT contain regular package object {}",
+                    specific_layer.name,
+                    obj
+                );
+            }
+        }
+
+        // VALIDATION PART 3: Regular package layers contain NO specific component objects
+        let regular_package_layers = &chunking.chunks[2..];
+        for regular_layer in regular_package_layers {
+            for obj in pkg1_objects.iter().chain(&pkg2_objects) {
+                assert!(
+                    !regular_layer.content.contains_key(*obj),
+                    "Regular package layer should NOT contain specific component object {}",
+                    obj
+                );
+            }
+        }
+
+        // VALIDATION PART 4: All regular package objects are in some regular layer
+        let mut found_regular_objects = BTreeSet::new();
+        for regular_layer in regular_package_layers {
+            for (obj, _) in &regular_layer.content {
+                found_regular_objects.insert(obj.as_ref());
+            }
+        }
+
+        for obj in pkg3_objects
+            .iter()
+            .chain(&pkg4_objects)
+            .chain(&pkg5_objects)
+        {
+            assert!(
+                found_regular_objects.contains(*obj),
+                "Regular package object {} should be in some regular layer",
+                obj
+            );
+        }
+
+        // VALIDATION PART 5: All objects moved from remainder
+        assert_eq!(
+            chunking.remainder.content.len(),
+            0,
+            "All objects should be moved from remainder"
+        );
+        assert_eq!(chunking.remainder.size, 0, "Remainder size should be 0");
+
+        Ok(())
+    }
 }