1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-18 04:18:10 +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);
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 {

View File

@ -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);
}
}

View File

@ -86,14 +86,16 @@ pub trait HomescreenDecompressor {
pub struct HomescreenJpeg<'i> {
pub output: BufferOutput,
pub jdec: Option<JDEC<'i, 'i>>,
pub input: BufferInput<'i>,
pub jdec: Option<JDEC<'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 {
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 {

View File

@ -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<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 {
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<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 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<i32, Error> {
fn bitext(&mut self, nbit: u32, input_func: &mut dyn JpegInput) -> Result<i32, Error> {
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<Self, Error> {
pub fn new(input_func: &mut dyn JpegInput, pool: &'p mut [u8]) -> Result<Self, Error> {
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 {