From 684ff50ef4931aeb8cfb15c5a2d62e55520f04a5 Mon Sep 17 00:00:00 2001
From: jacqueline <me@jacqueline.id.au>
Date: Thu, 21 Mar 2024 10:50:23 +1100
Subject: [PATCH] Add support for screens declaring that they can't be popped

Needed as prep for usb msc support; you really shouldn't leave the MSC
settings screen until you've disabled usb msc.
---
 src/lua/lua_screen.cpp           | 14 +++++++++-----
 src/ui/include/screen.hpp        |  2 ++
 src/ui/include/screen_lua.hpp    |  2 ++
 src/ui/include/screen_splash.hpp |  2 ++
 src/ui/screen_lua.cpp            | 19 +++++++++++++++++++
 src/ui/ui_fsm.cpp                |  3 +++
 6 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/lua/lua_screen.cpp b/src/lua/lua_screen.cpp
index 27843bc7..f17f6b1a 100644
--- a/src/lua/lua_screen.cpp
+++ b/src/lua/lua_screen.cpp
@@ -50,11 +50,15 @@ static auto screen_noop(lua_State* state) -> int {
   return 0;
 }
 
-static const struct luaL_Reg kScreenFuncs[] = {{"new", screen_new},
-                                               {"createUi", screen_noop},
-                                               {"onShown", screen_noop},
-                                               {"onHidden", screen_noop},
-                                               {NULL, NULL}};
+static auto screen_true(lua_State* state) -> int {
+  lua_pushboolean(state, true);
+  return 1;
+}
+
+static const struct luaL_Reg kScreenFuncs[] = {
+    {"new", screen_new},      {"createUi", screen_noop},
+    {"onShown", screen_noop}, {"onHidden", screen_noop},
+    {"canPop", screen_true},  {NULL, NULL}};
 
 static auto lua_screen(lua_State* state) -> int {
   luaL_newlib(state, kScreenFuncs);
diff --git a/src/ui/include/screen.hpp b/src/ui/include/screen.hpp
index 4241c712..40284fda 100644
--- a/src/ui/include/screen.hpp
+++ b/src/ui/include/screen.hpp
@@ -43,6 +43,8 @@ class Screen {
     return group_;
   }
 
+  virtual auto canPop() -> bool = 0;
+
  protected:
   lv_obj_t* const root_;
   lv_obj_t* content_;
diff --git a/src/ui/include/screen_lua.hpp b/src/ui/include/screen_lua.hpp
index 0ed3a508..41d97a1e 100644
--- a/src/ui/include/screen_lua.hpp
+++ b/src/ui/include/screen_lua.hpp
@@ -21,6 +21,8 @@ class Lua : public Screen {
   auto onShown() -> void override;
   auto onHidden() -> void override;
 
+  auto canPop() -> bool override;
+
   auto SetObjRef(lua_State*) -> void;
 
  private:
diff --git a/src/ui/include/screen_splash.hpp b/src/ui/include/screen_splash.hpp
index 1ee7dd89..6e746345 100644
--- a/src/ui/include/screen_splash.hpp
+++ b/src/ui/include/screen_splash.hpp
@@ -20,6 +20,8 @@ class Splash : public Screen {
   Splash();
   ~Splash();
 
+  auto canPop() -> bool override { return false; }
+
  private:
   lv_obj_t* container_;
   lv_obj_t* label_;
diff --git a/src/ui/screen_lua.cpp b/src/ui/screen_lua.cpp
index d6c7a26f..55eef119 100644
--- a/src/ui/screen_lua.cpp
+++ b/src/ui/screen_lua.cpp
@@ -58,6 +58,25 @@ auto Lua::onHidden() -> void {
   lua_pop(s_, 1);
 }
 
+auto Lua::canPop() -> bool {
+  if (!s_ || !obj_ref_) {
+    return true;
+  }
+  lua_rawgeti(s_, LUA_REGISTRYINDEX, *obj_ref_);
+  lua_pushliteral(s_, "canPop");
+
+  if (lua_gettable(s_, -2) == LUA_TFUNCTION) {
+    // If we got a callback instead of a value, then invoke it to turn it into
+    // value.
+    lua_pushvalue(s_, -2);
+    lua::CallProtected(s_, 1, 1);
+  }
+  bool ret = lua_toboolean(s_, -1);
+
+  lua_pop(s_, 2);
+  return ret;
+}
+
 auto Lua::SetObjRef(lua_State* s) -> void {
   assert(s_ == nullptr);
   s_ = s;
diff --git a/src/ui/ui_fsm.cpp b/src/ui/ui_fsm.cpp
index 0ed012a0..c11f66a7 100644
--- a/src/ui/ui_fsm.cpp
+++ b/src/ui/ui_fsm.cpp
@@ -610,6 +610,9 @@ auto Lua::QueuePrevious(lua_State*) -> int {
 }
 
 auto Lua::PopLuaScreen(lua_State* s) -> int {
+  if (!sCurrentScreen->canPop()) {
+    return 0;
+  }
   PopScreen();
   luavgl_set_root(s, sCurrentScreen->content());
   lv_group_set_default(sCurrentScreen->group());