1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-18 12:28:09 +00:00

refactor(core/rust): improve tjpgdec interface

[no changelog]
This commit is contained in:
cepetr 2024-01-22 16:09:21 +01:00 committed by cepetr
parent 12ff2e4570
commit 80d1ce25c6
4 changed files with 97 additions and 72 deletions

View File

@ -20,7 +20,7 @@ pub fn jpeg(data: &[u8], pos: Point, scale: u8) {
let mut inp = BufferInput(data); let mut inp = BufferInput(data);
if let Ok(mut jd) = JDEC::new(&mut inp, pool) { if let Ok(mut jd) = JDEC::new(&mut inp, pool) {
let _ = jd.set_scale(scale); 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 out = BlackHoleOutput;
let mut res = jd.decomp(&mut out); let mut res = jd.decomp(&mut inp, &mut out);
while res == Err(Error::Interrupted) { while res == Err(Error::Interrupted) {
res = jd.decomp(&mut out); res = jd.decomp(&mut inp, &mut out);
} }
res.is_ok() res.is_ok()
} else { } else {

View File

@ -215,9 +215,9 @@ impl Component for Homescreen {
if let Ok(data) = res { if let Ok(data) = res {
if is_image_jpeg(data.as_ref()) { 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 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( homescreen(
&mut hs_img, &mut hs_img,
&[text], &[text],
@ -241,9 +241,9 @@ impl Component for Homescreen {
} }
if show_default { if show_default {
let mut input = BufferInput(IMAGE_HOMESCREEN); let input = BufferInput(IMAGE_HOMESCREEN);
let mut pool = BufferJpegWork::get_cleared(); 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( homescreen(
&mut hs_img, &mut hs_img,
&[text], &[text],
@ -348,9 +348,9 @@ impl Component for Lockscreen<'_> {
if let Ok(data) = res { if let Ok(data) = res {
if is_image_jpeg(data.as_ref()) { 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 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); homescreen_blurred(&mut hs_img, texts);
show_default = false; show_default = false;
} else if is_image_toif(data.as_ref()) { } else if is_image_toif(data.as_ref()) {
@ -364,9 +364,9 @@ impl Component for Lockscreen<'_> {
} }
if show_default { if show_default {
let mut input = BufferInput(IMAGE_HOMESCREEN); let input = BufferInput(IMAGE_HOMESCREEN);
let mut pool = BufferJpegWork::get_cleared(); 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); homescreen_blurred(&mut hs_img, texts);
} }
} }

View File

@ -86,14 +86,16 @@ pub trait HomescreenDecompressor {
pub struct HomescreenJpeg<'i> { pub struct HomescreenJpeg<'i> {
pub output: BufferOutput, pub output: BufferOutput,
pub jdec: Option<JDEC<'i, 'i>>, pub input: BufferInput<'i>,
pub jdec: Option<JDEC<'i>>,
} }
impl<'i> HomescreenJpeg<'i> { 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 { Self {
output: BufferOutput::new(WIDTH, 16), 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) { 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 { fn get_data(&mut self) -> &mut BufferJpeg {

View File

@ -23,7 +23,8 @@
/ Jul 01, 2021 R0.03 Added JD_FASTDECODE option. / Jul 01, 2021 R0.03 Added JD_FASTDECODE option.
/ Some performance improvement. / Some performance improvement.
/ Jan 02, 2023 Rust version by Trezor Company, modified to meet our needs. / 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: Trezor modifications:
- included overflow detection from https://github.com/cmumford/TJpgDec - included overflow detection from https://github.com/cmumford/TJpgDec
- removed JD_FASTDECODE=0 option - removed JD_FASTDECODE=0 option
@ -31,6 +32,7 @@ Trezor modifications:
- allowed interrupted functionality - allowed interrupted functionality
- tighter integration into Trezor codebase by using our data structures - tighter integration into Trezor codebase by using our data structures
- removed generic input and output functions, replaced by our specific functionality - removed generic input and output functions, replaced by our specific functionality
- added better support for decompression resuming after an output function interruption
/----------------------------------------------------------------------------*/ /----------------------------------------------------------------------------*/
#![no_std] #![no_std]
@ -85,7 +87,7 @@ pub enum Error {
UnsupportedJpeg, UnsupportedJpeg,
} }
pub struct JDEC<'i, 'p> { pub struct JDEC<'p> {
dctr: usize, dctr: usize,
dptr: usize, dptr: usize,
inbuf: &'p mut [u8], inbuf: &'p mut [u8],
@ -113,8 +115,9 @@ pub struct JDEC<'i, 'p> {
hufflut_dc: [&'p mut [u8]; 2], hufflut_dc: [&'p mut [u8]; 2],
workbuf: &'p mut [i32], workbuf: &'p mut [i32],
mcubuf: &'p mut [i16], mcubuf: &'p mut [i16],
mcu_x: u16,
mcu_y: u16,
pool: &'p mut [u8], pool: &'p mut [u8],
input_func: &'i mut dyn JpegInput,
} }
/// Zigzag-order to raster-order conversion table /// 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), 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 /// Allocate a memory block from memory pool
/// `self`: decompressor object reference /// `self`: decompressor object reference
/// `ndata` number of `T` items to allocate /// `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<usize>, n_data: usize) -> usize { fn jpeg_in(&mut self, inbuf_offset: Option<usize>, n_data: usize, input_func: &mut dyn JpegInput) -> usize {
if let Some(offset) = inbuf_offset { if let Some(offset) = inbuf_offset {
let inbuf = &mut self.inbuf[offset..offset + n_data]; 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 { } 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 /// `self`: decompressor object reference
/// `id`: table ID (0:Y, 1:C) /// `id`: table ID (0:Y, 1:C)
/// `cls`: table class (0:DC, 1:AC) /// `cls`: table class (0:DC, 1:AC)
fn huffext(&mut self, id: usize, cls: usize) -> Result<i32, Error> { fn huffext(&mut self, id: usize, cls: usize, input_func: &mut dyn JpegInput) -> Result<i32, Error> {
let mut dc: usize = self.dctr; let mut dc: usize = self.dctr;
let mut dp: usize = self.dptr; let mut dp: usize = self.dptr;
let mut d: u32; let mut d: u32;
@ -373,7 +376,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
if dc == 0 { if dc == 0 {
// Buffer empty, re-fill input buffer // Buffer empty, re-fill input buffer
dp = 0; // Top of 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 { if dc == 0 {
// Err: read error or wrong stream termination // Err: read error or wrong stream termination
return Err(Error::Input); return Err(Error::Input);
@ -479,7 +482,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
/// Extract N bits from input stream /// Extract N bits from input stream
/// `self`: decompressor object reference /// `self`: decompressor object reference
/// `nbit`: number of bits to extract (1 to 16) /// `nbit`: number of bits to extract (1 to 16)
fn bitext(&mut self, nbit: u32) -> Result<i32, Error> { fn bitext(&mut self, nbit: u32, input_func: &mut dyn JpegInput) -> Result<i32, Error> {
let mut dc: usize = self.dctr; let mut dc: usize = self.dctr;
let mut dp: usize = self.dptr; let mut dp: usize = self.dptr;
let mut d: u32; let mut d: u32;
@ -494,7 +497,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
if dc == 0 { if dc == 0 {
// Buffer empty, re-fill input buffer // Buffer empty, re-fill input buffer
dp = 0; // Top of 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 { if dc == 0 {
// Err: read error or wrong stream termination // Err: read error or wrong stream termination
return Err(Error::Input); return Err(Error::Input);
@ -531,7 +534,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
/// Process restart interval /// Process restart interval
/// `self`: decompressor object reference /// `self`: decompressor object reference
/// `rstn`: expected restart sequence number /// `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 dp = self.dptr;
let mut dc: usize = self.dctr; let mut dc: usize = self.dctr;
let mut marker: u16; let mut marker: u16;
@ -546,7 +549,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
if dc == 0 { if dc == 0 {
// No input data is available, re-fill input buffer // No input data is available, re-fill input buffer
dp = 0; dp = 0;
dc = self.jpeg_in(Some(0), JD_SZBUF); dc = self.jpeg_in(Some(0), JD_SZBUF, input_func);
if dc == 0 { if dc == 0 {
return Err(Error::Input); return Err(Error::Input);
} }
@ -696,7 +699,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
/// Load all blocks in an MCU into working buffer /// Load all blocks in an MCU into working buffer
/// `self`: decompressor object reference /// `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 d: i32;
let mut e: i32; let mut e: i32;
let mut blk: u32; 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 id = if cmp != 0 { 1 } else { 0 }; // Huffman table ID of this component
// Extract a DC element from input stream // 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; bc = d as u32;
d = self.dcv[cmp as usize] as i32; // DC value of previous block d = self.dcv[cmp as usize] as i32; // DC value of previous block
if bc != 0 { if bc != 0 {
// If there is any difference from previous block // 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 bc = 1 << (bc - 1); // MSB position
if e as u32 & bc == 0 { if e as u32 & bc == 0 {
e -= ((bc << 1) - 1) as i32; // Restore negative value 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) z = 1; // Top of the AC elements (in zigzag-order)
loop { loop {
// Extract a huffman coded value (zero runs and bit length) // 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 { if d == 0 {
// EOB? // EOB?
break; break;
@ -765,7 +768,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
bc &= 0xf; bc &= 0xf;
if bc != 0 { if bc != 0 {
// Bit length? // Bit length?
d = self.bitext(bc)?; // Extract data bits d = self.bitext(bc, input_func)?; // Extract data bits
bc = 1 << (bc - 1); // MSB position bc = 1 << (bc - 1); // MSB position
if d as u32 & bc == 0 { if d as u32 & bc == 0 {
// Restore negative value if needed // Restore negative value if needed
@ -1112,7 +1115,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
} }
/// Analyze the JPEG image and Initialize decompressor object /// Analyze the JPEG image and Initialize decompressor object
pub fn new(input_func: &'i mut dyn JpegInput, pool: &'p mut [u8]) -> Result<Self, Error> { pub fn new(input_func: &mut dyn JpegInput, pool: &'p mut [u8]) -> Result<Self, Error> {
let mut jd = JDEC { let mut jd = JDEC {
dctr: 0, dctr: 0,
dptr: 0, dptr: 0,
@ -1142,7 +1145,8 @@ impl<'i, 'p> JDEC<'i, 'p> {
ncomp: 0, ncomp: 0,
nrst: 0, nrst: 0,
mcubuf: &mut [], mcubuf: &mut [],
input_func, mcu_x: 0,
mcu_y: 0,
}; };
let mut marker: u16; let mut marker: u16;
@ -1156,7 +1160,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
marker = 0; marker = 0;
ofs = marker as u32; ofs = marker as u32;
loop { loop {
if jd.jpeg_in(Some(0), 1) != 1 { if jd.jpeg_in(Some(0), 1, input_func) != 1 {
// Err: SOI was not detected // Err: SOI was not detected
return Err(Error::Input); return Err(Error::Input);
} }
@ -1169,7 +1173,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
loop { loop {
// Parse JPEG segments // Parse JPEG segments
// Get a JPEG marker // 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); return Err(Error::Input);
} }
// Marker // Marker
@ -1189,7 +1193,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
return Err(Error::MemoryInput); return Err(Error::MemoryInput);
} }
// Load segment data // 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); return Err(Error::Input);
} }
// Image width in unit of pixel // Image width in unit of pixel
@ -1235,7 +1239,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
return Err(Error::MemoryInput); return Err(Error::MemoryInput);
} }
// Load segment data // 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); return Err(Error::Input);
} }
// Get restart interval (MCUs) // Get restart interval (MCUs)
@ -1247,7 +1251,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
return Err(Error::MemoryInput); return Err(Error::MemoryInput);
} }
// Load segment data // 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); return Err(Error::Input);
} }
// Create huffman tables // Create huffman tables
@ -1259,7 +1263,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
return Err(Error::MemoryInput); return Err(Error::MemoryInput);
} }
// Load segment data // 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); return Err(Error::Input);
} }
// Create de-quantizer tables // Create de-quantizer tables
@ -1271,7 +1275,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
return Err(Error::MemoryInput); return Err(Error::MemoryInput);
} }
// Load segment data // 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); return Err(Error::Input);
} }
if jd.width == 0 || jd.height == 0 { if jd.width == 0 || jd.height == 0 {
@ -1323,7 +1327,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
// Align stream read offset to JD_SZBUF // Align stream read offset to JD_SZBUF
ofs %= JD_SZBUF as u32; ofs %= JD_SZBUF as u32;
if ofs != 0 { 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; jd.dptr = (ofs - (if JD_FASTDECODE != 0 { 0 } else { 1 })) as usize;
return Ok(jd); // Initialization succeeded. Ready to return Ok(jd); // Initialization succeeded. Ready to
@ -1338,7 +1342,7 @@ impl<'i, 'p> JDEC<'i, 'p> {
_ => { _ => {
// Unknown segment (comment, exif or etc..) // Unknown segment (comment, exif or etc..)
// Skip segment data (null pointer specifies to remove data from the stream) // 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); return Err(Error::Input);
} }
} }
@ -1346,17 +1350,17 @@ impl<'i, 'p> JDEC<'i, 'p> {
} }
} }
/// Start to decompress the JPEG picture /// Start/resume JPEG decompression
/// `scale`: output de-scaling factor (0 to 3) ///
pub fn decomp(&mut self, output_func: &mut dyn JpegOutput) -> Result<(), Error> { /// The function decompress the JPEG image in stream and calls
let mx = (self.msx as i32 * 8) as u32; // Size of the MCU (pixel) /// the output function for each decoded MCU.
let my = (self.msy as i32 * 8) as u32; // Size of the MCU (pixel) ///
let mut y = 0; /// If the output function returns `false`, the decompression is interrupted.
while y < self.height as u32 { /// It's possible later to call `decomp()` again to resume the decompression.
// Vertical loop of MCUs pub fn decomp(&mut self, input_func: &mut dyn JpegInput, output_func: &mut dyn JpegOutput) -> Result<(), Error> {
let mut x = 0; let mx = self.msx as u16 * 8; // Size of the MCU (pixel)
while x < self.width as u32 { let my = self.msy as u16 * 8; // Size of the MCU (pixel)
// Horizontal loop of MCUs while self.mcu_y < self.height {
if self.nrst != 0 && { if self.nrst != 0 && {
// Process restart interval if enabled // Process restart interval if enabled
let val = self.rst; let val = self.rst;
@ -1365,19 +1369,36 @@ impl<'i, 'p> JDEC<'i, 'p> {
} { } {
let val = self.rsc; let val = self.rsc;
self.rsc += 1; self.rsc += 1;
self.restart(val)?; self.restart(val, input_func)?;
self.rst = 1; self.rst = 1;
} }
// Load an MCU (decompress huffman coded stream, dequantize and apply IDCT) // Load an MCU (decompress huffman coded stream, dequantize and apply IDCT)
self.mcu_load()?; 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) // Output the MCU (YCbCr to RGB, scaling and output)
self.mcu_output(x, y, output_func)?; self.mcu_output(x, y, output_func)?;
x += mx;
}
y += my;
} }
Ok(()) 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 { pub trait JpegInput {