Convert station to abstraction

factorio-blueprint/src/abstraction.rs

@@ -1,9 +1,9 @@
 use factorio_core::{
     beltoptions::Beltspeed,
     pathfield::PathField,
-    prelude::{Direction, Position, PositionType},
+    prelude::{Direction, Position, PositionType, Transformable, Transformation},
 };
-use std::{collections::HashMap, str::FromStr, sync::atomic::AtomicUsize};
+use std::{collections::HashMap, sync::atomic::AtomicUsize};
 
 use crate::{BlueprintEntity, BlueprintPosition};
 
@@ -13,33 +13,220 @@ pub enum UndergroundType {
     Output,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum InserterType {
+    Burner,
+    Normal,
+    Fast,
+    Bulk,
+    Stack,
+}
+
+impl InserterType {
+    fn string(&self) -> String {
+        match self {
+            InserterType::Burner => "burner-inserter",
+            InserterType::Normal => "inserter",
+            InserterType::Fast => "fast-inserter",
+            InserterType::Bulk => "bulk-inserter",
+            InserterType::Stack => "stack-inserter",
+        }
+        .to_owned()
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum ElectricPoleType {
+    Small,
+    Medium,
+    Big,
+    Substation,
+}
+
+impl ElectricPoleType {
+    fn string(&self) -> String {
+        match self {
+            ElectricPoleType::Small => "small-electric-pole",
+            ElectricPoleType::Medium => "medium-electric-pole",
+            ElectricPoleType::Big => "big-electric-pole",
+            ElectricPoleType::Substation => "substation",
+        }
+        .to_owned()
+    }
+
+    fn size(&self) -> Position {
+        match self {
+            ElectricPoleType::Small | ElectricPoleType::Medium => Position::new(1, 1),
+            ElectricPoleType::Big | ElectricPoleType::Substation => Position::new(2, 2),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum Quality {
+    Normal,
+    Uncommon,
+    Rare,
+    Epic,
+    Legendary,
+}
+
 pub struct Entity {
-    name: String,
+    entity: EntityType,
     position: Position,
     direction: Direction,
-    size: Position,
-    underground_type: Option<UndergroundType>,
+    quality: Quality,
+}
+
+pub enum EntityType {
+    Belt(Beltspeed),
+    UndergroundBelt(Beltspeed, UndergroundType),
+    Splitter(Beltspeed),
+    ElectricPole(ElectricPoleType),
+    Inserter {
+        inserter_type: InserterType,
+        override_stack_size: Option<u8>,
+    },
+    Unknown {
+        name: String,
+        size: Position,
+        misc: HashMap<String, serde_json::Value>,
+    },
 }
 
 impl Entity {
-    pub fn new(
+    pub fn new(entity: EntityType, position: Position, direction: Direction) -> Self {
+        Self {
+            entity,
+            position,
+            direction,
+            quality: Quality::Normal,
+        }
+    }
+
+    pub fn new_belt(beltspeed: Beltspeed, position: Position, direction: Direction) -> Self {
+        Self::new(EntityType::Belt(beltspeed), position, direction)
+    }
+
+    pub fn new_underground_belt(
+        beltspeed: Beltspeed,
+        underground_type: UndergroundType,
+        position: Position,
+        direction: Direction,
+    ) -> Self {
+        Self::new(
+            EntityType::UndergroundBelt(beltspeed, underground_type),
+            position,
+            direction,
+        )
+    }
+
+    pub fn new_splitter(beltspeed: Beltspeed, position: Position, direction: Direction) -> Self {
+        Self::new(EntityType::Splitter(beltspeed), position, direction)
+    }
+
+    pub fn new_inserter(
+        inserter_type: InserterType,
+        override_stack_size: Option<u8>,
+        position: Position,
+        direction: Direction,
+    ) -> Self {
+        Self::new(
+            EntityType::Inserter {
+                inserter_type,
+                override_stack_size,
+            },
+            position,
+            direction,
+        )
+    }
+
+    pub fn new_electric_pole(electric_pole_type: ElectricPoleType, position: Position) -> Self {
+        Self::new(
+            EntityType::ElectricPole(electric_pole_type),
+            position,
+            Direction::Up,
+        )
+    }
+
+    pub fn new_unknown(
         name: impl AsRef<str>,
         position: Position,
         direction: Direction,
         size: Position,
     ) -> Self {
-        Self {
-            name: name.as_ref().to_owned(),
+        Self::new(
+            EntityType::Unknown {
+                name: name.as_ref().to_owned(),
+                size,
+                misc: HashMap::new(),
+            },
             position,
             direction,
-            size,
-            underground_type: None,
+        )
+    }
+
+    pub fn quality(mut self, quality: Quality) -> Self {
+        self.quality = quality;
+        self
+    }
+
+    pub fn get_name(&self) -> String {
+        match &self.entity {
+            EntityType::Belt(beltspeed) => beltspeed.string(),
+            EntityType::UndergroundBelt(beltspeed, _underground_type) => {
+                beltspeed.string_underground()
+            }
+            EntityType::Splitter(beltspeed) => beltspeed.string_splitter(),
+            EntityType::Unknown {
+                name,
+                size: _,
+                misc: _,
+            } => name.clone(),
+            EntityType::ElectricPole(electric_pole_type) => electric_pole_type.string(),
+            EntityType::Inserter {
+                inserter_type,
+                override_stack_size: _,
+            } => inserter_type.string(),
         }
     }
 
-    pub fn underground_type(mut self, underground_type: UndergroundType) -> Self {
-        self.underground_type = Some(underground_type);
-        self
+    pub fn get_maybe_underground_type_string(&self) -> Option<String> {
+        match &self.entity {
+            EntityType::UndergroundBelt(_, underground_type) => match underground_type {
+                UndergroundType::Input => Some(String::from("input")),
+                UndergroundType::Output => Some(String::from("output")),
+            },
+            _ => None,
+        }
+    }
+
+    pub fn get_maybe_override_stack_size(&self) -> Option<u8> {
+        match &self.entity {
+            EntityType::Inserter {
+                inserter_type: _,
+                override_stack_size,
+            } => *override_stack_size,
+            _ => None,
+        }
+    }
+
+    pub fn size(&self) -> Position {
+        match &self.entity {
+            EntityType::Splitter(_) => Position::new(2, 1),
+            EntityType::Unknown {
+                name: _,
+                size,
+                misc: _,
+            } => *size,
+            EntityType::ElectricPole(electric_pole_type) => electric_pole_type.size(),
+            _ => Position::new(1, 1),
+        }
+    }
+
+    pub fn transform(&mut self, transform: Transformation) {
+        self.position = self.position.transform(transform);
+        self.direction = self.direction.transform(transform);
     }
 }
 
@@ -51,6 +238,7 @@ static ENTITY_COUNTER: AtomicUsize = AtomicUsize::new(0);
 pub struct Blueprint {
     entities: Vec<(EntityKey, Entity)>,
     keys: HashMap<EntityKey, usize>,
+    wires: Vec<(EntityKey, u8, EntityKey, u8)>,
 }
 
 impl Blueprint {
@@ -58,6 +246,7 @@ impl Blueprint {
         Self {
             entities: Vec::new(),
             keys: HashMap::new(),
+            wires: Vec::new(),
         }
     }
 
@@ -79,49 +268,50 @@ impl Blueprint {
         for &p in path {
             match p {
                 PathField::Belt { pos, dir } => {
-                    self.add_entity(Entity::new(
-                        beltspeed.string(),
-                        pos,
-                        dir,
-                        Position::new(1, 1),
-                    ));
+                    self.add_entity(Entity::new_belt(beltspeed, pos, dir));
                 }
                 PathField::Underground { pos, dir, len } => {
-                    self.add_entity(
-                        Entity::new(
-                            beltspeed.string_underground(),
-                            pos,
-                            dir,
-                            Position::new(1, 1),
-                        )
-                        .underground_type(UndergroundType::Input),
-                    );
-                    self.add_entity(
-                        Entity::new(
-                            beltspeed.string_underground(),
-                            pos.in_direction(&dir, len as PositionType),
-                            dir,
-                            Position::new(1, 1),
-                        )
-                        .underground_type(UndergroundType::Output),
-                    );
+                    self.add_entity(Entity::new_underground_belt(
+                        beltspeed,
+                        UndergroundType::Input,
+                        pos,
+                        dir,
+                    ));
+                    self.add_entity(Entity::new_underground_belt(
+                        beltspeed,
+                        UndergroundType::Output,
+                        pos.in_direction(&dir, len as PositionType),
+                        dir,
+                    ));
                 }
             };
         }
     }
 
+    pub fn add_wire(&mut self, a: EntityKey, endpoint_a: u8, b: EntityKey, endpoint_b: u8) {
+        self.wires.push((a, endpoint_a, b, endpoint_b));
+    }
+
+    pub fn add_blueprint(&mut self, other: Self) {
+        self.entities.extend(other.entities);
+        self.keys.extend(other.keys);
+        self.wires.extend(other.wires);
+    }
+
     pub fn to_blueprint(&self) -> super::Blueprint {
         let entities = self
             .entities
             .iter()
             .enumerate()
             .map(|(i, (_, e))| {
+                let size = (e.size() - Position::new(1, 1))
+                    .transform(Transformation::new(e.direction, Position::new(0, 0)));
                 BlueprintEntity::builder(
-                    e.name.clone(),
+                    e.get_name(),
                     i as u32 + 1,
                     BlueprintPosition::new(
-                        e.position.x as f64 + 0.5 * e.size.x as f64,
-                        e.position.y as f64 + 0.5 * e.size.y as f64,
+                        e.position.x as f64 + 0.5 * size.x as f64 + 0.5,
+                        e.position.y as f64 + 0.5 * size.y as f64 + 0.5,
                     ),
                 )
                 .direction(match e.direction {
@@ -130,19 +320,37 @@ impl Blueprint {
                     Direction::Down => 8,
                     Direction::Left => 12,
                 })
-                .maybe_underground_type(e.underground_type.map(|u| match u {
-                    UndergroundType::Input => String::from("input"),
-                    UndergroundType::Output => String::from("output"),
-                }))
+                .maybe_underground_type(e.get_maybe_underground_type_string())
+                .maybe_override_stack_size(e.get_maybe_override_stack_size())
                 .build()
             })
             .collect();
 
+        let wires = self
+            .wires
+            .iter()
+            .map(|&(a, endpoint_a, b, endpoint_b)| {
+                [
+                    self.keys[&a] as u32 + 1,
+                    endpoint_a as u32,
+                    self.keys[&b] as u32 + 1,
+                    endpoint_b as u32,
+                ]
+            })
+            .collect();
+
         super::Blueprint::builder()
             .label(String::from("test"))
             .entities(entities)
+            .wires(wires)
             .build()
     }
+
+    pub fn transform(&mut self, transform: Transformation) {
+        for (_, e) in &mut self.entities {
+            e.transform(transform);
+        }
+    }
 }
 
 impl Default for Blueprint {