fix(core): ensure drop is called on layout objects

[no changelog]

core/embed/rust/librust_qstr.h

@@ -22,6 +22,7 @@ static void _librust_qstrs(void) {
   MP_QSTR_NORMAL;
   MP_QSTR_TR;
   MP_QSTR_TranslationsHeader;
+  MP_QSTR___del__;
   MP_QSTR___dict__;
   MP_QSTR___name__;
   MP_QSTR_account;

core/embed/rust/src/ui/component/base.rs

@@ -208,29 +208,45 @@ where
 /// Same as `Child` but also handles screen clearing when layout is first
 /// painted.
 pub struct Root<T> {
-    inner: Child<T>,
+    inner: Option<Child<T>>,
     marked_for_clear: bool,
 }
 
 impl<T> Root<T> {
     pub fn new(component: T) -> Self {
         Self {
-            inner: Child::new(component),
+            inner: Some(Child::new(component)),
             marked_for_clear: true,
         }
     }
 
+    pub fn inner_mut(&mut self) -> &mut Child<T> {
+        if let Some(ref mut c) = self.inner {
+            c
+        } else {
+            fatal_error!("deallocated", "Root object is deallocated")
+        }
+    }
+
     pub fn inner(&self) -> &Child<T> {
-        &self.inner
+        if let Some(ref c) = self.inner {
+            c
+        } else {
+            fatal_error!("deallocated", "Root object is deallocated")
+        }
     }
 
     pub fn skip_paint(&mut self) {
-        self.inner.skip_paint()
+        self.inner_mut().skip_paint();
     }
 
     pub fn clear_screen(&mut self) {
         self.marked_for_clear = true;
     }
+
+    pub fn delete(&mut self) {
+        self.inner = None;
+    }
 }
 
 impl<T> Component for Root<T>
@@ -240,15 +256,15 @@ where
     type Msg = T::Msg;
 
     fn place(&mut self, bounds: Rect) -> Rect {
-        self.inner.place(bounds)
+        self.inner_mut().place(bounds)
     }
 
     fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<Self::Msg> {
-        let msg = self.inner.event(ctx, event);
+        let msg = self.inner_mut().event(ctx, event);
         if ctx.needs_repaint_root() {
             self.marked_for_clear = true;
             let mut dummy_ctx = EventCtx::new();
-            let paint_msg = self.inner.event(&mut dummy_ctx, Event::RequestPaint);
+            let paint_msg = self.inner_mut().event(&mut dummy_ctx, Event::RequestPaint);
             assert!(paint_msg.is_none());
             assert!(dummy_ctx.timers.is_empty());
         }
@@ -256,7 +272,7 @@ where
     }
 
     fn paint(&mut self) {
-        if self.marked_for_clear && self.inner.will_paint() {
+        if self.marked_for_clear && self.inner().will_paint() {
             self.marked_for_clear = false;
             display::clear()
         }
@@ -264,12 +280,12 @@ where
     }
 
     fn render<'s>(&self, target: &mut impl Renderer<'s>) {
-        self.inner.render(target);
+        self.inner().render(target);
     }
 
     #[cfg(feature = "ui_bounds")]
     fn bounds(&self, sink: &mut dyn FnMut(Rect)) {
-        self.inner.bounds(sink)
+        self.inner().bounds(sink)
     }
 }
 
@@ -279,7 +295,7 @@ where
     T: crate::trace::Trace,
 {
     fn trace(&self, t: &mut dyn crate::trace::Tracer) {
-        self.inner.trace(t)
+        self.inner().trace(t);
     }
 }
 

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

@@ -54,6 +54,7 @@ pub trait ObjComponent: MaybeTrace {
     fn obj_bounds(&self, _sink: &mut dyn FnMut(Rect)) {}
     fn obj_skip_paint(&mut self) {}
     fn obj_request_clear(&mut self) {}
+    fn obj_delete(&mut self) {}
 }
 
 impl<T> ObjComponent for Root<T>
