Improve imgui texture management.

Keep track of all textures and recreate them as needed. This makes imgui
textures robust to render context reinitialization.

PiperOrigin-RevId: 840228061
Change-Id: Ia4ac35ff93a41ac74020f334adc438df4b030f7c

src/experimental/filament/filament/gui_view.cc

@@ -58,6 +58,9 @@ GuiView::~GuiView() {
   for (auto& instance : instances_) {
     engine_->destroy(instance);
   }
+  for (auto& texture : textures_) {
+    engine_->destroy(texture.second);
+  }
   engine_->destroyCameraComponent(camera_->getEntity());
   engine_->destroy(view_);
   engine_->destroy(scene_);
@@ -80,47 +83,47 @@ void GuiView::ResetRenderable() {
   buffers_.clear();
 }
 
-void GuiView::ProcessTexture(ImTextureData* data) {
+void GuiView::CreateTexture(ImTextureData* data) {
   filament::Engine* engine = object_mgr_->GetEngine();
+  if (data->Format != ImTextureFormat_RGBA32) {
+    mju_error("Unsupported texture format.");
+  }
 
-  if (data->Status == ImTextureStatus_OK) {
+  filament::Texture* texture =
-    return;
+      filament::Texture::Builder()
-  } else if (data->Status == ImTextureStatus_WantCreate) {
+          .width(data->Width)
-    if (data->Format != ImTextureFormat_RGBA32) {
+          .height(data->Height)
-      mju_error("Unsupported texture format.");
+          .levels(1)
-    }
+          .format(filament::Texture::InternalFormat::RGBA8)
+          .sampler(filament::Texture::Sampler::SAMPLER_2D)
+          .build(*engine);
 
-    filament::Texture* texture =
+  const uintptr_t tex_id = reinterpret_cast<uintptr_t>(texture);
-        filament::Texture::Builder()
+  textures_[tex_id] = texture;
-            .width(data->Width)
+  data->SetTexID((ImTextureID)tex_id);
-            .height(data->Height)
+  UpdateTexture(data);
-            .levels(1)
+}
-            .format(filament::Texture::InternalFormat::RGBA8)
-            .sampler(filament::Texture::Sampler::SAMPLER_2D)
-            .build(*engine);
 
-    const int size = data->Width * data->Height * 4;
+void GuiView::UpdateTexture(ImTextureData* data) {
-    filament::Texture::PixelBufferDescriptor pb(data->GetPixels(), size,
+  const int size = data->Width * data->Height * 4;
-                                                filament::Texture::Format::RGBA,
+  filament::Texture::PixelBufferDescriptor pb(data->GetPixels(), size,
-                                                filament::Texture::Type::UBYTE);
+                                              filament::Texture::Format::RGBA,
-    texture->setImage(*engine, 0, std::move(pb));
+                                              filament::Texture::Type::UBYTE);
+  auto iter = textures_.find(data->TexID);
+  if (iter == textures_.end()) {
+    mju_error("Texture not found: %llu", data->TexID);
+  }
 
-    data->SetTexID((ImTextureID)texture);
+  filament::Texture* texture = iter->second;
-    data->SetStatus(ImTextureStatus_OK);
+  texture->setImage(*object_mgr_->GetEngine(), 0, std::move(pb));
-  } else if (data->Status == ImTextureStatus_WantUpdates) {
+  data->SetStatus(ImTextureStatus_OK);
-    const int size = data->Width * data->Height * 4;
+}
-    filament::Texture::PixelBufferDescriptor pb(data->GetPixels(), size,
-                                                filament::Texture::Format::RGBA,
-                                                filament::Texture::Type::UBYTE);
-    filament::Texture* texture = (filament::Texture*)data->TexID;
-    texture->setImage(*engine, 0, std::move(pb));
-    data->SetStatus(ImTextureStatus_OK);
-  } else if (data->Status == ImTextureStatus_WantDestroy &&
-             data->UnusedFrames > 0) {
-    filament::Texture* texture =
-        reinterpret_cast<filament::Texture*>(data->TexID);
-    engine->destroy(texture);
 
+void GuiView::DestroyTexture(ImTextureData* data) {
+  auto iter = textures_.find(data->TexID);
+  if (iter != textures_.end()) {
+    object_mgr_->GetEngine()->destroy(iter->second);
+    textures_.erase(data->TexID);
     data->SetTexID(ImTextureID_Invalid);
     data->SetStatus(ImTextureStatus_Destroyed);
   }
@@ -166,7 +169,26 @@ bool GuiView::PrepareRenderable() {
 
   if (commands->Textures != nullptr) {
     for (ImTextureData* tex : *commands->Textures) {
-      ProcessTexture(tex);
+      if (tex->Status == ImTextureStatus_OK) {
+        // ImGui's lifecycle is independent of the filament context lifecycle.
+        // As such, it is possible to destroy and create a new filament context
+        // while ImGui is still expecting the "OK" textures to work. In this
+        // case, we simply recreate the texture.
+        if (textures_.find(tex->TexID) == textures_.end()) {
+          CreateTexture(tex);
+        }
+      } else if (tex->Status == ImTextureStatus_WantCreate) {
+        CreateTexture(tex);
+      } else if (tex->Status == ImTextureStatus_WantUpdates) {
+        if (textures_.find(tex->TexID) == textures_.end()) {
+          CreateTexture(tex);
+        } else {
+          UpdateTexture(tex);
+        }
+      } else if (tex->Status == ImTextureStatus_WantDestroy &&
+                 tex->UnusedFrames > 0) {
+        DestroyTexture(tex);
+      }
     }
   }
 
@@ -209,10 +231,10 @@ bool GuiView::PrepareRenderable() {
         clip_height = height;
       }
 
-      const intptr_t tex_id = static_cast<intptr_t>(command.GetTexID());
       mjrRect clip_rect{clip_left, clip_bottom, clip_width, clip_height};
-      builder.material(drawable_index,
+      builder.material(
-                       GetMaterialInstance(drawable_index, clip_rect, tex_id));
+          drawable_index,
+          GetMaterialInstance(drawable_index, clip_rect, command.GetTexID()));
       builder.geometry(drawable_index, kTriangles, buffer.vertex_buffer,
                       buffer.index_buffer, index_offset,
                       command.ElemCount);
@@ -232,13 +254,16 @@ bool GuiView::PrepareRenderable() {
 
 filament::MaterialInstance* GuiView::GetMaterialInstance(int index,
                                                          mjrRect rect,
-                                                         intptr_t texture_id) {
+                                                         uintptr_t texture_id) {
   while (index >= instances_.size()) {
-    const auto* texture = reinterpret_cast<filament::Texture*>(texture_id);
+    auto iter = textures_.find(texture_id);
+    if (iter == textures_.end()) {
+      mju_error("Texture not found: %lu", texture_id);
+    }
 
     filament::TextureSampler sampler;
     filament::MaterialInstance* instance = material_->createInstance();
-    instance->setParameter("glyph", texture, sampler);
+    instance->setParameter("glyph", iter->second, sampler);
     instances_.push_back(instance);
   }
 

src/experimental/filament/filament/gui_view.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_GUI_VIEW_H_
 
 #include <cstdint>
+#include <unordered_map>
 #include <vector>
 
 #include <imgui.h>
@@ -47,12 +48,14 @@ class GuiView {
   filament::View* PrepareRenderView();
 
  private:
-  void ProcessTexture(ImTextureData* data);
+  void CreateTexture(ImTextureData* data);
+  void UpdateTexture(ImTextureData* data);
+  void DestroyTexture(ImTextureData* data);
 
   // Returns the filament::MaterialInstance configured to draw into the given
   // scissor rect.
   filament::MaterialInstance* GetMaterialInstance(int index, mjrRect rect,
-                                                  intptr_t texture_id);
+                                                  uintptr_t texture_id);
 
   // Clears the filament::Scene of the UX renderable and releases all buffers.
   void ResetRenderable();
@@ -66,6 +69,7 @@ class GuiView {
   utils::Entity renderable_;
   std::vector<FilamentBuffers> buffers_;
   std::vector<filament::MaterialInstance*> instances_;
+  std::unordered_map<uintptr_t, filament::Texture*> textures_;
 };
 
 // Draws text at the given screen coordinates in clip space (i.e. [-1,-1,-1] to