Studio: Improve UI state saving and restoration

- Load settings prior to first frame so dockspace layout restores reliably.
- Save settings on exit in App::~App().
- Persist and restore open/closed state of collapsible sections.

PiperOrigin-RevId: 966657100
Change-Id: I64fb5d02912bb36d096d7ac26ce8e0de03e6617b

src/experimental/studio/app.cc

@@ -34,6 +34,7 @@
 #include <vector>
 
 #include <imgui.h>
+#include <imgui_internal.h>
 #include <implot.h>
 #include <mujoco/mujoco.h>
 #include "experimental/platform/hal/classic_renderer.h"
@@ -109,7 +110,13 @@ App::App(Config config)
       [this](const mjModel* m, mjData* d) { PostStep(m, d); });
 }
 
-App::~App() { mjv_freeScene(&plugin_scene_); }
+App::~App() {
+  // ImGui's autosave is on a timer and covers only state it tracks, so recent
+  // changes and plugin visibility would be lost. window_ owns the ImGui context
+  // and outlives this body.
+  SaveSettings();
+  mjv_freeScene(&plugin_scene_);
+}
 
 void App::SwitchGraphicsMode(int width, int height,
                              platform::GraphicsMode mode) {
@@ -428,6 +435,14 @@ void App::LoadHistory(int offset) {
 }
 
 bool App::Update() {
+  // Must precede the first NewFrame: the dockspace is built lazily on the frame
+  // that finds no root node, so the saved nodes have to be there already or the
+  // default layout wins.
+  if (tmp_.first_frame) {
+    LoadSettings();
+    tmp_.first_frame = false;
+  }
+
   const platform::Window::Status status = window_->NewFrame();
 
   HandleWindowEvents();
@@ -921,10 +936,41 @@ void App::LoadSettings() {
           plugin->active = std::stoi(it->second) != 0;
         }
       });
+
+      // Applied later: the owning windows do not exist until first submitted.
+      window_state_storage_ =
+          platform::ReadIniSection(settings, "[Studio][WindowStateStorage]");
     }
   }
 }
 
+// Key for one entry in state storage; the window name locates its owner on load.
+static std::string WindowStateStorageKey(const char* window_name, ImGuiID id) {
+  char id_str[16];
+  std::snprintf(id_str, sizeof(id_str), "%08X", id);
+  return std::string(window_name) + "/" + id_str;
+}
+
+void App::ApplyWindowStateStorage() {
+  for (auto it = window_state_storage_.begin(); it != window_state_storage_.end();) {
+    const std::string::size_type sep = it->first.rfind('/');
+    if (sep == std::string::npos) {
+      it = window_state_storage_.erase(it);
+      continue;
+    }
+    // Child windows are "<parent>/<child>_<id>": id follows the last '/'.
+    ImGuiWindow* window =
+        ImGui::FindWindowByName(it->first.substr(0, sep).c_str());
+    if (window == nullptr) {
+      ++it;  // Window not submitted yet; try again next frame.
+      continue;
+    }
+    const ImGuiID id = std::strtoul(it->first.c_str() + sep + 1, nullptr, 16);
+    window->StateStorage.SetInt(id, std::stoi(it->second) != 0);
+    it = window_state_storage_.erase(it);
+  }
+}
+
 void App::SaveSettings() {
   if (!ini_path_.empty()) {
     std::string settings = ImGui::SaveIniSettingsToMemory();
@@ -940,6 +986,21 @@ void App::SaveSettings() {
     });
     platform::AppendIniSection(settings, "[Studio][Plugins]", plugin_names);
 
+    // Pending entries first, so state for windows never opened this session is
+    // not dropped, then let the live windows override. Merge into a local copy:
+    // window_state_storage_ must stay pending-only, or a stale value could re-apply
+    // over a newer toggle.
+    platform::KeyValues window_state_storage = window_state_storage_;
+    ImGuiContext& g = *ImGui::GetCurrentContext();
+    for (ImGuiWindow* window : g.Windows) {
+      for (const ImGuiStoragePair& pair : window->StateStorage.Data) {
+        window_state_storage[WindowStateStorageKey(window->Name, pair.key)] =
+            std::to_string(pair.val_i);
+      }
+    }
+    platform::AppendIniSection(settings, "[Studio][WindowStateStorage]",
+                               window_state_storage);
+
     platform::SaveText(settings, ini_path_);
   }
 }
@@ -1177,11 +1238,11 @@ void App::BuildGui() {
     }
   });
 
+  // This frame's windows are submitted, so pending state storage can be restored.
+  // Must precede the save below, which would otherwise write the defaults.
+  ApplyWindowStateStorage();
+
   ImGuiIO& io = ImGui::GetIO();
-  if (tmp_.first_frame) {
-    LoadSettings();
-    tmp_.first_frame = false;
-  }
   if (io.WantSaveIniSettings) {
     SaveSettings();
     io.WantSaveIniSettings = false;

src/experimental/studio/app.h

@@ -35,6 +35,7 @@
 #include "experimental/platform/sim/step_control.h"
 #include "experimental/platform/ux/gui.h"
 #include "experimental/platform/ux/gui_spec.h"
+#include "experimental/platform/ux/imgui_widgets.h"
 #include "experimental/platform/ux/interaction.h"
 #include "experimental/platform/ux/picture_gui.h"
 #include "experimental/platform/ux/spec_editor.h"
@@ -208,6 +209,7 @@ class App {
 
   void LoadSettings();
   void SaveSettings();
+  void ApplyWindowStateStorage();
 
   void LoadHistory(int offset);
 
@@ -240,6 +242,10 @@ class App {
 
   std::string app_title_;
   std::string ini_path_;
+  // Window state storage (e.g. collapsing header open/closed state), keyed
+  // "<window name>/<id>", which ImGui does not serialize. Entries stay
+  // pending until their window is first created.
+  platform::KeyValues window_state_storage_;
   std::string model_name_;  // Used if model_kind_ is kModelFromBuffer.
   std::string model_path_;
   std::string load_error_;