diff --git a/core/embed/rust/build.rs b/core/embed/rust/build.rs index 0fc4e0416..d360844fe 100644 --- a/core/embed/rust/build.rs +++ b/core/embed/rust/build.rs @@ -95,6 +95,7 @@ fn generate_micropython_bindings() { .allowlist_type("mp_obj_list_t") .allowlist_function("mp_obj_new_list") .allowlist_function("mp_obj_list_append") + .allowlist_function("mp_obj_list_set_len") .allowlist_var("mp_type_list") // map .allowlist_type("mp_map_elem_t") diff --git a/core/embed/rust/src/micropython/list.rs b/core/embed/rust/src/micropython/list.rs index 4811459be..a34d3ceb5 100644 --- a/core/embed/rust/src/micropython/list.rs +++ b/core/embed/rust/src/micropython/list.rs @@ -1,4 +1,4 @@ -use core::convert::TryFrom; +use core::{convert::TryFrom, ptr}; use crate::error::Error; @@ -17,6 +17,31 @@ impl List { }) } + pub fn with_capacity(capacity: usize) -> Result, Error> { + // EXCEPTION: Will raise if allocation fails. + catch_exception(|| unsafe { + let list = ffi::mp_obj_new_list(capacity, ptr::null_mut()); + // By default, the new list will have its len set to n. We want to preallocate + // to a specific size and then use append() to add items, so we reset len to 0. + ffi::mp_obj_list_set_len(list, 0); + Gc::from_raw(list.as_ptr().cast()) + }) + } + + pub fn from_iter(iter: impl Iterator) -> Result, Error> + where + T: TryInto, + Error: From, + { + let max_size = iter.size_hint().1.unwrap_or(0); + let mut gc_list = List::with_capacity(max_size)?; + let list = unsafe { Gc::as_mut(&mut gc_list) }; + for value in iter { + list.append(value.try_into()?)?; + } + Ok(gc_list) + } + pub fn append(&mut self, value: Obj) -> Result<(), Error> { unsafe { let ptr = self as *mut Self; @@ -52,3 +77,32 @@ impl TryFrom for Gc { } } } + +#[cfg(test)] +mod tests { + use crate::micropython::{ + iter::{Iter, IterBuf}, + testutil::mpy_init, + }; + + use super::*; + use heapless::Vec; + + #[test] + fn list_from_iter() { + unsafe { mpy_init() }; + + // create an upy list of 5 elements + let vec: Vec = (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 = iter + .map(TryInto::try_into) + .collect::, Error>>() + .unwrap(); + assert_eq!(vec, retrieved_vec); + } +}