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

[no changelog]

core/embed/rust/librust_qstr.h

@@ -107,4 +107,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 flickering 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

@@ -62,8 +62,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>
@@ -82,13 +83,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
@@ -128,6 +135,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) {
@@ -175,8 +189,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.
@@ -186,7 +201,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
@@ -440,8 +455,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

@@ -1193,9 +1193,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: u8 = 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) }
@@ -1205,8 +1209,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) }
@@ -1217,6 +1225,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,
@@ -1227,6 +1236,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) }

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

@@ -48,6 +48,13 @@ class _RustLayout(ui.Layout):
         msg = self.layout.request_complete_repaint()
         assert msg is None
 
+    def _paint(self) -> None:
+        import storage.cache as storage_cache
+
+        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, ...]:
@@ -97,7 +104,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)
 
@@ -127,11 +134,11 @@ 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()
-        self.layout.paint()
+        self._paint()
 
         if __debug__ and self.should_notify_layout_change:
             from apps.debug import notify_layout_change
@@ -149,7 +156,7 @@ class _RustLayout(ui.Layout):
         from trezor import workflow
 
         touch = loop.wait(io.TOUCH)
-        self._before_render()
+        self._first_paint()
         # self.layout.bounds()
         while True:
             # Using `yield` instead of `await` to avoid allocations.
@@ -160,7 +167,7 @@ class _RustLayout(ui.Layout):
                 msg = self.layout.touch_event(event, x, y)
             if msg is not None:
                 raise ui.Result(msg)
-            self.layout.paint()
+            self._paint()
             # self.layout.bounds()
 
     def handle_timers(self) -> loop.Task:  # type: ignore [awaitable-is-generator]
@@ -170,7 +177,7 @@ class _RustLayout(ui.Layout):
             msg = self.layout.timer(token)
             if msg is not None:
                 raise ui.Result(msg)
-            self.layout.paint()
+            self._paint()
 
     def page_count(self) -> int:
         return self.layout.page_count()

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

@@ -8,7 +8,7 @@ import trezorui2
 from . import _RustLayout
 
 if TYPE_CHECKING:
-    from trezor import io
+    from trezor import loop
     from typing import Any, Tuple
 
 
@@ -28,6 +28,23 @@ class HomescreenBase(_RustLayout):
             storage_cache.homescreen_shown = None
             raise
 
+    def _first_paint(self) -> None:
+        from trezor import utils
+
+        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)
+
     # In __debug__ mode, ignore {confirm,swipe,input}_signal.
     def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
         return self.handle_timers(), self.handle_input_and_rendering()
@@ -51,12 +68,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,
             ),
         )
 
@@ -89,10 +108,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,
             ),
         )
 
@@ -107,11 +130,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,
             )
         )