arno/trezor-firmware
1
0
Fork 0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-22 22:38:08 +00:00
Code Issues Releases Wiki Activity

refactor(core/rust): improve ergonomy of IterBuf

Browse Source
This commit is contained in:
matejcik 2023-06-20 12:00:23 +02:00 committed by matejcik
parent b91e225076
commit 5a83a7171d
7 changed files with 63 additions and 88 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
core/embed/rust/src/micropython/iter.rs
View File

@ -6,6 +6,8 @@ use super::{ffi, runtime::catch_exception};
pub struct IterBuf {
iter_buf: ffi::mp_obj_iter_buf_t,
error_checking: bool,
caught_exception: Obj,
}
impl IterBuf {
@ -17,20 +19,41 @@ impl IterBuf {
},
buf: [Obj::const_null(), Obj::const_null(), Obj::const_null()],
},
error_checking: false,
caught_exception: Obj::const_null(),
}
}
pub fn new_fallible() -> Self {
let mut new = Self::new();
new.error_checking = true;
new
}
pub fn try_iterate(&mut self, o: Obj) -> Result<Iter, Error> {
Iter::try_from_obj_with_buf(o, self)
}
pub fn error(&self) -> Option<Obj> {
if !self.error_checking || self.caught_exception.is_null() {
None
} else {
Some(self.caught_exception)
}
}
}
#[allow(dead_code)] // iter_buf is not really used but needs to be there for SAFETY reasons
pub struct Iter<'a> {
iter: Obj,
// SAFETY:
// In the typical case, `iter` is literally a pointer to `iter_buf`. We need to hold
// an exclusive reference to `iter_buf` to avoid problems.
iter_buf: &'a mut IterBuf,
finished: bool,
caught_exception: Obj,
}
impl<'a> Iter<'a> {
pub fn try_from_obj_with_buf(o: Obj, iter_buf: &'a mut IterBuf) -> Result<Self, Error> {
fn try_from_obj_with_buf(o: Obj, iter_buf: &'a mut IterBuf) -> Result<Self, Error> {
// SAFETY:
// - In the common case, `ffi::mp_getiter` does not heap-allocate, but instead
// uses memory from the passed `iter_buf`. We maintain this invariant by
@ -43,7 +66,6 @@ impl<'a> Iter<'a> {
iter,
iter_buf,
finished: false,
caught_exception: Obj::const_null(),
})
}
}
@ -61,8 +83,8 @@ impl<'a> Iterator for Iter<'a> {
// EXCEPTION: Can raise from the underlying iterator.
let item = catch_exception(|| unsafe { ffi::mp_iternext(self.iter) });
match item {
Err(Error::CaughtException(exc)) => {
self.caught_exception = exc;
Err(Error::CaughtException(exc)) if self.iter_buf.error_checking => {
self.iter_buf.caught_exception = exc;
None
}
Err(_) => panic!("unexpected error"),

23
core/embed/rust/src/micropython/list.rs
View File

@ -144,10 +144,7 @@ impl TryFrom<Obj> for Gc<List> {
#[cfg(test)]
mod tests {
use crate::micropython::{
iter::{Iter, IterBuf},
testutil::mpy_init,
};
use crate::micropython::{iter::IterBuf, testutil::mpy_init};
use super::*;
use heapless::Vec;
@ -160,10 +157,10 @@ mod tests {
let vec: Vec<u8, 10> = (0..5).collect();
let list: Obj = List::from_iter(vec.iter().copied()).unwrap().into();
let mut buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(list, &mut buf).unwrap();
// collect the elements into a Vec of maximum length 10, through an iterator
let retrieved_vec: Vec<u8, 10> = iter
let retrieved_vec: Vec<u8, 10> = IterBuf::new()
.try_iterate(list)
.unwrap()
.map(TryInto::try_into)
.collect::<Result<Vec<u8, 10>, Error>>()
.unwrap();
@ -199,9 +196,9 @@ mod tests {
);
}
let mut buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(gc_list.into(), &mut buf).unwrap();
let retrieved_vec: Vec<u16, 17> = iter
let retrieved_vec: Vec<u16, 17> = IterBuf::new()
.try_iterate(gc_list.into())
.unwrap()
.map(TryInto::try_into)
.collect::<Result<Vec<u16, 17>, Error>>()
.unwrap();
@ -240,9 +237,9 @@ mod tests {
slice[i] = ((i + 10) as u16).into();
}
let mut buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(list.into(), &mut buf).unwrap();
let retrieved_vec: Vec<u16, 5> = iter
let retrieved_vec: Vec<u16, 5> = IterBuf::new()
.try_iterate(list.into())
.unwrap()
.map(TryInto::try_into)
.collect::<Result<Vec<u16, 5>, Error>>()
.unwrap();

8
core/embed/rust/src/micropython/obj.rs
View File

@ -447,11 +447,11 @@ impl Obj {
// SAFETY:
// Safe for pointers, for as long as MicroPython behaves sanely.
// We assume that:
// * The pointer is a valid MicroPython object, which has `ObjBase`
// as its first element.
// * The pointer is a valid MicroPython object, which has `ObjBase` as its first
// element.
// * The type pointer points to a valid type object.
// * The pointee has a 'static lifetime, i.e., either is
// ROM-based, or GC allocated.
// * The pointee has a 'static lifetime, i.e., either is ROM-based, or GC
// allocated.
let base = self.as_ptr() as *const ObjBase;
unsafe { (*base).type_.as_ref() }
} else {

20
core/embed/rust/src/protobuf/encode.rs
View File

@ -2,15 +2,7 @@ use core::convert::{TryFrom, TryInto};
use crate::{
error::Error,
micropython::{
buffer,
gc::Gc,
iter::{Iter, IterBuf},
list::List,
obj::Obj,
qstr::Qstr,
util,
},
micropython::{buffer, gc::Gc, iter::IterBuf, list::List, obj::Obj, qstr::Qstr, util},
};
use super::{
@ -80,9 +72,7 @@ impl Encoder {
};
if field.is_repeated() {
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(field_value, &mut iter_buf)?;
for iter_value in iter {
for iter_value in IterBuf::new().try_iterate(field_value)? {
stream.write_uvarint(field_key)?;
self.encode_field(stream, field, iter_value)?;
}
@ -126,8 +116,7 @@ impl Encoder {
// Serialize the total length of the buffer.
let mut len = 0;
let iter = Iter::try_from_obj_with_buf(value, &mut iter_buf)?;
for value in iter {
for value in iter_buf.try_iterate(value)? {
// SAFETY: buffer is dropped immediately.
let buffer = unsafe { buffer::get_buffer(value) }?;
len += buffer.len();
@ -135,8 +124,7 @@ impl Encoder {
stream.write_uvarint(len as u64)?;
// Serialize the buffers one-by-one.
let iter = Iter::try_from_obj_with_buf(value, &mut iter_buf)?;
for value in iter {
for value in iter_buf.try_iterate(value)? {
// SAFETY: buffer is dropped immediately.
let buffer = unsafe { buffer::get_buffer(value) }?;
stream.write(buffer)?;

5
core/embed/rust/src/ui/layout/util.rs
View File

@ -3,7 +3,7 @@ use crate::{
micropython::{
buffer::{hexlify_bytes, StrBuffer},
gc::Gc,
iter::{Iter, IterBuf},
iter::IterBuf,
list::List,
obj::Obj,
util::try_or_raise,
@ -42,8 +42,7 @@ where
Error: From<E>,
{
let mut vec = Vec::<T, N>::new();
let mut iter_buf = IterBuf::new();
for item in Iter::try_from_obj_with_buf(iterable, &mut iter_buf)? {
for item in IterBuf::new().try_iterate(iterable)? {
vec.push(item.try_into()?)
.map_err(|_| value_error!("Invalid iterable length"))?;
}

27
core/embed/rust/src/ui/model_tr/layout.rs
View File

@ -6,15 +6,8 @@ use crate::{
error::Error,
maybe_trace::MaybeTrace,
micropython::{
buffer::StrBuffer,
gc::Gc,
iter::{Iter, IterBuf},
list::List,
map::Map,
module::Module,
obj::Obj,
qstr::Qstr,
util,
buffer::StrBuffer, gc::Gc, iter::IterBuf, list::List, map::Map, module::Module, obj::Obj,
qstr::Qstr, util,
},
strutil::StringType,
ui::{
@ -367,9 +360,7 @@ extern "C" fn new_confirm_properties(n_args: usize, args: *const Obj, kwargs: *m
let mut paragraphs = ParagraphVecLong::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for para in iter {
for para in IterBuf::new().try_iterate(items)? {
let [key, value, is_data]: [Obj; 3] = iter_into_array(para)?;
let key = key.try_into_option::<StrBuffer>()?;
let value = value.try_into_option::<StrBuffer>()?;
@ -430,12 +421,10 @@ extern "C" fn new_show_address_details(n_args: usize, args: *const Obj, kwargs:
let path: Option<StrBuffer> = kwargs.get(Qstr::MP_QSTR_path)?.try_into_option()?;
let xpubs: Obj = kwargs.get(Qstr::MP_QSTR_xpubs)?;
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(xpubs, &mut iter_buf)?;
let mut ad = AddressDetails::new(address, case_sensitive, account, path)?;
for i in iter {
for i in IterBuf::new().try_iterate(xpubs)? {
let [xtitle, text]: [StrBuffer; 2] = iter_into_array(i)?;
ad.add_xpub(xtitle, text)?;
}
@ -897,9 +886,7 @@ extern "C" fn new_confirm_with_info(n_args: usize, args: *const Obj, kwargs: *mu
let mut paragraphs = ParagraphVecShort::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for para in iter {
for para in IterBuf::new().try_iterate(items)? {
let [font, text]: [Obj; 2] = iter_into_array(para)?;
let style: &TextStyle = theme::textstyle_number_bold_or_mono(font.try_into()?);
let text: StrBuffer = text.try_into()?;
@ -1062,10 +1049,8 @@ extern "C" fn new_show_checklist(n_args: usize, args: *const Obj, kwargs: *mut M
let active: usize = kwargs.get(Qstr::MP_QSTR_active)?.try_into()?;
let items: Obj = kwargs.get(Qstr::MP_QSTR_items)?;
let mut iter_buf = IterBuf::new();
let mut paragraphs = ParagraphVecLong::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for (i, item) in iter.enumerate() {
for (i, item) in IterBuf::new().try_iterate(items)?.enumerate() {
let style = match i.cmp(&active) {
Ordering::Less => &theme::TEXT_NORMAL,
Ordering::Equal => &theme::TEXT_BOLD,

34
core/embed/rust/src/ui/model_tt/layout.rs
View File

@ -6,7 +6,7 @@ use crate::{
micropython::{
buffer::{get_buffer, StrBuffer},
gc::Gc,
iter::{Iter, IterBuf},
iter::IterBuf,
list::List,
map::Map,
module::Module,
@ -466,10 +466,8 @@ extern "C" fn new_confirm_emphasized(n_args: usize, args: *const Obj, kwargs: *m
.try_into_option()?;
let items: Obj = kwargs.get(Qstr::MP_QSTR_items)?;
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
let mut ops = OpTextLayout::new(theme::TEXT_NORMAL);
for item in iter {
for item in IterBuf::new().try_iterate(items)? {
if item.is_str() {
ops = ops.text_normal(item.try_into()?)
} else {
@ -748,12 +746,10 @@ extern "C" fn new_show_address_details(n_args: usize, args: *const Obj, kwargs:
let path: Option<StrBuffer> = kwargs.get(Qstr::MP_QSTR_path)?.try_into_option()?;
let xpubs: Obj = kwargs.get(Qstr::MP_QSTR_xpubs)?;
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(xpubs, &mut iter_buf)?;
let mut ad = AddressDetails::new(address, case_sensitive, account, path)?;
for i in iter {
for i in IterBuf::new().try_iterate(xpubs)? {
let [xtitle, text]: [StrBuffer; 2] = iter_into_array(i)?;
ad.add_xpub(xtitle, text)?;
}
@ -833,9 +829,7 @@ extern "C" fn new_confirm_total(n_args: usize, args: *const Obj, kwargs: *mut Ma
let mut paragraphs = ParagraphVecShort::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for pair in iter {
for pair in IterBuf::new().try_iterate(items)? {
let [label, value]: [StrBuffer; 2] = iter_into_array(pair)?;
paragraphs.add(Paragraph::new(&theme::TEXT_NORMAL, label));
paragraphs.add(Paragraph::new(&theme::TEXT_MONO, value));
@ -1159,9 +1153,7 @@ extern "C" fn new_confirm_with_info(n_args: usize, args: *const Obj, kwargs: *mu
let mut paragraphs = ParagraphVecShort::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for para in iter {
for para in IterBuf::new().try_iterate(items)? {
let [font, text]: [Obj; 2] = iter_into_array(para)?;
let style: &TextStyle = theme::textstyle_number(font.try_into()?);
let text: StrBuffer = text.try_into()?;
@ -1191,9 +1183,7 @@ extern "C" fn new_confirm_more(n_args: usize, args: *const Obj, kwargs: *mut Map
let mut paragraphs = ParagraphVecLong::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for para in iter {
for para in IterBuf::new().try_iterate(items)? {
let [font, text]: [Obj; 2] = iter_into_array(para)?;
let style: &TextStyle = theme::textstyle_number(font.try_into()?);
let text: StrBuffer = text.try_into()?;
@ -1308,9 +1298,7 @@ extern "C" fn new_show_share_words(n_args: usize, args: *const Obj, kwargs: *mut
let pages: Obj = kwargs.get(Qstr::MP_QSTR_pages)?;
let mut paragraphs = ParagraphVecLong::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(pages, &mut iter_buf)?;
for page in iter {
for page in IterBuf::new().try_iterate(pages)? {
let text: StrBuffer = page.try_into()?;
paragraphs.add(Paragraph::new(&theme::TEXT_MONO, text).break_after());
}
@ -1360,10 +1348,8 @@ extern "C" fn new_show_checklist(n_args: usize, args: *const Obj, kwargs: *mut M
let active: usize = kwargs.get(Qstr::MP_QSTR_active)?.try_into()?;
let items: Obj = kwargs.get(Qstr::MP_QSTR_items)?;
let mut iter_buf = IterBuf::new();
let mut paragraphs = ParagraphVecLong::new();
let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
for (i, item) in iter.enumerate() {
for (i, item) in IterBuf::new().try_iterate(items)?.enumerate() {
let style = match i.cmp(&active) {
Ordering::Less => &theme::TEXT_CHECKLIST_DONE,
Ordering::Equal => &theme::TEXT_CHECKLIST_SELECTED,
@ -1489,9 +1475,7 @@ extern "C" fn new_show_remaining_shares(n_args: usize, args: *const Obj, kwargs:
let pages_iterable: Obj = kwargs.get(Qstr::MP_QSTR_pages)?;
let mut paragraphs = ParagraphVecLong::new();
let mut iter_buf = IterBuf::new();
let iter = Iter::try_from_obj_with_buf(pages_iterable, &mut iter_buf)?;
for page in iter {
for page in IterBuf::new().try_iterate(pages_iterable)? {
let [title, description]: [StrBuffer; 2] = iter_into_array(page)?;
paragraphs
.add(Paragraph::new(&theme::TEXT_DEMIBOLD, title))