From 9c14cae6569f1179fcdfc3b0e83d56293425ea23 Mon Sep 17 00:00:00 2001
From: Martin Milata <martin@martinmilata.cz>
Date: Thu, 23 May 2024 15:35:51 +0200
Subject: [PATCH] fix(core/ui): allow component to terminate flow on swipe

[skip_ci]
---
 .../rust/src/ui/component/button_request.rs   |  6 +--
 core/embed/rust/src/ui/component/map.rs       | 23 +++++++----
 core/embed/rust/src/ui/flow/base.rs           | 40 ++++++++++++++++---
 core/embed/rust/src/ui/flow/mod.rs            |  2 +-
 core/embed/rust/src/ui/flow/page.rs           | 30 +++++++++-----
 core/embed/rust/src/ui/flow/store.rs          | 26 ++++++++----
 core/embed/rust/src/ui/flow/swipe.rs          | 10 ++---
 .../src/ui/model_mercury/component/dialog.rs  |  2 +-
 .../src/ui/model_mercury/component/frame.rs   |  7 ++--
 .../model_mercury/component/prompt_screen.rs  |  2 +-
 .../ui/model_mercury/component/share_words.rs | 16 +++++---
 .../model_mercury/component/vertical_menu.rs  |  2 +-
 .../rust/src/ui/model_mercury/flow/util.rs    |  4 +-
 13 files changed, 117 insertions(+), 53 deletions(-)

