BLI_kdtree.c

Go to the documentation of this file.

/*
 * ***** BEGIN GPL LICENSE BLOCK *****
 *
 * 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.
 *
 * The Original Code is Copyright (C) 2001-2002 by NaN Holding BV.
 * All rights reserved.
 *
 * The Original Code is: none of this file.
 *
 * Contributor(s): Janne Karhu
 *                 Brecht Van Lommel
 *
 * ***** END GPL LICENSE BLOCK *****
 */

#include "MEM_guardedalloc.h"

#include "BLI_math.h"
#include "BLI_kdtree.h"

#ifndef SWAP
#define SWAP(type, a, b) { type sw_ap; sw_ap=(a); (a)=(b); (b)=sw_ap; }
#endif

typedef struct KDTreeNode {
    struct KDTreeNode *left, *right;
    float co[3], nor[3];
    int index;
    short d;
} KDTreeNode;

struct KDTree {
    KDTreeNode *nodes;
    int totnode;
    KDTreeNode *root;
};

KDTree *BLI_kdtree_new(int maxsize)
{
    KDTree *tree;

    tree= MEM_callocN(sizeof(KDTree), "KDTree");
    tree->nodes= MEM_callocN(sizeof(KDTreeNode)*maxsize, "KDTreeNode");
    tree->totnode= 0;

    return tree;
}

void BLI_kdtree_free(KDTree *tree)
{
    if(tree) {
        MEM_freeN(tree->nodes);
        MEM_freeN(tree);
    }
}

void BLI_kdtree_insert(KDTree *tree, int index, float *co, float *nor)
{
    KDTreeNode *node= &tree->nodes[tree->totnode++];

    node->index= index;
    copy_v3_v3(node->co, co);
    if(nor) copy_v3_v3(node->nor, nor);
}

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

    if(totnode <= 0)
        return NULL;
    else if(totnode == 1)
        return nodes;
    
    /* 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);
            while(nodes[--j].co[axis] > co && j>left);

            if(i >= j) break;
            SWAP(KDTreeNode, nodes[i], nodes[j]);
        }

        SWAP(KDTreeNode, nodes[i], 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;
    node->left= kdtree_balance(nodes, median, (axis+1)%3);
    node->right= kdtree_balance(nodes+median+1, (totnode-(median+1)), (axis+1)%3);

    return node;
}

void BLI_kdtree_balance(KDTree *tree)
{
    tree->root= kdtree_balance(tree->nodes, tree->totnode, 0);
}

static float squared_distance(float *v2, float *v1, float *n1, float *n2)
{
    float d[3], dist;
    (void)n1; /* unused */

    d[0]= v2[0]-v1[0];
    d[1]= v2[1]-v1[1];
    d[2]= v2[2]-v1[2];

    dist= d[0]*d[0] + d[1]*d[1] + d[2]*d[2];

    //if(n1 && n2 && n1[0]*n2[0] + n1[1]*n2[1] + n1[2]*n2[2] < 0.0f)
    if(n2 && d[0]*n2[0] + d[1]*n2[1] + d[2]*n2[2] < 0.0f)
        dist *= 10.0f;

    return dist;
}

int BLI_kdtree_find_nearest(KDTree *tree, float *co, float *nor, KDTreeNearest *nearest)
{
    KDTreeNode *root, *node, *min_node;
    KDTreeNode **stack, *defaultstack[100];
    float min_dist, cur_dist;
    int totstack, cur=0;

    if(!tree->root)
        return -1;

    stack= defaultstack;
    totstack= 100;

    root= tree->root;
    min_node= root;
    min_dist= squared_distance(root->co,co,root->nor,nor);

    if(co[root->d] < root->co[root->d]) {
        if(root->right)
            stack[cur++]=root->right;
        if(root->left)
            stack[cur++]=root->left;
    }
    else {
        if(root->left)
            stack[cur++]=root->left;
        if(root->right)
            stack[cur++]=root->right;
    }
    
    while(cur--){
        node=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=squared_distance(node->co,co,node->nor,nor);
                if(cur_dist<min_dist){
                    min_dist=cur_dist;
                    min_node=node;
                }
                if(node->left)
                    stack[cur++]=node->left;
            }
            if(node->right)
                stack[cur++]=node->right;
        }
        else{
            cur_dist= cur_dist*cur_dist;

            if(cur_dist<min_dist){
                cur_dist=squared_distance(node->co,co,node->nor,nor);
                if(cur_dist<min_dist){
                    min_dist=cur_dist;
                    min_node=node;
                }
                if(node->right)
                    stack[cur++]=node->right;
            }
            if(node->left)
                stack[cur++]=node->left;
        }
        if(cur+3 > totstack){
            KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
            memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
            if(stack != defaultstack)
                MEM_freeN(stack);
            stack=temp;
            totstack+=100;
        }
    }

    if(nearest) {
        nearest->index= min_node->index;
        nearest->dist= sqrt(min_dist);
        copy_v3_v3(nearest->co, min_node->co);
    }

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

    return min_node->index;
}

static void add_nearest(KDTreeNearest *ptn, int *found, int n, int index, float dist, float *co)
{
    int i;

    if(*found<n) (*found)++;

    for(i=*found-1; i>0; i--) {
        if(dist >= ptn[i-1].dist)
            break;
        else
            ptn[i]= ptn[i-1];
    }

    ptn[i].index= index;
    ptn[i].dist= dist;
    copy_v3_v3(ptn[i].co, co);
}

