Add: Add the library function "Collect From Kdtree (Box)" (#9)

Co-authored-by: 902D9 <[email protected]>

Author: 902D9

Kdtree/Source/Kdtree/Private/AsyncKdtreeBPLibrary.cpp

@@ -155,3 +155,75 @@ void UAsyncKdtreeBPLibrary::CollectFromKdtreeAsync(const UObject* WorldContextOb
 		}
 	}
 }
+
+struct FCollectFromKdtreeBoxTaskParams
+{
+	const FKdtree* Tree;
+	FBox Box;
+	TArray<int>* Indices;
+	TArray<FVector>* Data;
+};
+
+class FCollectFromKdtreeBoxTask : public FNonAbandonableTask
+{
+public:
+	FCollectFromKdtreeBoxTask(const FCollectFromKdtreeBoxTaskParams& InParams) : Params(InParams)
+	{
+	}
+
+	void DoWork()
+	{
+		KdtreeInternal::CollectFromKdtree(Params.Tree->Internal, Params.Box, Params.Indices);
+		for (int Index = 0; Index < Params.Indices->Num(); ++Index)
+		{
+			Params.Data->Add(Params.Tree->Internal.Data[(*Params.Indices)[Index]]);
+		}
+	}
+
+	FORCEINLINE TStatId GetStatId() const
+	{
+		RETURN_QUICK_DECLARE_CYCLE_STAT(FCollectFromKdtreeBoxTask, STATGROUP_ThreadPoolAsyncTasks);
+	}
+
+private:
+	FCollectFromKdtreeBoxTaskParams Params;
+};
+
+class FCollectFromKdtreeBoxAction : public FPendingLatentAction
+{
+public:
+	FLatentActionInfo LatentInfo;
+	FAsyncTask<FCollectFromKdtreeBoxTask>* Task;
+
+	FCollectFromKdtreeBoxAction(const FLatentActionInfo& InLatentInfo, const FKdtree* Tree, const FBox Box,
+		TArray<int>* Indices, TArray<FVector>* Data)
+		: LatentInfo(InLatentInfo), Task(nullptr)
+	{
+		FCollectFromKdtreeBoxTaskParams Params;
+		Params.Tree = Tree;
+		Params.Box = Box;
+		Params.Indices = Indices;
+		Params.Data = Data;
+		Task = new FAsyncTask<FCollectFromKdtreeBoxTask>(Params);
+		Task->StartBackgroundTask();
+	}
+
+	void UpdateOperation(FLatentResponse& Response) override
+	{
+		Response.FinishAndTriggerIf(Task->IsDone(), LatentInfo.ExecutionFunction, LatentInfo.Linkage, LatentInfo.CallbackTarget);
+	}
+};
+
+void UAsyncKdtreeBPLibrary::CollectFromKdtreeAsyncBox(const UObject* WorldContextObject, const FKdtree& Tree, const FBox Box,
+	TArray<int>& Indices, TArray<FVector>& Data, FLatentActionInfo LatentInfo)
+{
+	if (UWorld* World = GEngine->GetWorldFromContextObject(WorldContextObject, EGetWorldErrorMode::LogAndReturnNull))
+	{
+		FLatentActionManager& LatentManager = World->GetLatentActionManager();
+		if (LatentManager.FindExistingAction<FCollectFromKdtreeBoxAction>(LatentInfo.CallbackTarget, LatentInfo.UUID) == nullptr)
+		{
+			FCollectFromKdtreeBoxAction* NewAction = new FCollectFromKdtreeBoxAction(LatentInfo, &Tree, Box, &Indices, &Data);
+			LatentManager.AddNewAction(LatentInfo.CallbackTarget, LatentInfo.UUID, NewAction);
+		}
+	}
+}

Kdtree/Source/Kdtree/Private/KdtreeBPLibrary.cpp

@@ -36,6 +36,16 @@ void UKdtreeBPLibrary::CollectFromKdtree(
 	}
 }
 
+void UKdtreeBPLibrary::CollectFromKdtreeBox(const FKdtree& Tree, const FBox Box, TArray<int>& Indices,
+	TArray<FVector>& Data)
+{
+	KdtreeInternal::CollectFromKdtree(Tree.Internal, Box, &Indices);
+	for (int Index = 0; Index < Indices.Num(); ++Index)
+	{
+		Data.Add(Tree.Internal.Data[Indices[Index]]);
+	}
+}
+
 void UKdtreeBPLibrary::ValidateKdtree(const FKdtree& Tree)
 {
 	KdtreeInternal::ValidateKdtree(Tree.Internal);

Kdtree/Source/Kdtree/Private/KdtreeInternal.cpp

@@ -234,7 +234,32 @@ void CollectFromKdtree(
 		}
 	}
 }