diff --git a/core/embed/rust/src/ui/component/button_request.rs b/core/embed/rust/src/ui/component/button_request.rs
index a5fb1c592..f19eca1c3 100644
--- a/core/embed/rust/src/ui/component/button_request.rs
+++ b/core/embed/rust/src/ui/component/button_request.rs
@@ -65,15 +65,15 @@ pub trait ButtonRequestExt {
 impl<T: Component> ButtonRequestExt for T {}
 
 #[cfg(all(feature = "micropython", feature = "touch", feature = "new_rendering"))]
-impl<T> crate::ui::flow::Swipable for OneButtonRequest<T>
+impl<T> crate::ui::flow::Swipable<T::Msg> for OneButtonRequest<T>
 where
-    T: Component + crate::ui::flow::Swipable,
+    T: Component + crate::ui::flow::Swipable<T::Msg>,
 {
     fn swipe_start(
         &mut self,
         ctx: &mut EventCtx,
         direction: crate::ui::component::SwipeDirection,
-    ) -> bool {
+    ) -> crate::ui::flow::SwipableResult<T::Msg> {
         self.inner.swipe_start(ctx, direction)
     }
 
diff --git a/core/embed/rust/src/ui/component/map.rs b/core/embed/rust/src/ui/component/map.rs
index ce1486d7d..bc7fcf6cf 100644
--- a/core/embed/rust/src/ui/component/map.rs
+++ b/core/embed/rust/src/ui/component/map.rs
@@ -52,12 +52,17 @@ where
 }
 
 #[cfg(all(feature = "micropython", feature = "touch", feature = "new_rendering"))]
-impl<T, F> crate::ui::flow::Swipable for MsgMap<T, F>
+impl<T, F, U> crate::ui::flow::Swipable<U> for MsgMap<T, F>
 where
-    T: Component + crate::ui::flow::Swipable,
+    T: Component + crate::ui::flow::Swipable<T::Msg>,
+    F: Fn(T::Msg) -> Option<U>,
 {
-    fn swipe_start(&mut self, ctx: &mut EventCtx, direction: super::SwipeDirection) -> bool {
-        self.inner.swipe_start(ctx, direction)
+    fn swipe_start(
+        &mut self,
+        ctx: &mut EventCtx,
+        direction: super::SwipeDirection,
+    ) -> crate::ui::flow::SwipableResult<U> {
+        self.inner.swipe_start(ctx, direction).map(&self.func)
     }
 
     fn swipe_finished(&self) -> bool {
@@ -118,11 +123,15 @@ where
 }
 
 #[cfg(all(feature = "micropython", feature = "touch", feature = "new_rendering"))]
-impl<T, F> crate::ui::flow::Swipable for PageMap<T, F>
+impl<T, F> crate::ui::flow::Swipable<T::Msg> for PageMap<T, F>
 where
-    T: Component + crate::ui::flow::Swipable,
+    T: Component + crate::ui::flow::Swipable<T::Msg>,
 {
-    fn swipe_start(&mut self, ctx: &mut EventCtx, direction: super::SwipeDirection) -> bool {
+    fn swipe_start(
+        &mut self,
+        ctx: &mut EventCtx,
+        direction: super::SwipeDirection,
+    ) -> crate::ui::flow::SwipableResult<T::Msg> {
         self.inner.swipe_start(ctx, direction)
     }
 
diff --git a/core/embed/rust/src/ui/flow/base.rs b/core/embed/rust/src/ui/flow/base.rs
index d24f5982f..264525e52 100644
--- a/core/embed/rust/src/ui/flow/base.rs
+++ b/core/embed/rust/src/ui/flow/base.rs
@@ -4,12 +4,18 @@ use num_traits::ToPrimitive;
 /// Component must implement this trait in order to be part of swipe-based flow.
 ///
 /// Default implementation ignores every swipe.
-pub trait Swipable {
-    /// Attempt a swipe. Return false if component in its current state doesn't
-    /// accept a swipe in the given direction. Start a transition animation
-    /// if true is returned.
-    fn swipe_start(&mut self, _ctx: &mut EventCtx, _direction: SwipeDirection) -> bool {
-        false
+pub trait Swipable<T> {
+    /// Attempt a swipe. Return `Ignored` if the component in its current state
+    /// doesn't accept a swipe in that direction. Return `Animating` if
+    /// component accepted the swipe and started a transition animation. The
+    /// `Return(x)` variant indicates that the current flow should be terminated
+    /// with the result `x`.
+    fn swipe_start(
+        &mut self,
+        _ctx: &mut EventCtx,
+        _direction: SwipeDirection,
+    ) -> SwipableResult<T> {
+        SwipableResult::Ignored
     }
 
     /// Return true when transition animation is finished. SwipeFlow needs to
@@ -19,6 +25,28 @@ pub trait Swipable {
     }
 }
 
+pub enum SwipableResult<T> {
+    Ignored,
+    Animating,
+    Return(T),
+}
+
+impl<T> SwipableResult<T> {
+    pub fn map<U>(self, func: impl FnOnce(T) -> Option<U>) -> SwipableResult<U> {
+        match self {
+            SwipableResult::Ignored => SwipableResult::Ignored,
+            SwipableResult::Animating => SwipableResult::Animating,
+            SwipableResult::Return(x) => {
+                if let Some(res) = func(x) {
+                    SwipableResult::Return(res)
+                } else {
+                    SwipableResult::Ignored
+                }
+            }
+        }
+    }
+}
+
 /// Component::Msg for component parts of a flow. Converting results of
 /// different screens to a shared type makes things easier to work with.
 ///
diff --git a/core/embed/rust/src/ui/flow/mod.rs b/core/embed/rust/src/ui/flow/mod.rs
index 1de1888d6..91c69e4c5 100644
--- a/core/embed/rust/src/ui/flow/mod.rs
+++ b/core/embed/rust/src/ui/flow/mod.rs
@@ -3,7 +3,7 @@ pub mod page;
 mod store;
 mod swipe;
 
-pub use base::{FlowMsg, FlowState, Swipable};
+pub use base::{FlowMsg, FlowState, Swipable, SwipableResult};
 pub use page::{IgnoreSwipe, SwipePage};
 pub use store::{flow_store, FlowStore};
 pub use swipe::SwipeFlow;
diff --git a/core/embed/rust/src/ui/flow/page.rs b/core/embed/rust/src/ui/flow/page.rs
index 8f6f84fef..36395b50a 100644
--- a/core/embed/rust/src/ui/flow/page.rs
+++ b/core/embed/rust/src/ui/flow/page.rs
@@ -4,7 +4,7 @@ use crate::{
     ui::{
         animation::Animation,
         component::{Component, Event, EventCtx, Paginate, SwipeDirection},
-        flow::base::Swipable,
+        flow::base::{Swipable, SwipableResult},
         geometry::{Axis, Rect},
         shape::Renderer,
         util,
@@ -117,17 +117,27 @@ impl<T: Component + Paginate + Clone> Component for SwipePage<T> {
     }
 }
 
-impl<T: Component + Paginate + Clone> Swipable for SwipePage<T> {
-    fn swipe_start(&mut self, ctx: &mut EventCtx, direction: SwipeDirection) -> bool {
+impl<T: Component + Paginate + Clone> Swipable<T::Msg> for SwipePage<T> {
+    fn swipe_start(
+        &mut self,
+        ctx: &mut EventCtx,
+        direction: SwipeDirection,
+    ) -> SwipableResult<T::Msg> {
         match (self.axis, direction) {
             // Wrong direction
-            (Axis::Horizontal, SwipeDirection::Up | SwipeDirection::Down) => return false,
-            (Axis::Vertical, SwipeDirection::Left | SwipeDirection::Right) => return false,
+            (Axis::Horizontal, SwipeDirection::Up | SwipeDirection::Down) => {
+                return SwipableResult::Ignored
+            }
+            (Axis::Vertical, SwipeDirection::Left | SwipeDirection::Right) => {
+                return SwipableResult::Ignored
+            }
             // Begin
-            (_, SwipeDirection::Right | SwipeDirection::Down) if self.current == 0 => return false,
+            (_, SwipeDirection::Right | SwipeDirection::Down) if self.current == 0 => {
+                return SwipableResult::Ignored
+            }
             // End
             (_, SwipeDirection::Left | SwipeDirection::Up) if self.current + 1 >= self.pages => {
-                return false
+                return SwipableResult::Ignored;
             }
             _ => {}
         };
@@ -138,7 +148,7 @@ impl<T: Component + Paginate + Clone> Swipable for SwipePage<T> {
         if util::animation_disabled() {
             self.inner.change_page(self.current);
             ctx.request_paint();
-            return true;
+            return SwipableResult::Animating;
         }
         self.transition = Some(Transition {
             cloned: unwrap!(Gc::new(self.inner.clone())),
@@ -148,7 +158,7 @@ impl<T: Component + Paginate + Clone> Swipable for SwipePage<T> {
         self.inner.change_page(self.current);
         ctx.request_anim_frame();
         ctx.request_paint();
-        true
+        SwipableResult::Animating
     }
 
     fn swipe_finished(&self) -> bool {
@@ -195,7 +205,7 @@ impl<T: Component> Component for IgnoreSwipe<T> {
     }
 }
 
-impl<T> Swipable for IgnoreSwipe<T> {}
+impl<T: Component> Swipable<T::Msg> for IgnoreSwipe<T> {}
 
 #[cfg(feature = "ui_debug")]
 impl<T> crate::trace::Trace for IgnoreSwipe<T>
diff --git a/core/embed/rust/src/ui/flow/store.rs b/core/embed/rust/src/ui/flow/store.rs
index fef7e58a0..cf1bd2d60 100644
--- a/core/embed/rust/src/ui/flow/store.rs
+++ b/core/embed/rust/src/ui/flow/store.rs
@@ -29,10 +29,14 @@ pub trait FlowStore {
     fn trace(&self, i: usize, t: &mut dyn crate::trace::Tracer);
 
     /// Forward `Swipable` methods to i-th element.
-    fn map_swipable<T>(&mut self, i: usize, func: impl FnOnce(&mut dyn Swipable) -> T) -> T;
+    fn map_swipable<T>(
+        &mut self,
+        i: usize,
+        func: impl FnOnce(&mut dyn Swipable<FlowMsg>) -> T,
+    ) -> T;
 
     /// Add a Component to the end of a `FlowStore`.
-    fn add<E: Component<Msg = FlowMsg> + MaybeTrace + Swipable>(
+    fn add<E: Component<Msg = FlowMsg> + MaybeTrace + Swipable<FlowMsg>>(
         self,
         elem: E,
     ) -> Result<impl FlowStore, error::Error>
@@ -67,11 +71,15 @@ impl FlowStore for FlowEmpty {
         panic!()
     }
 
-    fn map_swipable<T>(&mut self, _i: usize, _func: impl FnOnce(&mut dyn Swipable) -> T) -> T {
+    fn map_swipable<T>(
+        &mut self,
+        _i: usize,
+        _func: impl FnOnce(&mut dyn Swipable<FlowMsg>) -> T,
+    ) -> T {
         panic!()
     }
 
-    fn add<E: Component<Msg = FlowMsg> + MaybeTrace + Swipable>(
+    fn add<E: Component<Msg = FlowMsg> + MaybeTrace + Swipable<FlowMsg>>(
         self,
         elem: E,
     ) -> Result<impl FlowStore, error::Error>
@@ -107,7 +115,7 @@ impl<E: Component<Msg = FlowMsg>, P> FlowComponent<E, P> {
 
 impl<E, P> FlowStore for FlowComponent<E, P>
 where
-    E: Component<Msg = FlowMsg> + MaybeTrace + Swipable,
+    E: Component<Msg = FlowMsg> + MaybeTrace + Swipable<FlowMsg>,
     P: FlowStore,
 {
     fn place(&mut self, bounds: Rect) -> Rect {
@@ -141,7 +149,11 @@ where
         }
     }
 
-    fn map_swipable<T>(&mut self, i: usize, func: impl FnOnce(&mut dyn Swipable) -> T) -> T {
+    fn map_swipable<T>(
+        &mut self,
+        i: usize,
+        func: impl FnOnce(&mut dyn Swipable<FlowMsg>) -> T,
+    ) -> T {
         if i == 0 {
             func(self.as_mut())
         } else {
@@ -149,7 +161,7 @@ where
         }
     }
 
-    fn add<F: Component<Msg = FlowMsg> + MaybeTrace + Swipable>(
+    fn add<F: Component<Msg = FlowMsg> + MaybeTrace + Swipable<FlowMsg>>(
         self,
         elem: F,
     ) -> Result<impl FlowStore, error::Error>
diff --git a/core/embed/rust/src/ui/flow/swipe.rs b/core/embed/rust/src/ui/flow/swipe.rs
index 7470e5e74..70bb21c1d 100644
--- a/core/embed/rust/src/ui/flow/swipe.rs
+++ b/core/embed/rust/src/ui/flow/swipe.rs
@@ -4,7 +4,7 @@ use crate::{
     ui::{
         animation::Animation,
         component::{Component, Event, EventCtx, Swipe, SwipeDirection},
-        flow::{base::Decision, FlowMsg, FlowState, FlowStore},
+        flow::{base::Decision, FlowMsg, FlowState, FlowStore, SwipableResult},
         geometry::Rect,
         shape::Renderer,
         util,
@@ -135,13 +135,13 @@ impl<Q: FlowState, S: FlowStore> SwipeFlow<Q, S> {
 
     fn handle_swipe_child(&mut self, ctx: &mut EventCtx, direction: SwipeDirection) -> Decision<Q> {
         let i = self.state.index();
-        if self
+        match self
             .store
             .map_swipable(i, |s| s.swipe_start(ctx, direction))
         {
-            Decision::Goto(self.state, direction)
-        } else {
-            Decision::Nothing
+            SwipableResult::Ignored => Decision::Nothing,
+            SwipableResult::Animating => Decision::Goto(self.state, direction),
+            SwipableResult::Return(x) => Decision::Return(x),
         }
     }
 
diff --git a/core/embed/rust/src/ui/model_mercury/component/dialog.rs b/core/embed/rust/src/ui/model_mercury/component/dialog.rs
index 55947f4ab..54fb32b15 100644
--- a/core/embed/rust/src/ui/model_mercury/component/dialog.rs
+++ b/core/embed/rust/src/ui/model_mercury/component/dialog.rs
@@ -230,4 +230,4 @@ where
 }
 
 #[cfg(feature = "micropython")]
-impl<U> crate::ui::flow::Swipable for IconDialog<U> {}
+impl<U: Component> crate::ui::flow::Swipable<U::Msg> for IconDialog<U> {}
diff --git a/core/embed/rust/src/ui/model_mercury/component/frame.rs b/core/embed/rust/src/ui/model_mercury/component/frame.rs
index c1035abc3..f558822d3 100644
--- a/core/embed/rust/src/ui/model_mercury/component/frame.rs
+++ b/core/embed/rust/src/ui/model_mercury/component/frame.rs
@@ -252,16 +252,17 @@ where
 }
 
 #[cfg(feature = "micropython")]
-impl<T> crate::ui::flow::Swipable for Frame<T>
+impl<T> crate::ui::flow::Swipable<FrameMsg<T::Msg>> for Frame<T>
 where
-    T: Component + crate::ui::flow::Swipable,
+    T: Component + crate::ui::flow::Swipable<T::Msg>,
 {
     fn swipe_start(
         &mut self,
         ctx: &mut EventCtx,
         direction: crate::ui::component::SwipeDirection,
-    ) -> bool {
+    ) -> crate::ui::flow::SwipableResult<FrameMsg<T::Msg>> {
         self.update_content(ctx, |ctx, inner| inner.swipe_start(ctx, direction))
+            .map(|x| Some(FrameMsg::Content(x)))
     }
 
     fn swipe_finished(&self) -> bool {
diff --git a/core/embed/rust/src/ui/model_mercury/component/prompt_screen.rs b/core/embed/rust/src/ui/model_mercury/component/prompt_screen.rs
index 1f9a99be3..432603726 100644
--- a/core/embed/rust/src/ui/model_mercury/component/prompt_screen.rs
+++ b/core/embed/rust/src/ui/model_mercury/component/prompt_screen.rs
@@ -134,7 +134,7 @@ impl Component for PromptScreen {
 }
 
 #[cfg(feature = "micropython")]
-impl crate::ui::flow::Swipable for PromptScreen {}
+impl crate::ui::flow::Swipable<()> for PromptScreen {}
 
 #[cfg(feature = "ui_debug")]
 impl crate::trace::Trace for PromptScreen {
diff --git a/core/embed/rust/src/ui/model_mercury/component/share_words.rs b/core/embed/rust/src/ui/model_mercury/component/share_words.rs
index e8aba1a7d..b21274b5f 100644
--- a/core/embed/rust/src/ui/model_mercury/component/share_words.rs
+++ b/core/embed/rust/src/ui/model_mercury/component/share_words.rs
@@ -6,7 +6,7 @@ use crate::{
     ui::{
         animation::Animation,
         component::{Component, Event, EventCtx, Never, SwipeDirection},
-        flow::Swipable,
+        flow::{Swipable, SwipableResult},
         geometry::{Alignment, Alignment2D, Insets, Offset, Rect},
         model_mercury::component::Footer,
         shape,
@@ -160,8 +160,12 @@ impl<'a> Component for ShareWords<'a> {
     fn bounds(&self, _sink: &mut dyn FnMut(Rect)) {}
 }
 
-impl<'a> Swipable for ShareWords<'a> {
-    fn swipe_start(&mut self, ctx: &mut EventCtx, direction: SwipeDirection) -> bool {
+impl<'a> Swipable<Never> for ShareWords<'a> {
+    fn swipe_start(
+        &mut self,
+        ctx: &mut EventCtx,
+        direction: SwipeDirection,
+    ) -> SwipableResult<Never> {
         match direction {
             SwipeDirection::Up if !self.is_final_page() => {
                 self.prev_index = self.page_index;
@@ -171,11 +175,11 @@ impl<'a> Swipable for ShareWords<'a> {
                 self.prev_index = self.page_index;
                 self.page_index = self.page_index.saturating_sub(1);
             }
-            _ => return false,
+            _ => return SwipableResult::Ignored,
         };
         if util::animation_disabled() {
             ctx.request_paint();
-            return true;
+            return SwipableResult::Animating;
         }
         self.animation = Some(Animation::new(
             0.0f32,
@@ -185,7 +189,7 @@ impl<'a> Swipable for ShareWords<'a> {
         ));
         ctx.request_anim_frame();
         ctx.request_paint();
-        true
+        SwipableResult::Animating
     }
 
     fn swipe_finished(&self) -> bool {
diff --git a/core/embed/rust/src/ui/model_mercury/component/vertical_menu.rs b/core/embed/rust/src/ui/model_mercury/component/vertical_menu.rs
index ac2b605f6..6773be7a8 100644
--- a/core/embed/rust/src/ui/model_mercury/component/vertical_menu.rs
+++ b/core/embed/rust/src/ui/model_mercury/component/vertical_menu.rs
@@ -151,4 +151,4 @@ impl crate::trace::Trace for VerticalMenu {
 }
 
 #[cfg(feature = "micropython")]
-impl crate::ui::flow::Swipable for VerticalMenu {}
+impl crate::ui::flow::Swipable<VerticalMenuChoiceMsg> for VerticalMenu {}
diff --git a/core/embed/rust/src/ui/model_mercury/flow/util.rs b/core/embed/rust/src/ui/model_mercury/flow/util.rs
index 41c41fa7b..7c5efb6e2 100644
--- a/core/embed/rust/src/ui/model_mercury/flow/util.rs
+++ b/core/embed/rust/src/ui/model_mercury/flow/util.rs
@@ -89,7 +89,7 @@ impl ConfirmBlobParams {
 
     pub fn into_layout(
         self,
-    ) -> Result<impl Component<Msg = FlowMsg> + Swipable + MaybeTrace, Error> {
+    ) -> Result<impl Component<Msg = FlowMsg> + Swipable<FlowMsg> + MaybeTrace, Error> {
         let paragraphs = ConfirmBlob {
             description: self.description.unwrap_or("".into()),
             extra: self.extra.unwrap_or("".into()),
@@ -187,7 +187,7 @@ impl ShowInfoParams {
 
     pub fn into_layout(
         self,
-    ) -> Result<impl Component<Msg = FlowMsg> + Swipable + MaybeTrace, Error> {
+    ) -> Result<impl Component<Msg = FlowMsg> + Swipable<FlowMsg> + MaybeTrace, Error> {
         let mut paragraphs = ParagraphVecShort::new();
         let mut first: bool = true;
         for item in self.items {