Zigzagging works, but uncommetning this quantization logic breaks things

Cargo.lock

@@ -119,6 +119,15 @@ dependencies = [
  "miniz_oxide",
 ]
 
+[[package]]
+name = "proc-macro2"
+version = "1.0.79"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e835ff2298f5721608eb1a980ecaee1aef2c132bf95ecc026a11b7bf3c01c02e"
+dependencies = [
+ "unicode-ident",
+]
+
 [[package]]
 name = "quickcheck"
 version = "1.0.3"
@@ -130,6 +139,15 @@ dependencies = [
  "rand",
 ]
 
+[[package]]
+name = "quote"
+version = "1.0.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "291ec9ab5efd934aaf503a6466c5d5251535d108ee747472c3977cc5acc868ef"
+dependencies = [
+ "proc-macro2",
+]
+
 [[package]]
 name = "rand"
 version = "0.8.5"
@@ -183,6 +201,7 @@ version = "0.1.0"
 dependencies = [
  "png",
  "quickcheck",
+ "thiserror",
 ]
 
 [[package]]
@@ -191,6 +210,43 @@ version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe"
 
+[[package]]
+name = "syn"
+version = "2.0.58"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44cfb93f38070beee36b3fef7d4f5a16f27751d94b187b666a5cc5e9b0d30687"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "thiserror"
+version = "1.0.58"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03468839009160513471e86a034bb2c5c0e4baae3b43f79ffc55c4a5427b3297"
+dependencies = [
+ "thiserror-impl",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.58"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
+
 [[package]]
 name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"

Cargo.toml

@@ -8,3 +8,4 @@ edition = "2021"
 [dependencies]
 png = "0.17.13"
 quickcheck = "1.0.3"
+thiserror = "1.0.58"

src/compression.rs

@@ -1,6 +1,6 @@
 use std::io::Write;
 
-use crate::{constants::{MAGIC, TILE_SZ}, transform};
+use crate::{constants::{MAGIC, TILE_SZ}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform};
 
 
 struct PixelTile {
@@ -11,14 +11,18 @@ struct CoefTile {
     coefs: [i32; TILE_SZ * TILE_SZ]
 }
 
+struct QuantTile {
+    quants: [i32; TILE_SZ * TILE_SZ]
+}
 
-pub fn compress(
+
-    width: u16,
+pub fn compress<W: Write>(
-    height: u16,
+    width: u32,
+    height: u32,
     // n_components: u16,
     layers: &[&[u8]],
-) -> Vec<u8> {
+    writer: &mut protocol::ProtocolWriter<W>
-    let mut output = vec![];
+) -> ProtocolWriterResult<()> {
 
     // validation
     for l in 0..layers.len() {
@@ -26,11 +30,10 @@ pub fn compress(
     }
 
     // write header
-    output.write(&MAGIC).unwrap();
+    writer.write_header(MAGIC)?;
-    output.write(&(width as u16).to_le_bytes()).unwrap();
+    writer.write_u32_wide(width)?;
-    output.write(&(height as u16).to_le_bytes()).unwrap();
+    writer.write_u32_wide(height)?;
-    // output.write(&n_components.to_le_bytes()).unwrap();
+    writer.write_u32_wide(layers.len() as u32)?;
-    output.write(&(layers.len() as u16).to_le_bytes()).unwrap();
 
     let mut tiles = vec![];
     for layer in layers.iter() {
@@ -41,19 +44,20 @@ pub fn compress(
                     width as usize, height as usize
                 );
                 let coef_tile = CoefTile::from_pixel_tile(&pixel_tile);
-                tiles.push(coef_tile);
+                let quant_tile = QuantTile::from_coef_tile(&coef_tile);
+                tiles.push(quant_tile);
             }
         }
     }
 
     for t in tiles.iter() {
-        t.write_zero(&mut output);
+        t.write_zero(writer)?;
     }
     for t in tiles.iter() {
-        t.write_rest(&mut output);
+        t.write_rest(writer)?;
     }
 
-    output
+    Ok(())
 }
 
 impl PixelTile {
@@ -98,14 +102,24 @@ impl CoefTile {
 
         return CoefTile { coefs }
     }
+}
 
-    fn write_zero(&self, output: &mut Vec<u8>) {
+impl QuantTile {
-        output.write(&self.coefs[0].to_le_bytes()).unwrap();
+    fn from_coef_tile(pt: &CoefTile) -> QuantTile {
-    }
+        QuantTile {
-
+            quants: quantization::to_quantized(pt.coefs)
-    fn write_rest(&self, output: &mut Vec<u8>) {
-        for i in 1..self.coefs.len() {
-            output.write(&self.coefs[i].to_le_bytes()).unwrap();
         }
     }
+
+    fn write_zero<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
+        writer.write_i32_packed(self.quants[0])?;
+        Ok(())
+    }
+
+    fn write_rest<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
+        for i in 1..self.quants.len() {
+            writer.write_i32_packed(self.quants[i])?;
+        }
+        Ok(())
+    }
 }

src/constants.rs

@@ -1,2 +1,7 @@
-pub const MAGIC: [u8; 4] = *b"rxi\0";
+use crate::protocol::Header;
+
+pub const MAGIC: Header = Header { 
+    bytes: *b"rxi",
+    version: 0
+};
 pub const TILE_SZ: usize = 8;

src/decompression.rs

@@ -1,6 +1,6 @@
-use std::io::{BufReader, Read};
+use std::io::Read;
 
-use crate::{constants::{MAGIC, TILE_SZ}, transform};
+use crate::{constants::{MAGIC, TILE_SZ}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform};
 
 struct PixelTile {
     // i32: representation that supports Walsh-Hadamard
@@ -12,35 +12,37 @@ struct CoefTile {
     coefs: [i32; TILE_SZ * TILE_SZ]
 }
 
-pub fn decompress(
+#[derive(Clone)]
-    input: &[u8]
+struct QuantTile {
-) -> (u16, u16, Vec<Vec<u8>>) {
+    quants: [i32; TILE_SZ * TILE_SZ]
-    let mut reader = BufReader::new(input);
+}
 
+pub fn decompress<R: Read>(
+    reader: &mut ProtocolReader<R>
+) -> ProtocolReaderResult<(u32, u32, Vec<Vec<u8>>)> {
     // read header
-    let mut magic_buffer = [0; MAGIC.len()];
+    let header = reader.read_header()?;
-    reader.read_exact(&mut magic_buffer).unwrap();
+    if header.bytes != MAGIC.bytes {
+        return Err(ProtocolReaderError::WrongHeader);
+    }
+    if header.version != MAGIC.version {
+        return Err(ProtocolReaderError::WrongVersion);
+    }
 
-    assert_eq!(MAGIC, magic_buffer);
+    let width = reader.read_u32_wide()?;
-
+    let height = reader.read_u32_wide()?;
-    let mut u16_buffer = [0; 2];
+    let n_layers = reader.read_u32_wide()?;
-    reader.read_exact(&mut u16_buffer).unwrap();
-    let width = u16::from_le_bytes(u16_buffer);
-    reader.read_exact(&mut u16_buffer).unwrap();
-    let height = u16::from_le_bytes(u16_buffer);
-    reader.read_exact(&mut u16_buffer).unwrap();
-    let n_layers = u16::from_le_bytes(u16_buffer);
 
     let width_in_tiles = (width as usize + TILE_SZ - 1) / TILE_SZ;
     let height_in_tiles = (height as usize + TILE_SZ - 1) / TILE_SZ;
 
     let n_tiles = width_in_tiles * height_in_tiles * (n_layers as usize);
-    let mut tiles = vec![CoefTile::new(); n_tiles];
+    let mut tiles = vec![QuantTile::new(); n_tiles];
     for i in 0..n_tiles {
-        tiles[i].load_zero(&mut reader);
+        tiles[i].load_zero(reader)?;
     }
     for i in 0..n_tiles {
-        tiles[i].load_rest(&mut reader);
+        tiles[i].load_rest(reader)?;
     }
 
     let mut tile_i = 0;
@@ -48,7 +50,7 @@ pub fn decompress(
     for layer in 0..n_layers {
         for x0 in (0..width).step_by(TILE_SZ) {
             for y0 in (0..height).step_by(TILE_SZ) {
-                let pixel_tile = tiles[tile_i].to_pixel_tile();
+                let pixel_tile = tiles[tile_i].to_coef_tile().to_pixel_tile();
                 tile_i += 1;
                 pixel_tile.to_layer(
                     x0 as usize, y0 as usize, &mut layers[layer as usize], 
@@ -57,28 +59,10 @@ pub fn decompress(
             }
         }
     }
-    return (width, height, layers)
+    return Ok((width, height, layers))
 }
 
 impl CoefTile {
-    fn new() -> CoefTile {
-        CoefTile { coefs: [0; TILE_SZ * TILE_SZ] }
-    }
-    
-    fn load_zero(&mut self, reader: &mut BufReader<&[u8]>) {
-        let mut coef_zero_buffer = [0; 4];
-        reader.read_exact(&mut coef_zero_buffer).unwrap();
-        self.coefs[0] = i32::from_le_bytes(coef_zero_buffer);
-    }
-
-    fn load_rest(&mut self, reader: &mut BufReader<&[u8]>) {
-        let mut coef_buffer = [0; 4];
-        for i in 1..self.coefs.len() {
-            reader.read_exact(&mut coef_buffer).unwrap();
-            self.coefs[i] = i32::from_le_bytes(coef_buffer);
-        }
-    }
-    
     fn to_pixel_tile(&self) -> PixelTile {
         let mut pixels = self.coefs.clone();
 
@@ -109,3 +93,28 @@ impl PixelTile {
         }
     }
 }
+
+impl QuantTile {
+    fn new() -> QuantTile {
+        QuantTile { quants: [0; TILE_SZ * TILE_SZ] }
+    }
+    
+    fn to_coef_tile(&self) -> CoefTile {
+        CoefTile {
+            coefs: quantization::from_quantized(self.quants)
+        }
+    }
+
+    fn load_zero<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> {
+        self.quants[0] = reader.read_i32_packed()?;
+        Ok(())
+    }
+
+    fn load_rest<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> {
+        for i in 1..self.quants.len() {
+            self.quants[i] = reader.read_i32_packed()?;
+        }
+        Ok(())
+    }
+    
+}

src/main.rs

@@ -1,10 +1,14 @@
-use std::fs::File;
+use std::{fs::File, io::{Cursor, Write}};
 
 use png::{BitDepth, ColorType};
 
+use crate::protocol::{ProtocolReader, ProtocolWriter};
+
 mod compression;
 mod constants;
 mod decompression;
+mod protocol;
+mod quantization;
 mod transform;
 
 
@@ -25,16 +29,19 @@ fn main() {
 }
 fn hard_main() {
     let (width, height, r, g, b) = load_image();
-    let compressed = 
+    let mut writer = ProtocolWriter::new(vec![]);
-        compression::compress(width as u16, height as u16, &[&r, &g, &b]);
+    compression::compress(
+        width as u32, height as u32, &[&r, &g, &b],
+        &mut writer
+    ).unwrap();
+    let compressed = writer.destroy();
 
-    /*
     let mut output_file = File::create("outputs/avatar2.rxi").unwrap();
     output_file.write_all(&compressed).unwrap();
-    */
 
+    let mut reader = ProtocolReader::new(Cursor::new(compressed));
     let (width2, height2, decompressed) = 
-        decompression::decompress(&compressed);
+        decompression::decompress(&mut reader).unwrap();
 
     assert_eq!(3, decompressed.len());
     save_image(width2 as usize, height2 as usize, decompressed);

src/protocol.rs

@@ -0,0 +1,138 @@
+use std::io::{self, ErrorKind, Read, Write};
+
+use thiserror::Error;
+
+#[derive(Clone, Copy)]
+pub struct Header {
+    pub bytes: [u8; 3],
+    pub version: u8
+}
+
+
+#[derive(Error, Debug)]
+pub enum ProtocolWriterError {
+    #[error("error using underlying writer")]
+    Io(#[from] io::Error)
+}
+
+pub type ProtocolWriterResult<T> = Result<T, ProtocolWriterError>;
+
+pub struct ProtocolWriter<W: Write> {
+    writer: W
+}
+
+#[derive(Error, Debug)]
+pub enum ProtocolReaderError {
+    // this is explicitly supported: rxi images are progressive
+    #[error("EOF before end of image")]  
+    EarlyEof,
+
+    #[error("wrong header")]
+    WrongHeader,
+
+    #[error("header has unsupported version")]
+    WrongVersion,
+
+    #[error("error using underlying reader")]
+    Io(#[from] io::Error),
+}
+
+pub type ProtocolReaderResult<T> = Result<T, ProtocolReaderError>;
+
+pub struct ProtocolReader<R: Read> {
+    reader: R
+}
+
+impl<W: Write> ProtocolWriter<W> {
+    pub fn new(writer: W) -> Self {
+        Self { writer }
+    }
+
+    pub fn destroy(self) -> W {
+        self.writer
+    }
+
+    pub fn write_header(&mut self, value: Header) -> ProtocolWriterResult<()> {
+        self.writer.write_all(&value.bytes)?;
+        self.writer.write_all(&[value.version])?;
+        Ok(())
+    }
+    pub fn write_u32_wide(&mut self, value: u32) -> ProtocolWriterResult<()> {
+        self.writer.write_all(&value.to_le_bytes())?;
+        Ok(())
+    }
+    pub fn write_i32_packed(&mut self, value: i32) -> ProtocolWriterResult<()> {
+        // We reserve i8::MAX and i16::MAX as signal values that `value` is too big
+        // for each respective type
+        if (i8::MIN as i32..i8::MAX as i32).contains(&value) {
+            self.writer.write_all(&(value as i8).to_le_bytes())?;
+            return Ok(())
+        } 
+        self.writer.write_all(&(i8::MAX).to_le_bytes())?;
+
+        if (i16::MIN as i32..i16::MAX as i32).contains(&value) {
+            self.writer.write_all(&(value as i16).to_le_bytes())?;
+            return Ok(())
+        } 
+        self.writer.write_all(&(i16::MAX).to_le_bytes())?;
+        self.writer.write_all(&value.to_le_bytes())?;
+        Ok(())
+    }
+}
+
+impl<R: Read> ProtocolReader<R> {
+    pub fn new(reader: R) -> Self {
+        Self { reader }
+    }
+
+    pub fn read_header(&mut self) -> ProtocolReaderResult<Header> {
+        let mut bytes_buf = [0; 3];
+        self.read_exact(&mut bytes_buf)?;
+
+        let mut version_buf = [0; 1];
+        self.read_exact(&mut version_buf)?;
+
+        Ok(Header {
+            bytes: bytes_buf,
+            version: version_buf[0],
+        })
+    }
+
+    pub fn read_u32_wide(&mut self) -> ProtocolReaderResult<u32> {
+        let mut u32_buf = [0; 4];
+        self.read_exact(&mut u32_buf)?;
+
+        Ok(u32::from_le_bytes(u32_buf))
+    }
+
+    pub fn read_i32_packed(&mut self) -> ProtocolReaderResult<i32> {
+        let mut i8_buf = [0; 1];
+        self.read_exact(&mut i8_buf)?;
+        let i8_value = i8::from_le_bytes(i8_buf);
+        if i8_value != i8::MAX { return Ok(i8_value as i32) }
+
+        let mut i16_buf = [0; 2];
+        self.read_exact(&mut i16_buf)?;
+        let i16_value = i16::from_le_bytes(i16_buf);
+        if i16_value != i16::MAX { return Ok(i16_value as i32) }
+
+        let mut i32_buf = [0; 4];
+        self.read_exact(&mut i32_buf)?;
+        let i32_value = i32::from_le_bytes(i32_buf);
+        Ok(i32_value)
+    }
+
+    // wrap UnexpectedEof, since that's an error for other Io clients,
+    // but it's specifically not an error for us
+    fn read_exact(&mut self, buf: &mut [u8]) -> ProtocolReaderResult<()> {
+        match self.reader.read_exact(buf) {
+            Ok(_) => Ok(()),
+            Err(e) => {
+                if e.kind() == ErrorKind::UnexpectedEof {
+                    return Err(ProtocolReaderError::EarlyEof)
+                }
+                return Err(ProtocolReaderError::Io(e));
+            }
+        }
+    }
+}

src/quantization.rs

@@ -0,0 +1,47 @@
+use crate::constants::TILE_SZ;
+
+const ZIGZAG: [u8; TILE_SZ * TILE_SZ] = [
+    0 , 1 , 5 , 6 , 14, 15, 27, 28,
+    2 , 4 , 7 , 13, 16, 26, 29, 42,
+    3 , 8 , 12, 17, 25, 30, 41, 43,
+    9 , 11, 18, 24, 31, 40, 44, 53,
+    10, 19, 23, 32, 39, 45, 52, 54,
+    20, 22, 33, 38, 46, 51, 55, 60,
+    21, 34, 37, 47, 50, 56, 59, 61,
+    35, 36, 48, 49, 57, 58, 62, 63,
+];
+const DIVISORS: [u8; TILE_SZ * TILE_SZ] = [
+    // source: https://en.wikipedia.org/wiki/Quantization_(image_processing)
+    16, 11, 10, 16, 24, 40, 51, 61,
+    12, 12, 14, 19, 26, 58, 60, 55,
+    14, 13, 16, 24, 40, 57, 69, 56,
+    14, 17, 22, 29, 51, 87, 80, 62,
+    18, 22, 37, 56, 68, 109, 103, 77,
+    24, 35, 55, 64, 81, 104, 113, 92,
+    49, 64, 78, 87, 103, 121, 120, 101,
+    72, 92, 95, 98, 112, 100, 103, 99,
+];
+
+pub fn to_quantized(
+    coefs: [i32; TILE_SZ * TILE_SZ]
+) -> [i32; TILE_SZ * TILE_SZ] {
+    let mut quant: [i32; TILE_SZ * TILE_SZ] = [0; TILE_SZ * TILE_SZ];
+
+    for cf_ix in 0..TILE_SZ * TILE_SZ {
+        quant[ZIGZAG[cf_ix] as usize] = 
+            coefs[cf_ix]
+            // (coefs[cf_ix] + DIVISORS[cf_ix] as i32 / 2) / 
+            // (DIVISORS[cf_ix] as i32);
+    }
+    quant
+}
+
+pub fn from_quantized(
+    quant: [i32; TILE_SZ* TILE_SZ]
+) -> [i32; TILE_SZ * TILE_SZ] {
+    let mut coefs: [i32; TILE_SZ * TILE_SZ] = [0; TILE_SZ * TILE_SZ];
+    for cf_ix in 0..TILE_SZ * TILE_SZ {
+        coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize]; // * DIVISORS[cf_ix] as i32;
+    }
+    coefs
+}