-}	 // namespace
+
+void CollectFromKdtree(
+	const FKdtreeInternal& Tree, const FKdtreeNode* Node, const FBox& Box, TArray<int>* Result)
+{
+	if (Node == nullptr)
+	{
+		return;
+	}
+
+	const FVector& Current = Tree.Data[Node->Index];
+	if (Box.IsInsideOrOn(Current))
+	{
+		Result->Add(Node->Index);
+	}
+
+	int Axis = Node->Axis;
+	if (Box.Min[Axis] < Current[Axis])
+	{
+		CollectFromKdtree(Tree, Node->ChildLeft, Box, Result);
+	}
+	if (Box.Max[Axis] > Current[Axis])
+	{
+		CollectFromKdtree(Tree, Node->ChildRight, Box, Result);
+	}
+}
+} // namespace
 
 void BuildKdtree(FKdtreeInternal* Tree, const TArray<FVector>& Data)
 {
@@ -262,6 +287,11 @@ void CollectFromKdtree(const FKdtreeInternal& Tree, const FVector& Center, float
 	CollectFromKdtree(Tree, Tree.Root, Center, Radius, Result);
 }
 
+void CollectFromKdtree(const FKdtreeInternal& Tree, const FBox& Box, TArray<int>* Result)
+{
+	CollectFromKdtree(Tree, Tree.Root, Box, Result);
+}
+
 void ValidateKdtree(const FKdtreeInternal& Tree)
 {
 	ValidateKdtree(Tree, Tree.Root, 0);

Kdtree/Source/Kdtree/Private/KdtreeInternal.h

@@ -24,6 +24,7 @@ struct FKdtreeNode
 void BuildKdtree(FKdtreeInternal* Tree, const TArray<FVector>& Data);
 void ClearKdtree(FKdtreeInternal* Tree);
 void CollectFromKdtree(const FKdtreeInternal& Tree, const FVector& Center, float Radius, TArray<int>* Result);
+void CollectFromKdtree(const FKdtreeInternal& Tree, const FBox& Box, TArray<int>* Result);
 void ValidateKdtree(const FKdtreeInternal& Tree);
 void DumpKdTree(const FKdtreeInternal& Tree);
 }	 // namespace KdtreeInternal

Kdtree/Source/Kdtree/Public/AsyncKdtreeBPLibrary.h

@@ -32,7 +32,14 @@ class KDTREE_API UAsyncKdtreeBPLibrary : public UBlueprintFunctionLibrary
 	UFUNCTION(BlueprintCallable,
 		meta = (WorldContextObject = "WorldContextObject", Latent, LatentInfo = "LatentInfo", HidePin = "WorldContextObject",
 			DefaultToSelf = "WorldContextObject"),
-		Category = "SpacialDataStructure|kd-tree")
+		Category = "SpacialDataStructure|kd-tree", DisplayName = "Collect From Kdtree Async (Sphere)")
 	static void CollectFromKdtreeAsync(const UObject* WorldContextObject, const FKdtree& Tree, const FVector Center, float Radius,
 		TArray<int>& Indices, TArray<FVector>& Data, FLatentActionInfo LatentInfo);
+
+	UFUNCTION(BlueprintCallable,
+	meta = (WorldContextObject = "WorldContextObject", Latent, LatentInfo = "LatentInfo", HidePin = "WorldContextObject",
+		DefaultToSelf = "WorldContextObject"),
+	Category = "SpacialDataStructure|kd-tree", DisplayName= "Collect From Kdtree Async (Box)")
+	static void CollectFromKdtreeAsyncBox(const UObject* WorldContextObject, const FKdtree& Tree, const FBox Box,
+		TArray<int>& Indices, TArray<FVector>& Data, FLatentActionInfo LatentInfo);
 };

Kdtree/Source/Kdtree/Public/KdtreeBPLibrary.h

@@ -28,10 +28,14 @@ class KDTREE_API UKdtreeBPLibrary : public UBlueprintFunctionLibrary
 	UFUNCTION(BluePrintCallable, Category = "SpacialDataStructure|kd-tree")
 	static void ClearKdtree(UPARAM(ref) FKdtree& Tree);
 
-	UFUNCTION(BluePrintCallable, Category = "SpacialDataStructure|kd-tree")
+	UFUNCTION(BlueprintCallable, Category = "SpacialDataStructure|kd-tree", DisplayName="Collect From Kdtree (Sphere)")
 	static void CollectFromKdtree(
 		const FKdtree& Tree, const FVector Center, float Radius, TArray<int>& Indices, TArray<FVector>& Data);
 
+	UFUNCTION(BlueprintCallable, Category = "SpacialDataStructure|kd-tree", DisplayName="Collect From Kdtree (Box)")
+	static void CollectFromKdtreeBox(
+		const FKdtree& Tree, const FBox Box, TArray<int>& Indices, TArray<FVector>& Data);
+
 	UFUNCTION(BluePrintCallable, Category = "SpacialDataStructure|kd-tree")
 	static void ValidateKdtree(const FKdtree& Tree);