refactor(core/rust): improve ergonomy of IterBuf

core/embed/rust/src/micropython/iter.rs

@@ -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)) => {
+            Err(Error::CaughtException(exc)) if self.iter_buf.error_checking => {
-                self.caught_exception = exc;
+                self.iter_buf.caught_exception = exc;
                 None
             }
             Err(_) => panic!("unexpected error"),

core/embed/rust/src/micropython/list.rs

@@ -144,10 +144,7 @@ impl TryFrom<Obj> for Gc<List> {
 
 #[cfg(test)]
 mod tests {
-    use crate::micropython::{
+    use crate::micropython::{iter::IterBuf, testutil::mpy_init};
-        iter::{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 retrieved_vec: Vec<u16, 17> = IterBuf::new()
-        let iter = Iter::try_from_obj_with_buf(gc_list.into(), &mut buf).unwrap();
+            .try_iterate(gc_list.into())
-        let retrieved_vec: Vec<u16, 17> = iter
+            .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 retrieved_vec: Vec<u16, 5> = IterBuf::new()
-        let iter = Iter::try_from_obj_with_buf(list.into(), &mut buf).unwrap();
+            .try_iterate(list.into())
-        let retrieved_vec: Vec<u16, 5> = iter
+            .unwrap()
             .map(TryInto::try_into)
             .collect::<Result<Vec<u16, 5>, Error>>()
             .unwrap();

core/embed/rust/src/micropython/obj.rs

@@ -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`
+            // * The pointer is a valid MicroPython object, which has `ObjBase` as its first
-            //   as its first element.
+            //   element.
             // * The type pointer points to a valid type object.
-            // * The pointee has a 'static lifetime, i.e., either is
+            // * The pointee has a 'static lifetime, i.e., either is ROM-based, or GC
-            //   ROM-based, or GC allocated.
+            //   allocated.
             let base = self.as_ptr() as *const ObjBase;
             unsafe { (*base).type_.as_ref() }
         } else {

core/embed/rust/src/protobuf/encode.rs

@@ -2,15 +2,7 @@ use core::convert::{TryFrom, TryInto};
 
 use crate::{
     error::Error,
-    micropython::{
+    micropython::{buffer, gc::Gc, iter::IterBuf, list::List, obj::Obj, qstr::Qstr, util},
-        buffer,
-        gc::Gc,
-        iter::{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();
+                for iter_value in IterBuf::new().try_iterate(field_value)? {
-                let iter = Iter::try_from_obj_with_buf(field_value, &mut iter_buf)?;
-                for iter_value in iter {
                     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_buf.try_iterate(value)? {
-                    for value in iter {
                         // 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_buf.try_iterate(value)? {
-                    for value in iter {
                         // SAFETY: buffer is dropped immediately.
                         let buffer = unsafe { buffer::get_buffer(value) }?;
                         stream.write(buffer)?;

core/embed/rust/src/ui/layout/util.rs

@@ -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 IterBuf::new().try_iterate(iterable)? {
-    for item in Iter::try_from_obj_with_buf(iterable, &mut iter_buf)? {
         vec.push(item.try_into()?)
             .map_err(|_| value_error!("Invalid iterable length"))?;
     }

core/embed/rust/src/ui/model_tr/layout.rs

@@ -6,15 +6,8 @@ use crate::{
     error::Error,
     maybe_trace::MaybeTrace,
     micropython::{
-        buffer::StrBuffer,
+        buffer::StrBuffer, gc::Gc, iter::IterBuf, list::List, map::Map, module::Module, obj::Obj,
-        gc::Gc,
+        qstr::Qstr, util,
-        iter::{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();
+        for para in IterBuf::new().try_iterate(items)? {
-        let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
-        for para in iter {
             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();
+        for para in IterBuf::new().try_iterate(items)? {
-        let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
-        for para in iter {
             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 IterBuf::new().try_iterate(items)?.enumerate() {
-        for (i, item) in iter.enumerate() {
             let style = match i.cmp(&active) {
                 Ordering::Less => &theme::TEXT_NORMAL,
                 Ordering::Equal => &theme::TEXT_BOLD,

core/embed/rust/src/ui/model_tt/layout.rs

@@ -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();
+        for pair in IterBuf::new().try_iterate(items)? {
-        let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
-        for pair in iter {
             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();
+        for para in IterBuf::new().try_iterate(items)? {
-        let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
-        for para in iter {
             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();
+        for para in IterBuf::new().try_iterate(items)? {
-        let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?;
-        for para in iter {
             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();
+        for page in IterBuf::new().try_iterate(pages)? {
-        let iter = Iter::try_from_obj_with_buf(pages, &mut iter_buf)?;
-        for page in iter {
             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 IterBuf::new().try_iterate(items)?.enumerate() {
-        for (i, item) in iter.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();
+        for page in IterBuf::new().try_iterate(pages_iterable)? {
-        let iter = Iter::try_from_obj_with_buf(pages_iterable, &mut iter_buf)?;
-        for page in iter {
             let [title, description]: [StrBuffer; 2] = iter_into_array(page)?;
             paragraphs
                 .add(Paragraph::new(&theme::TEXT_DEMIBOLD, title))