refactor(core/rust/ui): avoid homescreen flicker during workflow restarts

[no changelog]

core/embed/rust/librust_qstr.h

@@ -97,4 +97,5 @@ static void _librust_qstrs(void) {
   MP_QSTR_notification;
   MP_QSTR_notification_level;
   MP_QSTR_bootscreen;
+  MP_QSTR_skip_first_paint;
 }

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

@@ -105,6 +105,16 @@ impl<T> Child<T> {
         }
         result
     }
+
+    /// Do not draw on screen until an event requests paint. This is used by
+    /// homescreens to avoid flickerg when workflow restart happens.
+    pub fn skip_paint(&mut self) {
+        self.marked_for_paint = false;
+    }
+
+    pub fn will_paint(&self) -> bool {
+        self.marked_for_paint
+    }
 }
 
 impl<T> Component for Child<T>

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

@@ -70,8 +70,9 @@ use maybe_trace::MaybeTrace;
 pub trait ObjComponent: MaybeTrace {
     fn obj_place(&mut self, bounds: Rect) -> Rect;
     fn obj_event(&mut self, ctx: &mut EventCtx, event: Event) -> Result<Obj, Error>;
-    fn obj_paint(&mut self);
+    fn obj_paint(&mut self) -> bool;
     fn obj_bounds(&self, sink: &mut dyn FnMut(Rect));
+    fn obj_skip_paint(&mut self) {}
 }
 
 impl<T> ObjComponent for Child<T>
@@ -90,13 +91,19 @@ where
         }
     }
 
-    fn obj_paint(&mut self) {
+    fn obj_paint(&mut self) -> bool {
+        let will_paint = self.will_paint();
         self.paint();
+        will_paint
     }
 
     fn obj_bounds(&self, sink: &mut dyn FnMut(Rect)) {
         self.bounds(sink)
     }
+
+    fn obj_skip_paint(&mut self) {
+        self.skip_paint()
+    }
 }
 
 /// `LayoutObj` is a GC-allocated object exported to MicroPython, with type
@@ -136,6 +143,13 @@ impl LayoutObj {
         })
     }
 