/* finds the nearest n entries in tree to specified coordinates */
int BLI_kdtree_find_n_nearest(KDTree *tree, int n, float *co, float *nor, KDTreeNearest *nearest)
{
    KDTreeNode *root, *node= NULL;
    KDTreeNode **stack, *defaultstack[100];
    float cur_dist;
    int i, totstack, cur=0, found=0;

    if(!tree->root)
        return 0;

    stack= defaultstack;
    totstack= 100;

    root= tree->root;

    cur_dist= squared_distance(root->co,co,root->nor,nor);
    add_nearest(nearest,&found,n,root->index,cur_dist,root->co);
    
    if(co[root->d] < root->co[root->d]) {
        if(root->right)
            stack[cur++]=root->right;
        if(root->left)
            stack[cur++]=root->left;
    }
    else {
        if(root->left)
            stack[cur++]=root->left;
        if(root->right)
            stack[cur++]=root->right;
    }

    while(cur--){
        node=stack[cur];

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

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

            if(found<n || -cur_dist<nearest[found-1].dist){
                cur_dist=squared_distance(node->co,co,node->nor,nor);

                if(found<n || cur_dist<nearest[found-1].dist)
                    add_nearest(nearest,&found,n,node->index,cur_dist,node->co);

                if(node->left)
                    stack[cur++]=node->left;
            }
            if(node->right)
                stack[cur++]=node->right;
        }
        else{
            cur_dist= cur_dist*cur_dist;

            if(found<n || cur_dist<nearest[found-1].dist){
                cur_dist=squared_distance(node->co,co,node->nor,nor);
                if(found<n || cur_dist<nearest[found-1].dist)
                    add_nearest(nearest,&found,n,node->index,cur_dist,node->co);

                if(node->right)
                    stack[cur++]=node->right;
            }
            if(node->left)
                stack[cur++]=node->left;
        }
        if(cur+3 > totstack){
            KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
            memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
            if(stack != defaultstack)
                MEM_freeN(stack);
            stack=temp;
            totstack+=100;
        }
    }

    for(i=0; i<found; i++)
        nearest[i].dist= sqrt(nearest[i].dist);

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

    return found;
}

static int range_compare(const void * a, const void * b)
{
    const KDTreeNearest *kda = a;
    const KDTreeNearest *kdb = b;

    if(kda->dist < kdb->dist)
        return -1;
    else if(kda->dist > kdb->dist)
        return 1;
    else
        return 0;
}
static void add_in_range(KDTreeNearest **ptn, int found, int *totfoundstack, int index, float dist, float *co)
{
    KDTreeNearest *to;

    if(found+1 > *totfoundstack) {
        KDTreeNearest *temp=MEM_callocN((*totfoundstack+50)*sizeof(KDTreeNode), "psys_treefoundstack");
        memcpy(temp, *ptn, *totfoundstack * sizeof(KDTreeNearest));
        if(*ptn)
            MEM_freeN(*ptn);
        *ptn = temp;
        *totfoundstack+=50;
    }

    to = (*ptn) + found;

    to->index = index;
    to->dist = sqrt(dist);
    copy_v3_v3(to->co, co);
}
int BLI_kdtree_range_search(KDTree *tree, float range, float *co, float *nor, KDTreeNearest **nearest)
{
    KDTreeNode *root, *node= NULL;
    KDTreeNode **stack, *defaultstack[100];
    KDTreeNearest *foundstack=NULL;
    float range2 = range*range, dist2;
    int totstack, cur=0, found=0, totfoundstack=0;

    if(!tree || !tree->root)
        return 0;

    stack= defaultstack;
    totstack= 100;

    root= tree->root;

    if(co[root->d] + range < root->co[root->d]) {
        if(root->left)
            stack[cur++]=root->left;
    }
    else if(co[root->d] - range > root->co[root->d]) {
        if(root->right)
            stack[cur++]=root->right;
    }
    else {
        dist2 = squared_distance(root->co, co, root->nor, nor);
        if(dist2  <= range2)
            add_in_range(&foundstack, found++, &totfoundstack, root->index, dist2, root->co);

        if(root->left)
            stack[cur++]=root->left;
        if(root->right)
            stack[cur++]=root->right;
    }

    while(cur--) {
        node=stack[cur];

        if(co[node->d] + range < node->co[node->d]) {
            if(node->left)
                stack[cur++]=node->left;
        }
        else if(co[node->d] - range > node->co[node->d]) {
            if(node->right)
                stack[cur++]=node->right;
        }
        else {
            dist2 = squared_distance(node->co, co, node->nor, nor);
            if(dist2 <= range2)
                add_in_range(&foundstack, found++, &totfoundstack, node->index, dist2, node->co);

            if(node->left)
                stack[cur++]=node->left;
            if(node->right)
                stack[cur++]=node->right;
        }

        if(cur+3 > totstack){
            KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
            memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
            if(stack != defaultstack)
                MEM_freeN(stack);
            stack=temp;
            totstack+=100;
        }
    }

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

    if(found)
        qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);

    *nearest = foundstack;

    return found;
}