Improve the WingTypes registry

Plugins/UEWingman/Source/UEWingman/Private/WingTypes.cpp

@@ -3,7 +3,10 @@
 #include "Editor.h"
 #include "EdGraphSchema_K2.h"
 #include "Engine/Blueprint.h"
+#include "StructUtils/UserDefinedStruct.h"
+#include "Engine/UserDefinedEnum.h"
 #include "UObject/UObjectIterator.h"
+#include "AssetRegistry/AssetRegistryModule.h"
 
 // ---------------------------------------------------------------------------
 // Choose Short Name
@@ -29,42 +32,51 @@ void UWingTypes::ReserveShortName(FName Name)
 	ShortToPath.Add(NameStr.ToLower(), FString(TEXT("PRIMITIVE")));
 }
 
-FString UWingTypes::ChooseShortName(const UObject* Obj)
+FString UWingTypes::ChooseShortName(const FString &Proposal, const FString &FullObjectPath)
 {
-	if (!Cast<UScriptStruct>(Obj) && !Cast<UClass>(Obj) && !Cast<UEnum>(Obj))
+	// You must call this with an object path, not an asset path.
-		return FString();
+	check(FullObjectPath.Contains(TEXT(".")));
 
-	FString Path = Obj->GetPathName();
+	// Look to see if we've already assigned a short name to this path.
-	FString *OldShort = PathToShort.Find(Path);
+	FString *OldShort = PathToShort.Find(FullObjectPath);
 	if (OldShort != nullptr) return *OldShort;
 
-	FString Name = GetNameWithoutUnderscoreC(Obj);
+	// Check if the proposed name is available.
-
+	FString Lower = Proposal.ToLower();
-	FString Lower = Name.ToLower();
 	if (!ShortToPath.Contains(Lower))
 	{
-		ShortToPath.Add(Lower, Path);
+		ShortToPath.Add(Lower, FullObjectPath);
-		PathToShort.Add(Path, Name);
+		PathToShort.Add(FullObjectPath, Proposal);
-		return Name;
+		return Proposal;
 	}
 
+	// The proposed name is not available.  Add numbers until we
+	// find a name that is available.
 	for (int32 i = 2; ; ++i)
 	{
 		FString NumberedLower = FString::Printf(TEXT("%s%d"), *Lower, i);
 		if (!ShortToPath.Contains(NumberedLower))
 		{
-			FString NumberedName = FString::Printf(TEXT("%s_%d"), *Name, i);
+			FString NumberedProposal = FString::Printf(TEXT("%s_%d"), *Proposal, i);
-			ShortToPath.Add(NumberedLower, Path);
+			ShortToPath.Add(NumberedLower, FullObjectPath);
-			PathToShort.Add(Path, NumberedName);
+			PathToShort.Add(FullObjectPath, NumberedProposal);
-			return NumberedName;
+			return NumberedProposal;
 		}
 	}
 }
 
+FString UWingTypes::ChooseShortName(const UObject* Obj)
+{
+	if (!Cast<UScriptStruct>(Obj) && !Cast<UClass>(Obj) && !Cast<UEnum>(Obj))
+		return FString();
+
+	FString ProposedName = GetNameWithoutUnderscoreC(Obj);
+	return ChooseShortName(ProposedName, Obj->GetPathName());
+}
+
 void UWingTypes::ChooseShortNames(UPackage* Package)
 {
 	if (Package == nullptr) return;
-
 	ForEachObjectWithPackage(Package, [&](UObject* Obj)
 	{
 		ChooseShortName(Obj);
@@ -72,7 +84,25 @@ void UWingTypes::ChooseShortNames(UPackage* Package)
 	}, false);
 }
 
+void UWingTypes::ChooseShortName(const FAssetData &Data)
+{
+	FString AssetName = Data.AssetName.ToString();
+	FString PackageName = Data.PackageName.ToString();
 
+	FTopLevelAssetPath ClassPath = Data.AssetClassPath;
+	if (ClassPath == UBlueprint::StaticClass()->GetClassPathName())
+	{
+		// Blueprint: the generated class is AssetName_C
+		FString ObjectPath = FString::Printf(TEXT("%s.%s_C"), *PackageName, *AssetName);
+		ChooseShortName(AssetName, ObjectPath);
+	}
+	else if (ClassPath == UUserDefinedStruct::StaticClass()->GetClassPathName() ||
+			 ClassPath == UUserDefinedEnum::StaticClass()->GetClassPathName())
+	{
+		FString ObjectPath = FString::Printf(TEXT("%s.%s"), *PackageName, *AssetName);
+		ChooseShortName(AssetName, ObjectPath);
+	}
+}
 
 // ---------------------------------------------------------------------------
 // TypeToText
