Add power connection and steiner tree

Cargo.lock

@@ -544,6 +544,7 @@ dependencies = [
  "bon",
  "clap",
  "factorio-core",
+ "factorio-graph",
  "flate2",
  "serde",
  "serde_json",

factorio-blueprint-generator/src/factory.rs

@@ -533,10 +533,12 @@ pub fn generate_factory<L: Layouter, P: Pathfinder + Sync, R: Rng + SeedableRng
 
     b.transform(Transformation::new(
         Direction::Up,
-        Position::new(multistation_x_offset, multistation_y_offset - 2),
+        Position::new(multistation_x_offset + 1, multistation_y_offset - 2 + 1),
     ));
 
     b.add_blueprint(multistation_blueprint);
 
+    b.connect_power_networks();
+
     b
 }

factorio-blueprint/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2024"
 
 [dependencies]
 factorio-core = { path = "../factorio-core" }
+factorio-graph = { path = "../factorio-graph" }
 base64 = "0.22.1"
 bon = "3.3.2"
 clap = { version = "4.5.26", features = ["derive"] }

factorio-blueprint/src/abstraction.rs

@@ -1,11 +1,18 @@
 use factorio_core::{
-    aabb::AABB, beltoptions::Beltspeed, pathfield::PathField, prelude::*,
+    aabb::AABB,
+    beltoptions::Beltspeed,
+    pathfield::PathField,
+    prelude::*,
     quaterdirection::QuaterDirection,
+    visualize::{Color, Visualization},
 };
 use std::{collections::HashMap, sync::atomic::AtomicUsize};
 
 use crate::{BlueprintEntity, BlueprintPosition};
 
+mod power_connection;
+mod visualize;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum UndergroundType {
     Input,
@@ -94,15 +101,26 @@ impl RailType {
     }
 
     fn size(&self) -> Position {
-        todo!()
+        Position::new(12, 12)
     }
 }
 
+#[derive(Debug, Clone)]
 enum DirectionType {
     Dir(Direction),
     QuarterDir(QuaterDirection),
 }
 
