implemented ESH encoding, but found out there was a critical bug in WriteStream. Needs to be fixed ASAP

src/fot.rs

@@ -1,10 +1,10 @@
 mod decoder;
+mod esh;
 mod fstring;
 mod raw;
 pub mod save;
 mod sgd;
 mod ssg;
-mod esh;
 mod stream;
 mod tag;
 mod world;

src/fot/esh.rs

@@ -1,21 +1,21 @@
-use super::raw::Raw;
-use super::tag::Tag;
 use super::decoder::Decoder;
 use super::fstring::FString;
+use super::raw::Raw;
 use super::stream::{ReadStream, WriteStream};
+use super::tag::Tag;
 use anyhow::Result;
 use indexmap::IndexMap;
 
 #[derive(Debug)]
 pub struct ESHUnknown {
     pub data_type: u32,
-    pub data: Vec<u8>
+    pub data: Vec<u8>,
 }
 
 #[derive(Debug)]
 pub struct ESHEntityFlags {
     pub entity_id: u16,
-    pub flags: u16
+    pub flags: u16,
 }
 
 #[derive(Debug)]
@@ -23,7 +23,7 @@ pub struct ESHFrame {
     pub unk1: Vec<u8>,
     pub a: f32,
     pub b: f32,
-    pub c: f32
+    pub c: f32,
 }
 
 #[derive(Debug)]
@@ -31,7 +31,7 @@ pub struct ESHRect {
     pub top: i32,
     pub left: i32,
     pub right: i32,
-    pub bottom: i32
+    pub bottom: i32,
 }
 
 #[derive(Debug)]
@@ -45,7 +45,7 @@ pub enum ESHValue {
     Binary(Vec<u8>),
     EntityFlags(ESHEntityFlags),
     Frame(ESHFrame),
-    Rect(ESHRect)
+    Rect(ESHRect),
 }
 
 impl ESHValue {
@@ -67,7 +67,7 @@ impl Decoder for ESHValue {
         let mut rd = ReadStream::new(raw, offset);
         let data_type = rd.read_u32()?;
         let data_size = rd.read_u32()?;
-        
+
         Ok(match data_type {
             Self::TYPE_BOOL => ESHValue::Bool(rd.read_u8()? == 1),
             Self::TYPE_FLOAT => ESHValue::Float(rd.read_f32()?),
@@ -77,26 +77,31 @@ impl Decoder for ESHValue {
             Self::TYPE_ESBIN => {
                 let bin_size = rd.read_u32()?;
                 ESHValue::Binary(rd.read_bytes(bin_size as usize)?)
-            },
+            }
             Self::TYPE_ENTTITYFLAGS => {
                 let entity_id = rd.read_u16()?;
                 let flags = rd.read_u16()?;
                 ESHValue::EntityFlags(ESHEntityFlags { entity_id, flags })
-            },
+            }
             Self::TYPE_FRAME => {
                 let unk1 = rd.read_bytes(0x24)?;
                 let c = rd.read_f32()? * 4.;
                 let b = rd.read_f32()? * 4.;
                 let a = rd.read_f32()? * 4.;
                 ESHValue::Frame(ESHFrame { unk1, a, b, c })
-            },
+            }
             Self::TYPE_RECT => {
                 let top = rd.read_i32()?;
                 let left = rd.read_i32()?;
                 let right = rd.read_i32()?;
                 let bottom = rd.read_i32()?;
-                ESHValue::Rect(ESHRect { top, left, right, bottom })
-            },
+                ESHValue::Rect(ESHRect {
+                    top,
+                    left,
+                    right,
+                    bottom,
+                })
+            }
             _ => {
                 let data = rd.read_bytes(data_size as usize)?;
                 ESHValue::Unknown(ESHUnknown { data_type, data })
@@ -105,22 +110,94 @@ impl Decoder for ESHValue {
     }
 
     fn encode(&self) -> Result<Raw> {
-        todo!();
+        let mut wd = WriteStream::new(self.get_enc_size());
+        match self {
+            ESHValue::Unknown(unk) => {
+                wd.write_u32(unk.data_type)?;
+                wd.write_u32(unk.data.len() as u32)?;
+
+                wd.write_bytes(&unk.data);
+            }
+            ESHValue::Bool(val) => {
+                wd.write_u32(Self::TYPE_BOOL)?;
+                wd.write_u32(1)?;
+
+                wd.write_u8(*val as u8)?;
+            }
+            ESHValue::Float(val) => {
+                wd.write_u32(Self::TYPE_FLOAT)?;
+                wd.write_u32(4)?;
+
+                wd.write_f32(*val)?;
+            }
+            ESHValue::Int(val) => {
+                wd.write_u32(Self::TYPE_INT)?;
+                wd.write_u32(4)?;
+
+                wd.write_i32(*val)?;
+            }
+            ESHValue::String(str) => {
+                wd.write_u32(Self::TYPE_STRING)?;
+                wd.write_u32(str.get_enc_size() as u32)?;
+
+                wd.write(str)?;
+            }
+            ESHValue::Sprite(spr) => {
+                wd.write_u32(Self::TYPE_SPRITE)?;
+                wd.write_u32(spr.get_enc_size() as u32)?;
+
+                wd.write(spr)?;
+            }
+            ESHValue::Binary(bin) => {
+                wd.write_u32(Self::TYPE_ESBIN)?;
+                wd.write_u32(bin.len() as u32)?;
+
+                wd.write_bytes(bin);
+            }
+            ESHValue::EntityFlags(eflags) => {
+                wd.write_u32(Self::TYPE_ENTTITYFLAGS)?;
+                wd.write_u32(4)?;
+
+                wd.write_u16(eflags.entity_id)?;
+                wd.write_u16(eflags.flags)?;
+            }
+            ESHValue::Frame(frame) => {
+                wd.write_u32(Self::TYPE_FRAME)?;
+                wd.write_u32(0x24)?;
+
+                wd.write_bytes(&frame.unk1);
+                wd.write_f32(frame.c / 4.)?;
+                wd.write_f32(frame.b / 4.)?;
+                wd.write_f32(frame.a / 4.)?;
+            }
+            ESHValue::Rect(rect) => {
+                wd.write_u32(Self::TYPE_RECT)?;
+                wd.write_u32(0x10)?;
+
+                wd.write_i32(rect.top)?;
+                wd.write_i32(rect.left)?;
+                wd.write_i32(rect.right)?;
+                wd.write_i32(rect.bottom)?;
+            }
+        };
+
+        Ok(wd.into_raw(0, 0))
     }
 
     fn get_enc_size(&self) -> usize {
-        Self::HDR_SIZE + match self {
-            ESHValue::Unknown(unk) => unk.data.len(),
-            ESHValue::Bool(_) => 1,
-            ESHValue::Float(_) => 4,
-            ESHValue::Int(_) => 4,
-            ESHValue::String(str) => str.get_enc_size(),
-            ESHValue::Sprite(spr) => spr.get_enc_size(),
-            ESHValue::Binary(bin) => bin.len(),
-            ESHValue::EntityFlags(_) => 4,
-            ESHValue::Frame(_) => 0x30,
-            ESHValue::Rect(_) => 0x10
-        }
+        Self::HDR_SIZE
+            + match self {
+                ESHValue::Unknown(unk) => unk.data.len(),
+                ESHValue::Bool(_) => 1,
+                ESHValue::Float(_) => 4,
+                ESHValue::Int(_) => 4,
+                ESHValue::String(str) => str.get_enc_size(),
+                ESHValue::Sprite(spr) => spr.get_enc_size(),
+                ESHValue::Binary(bin) => bin.len(),
+                ESHValue::EntityFlags(_) => 4,
+                ESHValue::Frame(_) => 0x30,
+                ESHValue::Rect(_) => 0x10,
+            }
     }
 }
 
@@ -128,7 +205,7 @@ impl Decoder for ESHValue {
 pub struct ESH {
     pub tag: Tag,
     pub props: IndexMap<FString, ESHValue>,
-    enc_size: usize
+    enc_size: usize,
 }
 
 impl Decoder for ESH {
@@ -145,14 +222,27 @@ impl Decoder for ESH {
         }
 
         let enc_size = rd.offset() - offset;
-        Ok(ESH { tag, props, enc_size })
+        Ok(ESH {
+            tag,
+            props,
+            enc_size,
+        })
     }
 
     fn encode(&self) -> Result<Raw> {
-        todo!();
+        let mut wd = WriteStream::new(self.get_enc_size());
+        wd.write(&self.tag)?;
+
+        wd.write_u32(self.props.len() as u32)?;
+        for (name,value) in self.props.iter() {
+            wd.write(name)?;
+            wd.write(value)?;
+        }
+
+        Ok(wd.into_raw(0, 0))
     }
 
     fn get_enc_size(&self) -> usize {
         self.enc_size
     }
-}
+}

