From 80d1ce25c67bac6eaca4f3bf715ca22f73c0f552 Mon Sep 17 00:00:00 2001 From: cepetr Date: Mon, 22 Jan 2024 16:09:21 +0100 Subject: [PATCH] refactor(core/rust): improve tjpgdec interface [no changelog] --- core/embed/rust/src/ui/display/tjpgd.rs | 6 +- .../ui/model_tt/component/homescreen/mod.rs | 16 +-- .../model_tt/component/homescreen/render.rs | 12 +- rust/trezor-tjpgdec/src/lib.rs | 135 ++++++++++-------- 4 files changed, 97 insertions(+), 72 deletions(-) diff --git a/core/embed/rust/src/ui/display/tjpgd.rs b/core/embed/rust/src/ui/display/tjpgd.rs index 1fc6de3455..734cad50c9 100644 --- a/core/embed/rust/src/ui/display/tjpgd.rs +++ b/core/embed/rust/src/ui/display/tjpgd.rs @@ -20,7 +20,7 @@ pub fn jpeg(data: &[u8], pos: Point, scale: u8) { let mut inp = BufferInput(data); if let Ok(mut jd) = JDEC::new(&mut inp, pool) { let _ = jd.set_scale(scale); - let _ = jd.decomp(&mut out); + let _ = jd.decomp(&mut inp, &mut out); } } @@ -50,9 +50,9 @@ pub fn jpeg_test(data: &[u8]) -> bool { } let mut out = BlackHoleOutput; - let mut res = jd.decomp(&mut out); + let mut res = jd.decomp(&mut inp, &mut out); while res == Err(Error::Interrupted) { - res = jd.decomp(&mut out); + res = jd.decomp(&mut inp, &mut out); } res.is_ok() } else { diff --git a/core/embed/rust/src/ui/model_tt/component/homescreen/mod.rs b/core/embed/rust/src/ui/model_tt/component/homescreen/mod.rs index 1033139f5d..ec5769e776 100644 --- a/core/embed/rust/src/ui/model_tt/component/homescreen/mod.rs +++ b/core/embed/rust/src/ui/model_tt/component/homescreen/mod.rs @@ -215,9 +215,9 @@ impl Component for Homescreen { if let Ok(data) = res { if is_image_jpeg(data.as_ref()) { - let mut input = BufferInput(data.as_ref()); + let input = BufferInput(data.as_ref()); let mut pool = BufferJpegWork::get_cleared(); - let mut hs_img = HomescreenJpeg::new(&mut input, pool.buffer.as_mut_slice()); + let mut hs_img = HomescreenJpeg::new(input, pool.buffer.as_mut_slice()); homescreen( &mut hs_img, &[text], @@ -241,9 +241,9 @@ impl Component for Homescreen { } if show_default { - let mut input = BufferInput(IMAGE_HOMESCREEN); + let input = BufferInput(IMAGE_HOMESCREEN); let mut pool = BufferJpegWork::get_cleared(); - let mut hs_img = HomescreenJpeg::new(&mut input, pool.buffer.as_mut_slice()); + let mut hs_img = HomescreenJpeg::new(input, pool.buffer.as_mut_slice()); homescreen( &mut hs_img, &[text], @@ -348,9 +348,9 @@ impl Component for Lockscreen<'_> { if let Ok(data) = res { if is_image_jpeg(data.as_ref()) { - let mut input = BufferInput(data.as_ref()); + let input = BufferInput(data.as_ref()); let mut pool = BufferJpegWork::get_cleared(); - let mut hs_img = HomescreenJpeg::new(&mut input, pool.buffer.as_mut_slice()); + let mut hs_img = HomescreenJpeg::new(input, pool.buffer.as_mut_slice()); homescreen_blurred(&mut hs_img, texts); show_default = false; } else if is_image_toif(data.as_ref()) { @@ -364,9 +364,9 @@ impl Component for Lockscreen<'_> { } if show_default { - let mut input = BufferInput(IMAGE_HOMESCREEN); + let input = BufferInput(IMAGE_HOMESCREEN); let mut pool = BufferJpegWork::get_cleared(); - let mut hs_img = HomescreenJpeg::new(&mut input, pool.buffer.as_mut_slice()); + let mut hs_img = HomescreenJpeg::new(input, pool.buffer.as_mut_slice()); homescreen_blurred(&mut hs_img, texts); } } diff --git a/core/embed/rust/src/ui/model_tt/component/homescreen/render.rs b/core/embed/rust/src/ui/model_tt/component/homescreen/render.rs index 458f5aa355..b50b5f28b6 100644 --- a/core/embed/rust/src/ui/model_tt/component/homescreen/render.rs +++ b/core/embed/rust/src/ui/model_tt/component/homescreen/render.rs @@ -86,14 +86,16 @@ pub trait HomescreenDecompressor { pub struct HomescreenJpeg<'i> { pub output: BufferOutput, - pub jdec: Option>, + pub input: BufferInput<'i>, + pub jdec: Option>, } impl<'i> HomescreenJpeg<'i> { - pub fn new(input: &'i mut BufferInput<'i>, pool: &'i mut [u8]) -> Self { + pub fn new(mut input: BufferInput<'i>, pool: &'i mut [u8]) -> Self { Self { output: BufferOutput::new(WIDTH, 16), - jdec: JDEC::new(input, pool).ok(), + jdec: JDEC::new(&mut input, pool).ok(), + input, } } } @@ -107,7 +109,9 @@ impl<'i> HomescreenDecompressor for HomescreenJpeg<'i> { } fn decompress(&mut self) { - self.jdec.as_mut().map(|dec| dec.decomp(&mut self.output)); + self.jdec + .as_mut() + .map(|dec| dec.decomp(&mut self.input, &mut self.output)); } fn get_data(&mut self) -> &mut BufferJpeg { diff --git a/rust/trezor-tjpgdec/src/lib.rs b/rust/trezor-tjpgdec/src/lib.rs index 39a2c890e3..9333352e33 100644 --- a/rust/trezor-tjpgdec/src/lib.rs +++ b/rust/trezor-tjpgdec/src/lib.rs @@ -23,7 +23,8 @@ / Jul 01, 2021 R0.03 Added JD_FASTDECODE option. / Some performance improvement. / Jan 02, 2023 Rust version by Trezor Company, modified to meet our needs. - +/ May 14, 2024 Added better support for decompression resuming after an +/ output function interruption. Trezor modifications: - included overflow detection from https://github.com/cmumford/TJpgDec - removed JD_FASTDECODE=0 option @@ -31,6 +32,7 @@ Trezor modifications: - allowed interrupted functionality - tighter integration into Trezor codebase by using our data structures - removed generic input and output functions, replaced by our specific functionality + - added better support for decompression resuming after an output function interruption /----------------------------------------------------------------------------*/ #![no_std] @@ -85,7 +87,7 @@ pub enum Error { UnsupportedJpeg, } -pub struct JDEC<'i, 'p> { +pub struct JDEC<'p> { dctr: usize, dptr: usize, inbuf: &'p mut [u8], @@ -113,8 +115,9 @@ pub struct JDEC<'i, 'p> { hufflut_dc: [&'p mut [u8]; 2], workbuf: &'p mut [i32], mcubuf: &'p mut [i16], + mcu_x: u16, + mcu_y: u16, pool: &'p mut [u8], - input_func: &'i mut dyn JpegInput, } /// Zigzag-order to raster-order conversion table @@ -146,7 +149,7 @@ const IPSF: [u16; 64] = [ f!(0.27590), f!(0.38268), f!(0.36048), f!(0.32442), f!(0.27590), f!(0.21678), f!(0.14932), f!(0.07612), ]; -impl<'i, 'p> JDEC<'i, 'p> { +impl<'p> JDEC<'p> { /// Allocate a memory block from memory pool /// `self`: decompressor object reference /// `ndata` number of `T` items to allocate @@ -175,12 +178,12 @@ impl<'i, 'p> JDEC<'i, 'p> { } } - fn jpeg_in(&mut self, inbuf_offset: Option, n_data: usize) -> usize { + fn jpeg_in(&mut self, inbuf_offset: Option, n_data: usize, input_func: &mut dyn JpegInput) -> usize { if let Some(offset) = inbuf_offset { let inbuf = &mut self.inbuf[offset..offset + n_data]; - self.input_func.read(Some(inbuf), n_data) + input_func.read(Some(inbuf), n_data) } else { - self.input_func.read(None, n_data) + input_func.read(None, n_data) } } @@ -355,7 +358,7 @@ impl<'i, 'p> JDEC<'i, 'p> { /// `self`: decompressor object reference /// `id`: table ID (0:Y, 1:C) /// `cls`: table class (0:DC, 1:AC) - fn huffext(&mut self, id: usize, cls: usize) -> Result { + fn huffext(&mut self, id: usize, cls: usize, input_func: &mut dyn JpegInput) -> Result { let mut dc: usize = self.dctr; let mut dp: usize = self.dptr; let mut d: u32; @@ -373,7 +376,7 @@ impl<'i, 'p> JDEC<'i, 'p> { if dc == 0 { // Buffer empty, re-fill input buffer dp = 0; // Top of input buffer - dc = self.jpeg_in(Some(0), JD_SZBUF); + dc = self.jpeg_in(Some(0), JD_SZBUF, input_func); if dc == 0 { // Err: read error or wrong stream termination return Err(Error::Input); @@ -479,7 +482,7 @@ impl<'i, 'p> JDEC<'i, 'p> { /// Extract N bits from input stream /// `self`: decompressor object reference /// `nbit`: number of bits to extract (1 to 16) - fn bitext(&mut self, nbit: u32) -> Result { + fn bitext(&mut self, nbit: u32, input_func: &mut dyn JpegInput) -> Result { let mut dc: usize = self.dctr; let mut dp: usize = self.dptr; let mut d: u32; @@ -494,7 +497,7 @@ impl<'i, 'p> JDEC<'i, 'p> { if dc == 0 { // Buffer empty, re-fill input buffer dp = 0; // Top of input buffer - dc = self.jpeg_in(Some(0), JD_SZBUF); + dc = self.jpeg_in(Some(0), JD_SZBUF, input_func); if dc == 0 { // Err: read error or wrong stream termination return Err(Error::Input); @@ -531,7 +534,7 @@ impl<'i, 'p> JDEC<'i, 'p> { /// Process restart interval /// `self`: decompressor object reference /// `rstn`: expected restart sequence number - fn restart(&mut self, rstn: u16) -> Result<(), Error> { + fn restart(&mut self, rstn: u16, input_func: &mut dyn JpegInput) -> Result<(), Error> { let mut dp = self.dptr; let mut dc: usize = self.dctr; let mut marker: u16; @@ -546,7 +549,7 @@ impl<'i, 'p> JDEC<'i, 'p> { if dc == 0 { // No input data is available, re-fill input buffer dp = 0; - dc = self.jpeg_in(Some(0), JD_SZBUF); + dc = self.jpeg_in(Some(0), JD_SZBUF, input_func); if dc == 0 { return Err(Error::Input); } @@ -696,7 +699,7 @@ impl<'i, 'p> JDEC<'i, 'p> { /// Load all blocks in an MCU into working buffer /// `self`: decompressor object reference - fn mcu_load(&mut self) -> Result<(), Error> { + fn mcu_load(&mut self, input_func: &mut dyn JpegInput) -> Result<(), Error> { let mut d: i32; let mut e: i32; let mut blk: u32; @@ -720,12 +723,12 @@ impl<'i, 'p> JDEC<'i, 'p> { id = if cmp != 0 { 1 } else { 0 }; // Huffman table ID of this component // Extract a DC element from input stream - d = self.huffext(id as usize, 0)?; // Extract a huffman coded data (bit length) + d = self.huffext(id as usize, 0, input_func)?; // Extract a huffman coded data (bit length) bc = d as u32; d = self.dcv[cmp as usize] as i32; // DC value of previous block if bc != 0 { // If there is any difference from previous block - e = self.bitext(bc)?; // Extract data bits + e = self.bitext(bc, input_func)?; // Extract data bits bc = 1 << (bc - 1); // MSB position if e as u32 & bc == 0 { e -= ((bc << 1) - 1) as i32; // Restore negative value @@ -751,7 +754,7 @@ impl<'i, 'p> JDEC<'i, 'p> { z = 1; // Top of the AC elements (in zigzag-order) loop { // Extract a huffman coded value (zero runs and bit length) - d = self.huffext(id as usize, 1)?; + d = self.huffext(id as usize, 1, input_func)?; if d == 0 { // EOB? break; @@ -765,7 +768,7 @@ impl<'i, 'p> JDEC<'i, 'p> { bc &= 0xf; if bc != 0 { // Bit length? - d = self.bitext(bc)?; // Extract data bits + d = self.bitext(bc, input_func)?; // Extract data bits bc = 1 << (bc - 1); // MSB position if d as u32 & bc == 0 { // Restore negative value if needed @@ -1112,7 +1115,7 @@ impl<'i, 'p> JDEC<'i, 'p> { } /// Analyze the JPEG image and Initialize decompressor object - pub fn new(input_func: &'i mut dyn JpegInput, pool: &'p mut [u8]) -> Result { + pub fn new(input_func: &mut dyn JpegInput, pool: &'p mut [u8]) -> Result { let mut jd = JDEC { dctr: 0, dptr: 0, @@ -1142,7 +1145,8 @@ impl<'i, 'p> JDEC<'i, 'p> { ncomp: 0, nrst: 0, mcubuf: &mut [], - input_func, + mcu_x: 0, + mcu_y: 0, }; let mut marker: u16; @@ -1156,7 +1160,7 @@ impl<'i, 'p> JDEC<'i, 'p> { marker = 0; ofs = marker as u32; loop { - if jd.jpeg_in(Some(0), 1) != 1 { + if jd.jpeg_in(Some(0), 1, input_func) != 1 { // Err: SOI was not detected return Err(Error::Input); } @@ -1169,7 +1173,7 @@ impl<'i, 'p> JDEC<'i, 'p> { loop { // Parse JPEG segments // Get a JPEG marker - if jd.jpeg_in(Some(0), 4) != 4 { + if jd.jpeg_in(Some(0), 4, input_func) != 4 { return Err(Error::Input); } // Marker @@ -1189,7 +1193,7 @@ impl<'i, 'p> JDEC<'i, 'p> { return Err(Error::MemoryInput); } // Load segment data - if jd.jpeg_in(Some(0), len) != len { + if jd.jpeg_in(Some(0), len, input_func) != len { return Err(Error::Input); } // Image width in unit of pixel @@ -1235,7 +1239,7 @@ impl<'i, 'p> JDEC<'i, 'p> { return Err(Error::MemoryInput); } // Load segment data - if jd.jpeg_in(Some(0), len) != len { + if jd.jpeg_in(Some(0), len, input_func) != len { return Err(Error::Input); } // Get restart interval (MCUs) @@ -1247,7 +1251,7 @@ impl<'i, 'p> JDEC<'i, 'p> { return Err(Error::MemoryInput); } // Load segment data - if jd.jpeg_in(Some(0), len) != len { + if jd.jpeg_in(Some(0), len, input_func) != len { return Err(Error::Input); } // Create huffman tables @@ -1259,7 +1263,7 @@ impl<'i, 'p> JDEC<'i, 'p> { return Err(Error::MemoryInput); } // Load segment data - if jd.jpeg_in(Some(0), len) != len { + if jd.jpeg_in(Some(0), len, input_func) != len { return Err(Error::Input); } // Create de-quantizer tables @@ -1271,7 +1275,7 @@ impl<'i, 'p> JDEC<'i, 'p> { return Err(Error::MemoryInput); } // Load segment data - if jd.jpeg_in(Some(0), len) != len { + if jd.jpeg_in(Some(0), len, input_func) != len { return Err(Error::Input); } if jd.width == 0 || jd.height == 0 { @@ -1323,7 +1327,7 @@ impl<'i, 'p> JDEC<'i, 'p> { // Align stream read offset to JD_SZBUF ofs %= JD_SZBUF as u32; if ofs != 0 { - jd.dctr = jd.jpeg_in(Some(ofs as usize), (JD_SZBUF as u32 - ofs) as usize); + jd.dctr = jd.jpeg_in(Some(ofs as usize), (JD_SZBUF as u32 - ofs) as usize, input_func); } jd.dptr = (ofs - (if JD_FASTDECODE != 0 { 0 } else { 1 })) as usize; return Ok(jd); // Initialization succeeded. Ready to @@ -1338,7 +1342,7 @@ impl<'i, 'p> JDEC<'i, 'p> { _ => { // Unknown segment (comment, exif or etc..) // Skip segment data (null pointer specifies to remove data from the stream) - if jd.jpeg_in(None, len) != len { + if jd.jpeg_in(None, len, input_func) != len { return Err(Error::Input); } } @@ -1346,38 +1350,55 @@ impl<'i, 'p> JDEC<'i, 'p> { } } - /// Start to decompress the JPEG picture - /// `scale`: output de-scaling factor (0 to 3) - pub fn decomp(&mut self, output_func: &mut dyn JpegOutput) -> Result<(), Error> { - let mx = (self.msx as i32 * 8) as u32; // Size of the MCU (pixel) - let my = (self.msy as i32 * 8) as u32; // Size of the MCU (pixel) - let mut y = 0; - while y < self.height as u32 { - // Vertical loop of MCUs - let mut x = 0; - while x < self.width as u32 { - // Horizontal loop of MCUs - if self.nrst != 0 && { - // Process restart interval if enabled - let val = self.rst; - self.rst += 1; - val == self.nrst - } { - let val = self.rsc; - self.rsc += 1; - self.restart(val)?; - self.rst = 1; - } - // Load an MCU (decompress huffman coded stream, dequantize and apply IDCT) - self.mcu_load()?; - // Output the MCU (YCbCr to RGB, scaling and output) - self.mcu_output(x, y, output_func)?; - x += mx; + /// Start/resume JPEG decompression + /// + /// The function decompress the JPEG image in stream and calls + /// the output function for each decoded MCU. + /// + /// If the output function returns `false`, the decompression is interrupted. + /// It's possible later to call `decomp()` again to resume the decompression. + pub fn decomp(&mut self, input_func: &mut dyn JpegInput, output_func: &mut dyn JpegOutput) -> Result<(), Error> { + let mx = self.msx as u16 * 8; // Size of the MCU (pixel) + let my = self.msy as u16 * 8; // Size of the MCU (pixel) + while self.mcu_y < self.height { + if self.nrst != 0 && { + // Process restart interval if enabled + let val = self.rst; + self.rst += 1; + val == self.nrst + } { + let val = self.rsc; + self.rsc += 1; + self.restart(val, input_func)?; + self.rst = 1; } - y += my; + + // Load an MCU (decompress huffman coded stream, dequantize and apply IDCT) + self.mcu_load(input_func)?; + + let x = self.mcu_x as u32; + let y = self.mcu_y as u32; + + self.mcu_x += mx; + if self.mcu_x >= self.width { + self.mcu_x = 0; + self.mcu_y += my; + } + + // Output the MCU (YCbCr to RGB, scaling and output) + self.mcu_output(x, y, output_func)?; } Ok(()) } + + /// Returns pixel coordinates (top-left) of the next decoded MCU + /// + /// The function is useful when the decompression is interrupted + /// and later resumed by `decomp()`. + pub fn next_mcu(&self) -> (u16, u16) { + (self.mcu_x, self.mcu_y) + } + } pub trait JpegInput {