From be09655e605fed79e9ee5d7918fedfb8d11a2041 Mon Sep 17 00:00:00 2001
From: mgi388 <135186256+mgi388@users.noreply.github.com>
Date: Tue, 2 Jul 2024 20:51:18 +1000
Subject: [PATCH] add army race enum

---
 src/army/decoder.rs |  7 ++++++-
 src/army/encoder.rs |  4 +++-
 src/army/mod.rs     | 14 +++++++++++++-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/army/decoder.rs b/src/army/decoder.rs
index b1bb45f..b1417d7 100644
--- a/src/army/decoder.rs
+++ b/src/army/decoder.rs
@@ -12,6 +12,7 @@ pub enum DecodeError {
     IoError(IoError),
     InvalidFormat(u32),
     InvalidString,
+    InvalidArmyRace(u8),
     InvalidRegimentAlignment(u8),
     InvalidRegimentClass(u8),
 }
@@ -30,6 +31,7 @@ impl fmt::Display for DecodeError {
             DecodeError::IoError(e) => write!(f, "IO error: {}", e),
             DecodeError::InvalidFormat(format) => write!(f, "invalid format: {}", format),
             DecodeError::InvalidString => write!(f, "invalid string"),
+            DecodeError::InvalidArmyRace(v) => write!(f, "invalid army race: {}", v),
             DecodeError::InvalidRegimentAlignment(v) => {
                 write!(f, "invalid regiment alignment: {}", v)
             }
@@ -88,10 +90,13 @@ impl<R: Read + Seek> Decoder<R> {
 
         let header = self.read_header(start_pos)?;
 
+        let race = ArmyRace::try_from(header.race)
+            .map_err(|_| DecodeError::InvalidArmyRace(header.race))?;
+
         let regiments = self.read_regiments(&header)?;
 
         Ok(Army {
-            race: header.race,
+            race,
             unknown1: header.unknown1.to_vec(),
             unknown2: header.unknown2.to_vec(),
             regiments,
diff --git a/src/army/encoder.rs b/src/army/encoder.rs
index 071df79..19356d7 100644
--- a/src/army/encoder.rs
+++ b/src/army/encoder.rs
@@ -52,12 +52,14 @@ impl<W: Write> Encoder<W> {
     fn write_header(&mut self, army: &Army) -> Result<(), EncodeError> {
         // TODO: Ignoring save file header.
 
+        let race: u8 = army.race.into();
+
         self.writer.write_all(&FORMAT.to_le_bytes())?;
         self.writer
             .write_all(&(army.regiments.len() as u32).to_le_bytes())?;
         self.writer
             .write_all(&(REGIMENT_BLOCK_SIZE as u32).to_le_bytes())?;
-        self.writer.write_all(&[army.race])?;
+        self.writer.write_all(&[race])?;
         self.writer.write_all(&army.unknown1)?;
         self.writer.write_all(&army.unknown2)?;
         self.write_string(&army.small_banner_path)?;
diff --git a/src/army/mod.rs b/src/army/mod.rs
index 672573c..5fb7533 100644
--- a/src/army/mod.rs
+++ b/src/army/mod.rs
@@ -10,7 +10,7 @@ pub use encoder::{EncodeError, Encoder};
 
 #[derive(Debug, Clone, Serialize)]
 pub struct Army {
-    pub race: u8,
+    pub race: ArmyRace,
     unknown1: Vec<u8>,
     unknown2: Vec<u8>,
     pub regiments: Vec<Regiment>,
@@ -32,6 +32,17 @@ pub struct Army {
     unknown3: Vec<u8>,
 }
 
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, IntoPrimitive, PartialEq, Serialize, TryFromPrimitive)]
+pub enum ArmyRace {
+    Empire = 0,
+    EmpireMultiplayer = 1,
+    Greenskin = 2,
+    GreenskinMultiplayer = 3,
+    Undead = 4,
+    UndeadMultiplayer = 5,
+}
+
 #[derive(Debug, Clone, Serialize)]
 pub struct Regiment {
     pub status: [u8; 2],
@@ -309,6 +320,7 @@ mod tests {
         let file = File::open(d).unwrap();
         let a = Decoder::new(file).decode().unwrap();
 
+        assert_eq!(a.race, ArmyRace::Empire);
         assert_eq!(a.small_banner_path, "[BOOKS]\\hshield.spr");
         assert_eq!(a.small_banner_disabled_path, "[BOOKS]\\hgban.spr");
         assert_eq!(a.large_banner_path, "[BOOKS]\\hlban.spr");