+impl DirectionType {
+    fn unwrap_dir(&self) -> Direction {
+        match self {
+            DirectionType::Dir(direction) => *direction,
+            DirectionType::QuarterDir(_quater_direction) => panic!(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
 pub struct Entity {
     entity: EntityType,
     position: Position,
@@ -110,6 +128,7 @@ pub struct Entity {
     quality: Quality,
 }
 
+#[derive(Debug, Clone)]
 pub enum EntityType {
     Belt(Beltspeed),
     UndergroundBelt(Beltspeed, UndergroundType),
@@ -335,6 +354,96 @@ impl Entity {
             }
         };
     }
+
+    pub fn get_aabb(&self) -> AABB {
+        let halve_size = self.size() / 2;
+        match self.direction {
+            DirectionType::Dir(direction) => match direction {
+                Direction::Up | Direction::Down => {
+                    AABB::new(self.position - halve_size, self.position + halve_size)
+                }
+                Direction::Right | Direction::Left => AABB::new(
+                    self.position - Position::new(halve_size.y, halve_size.x),
+                    self.position + Position::new(halve_size.y, halve_size.x),
+                ),
+            },
+            DirectionType::QuarterDir(_) => {
+                AABB::new(self.position - halve_size, self.position + halve_size)
+            }
+        }
+    }
+    fn visualize(&self, v: &mut Visualization, offset: Position) {
+        match &self.entity {
+            EntityType::Belt(beltspeed) => {
+                v.add_symbol(
+                    (self.position - Position::new(1, 1)) / 2 + offset,
+                    factorio_core::visualize::Symbol::Arrow(self.direction.unwrap_dir()),
+                    Some(factorio_core::visualize::Color::white()),
+                    None,
+                );
+            }
+            EntityType::UndergroundBelt(beltspeed, underground_type) => match underground_type {
+                UndergroundType::Input => {
+                    v.add_symbol(
+                        (self.position - Position::new(1, 1)) / 2 + offset,
+                        factorio_core::visualize::Symbol::ArrowEnter(self.direction.unwrap_dir()),
+                        Some(factorio_core::visualize::Color::white()),
+                        None,
+                    );
+                }
+                UndergroundType::Output => {
+                    v.add_symbol(
+                        (self.position - Position::new(1, 1)) / 2 + offset,
+                        factorio_core::visualize::Symbol::ArrowExit(self.direction.unwrap_dir()),
+                        Some(factorio_core::visualize::Color::white()),
+                        None,
+                    );
+                }
+            },
+            EntityType::Splitter(beltspeed) => (),
+            EntityType::ElectricPole(electric_pole_type) => match electric_pole_type {
+                ElectricPoleType::Small => v.add_symbol(
+                    (self.position - Position::new(1, 1)) / 2 + offset,
+                    factorio_core::visualize::Symbol::Char('s'),
+                    Some(Color::cyan()),
+                    None,
+                ),
+                ElectricPoleType::Medium => v.add_symbol(
+                    (self.position - Position::new(1, 1)) / 2 + offset,
+                    factorio_core::visualize::Symbol::Char('m'),
+                    Some(Color::cyan()),
+                    None,
+                ),
+                ElectricPoleType::Big => {
+                    for (dx, dy) in [(-1, -1), (1, -1), (-1, 1), (1, 1)] {
+                        v.add_symbol(
+                            (self.position + Position::new(dx, dy)) / 2 + offset,
+                            factorio_core::visualize::Symbol::Char('l'),
+                            Some(Color::cyan()),
+                            None,
+                        )
+                    }
+                }
+                ElectricPoleType::Substation => {
+                    for (dx, dy) in [(-1, -1), (1, -1), (-1, 1), (1, 1)] {
+                        v.add_symbol(
+                            (self.position + Position::new(dx, dy)) / 2 + offset,
+                            factorio_core::visualize::Symbol::Char('S'),
+                            Some(Color::cyan()),
+                            None,
+                        )
+                    }
+                }
+            },
+            EntityType::Inserter {
+                inserter_type,
+                override_stack_size,
+            } => (),
+            EntityType::Production { name, recipe, size } => (),
+            EntityType::Rail { rail_type } => (),
+            EntityType::Unknown { name, size, misc } => (),
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -459,6 +568,13 @@ impl Blueprint {
             .build()
     }
 
+    pub fn get_aabb(&self) -> Option<AABB> {
+        self.entities
+            .iter()
+            .map(|(_, e)| e.get_aabb())
+            .reduce(AABB::combine)
+    }
+
     pub fn transform(&mut self, transform: Transformation) {
         for (_, e) in &mut self.entities {
             e.transform(transform);

factorio-blueprint/src/abstraction/power_connection.rs

@@ -0,0 +1,122 @@
+use factorio_core::visualize::Visualize;
+use factorio_graph::{
+    priority_queue::binary_heap::FastBinaryHeap,
+    wheighted_graph::{WheightedGraph, steiner_tree},
+};
+
+use super::*;
+
+#[derive(Debug)]
+struct PowerGraph {
+    nodes: HashMap<Position, Vec<(Position, f64)>>,
+}
+
+impl WheightedGraph for PowerGraph {
+    type Node = Position;
+
+    fn edge(&self, node: &Self::Node, num: usize) -> Option<(Self::Node, f64)> {
+        self.nodes.get(node).and_then(|v| v.get(num).cloned())
+    }
+}
+
+impl Blueprint {
+    pub fn connect_power_networks(&mut self) {
+        let power_poles = self
+            .entities
+            .iter()
+            .filter(|&(_, e)| matches!(e.entity, EntityType::ElectricPole(_)))
+            .cloned()
+            .collect::<Vec<_>>();
+
+        let aabb = self.get_aabb().unwrap();
+        let size = (aabb.size() - Position::new(1, 1)) / 2;
+
+        let offset = aabb.min() / -2;
+
+        let mut blocked =
+            factorio_core::misc::Map::new_with(size.x as usize, size.y as usize, Some(1.0));
+
+        for (_, e) in &self.entities {
+            let entity_aabb = e.get_aabb();
+
+            let entity_min = entity_aabb.min() / 2;
+            let entity_max = entity_aabb.max() / 2;
+            for y in entity_min.y..entity_max.y {
+                for x in entity_min.x..entity_max.x {
+                    *blocked.get_mut((x + offset.x) as usize, (y + offset.y) as usize) = None;
+                }
+            }
+        }
+
+        for (_, p) in &power_poles {
+            let pos = (p.position - Position::new(1, 1)) / 2 + offset;
+            *blocked.get_mut(pos.x as usize, pos.y as usize) = Some(0.0);
+        }
+
+        let mut graph = PowerGraph {
+            nodes: HashMap::new(),
+        };
+
+        for y_source in 0..blocked.height {
+            for x_source in 0..blocked.width {
+                if blocked.get(x_source, y_source).is_some() {
+                    let pos = Position::new(x_source as PositionType, y_source as PositionType);
+                    let mut edges = Vec::new();
+                    for (xx, yy) in (-9..=9)
+                        .flat_map(|dx| (-9..=9).map(move |dy| (dx, dy)))
+                        .filter(|&(dx, dy)| dx * dx + dy * dy <= 9 * 9)
+                        .filter_map(|(dx, dy)| {
+                            x_source
+                                .checked_add_signed(dx)
+                                .filter(|&x| x < blocked.width)
+                                .zip(
+                                    y_source
+                                        .checked_add_signed(dy)
+                                        .filter(|&y| y < blocked.height),
+                                )
+                        })
+                    {
+                        if let &Some(w) = blocked.get(xx, yy) {
+                            edges.push((Position::new(xx as PositionType, yy as PositionType), w));
+                        }
+                    }
+                    if !edges.is_empty() {
+                        graph.nodes.insert(pos, edges);
+                    }
+                }
+            }
+        }
+
+        let pole_positions = power_poles
+            .iter()
+            .map(|(_, e)| (e.position - Position::new(1, 1)) / 2 + offset)
+            .collect::<Vec<_>>();
+
+        if let Some(res) =
+            steiner_tree::takaheshi_matsuyama::<_, FastBinaryHeap<_>>(&graph, &pole_positions)
+        {
+            let mut power_pole_map = power_poles
+                .iter()
+                .map(|(k, e)| ((e.position - Position::new(1, 1)) / 2 + offset, *k))
+                .collect::<HashMap<_, _>>();
+
+            for path in res {
+                for &p in &path {
+                    match power_pole_map.entry(p) {
+                        std::collections::hash_map::Entry::Occupied(_occupied_entry) => (),
+                        std::collections::hash_map::Entry::Vacant(vacant_entry) => {
+                            let k = self.add_entity(Entity::new_electric_pole(
+                                ElectricPoleType::Medium,
+                                (p - offset) * 2 + Position::new(1, 1),
+                            ));
+                            vacant_entry.insert(k);
+                        }
+                    }
+                }
+                for (p, n) in path.iter().zip(path[1..].iter()) {
+                    self.add_wire(power_pole_map[n], 5, power_pole_map[p], 5);
+                }
+            }
+        }
+    }
+}

factorio-blueprint/src/abstraction/visualize.rs

@@ -0,0 +1,31 @@
+use factorio_core::{
+    prelude::*,
+    visualize::{Color, Visualization, Visualize},
+};
+
+impl Visualize for super::Blueprint {
+    fn visualize(&self) -> factorio_core::visualize::Visualization {
+        let aabb = self.get_aabb().unwrap();
+        let mut v = Visualization::new((aabb.size() - Position::new(1, 1)) / 2);
+
+        let offset = aabb.min() / -2;
+
+        for (_, e) in &self.entities {
+            e.visualize(&mut v, offset);
+        }
+
+        for (_, e) in &self.entities {
+            let entity_aabb = e.get_aabb();
+
+            let entity_min = entity_aabb.min() / 2;
+            let entity_max = entity_aabb.max() / 2;
+            for y in entity_min.y..entity_max.y {
+                for x in entity_min.x..entity_max.x {
+                    v.overwrite_background(Position::new(x, y) + offset, Some(Color::gray(32)));
+                }
+            }
+        }
+
+        v
+    }
+}

factorio-graph/src/wheighted_graph/steiner_tree.rs

@@ -2,16 +2,72 @@ use crate::priority_queue::PriorityQueue;
 use std::hash::Hash;
 use std::{collections::HashSet, fmt::Debug};
 
+use super::shortest_path::dijkstra;
 use super::{WheightedGraph, shortest_path::QueueObject};
 
-pub fn takaheshi_matsuyama<G, P>(graph: G, nodes: &[G::Node]) -> Option<Vec<Vec<G::Node>>>
+struct MultistartWrapper<'a, 'b, G>
 where
-    P: PriorityQueue<QueueObject<G::Node>> + Debug,
-    P::Handle: Debug,
-    G::Node: Eq + Hash + Clone,
     G: WheightedGraph,
 {
-    let end_nodes: HashSet<G::Node> = HashSet::from_iter(nodes.iter().cloned());
-
-    todo!()
+    graph: &'a G,
+    start_nodes: &'b [G::Node],
+}
+
+impl<G> WheightedGraph for MultistartWrapper<'_, '_, G>
+where
+    G: WheightedGraph,
+    G::Node: Clone,
+{
+    type Node = Option<G::Node>;
+
+    fn edge(&self, node: &Self::Node, num: usize) -> Option<(Self::Node, f64)> {
+        match node {
+            Some(n) => self.graph.edge(n, num).map(|(n, v)| (Some(n), v)),
+            None => self.start_nodes.get(num).cloned().map(|n| (Some(n), 0.0)),
+        }
+    }
+}
+
+pub fn takaheshi_matsuyama<G, P>(graph: &G, nodes: &[G::Node]) -> Option<Vec<Vec<G::Node>>>
+where
+    P: PriorityQueue<QueueObject<Option<G::Node>>> + Debug,
+    P::Handle: Debug,
+    G::Node: Eq + Hash + Clone + Debug,
+    G: WheightedGraph,
+{
+    if nodes.is_empty() {
+        return Some(Vec::new());
+    }
+
+    let mut end_nodes: HashSet<G::Node> = HashSet::from_iter(nodes.iter().cloned());
+
+    let mut start_nodes = vec![end_nodes.iter().next().cloned().unwrap()];
+
+    end_nodes.remove(&start_nodes[0]);
+
+    let mut paths = Vec::new();
+
+    while !end_nodes.is_empty() {
+        let wrapper = MultistartWrapper {
+            graph,
+            start_nodes: &start_nodes,
+        };
+
+        if let Some(p) = dijkstra::<_, P, _>(&wrapper, None, |n| {
+            n.as_ref().is_some_and(|n| end_nodes.contains(n))
+        }) {
+            let p = p.into_iter().skip(1).map(|n| n.unwrap()).collect();
+
+            for n in &p {
+                end_nodes.remove(n);
+                start_nodes.push(n.clone());
+            }
+
+            paths.push(p);
+        } else {
+            return None;
+        }
+    }
+
+    Some(paths)
 }