diff --git a/src/heatmap.rs b/src/heatmap.rs index 2698321..10a1bec 100644 --- a/src/heatmap.rs +++ b/src/heatmap.rs @@ -12,6 +12,7 @@ pub struct Heatmap { data: Vec, cost: Vec, + raw_cost: Vec, literal_index: Vec, } @@ -20,6 +21,7 @@ impl Heatmap { Heatmap { data: Vec::new(), cost: Vec::new(), + raw_cost: Vec::new(), literal_index: Vec::new(), } } @@ -41,6 +43,8 @@ impl Heatmap { } pub(crate) fn finish(&mut self) { + self.raw_cost = self.cost.clone(); + let mut ref_count = vec![0usize; self.literal_index.len()]; for &index in &self.literal_index { ref_count[index] += 1; @@ -78,11 +82,18 @@ impl Heatmap { self.literal_index[index] == index } - /// Returns the cost of encoding the byte at `index` in (fractional) bits + /// Returns the cost of encoding the byte at `index` in (fractional) bits. + /// The cost of literal bytes is spread across the matches that reference it. + /// See `raw_cost` for the raw encoding cost of each byte. pub fn cost(&self, index: usize) -> f32 { self.cost[index] } + /// Returns the raw cost of encoding the byte at `index` in (fractional) bits + pub fn raw_cost(&self, index: usize) -> f32 { + self.raw_cost[index] + } + /// Returns the uncompressed data byte at `index` pub fn byte(&self, index: usize) -> u8 { self.data[index] @@ -91,6 +102,17 @@ impl Heatmap { #[cfg(feature = "crossterm")] /// Print the heatmap as a colored hexdump pub fn print_as_hex(&self) -> std::io::Result<()> { + self.print_as_hex_internal(false) + } + + #[cfg(feature = "crossterm")] + /// Print the heatmap as a colored hexdump, based on `raw_cost`. + pub fn print_as_hex_raw_cost(&self) -> std::io::Result<()> { + self.print_as_hex_internal(true) + } + + #[cfg(feature = "crossterm")] + fn print_as_hex_internal(&self, report_raw_cost: bool) -> std::io::Result<()> { use crossterm::{ style::{Attribute, Color, Print, SetAttribute, SetBackgroundColor}, QueueableCommand, @@ -102,8 +124,13 @@ impl Heatmap { heatmap: &Heatmap, index: usize, num_colors: u16, + report_raw_cost: bool, ) -> std::io::Result<()> { - let cost = heatmap.cost(index); + let cost = if report_raw_cost { + heatmap.raw_cost(index) + } else { + heatmap.cost(index) + }; if num_colors < 256 { let colors = [ Color::Red, @@ -149,7 +176,7 @@ impl Heatmap { stdout.queue(Print(&format!("{:04x} ", row_start)))?; for i in row_range.clone() { - set_color(&mut stdout, self, i, num_colors)?; + set_color(&mut stdout, self, i, num_colors, report_raw_cost)?; stdout.queue(Print(&format!("{:02x} ", self.data[i])))?; } @@ -160,7 +187,7 @@ impl Heatmap { .queue(Print(&gap))?; for i in row_range.clone() { - set_color(&mut stdout, self, i, num_colors)?; + set_color(&mut stdout, self, i, num_colors, report_raw_cost)?; let byte = self.data[i]; if byte >= 32 && byte < 127 { stdout.queue(Print(format!("{}", byte as char)))?; diff --git a/src/main.rs b/src/main.rs index 430f209..c876e43 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ fn main() -> Result<()> { let mut unpack = false; let mut calculate_margin = false; let mut create_heatmap = false; + let mut report_raw_cost = false; #[allow(unused_mut)] let mut do_hexdump = false; let mut level = 2; @@ -62,6 +63,7 @@ fn main() -> Result<()> { Short('u') | Long("unpack") | Short('d') | Long("decompress") => unpack = true, Long("margin") => calculate_margin = true, Long("heatmap") => create_heatmap = true, + Long("raw-cost") => report_raw_cost = true, #[cfg(feature = "crossterm")] Long("hexdump") => do_hexdump = true, Short('l') | Long("level") => level = parser.value()?.parse()?, @@ -141,14 +143,22 @@ fn main() -> Result<()> { } match do_hexdump { #[cfg(feature = "crossterm")] - true => heatmap.print_as_hex()?, + true => { + if report_raw_cost { + heatmap.print_as_hex_raw_cost()? + } else { + heatmap.print_as_hex()? + } + } _ => { let mut heatmap_bin = Vec::with_capacity(heatmap.len()); for i in 0..heatmap.len() { - let cost = (heatmap.cost(i).log2() * 8. + 64.) - .round() - .max(0.) - .min(127.) as u8; + let cost = if report_raw_cost { + heatmap.raw_cost(i) + } else { + heatmap.cost(i) + }; + let cost = (cost.log2() * 8. + 64.).round().max(0.).min(127.) as u8; heatmap_bin.push((cost << 1) | heatmap.is_literal(i) as u8); } outfile(OutFileType::Heatmap).write(&heatmap_bin)?;