Working on the belt extension

src/belt.rs

@@ -199,11 +199,19 @@ where
         }
     }
 
+    pub fn set_extension(&mut self, extension_px: u32) -> Option<Vec<Entity>> {
+        let ents = self.t.set_extension(extension_px);
+        self.update_t_origin();
+        ents
+    }
+
     pub fn extend_by_belt(&mut self, mut other: Belt<Entity>) -> bool {
         if self.dir != other.dir {
             return false;
         }
 
+        println!("EXT BY BELT: {:?} {:?}", self, other);
+
         if !self.tile_on_same_axis(other.front_tile) {
             return false;
         }
@@ -867,6 +875,7 @@ 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 {
@@ -877,6 +886,8 @@ 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 {
@@ -890,6 +901,7 @@ 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

@@ -29,6 +29,7 @@ where
     Entity: Copy + Clone + std::fmt::Debug + std::cmp::PartialEq,
 {
     length_px: u32,
+    ext_px: u32,
     slot_size: u32,
     // contents: (pos, free, entity)
     contents: VecDeque<(u32, u32, Entity)>,
@@ -75,6 +76,7 @@ where
         Self {
             slot_size,
             length_px: 0,
+            ext_px: 0,
             contents: VecDeque::new(),
             probes: vec![],
         }
@@ -454,6 +456,32 @@ where
         self.length_px.max(start_of_next_elem)
     }
 
+    pub fn extension_len(&self) -> u32 {
+        self.ext_px
+    }
+
+    pub fn set_extension(&mut self, ext_px: u32) -> Option<Vec<Entity>> {
+        let entities = if ext_px < self.ext_px {
+            let from = self.length_px - self.ext_px;
+            let to = self.length_px - 1;
+            self.probe(from, to, |_pos, _ent| ProbeAction::TakeRange)
+        } else {
+            None
+        };
+
+        self.length_px -= self.ext_px;
+        self.ext_px = ext_px;
+        self.length_px += self.ext_px;
+
+        if self.contents.len() > 0 {
+            self.recalc_free_at(self.contents.len() - 1);
+        }
+
+        self.clear_all_probes();
+
+        entities
+    }
+
     pub fn append(&mut self, add_px: u32) {
         self.length_px += add_px;
 
@@ -474,6 +502,7 @@ where
         }
 
         self.length_px += other.length_px;
+        self.ext_px += other.ext_px;
 
         if prev_len > 0 && (prev_len - 1) < self.contents.len() {
             self.recalc_free_at(prev_len - 1);
@@ -506,9 +535,21 @@ where
         }
 
         let to = to.min(self.length_px - 1);
+        let (from, to) = if from > to { (to, from) } else { (from, to) };
+
         let split_len = (to - from) + 1;
 
-        let (from, to) = if from > to { (to, from) } else { (from, to) };
+        let to = if self.ext_px > 0 {
+            // If we have an extension at the end, we need to extend the "to" range
+            // if it matches the actual end of the belt without the extension.
+            if to == ((self.length_px - 1) - self.ext_px) {
+                to + self.ext_px
+            } else {
+                to
+            }
+        } else {
+            to
+        };
 
         let entities = self.probe(from, to, |_pos, _ent| ProbeAction::TakeRange);
 
@@ -2642,7 +2683,13 @@ mod test {
         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), (64, 0, 50)]
+            vec![
+                (0, 0, 10),
+                (16, 0, 20),
+                (32, 0, 30),
+                (48, 0, 40),
+                (64, 0, 50)
+            ]
         );
     }
 
@@ -2660,7 +2707,13 @@ mod test {
         assert!(t.insert_at_probe(p1, 50));
         assert_eq!(
             t.dump_contents(),
-            vec![(0, 0, 10), (16, 0, 20), (32, 0, 30), (48, 0, 40), (64, 0, 50)]
+            vec![
+                (0, 0, 10),
+                (16, 0, 20),
+                (32, 0, 30),
+                (48, 0, 40),
+                (64, 0, 50)
+            ]
         );
     }