Refactor power connection to work with different electric poles

factorio-blueprint-generator/src/bin/multistation.rs

@@ -2,6 +2,7 @@ use clap::Parser;
 use factorio_blueprint::{BlueprintString, encode};
 use factorio_blueprint_generator::multistation::{StationSpec, multistation};
 use factorio_core::beltoptions::Beltspeed;
+use factorio_core::visualize::Visualize;
 
 #[derive(Parser)]
 struct Args {
@@ -35,59 +36,60 @@ fn main() {
             })
         })
         .collect();
+    let mut b = multistation(
+        &stations,
+        // &[
+        //     StationSpec {
+        //         locomotives: 2,
+        //         wagons: 4,
+        //         load: false,
+        //         beltspeed: Beltspeed::Normal,
+        //         lanes: 4,
+        //     },
+        //     StationSpec {
+        //         locomotives: 3,
+        //         wagons: 8,
+        //         load: false,
+        //         beltspeed: Beltspeed::Turbo,
+        //         lanes: 8,
+        //     },
+        //     StationSpec {
+        //         locomotives: 3,
+        //         wagons: 8,
+        //         load: false,
+        //         beltspeed: Beltspeed::Turbo,
+        //         lanes: 4,
+        //     },
+        //     StationSpec {
+        //         locomotives: 3,
+        //         wagons: 8,
+        //         load: false,
+        //         beltspeed: Beltspeed::Turbo,
+        //         lanes: 2,
+        //     },
+        //     StationSpec {
+        //         locomotives: 3,
+        //         wagons: 8,
+        //         load: false,
+        //         beltspeed: Beltspeed::Turbo,
+        //         lanes: 1,
+        //     },
+        //     StationSpec {
+        //         locomotives: 1,
+        //         wagons: 1,
+        //         load: false,
+        //         beltspeed: Beltspeed::Turbo,
+        //         lanes: 1,
+        //     },
+        // ],
+        8,
+    )
+    .0;
 
-    let b = BlueprintString::Blueprint(
+    b.connect_power_networks();
-        multistation(
+
-            &stations,
+    b.print_visualization();
-            // &[
+    let b = BlueprintString::Blueprint(b.to_blueprint());
-            //     StationSpec {
-            //         locomotives: 2,
-            //         wagons: 4,
-            //         load: false,
-            //         beltspeed: Beltspeed::Normal,
-            //         lanes: 4,
-            //     },
-            //     StationSpec {
-            //         locomotives: 3,
-            //         wagons: 8,
-            //         load: false,
-            //         beltspeed: Beltspeed::Turbo,
-            //         lanes: 8,
-            //     },
-            //     StationSpec {
-            //         locomotives: 3,
-            //         wagons: 8,
-            //         load: false,
-            //         beltspeed: Beltspeed::Turbo,
-            //         lanes: 4,
-            //     },
-            //     StationSpec {
-            //         locomotives: 3,
-            //         wagons: 8,
-            //         load: false,
-            //         beltspeed: Beltspeed::Turbo,
-            //         lanes: 2,
-            //     },
-            //     StationSpec {
-            //         locomotives: 3,
-            //         wagons: 8,
-            //         load: false,
-            //         beltspeed: Beltspeed::Turbo,
-            //         lanes: 1,
-            //     },
-            //     StationSpec {
-            //         locomotives: 1,
-            //         wagons: 1,
-            //         load: false,
-            //         beltspeed: Beltspeed::Turbo,
-            //         lanes: 1,
-            //     },
-            // ],
-            8,
-        )
-        .0
-        .to_blueprint(),
-    );
 
     if args.json {
         println!("{}", serde_json::to_string_pretty(&b).unwrap());

factorio-blueprint-generator/src/bin/power_connection_test.rs

@@ -0,0 +1,19 @@
+use factorio_blueprint::abstraction::{Blueprint, Entity};
+use factorio_core::prelude::*;
+use factorio_core::visualize::Visualize;
+
+fn main() {
+    let mut b = Blueprint::new();
+
+    b.add_entity(Entity::new_electric_pole(
+        factorio_blueprint::abstraction::ElectricPoleType::Medium,
+        Position::new(1, 1),
+    ));
+    b.add_entity(Entity::new_electric_pole(
+        factorio_blueprint::abstraction::ElectricPoleType::Medium,
+        Position::new(101, 3),
+    ));
+
+    b.connect_power_networks();
+    b.print_visualization();
+}

factorio-blueprint/src/abstraction.rs

@@ -1,12 +1,16 @@
 use factorio_core::{
     aabb::AABB,
     beltoptions::Beltspeed,
+    misc::PositionMap,
     pathfield::PathField,
     prelude::*,
     quaterdirection::QuaterDirection,
     visualize::{Color, Visualization},
 };
-use std::{collections::HashMap, sync::atomic::AtomicUsize};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::atomic::AtomicUsize,
+};
 
 use crate::{BlueprintEntity, BlueprintPosition};
 
