improve code functionality

src/fot/attributes.rs

@@ -4,6 +4,7 @@ use super::tag::Tag;
 use anyhow::{anyhow, Result};
 use indexmap::IndexMap;
 
+const MAX_STATS: usize = 7;
 const STATS: [&str; 7] = [
     "strength",
     "perception",
@@ -14,6 +15,7 @@ const STATS: [&str; 7] = [
     "luck",
 ];
 
+const MAX_TRAITS: usize = 11;
 const TRAITS: [&str; 11] = [
     "experience",
     "skillPoints",
@@ -28,6 +30,7 @@ const TRAITS: [&str; 11] = [
     "race",
 ];
 
+const MAX_DERIVED: usize = 26;
 const DERIVED: [&str; 26] = [
     "maxHitPoints",
     "maxCarryWeight",
@@ -57,6 +60,7 @@ const DERIVED: [&str; 26] = [
     "levelsPerPerk",
 ];
 
+const MAX_SKILLS: usize = 18;
 const SKILLS: [&str; 18] = [
     "smallGuns",
     "bigGuns",
@@ -78,6 +82,7 @@ const SKILLS: [&str; 18] = [
     "outdoorsman",
 ];
 
+const MAX_OPT_TRAITS: usize = 38;
 const OPT_TRAITS: [&str; 38] = [
     "fastMetabolism",
     "bruiser",
@@ -119,6 +124,7 @@ const OPT_TRAITS: [&str; 38] = [
     "doNightPerson",
 ];
 
+const MAX_PERKS: usize = 111;
 const PERKS: [&str; 111] = [
     "awareness",
     "bonusHtHAttacks",
@@ -233,6 +239,7 @@ const PERKS: [&str; 111] = [
     "unk10",
 ];
 
+const MAX_ADDICTIONS: usize = 10;
 const ADDICTIONS: [&str; 10] = [
     "buffoutAddiction",
     "afterburnerAddiction",
@@ -284,28 +291,28 @@ impl Attributes {
             let _ = rd.read_u32()?;
             let _: Tag = rd.read()?;
 
-            for i in 0..7 {
+            for i in 0..MAX_STATS {
                 stats.insert(STATS[i], rd.read_u32()?);
             }
-            for i in 0..11 {
+            for i in 0..MAX_TRAITS {
                 traits.insert(TRAITS[i], rd.read_u32()?);
             }
-            for i in 0..26 {
+            for i in 0..MAX_DERIVED {
                 derived.insert(DERIVED[i], rd.read_u32()?);
             }
-            for i in 0..18 {
+            for i in 0..MAX_SKILLS {
                 skills.insert(SKILLS[i], rd.read_u32()?);
             }
-            for i in 0..18 {
+            for i in 0..MAX_SKILLS {
                 skill_tags.insert(SKILLS[i], rd.read_bool()?);
             }
-            for i in 0..38 {
+            for i in 0..MAX_OPT_TRAITS {
                 opt_traits.insert(OPT_TRAITS[i], rd.read_bool()?);
             }
-            for i in 0..111 {
+            for i in 0..MAX_PERKS {
                 perks.insert(PERKS[i], rd.read_u32()?);
             }
-            for i in 0..10 {
+            for i in 0..MAX_ADDICTIONS {
                 addictions.insert(ADDICTIONS[i], rd.read_u32()?);
             }
 

src/fot/entity.rs

@@ -1,8 +1,9 @@
+use super::attributes::Attributes;
 use super::decoder::DecoderCtx;
 use super::entitylist::{EntityEncoding, EntityList};
-use super::esh::ESH;
+use super::esh::{ESHValue, ESH};
 use super::stream::{ReadStream, WriteStream};
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 
 pub const NO_FLAGS: u32 = 0;
 pub const NO_ESH: usize = 0xFFFF;
@@ -14,6 +15,35 @@ pub struct Entity {
     enc_size: usize,
 }
 
+impl Entity {
+    pub fn get_esh(&self) -> Result<&ESH> {
+        match &self.esh {
+            Some(esh) => Ok(esh),
+            None => Err(anyhow!("Entity has no ESH")),
+        }
+    }
+
+    pub fn get_esh_mut(&mut self) -> Result<&mut ESH> {
+        match &mut self.esh {
+            Some(esh) => Ok(esh),
+            None => Err(anyhow!("Entity has no ESH")),
+        }
+    }
+
+    pub fn get_attributes(&self) -> Result<Attributes> {
+        let value = match self.get_esh()?.get("Attributes") {
+            Some(value) => value,
+            None => return Err(anyhow!("Entity has no Attributes")),
+        };
+
+        if let ESHValue::Binary(bin) = value {
+            Ok(Attributes::from_binary(&bin)?)
+        } else {
+            Err(anyhow!("Attributes is not binary"))
+        }
+    }
+}
+
 impl DecoderCtx<&mut EntityList, &EntityList> for Entity {
     fn decode<'a>(rd: &mut ReadStream<'a>, ctx: &mut EntityList) -> Result<Self> {
         let offset = rd.offset();

src/fot/esh.rs

@@ -255,6 +255,16 @@ pub struct ESH {
     enc_size: usize,
 }
 
+impl ESH {
+    pub fn get(&self, name: &str) -> Option<&ESHValue> {
+        self.props.get(name)
+    }
+
+    pub fn set(&mut self, name: &str, value: ESHValue) {
+        self.props[name] = value;
+    }
+}
+
 impl Decoder for ESH {
     fn decode<'a>(rd: &mut ReadStream<'a>) -> Result<Self> {
         let offset = rd.offset();

src/fot/world.rs

@@ -13,8 +13,8 @@ use anyhow::Result;
 use deflate::deflate_bytes_zlib;
 use inflate::inflate_bytes_zlib;
 
-use super::esh::{ESHValue, ESH};
 use super::attributes::Attributes;
+use super::esh::{ESHValue, ESH};
 use std::path::Path;
 
 pub struct World {
@@ -39,17 +39,15 @@ impl World {
     pub fn test(&mut self) -> Result<()> {
         //let actor_type = self.entlist.get_type_idx("Actor").unwrap();
         let ent = self.entlist.get_entity(2122);
-        let esh = ent.esh.as_ref().unwrap();
+        let esh = ent.get_esh()?;
         for (name, value) in &esh.props {
             println!("{} {}", name, value);
         }
         //self.entlist.dump_to_entfile(ent, Path::new("D:\\actor.ent"))?;
 
         println!("");
-        if let ESHValue::Binary(binary) = &esh.props["Attributes"] {
-            let attributes = Attributes::from_binary(&binary)?;
+        let attributes = ent.get_attributes()?;
         dbg!(attributes);
-        }
 
         Ok(())
     }