Add async wrappers for UAssetManager::LoadPrimaryAsset(s), change all…

… of them to re-resolve their inputs once completed

Source/UE5Coro/Private/LatentAwaiters_AsyncLoad.cpp

@@ -29,73 +29,147 @@
 // OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 // ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+#include "Engine/AssetManager.h"
 #include "UE5Coro/LatentAwaiters.h"
 
 using namespace UE5Coro;
 using namespace UE5Coro::Private;
 
 namespace
 {
-struct FLatentLoader
+template<typename T, typename Item>
+struct TLatentLoader
 {
-	FStreamableManager Manager;
+	T Manager;
+	TArray<Item> Sources;
 	TSharedPtr<FStreamableHandle> Handle;
 
-	explicit FLatentLoader(TArray<FSoftObjectPath> Paths,
+	explicit TLatentLoader(TArray<Item> Paths,
 	                       TAsyncLoadPriority Priority)
+#if UE5CORO_CPP20
+		requires std::is_same_v<Item, FSoftObjectPath>
+#endif
+		: Sources(std::move(Paths))
 	{
-		Handle = Manager.RequestAsyncLoad(
-			std::move(Paths), FStreamableDelegate(), Priority);
+		static_assert(std::is_same_v<T, FStreamableManager>);
+		checkf(IsInGameThread(),
+		       TEXT("Latent awaiters may only be used on the game thread"));
+		Handle = Manager.RequestAsyncLoad(Sources, FStreamableDelegate(),
+		                                  Priority);
 	}
 
-	~FLatentLoader()
+	explicit TLatentLoader(TArray<Item> AssetIds,
+	                       const TArray<FName>& Bundles,
+	                       TAsyncLoadPriority Priority)
+#if UE5CORO_CPP20
+		requires std::is_same_v<Item, FPrimaryAssetId>
+#endif
+		: Manager(UAssetManager::Get()), Sources(std::move(AssetIds))
+	{
+		static_assert(std::is_same_v<T, UAssetManager&>);
+		checkf(IsInGameThread(),
+		       TEXT("Latent awaiters may only be used on the game thread"));
+		Handle = Manager.LoadPrimaryAssets(Sources, Bundles,
+		                                   FStreamableDelegate(), Priority);
+	}
+
+	~TLatentLoader()
 	{
+		checkf(IsInGameThread(), TEXT("Unexpected cleanup off the game thread"));
 		if (Handle)
 			Handle->ReleaseHandle();
 	}
+
+	TArray<UObject*> ResolveItems()
+	{
+		checkf(IsInGameThread(),
+		       TEXT("Unexpected object resolve request off the game thread"));
+		// Handle->GetLoadedAssets() is unreliable, the async loading BP nodes
+		// re-resolve the sources instead once loading is done. Let's do that.
+		TArray<UObject*> Items;
+		for (auto& i : Sources)
+		{
+			UObject* Obj = nullptr;
+			if constexpr (std::is_same_v<Item, FSoftObjectPath>)
+				Obj = i.ResolveObject();
+			else if constexpr (std::is_same_v<Item, FPrimaryAssetId>)
+				Obj = Manager.GetPrimaryAssetObject(i);
+			else
+				// This needs to depend on a template parameter
+				static_assert(false && std::is_void_v<Item>, "Unknown type");
+
+			// Null filtering matches how the array BP nodes behave
+			if (IsValid(Obj))
+				Items.Add(Obj);
+		}
+		return Items;
+	}
 };
+using FLatentLoader = TLatentLoader<FStreamableManager, FSoftObjectPath>;
+using FPrimaryLoader = TLatentLoader<UAssetManager&, FPrimaryAssetId>;
 
+template<typename T>
 bool ShouldResume(void*& Loader, bool bCleanup)
 {
-	auto* This = static_cast<FLatentLoader*>(Loader);
+	auto* This = static_cast<T*>(Loader);
 
 	if (UNLIKELY(bCleanup))
 	{
 		delete This;
 		return false;
 	}
 
-	// This is the same logic that FLoadAssetActionBase::UpdateOperation() uses
+	// This is the same logic that FLoadAssetActionBase::UpdateOperation() uses.
+	// !Handle is how UAssetManager communicates an instant/synchronous finish.
 	auto& Handle = This->Handle;
 	return !Handle || Handle->HasLoadCompleted() || Handle->WasCanceled();
 }
 }
 
+template<int HiddenType>
 TArray<UObject*> AsyncLoad::InternalResume(void* State)
 {
-	checkf(ShouldResume(State, false), TEXT("Internal error"));
+	using T = std::conditional_t<HiddenType == 0, FLatentLoader, FPrimaryLoader>;
+	checkf(ShouldResume<T>(State, false), TEXT("Internal error"));
 
-	TArray<UObject*> Assets;
-	if (auto* This = static_cast<FLatentLoader*>(State); This->Handle)
-		This->Handle->GetLoadedAssets(Assets);
-	return Assets;
+	return static_cast<T*>(State)->ResolveItems();
 }
+template UE5CORO_API TArray<UObject*> AsyncLoad::InternalResume<0>(void*);
+template UE5CORO_API TArray<UObject*> AsyncLoad::InternalResume<1>(void*);
 
 FLatentAwaiter Latent::AsyncLoadObjects(TArray<FSoftObjectPath> Paths,
                                         TAsyncLoadPriority Priority)
 {
 	return FLatentAwaiter(new FLatentLoader(std::move(Paths), Priority),
-		&ShouldResume);
+	                      &ShouldResume<FLatentLoader>);
+}
+
+FLatentAwaiter Latent::AsyncLoadPrimaryAsset(
+	const FPrimaryAssetId& AssetToLoad,
+	const TArray<FName>& LoadBundles,
+	TAsyncLoadPriority Priority)
+{
+	return AsyncLoadPrimaryAssets(TArray{AssetToLoad}, LoadBundles, Priority);
+}
+
+FLatentAwaiter Latent::AsyncLoadPrimaryAssets(
+	TArray<FPrimaryAssetId> AssetsToLoad,
+	const TArray<FName>& LoadBundles,
+	TAsyncLoadPriority Priority)
+{
+	return FLatentAwaiter(
+		new FPrimaryLoader(std::move(AssetsToLoad), LoadBundles, Priority),
+		&ShouldResume<FPrimaryLoader>);
 }
 
