Unifying label formats

4 files changed, 124 insertions, 131 deletions

diff --git a/src/binfile.rs b/src/binfile.rs
index 5b0b783..b61945a 100644
--- a/src/binfile.rs
+++ b/src/binfile.rs
@@ -1,18 +1,5 @@
 use std::io::{Error,ErrorKind};
 
-pub fn dump_u32_vec(file_name: &str, data: Vec<u32>) -> Result<(), Error> {
-    use std::fs::File;
-    use std::io::Write;
-    if let Ok(mut file) = File::create(file_name) {
-        for value in data {
-            file.write_all(&value.to_le_bytes())?;
-        }
-    } else {
-        return Err(Error::new(ErrorKind::Other, "Error creating file"));
-    }
-    Ok(())
-}
-
 pub fn dump_u16_vec(file_name: &str, data: Vec<u16>) -> Result<(), Error> {
     use std::fs::File;
     use std::io::Write;
diff --git a/src/label_format.rs b/src/label_format.rs
index 38c5d19..c127e81 100644
--- a/src/label_format.rs
+++ b/src/label_format.rs
@@ -1,12 +1,65 @@
-use crate::flood_u16;
+use core::cmp::PartialEq;
+use core::marker::Copy;
+use crate::flood;
+use crate::color;
 
-pub struct LabelFormat {
-    pub buffer: Vec<u16>,
+pub trait LabelU16 {
+    fn zero(&self) -> Self;
+}
+
+impl LabelU16 for u16 {
+    fn zero(&self) -> Self {
+        0u16
+    }
+}
+
+impl LabelU16 for u32 {
+    fn zero(&self) -> Self {
+        0u32
+    }
+}
+
+pub struct LabelFormat<T: LabelU16 + PartialEq> {
+    pub buffer: Vec<T>,
     pub width: usize,
     pub height: usize,
 }
 
-impl LabelFormat {
+impl<T: LabelU16 + PartialEq> std::fmt::Debug for LabelFormat<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LabelFormat")
+            .field("width", &self.width)
+            .field("height", &self.height)
+            .finish()
+    }
+}
+
+impl<T: LabelU16 + PartialEq + Copy> LabelFormat<T> {
+    pub fn compress(&self) -> LabelFormat<u16> {
+        let mut label: u16 = 1;
+        let mut output_buffer: Vec<u16> = vec![0u16; self.buffer.len()];
+        for y in 0..self.height {
+            for x in 0..self.width {
+                let index = x + y*self.width;
+                if self.buffer[index] == self.buffer[index].zero() {
+                    continue;
+                }
+                if output_buffer[index] == 0 {
+                    let color = self.buffer[index];
+                    flood(&self, &mut output_buffer, x, y, color, label);
+                    label += 1;
+                }
+            }
+        }
+        LabelFormat::<u16> {
+            buffer: output_buffer,
+            width: self.width,
+            height: self.height,
+        }
+    }
+}
+
+impl LabelFormat<u16> {
     pub fn get_left(&self, x: usize, y: usize) -> Option<u16> {
         if x > 0 {
             return Some(self.buffer[(x-1) + y * self.width]);
@@ -52,7 +105,7 @@ impl LabelFormat {
         self.buffer = dilated.buffer;
     }
 
-    pub fn dilate(&self) -> LabelFormat {
+    pub fn dilate(&self) -> LabelFormat<u16> {
         let width = self.width;
         let height = self.height;
         let mut output_buffer = vec![0u16; width*height];
@@ -90,7 +143,7 @@ impl LabelFormat {
                 }
             }
         }
-        LabelFormat {
+        LabelFormat::<u16> {
             buffer: output_buffer,
             width,
             height,
@@ -102,7 +155,7 @@ impl LabelFormat {
         self.buffer = eroded.buffer;
     }
 
-    pub fn erode(&self) -> LabelFormat {
+    pub fn erode(&self) -> LabelFormat<u16> {
         let width = self.width;
         let height = self.height;
         let mut output_buffer = vec![0u16; width*height];
@@ -136,7 +189,7 @@ impl LabelFormat {
                 output_buffer[index] = current_color;
             }
         }
-        LabelFormat {
+        LabelFormat::<u16> {
             buffer: output_buffer,
             width,
             height,
@@ -152,7 +205,7 @@ impl LabelFormat {
         }
     }
 
-    pub fn closeup(&self, n: usize) -> LabelFormat {
+    pub fn closeup(&self, n: usize) -> LabelFormat<u16> {
         let mut x = self.erode();
         x.icloseup(n-1);
         x.idilate();
@@ -168,7 +221,7 @@ impl LabelFormat {
         }
     }
 
-    pub fn fill(&self, n: usize) -> LabelFormat {
+    pub fn fill(&self, n: usize) -> LabelFormat<u16> {
         let mut x = self.dilate();
         x.icloseup(n-1);
         x.ierode();
@@ -179,7 +232,7 @@ impl LabelFormat {
         crate::binfile::dump_u16_vec(filename, self.buffer.clone())
     }
 
-    pub fn combine(&self, other: &LabelFormat) -> Option<LabelFormat> {
+    pub fn combine(&self, other: &LabelFormat<u16>) -> Option<LabelFormat<u16>> {
         let width = self.width;
         let height = self.height;
         let rwidth = other.width;
@@ -187,7 +240,7 @@ impl LabelFormat {
         if (width != rwidth) || (height != rheight) {
             return None;
         }
-        let mut output = LabelFormat {
+        let mut output = LabelFormat::<u16> {
             buffer: vec![0u16; width * height],
             width,
             height,
@@ -204,7 +257,7 @@ impl LabelFormat {
         return Some(output);
     }
 
-    pub fn refresh(&self) -> LabelFormat {
+    pub fn refresh(&self) -> LabelFormat<u16> {
         let mut label: u16 = 1;
         let mut output_buffer: Vec<u16> = vec![0u16; self.buffer.len()];
         for y in 0..self.height {
@@ -215,12 +268,12 @@ impl LabelFormat {
                 }
                 if output_buffer[index] == 0 {
                     let color = self.buffer[index];
-                    flood_u16(self, &mut output_buffer, x, y, color, label);
+                    flood(self, &mut output_buffer, x, y, color, label);
                     label += 1;
                 }
             }
         }
-        return LabelFormat {
+        return LabelFormat::<u16> {
             buffer: output_buffer,
             width: self.width,
             height: self.height,
@@ -241,35 +294,41 @@ impl LabelFormat {
 
 }
 
-mod color {
-    pub(crate) fn reset_color() {
-        print!("\x1b[0m");
-    }
-
-    // Set background color from a number
-    pub(crate) fn set_color(color: usize) {
-        if color == 0 {
-            reset_color();
-            return;
-        }
-        let paint_color = color-1;
-        let paint_color = paint_color % 13;
-        if paint_color < 7 {
-            print!("\x1b[{}m", 40 + (paint_color+1));
-        } else {
-            print!("\x1b[{}m", 100 + (paint_color+1-7));
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
+    fn label_compress_test() {
+        const DIM: usize = 6;
+        let mut test_data = LabelFormat::<u32> {
+            buffer: vec![0u32; DIM*DIM],
+            width: DIM,
+            height: DIM,
+        };
+
+        let expected: [Vec<u16>;1] = [
+            [
+                0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0,
+                0, 0, 1, 1, 2, 0,
+                0, 0, 0, 0, 0, 0,
+                0, 0, 0, 0, 0, 0,
+            ].to_vec(),
+        ];
+
+        test_data.buffer[2+3*DIM] = 2;
+        test_data.buffer[3+3*DIM] = 2;
+        test_data.buffer[4+3*DIM] = 1;
+        let compressed_data = test_data.compress();
+        assert_eq!(compressed_data.buffer, expected[0]);
+    }
+
+    #[test]
     fn label_morphology_test() {
         const DIM: usize = 6;
-        let mut test_data = LabelFormat {
+        let mut test_data = LabelFormat::<u16> {
             buffer: vec![0u16; DIM*DIM],
             width: DIM,
             height: DIM,
@@ -317,7 +376,7 @@ mod tests {
     #[test]
     fn mutable_morphology_test() {
         const DIM: usize = 6;
-        let mut test_data = LabelFormat {
+        let mut test_data = LabelFormat::<u16> {
             buffer: vec![0u16; DIM*DIM],
             width: DIM,
             height: DIM,
@@ -365,7 +424,7 @@ mod tests {
     #[test]
     fn combination_test() {
         const DIM: usize = 6;
-        let mut test_data_1 = LabelFormat {
+        let mut test_data_1 = LabelFormat::<u16> {
             buffer: vec![0u16; DIM*DIM],
             width: DIM,
             height: DIM,
@@ -410,7 +469,7 @@ mod tests {
         test_data_1.buffer[3+3*DIM] = 1;
         test_data_1.buffer[4+3*DIM] = 2;
 
-        let mut test_data_2 = LabelFormat {
+        let mut test_data_2 = LabelFormat::<u16> {
             buffer: vec![0u16; DIM*DIM],
             width: DIM,
             height: DIM,
diff --git a/src/large_label_format.rs b/src/large_label_format.rs
deleted file mode 100644
index 0a5c458..0000000
--- a/src/large_label_format.rs
+++ /dev/null
@@ -1,45 +0,0 @@
-use crate::{LabelFormat,flood_u32};
-
-pub struct LargeLabelFormat {
-    pub buffer: Vec<u32>,
-    pub width: usize,
-    pub height: usize,
-}
-
-impl std::fmt::Debug for LargeLabelFormat {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("LargeLabelFormat")
-            .field("width", &self.width)
-            .field("height", &self.height)
-            .finish()
-    }
-}
-
-impl LargeLabelFormat {
-    pub fn compress(&self) -> LabelFormat {
-        let mut label: u16 = 1;
-        let mut output_buffer: Vec<u16> = vec![0u16; self.buffer.len()];
-        for y in 0..self.height {
-            for x in 0..self.width {
-                let index = x + y*self.width;
-                if self.buffer[index] == 0 {
-                    continue;
-                }
-                if output_buffer[index] == 0 {
-                    let color = self.buffer[index];
-                    flood_u32(&self, &mut output_buffer, x, y, color, label);
-                    label += 1;
-                }
-            }
-        }
-        return LabelFormat {
-            buffer: output_buffer,
-            width: self.width,
-            height: self.height,
-        };
-    }
-
-    pub fn dump(&self, filename: &str) -> Result<(), std::io::Error> {
-        crate::binfile::dump_u32_vec(filename, self.buffer.clone())
-    }
-}
diff --git a/src/lib.rs b/src/lib.rs
index 46c7cd6..7278c45 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,16 +1,14 @@
-mod large_label_format;
 mod label_format;
 pub mod binfile;
 
-pub use large_label_format::LargeLabelFormat;
-pub use label_format::LabelFormat;
+use crate::label_format::{LabelFormat,LabelU16};
 
-pub(crate) fn flood_u32(source: &LargeLabelFormat, destination: &mut Vec<u16>, 
+pub(crate) fn flood<T: LabelU16 + PartialEq + Copy, R: LabelU16 + PartialEq + Copy>(source: &LabelFormat<T>, destination: &mut Vec<R>, 
     x: usize, y: usize,
-    from_color: u32, to_color: u16) {
+    from_color: T, to_color: R) {
     let width = source.width;
     let destination_color = destination[x + y * width];
-    if destination_color != 0 {
+    if destination_color != destination_color.zero() {
         return;
     }
     let source_color = source.buffer[x + y * width];
@@ -19,42 +17,36 @@ pub(crate) fn flood_u32(source: &LargeLabelFormat, destination: &mut Vec<u16>,
     }
     destination[x + y * width] = to_color;
     if x > 0 {
-        flood_u32(source, destination, x-1, y, from_color, to_color);
+        flood(source, destination, x-1, y, from_color, to_color);
     }
     if (x+1) < width {
-        flood_u32(source, destination, x+1, y, from_color, to_color);
+        flood(source, destination, x+1, y, from_color, to_color);
     }
     if y > 0 {
-        flood_u32(source, destination, x, y-1, from_color, to_color);
+        flood(source, destination, x, y-1, from_color, to_color);
     }
     if (y+1) < source.height {
-        flood_u32(source, destination, x, y+1, from_color, to_color);
+        flood(source, destination, x, y+1, from_color, to_color);
     }
 }
 
-pub(crate) fn flood_u16(source: &LabelFormat, destination: &mut Vec<u16>, 
-    x: usize, y: usize,
-    from_color: u16, to_color: u16) {
-    let width = source.width;
-    let destination_color = destination[x + y * width];
-    if destination_color != 0 {
-        return;
+pub(crate) mod color {
+    pub(crate) fn reset_color() {
+        print!("\x1b[0m");
     }
-    let source_color = source.buffer[x + y * width];
-    if source_color != from_color {
-        return;
-    }
-    destination[x + y * width] = to_color;
-    if x > 0 {
-        flood_u16(source, destination, x-1, y, from_color, to_color);
-    }
-    if (x+1) < width {
-        flood_u16(source, destination, x+1, y, from_color, to_color);
-    }
-    if y > 0 {
-        flood_u16(source, destination, x, y-1, from_color, to_color);
-    }
-    if (y+1) < source.height {
-        flood_u16(source, destination, x, y+1, from_color, to_color);
+
+    // Set background color from a number
+    pub(crate) fn set_color(color: usize) {
+        if color == 0 {
+            reset_color();
+            return;
+        }
+        let paint_color = color-1;
+        let paint_color = paint_color % 13;
+        if paint_color < 7 {
+            print!("\x1b[{}m", 40 + (paint_color+1));
+        } else {
+            print!("\x1b[{}m", 100 + (paint_color+1-7));
+        }
     }
 }