Refactoring and testing the extension code a little bit

src/belt.rs

@@ -199,12 +199,26 @@ where
         }
     }
 
-    pub fn set_extension(&mut self, extension_px: u32) -> Option<Vec<Entity>> {
-        let ents = self.t.set_extension(extension_px);
+    pub fn enable_extension(&mut self) {
+        let ents = self.t.set_extension((BELT_TILE_SIZE_PX / 2) as u32);
+        assert!(
+            ents.is_none(),
+            "Enabling the extension should not lead to items drop off, \
+            unless there is a grave bug in the belt logic!"
+        );
+        self.update_t_origin();
+    }
+
+    pub fn disable_extension(&mut self) -> Option<Vec<Entity>> {
+        let ents = self.t.set_extension(0);
         self.update_t_origin();
         ents
     }
 
+    pub fn is_extended(&self) -> bool {
+        self.t.extension_len() > 0
+    }
+
     pub fn extend_by_belt(&mut self, mut other: Belt<Entity>) -> bool {
         if self.dir != other.dir {
             return false;
@@ -875,7 +889,6 @@ where
                 .expect("Belt must be present here!");
             //            println!("BELT PRE EXT: {:?}", belt);
             belt.extend_to(tile_pos);
-            assert!(belt.set_extension(0).is_none(), "Should have no entities at the end here!");
             //            println!("BELT POST EXT: {:?}", belt);
 
             if let Some(adj_2_id) = adj_2_id {
@@ -886,8 +899,6 @@ where
             }
             //            println!("BELT POST EXTBB: {:?}", belt);
 
-            belt.set_extension((BELT_TILE_SIZE_PX / 2) as u32);
-
             let belt_id = self.mount_belt(belt);
 
             if let Some(entities) = entities {
@@ -901,7 +912,6 @@ where
 
         let mut belt = Belt::new(0, dir);
         belt.extend(1);
-        belt.set_extension((BELT_TILE_SIZE_PX / 2) as u32);
         belt.set_front(tile_pos);
 
         let belt_id = self.mount_belt(belt);

src/transport.rs

@@ -2738,4 +2738,50 @@ mod test {
             vec![(0, 0, 10), (16, 0, 20), (32, 0, 30), (48, 0, 40)]
         );
     }
+
+    #[test]
+    fn check_transport_extension_tprobe_on_all_entities_fit() {
+        let mut t = Transport::<u64>::new(16);
+        t.append(1);
+        t.set_extension(63);
+        t.append_entity(10);
+        t.append_entity(20);
+        t.append_entity(30);
+
+        let p1 = t.add_probe(48);
+
+        assert_eq!(
+            t.dump_contents(),
+            vec![(0, 0, 10), (16, 0, 20), (32, 16, 30)]
+        );
+        assert!(t.insert_at_probe(p1, 40));
+
+        assert_eq!(
+            t.dump_contents(),
+            vec![(0, 0, 10), (16, 0, 20), (32, 0, 30), (48, 0, 40)]
+        );
+    }
+
+    #[test]
+    fn check_transport_extension_merge_end_full() {
+        let mut t1 = Transport::new(16);
+        t1.append(16);
+        t1.append_entity(941);
+        t1.set_extension(16);
+        t1.append_entity(942);
+        let mut t2 = Transport::new(16);
+        t2.append(16);
+        t2.append_entity(841);
+        t2.set_extension(16);
+        t2.append_entity(842);
+
+        assert_eq!(t1.dump_contents(), vec![(0, 0, 941), (16, 0, 942)]);
+        t1.append_transport(&t2);
+        assert_eq!(
+            t1.dump_contents(),
+            vec![(0, 0, 941), (16, 0, 942), (32, 0, 841), (48, 0, 842)]
+        );
+        assert_eq!(t1.length(), 64);
+        assert_eq!(t1.extension_len(), 32);
+    }
 }