Improve power connection

factorio-blueprint/src/abstraction/power_connection.rs

@@ -8,14 +8,35 @@ use super::*;
 
 #[derive(Debug)]
 struct PowerGraph {
-    nodes: HashMap<Position, Vec<(Position, f64)>>,
+    nodes: HashMap<Position, (Vec<Position>, f64)>,
+}
+
+#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
+enum NodeType {
+    In,
+    Out,
 }
 
 impl WheightedGraph for PowerGraph {
-    type Node = Position;
+    type Node = (Position, NodeType);
 
     fn edge(&self, node: &Self::Node, num: usize) -> Option<(Self::Node, f64)> {
-        self.nodes.get(node).and_then(|v| v.get(num).cloned())
+        match node.1 {
+            NodeType::In => {
+                if num == 0 {
+                    self.nodes
+                        .get(&node.0)
+                        .map(|v| ((node.0, NodeType::Out), v.1))
+                } else {
+                    None
+                }
+            }
+            NodeType::Out => self
+                .nodes
+                .get(&node.0)
+                .and_then(|v| v.0.get(num).cloned())
+                .map(|v| ((v, NodeType::In), 0.0)),
+        }
     }
 }
 
@@ -59,7 +80,7 @@ impl Blueprint {
 
         for y_source in 0..blocked.height {
             for x_source in 0..blocked.width {
-                if blocked.get(x_source, y_source).is_some() {
+                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)
@@ -77,11 +98,11 @@ impl Blueprint {
                         })
                     {
                         if let &Some(w) = blocked.get(xx, yy) {
-                            edges.push((Position::new(xx as PositionType, yy as PositionType), w));
+                            edges.push(Position::new(xx as PositionType, yy as PositionType));
                         }
                     }
                     if !edges.is_empty() {
-                        graph.nodes.insert(pos, edges);
+                        graph.nodes.insert(pos, (edges, w));
                     }
                 }
             }
@@ -89,7 +110,12 @@ impl Blueprint {
 
         let pole_positions = power_poles
             .iter()
-            .map(|(_, e)| (e.position - Position::new(1, 1)) / 2 + offset)
+            .map(|(_, e)| {
+                (
+                    (e.position - Position::new(1, 1)) / 2 + offset,
+                    NodeType::Out,
+                )
+            })
             .collect::<Vec<_>>();
 
         if let Some(res) =
@@ -101,7 +127,11 @@ impl Blueprint {
                 .collect::<HashMap<_, _>>();
 
             for path in res {
-                for &p in &path {
+                let path_iter = path.iter().filter_map(|(p, n)| match n {
+                    NodeType::In => None,
+                    NodeType::Out => Some(p),
+                });
+                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) => {
@@ -113,7 +143,7 @@ impl Blueprint {
                         }
                     }
                 }
-                for (p, n) in path.iter().zip(path[1..].iter()) {
+                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);
                 }
             }