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!

9 files changed, 97 insertions, 58 deletions

diff --git a/lua/backstack.lua b/lua/backstack.lua
deleted file mode 100644
index c54fbac4..00000000
--- a/lua/backstack.lua
+++ /dev/null
@@ -1,37 +0,0 @@
-local lvgl = require("lvgl")
-
-local backstack = {
-  root = lvgl.Object(nil, {
-    w = lvgl.HOR_RES(),
-    h = lvgl.VER_RES(),
-  }),
-  stack = {},
-}
-
-function backstack:Top()
-  return self.stack[#self.stack]
-end
-
-function backstack:SetTopParent(parent)
-  local top = self:Top()
-  if top and top.root then
-    top.root:set_parent(parent)
-  end
-end
-
-function backstack:Push(screen)
-  self:SetTopParent(nil)
-  table.insert(self.stack, screen)
-  self:SetTopParent(self.root)
-end
-
-function backstack:Pop(num)
-  num = num or 1
-  for _ = 1, num do
-    local removed = table.remove(self.stack)
-    removed.root:delete()
-  end
-  self:SetTopParent(self.root)
-end
-
-return backstack
diff --git a/lua/main.lua b/lua/main.lua
index ce9596af..106f11c7 100644
--- a/lua/main.lua
+++ b/lua/main.lua
@@ -1,3 +1,4 @@
 local backstack = require("backstack")
-local main_menu = require("main_menu"):Create(backstack.root)
-backstack:Push(main_menu)
+local main_menu = require("main_menu")
+
+backstack.push(main_menu)
diff --git a/lua/main_menu.lua b/lua/main_menu.lua
index f0be33de..3b88111d 100644
--- a/lua/main_menu.lua
+++ b/lua/main_menu.lua
@@ -3,11 +3,9 @@ local widgets = require("widgets")
 local legacy_ui = require("legacy_ui")
 local database = require("database")
 
-local main_menu = {}
-
-function main_menu:Create(parent)
+return function()
   local menu = {}
-  menu.root = lvgl.Object(parent, {
+  menu.root = lvgl.Object(nil, {
     flex = {
       flex_direction = "column",
       flex_wrap = "wrap",
@@ -46,5 +44,3 @@ function main_menu:Create(parent)
 
   return menu
 end
-
-return main_menu
diff --git a/src/lua/property.cpp b/src/lua/property.cpp
index 353e01ae..f0383dd8 100644
--- a/src/lua/property.cpp
+++ b/src/lua/property.cpp
@@ -5,13 +5,10 @@
  */
 
 #include "property.hpp"
-#include <sys/_stdint.h>
 
 #include <memory>
 #include <string>
 
-#include "lauxlib.h"
-#include "lua.h"
 #include "lua.hpp"
 #include "lvgl.h"
 #include "service_locator.hpp"
@@ -19,7 +16,7 @@
 namespace lua {
 
 static const char kPropertyMetatable[] = "property";
-static const char kFunctionMetatable[] = "function";
+static const char kFunctionMetatable[] = "c_func";
 static const char kBindingsTable[] = "bindings";
 static const char kBinderKey[] = "binder";
 
@@ -63,7 +60,7 @@ static auto property_bind(lua_State* state) -> int {
 
   p->AddLuaBinding(state, ref);
 
-  // Pop the bindings table, leaving one of the copiesw of the callback fn at
+  // Pop the bindings table, leaving one of the copies of the callback fn at
   // the top of the stack.
   lua_pop(state, 1);
 
@@ -84,6 +81,11 @@ static auto generic_function_cb(lua_State* state) -> int {
   size_t* index =
       reinterpret_cast<size_t*>(luaL_checkudata(state, 1, kFunctionMetatable));
   const LuaFunction& fn = binder->GetFunction(*index);
+
+  // Ensure the C++ function is called with a clean stack; we don't want it to
+  // see the index we just used.
+  lua_remove(state, 1);
+
   return std::invoke(fn, state);
 }
 
diff --git a/src/lua/stubs/backstack.lua b/src/lua/stubs/backstack.lua
new file mode 100644
index 00000000..d4807d37
--- /dev/null
+++ b/src/lua/stubs/backstack.lua
@@ -0,0 +1,13 @@
+--- Module for adding and removing screens from the system's backstack.
+-- @module backstack
+
+local backstack = {}
+
+--- Pushes a new screen onto the backstack.
+-- @tparam function constructor Called to create the UI for the new screen. A new default root object and group will be set before calling this function. The function provided should return a table holding any bindings used by this screen; the returned value is retained so long as this screen is present in the backstack.
+function backstack.push(constructor) end
+
+--- Removes the currently active screen, and instead shows the screen underneath it on the backstack. Does nothing if this is the only existing screen.
+function backstack.pop() end
+
+return backstack
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) {