@@ -105,6 +106,10 @@ where
     fn obj_request_clear(&mut self) {
         self.clear_screen()
     }
+
+    fn obj_delete(&mut self) {
+        self.delete()
+    }
 }
 
 /// `LayoutObj` is a GC-allocated object exported to MicroPython, with type
@@ -133,15 +138,25 @@ impl LayoutObj {
         let root =
             unsafe { Gc::from_raw(Gc::into_raw(Gc::new(wrapped_root)?) as *mut dyn ObjComponent) };
 
-        Gc::new(Self {
-            base: Self::obj_type().as_base(),
-            inner: RefCell::new(LayoutObjInner {
-                root,
-                event_ctx: EventCtx::new(),
-                timer_fn: Obj::const_none(),
-                page_count: 1,
-            }),
-        })
+        // SAFETY: This is a Python object and hase a base as first element
+        unsafe {
+            Gc::new_with_custom_finaliser(Self {
+                base: Self::obj_type().as_base(),
+                inner: RefCell::new(LayoutObjInner {
+                    root,
+                    event_ctx: EventCtx::new(),
+                    timer_fn: Obj::const_none(),
+                    page_count: 1,
+                }),
+            })
+        }
+    }
+
+    pub fn obj_delete(&self) {
+        let mut inner = self.inner.borrow_mut();
+
+        // SAFETY: `inner.root` is unique because of the `inner.borrow_mut()`.
+        unsafe { Gc::as_mut(&mut inner.root) }.obj_delete();
     }
 
     pub fn skip_first_paint(&self) {
@@ -291,6 +306,7 @@ impl LayoutObj {
                 Qstr::MP_QSTR_request_complete_repaint => obj_fn_1!(ui_layout_request_complete_repaint).as_obj(),
                 Qstr::MP_QSTR_trace => obj_fn_2!(ui_layout_trace).as_obj(),
                 Qstr::MP_QSTR_bounds => obj_fn_1!(ui_layout_bounds).as_obj(),
+                Qstr::MP_QSTR___del__ => obj_fn_1!(ui_layout_delete).as_obj(),
                 Qstr::MP_QSTR_page_count => obj_fn_1!(ui_layout_page_count).as_obj(),
                 Qstr::MP_QSTR_button_request => obj_fn_1!(ui_layout_button_request).as_obj(),
             }),
@@ -526,3 +542,12 @@ extern "C" fn ui_layout_bounds(this: Obj) -> Obj {
 extern "C" fn ui_layout_bounds(_this: Obj) -> Obj {
     Obj::const_none()
 }
+
+extern "C" fn ui_layout_delete(this: Obj) -> Obj {
+    let block = || {
+        let this: Gc<LayoutObj> = this.try_into()?;
+        this.obj_delete();
+        Ok(Obj::const_none())
+    };
+    unsafe { util::try_or_raise(block) }
+}

core/embed/rust/src/ui/model_mercury/component/homescreen.rs

@@ -14,7 +14,6 @@ use crate::{
         shape::{self, Renderer},
     },
 };
-use core::mem;
 
 use crate::ui::{
     component::Label,
@@ -344,10 +343,6 @@ impl Component for Lockscreen<'_> {
         }
 
         if let Event::Touch(TouchEvent::TouchEnd(_)) = event {
-            let bg_img = mem::replace(&mut self.bg_image, None);
-            if let Some(bg_img) = bg_img {
-                drop(bg_img);
-            }
             return Some(HomescreenMsg::Dismissed);
         }
 

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

@@ -1532,6 +1532,9 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     def page_count(self) -> int:
     ///         """Return the number of pages in the layout object."""
     ///
+    ///     def __del__(self) -> None:
+    ///         """Calls drop on contents of the root component."""
+    ///
     /// class UiResult:
     ///    """Result of a UI operation."""
     ///    pass

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

@@ -1679,6 +1679,9 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     def button_request(self) -> tuple[int, str] | None:
     ///         """Return (code, type) of button request made during the last event or timer pass."""
     ///