@@ -212,11 +242,26 @@ void UWingTypes::Initialize(FSubsystemCollectionBase& Collection)
 	for (UPackage* Pkg : Packages)
 		ChooseShortNames(Pkg);
 
+	// Scan the asset registry for unloaded assets.
+	IAssetRegistry& AssetRegistry = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry")).Get();
+	TArray<FAssetData> AllAssets;
+	AssetRegistry.GetAllAssets(AllAssets, true);
+	for (const FAssetData& Data : AllAssets)
+		ChooseShortName(Data);
+
+	// Register for future asset discoveries.
+	AssetRegistry.OnAssetAdded().AddUObject(this, &UWingTypes::ChooseShortName);
+
 	UE_LOG(LogTemp, Display, TEXT("WingTypes: Registered %d types"), ShortToPath.Num());
 }
 
 void UWingTypes::Deinitialize()
 {
+	if (FModuleManager::Get().IsModuleLoaded(TEXT("AssetRegistry")))
+	{
+		IAssetRegistry& AssetRegistry = FModuleManager::GetModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry")).Get();
+		AssetRegistry.OnAssetAdded().RemoveAll(this);
+	}
 	Super::Deinitialize();
 }
 
@@ -259,30 +304,43 @@ void UWingTypes::Tokenize(const FString& Input)
 }
 
 // ---------------------------------------------------------------------------
-// Path to Object Conversion
+// Short Name Resolution
 // ---------------------------------------------------------------------------
 
-bool UWingTypes::ResolvePath(const FString &Name, const FString &Path, FEdGraphPinType &OutType)
+bool UWingTypes::ResolveShortName(const FString &Name, FEdGraphPinType &OutType)
 {
-	// Load the object.
+	FString *Path = ShortToPath.Find(Name.ToLower());
-	UObject* Obj = LoadObject<UObject>(nullptr, *Path);
+	if (!Path)
-	if (!Obj)
 	{
-		Error = FString::Printf(TEXT("Failed to load type '%s' at path '%s'"), *Name, *Path);
+		Error = FString::Printf(TEXT("Unrecognized type '%s'"), *Name);
 		return false;
 	}
 
-	// If it's a blueprint, use its generated class.
+	// Primitives (int, float, etc.) are registered with the path "PRIMITIVE".
-	if (UBlueprint* BP = Cast<UBlueprint>(Obj))
+	if (*Path == TEXT("PRIMITIVE"))
 	{
-		Obj = BP->GeneratedClass;
+		OutType.PinCategory = FName(*Name);
-		if (!Obj)
+		if ((OutType.PinCategory == UEdGraphSchema_K2::PC_Double) ||
+			(OutType.PinCategory == UEdGraphSchema_K2::PC_Float))
 		{
-			Error = FString::Printf(TEXT("Blueprint '%s' has no generated class"), *Name);
+			OutType.PinSubCategory = OutType.PinCategory;
-			return false;
+			OutType.PinCategory = UEdGraphSchema_K2::PC_Real;
 		}
+		return true;
 	}
 
+	// Load the object.
+	UObject* Obj = LoadObject<UObject>(nullptr, **Path);
+	if (!Obj)
+	{
+		Error = FString::Printf(TEXT("Failed to load type '%s' at path '%s'"), *Name, **Path);
+		return false;
+	}
+
+	// The short name registry only contains UClass, UScriptStruct, and UEnum.
+	checkf(Cast<UClass>(Obj) || Cast<UScriptStruct>(Obj) || Cast<UEnum>(Obj),
+		TEXT("Short name '%s' resolved to unexpected type '%s'"), *Name, *Obj->GetClass()->GetName());
+
 	// Determine the category from the object type.
 	if (Cast<UScriptStruct>(Obj))
 	{
@@ -295,15 +353,9 @@ bool UWingTypes::ResolvePath(const FString &Name, const FString &Path, FEdGraphP
 		else
 			OutType.PinCategory = UEdGraphSchema_K2::PC_Object;
 	}
-	else if (Cast<UEnum>(Obj))
-	{
-		OutType.PinCategory = UEdGraphSchema_K2::PC_Byte;
-	}
 	else
 	{
-		// This really shouldn't happen.
+		OutType.PinCategory = UEdGraphSchema_K2::PC_Byte;
-		Error = FString::Printf(TEXT("'%s' is not a struct, class, enum, or interface"), *Name);
-		return false;
 	}
 
 	OutType.PinSubCategoryObject = Obj;
@@ -370,28 +422,7 @@ bool UWingTypes::ParsePlainIdentifier(FEdGraphPinType& OutType)
 		return false;
 	}
 	FString Name = Tokens[Cursor++];
-	FString *Path = ShortToPath.Find(Name.ToLower());
+	return ResolveShortName(Name, OutType);
-	if (Path == nullptr)
-	{
-		Error = TEXT("Unrecognized Type");
-		return false;
-	}
-
-	if (*Path == TEXT("PRIMITIVE"))
-	{
-		OutType.PinCategory = FName(*Name);
-		if ((OutType.PinCategory == UEdGraphSchema_K2::PC_Double) ||
-			(OutType.PinCategory == UEdGraphSchema_K2::PC_Float))
-		{
-			OutType.PinSubCategory = OutType.PinCategory;
-			OutType.PinCategory = UEdGraphSchema_K2::PC_Real;
-		}
-		return true;
-	}
-	else
-	{
-		return ResolvePath(Name, *Path, OutType);
-	}
 }
 
 bool UWingTypes::ParseWrapped(FName Wrapper, FEdGraphPinType& OutType)

