Improve station generator

src/bin/station.rs

@@ -1,7 +1,25 @@
-use factorio_blueprint::blueprint::{station::basic_unload_station, BlueprintString};
+use clap::Parser;
+use factorio_blueprint::blueprint::{
+    belt::Beltspeed, station::basic_unload_station, BlueprintString,
+};
+
+#[derive(Parser)]
+struct Args {
+    locomotives: usize,
+    length: usize,
+    outputs: usize,
+    beltspeed: Beltspeed,
+}
 
 fn main() {
-    let b = BlueprintString::Blueprint(basic_unload_station());
+    let args = Args::parse();
+
+    let b = BlueprintString::Blueprint(basic_unload_station(
+        args.locomotives,
+        args.length,
+        args.outputs,
+        args.beltspeed,
+    ));
 
     println!("{}", serde_json::to_string_pretty(&b).unwrap());
 

src/blueprint/belt.rs

@@ -1,7 +1,10 @@
+use clap::ValueEnum;
+
 use crate::{belt_finding::common::PathField, prelude::PositionType};
 
 use super::{Blueprint, BlueprintEntity, BlueprintPosition};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
 pub enum Beltspeed {
     Normal,
     Fast,
@@ -10,7 +13,7 @@ pub enum Beltspeed {
 }
 
 impl Beltspeed {
-    pub fn string(&self) -> String {
+    pub fn string(self) -> String {
         match self {
             Beltspeed::Normal => "transport-belt",
             Beltspeed::Fast => "fast-transport-belt",
@@ -20,7 +23,7 @@ impl Beltspeed {
         .to_owned()
     }
 
-    pub fn string_underground(&self) -> String {
+    pub fn string_underground(self) -> String {
         match self {
             Beltspeed::Normal => "underground-belt",
             Beltspeed::Fast => "fast-underground-belt",
@@ -30,7 +33,7 @@ impl Beltspeed {
         .to_owned()
     }
 
-    pub fn string_splitter(&self) -> String {
+    pub fn string_splitter(self) -> String {
         match self {
             Beltspeed::Normal => "splitter",
             Beltspeed::Fast => "fast-splitter",
@@ -40,7 +43,7 @@ impl Beltspeed {
         .to_owned()
     }
 
-    pub fn halve(&self) -> Self {
+    pub fn halve(self) -> Self {
         match self {
             Beltspeed::Normal => Beltspeed::Normal,
             Beltspeed::Fast => Beltspeed::Normal,
@@ -48,6 +51,14 @@ impl Beltspeed {
             Beltspeed::Turbo => Beltspeed::Fast,
         }
     }
+
+    pub fn halvings(self, i: usize) -> Self {
+        let mut s = self;
+        for _ in 0..i {
+            s = s.halve();
+        }
+        s
+    }
 }
 
 pub fn convert_belt_to_blueprint(

src/blueprint/station.rs

@@ -7,16 +7,20 @@ pub fn station_unload(
     locomotives: usize,
     unloader: &Vec<BlueprintEntity>,
     beltspeed: Beltspeed,
-    double: bool,
+    outputs: usize,
     output_x: f64,
     output_y: f64,
 ) -> Blueprint {
-    assert!(length.is_power_of_two());
+    let section_size = length / outputs;
+    assert!(length % outputs == 0);
+    assert!(section_size.is_power_of_two());
+    assert!(outputs <= length);
 
     let mut e = Vec::new();
 
     let global_x_offset = locomotives * 7;
 
+    // electric poles
     for l in 1..=(length + locomotives) {
         e.push(
             BlueprintEntity::builder(
@@ -28,6 +32,7 @@ pub fn station_unload(
         );
     }
 
+    // unloader
     let offset = e.len();
     for l in 0..length {
         e.extend(unloader.iter().cloned().map(|mut e| {
@@ -37,6 +42,7 @@ pub fn station_unload(
         }));
     }
 
+    // train stop
     let offset = e.len();
     e.push(
         BlueprintEntity::builder(
@@ -48,6 +54,7 @@ pub fn station_unload(
         .build(),
     );
 
+    // rails
     for l in 0..((length * 7 + global_x_offset + 1) / 2) {
         let offset = e.len() as u32;
         e.push(
@@ -61,110 +68,96 @@ pub fn station_unload(
         );
     }
 
-    for i in 0..length {
+    // output and merging
-        let depth = (i / 2).count_ones();
+    for o in 0..outputs {
+        // stubs
+        let stubspeed = beltspeed.halvings(section_size.ilog2() as usize);
+        for i in 0..section_size {
+            let depth = o + i.count_ones() as usize;
+
+            for j in 0..depth {
+                let offset = e.len() as u32;
+                e.push(
+                    BlueprintEntity::builder(
+                        stubspeed.string(),
+                        offset,
+                        BlueprintPosition::new(
+                            (7 * (i + o * section_size) + global_x_offset) as f64 + output_x,
+                            output_y - j as f64,
+                        ),
+                    )
+                    .build(),
+                );
+            }
 
-        for j in 0..depth {
             let offset = e.len() as u32;
             e.push(
                 BlueprintEntity::builder(
-                    beltspeed.string(),
+                    stubspeed.string(),
                     offset,
                     BlueprintPosition::new(
-                        (7 * i + global_x_offset) as f64 + output_x,
+                        (7 * (i + o * section_size) + global_x_offset) as f64 + output_x,
-                        output_y - j as f64,
-                    ),
-                )
-                .build(),
-            );
-        }
-
-        if i % 2 == 0 {
-            let offset = e.len() as u32;
-            e.push(
-                BlueprintEntity::builder(
-                    beltspeed.string(),
-                    offset,
-                    BlueprintPosition::new(
-                        (7 * i + global_x_offset) as f64 + output_x,
                         output_y - depth as f64,
                     ),
                 )
                 .direction(12)
                 .build(),
             );
-            let offset = e.len() as u32;
-            e.push(
-                BlueprintEntity::builder(
-                    beltspeed.string_splitter(),
-                    offset,
-                    BlueprintPosition::new(
-                        (7 * i + global_x_offset) as f64 + output_x - 1.0,
-                        output_y - depth as f64 - 0.5,
-                    ),
-                )
-                .direction(12)
-                .build(),
-            );
-        } else {
-            let offset = e.len() as u32;
-            e.extend((0..=7).map(|j| {
-                BlueprintEntity::builder(
-                    beltspeed.string(),
-                    offset + j as u32,
-                    BlueprintPosition::new(
-                        (7 * i - j + global_x_offset) as f64 + output_x,
-                        output_y - depth as f64 - 1.0,
-                    ),
-                )
-                .direction(12)
-                .build()
-            }));
-            let offset = e.len() as u32;
-            e.push(
-                BlueprintEntity::builder(
-                    beltspeed.string(),
-                    offset,
-                    BlueprintPosition::new(
-                        (7 * i + global_x_offset) as f64 + output_x,
-                        output_y - depth as f64,
-                    ),
-                )
-                .build(),
-            );
         }
-    }
 
-    for i in 2..(length.ilog2() as usize + 1) {
+        // merger
-        let p = 1 << i;
+        for i in 1..(section_size.ilog2() as usize + 1) {
-        for j in 0..(length / p) {
+            let p = 1 << i;
-            let depth = j.count_ones();
+            let mergespeed = beltspeed.halvings((section_size.ilog2() as usize) - i);
-            let offset = e.len() as u32;
+            for j in 0..(section_size / p) {
-            e.push(
+                let depth = o + j.count_ones() as usize;
-                BlueprintEntity::builder(
+                let offset = e.len() as u32;
-                    beltspeed.string_splitter(),
+                e.push(
-                    offset,
+                    BlueprintEntity::builder(
-                    BlueprintPosition::new(
+                        mergespeed.string_splitter(),
-                        (7 * j * p + global_x_offset) as f64 - i as f64 + output_x,
+                        offset,
-                        output_y - i as f64 + 0.5 - depth as f64,
+                        BlueprintPosition::new(
-                    ),
+                            (7 * (j * p + o * section_size) + global_x_offset) as f64 - i as f64
-                )
+                                + output_x,
-                .direction(12)
+                            output_y - i as f64 + 0.5 - depth as f64,
-                .build(),
+                        ),
-            );
+                    )
-            e.extend((0..(7 * p / 2)).map(|l| {
+                    .direction(12)
-                BlueprintEntity::builder(
+                    .build(),
-                    beltspeed.string(),
+                );
-                    offset + l as u32,
+                e.extend((0..(7 * p / 2)).map(|l| {
-                    BlueprintPosition::new(
+                    BlueprintEntity::builder(
-                        (7 * j * p + global_x_offset) as f64 - i as f64 + output_x + l as f64 + 1.0,
+                        mergespeed.halve().string(),
-                        output_y - i as f64 - depth as f64,
+                        offset + 1 + l as u32,
-                    ),
+                        BlueprintPosition::new(
-                )
+                            (7 * (j * p + o * section_size) + global_x_offset) as f64 - i as f64
-                .direction(12)
+                                + output_x
-                .build()
+                                + l as f64
-            }))
+                                + 1.0,
+                            output_y - i as f64 - depth as f64,
+                        ),
+                    )
+                    .direction(12)
+                    .build()
+                }));
+            }
         }
+
+        // connect
+        let offset = e.len() as u32;
+        let step = o + section_size.ilog2() as usize;
+        e.extend((0..(7 * o * section_size)).map(|l| {
+            BlueprintEntity::builder(
+                beltspeed.string(),
+                offset + l as u32,
+                BlueprintPosition::new(
+                    (l + global_x_offset) as f64 + output_x - section_size.ilog2() as f64,
+                    output_y - step as f64,
+                ),
+            )
+            .direction(12)
+            .build()
+        }));
     }
 
     Blueprint::builder()
@@ -178,7 +171,12 @@ pub fn station_unload(
         .build()
 }
 
-pub fn basic_unload_station() -> Blueprint {
+pub fn basic_unload_station(
+    locomotives: usize,
+    length: usize,
+    outputs: usize,
+    beltspeed: Beltspeed,
+) -> Blueprint {
     let top = vec![
         BlueprintEntity::builder(
             "steel-chest".to_owned(),
@@ -286,5 +284,5 @@ pub fn basic_unload_station() -> Blueprint {
         .build(),
     ];
 
-    station_unload(64, 3, &full, Beltspeed::Normal, false, 3.5, -2.5)
+    station_unload(length, locomotives, &full, beltspeed, outputs, 3.5, -2.5)
 }