+    pub fn skip_first_paint(&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_skip_paint();
+    }
+
     /// Timer callback is expected to be a callable object of the following
     /// form: `def timer(token: int, deadline_in_ms: int)`.
     fn obj_set_timer_fn(&self, timer_fn: Obj) {
@@ -183,8 +197,9 @@ impl LayoutObj {
         Ok(msg)
     }
 
-    /// Run a paint pass over the component tree.
-    fn obj_paint_if_requested(&self) {
+    /// Run a paint pass over the component tree. Returns true if any component
+    /// actually requested painting since last invocation of the function.
+    fn obj_paint_if_requested(&self) -> bool {
         let mut inner = self.inner.borrow_mut();
 
         // Place the root component on the screen in case it was previously requested.
@@ -194,7 +209,7 @@ impl LayoutObj {
         }
 
         // SAFETY: `inner.root` is unique because of the `inner.borrow_mut()`.
-        unsafe { Gc::as_mut(&mut inner.root) }.obj_paint();
+        unsafe { Gc::as_mut(&mut inner.root) }.obj_paint()
     }
 
     /// Run a tracing pass over the component tree. Passed `callback` is called
@@ -431,8 +446,8 @@ extern "C" fn ui_layout_timer(this: Obj, token: Obj) -> Obj {
 extern "C" fn ui_layout_paint(this: Obj) -> Obj {
     let block = || {
         let this: Gc<LayoutObj> = this.try_into()?;
-        this.obj_paint_if_requested();
-        Ok(Obj::const_true())
+        let painted = this.obj_paint_if_requested().into();
+        Ok(painted)
     };
     unsafe { util::try_or_raise(block) }
 }

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

@@ -1011,9 +1011,13 @@ extern "C" fn new_show_homescreen(n_args: usize, args: *const Obj, kwargs: *mut
             kwargs.get(Qstr::MP_QSTR_notification)?.try_into_option()?;
         let notification_level: u32 = kwargs.get_or(Qstr::MP_QSTR_notification_level, 0)?;
         let hold: bool = kwargs.get(Qstr::MP_QSTR_hold)?.try_into()?;
+        let skip_first_paint: bool = kwargs.get(Qstr::MP_QSTR_skip_first_paint)?.try_into()?;
 
         let notification = notification.map(|w| (w, notification_level));
         let obj = LayoutObj::new(Homescreen::new(label, notification, hold))?;
+        if skip_first_paint {
+            obj.skip_first_paint();
+        }
         Ok(obj.into())
     };
     unsafe { util::try_with_args_and_kwargs(n_args, args, kwargs, block) }
@@ -1023,8 +1027,12 @@ extern "C" fn new_show_lockscreen(n_args: usize, args: *const Obj, kwargs: *mut
     let block = move |_args: &[Obj], kwargs: &Map| {
         let label: StrBuffer = kwargs.get(Qstr::MP_QSTR_label)?.try_into()?;
         let bootscreen: bool = kwargs.get(Qstr::MP_QSTR_bootscreen)?.try_into()?;
+        let skip_first_paint: bool = kwargs.get(Qstr::MP_QSTR_skip_first_paint)?.try_into()?;
 
         let obj = LayoutObj::new(Lockscreen::new(label, bootscreen))?;
+        if skip_first_paint {
+            obj.skip_first_paint();
+        }
         Ok(obj.into())
     };
     unsafe { util::try_with_args_and_kwargs(n_args, args, kwargs, block) }
@@ -1035,6 +1043,7 @@ extern "C" fn new_show_busyscreen(n_args: usize, args: *const Obj, kwargs: *mut
         let title: StrBuffer = kwargs.get(Qstr::MP_QSTR_title)?.try_into()?;
         let description: StrBuffer = kwargs.get(Qstr::MP_QSTR_description)?.try_into()?;
         let time_ms: u32 = kwargs.get(Qstr::MP_QSTR_time_ms)?.try_into()?;
+        let skip_first_paint: bool = kwargs.get(Qstr::MP_QSTR_skip_first_paint)?.try_into()?;
 
         let obj = LayoutObj::new(Frame::new(
             title,
@@ -1047,6 +1056,9 @@ extern "C" fn new_show_busyscreen(n_args: usize, args: *const Obj, kwargs: *mut
                 }),
             ),
         ))?;
+        if skip_first_paint {
+            obj.skip_first_paint();
+        }
         Ok(obj.into())
     };
     unsafe { util::try_with_args_and_kwargs(n_args, args, kwargs, block) }
@@ -1337,6 +1349,7 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     hold: bool,
     ///     notification: str | None,
     ///     notification_level: int = 0,
+    ///     skip_first_paint: bool,
     /// ) -> trezorui2.CANCELLED:
     ///     """Idle homescreen."""
     Qstr::MP_QSTR_show_homescreen => obj_fn_kw!(0, new_show_homescreen).as_obj(),
@@ -1345,6 +1358,7 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     *,
     ///     label: str,
     ///     bootscreen: bool,
+    ///     skip_first_paint: bool,
     /// ) -> trezorui2.CANCELLED:
     ///     """Homescreen for locked device."""
     Qstr::MP_QSTR_show_lockscreen => obj_fn_kw!(0, new_show_lockscreen).as_obj(),
@@ -1354,6 +1368,7 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     title: str,
     ///     description: str,
     ///     time_ms: int,
+    ///     skip_first_paint: bool,
     /// ) -> trezorui2.CANCELLED:
     ///     """Homescreen used for indicating coinjoin in progress."""
     Qstr::MP_QSTR_show_busyscreen => obj_fn_kw!(0, new_show_busyscreen).as_obj(),

core/mocks/generated/trezorui2.pyi

@@ -357,6 +357,7 @@ def show_homescreen(
     hold: bool,
     notification: str | None,
     notification_level: int = 0,
+    skip_first_paint: bool,
 ) -> trezorui2.CANCELLED:
     """Idle homescreen."""
 
@@ -366,6 +367,7 @@ def show_lockscreen(
     *,
     label: str,
     bootscreen: bool,
+    skip_first_paint: bool,
 ) -> trezorui2.CANCELLED:
     """Homescreen for locked device."""
 
@@ -376,5 +378,6 @@ def show_busyscreen(
     title: str,
     description: str,
     time_ms: int,
+    skip_first_paint: bool,
 ) -> trezorui2.CANCELLED:
     """Homescreen used for indicating coinjoin in progress."""

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

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING
 from ubinascii import hexlify
 
+import storage.cache
 from trezor import io, log, loop, ui, wire, workflow
 from trezor.enums import ButtonRequestType
 
@@ -32,6 +33,11 @@ class _RustLayout(ui.Layout):
         msg = self.layout.request_complete_repaint()
         assert msg is None
 
+    def _paint(self) -> None:
+        painted = self.layout.paint()
+        if storage.cache.homescreen_shown is not None and painted:
+            storage.cache.homescreen_shown = None
+
     if __debug__:
 
         def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
@@ -81,7 +87,7 @@ class _RustLayout(ui.Layout):
                     (io.TOUCH_END, orig_x + 2 * off_x, orig_y + 2 * off_y),
                 ):
                     msg = self.layout.touch_event(event, x, y)
-                    self.layout.paint()
+                    self._paint()
                     if msg is not None:
                         raise ui.Result(msg)
 
@@ -111,7 +117,7 @@ class _RustLayout(ui.Layout):
         def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
             return self.handle_timers(), self.handle_input_and_rendering()
 
-    def _before_render(self) -> None:
+    def _first_paint(self) -> None:
         # Clear the screen of any leftovers.
         ui.backlight_fade(ui.style.BACKLIGHT_DIM)
         ui.display.clear()
@@ -127,11 +133,11 @@ class _RustLayout(ui.Layout):
 
         # Turn the brightness on again.
         ui.backlight_fade(self.BACKLIGHT_LEVEL)
+        self._paint()
 
     def handle_input_and_rendering(self) -> loop.Task:  # type: ignore [awaitable-is-generator]
         touch = loop.wait(io.TOUCH)
-        self._before_render()
-        self.layout.paint()
+        self._first_paint()
         # self.layout.bounds()
         while True:
             # Using `yield` instead of `await` to avoid allocations.
@@ -140,7 +146,7 @@ class _RustLayout(ui.Layout):
             msg = None
             if event in (io.TOUCH_START, io.TOUCH_MOVE, io.TOUCH_END):
                 msg = self.layout.touch_event(event, x, y)
-            self.layout.paint()
+            self._paint()
             # self.layout.bounds()
             if msg is not None:
                 raise ui.Result(msg)
@@ -150,7 +156,7 @@ class _RustLayout(ui.Layout):
             # Using `yield` instead of `await` to avoid allocations.
             token = yield self.timer
             msg = self.layout.timer(token)
-            self.layout.paint()
+            self._paint()
             if msg is not None:
                 raise ui.Result(msg)
 

core/src/trezor/ui/layouts/tt_v2/homescreen.py

@@ -2,7 +2,7 @@ from typing import Any, Tuple
 
 import storage.cache
 import storage.device
-from trezor import io, loop, ui
+from trezor import io, loop, ui, utils
 
 import trezorui2
 from apps.base import set_homescreen
@@ -26,6 +26,21 @@ class HomescreenBase(_RustLayout):
             storage.cache.homescreen_shown = None
             raise
 
+    def _first_paint(self) -> None:
+        if storage.cache.homescreen_shown is not self.RENDER_INDICATOR:
+            super()._first_paint()
+            storage.cache.homescreen_shown = self.RENDER_INDICATOR
+
+        # - RENDER_INDICATOR is set -> USB warning is not displayed
+        # - RENDER_INDICATOR is not set -> initially homescreen does not display warning
+        # - usb_checker_task only handles state changes
+        # Here we need to handle the case when homescreen is started with USB disconnected.
+        if not utils.usb_data_connected():
+            msg = self.layout.usb_event(False)
+            self._paint()
+            if msg is not None:
+                raise ui.Result(msg)
+
 
 class Homescreen(HomescreenBase):
     RENDER_INDICATOR = storage.cache.HOMESCREEN_ON
@@ -45,12 +60,14 @@ class Homescreen(HomescreenBase):
             elif notification_is_error:
                 level = 0
 
+        skip = storage.cache.homescreen_shown is self.RENDER_INDICATOR
         super().__init__(
             layout=trezorui2.show_homescreen(
                 label=label or "My Trezor",
                 notification=notification,
                 notification_level=level,
                 hold=hold_to_lock,
+                skip_first_paint=skip,
             ),
         )
 
@@ -81,10 +98,14 @@ class Lockscreen(HomescreenBase):
         if bootscreen:
             self.BACKLIGHT_LEVEL = ui.BACKLIGHT_NORMAL
 
+        skip = (
+            not bootscreen and storage.cache.homescreen_shown is self.RENDER_INDICATOR
+        )
         super().__init__(
             layout=trezorui2.show_lockscreen(
                 label=label or "My Trezor",
                 bootscreen=bootscreen,
+                skip_first_paint=skip,
             ),
         )
 
@@ -99,11 +120,13 @@ class Busyscreen(HomescreenBase):
     RENDER_INDICATOR = storage.cache.BUSYSCREEN_ON
 
     def __init__(self, delay_ms: int) -> None:
+        skip = storage.cache.homescreen_shown is self.RENDER_INDICATOR
         super().__init__(
             layout=trezorui2.show_busyscreen(
                 title="PLEASE WAIT",
                 description="CoinJoin in progress.\n\nDo not disconnect your\nTrezor.",
                 time_ms=delay_ms,
+                skip_first_paint=skip,
             )
         )