@@ -68,6 +72,13 @@ impl ElectricPoleType {
             ElectricPoleType::Big | ElectricPoleType::Substation => Position::new(4, 4),
         }
     }
+
+    fn alignment(&self) -> PositionType {
+        match self {
+            ElectricPoleType::Small | ElectricPoleType::Medium => 1,
+            ElectricPoleType::Big | ElectricPoleType::Substation => 0,
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -359,12 +370,13 @@ impl Entity {
         let halve_size = self.size() / 2;
         match self.direction {
             DirectionType::Dir(direction) => match direction {
-                Direction::Up | Direction::Down => {
+                Direction::Up | Direction::Down => AABB::new(
-                    AABB::new(self.position - halve_size, self.position + halve_size)
+                    self.position - halve_size,
-                }
+                    self.position + halve_size - Position::new(1, 1),
+                ),
                 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),
+                    self.position + Position::new(halve_size.y, halve_size.x) - Position::new(1, 1),
                 ),
             },
             DirectionType::QuarterDir(_) => {
@@ -415,7 +427,7 @@ impl Entity {
                     None,
                 ),
                 ElectricPoleType::Big => {
-                    for (dx, dy) in [(-1, -1), (1, -1), (-1, 1), (1, 1)] {
+                    for (dx, dy) in [(-2, -2), (0, -2), (-2, 0), (0, 0)] {
                         v.add_symbol(
                             (self.position + Position::new(dx, dy)) / 2 + offset,
                             factorio_core::visualize::Symbol::Char('l'),
@@ -588,6 +600,51 @@ impl Blueprint {
             .reduce(|a, b| a.combine(b))
             .unwrap()
     }
+
+    pub fn placeable(&self, pos: Position, size: Position) -> bool {
+        let aabb = AABB::new(pos - size / 2, pos + size / 2 - Position::new(1, 1));
+        !self
+            .entities
+            .iter()
+            .any(|(_, e)| AABB::collision(e.get_aabb(), aabb))
+    }
+
+    pub fn placibility_map(&self) -> Placibility {
+        let mut blocked = PositionMap::new(self.get_aabb().unwrap(), &false);
+        // dbg!(self.get_aabb(), self.get_aabb().unwrap().size());
+
+        for (_, e) in self.entities.iter() {
+            let aabb = e.get_aabb();
+            for y in aabb.min().y..=aabb.max().y {
+                for x in aabb.min().x..=aabb.max().x {
+                    blocked[Position::new(x, y)] = true;
+                }
+            }
+        }
+
+        Placibility { blocked }
+    }
+}
+
+#[derive(Debug)]
+struct Placibility {
+    blocked: PositionMap<bool>,
+}
+
+impl Placibility {
+    pub fn placeable(&self, pos: Position, size: Position) -> bool {
+        let aabb = AABB::new(pos - size / 2, pos + size / 2 - Position::new(1, 1));
+        for y in aabb.min().y..=aabb.max().y {
+            for x in aabb.min().x..=aabb.max().x {
+                if !self.blocked.get_aabb().contains_pos(Position::new(x, y))
+                    || self.blocked[Position::new(x, y)] == true
+                {
+                    return false;
+                }
+            }
+        }
+        true
+    }
 }
 
 impl Default for Blueprint {

factorio-blueprint/src/abstraction/power_connection.rs

@@ -8,7 +8,7 @@ use super::*;
 
 #[derive(Debug)]
 struct PowerGraph {
-    nodes: HashMap<Position, (Vec<Position>, f64)>,
+    nodes: HashMap<(Position, ElectricPoleType), (Vec<(Position, ElectricPoleType)>, f64)>,
 }
 
 #[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
@@ -17,134 +17,215 @@ enum NodeType {
     Out,
 }
 
+#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
+struct Node {
+    pos: Position,
+    electric_pole_type: ElectricPoleType,
+    node_type: NodeType,
+}
+
 impl WheightedGraph for PowerGraph {
-    type Node = (Position, NodeType);
+    type Node = Node;
 
     fn edge(&self, node: &Self::Node, num: usize) -> Option<(Self::Node, f64)> {
-        match node.1 {
+        match node.node_type {
             NodeType::In => {
                 if num == 0 {
                     self.nodes
-                        .get(&node.0)
+                        .get(&(node.pos, node.electric_pole_type))
-                        .map(|v| ((node.0, NodeType::Out), v.1))
+                        .map(|v| {
+                            (
+                                Node {
+                                    pos: node.pos,
+                                    electric_pole_type: node.electric_pole_type,
+                                    node_type: NodeType::Out,
+                                },
+                                v.1,
+                            )
+                        })
                 } else {
                     None
                 }
             }
             NodeType::Out => self
                 .nodes
-                .get(&node.0)
+                .get(&(node.pos, node.electric_pole_type))
-                .and_then(|v| v.0.get(num).cloned())
+                .and_then(|v| {
-                .map(|v| ((v, NodeType::In), 0.0)),
+                    v.0.get(num).cloned().map(|(pos, electric_pole_type)| {
+                        (
+                            Node {
+                                pos,
+                                electric_pole_type,
+                                node_type: NodeType::In,
+                            },
+                            0.0,
+                        )
+                    })
+                }),
         }
     }
 }
 
+fn wire_reach(electric_pole_type: ElectricPoleType) -> PositionType {
+    match electric_pole_type {
+        ElectricPoleType::Small => 15,
+        ElectricPoleType::Medium => 18,
+        ElectricPoleType::Big => 64,
+        ElectricPoleType::Substation => todo!(),
+    }
+}
+
+fn power_connections(
+    pos: Position,
+    electric_pole_type: ElectricPoleType,
+) -> impl Iterator<Item = (Position, ElectricPoleType)> {
+    [
+        ElectricPoleType::Small,
+        ElectricPoleType::Medium,
+        ElectricPoleType::Big,
+    ]
+    .into_iter()
+    .flat_map(move |other_electric_pole_type| {
+        let reach = PositionType::min(
+            wire_reach(electric_pole_type),
+            wire_reach(other_electric_pole_type),
+        );
+
+        let alignment = other_electric_pole_type.alignment();
+
+        ((pos.x - reach)..=(pos.x + reach))
+            .filter(move |x| x.rem_euclid(2) == alignment)
+            .flat_map(move |x| {
+                ((pos.y - reach)..=(pos.y + reach))
+                    .filter(move |y| y.rem_euclid(2) == alignment)
+                    .map(move |y| (x, y))
+            })
+            .filter(move |(x, y)| {
+                (x - pos.x) * (x - pos.x) + (y - pos.y) * (y - pos.y) <= reach * reach
+            })
+            .map(move |(x, y)| (Position::new(x, y), other_electric_pole_type))
+    })
+}
+
 impl Blueprint {
     pub fn connect_power_networks(&mut self) {
-        let power_poles = self
+        let mut power_pole_map = self
             .entities
             .iter()
-            .filter(|&(_, e)| matches!(e.entity, EntityType::ElectricPole(_)))
+            .filter_map(|(k, e)| match e.entity {
-            .cloned()
+                EntityType::ElectricPole(electric_pole_type) => {
-            .collect::<Vec<_>>();
+                    Some(((e.position, electric_pole_type), *k))
+                }
+                _ => None,
+            })
+            .collect::<HashMap<_, _>>();
 
         let aabb = self.get_aabb().unwrap();
-        let size = (aabb.size() - Position::new(1, 1)) / 2;
 
-        let offset = aabb.min() / -2;
+        let mut nodes = HashMap::new();
 
-        let mut blocked =
+        let placibility = self.placibility_map();
-            factorio_core::misc::Map::new_with(size.x as usize, size.y as usize, Some(1.0));
+        // dbg!(&placibility);
 
-        for (_, e) in &self.entities {
+        for y in aabb.min().y..=aabb.max().y {
-            let entity_aabb = e.get_aabb();
+            for x in aabb.min().x..=aabb.max().x {
+                for electric_pole_type in [
+                    ElectricPoleType::Small,
+                    ElectricPoleType::Medium,
+                    ElectricPoleType::Big,
+                ] {
+                    // dbg!(x, y, electric_pole_type);
+                    let alignment = electric_pole_type.alignment();
+                    if x.rem_euclid(2) == alignment && y.rem_euclid(2) == alignment {
+                        let cost = if placibility
+                            .placeable(Position::new(x, y), electric_pole_type.size())
+                        {
+                            Some(match electric_pole_type {
+                                ElectricPoleType::Small => 0.8,
+                                ElectricPoleType::Medium => 1.0,
+                                ElectricPoleType::Big => 1.5,
+                                ElectricPoleType::Substation => todo!(),
+                            })
+                        } else if power_pole_map
+                            .contains_key(&(Position::new(x, y), electric_pole_type))
+                        {
+                            Some(0.0)
+                        } else {
+                            None
+                        };
 
-            let entity_min = entity_aabb.min() / 2;
+                        if let Some(cost) = cost {
-            let entity_max = entity_aabb.max() / 2;
+                            let v = power_connections(Position::new(x, y), electric_pole_type)
-            for y in entity_min.y..entity_max.y {
+                                // .filter(|&(pos, _)| aabb.contains_pos(pos))
-                for x in entity_min.x..entity_max.x {
+                                // .inspect(|d| {
-                    *blocked.get_mut((x + offset.x) as usize, (y + offset.y) as usize) = None;
+                                //     if y == 261 {
-                }
+                                //         dbg!(d);
-            }
+                                //     }
-        }
+                                // })
-
+                                // .filter(|&(pos, electric_pole_type)| {
-        for (_, p) in &power_poles {
+                                //     placibility.placeable(pos, electric_pole_type.size())
-            let pos = (p.position - Position::new(1, 1)) / 2 + offset;
+                                //         || power_pole_map.contains_key(&(pos, electric_pole_type))
-            *blocked.get_mut(pos.x as usize, pos.y as usize) = Some(0.0);
+                                // })
-        }
+                                .collect();
-
+                            nodes.insert((Position::new(x, y), electric_pole_type), (v, cost));
-        let mut graph = PowerGraph {
-            nodes: HashMap::new(),
-        };
-
-        for y_source in 0..blocked.height {
-            for x_source in 0..blocked.width {
-                if let &Some(w) = blocked.get(x_source, y_source) {
-                    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));
                         }
                     }
-                    if !edges.is_empty() {
-                        graph.nodes.insert(pos, (edges, w));
-                    }
                 }
             }
         }
 
-        let pole_positions = power_poles
+        let graph = PowerGraph { nodes };
-            .iter()
+        // dbg!(&graph);
-            .map(|(_, e)| {
+        let pole_positions = power_pole_map
-                (
+            .keys()
-                    (e.position - Position::new(1, 1)) / 2 + offset,
+            .map(|&(pos, electric_pole_type)| Node {
-                    NodeType::Out,
+                pos,
-                )
+                electric_pole_type,
+                node_type: NodeType::Out,
             })
             .collect::<Vec<_>>();
 
+        // dbg!(&pole_positions);
+
+        // let pole_positions = [
+        //     Node {
+        //         pos: Position::new(-25, 279),
+        //         electric_pole_type: ElectricPoleType::Medium,
+        //         node_type: NodeType::Out,
+        //     },
+        //     Node {
+        //         pos: Position::new(239, 257),
+        //         electric_pole_type: ElectricPoleType::Medium,
+        //         node_type: NodeType::Out,
+        //     },
+        // ];
+
         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 {
-                let path_iter = path.iter().filter_map(|(p, n)| match n {
+                let path_iter = path.iter().filter_map(
-                    NodeType::In => None,
+                    |Node {
-                    NodeType::Out => Some(p),
+                         pos,
-                });
+                         electric_pole_type,
-                for &p in path_iter.clone() {
+                         node_type,
+                     }| match node_type {
+                        NodeType::In => None,
+                        NodeType::Out => Some((*pos, *electric_pole_type)),
+                    },
+                );
+
+                for p in path_iter.clone() {
                     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(
+                            let k = self.add_entity(Entity::new_electric_pole(p.1, p.0));
-                                ElectricPoleType::Medium,
-                                (p - offset) * 2 + Position::new(1, 1),
-                            ));
                             vacant_entry.insert(k);
                         }
                     }
                 }
                 for (p, n) in path_iter.clone().zip(path_iter.skip(1)) {
-                    self.add_wire(power_pole_map[n], 5, power_pole_map[p], 5);
+                    self.add_wire(power_pole_map[&n], 5, power_pole_map[&p], 5);
                 }
             }
         }

factorio-blueprint/src/abstraction/visualize.rs

@@ -6,7 +6,7 @@ use factorio_core::{
 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 mut v = Visualization::new((aabb.size() + Position::new(1, 1)) / 2);
 
         let offset = aabb.min() / -2;
 
@@ -18,9 +18,9 @@ impl Visualize for super::Blueprint {
             let entity_aabb = e.get_aabb();
 
             let entity_min = entity_aabb.min() / 2;
-            let entity_max = entity_aabb.max() / 2;
+            let entity_max = (entity_aabb.max() - Position::new(1, 1)) / 2;
-            for y in entity_min.y..entity_max.y {
+            for y in entity_min.y..=entity_max.y {
-                for x in entity_min.x..entity_max.x {
+                for x in entity_min.x..=entity_max.x {
                     v.overwrite_background(Position::new(x, y) + offset, Some(Color::gray(32)));
                 }
             }
@@ -29,3 +29,26 @@ impl Visualize for super::Blueprint {
         v
     }
 }
+
+impl Visualize for super::Placibility {
+    fn visualize(&self) -> Visualization {
+        let aabb = self.blocked.get_aabb();
+
+        let mut v = Visualization::new(aabb.size());
+
+        for y in aabb.min().y..=aabb.max().y {
+            for x in aabb.min().x..=aabb.max().x {
+                if self.blocked[Position::new(x, y)] {
+                    v.add_symbol(
+                        Position::new(x, y) - aabb.min(),
+                        factorio_core::visualize::Symbol::Block,
+                        Some(Color::white()),
+                        None,
+                    );
+                }
+            }
+        }
+
+        v
+    }
+}

factorio-core/src/aabb.rs

@@ -45,6 +45,10 @@ impl AABB {
             && self.min.y <= other.max.y
             && self.max.y >= other.min.y
     }
+
+    pub fn contains_pos(self, pos: Position) -> bool {
+        self.min.x <= pos.x && self.max.x >= pos.x && self.min.y <= pos.y && self.max.y >= pos.y
+    }
 }
 
 #[cfg(test)]

factorio-core/src/misc/mod.rs

@@ -1,5 +1,7 @@
 pub mod arena;
 pub mod map;
+pub mod position_map;
 
 pub use arena::*;
 pub use map::*;
+pub use position_map::*;

factorio-core/src/misc/position_map.rs

@@ -0,0 +1,62 @@
+use std::ops::{Index, IndexMut};
+
+use crate::{aabb::AABB, prelude::Position};
+
+#[derive(Debug)]
+pub struct PositionMap<T> {
+    data: Vec<T>,
+    aabb: AABB,
+}
+
+impl<T> PositionMap<T>
+where
+    T: Clone,
+{
+    pub fn new(aabb: AABB, value: &T) -> Self {
+        Self::new_internal(aabb, || value.clone())
+    }
+}
+
+impl<T> PositionMap<T>
+where
+    T: Default,
+{
+    pub fn new_with_default(aabb: AABB) -> Self {
+        Self::new_internal(aabb, || T::default())
+    }
+}
+
+impl<T> PositionMap<T> {
+    fn new_internal(aabb: AABB, f: impl Fn() -> T) -> Self {
+        let len = aabb.size().x as usize * aabb.size().y as usize;
+        Self {
+            data: (0..len).map(|_| f()).collect(),
+            aabb,
+        }
+    }
+
+    fn get_index(&self, pos: Position) -> usize {
+        let p = pos - self.aabb.min();
+
+        p.y as usize * self.aabb.size().x as usize + p.x as usize
+    }
+
+    pub fn get_aabb(&self) -> AABB {
+        self.aabb
+    }
+}
+
+impl<T> Index<Position> for PositionMap<T> {
+    type Output = T;
+
+    fn index(&self, index: Position) -> &Self::Output {
+        &self.data[self.get_index(index)]
+    }
+}
+
+impl<T> IndexMut<Position> for PositionMap<T> {
+    fn index_mut(&mut self, index: Position) -> &mut Self::Output {
+        let i = self.get_index(index);
+        &mut self.data[i]
+    }
+}