#56 crash fix + optimizations

MothCocoon · Oct 9, 2021 · 9e6b57d · 9e6b57d
1 parent 5ee34b8
commit 9e6b57d
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 23 deletions.
diff --git a/Source/FlowEditor/Private/Graph/FlowGraphSchema.cpp b/Source/FlowEditor/Private/Graph/FlowGraphSchema.cpp
@@ -39,7 +39,7 @@ void UFlowGraphSchema::SubscribeToAssetChanges()
 	const FAssetRegistryModule& AssetRegistry = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(AssetRegistryConstants::ModuleName);
 	AssetRegistry.Get().OnFilesLoaded().AddStatic(&UFlowGraphSchema::GatherFlowNodes);
 	AssetRegistry.Get().OnAssetAdded().AddStatic(&UFlowGraphSchema::OnAssetAdded);
-	AssetRegistry.Get().OnAssetRemoved().AddStatic(&UFlowGraphSchema::RemoveAsset);
+	AssetRegistry.Get().OnAssetRemoved().AddStatic(&UFlowGraphSchema::OnAssetRemoved);
 
 	FCoreUObjectDelegates::ReloadCompleteDelegate.AddStatic(&UFlowGraphSchema::OnHotReload);
 
@@ -203,14 +203,16 @@ void UFlowGraphSchema::GetFlowNodeActions(FGraphActionMenuBuilder& ActionMenuBui
 	TArray<UFlowNode*> FlowNodes;
 	FlowNodes.Reserve(NativeFlowNodes.Num() + BlueprintFlowNodes.Num());
 
-	for (UClass* FlowNodeClass : NativeFlowNodes)
+	for (const UClass* FlowNodeClass : NativeFlowNodes)
 	{
 		FlowNodes.Emplace(FlowNodeClass->GetDefaultObject<UFlowNode>());
 	}
 	for (const TPair<FName, FAssetData>& AssetData : BlueprintFlowNodes)
 	{
-		UBlueprint* Blueprint = GetNodeBlueprint(AssetData.Value);
-		FlowNodes.Emplace(Blueprint->GeneratedClass->GetDefaultObject<UFlowNode>());
+		if (const UBlueprint* Blueprint = GetNodeBlueprint(AssetData.Value))
+		{
+			FlowNodes.Emplace(Blueprint->GeneratedClass->GetDefaultObject<UFlowNode>());
+		}
 	}
 
 	for (const UFlowNode* FlowNode : FlowNodes)
@@ -252,34 +254,36 @@ void UFlowGraphSchema::GatherFlowNodes()
 	// collect C++ nodes once per editor session
 	if (NativeFlowNodes.Num() == 0)
 	{
-		for (TObjectIterator<UClass> It; It; ++It)
+		TArray<UClass*> FlowNodes;
+		GetDerivedClasses(UFlowNode::StaticClass(), FlowNodes);
+		for (UClass* Class : FlowNodes)
 		{
-			if (It->IsChildOf(UFlowNode::StaticClass()))
+			if (Class->ClassGeneratedBy == nullptr && IsFlowNodePlaceable(Class))
 			{
-				if (It->ClassGeneratedBy == nullptr && IsFlowNodePlaceable(*It))
-				{
-					NativeFlowNodes.Emplace(*It);
+				NativeFlowNodes.Emplace(Class);
 
-					const UFlowNode* DefaultObject = It->GetDefaultObject<UFlowNode>();
-					UnsortedCategories.Emplace(DefaultObject->GetNodeCategory());
-				}
+				const UFlowNode* DefaultObject = Class->GetDefaultObject<UFlowNode>();
+				UnsortedCategories.Emplace(DefaultObject->GetNodeCategory());
 			}
-			else if (It->IsChildOf(UFlowGraphNode::StaticClass()))
+		}
+
+		TArray<UClass*> GraphNodes;
+		GetDerivedClasses(UFlowGraphNode::StaticClass(), GraphNodes);
+		for (UClass* Class : GraphNodes)
+		{
+			const UFlowGraphNode* DefaultObject = Class->GetDefaultObject<UFlowGraphNode>();
+			for (UClass* AssignedClass : DefaultObject->AssignedNodeClasses)
 			{
-				const UFlowGraphNode* DefaultObject = It->GetDefaultObject<UFlowGraphNode>();
-				for (UClass* AssignedClass : DefaultObject->AssignedNodeClasses)
+				if (AssignedClass->IsChildOf(UFlowNode::StaticClass()))
 				{
-					if (AssignedClass->IsChildOf(UFlowNode::StaticClass()))
-					{
-						AssignedGraphNodeClasses.Emplace(AssignedClass, *It);
-					}
+					AssignedGraphNodeClasses.Emplace(AssignedClass, Class);
 				}
 			}
 		}
 	}
 
 	// retrieve all blueprint nodes
-	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(AssetRegistryConstants::ModuleName);
+	const FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(AssetRegistryConstants::ModuleName);
 
 	FARFilter Filter;
 	Filter.ClassNames.Add(UBlueprint::StaticClass()->GetFName());
@@ -310,7 +314,7 @@ void UFlowGraphSchema::AddAsset(const FAssetData& AssetData, const bool bBatch)
 {
 	if (!BlueprintFlowNodes.Contains(AssetData.PackageName))
 	{
-		FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(AssetRegistryConstants::ModuleName);
+		const FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(AssetRegistryConstants::ModuleName);
 		if (AssetRegistryModule.Get().IsLoadingAssets())
 		{
 			return;
@@ -350,7 +354,7 @@ void UFlowGraphSchema::AddAsset(const FAssetData& AssetData, const bool bBatch)
 	}
 }
 
-void UFlowGraphSchema::RemoveAsset(const FAssetData& AssetData)
+void UFlowGraphSchema::OnAssetRemoved(const FAssetData& AssetData)
 {
 	if (BlueprintFlowNodes.Contains(AssetData.PackageName))
 	{

diff --git a/Source/FlowEditor/Public/Graph/FlowGraphSchema.h b/Source/FlowEditor/Public/Graph/FlowGraphSchema.h
@@ -49,7 +49,8 @@ class FLOWEDITOR_API UFlowGraphSchema : public UEdGraphSchema
 
 	static void OnAssetAdded(const FAssetData& AssetData);
 	static void AddAsset(const FAssetData& AssetData, const bool bBatch);
-	static void RemoveAsset(const FAssetData& AssetData);
+	static void OnAssetRemoved(const FAssetData& AssetData);
+
 	static void RefreshNodeList();
 
 public: