From ef8b1f7ba5cf9f2607c43467855096845b7be1cc Mon Sep 17 00:00:00 2001 From: cepetr Date: Mon, 22 Jan 2024 16:09:21 +0100 Subject: [PATCH] refactor(core): improve tjpg interface --- 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 | 105 ++++++++++++------ 4 files changed, 93 insertions(+), 46 deletions(-) diff --git a/core/embed/rust/src/ui/display/tjpgd.rs b/core/embed/rust/src/ui/display/tjpgd.rs index 1fc6de345..734cad50c 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 ed188815e..7937472eb 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 3323a1ca8..32798115d 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 39a2c890e..29671db58 100644 --- a/rust/trezor-tjpgdec/src/lib.rs +++ b/rust/trezor-tjpgdec/src/lib.rs @@ -85,7 +85,7 @@ pub enum Error { UnsupportedJpeg, } -pub struct JDEC<'i, 'p> { +pub struct JDEC<'p> { dctr: usize, dptr: usize, inbuf: &'p mut [u8], @@ -113,8 +113,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 +147,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 +176,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 +356,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 +374,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 +480,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 +495,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 +532,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 +547,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 +697,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 +721,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 +752,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 +766,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 +1113,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 +1143,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 +1158,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 +1171,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 +1191,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 +1237,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 +1249,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 +1261,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 +1273,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 +1325,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 +1340,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); } } @@ -1348,7 +1350,7 @@ 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> { + pub fn decomp(&mut self, input_func: &mut dyn JpegInput, 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; @@ -1365,11 +1367,11 @@ impl<'i, 'p> JDEC<'i, 'p> { } { let val = self.rsc; self.rsc += 1; - self.restart(val)?; + self.restart(val, input_func)?; self.rst = 1; } // Load an MCU (decompress huffman coded stream, dequantize and apply IDCT) - self.mcu_load()?; + self.mcu_load(input_func)?; // Output the MCU (YCbCr to RGB, scaling and output) self.mcu_output(x, y, output_func)?; x += mx; @@ -1378,6 +1380,47 @@ impl<'i, 'p> JDEC<'i, 'p> { } Ok(()) } + + /// Start to decompress the JPEG picture + /// `scale`: output de-scaling factor (0 to 3) + pub fn decomp2(&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; + } + + // 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(()) + } + + pub fn next_mcu(&self) -> (u16, u16) { + (self.mcu_x, self.mcu_y) + } + } pub trait JpegInput {