From b3b3e0efa4f74f24ed3242acd8a087bdd17005ae Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Wed, 27 Apr 2022 11:57:45 +0200
Subject: [PATCH] tests(core/rust): ability to set up MicroPython env for
 testing

warning: super hacky!
[no changelog]
---
 core/embed/rust/build.rs                    |  6 ++++
 core/embed/rust/src/micropython/mod.rs      |  3 ++
 core/embed/rust/src/micropython/testutil.rs | 33 +++++++++++++++++++++
 3 files changed, 42 insertions(+)
 create mode 100644 core/embed/rust/src/micropython/testutil.rs

diff --git a/core/embed/rust/build.rs b/core/embed/rust/build.rs
index 9c2c3add1..0fc4e0416 100644
--- a/core/embed/rust/build.rs
+++ b/core/embed/rust/build.rs
@@ -221,6 +221,12 @@ fn link_core_objects() {
         let obj = obj.unwrap();
         cc.object(obj);
     }
+
+    // Add frozen modules, if present.
+    for obj in glob::glob(&format!("{}/*.o", build_path)).unwrap() {
+        cc.object(obj.unwrap());
+    }
+
     // Compile all the objects into a static library and link it in automatically.
     cc.compile("core_lib");
 
diff --git a/core/embed/rust/src/micropython/mod.rs b/core/embed/rust/src/micropython/mod.rs
index dff3f9fa3..4849d57f5 100644
--- a/core/embed/rust/src/micropython/mod.rs
+++ b/core/embed/rust/src/micropython/mod.rs
@@ -16,3 +16,6 @@ pub mod qstr;
 pub mod runtime;
 pub mod time;
 pub mod typ;
+
+#[cfg(test)]
+pub mod testutil;
diff --git a/core/embed/rust/src/micropython/testutil.rs b/core/embed/rust/src/micropython/testutil.rs
new file mode 100644
index 000000000..333eee468
--- /dev/null
+++ b/core/embed/rust/src/micropython/testutil.rs
@@ -0,0 +1,33 @@
+extern "C" {
+    fn gc_init(start: *mut cty::c_void, end: *mut cty::c_void);
+    fn mp_stack_ctrl_init();
+    fn mp_stack_set_limit(limit: usize);
+    fn mp_init();
+}
+
+static mut HEAP: [u8; 20 * 1024 * 1024] = [0; 20 * 1024 * 1024];
+static mut MPY_INITIALIZED: bool = false;
+
+/// Initialize the MicroPython environment.
+///
+/// This is very hacky, in no way safe, and should not be used in production.
+/// The stack is configured on a best-effort basis and depending on from where
+/// this is called, you might get errorneous "recursion exceeded" problems.
+///
+/// This should only be called at start of your test function.
+#[cfg(test)]
+pub unsafe fn mpy_init() {
+    unsafe {
+        if MPY_INITIALIZED {
+            return;
+        }
+        mp_stack_ctrl_init();
+        mp_stack_set_limit(6000000);
+        gc_init(
+            HEAP.as_mut_ptr().cast(),
+            HEAP.as_mut_ptr().add(HEAP.len()).cast(),
+        );
+        mp_init();
+        MPY_INITIALIZED = true;
+    }
+}