+    ///     def __del__(self) -> None:
+    ///         """Calls drop on contents of the root component."""
+    ///
     /// class UiResult:
     ///    """Result of a UI operation."""
     ///    pass

core/mocks/generated/trezorui2.pyi

@@ -49,6 +49,8 @@ class LayoutObj(Generic[T]):
             """Paint bounds of individual components on screen."""
     def page_count(self) -> int:
         """Return the number of pages in the layout object."""
+    def __del__(self) -> None:
+        """Calls drop on contents of the root component."""
 
 
 # rust/src/ui/model_mercury/layout.rs
@@ -1122,6 +1124,8 @@ class LayoutObj(Generic[T]):
         """Return the number of pages in the layout object."""
     def button_request(self) -> tuple[int, str] | None:
         """Return (code, type) of button request made during the last event or timer pass."""
+    def __del__(self) -> None:
+        """Calls drop on contents of the root component."""
 
 
 # rust/src/ui/model_tt/layout.rs

core/src/apps/homescreen/__init__.py

@@ -12,7 +12,11 @@ from apps.common.authorization import is_set_any_session
 
 
 async def busyscreen() -> None:
-    await Busyscreen(busy_expiry_ms())
+    obj = Busyscreen(busy_expiry_ms())
+    try:
+        await obj
+    finally:
+        obj.__del__()
 
 
 async def homescreen() -> None:
@@ -42,12 +46,17 @@ async def homescreen() -> None:
     elif storage.device.get_experimental_features():
         notification = TR.homescreen__title_experimental_mode
 
-    await Homescreen(
+    obj = Homescreen(
         label=label,
         notification=notification,
         notification_is_error=notification_is_error,
         hold_to_lock=config.has_pin(),
     )
+    try:
+        await obj
+    finally:
+        obj.__del__()
+
     lock_device()
 
 
@@ -58,10 +67,14 @@ async def _lockscreen(screensaver: bool = False) -> None:
     # Only show the lockscreen UI if the device can in fact be locked, or if it is
     # and OLED device (in which case the lockscreen is a screensaver).
     if can_lock_device() or screensaver:
-        await Lockscreen(
+        obj = Lockscreen(
             label=storage.device.get_label(),
             coinjoin_authorized=is_set_any_session(MessageType.AuthorizeCoinJoin),
         )
+        try:
+            await obj
+        finally:
+            obj.__del__()
     # Otherwise proceed directly to unlock() call. If the device is already unlocked,
     # it should be a no-op storage-wise, but it resets the internal configuration
     # to an unlocked state.

core/src/boot.py

@@ -62,6 +62,7 @@ async def bootscreen() -> None:
                     label=storage.device.get_label(), bootscreen=True
                 )
                 await lockscreen
+                lockscreen.__del__()
                 await verify_user_pin()
                 storage.init_unlocked()
                 allow_all_loader_messages()

core/src/trezor/ui/layouts/mercury/__init__.py

@@ -40,6 +40,9 @@ class RustLayout(ui.Layout):
         self._send_button_request()
         self.backlight_level = ui.BacklightLevels.NORMAL
 
+    def __del__(self):
+        self.layout.__del__()
+
     def set_timer(self, token: int, deadline: int) -> None:
         self.timer.schedule(deadline, token)
 

core/src/trezor/ui/layouts/tr/__init__.py

@@ -44,6 +44,9 @@ class RustLayout(LayoutParentType[T]):
         self.layout.attach_timer_fn(self.set_timer)
         self._send_button_request()
 
+    def __del__(self):
+        self.layout.__del__()
+
     def set_timer(self, token: int, deadline: int) -> None:
         self.timer.schedule(deadline, token)
 

core/src/trezor/ui/layouts/tt/__init__.py

@@ -46,6 +46,9 @@ class RustLayout(LayoutParentType[T]):
         self._send_button_request()
         self.backlight_level = ui.BacklightLevels.NORMAL
 
+    def __del__(self):
+        self.layout.__del__()
+
     def set_timer(self, token: int, deadline: int) -> None:
         self.timer.schedule(deadline, token)