Differential D4535 Diff 14226 source/blender/blenlib/intern/BLI_kdtree_nd.c

Changeset View

Standalone View

source/blender/blenlib/intern/BLI_kdtree_nd.c

This file was added.

				/*
				* This program is free software; you can redistribute it and/or
				* modify it under the terms of the GNU General Public License
				* as published by the Free Software Foundation; either version 2
				* of the License, or (at your option) any later version.
				*
				* This program is distributed in the hope that it will be useful,
				* but WITHOUT ANY WARRANTY; without even the implied warranty of
				* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
				* GNU General Public License for more details.
				*
				* You should have received a copy of the GNU General Public License
				* along with this program; if not, write to the Free Software Foundation,
				* Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
				*/

				/** \file
				* \ingroup bli
				*/

				#include "MEM_guardedalloc.h"

				#include "BLI_math.h"
				#include "BLI_kdtree_nd.h"
				#include "BLI_utildefines.h"
				#include "BLI_strict_flags.h"

				BLI_INLINE float len_squared_vnvn(const float v0[], const float v1[], const uint dims)
				{
				float d = 0.0;
				for (uint j = 0; j < dims; j++) {
				d += SQUARE(v0[j] - v1[j]);
				}
				return d;
				}

				typedef struct KDTreeNode_head {
				uint left, right;
				const float *co;
				int index;
				} KDTreeNode_head;

				typedef struct KDTreeNode {
				uint left, right;
				const float *co;
				int index;
				uint d; /* range is only (0-2) */
				} KDTreeNode;

				struct KDTreeND {
				KDTreeNode *nodes;
				uint totnode;
				uint root;
				#ifdef DEBUG
				bool is_balanced; /* ensure we call balance first */
				uint maxsize; /* max size of the tree */
				#endif
				uint dims;
				float *coords;
				};

				#define KD_STACK_INIT 100 /* initial size for array (on the stack) */
				#define KD_NEAR_ALLOC_INC 100 /* alloc increment for collecting nearest */
				#define KD_FOUND_ALLOC_INC 50 /* alloc increment for collecting nearest */

				#define KD_NODE_UNSET ((uint)-1)

				/** When set we know all values are unbalanced, otherwise clear them when re-balancing: see T62210. */
				#define KD_NODE_ROOT_IS_INIT ((uint)-2)

				/**
				* Creates or free a kdtree
				*/
				KDTreeND *BLI_kdtree_nd_new(uint maxsize, uint dims)
				{
				KDTreeND *tree;

				tree = MEM_mallocN(sizeof(KDTreeND), "KDTreeND");
				tree->nodes = MEM_mallocN(sizeof(KDTreeNode) * maxsize, "KDTreeNode");
				tree->coords = MEM_mallocN(sizeof(float) * dims * maxsize, "KDTreeNode coords");
				tree->totnode = 0;
				tree->root = KD_NODE_ROOT_IS_INIT;

				#ifdef DEBUG
				tree->is_balanced = false;
				tree->maxsize = maxsize;
				#endif
				tree->dims = dims;

				return tree;
				}

				void BLI_kdtree_nd_free(KDTreeND *tree)
				{
				if (tree) {
				MEM_freeN(tree->nodes);
				MEM_freeN(tree->coords);
				MEM_freeN(tree);
				}
				}

				/**
				* Construction: first insert points, then call balance. Normal is optional.
				*/
				void BLI_kdtree_nd_insert(KDTreeND tree, int index, const float co)
				{
				float node_co = &tree->coords[tree->totnode tree->dims];
				KDTreeNode *node = &tree->nodes[tree->totnode++];

				#ifdef DEBUG
				BLI_assert(tree->totnode <= tree->maxsize);
				#endif

				/* note, array isn't calloc'd,
				* need to initialize all struct members */

				node->left = node->right = KD_NODE_UNSET;
				memcpy(node_co, co, sizeof(node_co) tree->dims);
				node->co = node_co;
				node->index = index;
				node->d = 0;

				#ifdef DEBUG
				tree->is_balanced = false;
				#endif
				}

				static uint kdtree_balance(KDTreeNode *nodes, uint totnode, uint axis, const uint ofs, const uint dims)
				{
				KDTreeNode *node;
				float co;
				uint left, right, median, i, j;

				if (totnode <= 0) {
				return KD_NODE_UNSET;
				}
				else if (totnode == 1) {
				return 0 + ofs;
				}

				/* quicksort style sorting around median */
				left = 0;
				right = totnode - 1;
				median = totnode / 2;

				while (right > left) {
				co = nodes[right].co[axis];
				i = left - 1;
				j = right;

				while (1) {
				while (nodes[++i].co[axis] < co) { /* pass */ }
				while (nodes[--j].co[axis] > co && j > left) { /* pass */ }

				if (i >= j) {
				break;
				}

				SWAP(KDTreeNode_head, (KDTreeNode_head )&nodes[i], (KDTreeNode_head )&nodes[j]);
				}

				SWAP(KDTreeNode_head, (KDTreeNode_head )&nodes[i], (KDTreeNode_head )&nodes[right]);
				if (i >= median) {
				right = i - 1;
				}
				if (i <= median) {
				left = i + 1;
				}
				}

				/* set node and sort subnodes */
				node = &nodes[median];
				node->d = axis;
				axis = (axis + 1) % dims;
				node->left = kdtree_balance(nodes, median, axis, ofs, dims);
				node->right = kdtree_balance(nodes + median + 1, (totnode - (median + 1)), axis, (median + 1) + ofs, dims);

				return median + ofs;
				}

				void BLI_kdtree_nd_balance(KDTreeND *tree)
				{
				if (tree->root != KD_NODE_ROOT_IS_INIT) {
				for (uint i = 0; i < tree->totnode; i++) {
				tree->nodes[i].left = KD_NODE_UNSET;
				tree->nodes[i].right = KD_NODE_UNSET;
				}
				}

				tree->root = kdtree_balance(tree->nodes, tree->totnode, 0, 0, tree->dims);

				#ifdef DEBUG
				tree->is_balanced = true;
				#endif
				}

				static uint realloc_nodes(uint stack, uint *totstack, const bool is_alloc)
				{
				uint stack_new = MEM_mallocN((totstack + KD_NEAR_ALLOC_INC) * sizeof(uint), "KDTreeND.treestack");
				memcpy(stack_new, stack, totstack sizeof(uint));
				// memset(stack_new + totstack, 0, sizeof(uint) KD_NEAR_ALLOC_INC);
				if (is_alloc) {
				MEM_freeN(stack);
				}
				*totstack += KD_NEAR_ALLOC_INC;
				return stack_new;
				}

				/**
				* Find nearest returns index, and -1 if no node is found.
				*/
				int BLI_kdtree_nd_find_nearest(
				const KDTreeND tree, const float co,
				KDTreeNearestND_Base *r_nearest)
				{
				const KDTreeNode *nodes = tree->nodes;
				const KDTreeNode root, min_node;
				uint *stack, defaultstack[KD_STACK_INIT];
				float min_dist, cur_dist;
				uint totstack, cur = 0;

				#ifdef DEBUG
				BLI_assert(tree->is_balanced == true);
				#endif

				if (UNLIKELY(tree->root == KD_NODE_UNSET)) {
				return -1;
				}

				stack = defaultstack;
				totstack = KD_STACK_INIT;

				root = &nodes[tree->root];
				min_node = root;
				min_dist = len_squared_vnvn(root->co, co, tree->dims);

				if (co[root->d] < root->co[root->d]) {
				if (root->right != KD_NODE_UNSET) {
				stack[cur++] = root->right;
				}
				if (root->left != KD_NODE_UNSET) {
				stack[cur++] = root->left;
				}
				}
				else {
				if (root->left != KD_NODE_UNSET) {
				stack[cur++] = root->left;
				}
				if (root->right != KD_NODE_UNSET) {
				stack[cur++] = root->right;
				}
				}

				while (cur--) {
				const KDTreeNode *node = &nodes[stack[cur]];

				cur_dist = node->co[node->d] - co[node->d];

				if (cur_dist < 0.0f) {
				cur_dist = -cur_dist * cur_dist;

				if (-cur_dist < min_dist) {
				cur_dist = len_squared_vnvn(node->co, co, tree->dims);
				if (cur_dist < min_dist) {
				min_dist = cur_dist;
				min_node = node;
				}
				if (node->left != KD_NODE_UNSET) {
				stack[cur++] = node->left;
				}
				}
				if (node->right != KD_NODE_UNSET) {
				stack[cur++] = node->right;
				}
				}
				else {
				cur_dist = cur_dist * cur_dist;

				if (cur_dist < min_dist) {
				cur_dist = len_squared_vnvn(node->co, co, tree->dims);
				if (cur_dist < min_dist) {
				min_dist = cur_dist;
				min_node = node;
				}
				if (node->right != KD_NODE_UNSET) {
				stack[cur++] = node->right;
				}
				}
				if (node->left != KD_NODE_UNSET) {
				stack[cur++] = node->left;
				}
				}
				if (UNLIKELY(cur + tree->dims > totstack)) {
				stack = realloc_nodes(stack, &totstack, defaultstack != stack);
				}
				}

				if (r_nearest) {
				r_nearest->index = min_node->index;
				r_nearest->dist = sqrtf(min_dist);
				memcpy(r_nearest->co, min_node->co, sizeof(float) * tree->dims);
				}

				if (stack != defaultstack) {
				MEM_freeN(stack);
				}

				return min_node->index;
				}