refactor(core/rust): introduce layout lifecycle states on Rust side

core/embed/rust/librust_qstr.h

@@ -6,15 +6,18 @@
 
 static void _librust_qstrs(void) {
   MP_QSTR_;
+  MP_QSTR_ATTACHED;
   MP_QSTR_AttachType;
   MP_QSTR_BacklightLevels;
   MP_QSTR_CANCELLED;
   MP_QSTR_CONFIRMED;
   MP_QSTR_DIM;
+  MP_QSTR_DONE;
   MP_QSTR_INFO;
   MP_QSTR_INITIAL;
   MP_QSTR_LOW;
   MP_QSTR_LayoutObj;
+  MP_QSTR_LayoutState;
   MP_QSTR_MAX;
   MP_QSTR_MESSAGE_NAME;
   MP_QSTR_MESSAGE_WIRE_TYPE;
@@ -28,6 +31,7 @@ static void _librust_qstrs(void) {
   MP_QSTR_SWIPE_RIGHT;
   MP_QSTR_SWIPE_UP;
   MP_QSTR_TR;
+  MP_QSTR_TRANSITIONING;
   MP_QSTR_TranslationsHeader;
   MP_QSTR___del__;
   MP_QSTR___dict__;
@@ -581,6 +585,7 @@ static void _librust_qstrs(void) {
   MP_QSTR_reset__wrong_word_selected;
   MP_QSTR_reset__you_need_one_share;
   MP_QSTR_reset__your_backup_is_done;
+  MP_QSTR_return_value;
   MP_QSTR_reverse;
   MP_QSTR_rotation__change_template;
   MP_QSTR_rotation__east;

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

@@ -79,9 +79,9 @@ macro_rules! obj_fn_kw {
 /// Construct fixed static const `Map` from `key` => `val` pairs.
 macro_rules! obj_map {
     ($($key:expr => $val:expr),*) => ({
-        Map::from_fixed_static(&[
+        $crate::micropython::map::Map::from_fixed_static(&[
             $(
-                Map::at($key, $val),
+                $crate::micropython::map::Map::at($key, $val),
             )*
         ])
     });
@@ -110,6 +110,7 @@ macro_rules! obj_dict {
 /// Compose a `Type` object definition.
 macro_rules! obj_type {
     (name: $name:expr,
+     $(base: $base:expr,)?
      $(locals: $locals:expr,)?
      $(make_new_fn: $make_new_fn:path,)?
      $(attr_fn: $attr_fn:path,)?
@@ -121,6 +122,11 @@ macro_rules! obj_type {
 
             let name = $name.to_u16();
 
+            #[allow(unused_mut)]
+            #[allow(unused_assignments)]
+            let mut base_type: &'static ffi::mp_obj_type_t = &ffi::mp_type_type;
+            $(base_type = &$base;)?
+
             #[allow(unused_mut)]
             #[allow(unused_assignments)]
             let mut attr: ffi::mp_attr_fun_t = None;
@@ -146,7 +152,7 @@ macro_rules! obj_type {
 
             ffi::mp_obj_type_t {
                 base: ffi::mp_obj_base_t {
-                    type_: &ffi::mp_type_type,
+                    type_: base_type,
                 },
                 flags: 0,
                 name,

core/embed/rust/src/ui/button_request.rs

@@ -3,7 +3,7 @@ use num_traits::FromPrimitive;
 
 // ButtonRequestType from messages-common.proto
 // Eventually this should be generated
-#[derive(Clone, Copy, FromPrimitive)]
+#[derive(Clone, Copy, FromPrimitive, PartialEq, Eq)]
 #[repr(u16)]
 pub enum ButtonRequestCode {
     Other = 1,
@@ -41,7 +41,7 @@ impl ButtonRequestCode {
     }
 }
 
-#[derive(Clone)]
+#[derive(Copy, Clone, PartialEq, Eq)]
 pub struct ButtonRequest {
     pub code: ButtonRequestCode,
     pub name: TString<'static>,

core/embed/rust/src/ui/flow/swipe.rs

@@ -14,9 +14,11 @@ use crate::{
         event::SwipeEvent,
         flow::{base::Decision, FlowController},
         geometry::{Direction, Rect},
-        layout::obj::ObjComponent,
+        layout::base::{Layout, LayoutState},
         shape::{render_on_display, ConcreteRenderer, Renderer, ScopedRenderer},
+        ui_features::ModelUI,
         util::animation_disabled,
+        UIFeaturesCommon,
     },
 };
 
@@ -108,7 +110,11 @@ pub struct SwipeFlow {
     internal_pages: u16,
     /// If triggering swipe by event, make this decision instead of default
     /// after the swipe.
-    decision_override: Option<Decision>,
+    pending_decision: Option<Decision>,
+    /// Layout lifecycle state.
+    lifecycle_state: LayoutState,
+    /// Returned value from latest transition, stored as Obj.
+    returned_value: Option<Result<Obj, Error>>,
 }
 
 impl SwipeFlow {
@@ -120,7 +126,9 @@ impl SwipeFlow {
             allow_swipe: true,
             internal_page_idx: 0,
             internal_pages: 1,
-            decision_override: None,
+            pending_decision: None,
+            lifecycle_state: LayoutState::Initial,
+            returned_value: None,
         })
     }
 
@@ -169,6 +177,7 @@ impl SwipeFlow {
 
         // reset and unlock swipe config
         self.swipe = SwipeDetect::new();
+        // unlock swipe events
         self.allow_swipe = true;
 
         // send an Attach event to the new page
@@ -197,7 +206,7 @@ impl SwipeFlow {
         }
     }
 
-    fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<FlowMsg> {
+    fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<LayoutState> {
         let mut decision = Decision::Nothing;
         let mut return_transition: AttachType = AttachType::Initial;
 
@@ -215,20 +224,21 @@ impl SwipeFlow {
 
             match self.swipe.event(ctx, event, config) {
                 Some(SwipeEvent::End(dir)) => {
-                    if let Some(override_decision) = self.decision_override.take() {
-                        decision = override_decision;
-                    } else {
-                        decision = self.handle_swipe_child(ctx, dir);
-                    }
-
                     return_transition = AttachType::Swipe(dir);
 
                     let new_internal_page_idx =
                         config.paging_event(dir, self.internal_page_idx, self.internal_pages);
                     if new_internal_page_idx != self.internal_page_idx {
+                        // internal paging event
                         self.internal_page_idx = new_internal_page_idx;
                         decision = Decision::Nothing;
                         attach = true;
+                    } else if let Some(override_decision) = self.pending_decision.take() {
+                        // end of simulated swipe, applying original decision
+                        decision = override_decision;
+                    } else {
+                        // normal end-of-swipe event handling
+                        decision = self.handle_swipe_child(ctx, dir);
                     }
                     Event::Swipe(SwipeEvent::End(dir))
                 }
@@ -265,31 +275,36 @@ impl SwipeFlow {
 
                 if let Decision::Transition(_, Swipe(direction)) = decision {
                     if config.is_allowed(direction) {
+                        self.allow_swipe = true;
                         if !animation_disabled() {
                             self.swipe.trigger(ctx, direction, config);
-                            self.decision_override = Some(decision);
+                            self.pending_decision = Some(decision);
-                            decision = Decision::Nothing;
+                            return Some(LayoutState::Transitioning(return_transition));
                         }
-                        self.allow_swipe = true;
                     }
                 }
             }
             _ => {
                 //ignore message, we are already transitioning
-                self.current_page_mut().event(ctx, event);
+                let msg = self.current_page_mut().event(ctx, event);
+                assert!(msg.is_none());
             }
         }
 
         match decision {
             Decision::Transition(new_state, attach) => {
                 self.goto(ctx, new_state, attach);
-                None
+                Some(LayoutState::Attached(ctx.button_request().take()))
             }
             Decision::Return(msg) => {
                 ctx.set_transition_out(return_transition);
                 self.swipe.reset();
                 self.allow_swipe = true;
-                Some(msg)
+                self.returned_value = Some(msg.try_into());
+                Some(LayoutState::Done)
+            }
+            Decision::Nothing if matches!(event, Event::Attach(_)) => {
+                Some(LayoutState::Attached(ctx.button_request().take()))
             }
             _ => None,
         }
@@ -307,23 +322,22 @@ impl SwipeFlow {
 /// This implementation relies on the fact that swipe components always return
 /// `FlowMsg` as their `Component::Msg` (provided by `impl FlowComponentTrait`
 /// earlier in this file).
-#[cfg(feature = "micropython")]
+impl Layout<Result<Obj, Error>> for SwipeFlow {
-impl ObjComponent for SwipeFlow {
+    fn place(&mut self) {
-    fn obj_place(&mut self, bounds: Rect) -> Rect {
         for elem in self.store.iter_mut() {
-            elem.place(bounds);
+            elem.place(ModelUI::SCREEN);
-        }
-        bounds
-    }
-
-    fn obj_event(&mut self, ctx: &mut EventCtx, event: Event) -> Result<Obj, Error> {
-        match self.event(ctx, event) {
-            None => Ok(Obj::const_none()),
-            Some(msg) => msg.try_into(),
         }
     }
 
-    fn obj_paint(&mut self) {
+    fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<LayoutState> {
+        self.event(ctx, event)
+    }
+
+    fn value(&self) -> Option<&Result<Obj, Error>> {
+        self.returned_value.as_ref()
+    }
+
+    fn paint(&mut self) {
         render_on_display(None, Some(Color::black()), |target| {
             self.render_state(self.state.index(), target);
         });

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

@@ -0,0 +1,84 @@
+use crate::ui::{
+    button_request::ButtonRequest,
+    component::{base::AttachType, Event, EventCtx},
+};
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum LayoutState {
+    Initial,
+    Attached(Option<ButtonRequest>),
+    Transitioning(AttachType),
+    Done,
+}
+
+pub trait Layout<T> {
+    //fn attach(&mut self, ctx: &mut EventCtx, attach_type: AttachType);
+    fn place(&mut self);
+    fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<LayoutState>;
+    fn value(&self) -> Option<&T>;
+    fn paint(&mut self);
+}
+
+#[cfg(feature = "micropython")]
+mod micropython {
+    use crate::micropython::{
+        macros::{obj_dict, obj_map, obj_type},
+        obj::Obj,
+        qstr::Qstr,
+        simple_type::SimpleTypeObj,
+        typ::Type,
+    };
+
+    use super::LayoutState;
+
+    static STATE_INITIAL_TYPE: Type = obj_type! {
+        name: Qstr::MP_QSTR_INITIAL,
+        base: LAYOUT_STATE_TYPE,
+    };
+
+    static STATE_ATTACHED_TYPE: Type = obj_type! {
+        name: Qstr::MP_QSTR_ATTACHED,
+        base: LAYOUT_STATE_TYPE,
+    };
+
+    static STATE_TRANSITIONING_TYPE: Type = obj_type! {
+        name: Qstr::MP_QSTR_TRANSITIONING,
+        base: LAYOUT_STATE_TYPE,
+    };
+
+    static STATE_DONE_TYPE: Type = obj_type! {
+        name: Qstr::MP_QSTR_DONE,
+        base: LAYOUT_STATE_TYPE,
+    };
+
+    pub static STATE_INITIAL: SimpleTypeObj = SimpleTypeObj::new(&STATE_INITIAL_TYPE);
+    pub static STATE_ATTACHED: SimpleTypeObj = SimpleTypeObj::new(&STATE_ATTACHED_TYPE);
+    pub static STATE_TRANSITIONING: SimpleTypeObj = SimpleTypeObj::new(&STATE_TRANSITIONING_TYPE);
+    pub static STATE_DONE: SimpleTypeObj = SimpleTypeObj::new(&STATE_DONE_TYPE);
+
+    static LAYOUT_STATE_TYPE: Type = obj_type! {
+        name: Qstr::MP_QSTR_LayoutState,
+        locals: &obj_dict! { obj_map! {
+            Qstr::MP_QSTR_INITIAL => STATE_INITIAL.as_obj(),
+            Qstr::MP_QSTR_ATTACHED => STATE_ATTACHED.as_obj(),
+            Qstr::MP_QSTR_TRANSITIONING => STATE_TRANSITIONING.as_obj(),
+            Qstr::MP_QSTR_DONE => STATE_DONE.as_obj(),
+        } },
+    };
+
+    pub static LAYOUT_STATE: SimpleTypeObj = SimpleTypeObj::new(&LAYOUT_STATE_TYPE);
+
+    impl From<LayoutState> for Obj {
+        fn from(state: LayoutState) -> Self {
+            match state {
+                LayoutState::Initial => STATE_INITIAL.as_obj(),
+                LayoutState::Attached(_) => STATE_ATTACHED.as_obj(),
+                LayoutState::Transitioning(_) => STATE_TRANSITIONING.as_obj(),
+                LayoutState::Done => STATE_DONE.as_obj(),
+            }
+        }
+    }
+}
+
+#[cfg(feature = "micropython")]
+pub use micropython::*;

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

@@ -1,3 +1,5 @@
+pub mod base;
+
 #[cfg(feature = "micropython")]
 pub mod obj;
 

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

@@ -1,6 +1,7 @@
 use core::{
     cell::{RefCell, RefMut},
     convert::{TryFrom, TryInto},
+    marker::PhantomData,
     ops::{Deref, DerefMut},
 };
 use num_traits::{FromPrimitive, ToPrimitive};
@@ -8,7 +9,7 @@ use num_traits::{FromPrimitive, ToPrimitive};
 #[cfg(feature = "button")]
 use crate::ui::event::ButtonEvent;
 #[cfg(feature = "new_rendering")]
-use crate::ui::{display::Color, shape::render_on_display};
+use crate::ui::{display::Color, shape::render_on_display, shape::Renderer};
 #[cfg(feature = "touch")]
 use crate::ui::{event::TouchEvent, geometry::Direction};
 use crate::{
@@ -32,9 +33,13 @@ use crate::{
         constant, display,
         event::USBEvent,
         geometry::Rect,
+        ui_features::ModelUI,
+        UIFeaturesCommon,
     },
 };
 
+use super::base::{Layout, LayoutState};
+
 impl AttachType {
     fn to_obj(self) -> Obj {
         match self {
@@ -89,44 +94,119 @@ pub trait ComponentMsgObj: Component {
 /// Note: we need to use an object-safe trait in order to store it in a `Gc<dyn
 /// T>` field. `Component` itself is not object-safe because of `Component::Msg`
 /// associated type.
-pub trait ObjComponent: MaybeTrace {
+// pub trait ObjComponent: MaybeTrace {
-    fn obj_place(&mut self, bounds: Rect) -> Rect;
+//     fn obj_place(&mut self, bounds: Rect) -> Rect;
-    fn obj_event(&mut self, ctx: &mut EventCtx, event: Event) -> Result<Obj, Error>;
+//     fn obj_event(&mut self, ctx: &mut EventCtx, event: Event) -> Result<Obj,
-    fn obj_paint(&mut self);
+// Error>;     fn obj_paint(&mut self);
-    fn obj_bounds(&self, _sink: &mut dyn FnMut(Rect)) {}
+//     fn obj_bounds(&self, _sink: &mut dyn FnMut(Rect)) {}
+// }
+
+// impl<T> ObjComponent for T
+// where
+//     T: Component + ComponentMsgObj + MaybeTrace,
+// {
+//     fn obj_place(&mut self, bounds: Rect) -> Rect {
+//         self.place(bounds)
+//     }
+
+//     fn obj_event(&mut self, ctx: &mut EventCtx, event: Event) -> Result<Obj,
+// Error> {         if let Some(msg) = self.event(ctx, event) {
+//             self.msg_try_into_obj(msg)
+//         } else {
+//             Ok(Obj::const_none())
+//         }
+//     }
+
+//     fn obj_paint(&mut self) {
+//         #[cfg(not(feature = "new_rendering"))]
+//         {
+//             self.paint();
+//         }
+
+//         #[cfg(feature = "new_rendering")]
+//         {
+//             render_on_display(None, Some(Color::black()), |target| {
+//                 self.render(target);
+//             });
+//         }
+//     }
+// }
+
+trait ComponentMaybeTrace: Component + ComponentMsgObj + MaybeTrace {}
+impl<T> ComponentMaybeTrace for T where T: Component + ComponentMsgObj + MaybeTrace {}
+
+struct RootComponent<T, M>
+where
+    T: Component,
+    M: UIFeaturesCommon,
+{
+    inner: T,
+    returned_value: Option<Result<Obj, Error>>,
+    _features: PhantomData<M>,
 }
 
-impl<T> ObjComponent for T
+impl<T, M> RootComponent<T, M>
 where
-    T: Component + ComponentMsgObj + MaybeTrace,
+    T: ComponentMaybeTrace,
+    M: UIFeaturesCommon,
 {
-    fn obj_place(&mut self, bounds: Rect) -> Rect {
+    pub fn new(component: T) -> Self {
-        self.place(bounds)
+        Self {
+            inner: component,
+            returned_value: None,
+            _features: PhantomData,
+        }
+    }
+}
+
+impl<T> Layout<Result<Obj, Error>> for RootComponent<T, ModelUI>
+where
+    T: Component + ComponentMsgObj,
+{
+    fn place(&mut self) {
+        self.inner.place(ModelUI::SCREEN);
     }
 
-    fn obj_event(&mut self, ctx: &mut EventCtx, event: Event) -> Result<Obj, Error> {
+    fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<LayoutState> {
-        if let Some(msg) = self.event(ctx, event) {
+        if let Some(msg) = self.inner.event(ctx, event) {
-            self.msg_try_into_obj(msg)
+            self.returned_value = Some(self.inner.msg_try_into_obj(msg));
+            Some(LayoutState::Done)
+        } else if matches!(event, Event::Attach(_)) {
+            Some(LayoutState::Attached(ctx.button_request().take()))
         } else {
-            Ok(Obj::const_none())
+            None
         }
     }
 
-    fn obj_paint(&mut self) {
+    fn value(&self) -> Option<&Result<Obj, Error>> {
-        #[cfg(not(feature = "new_rendering"))]
+        self.returned_value.as_ref()
-        {
+    }
-            self.paint();
-        }
 
+    fn paint(&mut self) {
+        #[cfg(not(feature = "new_rendering"))]
+        self.inner.paint();
         #[cfg(feature = "new_rendering")]
         {
             render_on_display(None, Some(Color::black()), |target| {
-                self.render(target);
+                self.inner.render(target);
             });
         }
     }
 }
 
+#[cfg(feature = "ui_debug")]
+impl<T> crate::trace::Trace for RootComponent<T, ModelUI>
+where
+    T: Component + crate::trace::Trace,
+{
+    fn trace(&self, t: &mut dyn crate::trace::Tracer) {
+        self.inner.trace(t);
+    }
+}
+
+trait LayoutMaybeTrace: Layout<Result<Obj, Error>> + MaybeTrace {}
+impl<T> LayoutMaybeTrace for T where T: Layout<Result<Obj, Error>> + MaybeTrace {}
+
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "debug", derive(ufmt::derive::uDebug))]
 enum Repaint {
@@ -145,31 +225,33 @@ pub struct LayoutObj {
 }
 
 struct LayoutObjInner {
-    root: Option<GcBox<dyn ObjComponent>>,
+    root: Option<GcBox<dyn LayoutMaybeTrace>>,
     event_ctx: EventCtx,
     timer_fn: Obj,
     page_count: u16,
     repaint: Repaint,
     transition_out: AttachType,
+    button_request: Option<ButtonRequest>,
 }
 
 impl LayoutObjInner {
     /// Create a new `LayoutObj`, wrapping a root component.
     #[inline(never)]
-    pub fn new(root: impl ObjComponent + 'static) -> Result<Self, Error> {
+    pub fn new(root: impl LayoutMaybeTrace + 'static) -> Result<Self, Error> {
         let root = GcBox::new(root)?;
 
         let mut new = Self {
-            root: Some(gc::coerce!(ObjComponent, root)),
+            root: Some(gc::coerce!(LayoutMaybeTrace, root)),
             event_ctx: EventCtx::new(),
             timer_fn: Obj::const_none(),
             page_count: 1,
             repaint: Repaint::Full,
             transition_out: AttachType::Initial,
+            button_request: None,
         };
 
         // invoke the initial placement
-        new.root_mut().obj_place(constant::screen());
+        new.root_mut().place();
         // cause a repaint pass to update the number of pages
         let msg = new.obj_event(Event::RequestPaint);
         assert!(matches!(msg, Ok(s) if s == Obj::const_none()));
@@ -187,22 +269,20 @@ impl LayoutObjInner {
         self.timer_fn = timer_fn;
     }
 
-    fn root(&self) -> &impl Deref<Target = dyn ObjComponent> {
+    fn root(&self) -> &impl Deref<Target = dyn LayoutMaybeTrace> {
         unwrap!(self.root.as_ref())
     }
 
-    fn root_mut(&mut self) -> &mut impl DerefMut<Target = dyn ObjComponent> {
+    fn root_mut(&mut self) -> &mut impl DerefMut<Target = dyn LayoutMaybeTrace> {
         unwrap!(self.root.as_mut())
     }
 
     fn obj_request_repaint(&mut self) {
         self.repaint = Repaint::Full;
         let mut event_ctx = EventCtx::new();
-        let paint_msg = self
+        let paint_msg = self.root_mut().event(&mut event_ctx, Event::RequestPaint);
-            .root_mut()
+        // paint_msg must not change the state
-            .obj_event(&mut event_ctx, Event::RequestPaint);
+        assert!(paint_msg.is_none());
-        // paint_msg must not be an error and it must not return a result
-        assert!(matches!(paint_msg, Ok(s) if s == Obj::const_none()));
         // there must be no timers set
         assert!(event_ctx.pop_timer().is_none());
     }
@@ -217,12 +297,22 @@ impl LayoutObjInner {
         // Get the event context ready for a new event
         self.event_ctx.clear();
 
-        // Send the event down the component tree. Bail out in case of failure.
+        // Send the event down the component tree.
-        let msg = root.obj_event(&mut self.event_ctx, event)?;
+        let msg = root.event(&mut self.event_ctx, event);
+
+        match msg {
+            Some(LayoutState::Done) => return Ok(msg.into()), // short-circuit
+            Some(LayoutState::Attached(br)) => {
+                assert!(self.button_request.is_none());
+                self.button_request = br;
+            }
+            Some(LayoutState::Transitioning(t)) => self.transition_out = t,
+            _ => (),
+        };
 
         // Place the root component on the screen in case it was requested.
         if self.event_ctx.needs_place() {
-            root.obj_place(constant::screen());
+            root.place();
         }
 
         // Check if we should repaint next time
@@ -251,12 +341,7 @@ impl LayoutObjInner {
             self.page_count = count as u16;
         }
 
-        // Update outgoing transition if set
+        Ok(msg.into())
-        if let Some(t) = self.event_ctx.get_transition_out() {
-            self.transition_out = t;
-        }
-
-        Ok(msg)
     }
 
     /// Run a paint pass over the component tree. Returns true if any component
@@ -270,7 +355,7 @@ impl LayoutObjInner {
 
         if self.repaint != Repaint::None {
             self.repaint = Repaint::None;
-            self.root_mut().obj_paint();
+            self.root_mut().paint();
             true
         } else {
             false
@@ -291,10 +376,10 @@ impl LayoutObjInner {
         // For Reasons(tm), we must pass a closure in which we call `root.trace(t)`,
         // instead of passing `root` into the tracer.
 
-        // (The Reasons being, root is a `Gc<dyn ObjComponent>`, and `Gc` does not
+        // (The Reasons being, root is a `Gc<dyn LayoutMaybeTrace>`, and `Gc` does not
-        // implement `Trace`, and `dyn ObjComponent` is unsized so we can't deref it to
+        // implement `Trace`, and `dyn LayoutMaybeTrace` is unsized so we can't deref it
-        // claim that it implements `Trace`, and we also can't upcast it to `&dyn Trace`
+        // to claim that it implements `Trace`, and we also can't upcast it to
-        // because trait upcasting is unstable.
+        // `&dyn Trace` because trait upcasting is unstable.
         // Luckily, calling `root.trace()` works perfectly fine in spite of the above.)
         tracer.root(&|t| {
             self.root().trace(t);
@@ -306,7 +391,7 @@ impl LayoutObjInner {
     }
 
     fn obj_button_request(&mut self) -> Result<Obj, Error> {
-        match self.event_ctx.button_request() {
+        match self.button_request.take() {
             None => Ok(Obj::const_none()),
             Some(ButtonRequest { code, name }) => (code.num().into(), name.try_into()?).try_into(),
         }
@@ -315,11 +400,23 @@ impl LayoutObjInner {
     fn obj_get_transition_out(&self) -> Obj {
         self.transition_out.to_obj()
     }
+
+    fn obj_return_value(&self) -> Result<Obj, Error> {
+        self.root()
+            .value()
+            .cloned()
+            .unwrap_or(Ok(Obj::const_none()))
+    }
 }
 
 impl LayoutObj {
     /// Create a new `LayoutObj`, wrapping a root component.
-    pub fn new(root: impl ObjComponent + 'static) -> Result<Gc<Self>, Error> {
+    pub fn new<T: ComponentMaybeTrace + 'static>(root: T) -> Result<Gc<Self>, Error> {
+        let root_component = RootComponent::new(root);
+        Self::new_root(root_component)
+    }
+
+    pub fn new_root(root: impl LayoutMaybeTrace + 'static) -> Result<Gc<Self>, Error> {
         // SAFETY: This is a Python object and hase a base as first element
         unsafe {
             Gc::new_with_custom_finaliser(Self {
@@ -350,6 +447,7 @@ impl LayoutObj {
                 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(),
                 Qstr::MP_QSTR_get_transition_out => obj_fn_1!(ui_layout_get_transition_out).as_obj(),
+                Qstr::MP_QSTR_return_value => obj_fn_1!(ui_layout_return_value).as_obj(),
             }),
         };
         &TYPE
@@ -427,9 +525,8 @@ extern "C" fn ui_layout_attach_timer_fn(this: Obj, timer_fn: Obj, attach_type: O
 
         let msg = this
             .inner_mut()
-            .obj_event(Event::Attach(AttachType::try_from_obj(attach_type)?))?;
+            .obj_event(Event::Attach(AttachType::try_from_obj(attach_type)?));
-        assert!(msg == Obj::const_none());
+        msg
-        Ok(Obj::const_none())
     };
     unsafe { util::try_or_raise(block) }
 }
@@ -560,6 +657,15 @@ extern "C" fn ui_layout_get_transition_out(this: Obj) -> Obj {
     unsafe { util::try_or_raise(block) }
 }
 
+extern "C" fn ui_layout_return_value(this: Obj) -> Obj {
+    let block = || {
+        let this: Gc<LayoutObj> = this.try_into()?;
+        let value = this.inner_mut().obj_return_value();
+        value
+    };
+    unsafe { util::try_or_raise(block) }
+}
+
 #[cfg(feature = "ui_debug")]
 #[no_mangle]
 pub extern "C" fn ui_debug_layout_type() -> &'static Type {

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

@@ -1,4 +1,9 @@
-use crate::micropython::{macros::obj_type, qstr::Qstr, simple_type::SimpleTypeObj, typ::Type};
+use crate::micropython::{
+    macros::{obj_dict, obj_map, obj_type},
+    qstr::Qstr,
+    simple_type::SimpleTypeObj,
+    typ::Type,
+};
 
 static CONFIRMED_TYPE: Type = obj_type! { name: Qstr::MP_QSTR_CONFIRMED, };
 static CANCELLED_TYPE: Type = obj_type! { name: Qstr::MP_QSTR_CANCELLED, };

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

@@ -13,7 +13,6 @@ pub enum BinarySelectionMsg {
 
 /// Component presenting a binary choice represented as two buttons, left and
 /// right. Both buttons are parameterized with content and style.
-#[derive(Clone)]
 pub struct BinarySelection {
     buttons_area: Rect,
     button_left: Button,

core/embed/rust/src/ui/model_mercury/flow/confirm_action.rs

@@ -271,7 +271,7 @@ fn new_confirm_action_uni<T: Component + MaybeTrace + 'static>(
 
     let flow = create_confirm(flow, strings.subtitle, hold, prompt_screen)?;
 
-    Ok(LayoutObj::new(flow)?.into())
+    Ok(LayoutObj::new_root(flow)?.into())
 }
 
 fn create_flow(

core/embed/rust/src/ui/model_mercury/flow/confirm_fido.rs

@@ -214,6 +214,6 @@ impl ConfirmFido {
             .with_page(&ConfirmFido::Details, content_details)?
             .with_page(&ConfirmFido::Tap, content_tap)?
             .with_page(&ConfirmFido::Menu, content_menu)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/confirm_firmware_update.rs

@@ -137,6 +137,6 @@ impl ConfirmFirmwareUpdate {
             .with_page(&ConfirmFirmwareUpdate::Menu, content_menu)?
             .with_page(&ConfirmFirmwareUpdate::Fingerprint, content_fingerprint)?
             .with_page(&ConfirmFirmwareUpdate::Confirm, content_confirm)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/confirm_output.rs

@@ -450,5 +450,5 @@ fn new_confirm_output_obj(_args: &[Obj], kwargs: &Map) -> Result<Obj, error::Err
             .with_page(&ConfirmOutput::CancelTap, get_cancel_page())?
     };
 
-    Ok(LayoutObj::new(res)?.into())
+    Ok(LayoutObj::new_root(res)?.into())
 }

core/embed/rust/src/ui/model_mercury/flow/confirm_reset.rs

@@ -166,5 +166,5 @@ fn new_confirm_reset_obj(_args: &[Obj], kwargs: &Map) -> Result<Obj, error::Erro
             .with_page(&ConfirmResetCreate::Menu, content_menu)?
             .with_page(&ConfirmResetCreate::Confirm, content_confirm)?
     };
-    Ok(LayoutObj::new(res)?.into())
+    Ok(LayoutObj::new_root(res)?.into())
 }

core/embed/rust/src/ui/model_mercury/flow/confirm_set_new_pin.rs

@@ -138,6 +138,6 @@ impl SetNewPin {
             .with_page(&SetNewPin::Menu, content_menu)?
             .with_page(&SetNewPin::CancelPinIntro, content_cancel_intro)?
             .with_page(&SetNewPin::CancelPinConfirm, content_cancel_confirm)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/confirm_summary.rs

@@ -203,6 +203,6 @@ impl ConfirmSummary {
             .with_page(&ConfirmSummary::AccountInfo, content_account)?
             .with_page(&ConfirmSummary::CancelTap, content_cancel_tap)?;
 
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/continue_recovery.rs

@@ -360,5 +360,5 @@ fn new_obj(_args: &[Obj], kwargs: &Map) -> Result<Obj, error::Error> {
                 content_remaining_shares,
             )?
     };
-    Ok(LayoutObj::new(res)?.into())
+    Ok(LayoutObj::new_root(res)?.into())
 }

core/embed/rust/src/ui/model_mercury/flow/get_address.rs

@@ -233,6 +233,6 @@ impl GetAddress {
             .with_page(&GetAddress::AccountInfo, content_account)?
             .with_page(&GetAddress::Cancel, content_cancel_info)?
             .with_page(&GetAddress::CancelTap, content_cancel_tap)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/prompt_backup.rs

@@ -141,6 +141,6 @@ impl PromptBackup {
             .with_page(&PromptBackup::Menu, content_menu)?
             .with_page(&PromptBackup::SkipBackupIntro, content_skip_intro)?
             .with_page(&PromptBackup::SkipBackupConfirm, content_skip_confirm)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/request_number.rs

@@ -132,6 +132,6 @@ impl RequestNumber {
             .with_page(&RequestNumber::Number, content_number_input)?
             .with_page(&RequestNumber::Menu, content_menu)?
             .with_page(&RequestNumber::Info, content_info)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/request_passphrase.rs

@@ -81,6 +81,6 @@ impl RequestPassphrase {
         let res = SwipeFlow::new(&RequestPassphrase::Keypad)?
             .with_page(&RequestPassphrase::Keypad, content_keypad)?
             .with_page(&RequestPassphrase::ConfirmEmpty, content_confirm_empty)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/set_brightness.rs

@@ -139,6 +139,6 @@ impl SetBrightness {
             .with_page(&SetBrightness::Confirm, content_confirm)?
             .with_page(&SetBrightness::Confirmed, content_confirmed)?;
 
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/show_share_words.rs

@@ -160,6 +160,6 @@ impl ShowShareWords {
                 &ShowShareWords::CheckBackupIntro,
                 content_check_backup_intro,
             )?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/show_tutorial.rs

@@ -202,6 +202,6 @@ impl ShowTutorial {
             .with_page(&ShowTutorial::Menu, content_menu)?
             .with_page(&ShowTutorial::DidYouKnow, content_did_you_know)?
             .with_page(&ShowTutorial::HoldToExit, content_hold_to_exit)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

core/embed/rust/src/ui/model_mercury/flow/warning_hi_prio.rs

@@ -117,6 +117,6 @@ impl WarningHiPrio {
             .with_page(&WarningHiPrio::Message, content_message)?
             .with_page(&WarningHiPrio::Menu, content_menu)?
             .with_page(&WarningHiPrio::Cancelled, content_cancelled)?;
-        Ok(LayoutObj::new(res)?.into())
+        Ok(LayoutObj::new_root(res)?.into())
     }
 }

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

@@ -44,9 +44,7 @@ use crate::{
         flow::Swipable,
         geometry::{self, Direction},
         layout::{
-            obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ},
+            base::LAYOUT_STATE, obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ}, result::{CANCELLED, CONFIRMED, INFO}, util::{upy_disable_animation, ConfirmBlob, PropsList, RecoveryType}
-            result::{CANCELLED, CONFIRMED, INFO},
-            util::{upy_disable_animation, ConfirmBlob, PropsList, RecoveryType},
         },
         model_mercury::{
             component::{check_homescreen_format, SwipeContent},
@@ -883,7 +881,7 @@ extern "C" fn new_confirm_with_info(n_args: usize, args: *const Obj, kwargs: *mu
 
         let flow =
             confirm_with_info::new_confirm_with_info(title, button, info_button, paragraphs)?;
-        Ok(LayoutObj::new(flow)?.into())
+        Ok(LayoutObj::new_root(flow)?.into())
     };
     unsafe { util::try_with_args_and_kwargs(n_args, args, kwargs, block) }
 }
@@ -1206,82 +1204,6 @@ extern "C" fn new_confirm_fido(n_args: usize, args: *const Obj, kwargs: *mut Map
 #[no_mangle]
 pub static mp_module_trezorui2: Module = obj_module! {
     /// from trezor import utils
-    ///
-    /// T = TypeVar("T")
-    ///
-    /// class LayoutObj(Generic[T]):
-    ///     """Representation of a Rust-based layout object.
-    ///     see `trezor::ui::layout::obj::LayoutObj`.
-    ///     """
-    ///
-    ///     def attach_timer_fn(self, fn: Callable[[int, int], None], attach_type: AttachType | None) -> None:
-    ///         """Attach a timer setter function.
-    ///
-    ///         The layout object can call the timer setter with two arguments,
-    ///         `token` and `duration_ms`. When `duration_ms` is reached, the layout object
-    ///         expects a callback to `self.timer(token)`.
-    ///         """
-    ///
-    ///     if utils.USE_TOUCH:
-    ///         def touch_event(self, event: int, x: int, y: int) -> T | None:
-    ///             """Receive a touch event `event` at coordinates `x`, `y`."""
-    ///
-    ///     if utils.USE_BUTTON:
-    ///         def button_event(self, event: int, button: int) -> T | None:
-    ///             """Receive a button event `event` for button `button`."""
-    ///
-    ///     def progress_event(self, value: int, description: str) -> T | None:
-    ///         """Receive a progress event."""
-    ///
-    ///     def usb_event(self, connected: bool) -> T | None:
-    ///         """Receive a USB connect/disconnect event."""
-    ///
-    ///     def timer(self, token: int) -> T | None:
-    ///         """Callback for the timer set by `attach_timer_fn`.
-    ///
-    ///         This function should be called by the executor after the corresponding
-    ///         duration has expired.
-    ///         """
-    ///
-    ///     def paint(self) -> bool:
-    ///         """Paint the layout object on screen.
-    ///
-    ///         Will only paint updated parts of the layout as required.
-    ///         Returns True if any painting actually happened.
-    ///         """
-    ///
-    ///     def request_complete_repaint(self) -> None:
-    ///         """Request a complete repaint of the screen.
-    ///
-    ///         Does not repaint the screen, a subsequent call to `paint()` is required.
-    ///         """
-    ///
-    ///     if __debug__:
-    ///         def trace(self, tracer: Callable[[str], None]) -> None:
-    ///             """Generate a JSON trace of the layout object.
-    ///
-    ///             The JSON can be emitted as a sequence of calls to `tracer`, each of
-    ///             which is not necessarily a valid JSON chunk. The caller must
-    ///             reassemble the chunks to get a sensible result.
-    ///             """
-    ///
-    ///         def bounds(self) -> None:
-    ///             """Paint bounds of individual components on screen."""
-    ///
-    ///     def page_count(self) -> int:
-    ///         """Return the number of pages in the layout object."""
-    ///
-    ///     def get_transition_out(self) -> AttachType:
-    ///         """Return the transition type."""
-    ///
-    ///     def __del__(self) -> None:
-    ///         """Calls drop on contents of the root component."""
-    ///
-    /// class UiResult:
-    ///    """Result of a UI operation."""
-    ///    pass
-    ///
-    /// mock:global
     Qstr::MP_QSTR___name__ => Qstr::MP_QSTR_trezorui2.to_obj(),
 
     /// CONFIRMED: UiResult
@@ -1818,4 +1740,11 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     SWIPE_RIGHT: ClassVar[int]
     Qstr::MP_QSTR_AttachType => ATTACH_TYPE_OBJ.as_obj(),
 
+    /// class LayoutState:
+    ///     """Layout state."""
+    ///     INITIAL: "ClassVar[LayoutState]"
+    ///     ATTACHED: "ClassVar[LayoutState]"
+    ///     TRANSITIONING: "ClassVar[LayoutState]"
+    ///     DONE: "ClassVar[LayoutState]"
+    Qstr::MP_QSTR_LayoutState => LAYOUT_STATE.as_obj(),
 };

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

@@ -47,9 +47,7 @@ use crate::{
         },
         geometry,
         layout::{
-            obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ},
+            base::LAYOUT_STATE, obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ}, result::{CANCELLED, CONFIRMED, INFO}, util::{upy_disable_animation, ConfirmBlob, RecoveryType}
-            result::{CANCELLED, CONFIRMED, INFO},
-            util::{upy_disable_animation, ConfirmBlob, RecoveryType},
         },
         model_tr::component::check_homescreen_format,
     },
@@ -2144,4 +2142,12 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     SWIPE_LEFT: ClassVar[int]
     ///     SWIPE_RIGHT: ClassVar[int]
     Qstr::MP_QSTR_AttachType => ATTACH_TYPE_OBJ.as_obj(),
+
+    /// class LayoutState:
+    ///     """Layout state."""
+    ///     INITIAL: "ClassVar[LayoutState]"
+    ///     ATTACHED: "ClassVar[LayoutState]"
+    ///     TRANSITIONING: "ClassVar[LayoutState]"
+    ///     DONE: "ClassVar[LayoutState]"
+    Qstr::MP_QSTR_LayoutState => LAYOUT_STATE.as_obj(),
 };

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

@@ -50,9 +50,7 @@ use crate::{
         },
         geometry,
         layout::{
-            obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ},
+            base::LAYOUT_STATE, obj::{ComponentMsgObj, LayoutObj, ATTACH_TYPE_OBJ}, result::{CANCELLED, CONFIRMED, INFO}, util::{upy_disable_animation, ConfirmBlob, PropsList, RecoveryType}
-            result::{CANCELLED, CONFIRMED, INFO},
-            util::{upy_disable_animation, ConfirmBlob, PropsList, RecoveryType},
         },
         model_tt::component::check_homescreen_format,
     },
@@ -1660,7 +1658,7 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     see `trezor::ui::layout::obj::LayoutObj`.
     ///     """
     ///
-    ///     def attach_timer_fn(self, fn: Callable[[int, int], None], attach_type: AttachType | None) -> None:
+    ///     def attach_timer_fn(self, fn: Callable[[int, int], None], attach_type: AttachType | None) -> LayoutState | None:
     ///         """Attach a timer setter function.
     ///
     ///         The layout object can call the timer setter with two arguments,
@@ -1669,20 +1667,20 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///         """
     ///
     ///     if utils.USE_TOUCH:
-    ///         def touch_event(self, event: int, x: int, y: int) -> T | None:
+    ///         def touch_event(self, event: int, x: int, y: int) -> LayoutState | None:
     ///             """Receive a touch event `event` at coordinates `x`, `y`."""
     ///
     ///     if utils.USE_BUTTON:
-    ///         def button_event(self, event: int, button: int) -> T | None:
+    ///         def button_event(self, event: int, button: int) -> LayoutState | None:
     ///             """Receive a button event `event` for button `button`."""
     ///
-    ///     def progress_event(self, value: int, description: str) -> T | None:
+    ///     def progress_event(self, value: int, description: str) -> LayoutState | None:
     ///         """Receive a progress event."""
     ///
-    ///     def usb_event(self, connected: bool) -> T | None:
+    ///     def usb_event(self, connected: bool) -> LayoutState | None:
     ///         """Receive a USB connect/disconnect event."""
     ///
-    ///     def timer(self, token: int) -> T | None:
+    ///     def timer(self, token: int) -> LayoutState | None:
     ///         """Callback for the timer set by `attach_timer_fn`.
     ///
     ///         This function should be called by the executor after the corresponding
@@ -1722,6 +1720,9 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///
     ///     def get_transition_out(self) -> AttachType:
     ///         """Return the transition type."""
+    /// 
+    ///     def return_value(self) -> T:
+    ///         """Retrieve the return value of the layout object."""
     ///
     ///     def __del__(self) -> None:
     ///         """Calls drop on contents of the root component."""
@@ -2199,6 +2200,14 @@ pub static mp_module_trezorui2: Module = obj_module! {
     ///     SWIPE_LEFT: ClassVar[int]
     ///     SWIPE_RIGHT: ClassVar[int]
     Qstr::MP_QSTR_AttachType => ATTACH_TYPE_OBJ.as_obj(),
+
+    /// class LayoutState:
+    ///     """Layout state."""
+    ///     INITIAL: "ClassVar[LayoutState]"
+    ///     ATTACHED: "ClassVar[LayoutState]"
+    ///     TRANSITIONING: "ClassVar[LayoutState]"
+    ///     DONE: "ClassVar[LayoutState]"
+    Qstr::MP_QSTR_LayoutState => LAYOUT_STATE.as_obj(),
 };
 
 #[cfg(test)]

core/mocks/generated/trezorui2.pyi

@@ -1,64 +1,5 @@
 from typing import *
 from trezor import utils
-T = TypeVar("T")
-
-
-# rust/src/ui/model_mercury/layout.rs
-class LayoutObj(Generic[T]):
-    """Representation of a Rust-based layout object.
-    see `trezor::ui::layout::obj::LayoutObj`.
-    """
-    def attach_timer_fn(self, fn: Callable[[int, int], None], attach_type: AttachType | None) -> None:
-        """Attach a timer setter function.
-        The layout object can call the timer setter with two arguments,
-        `token` and `duration_ms`. When `duration_ms` is reached, the layout object
-        expects a callback to `self.timer(token)`.
-        """
-    if utils.USE_TOUCH:
-        def touch_event(self, event: int, x: int, y: int) -> T | None:
-            """Receive a touch event `event` at coordinates `x`, `y`."""
-    if utils.USE_BUTTON:
-        def button_event(self, event: int, button: int) -> T | None:
-            """Receive a button event `event` for button `button`."""
-    def progress_event(self, value: int, description: str) -> T | None:
-        """Receive a progress event."""
-    def usb_event(self, connected: bool) -> T | None:
-        """Receive a USB connect/disconnect event."""
-    def timer(self, token: int) -> T | None:
-        """Callback for the timer set by `attach_timer_fn`.
-        This function should be called by the executor after the corresponding
-        duration has expired.
-        """
-    def paint(self) -> bool:
-        """Paint the layout object on screen.
-        Will only paint updated parts of the layout as required.
-        Returns True if any painting actually happened.
-        """
-    def request_complete_repaint(self) -> None:
-        """Request a complete repaint of the screen.
-        Does not repaint the screen, a subsequent call to `paint()` is required.
-        """
-    if __debug__:
-        def trace(self, tracer: Callable[[str], None]) -> None:
-            """Generate a JSON trace of the layout object.
-            The JSON can be emitted as a sequence of calls to `tracer`, each of
-            which is not necessarily a valid JSON chunk. The caller must
-            reassemble the chunks to get a sensible result.
-            """
-        def bounds(self) -> None:
-            """Paint bounds of individual components on screen."""
-    def page_count(self) -> int:
-        """Return the number of pages in the layout object."""
-    def get_transition_out(self) -> AttachType:
-        """Return the transition type."""
-    def __del__(self) -> None:
-        """Calls drop on contents of the root component."""
-
-
-# rust/src/ui/model_mercury/layout.rs
-class UiResult:
-   """Result of a UI operation."""
-   pass
 CONFIRMED: UiResult
 CANCELLED: UiResult
 INFO: UiResult
@@ -634,6 +575,15 @@ class AttachType:
     SWIPE_DOWN: ClassVar[int]
     SWIPE_LEFT: ClassVar[int]
     SWIPE_RIGHT: ClassVar[int]
+
+
+# rust/src/ui/model_mercury/layout.rs
+class LayoutState:
+    """Layout state."""
+    INITIAL: "ClassVar[LayoutState]"
+    ATTACHED: "ClassVar[LayoutState]"
+    TRANSITIONING: "ClassVar[LayoutState]"
+    DONE: "ClassVar[LayoutState]"
 CONFIRMED: UiResult
 CANCELLED: UiResult
 INFO: UiResult
@@ -1123,6 +1073,15 @@ class AttachType:
     SWIPE_DOWN: ClassVar[int]
     SWIPE_LEFT: ClassVar[int]
     SWIPE_RIGHT: ClassVar[int]
+
+
+# rust/src/ui/model_tr/layout.rs
+class LayoutState:
+    """Layout state."""
+    INITIAL: "ClassVar[LayoutState]"
+    ATTACHED: "ClassVar[LayoutState]"
+    TRANSITIONING: "ClassVar[LayoutState]"
+    DONE: "ClassVar[LayoutState]"
 from trezor import utils
 T = TypeVar("T")
 
@@ -1139,16 +1098,16 @@ class LayoutObj(Generic[T]):
         expects a callback to `self.timer(token)`.
         """
     if utils.USE_TOUCH:
-        def touch_event(self, event: int, x: int, y: int) -> T | None:
+        def touch_event(self, event: int, x: int, y: int) -> LayoutState | None:
             """Receive a touch event `event` at coordinates `x`, `y`."""
     if utils.USE_BUTTON:
-        def button_event(self, event: int, button: int) -> T | None:
+        def button_event(self, event: int, button: int) -> LayoutState | None:
             """Receive a button event `event` for button `button`."""
-    def progress_event(self, value: int, description: str) -> T | None:
+    def progress_event(self, value: int, description: str) -> LayoutState | None:
         """Receive a progress event."""
-    def usb_event(self, connected: bool) -> T | None:
+    def usb_event(self, connected: bool) -> LayoutState | None:
         """Receive a USB connect/disconnect event."""
-    def timer(self, token: int) -> T | None:
+    def timer(self, token: int) -> LayoutState | None:
         """Callback for the timer set by `attach_timer_fn`.
         This function should be called by the executor after the corresponding
         duration elapses.
@@ -1177,6 +1136,9 @@ class LayoutObj(Generic[T]):
         """Return (code, type) of button request made during the last event or timer pass."""
     def get_transition_out(self) -> AttachType:
         """Return the transition type."""
+
+    def return_value(self) -> T:
+        """Retrieve the return value of the layout object."""
     def __del__(self) -> None:
         """Calls drop on contents of the root component."""
 
@@ -1689,3 +1651,12 @@ class AttachType:
     SWIPE_DOWN: ClassVar[int]
     SWIPE_LEFT: ClassVar[int]
     SWIPE_RIGHT: ClassVar[int]
+
+
+# rust/src/ui/model_tt/layout.rs
+class LayoutState:
+    """Layout state."""
+    INITIAL: "ClassVar[LayoutState]"
+    ATTACHED: "ClassVar[LayoutState]"
+    TRANSITIONING: "ClassVar[LayoutState]"
+    DONE: "ClassVar[LayoutState]"