Fix #134264: Cannot undo asset catalog changes to before initial save

When saving a file, move the catalog service from the runtime asset library into the new on-disk
asset library. This makes all catalog data like the undo history and deleted catalogs be preserved.
The fact that a new asset library type gets allocated is an implementation detail that shouldn't
affect behavior.

Once the service is moved, an undo push is added (so this state can be restored to), and if the new
.blend file location can be associated with a catalog definition file, its catalogs are merged in.

Includes unit tests.

Pull Request: https://projects.blender.org/blender/blender/pulls/135589

source/blender/asset_system/AS_asset_catalog.hh

@@ -50,8 +50,11 @@ class AssetCatalogService {
   Vector<std::unique_ptr<AssetCatalogCollection>> undo_snapshots_;
   Vector<std::unique_ptr<AssetCatalogCollection>> redo_snapshots_;
 
-  const CatalogFilePath asset_library_root_;
-  const bool is_read_only_ = false;
+  CatalogFilePath asset_library_root_;
+  bool is_read_only_ = false;
+
+  friend class AssetLibraryService;
+  friend class AssetLibrary;
 
  public:
   static const CatalogFilePath DEFAULT_CATALOG_FILENAME;

source/blender/asset_system/AS_asset_library.hh

@@ -128,7 +128,7 @@ class AssetLibrary {
    */
   virtual std::optional<AssetLibraryReference> library_reference() const = 0;
 
-  void load_catalogs();
+  void load_or_reload_catalogs();
 
   AssetCatalogService &catalog_service() const;
 

source/blender/asset_system/intern/asset_library.cc

@@ -26,6 +26,7 @@
 #include "asset_catalog_collection.hh"
 #include "asset_catalog_definition_file.hh"
 #include "asset_library_service.hh"
+#include "runtime_library.hh"
 #include "utils.hh"
 
 using namespace blender;
@@ -166,7 +167,7 @@ AssetLibrary::AssetLibrary(eAssetLibraryType library_type, StringRef name, Strin
     : library_type_(library_type),
       name_(name),
       root_path_(std::make_shared<std::string>(utils::normalize_directory_path(root_path))),
-      catalog_service_(std::make_unique<AssetCatalogService>())
+      catalog_service_(std::make_unique<AssetCatalogService>(*root_path_))
 {
 }
 
@@ -184,12 +185,26 @@ void AssetLibrary::foreach_loaded(FunctionRef<void(AssetLibrary &)> fn,
   service->foreach_loaded_asset_library(fn, include_all_library);
 }
 
-void AssetLibrary::load_catalogs()
+void AssetLibrary::load_or_reload_catalogs()
 {
-  auto catalog_service = std::make_unique<AssetCatalogService>(root_path());
-  catalog_service->load_from_disk();
-  std::lock_guard lock{catalog_service_mutex_};
-  catalog_service_ = std::move(catalog_service);
+  {
+    std::lock_guard lock{catalog_service_mutex_};
+    /* Should never actually be the case, catalog service gets allocated with the asset library. */
+    if (catalog_service_ == nullptr) {
+      auto catalog_service = std::make_unique<AssetCatalogService>(root_path());
+      catalog_service->load_from_disk();
+      catalog_service_ = std::move(catalog_service);
+      return;
+    }
+  }
+
+  /* The catalog service was created before without being associated with a definition file. */
+  if (catalog_service_->get_catalog_definition_file() == nullptr) {
+    catalog_service_->load_from_disk();
+  }
+  else {
+    this->refresh_catalogs();
+  }
 }
 
 AssetCatalogService &AssetLibrary::catalog_service() const
@@ -253,11 +268,25 @@ void asset_library_on_save_post(Main *bmain,
                                 void *arg)
 {
   AssetLibrary *asset_lib = static_cast<AssetLibrary *>(arg);
-  asset_lib->on_blend_save_post(bmain, pointers, num_pointers);
 
-  if (asset_lib->library_type() == ASSET_LIBRARY_LOCAL) {
-    AssetLibraryService::destroy_runtime_current_file_library();
+  /* Transform 'runtime' current file library into 'on-disk' current file library. */
+  if (asset_lib->library_type() == ASSET_LIBRARY_LOCAL && asset_lib->root_path().is_empty()) {
+    BLI_assert(dynamic_cast<RuntimeAssetLibrary *>(asset_lib) != nullptr);
+
+    if (AssetLibrary *on_disk_lib =
+            AssetLibraryService::move_runtime_current_file_into_on_disk_library(*bmain))
+    {
+      /* Allow undoing to the state before merging in catalogs from disk. */
+      on_disk_lib->catalog_service().undo_push();
+
+      /* Force refresh to merge on-disk catalogs with the ones stolen from the runtime library. */
+      asset_lib = AssetLibraryService::get()->get_asset_library_on_disk_builtin(
+          ASSET_LIBRARY_LOCAL, on_disk_lib->root_path());
+      BLI_assert(asset_lib == on_disk_lib);
+    }
   }
+
+  asset_lib->on_blend_save_post(bmain, pointers, num_pointers);
 }
 
 }  // namespace

source/blender/asset_system/intern/asset_library_service.cc

@@ -116,22 +116,19 @@ AssetLibrary *AssetLibraryService::get_asset_library(
 
 AssetLibrary *AssetLibraryService::get_asset_library_on_disk(eAssetLibraryType library_type,
                                                              StringRef name,
-                                                             StringRefNull root_path)
+                                                             StringRefNull root_path,
+                                                             const bool load_catalogs)
 {
-  BLI_assert_msg(!root_path.is_empty(),
-                 "top level directory must be given for on-disk asset library");
-
-  std::string normalized_root_path = utils::normalize_directory_path(root_path);
-
-  std::unique_ptr<OnDiskAssetLibrary> *lib_uptr_ptr = on_disk_libraries_.lookup_ptr(
-      {library_type, normalized_root_path});
-  if (lib_uptr_ptr != nullptr) {
-    CLOG_INFO(&LOG, 2, "get \"%s\" (cached)", normalized_root_path.c_str());
-    AssetLibrary *lib = lib_uptr_ptr->get();
-    lib->refresh_catalogs();
+  if (OnDiskAssetLibrary *lib = this->lookup_on_disk_library(library_type, root_path)) {
+    CLOG_INFO(&LOG, 2, "get \"%s\" (cached)", root_path.c_str());
+    if (load_catalogs) {
+      lib->load_or_reload_catalogs();
+    }
     return lib;
   }
 
+  const std::string normalized_root_path = utils::normalize_directory_path(root_path);
+
   std::unique_ptr<OnDiskAssetLibrary> lib_uptr;
   switch (library_type) {
     case ASSET_LIBRARY_CUSTOM:
@@ -147,7 +144,9 @@ AssetLibrary *AssetLibraryService::get_asset_library_on_disk(eAssetLibraryType l
 
   AssetLibrary *lib = lib_uptr.get();
 
-  lib->load_catalogs();
+  if (load_catalogs) {
+    lib->load_or_reload_catalogs();
+  }
 
   on_disk_libraries_.add_new({library_type, normalized_root_path}, std::move(lib_uptr));
   CLOG_INFO(&LOG, 2, "get \"%s\" (loaded)", normalized_root_path.c_str());
@@ -205,10 +204,50 @@ void AssetLibraryService::reload_all_library_catalogs_if_dirty()
   }
 }
 
-void AssetLibraryService::destroy_runtime_current_file_library()
+AssetLibrary *AssetLibraryService::move_runtime_current_file_into_on_disk_library(
+    const Main &bmain)
 {
   AssetLibraryService &library_service = *AssetLibraryService::get();
+
+  const std::string root_path = AS_asset_library_find_suitable_root_path_from_main(&bmain);
+  if (root_path.empty()) {
+    return nullptr;
+  }
+
+  BLI_assert_msg(!library_service.lookup_on_disk_library(ASSET_LIBRARY_LOCAL, root_path),
+                 "On-disk \"Current File\" asset library shouldn't exist yet, it should only be "
+                 "created now in response to initially saving the file - catalog service "
+                 "will be overridden");
+
+  /* Create on disk library without loading catalogs. We'll steal the catalog service from the
+   * runtime library below. */
+  AssetLibrary *on_disk_library = library_service.get_asset_library_on_disk(
+      ASSET_LIBRARY_LOCAL,
+      {},
+      root_path,
+      /*load_catalogs=*/false);
+
+  {
+    /* These should always be completely separate, just sanity check since it would cause a
+     * deadlock below. */
+    BLI_assert(on_disk_library != library_service.current_file_library_.get());
+
+    std::lock_guard lock_on_disk{on_disk_library->catalog_service_mutex_};
+    std::lock_guard lock_runtime{library_service.current_file_library_->catalog_service_mutex_};
+    on_disk_library->catalog_service_.swap(
+        library_service.current_file_library_->catalog_service_);
+  }
+
+  on_disk_library->catalog_service().asset_library_root_ = on_disk_library->root_path();
+  /* The catalogs are not stored on disk, so there should not be any CDF. Otherwise, we'd have to
+   * remap their stored file-path too (#AssetCatalogDefinitionFile.file_path). */
+  BLI_assert_msg(on_disk_library->catalog_service().get_catalog_definition_file() == nullptr,
+                 "new on-disk library shouldn't have catalog definition files - root path "
+                 "changed, so they would have to be relocated");
+
   library_service.current_file_library_ = nullptr;
+
+  return on_disk_library;
 }
 
 AssetLibrary *AssetLibraryService::get_asset_library_all(const Main *bmain)
@@ -238,6 +277,19 @@ AssetLibrary *AssetLibraryService::get_asset_library_all(const Main *bmain)
   return all_library_.get();
 }
 
+OnDiskAssetLibrary *AssetLibraryService::lookup_on_disk_library(eAssetLibraryType library_type,
+                                                                StringRefNull root_path)
+{
+  BLI_assert_msg(!root_path.is_empty(),
+                 "top level directory must be given for on-disk asset library");
+
+  std::string normalized_root_path = utils::normalize_directory_path(root_path);
+
+  std::unique_ptr<OnDiskAssetLibrary> *lib_uptr_ptr = on_disk_libraries_.lookup_ptr(
+      {library_type, normalized_root_path});
+  return lib_uptr_ptr ? lib_uptr_ptr->get() : nullptr;
+}
+
 bUserAssetLibrary *AssetLibraryService::find_custom_preferences_asset_library_from_asset_weak_ref(
     const AssetWeakReference &asset_reference)
 {

source/blender/asset_system/intern/asset_library_service.hh

@@ -80,10 +80,16 @@ class AssetLibraryService {
   static bUserAssetLibrary *find_custom_preferences_asset_library_from_asset_weak_ref(
       const AssetWeakReference &asset_reference);
   /**
-   * Call when the .blend file is saved or unloaded to destroy the runtime library. It's not
-   * represented as a on disk library.
+   * Turn the runtime current file library into a on-disk current file library, preserving catalog
+   * data like undo/redo history, deleted catalog info, catalog saving state, etc. Note that this
+   * creates a new on-disk asset library and destroys the runtime one.
+   *
+   * Call when the .blend file is saved to disk.
+   *
+   * \return the new on-disk current file asset library (null in case of failure to find a path to
+   * store the library in, based on the #Main.filepath from \a main).
    */
-  static void destroy_runtime_current_file_library();
+  static AssetLibrary *move_runtime_current_file_into_on_disk_library(const Main &bmain);
 
   AssetLibrary *get_asset_library(const Main *bmain,
                                   const AssetLibraryReference &library_reference);
@@ -161,6 +167,8 @@ class AssetLibraryService {
   /** Allocate a new instance of the service and assign it to `instance_`. */
   static void allocate_service_instance();
 
+  OnDiskAssetLibrary *lookup_on_disk_library(eAssetLibraryType type, StringRefNull root_path);
+
   AssetLibrary *find_loaded_on_disk_asset_library_from_name(StringRef name) const;
 
   /**
@@ -170,7 +178,8 @@ class AssetLibraryService {
    */
   AssetLibrary *get_asset_library_on_disk(eAssetLibraryType library_type,
                                           StringRef name,
-                                          StringRefNull root_path);
+                                          StringRefNull root_path,
+                                          bool load_catalogs = true);
   /**
    * Ensure the AssetLibraryService instance is destroyed before a new blend file is loaded.
    * This makes memory management simple, and ensures a fresh start for every blend file. */

source/blender/asset_system/tests/asset_library_service_test.cc

@@ -240,4 +240,145 @@ TEST_F(AssetLibraryServiceTest, has_any_unsaved_catalogs_after_write)
   EXPECT_FALSE(cat->flags.has_unsaved_changes);
 }
 
+/** Call #AssetLibraryService::move_runtime_current_file_into_on_disk_library() with a on disk
+ * location that contains no existing asset catalog definition file. */
+TEST_F(AssetLibraryServiceTest, move_runtime_current_file_into_on_disk_library__empty_directory)
+{
+  AssetLibraryService *service = AssetLibraryService::get();
+
+  AssetLibrary *runtime_lib = service->get_asset_library_current_file();
+  AssetCatalogService &runtime_catservice = runtime_lib->catalog_service();
+
+  /* Catalog created in the runtime lib that should be moved to the on-disk lib. */
+  AssetCatalog *catalog = runtime_catservice.create_catalog("Some/Catalog/Path");
+  runtime_catservice.undo_push();
+
+  {
+    EXPECT_TRUE(catalog->flags.has_unsaved_changes);
+
+    EXPECT_EQ(nullptr, runtime_catservice.find_catalog(UUID_POSES_ELLIE))
+        << "Catalog not expected in the runtime asset library.";
+  }
+
+  {
+    Main dummy_main{};
+    std::string dummy_filepath = create_temp_path() + "dummy.blend";
+    STRNCPY(dummy_main.filepath, dummy_filepath.c_str());
+
+    AssetLibraryService::move_runtime_current_file_into_on_disk_library(dummy_main);
+
+    AssetLibraryReference ref{};
+    ref.type = ASSET_LIBRARY_LOCAL;
+
+    /* Loads and merges the catalogs from disk. */
+    AssetLibrary *on_disk_lib = service->get_asset_library(&dummy_main, ref);
+    AssetCatalogService &on_disk_catservice = on_disk_lib->catalog_service();
+
+    /* Can only test the pointer equality here because the implementation keeps the runtime library
+     * alive until all its contents are moved to the on-disk library. Otherwise the allocator might
+     * choose the same address for the new on-disk library. Useful for testing, though not
+     * required. */
+    EXPECT_NE(on_disk_lib, runtime_lib);
+    EXPECT_EQ(on_disk_lib->root_path(), temp_library_path_);
+
+    /* Check if catalog was moved correctly .*/
+    {
+      EXPECT_EQ(on_disk_catservice.find_catalog(catalog->catalog_id)->path, catalog->path);
+      /* Compare catalog by pointer. #move_runtime_current_file_into_on_disk_library() doesn't
+       * guarantee publicly that catalog pointers remain unchanged, but practically code might rely
+       * on it. Good to know if this breaks. */
+      EXPECT_EQ(on_disk_catservice.find_catalog(catalog->catalog_id), catalog);
+      /* No writing happened, just merging. */
+      EXPECT_TRUE(on_disk_catservice.find_catalog(catalog->catalog_id)->flags.has_unsaved_changes);
+    }
+
+    EXPECT_EQ(nullptr, runtime_catservice.find_catalog(UUID_POSES_ELLIE))
+        << "Catalog not expected in the on disk asset library.";
+
+    /* Check if undo stack was moved correctly. */
+    {
+      on_disk_catservice.undo();
+      const AssetCatalog *ellie_catalog = on_disk_catservice.find_catalog(UUID_POSES_ELLIE);
+      EXPECT_EQ(nullptr, ellie_catalog) << "This catalog should not be present after undo";
+      EXPECT_EQ(on_disk_catservice.find_catalog(catalog->catalog_id)->path, catalog->path)
+          << "This catalog should still be present after undo";
+    }
+
+    /* Force a new current file runtime library to be created. */
+    EXPECT_NE(service->get_asset_library_current_file(), on_disk_lib);
+  }
+}
+
+/** Call #AssetLibraryService::move_runtime_current_file_into_on_disk_library() with a on disk
+ * location that contains an existing asset catalog definition file. Result should be merged
+ * libraries. */
+TEST_F(AssetLibraryServiceTest,
+       move_runtime_current_file_into_on_disk_library__directory_with_catalogs)
+{
+  AssetLibraryService *service = AssetLibraryService::get();
+
+  AssetLibrary *runtime_lib = service->get_asset_library_current_file();
+  AssetCatalogService &runtime_catservice = runtime_lib->catalog_service();
+
+  /* Catalog created in the runtime lib that should be moved to the on-disk lib. */
+  AssetCatalog *catalog = runtime_catservice.create_catalog("Some/Catalog/Path");
+  runtime_catservice.undo_push();
+
+  {
+    EXPECT_TRUE(catalog->flags.has_unsaved_changes);
+
+    EXPECT_EQ(nullptr, runtime_catservice.find_catalog(UUID_POSES_ELLIE))
+        << "Catalog not expected in the runtime asset library.";
+  }
+
+  {
+    Main dummy_main{};
+    std::string dummy_filepath = asset_library_root_ + SEP + "dummy.blend";
+    STRNCPY(dummy_main.filepath, dummy_filepath.c_str());
+
+    AssetLibraryService::move_runtime_current_file_into_on_disk_library(dummy_main);
+
+    AssetLibraryReference ref{};
+    ref.type = ASSET_LIBRARY_LOCAL;
+
+    /* Loads and merges the catalogs from disk. */
+    AssetLibrary *on_disk_lib = service->get_asset_library(&dummy_main, ref);
+    AssetCatalogService &on_disk_catservice = on_disk_lib->catalog_service();
+
+    EXPECT_NE(on_disk_lib, runtime_lib);
+    EXPECT_EQ(on_disk_lib->root_path(), asset_library_root_ + SEP);
+
+    /* Check if catalog was moved correctly .*/
+    {
+      EXPECT_EQ(on_disk_catservice.find_catalog(catalog->catalog_id)->path, catalog->path);
+      /* Compare catalog by pointer. #move_runtime_current_file_into_on_disk_library() doesn't
+       * guarantee publicly that catalog pointers remain unchanged, but practically code might rely
+       * on it. Good to know if this breaks. */
+      EXPECT_EQ(on_disk_catservice.find_catalog(catalog->catalog_id), catalog);
+      /* No writing happened, just merging. */
+      EXPECT_TRUE(on_disk_catservice.find_catalog(catalog->catalog_id)->flags.has_unsaved_changes);
+    }
+
+    /* Check if catalogs have been merged in from disk correctly (by #get_asset_library()). */
+    {
+      const AssetCatalog *ellie_catalog = on_disk_catservice.find_catalog(UUID_POSES_ELLIE);
+      EXPECT_NE(nullptr, ellie_catalog)
+          << "Catalogs should be loaded after getting an asset library from disk.";
+      EXPECT_FALSE(ellie_catalog->flags.has_unsaved_changes);
+    }
+
+    /* Check if undo stack was moved correctly. */
+    {
+      on_disk_catservice.undo();
+      const AssetCatalog *ellie_catalog = on_disk_catservice.find_catalog(UUID_POSES_ELLIE);
+      EXPECT_EQ(nullptr, ellie_catalog) << "This catalog should not be present after undo";
+      EXPECT_EQ(on_disk_catservice.find_catalog(catalog->catalog_id)->path, catalog->path)
+          << "This catalog should still be present after undo";
+    }
+
+    /* Force a new current file runtime library to be created. */
+    EXPECT_NE(service->get_asset_library_current_file(), on_disk_lib);
+  }
+}
+
 }  // namespace blender::asset_system::tests