From af625cf905a1500950131288c1f83a16f286c821 Mon Sep 17 00:00:00 2001 From: hal8174 Date: Mon, 3 Mar 2025 22:45:24 +0100 Subject: [PATCH] Add power connection and steiner tree --- Cargo.lock | 1 + factorio-blueprint-generator/src/factory.rs | 4 +- factorio-blueprint/Cargo.toml | 1 + factorio-blueprint/src/abstraction.rs | 120 ++++++++++++++++- .../src/abstraction/power_connection.rs | 122 ++++++++++++++++++ .../src/abstraction/visualize.rs | 31 +++++ .../src/wheighted_graph/steiner_tree.rs | 70 +++++++++- 7 files changed, 339 insertions(+), 10 deletions(-) create mode 100644 factorio-blueprint/src/abstraction/power_connection.rs create mode 100644 factorio-blueprint/src/abstraction/visualize.rs diff --git a/Cargo.lock b/Cargo.lock index 06d5d5e..88286c8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -544,6 +544,7 @@ dependencies = [ "bon", "clap", "factorio-core", + "factorio-graph", "flate2", "serde", "serde_json", diff --git a/factorio-blueprint-generator/src/factory.rs b/factorio-blueprint-generator/src/factory.rs index 865e36f..28a6019 100644 --- a/factorio-blueprint-generator/src/factory.rs +++ b/factorio-blueprint-generator/src/factory.rs @@ -533,10 +533,12 @@ pub fn generate_factory 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 { + 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); diff --git a/factorio-blueprint/src/abstraction/power_connection.rs b/factorio-blueprint/src/abstraction/power_connection.rs new file mode 100644 index 0000000..280d425 --- /dev/null +++ b/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>, +} + +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::>(); + + 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::>(); + + 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::>(); + + 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); + } + } + } + } +} diff --git a/factorio-blueprint/src/abstraction/visualize.rs b/factorio-blueprint/src/abstraction/visualize.rs new file mode 100644 index 0000000..87ffc03 --- /dev/null +++ b/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 + } +} diff --git a/factorio-graph/src/wheighted_graph/steiner_tree.rs b/factorio-graph/src/wheighted_graph/steiner_tree.rs index f365f52..54bcc1e 100644 --- a/factorio-graph/src/wheighted_graph/steiner_tree.rs +++ b/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(graph: G, nodes: &[G::Node]) -> Option>> +struct MultistartWrapper<'a, 'b, G> where - P: PriorityQueue> + Debug, - P::Handle: Debug, - G::Node: Eq + Hash + Clone, G: WheightedGraph, { - let end_nodes: HashSet = HashSet::from_iter(nodes.iter().cloned()); - - todo!() + graph: &'a G, + start_nodes: &'b [G::Node], +} + +impl WheightedGraph for MultistartWrapper<'_, '_, G> +where + G: WheightedGraph, + G::Node: Clone, +{ + type Node = Option; + + 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(graph: &G, nodes: &[G::Node]) -> Option>> +where + P: PriorityQueue>> + 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 = 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) }