From c5c13217886d24cc3c0e61e2118ab88f8ac4566c Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Sat, 16 Mar 2024 20:42:54 +0100
Subject: [PATCH] feat(core/rust): some useful traits for TString

string equality comparison and SkipPrefix
---
 core/embed/rust/src/strutil.rs                | 42 +++++++++++++++++--
 .../generated/translated_string.rs            |  2 +-
 .../generated/translated_string.rs.mako       |  2 +-
 .../src/translations/translated_string.rs     |  2 +-
 4 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/core/embed/rust/src/strutil.rs b/core/embed/rust/src/strutil.rs
index 576115b40..4eb200720 100644
--- a/core/embed/rust/src/strutil.rs
+++ b/core/embed/rust/src/strutil.rs
@@ -93,13 +93,20 @@ pub enum TString<'a> {
     #[cfg(feature = "micropython")]
     Allocated(StrBuffer),
     #[cfg(feature = "translations")]
-    Translation(TR),
+    Translation {
+        tr: TR,
+        offset: u16,
+    },
     Str(&'a str),
 }
 
 impl TString<'_> {
+    pub fn len(&self) -> usize {
+        self.map(|s| s.len())
+    }
+
     pub fn is_empty(&self) -> bool {
-        self.map(|s| s.is_empty())
+        self.len() == 0
     }
 
     pub fn map<F, T>(&self, fun: F) -> T
@@ -111,7 +118,7 @@ impl TString<'_> {
             #[cfg(feature = "micropython")]
             Self::Allocated(buf) => fun(buf.as_ref()),
             #[cfg(feature = "translations")]
-            Self::Translation(tr) => tr.map_translated(fun),
+            Self::Translation { tr, offset } => tr.map_translated(|s| fun(&s[*offset as usize..])),
             Self::Str(s) => fun(s),
         }
     }
@@ -120,7 +127,7 @@ impl TString<'_> {
 impl TString<'static> {
     #[cfg(feature = "translations")]
     pub const fn from_translation(tr: TR) -> Self {
-        Self::Translation(tr)
+        Self::Translation { tr, offset: 0 }
     }
 
     #[cfg(feature = "micropython")]
@@ -176,3 +183,30 @@ impl<'a> TryFrom<TString<'a>> for Obj {
         s.map(|t| t.try_into())
     }
 }
+
+impl<'a, 'b> PartialEq<TString<'a>> for TString<'b> {
+    fn eq(&self, other: &TString<'a>) -> bool {
+        self.map(|s| other.map(|o| s == o))
+    }
+}
+
+impl Eq for TString<'_> {}
+
+impl SkipPrefix for TString<'_> {
+    fn skip_prefix(&self, skip_bytes: usize) -> Self {
+        self.map(|s| {
+            assert!(skip_bytes <= s.len());
+            assert!(s.is_char_boundary(skip_bytes));
+        });
+        match self {
+            #[cfg(feature = "micropython")]
+            Self::Allocated(s) => Self::Allocated(s.skip_prefix(skip_bytes)),
+            #[cfg(feature = "translations")]
+            Self::Translation { tr, offset } => Self::Translation {
+                tr: *tr,
+                offset: offset + skip_bytes as u16,
+            },
+            Self::Str(s) => Self::Str(&s[skip_bytes..]),
+        }
+    }
+}
diff --git a/core/embed/rust/src/translations/generated/translated_string.rs b/core/embed/rust/src/translations/generated/translated_string.rs
index d6af2d020..974dae100 100644
--- a/core/embed/rust/src/translations/generated/translated_string.rs
+++ b/core/embed/rust/src/translations/generated/translated_string.rs
@@ -6,7 +6,7 @@
 #[cfg(feature = "micropython")]
 use crate::micropython::qstr::Qstr;
 
-#[derive(Debug, Copy, Clone, FromPrimitive)]
+#[derive(Debug, Copy, Clone, FromPrimitive, PartialEq, Eq)]
 #[repr(u16)]
 #[allow(non_camel_case_types)]
 pub enum TranslatedString {
diff --git a/core/embed/rust/src/translations/generated/translated_string.rs.mako b/core/embed/rust/src/translations/generated/translated_string.rs.mako
index 856b31793..74acf6757 100644
--- a/core/embed/rust/src/translations/generated/translated_string.rs.mako
+++ b/core/embed/rust/src/translations/generated/translated_string.rs.mako
@@ -35,7 +35,7 @@ en_data = json.loads(en_file.read_text())["translations"]
 #[cfg(feature = "micropython")]
 use crate::micropython::qstr::Qstr;
 
-#[derive(Debug, Copy, Clone, FromPrimitive)]
+#[derive(Debug, Copy, Clone, FromPrimitive, PartialEq, Eq)]
 #[repr(u16)]
 #[allow(non_camel_case_types)]
 pub enum TranslatedString {
diff --git a/core/embed/rust/src/translations/translated_string.rs b/core/embed/rust/src/translations/translated_string.rs
index 3127868b1..c7e90e5db 100644
--- a/core/embed/rust/src/translations/translated_string.rs
+++ b/core/embed/rust/src/translations/translated_string.rs
@@ -23,7 +23,7 @@ impl TranslatedString {
     }
 
     pub const fn as_tstring(self) -> TString<'static> {
-        TString::Translation(self)
+        TString::from_translation(self)
     }
 }