src/fot/raw.rs

@@ -88,4 +88,16 @@ impl Raw {
         file.flush()?;
         Ok(())
     }
+
+    pub fn dump(&self, path: &Path) -> Result<()> {
+        let mut file = BufWriter::new(
+            OpenOptions::new()
+                .create(true)
+                .truncate(true)
+                .write(true)
+                .open(path)?,
+        );
+        file.write_all(&self.mem)?;
+        Ok(())
+    }
 }

src/fot/stream.rs

@@ -93,12 +93,18 @@ impl WriteStream {
         self.buf.position() as usize
     }
 
+    pub fn skip(&mut self, size: usize) {
+        self.buf.set_position(self.buf.position() + size as u64);
+    }
+
     pub fn write_bytes(&mut self, bytes: &[u8]) {
         self.buf.get_mut().extend(bytes.iter());
+        self.skip(bytes.len());
     }
 
     pub fn write<T: Decoder>(&mut self, val: &T) -> Result<()> {
         let mut raw = val.encode()?;
+        self.skip(raw.mem.len());
         self.buf.get_mut().append(&mut raw.mem);
         Ok(())
     }

src/fot/world.rs

@@ -1,9 +1,9 @@
 use super::decoder::Decoder;
+use super::esh::ESH;
 use super::fstring::FString;
 use super::raw::Raw;
 use super::sgd::SGD;
 use super::ssg::SSG;
-use super::esh::ESH;
 use super::stream::ReadStream;
 use super::tag::Tag;
 use anyhow::anyhow;
@@ -12,6 +12,7 @@ use byteorder::{LittleEndian, WriteBytesExt};
 use deflate::deflate_bytes_zlib;
 use inflate::inflate_bytes_zlib;
 use std::io::Cursor;
+use std::path::Path;
 
 #[derive(Debug)]
 pub struct World {
@@ -34,6 +35,10 @@ impl World {
         let esh: ESH = rd.read(0)?;
         dbg!(&esh);
 
+        let esh2 = esh.encode()?;
+        esh2.dump(Path::new("esh2.bin"))?;
+        assert_eq!(&self.data.mem[0x14AC..0x14AC+esh.get_enc_size()], &esh2.mem, "ESH encoding test passed");
+
         Ok(())
     }
 }