Always save comments

src/main.rs

@@ -2,10 +2,14 @@ use std::{
     collections::{hash_map::Entry, HashMap},
     env,
     fs::File,
-    io::{prelude::*, BufReader},
+    io::{prelude::*, BufReader, Error, ErrorKind, Result},
     net::{TcpListener, TcpStream},
 };
 
+use data_stream::{
+    collections::SizeSettings, default_settings::PortableSettings, from_stream, to_stream,
+    FromStream, ToStream,
+};
 use maud::html;
 use percent_encoding::percent_decode_str;
 use pukram2html::{convert, convert_extended, convert_subheader, Settings};
@@ -14,7 +18,12 @@ mod request;
 use request::Request;
 
 fn main() {
-    let mut context = Context::default();
+    let mut context = File::open("context.dat")
+        .map(|mut file| {
+            from_stream::<PortableSettings, Context, _>(&mut file)
+                .unwrap_or_else(|_| panic!("Context file corrupted. Delete or repair it!"))
+        })
+        .unwrap_or_default();
     let mut args = env::args();
     args.next();
     let address = args.next().unwrap_or("127.0.0.1:8080".to_string());
@@ -56,6 +65,81 @@ struct Context {
     infos: HashMap<Box<str>, SiteInfo>,
 }
 
+impl<S: SizeSettings> ToStream<S> for Context {
+    fn to_stream<W: Write>(&self, stream: &mut W) -> Result<()> {
+        S::size_to_stream(self.infos.len(), stream)?;
+        for (key, info) in &self.infos {
+            let key_bytes = key.as_bytes();
+            S::size_to_stream(key_bytes.len(), stream)?;
+            stream.write_all(key_bytes)?;
+
+            S::size_to_stream(info.comments.len(), stream)?;
+            for Comment { name, text } in &info.comments {
+                let name_bytes = name.as_bytes();
+                S::size_to_stream(name_bytes.len(), stream)?;
+                stream.write_all(name_bytes)?;
+
+                let text_bytes = text.as_bytes();
+                S::size_to_stream(text_bytes.len(), stream)?;
+                stream.write_all(text_bytes)?;
+            }
+
+            S::size_to_stream(info.visits, stream)?;
+            S::size_to_stream(info.up, stream)?;
+            S::size_to_stream(info.down, stream)?;
+        }
+
+        Ok(())
+    }
+}
+
+impl<S: SizeSettings> FromStream<S> for Context {
+    fn from_stream<R: Read>(stream: &mut R) -> Result<Self> {
+        let size = S::size_from_stream(stream)?;
+        let infos = (0..size)
+            .map(|_| {
+                let key_bytes = <Vec<_> as FromStream<S>>::from_stream(stream)?;
+                let key = std::str::from_utf8(&key_bytes)
+                    .map_err(|e| Error::new(ErrorKind::InvalidData, e))?
+                    .into();
+
+                let size = S::size_from_stream(stream)?;
+                let comments = (0..size)
+                    .map(|_| {
+                        let name_bytes = <Vec<_> as FromStream<S>>::from_stream(stream)?;
+                        let name = std::str::from_utf8(&name_bytes)
+                            .map_err(|e| Error::new(ErrorKind::InvalidData, e))?
+                            .into();
+
+                        let text_bytes = <Vec<_> as FromStream<S>>::from_stream(stream)?;
+                        let text = std::str::from_utf8(&text_bytes)
+                            .map_err(|e| Error::new(ErrorKind::InvalidData, e))?
+                            .into();
+
+                        Ok(Comment { name, text })
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+
+                let visits = S::size_from_stream(stream)?;
+                let up = S::size_from_stream(stream)?;
+                let down = S::size_from_stream(stream)?;
+
+                Ok((
+                    key,
+                    SiteInfo {
+                        comments,
+                        visits,
+                        up,
+                        down,
+                    },
+                ))
+            })
+            .collect::<Result<HashMap<_, _>>>()?;
+
+        Ok(Self { infos })
+    }
+}
+
 impl Context {
     fn handle_connection(&mut self, mut stream: TcpStream) {
         let Some(request) = Request::from(&stream) else {
@@ -294,5 +378,12 @@ impl Context {
 
         let _ = writeln!(stream, "<hr>");
         let _ = writeln!(stream, "<a href=\"/{parent_path}\">&lt;&lt; Back</a>");
+
+        if let Ok(mut file) = File::create("context.dat") {
+            if to_stream::<PortableSettings, _, _>(self, &mut file).is_err() {
+                eprintln!("Error saving data!");
+                eprintln!();
+            }
+        }
     }
 }