Plugins/UEWingman/Source/UEWingman/Public/WingTypes.h

@@ -46,7 +46,7 @@ public:
 private:
 	FString TypeToTextInner(FName Category, FName SubCategory, UObject* SubCategoryObject);
 
-	bool ResolvePath(const FString &Name, const FString &Path, FEdGraphPinType &OutType);
+	bool ResolveShortName(const FString &Name, FEdGraphPinType &OutType);
 
 	// Get the object's name, not including the _C generated
 	// by the blueprint compiler for generated classes.
@@ -56,15 +56,27 @@ private:
 	// The value stored in the map is just "PRIMITIVE".
 	void ReserveShortName(FName Name);
 
-	// Chooses a short name for the specified type, making sure not to
+	// Choose a short name for the type at the specified
-	// choose the same name for two classes.  If the object already has
+	// full object path. Never reuses a short name, so every
-	// a short name, return it.  If it's not a class, struct, enum, or
+	// type will have a unique name, at least as long as the
-	// interface, return empty string.
+	// editor is up. The path must be the complete path of
+	// an object, including the .FOO_C at the end.  The name
+	// chosen will either be the proposed name, or some
+	// small variation of that.
+	FString ChooseShortName(const FString &Proposed, const FString &FullObjectPath);
+
+	// Chooses a short name for the specified type.  If the 
+	// object is not a class, struct, interface, or enum, returns
+	// null string.
 	FString ChooseShortName(const UObject* Obj);
 	
 	// Choose short names for every type in the package.
 	void ChooseShortNames(UPackage* Package);
 
+	// Choose a name for the asset's primary object, without
+	// loading the object.
+	void ChooseShortName(const FAssetData &Data);
+
 	// Short name registry: bidirectional mapping between short names and full paths.
 	TMap<FString, FString> ShortToPath;  // e.g. "vector" -> "/Script/CoreUObject.Vector"
 	TMap<FString, FString> PathToShort;  // e.g. "/Script/CoreUObject.Vector" -> "Vector"

Plugins/UEWingman/Source/UEWingman/Public/WingUtils.h

@@ -98,7 +98,7 @@ public:
 	}
 
 	template<typename T>
-	T* FindExactlyOneNamed(const FString &Name, const TArray<T*> &Array)
+	static T* FindExactlyOneNamed(const FString &Name, const TArray<T*> &Array)
 	{
 		int Count = 0;
 		T* Result = nullptr;
@@ -108,7 +108,7 @@ public:
 	}
 
 	template<typename T>
-	bool FindExactlyNoneNamed(const FString &Name, const TArray<T*> &Array)
+	static bool FindExactlyNoneNamed(const FString &Name, const TArray<T*> &Array)
 	{
 		for (T* Elt: Array) if (Identifies(Name, Elt))
 		{
@@ -199,10 +199,10 @@ public:
 	static void FormatCommandHelp(UClass* HandlerClass);
 
 	// ----- Common Error Reporting -----
-	bool CheckExactlyOneNamed(int Count, const FString &Kind, const FString &Name);
+	static bool CheckExactlyOneNamed(int Count, const FString &Kind, const FString &Name);
-	bool CheckExactlyOneNamed(int Count, UClass *Class, const FString &Name);
+	static bool CheckExactlyOneNamed(int Count, UClass *Class, const FString &Name);
-	bool CheckExactlyNoneNamed(int Count, const FString &Kind, const FString &Name);
+	static bool CheckExactlyNoneNamed(int Count, const FString &Kind, const FString &Name);
-	bool CheckExactlyNoneNamed(int Count, UClass *Class, const FString &Name);
+	static bool CheckExactlyNoneNamed(int Count, UClass *Class, const FString &Name);
 
 private:
 	static void AppendNumericSuffix(FString &Name, int32 N);