Use C functions for the backstack, instead of a lua module

Working with the default group and root kinda sucks if you have to do it
from lua!

4 files changed, 72 insertions, 8 deletions

diff --git a/src/ui/include/screen_lua.hpp b/src/ui/include/screen_lua.hpp
index df83ea8b..ee9f6813 100644
--- a/src/ui/include/screen_lua.hpp
+++ b/src/ui/include/screen_lua.hpp
@@ -15,7 +15,14 @@ namespace screens {
 
 class Lua : public Screen {
  public:
-  explicit Lua(lua_State* l);
+  Lua();
+  ~Lua();
+
+  auto SetObjRef(lua_State*) -> void;
+
+ private:
+  lua_State* s_;
+  std::optional<int> obj_ref_;
 };
 
 }  // namespace screens
diff --git a/src/ui/include/ui_fsm.hpp b/src/ui/include/ui_fsm.hpp
index 39fae4b0..d3ea7eac 100644
--- a/src/ui/include/ui_fsm.hpp
+++ b/src/ui/include/ui_fsm.hpp
@@ -135,6 +135,9 @@ class Lua : public UiState {
   using UiState::react;
 
  private:
+  auto PushLuaScreen(lua_State*) -> int;
+  auto PopLuaScreen(lua_State*) -> int;
+
   std::shared_ptr<lua::Property> battery_pct_;
   std::shared_ptr<lua::Property> battery_mv_;
   std::shared_ptr<lua::Property> battery_charging_;
diff --git a/src/ui/screen_lua.cpp b/src/ui/screen_lua.cpp
index cee5186e..ae49ffd5 100644
--- a/src/ui/screen_lua.cpp
+++ b/src/ui/screen_lua.cpp
@@ -6,13 +6,26 @@
 
 #include "screen_lua.hpp"
 
+#include "lauxlib.h"
+#include "lua.h"
 #include "lua.hpp"
 #include "luavgl.h"
 
 namespace ui {
 namespace screens {
 
-Lua::Lua(lua_State* l) {
+Lua::Lua() : s_(nullptr), obj_ref_() {}
+
+Lua::~Lua() {
+  if (s_ && obj_ref_) {
+    luaL_unref(s_, LUA_REGISTRYINDEX, *obj_ref_);
+  }
+}
+
+auto Lua::SetObjRef(lua_State* s) -> void {
+  assert(s_ == nullptr);
+  s_ = s;
+  obj_ref_ = luaL_ref(s, LUA_REGISTRYINDEX);
 }
 
 }  // namespace screens
diff --git a/src/ui/ui_fsm.cpp b/src/ui/ui_fsm.cpp
index 9ecc9b7c..d5de53f0 100644
--- a/src/ui/ui_fsm.cpp
+++ b/src/ui/ui_fsm.cpp
@@ -8,6 +8,9 @@
 
 #include <memory>
 
+#include "lua.h"
+#include "lua.hpp"
+
 #include "audio_fsm.hpp"
 #include "battery.hpp"
 #include "core/lv_group.h"
@@ -17,7 +20,6 @@
 #include "esp_heap_caps.h"
 #include "haptics.hpp"
 #include "lauxlib.h"
-#include "lua.hpp"
 #include "lua_thread.hpp"
 #include "luavgl.h"
 #include "misc/lv_gc.h"
@@ -181,9 +183,6 @@ void Splash::react(const system_fsm::StorageMounted&) {
 
 void Lua::entry() {
   if (!sLua) {
-    sCurrentScreen.reset(new Screen());
-    lv_group_set_default(sCurrentScreen->group());
-
     auto bat =
         sServices->battery().State().value_or(battery::Battery::BatteryState{});
     battery_pct_ =
@@ -213,13 +212,55 @@ void Lua::entry() {
                                          {"playing", playback_playing_},
                                          {"track", playback_track_},
                                      });
+    sLua->bridge().AddPropertyModule(
+        "backstack",
+        {
+            {"push", [&](lua_State* s) { return PushLuaScreen(s); }},
+            {"pop", [&](lua_State* s) { return PopLuaScreen(s); }},
+        });
 
     sLua->RunScript("/lua/main.lua");
-
-    lv_group_set_default(NULL);
   }
 }
 
+auto Lua::PushLuaScreen(lua_State* s) -> int {
+  // Ensure the arg looks right before continuing.
+  luaL_checktype(s, 1, LUA_TFUNCTION);
+
+  // First, create a new plain old Screen object. We will use its root and
+  // group for the Lua screen.
+  auto new_screen = std::make_shared<screens::Lua>();
+
+  // Tell lvgl about the new roots.
+  luavgl_set_root(s, new_screen->root());
+  lv_group_set_default(new_screen->group());
+
+  // Call the constructor for this screen.
+  lua_settop(s, 1);  // Make sure the function is actually at top of stack
+  // FIXME: This should ideally be lua_pcall, for safety.
+  lua_call(s, 0, 1);
+
+  // Store the reference for the table the constructor returned.
+  new_screen->SetObjRef(s);
+
+  // Ensure that we don't pollute the new screen's group. We leave the luavgl
+  // root alone.
+  // FIXME: maybe we should set the luavgl root to some catch-all that throws
+  // when anything is added to it? this may help catch bugs!
+  lv_group_set_default(NULL);
+
+  // Finally, push the now-initialised screen as if it were a regular C++
+  // screen.
+  PushScreen(new_screen);
+
+  return 0;
+}
+
+auto Lua::PopLuaScreen(lua_State*) -> int {
+  PopScreen();
+  return 0;
+}
+
 void Lua::exit() {}
 
 void Lua::react(const internal::IndexSelected& ev) {