Fix xor_compress algorithm bug for single-byte input

Do use read_to_end for Rust in combination with map_err

3 files changed, 40 insertions, 42 deletions

diff --git a/tools/xor_compress.c b/tools/xor_compress.c
index 2e1f4fe..0e1c5e2 100644
--- a/tools/xor_compress.c
+++ b/tools/xor_compress.c
@@ -67,17 +67,10 @@ int write_compressed(const char *filename, unsigned char *data, size_t n, bool v
     for (size_t i = 0; i < n; runs++) {
         unsigned char byte = data[i++];
         unsigned char size = 0;
-        if (data[i] == v) {
-            // Alternating (>= 0x80)
-            // Run stops at 0x80 bytes or when the values stop alternating
-            for (; i < n && size < 0x80 && data[i] == ((size % 2) ? byte : v); size++, i++);
-            fputc(size + 0x7f, f);
-            fputc(v ^ byte, f);
-            if (size % 2 == 0) v = byte;
-        } else {
+        if (i == n || data[i] != v) {
             // Sequential (< 0x80)
             // Run stops at 0x80 bytes or when the value two ahead is equal to v
-            unsigned char buffer[256];
+            unsigned char buffer[0x80];
             buffer[size++] = v ^ byte;
             for (; i < n; i++) {
                 v = byte;
@@ -87,6 +80,13 @@ int write_compressed(const char *filename, unsigned char *data, size_t n, bool v
             }
             fputc(size - 1, f);
             fwrite(buffer, 1, size, f);
+        } else {
+            // Alternating (>= 0x80)
+            // Run stops at 0x80 bytes or when the values stop alternating
+            for (; i < n && size < 0x80 && data[i] == ((size % 2) ? byte : v); size++, i++);
+            fputc(size + 0x7f, f);
+            fputc(v ^ byte, f);
+            if (size % 2 == 0) v = byte;
         }
     }
 
diff --git a/tools/xor_compress.py b/tools/xor_compress.py
index 1c3a2f1..42d091f 100644
--- a/tools/xor_compress.py
+++ b/tools/xor_compress.py
@@ -32,19 +32,7 @@ while i < n:
     i += 1
     runs += 1
 
-    if data[i] == v:
-        # Alternating (>= 0x80)
-        # Run stops at 0x80 bytes or when the values stop alternating
-        size = 0
-        while i < n and size < 0x80 and data[i] == (byte if size % 2 else v):
-            size += 1
-            i += 1
-        output.append(size + 0x7f)
-        output.append(v ^ byte)
-        if not size % 2:
-            v = byte
-
-    else:
+    if i == n or data[i] != v:
         # Sequential (< 0x80)
         # Run stops at 0x80 bytes or when the value two ahead is equal to v
         buffer = [v ^ byte]
@@ -57,6 +45,17 @@ while i < n:
             i += 1
         output.append(len(buffer) - 1)
         output.extend(buffer)
+    else:
+        # Alternating (>= 0x80)
+        # Run stops at 0x80 bytes or when the values stop alternating
+        size = 0
+        while i < n and size < 0x80 and data[i] == (byte if size % 2 else v):
+            size += 1
+            i += 1
+        output.append(size + 0x7f)
+        output.append(v ^ byte)
+        if not size % 2:
+            v = byte
 
 with open(out_filename, 'wb') as f:
     f.write(output)
diff --git a/tools/xor_compress.rs b/tools/xor_compress.rs
index c6b78cf..3a41b16 100644
--- a/tools/xor_compress.rs
+++ b/tools/xor_compress.rs
@@ -1,17 +1,16 @@
 use std::env::args;
 use std::process::exit;
-use std::fs::{read, write};
-use std::io::Error;
+use std::fs::{File, write};
+use std::io::{Read, Error};
 
 const PROGRAM_NAME: &str = "xor-compress";
 
 fn read_files(filenames: &[String]) -> Result<Vec<u8>, (&String, Error)> {
     let mut data = Vec::new();
     for filename in filenames {
-        match read(filename) {
-            Ok(bytes) => data.extend(&bytes),
-            Err(err) => return Err((filename, err)),
-        }
+        (|| {
+            Ok(File::open(filename)?.read_to_end(&mut data)?)
+        })().map_err(|err| (filename, err))?;
     }
     Ok(data)
 }
@@ -29,20 +28,7 @@ fn write_compressed(filename: &str, data: &[u8]) -> Result<u32, Error> {
         i += 1;
         runs += 1;
 
-        if data[i] == v {
-            // Alternating (>= 0x80)
-            // Run stops at 0x80 bytes or when the values stop alternating
-            let mut size = 0;
-            while i < n && size < 0x80 && data[i] == (if size % 2 == 0 { v } else { byte }) {
-                size += 1;
-                i += 1;
-            }
-            output.push(size + 0x7f);
-            output.push(v ^ byte);
-            if size % 2 == 0 {
-                v = byte;
-            }
-        } else {
+        if i == n || data[i] != v {
             // Sequential (< 0x80)
             // Run stops at 0x80 bytes or when the value two ahead is equal to v
             let mut buffer = vec![v ^ byte];
@@ -57,6 +43,19 @@ fn write_compressed(filename: &str, data: &[u8]) -> Result<u32, Error> {
             }
             output.push((buffer.len() - 1) as u8);
             output.extend(buffer);
+        } else {
+            // Alternating (>= 0x80)
+            // Run stops at 0x80 bytes or when the values stop alternating
+            let mut size = 0;
+            while i < n && size < 0x80 && data[i] == (if size % 2 == 0 { v } else { byte }) {
+                size += 1;
+                i += 1;
+            }
+            output.push(size + 0x7f);
+            output.push(v ^ byte);
+            if size % 2 == 0 {
+                v = byte;
+            }
         }
     }