Improve object reader/writer. Add compression.

libbsv/src/simple_db/object.rs

@@ -22,6 +22,12 @@ use std::io::{Read, Write, copy};
 
 use digest::Digest;
 
+use flate2::{
+    Compression,
+    write::GzEncoder,
+    read::GzDecoder,
+};
+
 use crate::core::error::*;
 use crate::core::ObjectId;
 
@@ -32,37 +38,43 @@ pub const TYPE_BLOB: &ObjectType = b"blob";
 pub const TYPE_TREE: &ObjectType = b"tree";
 
 
-pub struct ObjectWriter<'a, W: Write, D: Digest> {
-    writer: &'a mut W,
+pub struct ObjectWriter<W: Write, D: Digest> {
+    writer: GzEncoder<W>,
     digest: D,
 }
 
 
-impl<'a, W: Write, D: Digest> ObjectWriter<'a, W, D> {
-    pub fn new(writer: &'a mut W, otype: &ObjectType, size: u64)
-        -> Result<ObjectWriter<'a, W, D>> {
+impl<W: Write, D: Digest> ObjectWriter<W, D> {
+    pub fn new(writer: W, otype: &ObjectType, size: u64)
+        -> Result<ObjectWriter<W, D>> {
 
         let mut digest = D::new();
 
         digest.update(otype);
         digest.update(&size.to_be_bytes());
 
-        writer.write_all(otype)?;
-        writer.write_all(&size.to_be_bytes())?;
+        let mut zwriter = GzEncoder::new(writer, Compression::default());
+
+        zwriter.write_all(otype)?;
+        zwriter.write_all(&size.to_be_bytes())?;
 
         Ok(ObjectWriter {
-            writer,
+            writer: zwriter,
             digest,
         })
     }
 
-    pub fn close(self) -> Result<ObjectId> {
-        Ok(ObjectId::new(&self.digest.finalize()))
+    pub fn finish(mut self) -> Result<(ObjectId, W)> {
+        self.writer.try_finish()?;
+        Ok((
+            ObjectId::new(&self.digest.finalize()),
+            self.writer.finish()?,
+        ))
     }
 }
 
 
-impl<'a, W: Write, D: Digest> Write for ObjectWriter<'a, W, D> {
+impl<W: Write, D: Digest> Write for ObjectWriter<W, D> {
     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
         let size = self.writer.write(buf)?;
         self.digest.update(&buf[..size]);
@@ -75,9 +87,9 @@ impl<'a, W: Write, D: Digest> Write for ObjectWriter<'a, W, D> {
 }
 
 
-pub fn write_object<'a, R, W>(reader: &mut R, writer: &'a mut W,
+pub fn write_object<R, W>(reader: &mut R, writer: W,
         otype: &ObjectType, size: u64)
-    -> Result<ObjectId> where
+    -> Result<(ObjectId, W)> where
         R: Read,
         W: Write {
     let mut owriter: ObjectWriter<W, sha2::Sha224>
@@ -85,21 +97,23 @@ pub fn write_object<'a, R, W>(reader: &mut R, writer: &'a mut W,
 
     copy(reader, &mut owriter)?;
 
-    owriter.close()
+    owriter.finish()
 }
 
 
 pub struct ObjectReader<R: Read> {
-    reader: R,
+    reader: GzDecoder<R>,
 }
 
 
 impl<R: Read> ObjectReader<R> {
-    pub fn new(mut reader: R)
+    pub fn new(reader: R)
         -> Result<(ObjectType, u64, ObjectReader<R>)> {
 
+        let mut zreader = GzDecoder::new(reader);
+
         let mut buffer = [0u8; 12];
-        reader.read_exact(&mut buffer)?;
+        zreader.read_exact(&mut buffer)?;
 
         let otype = {
             let mut otype = [0; 4];
@@ -116,13 +130,13 @@ impl<R: Read> ObjectReader<R> {
             otype,
             size,
             ObjectReader {
-                reader,
+                reader: zreader,
             }
         ))
     }
 
-    pub fn close(self) -> Result<()> {
-        Ok(())
+    pub fn close(self) -> Result<R> {
+        Ok(self.reader.into_inner())
     }
 }
 
@@ -147,9 +161,12 @@ mod tests {
         let mut fake_file = Cursor::new(vec![]);
         let mut output = vec![];
 
-        write_object(&mut source, &mut fake_file, TYPE_BLOB, payload.len() as u64).unwrap();
-        fake_file.seek(SeekFrom::Start(0)).unwrap();
+        let (oid, _) = write_object(&mut source, &mut fake_file,
+            TYPE_BLOB, payload.len() as u64).unwrap();
 
+        assert_eq!(oid, ObjectId::from_str("c3b4032160b015b2261530532a6c49f2bdadbe0687fb1f5a6a32e083").unwrap());
+
+        fake_file.seek(SeekFrom::Start(0)).unwrap();
         let (otype, size, mut reader) = ObjectReader::new(fake_file).unwrap();
         let read_size = reader.read_to_end(&mut output).unwrap();