-TAsyncLoadAwaiter<UClass*> Latent::AsyncLoadClass(TSoftClassPtr<UObject> Ptr,
-                                                  TAsyncLoadPriority Priority)
+TAsyncLoadAwaiter<UClass*, 0> Latent::AsyncLoadClass(TSoftClassPtr<UObject> Ptr,
+                                                     TAsyncLoadPriority Priority)
 {
-	return TAsyncLoadAwaiter<UClass*>(
+	return TAsyncLoadAwaiter<UClass*, 0>(
 		AsyncLoadObjects(TArray{Ptr.ToSoftObjectPath()}, Priority));
 }
 
-TAsyncLoadAwaiter<TArray<UClass*>> Latent::AsyncLoadClasses(
+TAsyncLoadAwaiter<TArray<UClass*>, 0> Latent::AsyncLoadClasses(
 	const TArray<TSoftClassPtr<UObject>>& Ptrs,
 	TAsyncLoadPriority Priority)
 {
@@ -104,7 +178,7 @@ TAsyncLoadAwaiter<TArray<UClass*>> Latent::AsyncLoadClasses(
 	for (const auto& Ptr : Ptrs)
 		Paths.Add(Ptr.ToSoftObjectPath());
 
-	return TAsyncLoadAwaiter<TArray<UClass*>>(
+	return TAsyncLoadAwaiter<TArray<UClass*>, 0>(
 		AsyncLoadObjects(std::move(Paths), Priority));
 }
 

Source/UE5Coro/Public/UE5Coro/LatentAwaiters.h

@@ -45,7 +45,7 @@ class FLatentCancellation;
 class FLatentChainAwaiter;
 class FLatentPromise;
 class FPackageLoadAwaiter;
-template<typename> class TAsyncLoadAwaiter;
+template<typename, int> class TAsyncLoadAwaiter;
 template<typename> class TAsyncQueryAwaiter;
 }
 
@@ -131,15 +131,15 @@ Private::FLatentChainAwaiter ChainEx(F&& Function, A&&... Args);
 /** Asynchronously starts loading the object, resumes once it's loaded.<br>
  *  The result of the co_await expression is the loaded T*. */
 template<typename T>
-std::enable_if_t<std::is_base_of_v<UObject, T>, Private::TAsyncLoadAwaiter<T*>>
+std::enable_if_t<std::is_base_of_v<UObject, T>, Private::TAsyncLoadAwaiter<T*, 0>>
 AsyncLoadObject(TSoftObjectPtr<T>,
                 TAsyncLoadPriority = FStreamableManager::DefaultAsyncLoadPriority);
 
 /** Asynchronously starts loading the objects, resumes once they're loaded.<br>
  *  The result of the co_await expression is TArray<T*>. */
 template<typename T>
 std::enable_if_t<std::is_base_of_v<UObject, T>,
-                 Private::TAsyncLoadAwaiter<TArray<T*>>>
+                 Private::TAsyncLoadAwaiter<TArray<T*>, 0>>
 AsyncLoadObjects(const TArray<TSoftObjectPtr<T>>&,
                  TAsyncLoadPriority = FStreamableManager::DefaultAsyncLoadPriority);
 
@@ -149,15 +149,54 @@ UE5CORO_API Private::FLatentAwaiter AsyncLoadObjects(
 	TArray<FSoftObjectPath>,
 	TAsyncLoadPriority = FStreamableManager::DefaultAsyncLoadPriority);
 
+/** Asynchronously starts loading the primary asset with any bundles specified,
+ *  resumes once they're loaded.<br>
+ *  The asset will stay in memory until explicitly unloaded. */
+UE5CORO_API Private::FLatentAwaiter AsyncLoadPrimaryAsset(
+	const FPrimaryAssetId& AssetToLoad,
+	const TArray<FName>& LoadBundles = {},
+	TAsyncLoadPriority Priority = FStreamableManager::DefaultAsyncLoadPriority);
+
+/** Asynchronously starts loading the primary asset of the given type with any
+ *  bundles specified, resumes once they're loaded.<br>
+ *  The asset will stay in memory until explicitly unloaded.<br>
+ *  The result of the co_await expression is the loaded T* or nullptr. */
+template<typename T>
+std::enable_if_t<std::is_base_of_v<UObject, T>, Private::TAsyncLoadAwaiter<T*, 1>>
+AsyncLoadPrimaryAsset(
+	FPrimaryAssetId AssetToLoad,
+	const TArray<FName>& LoadBundles = {},
+	TAsyncLoadPriority Priority = FStreamableManager::DefaultAsyncLoadPriority);
+
+/** Asynchronously starts loading the primary assets with any bundles specified,
+ *  resumes once they're loaded.<br>
+ *  The assets will stay in memory until explicitly unloaded. */
+UE5CORO_API Private::FLatentAwaiter AsyncLoadPrimaryAssets(
+	TArray<FPrimaryAssetId> AssetsToLoad,
+	const TArray<FName>& LoadBundles = {},
+	TAsyncLoadPriority Priority = FStreamableManager::DefaultAsyncLoadPriority);
+
+/** Asynchronously starts loading the primary assets of the given type with any
+ *  bundles specified, resumes once they're loaded.<br>
+ *  The assets will stay in memory until explicitly unloaded.<br>
+ *  The result of the co_await expression is the loaded and filtered TArray<T*>. */
+template<typename T>
+std::enable_if_t<std::is_base_of_v<UObject, T>,
+                 Private::TAsyncLoadAwaiter<TArray<T*>, 1>>
+AsyncLoadPrimaryAssets(
+	TArray<FPrimaryAssetId> AssetsToLoad,
+	const TArray<FName>& LoadBundles = {},
+	TAsyncLoadPriority Priority = FStreamableManager::DefaultAsyncLoadPriority);
+
 /** Asynchronously starts loading the class, resumes once it's loaded.<br>
  *  The result of the co_await expression is the loaded UClass*. */
-UE5CORO_API Private::TAsyncLoadAwaiter<UClass*> AsyncLoadClass(
+UE5CORO_API Private::TAsyncLoadAwaiter<UClass*, 0> AsyncLoadClass(
 	TSoftClassPtr<UObject>,
 	TAsyncLoadPriority = FStreamableManager::DefaultAsyncLoadPriority);
 
 /** Asynchronously starts loading the classes, resumes once they're loaded.<br>
  *  The result of the co_await expression is TArray<UClass*>. */
-UE5CORO_API Private::TAsyncLoadAwaiter<TArray<UClass*>> AsyncLoadClasses(
+UE5CORO_API Private::TAsyncLoadAwaiter<TArray<UClass*>, 0> AsyncLoadClasses(
 	const TArray<TSoftClassPtr<UObject>>&,
 	TAsyncLoadPriority = FStreamableManager::DefaultAsyncLoadPriority);
 
@@ -266,10 +305,11 @@ class [[nodiscard]] UE5CORO_API FLatentAwaiter
 
 namespace AsyncLoad
 {
+template<int> // Switches between non-exported types
 UE5CORO_API TArray<UObject*> InternalResume(void*);
 }
 
-template<typename T>
+template<typename T, int HiddenType>
 class [[nodiscard]] TAsyncLoadAwaiter : public FLatentAwaiter
 {
 public:
@@ -278,16 +318,16 @@ class [[nodiscard]] TAsyncLoadAwaiter : public FLatentAwaiter
 
 	T await_resume()
 	{
-		TArray<UObject*> Assets = AsyncLoad::InternalResume(State);
+		TArray<UObject*> Assets = AsyncLoad::InternalResume<HiddenType>(State);
 		if constexpr (TIsTArray<T>::Value)
 		{
 			static_assert(std::is_pointer_v<typename T::ElementType>);
 			using V = std::remove_pointer_t<typename T::ElementType>;
 			static_assert(std::is_base_of_v<UObject, V>);
-			checkCode(
-				for (auto* Ptr : Assets)
-					check(Ptr->IsA<V>());
-			);
+			[[maybe_unused]] int OldNum = Assets.Num();
+			Assets.RemoveAll([](UObject* Ptr) { return !Ptr->IsA<V>(); });
+			if constexpr (HiddenType == 0) // These come from typed soft ptrs
+				check(Assets.Num() == OldNum); // Strict check
 			return reinterpret_cast<T&&>(Assets);
 		}
 		else
@@ -302,8 +342,10 @@ class [[nodiscard]] TAsyncLoadAwaiter : public FLatentAwaiter
 	}
 };
 
-static_assert(sizeof(FLatentAwaiter) == sizeof(TAsyncLoadAwaiter<UObject*>));
-static_assert(sizeof(FLatentAwaiter) == sizeof(TAsyncLoadAwaiter<TArray<UObject*>>));
+static_assert(sizeof(FLatentAwaiter) ==
+              sizeof(TAsyncLoadAwaiter<UObject*, 0>));
+static_assert(sizeof(FLatentAwaiter) ==
+              sizeof(TAsyncLoadAwaiter<TArray<UObject*>, 0>));
 
 class [[nodiscard]] UE5CORO_API FPackageLoadAwaiter
 {
@@ -352,17 +394,17 @@ inline UE5Coro::Private::FLatentCancellation UE5Coro::Latent::Cancel()
 
 template<typename T>
 std::enable_if_t<std::is_base_of_v<UObject, T>,
-                 UE5Coro::Private::TAsyncLoadAwaiter<T*>>
+                 UE5Coro::Private::TAsyncLoadAwaiter<T*, 0>>
 UE5Coro::Latent::AsyncLoadObject(TSoftObjectPtr<T> Ptr,
                                  TAsyncLoadPriority Priority)
 {
-	return Private::TAsyncLoadAwaiter<T*>(
+	return Private::TAsyncLoadAwaiter<T*, 0>(
 		AsyncLoadObjects(TArray{Ptr.ToSoftObjectPath()}, Priority));
 }
 
 template<typename T>
 std::enable_if_t<std::is_base_of_v<UObject, T>,
-                 UE5Coro::Private::TAsyncLoadAwaiter<TArray<T*>>>
+                 UE5Coro::Private::TAsyncLoadAwaiter<TArray<T*>, 0>>
 UE5Coro::Latent::AsyncLoadObjects(const TArray<TSoftObjectPtr<T>>& Ptrs,
                                   TAsyncLoadPriority Priority)
 {
@@ -371,8 +413,30 @@ UE5Coro::Latent::AsyncLoadObjects(const TArray<TSoftObjectPtr<T>>& Ptrs,
 	for (const auto& Ptr : Ptrs)
 		Paths.Add(Ptr.ToSoftObjectPath());
 
-	return Private::TAsyncLoadAwaiter<TArray<T*>>(
+	return Private::TAsyncLoadAwaiter<TArray<T*>, 0>(
 		AsyncLoadObjects(std::move(Paths), Priority));
 }
 
+template<typename T>
+std::enable_if_t<std::is_base_of_v<UObject, T>,
+                 UE5Coro::Private::TAsyncLoadAwaiter<T*, 1>>
+UE5Coro::Latent::AsyncLoadPrimaryAsset(FPrimaryAssetId AssetToLoad,
+                                       const TArray<FName>& LoadBundles,
+                                       TAsyncLoadPriority Priority)
+{
+	return Private::TAsyncLoadAwaiter<T*, 1>(
+		AsyncLoadPrimaryAsset(std::move(AssetToLoad), LoadBundles, Priority));
+}
+
+template<typename T>
+std::enable_if_t<std::is_base_of_v<UObject, T>,
+                 UE5Coro::Private::TAsyncLoadAwaiter<TArray<T*>, 1>>
+UE5Coro::Latent::AsyncLoadPrimaryAssets(TArray<FPrimaryAssetId> AssetsToLoad,
+                                        const TArray<FName>& LoadBundles,
+                                        TAsyncLoadPriority Priority)
+{
+	return Private::TAsyncLoadAwaiter<TArray<T*>, 1>(
+		AsyncLoadPrimaryAssets(std::move(AssetsToLoad), LoadBundles, Priority));
+}
+
 #include "LatentChain.inl"

UE5Coro.uplugin

@@ -1,7 +1,7 @@
 {
 	"FileVersion": 3,
 	"Version": 1,
-	"VersionName": "1.5",
+	"VersionName": "1.6",
 	"FriendlyName": "UE5Coro",
 	"Description": "C++17/20 coroutine implementation for Unreal Engine",
 	"Category": "Programming",