d48ab8037400b6987ca10b36d7e25826b8ca55a5
[blender.git] / source / blender / python / mathutils / mathutils_kdtree.c
1 /*
2  * ***** BEGIN GPL LICENSE BLOCK *****
3  *
4  * This program is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU General Public License
6  * as published by the Free Software Foundation; either version 2
7  * of the License, or (at your option) any later version.
8  *
9  * This program is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12  * GNU General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License
15  * along with this program; if not, write to the Free Software Foundation,
16  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
17  *
18  * Contributor(s): Dan Eicher, Campbell Barton
19  *
20  * ***** END GPL LICENSE BLOCK *****
21  */
22
23 /** \file blender/python/mathutils/mathutils_kdtree.c
24  *  \ingroup mathutils
25  *
26  * This file defines the 'mathutils.kdtree' module, a general purpose module to access
27  * blenders kdtree for 3d spatial lookups.
28  */
29
30 #include <Python.h>
31
32 #include "MEM_guardedalloc.h"
33
34 #include "BLI_utildefines.h"
35 #include "BLI_kdtree.h"
36
37 #include "../generic/py_capi_utils.h"
38
39 #include "mathutils.h"
40 #include "mathutils_kdtree.h"  /* own include */
41
42 #include "BLI_strict_flags.h"
43
44 typedef struct {
45         PyObject_HEAD
46         KDTree *obj;
47         unsigned int maxsize;
48         unsigned int count;
49         unsigned int count_balance;  /* size when we last balanced */
50 } PyKDTree;
51
52
53 /* -------------------------------------------------------------------- */
54 /* Utility helper functions */
55
56 static void kdtree_nearest_to_py_tuple(const KDTreeNearest *nearest, PyObject *py_retval)
57 {
58         BLI_assert(nearest->index >= 0);
59         BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);
60
61         PyTuple_SET_ITEM(py_retval, 0, Vector_CreatePyObject((float *)nearest->co, 3, Py_NEW, NULL));
62         PyTuple_SET_ITEM(py_retval, 1, PyLong_FromLong(nearest->index));
63         PyTuple_SET_ITEM(py_retval, 2, PyFloat_FromDouble(nearest->dist));
64 }
65
66 static PyObject *kdtree_nearest_to_py(const KDTreeNearest *nearest)
67 {
68         PyObject *py_retval;
69
70         py_retval = PyTuple_New(3);
71
72         kdtree_nearest_to_py_tuple(nearest, py_retval);
73
74         return py_retval;
75 }
76
77 static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest *nearest)
78 {
79         PyObject *py_retval;
80
81         py_retval = PyTuple_New(3);
82
83         if (nearest->index != -1) {
84                 kdtree_nearest_to_py_tuple(nearest, py_retval);
85         }
86         else {
87                 PyC_Tuple_Fill(py_retval, Py_None);
88         }
89
90         return py_retval;
91 }
92
93
94 /* -------------------------------------------------------------------- */
95 /* KDTree */
96
97 /* annoying since arg parsing won't check overflow */
98 #define UINT_IS_NEG(n) ((n) > INT_MAX)
99
100 static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
101 {
102         unsigned int maxsize;
103         const char *keywords[] = {"size", NULL};
104
105         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *)"I:KDTree", (char **)keywords, &maxsize)) {
106                 return -1;
107         }
108
109         if (UINT_IS_NEG(maxsize)) {
110                 PyErr_SetString(PyExc_ValueError, "negative 'size' given");
111                 return -1;
112         }
113
114         self->obj = BLI_kdtree_new(maxsize);
115         self->maxsize = maxsize;
116         self->count = 0;
117         self->count_balance = 0;
118
119         return 0;
120 }
121
122 static void PyKDTree__tp_dealloc(PyKDTree *self)
123 {
124         BLI_kdtree_free(self->obj);
125         Py_TYPE(self)->tp_free((PyObject *)self);
126 }
127
128 PyDoc_STRVAR(py_kdtree_insert_doc,
129 ".. method:: insert(index, co)\n"
130 "\n"
131 "   Insert a point into the KDTree.\n"
132 "\n"
133 "   :arg co: Point 3d position.\n"
134 "   :type co: float triplet\n"
135 "   :arg index: The index of the point.\n"
136 "   :type index: int\n"
137 );
138 static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
139 {
140         PyObject *py_co;
141         float co[3];
142         int index;
143         const char *keywords[] = {"co", "index", NULL};
144
145         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Oi:insert", (char **)keywords,
146                                          &py_co, &index))
147         {
148                 return NULL;
149         }
150
151         if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1)
152                 return NULL;
153
154         if (index < 0) {
155                 PyErr_SetString(PyExc_ValueError, "negative index given");
156                 return NULL;
157         }
158
159         if (self->count >= self->maxsize) {
160                 PyErr_SetString(PyExc_RuntimeError, "Trying to insert more items than KDTree has room for");
161                 return NULL;
162         }
163
164         BLI_kdtree_insert(self->obj, index, co, NULL);
165         self->count++;
166
167         Py_RETURN_NONE;
168 }
169
170 PyDoc_STRVAR(py_kdtree_balance_doc,
171 ".. method:: balance()\n"
172 "\n"
173 "   Balance the tree.\n"
174 );
175 static PyObject *py_kdtree_balance(PyKDTree *self)
176 {
177         BLI_kdtree_balance(self->obj);
178         self->count_balance = self->count;
179         Py_RETURN_NONE;
180 }
181
182 PyDoc_STRVAR(py_kdtree_find_doc,
183 ".. method:: find(co)\n"
184 "\n"
185 "   Find nearest point to ``co``.\n"
186 "\n"
187 "   :arg co: 3d coordinates.\n"
188 "   :type co: float triplet\n"
189 "   :return: Returns (:class:`Vector`, index, distance).\n"
190 "   :rtype: :class:`tuple`\n"
191 );
192 static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
193 {
194         PyObject *py_co;
195         float co[3];
196         KDTreeNearest nearest;
197         const char *keywords[] = {"co", NULL};
198
199         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "O:find", (char **)keywords,
200                                          &py_co))
201         {
202                 return NULL;
203         }
204
205         if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1)
206                 return NULL;
207
208         if (self->count != self->count_balance) {
209                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find()");
210                 return NULL;
211         }
212
213
214         nearest.index = -1;
215
216         BLI_kdtree_find_nearest(self->obj, co, NULL, &nearest);
217
218         return kdtree_nearest_to_py_and_check(&nearest);
219 }
220
221 PyDoc_STRVAR(py_kdtree_find_n_doc,
222 ".. method:: find_n(co, n)\n"
223 "\n"
224 "   Find nearest ``n`` points to ``co``.\n"
225 "\n"
226 "   :arg co: 3d coordinates.\n"
227 "   :type co: float triplet\n"
228 "   :arg n: Number of points to find.\n"
229 "   :type n: int\n"
230 "   :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
231 "   :rtype: :class:`list`\n"
232 );
233 static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
234 {
235         PyObject *py_list;
236         PyObject *py_co;
237         float co[3];
238         KDTreeNearest *nearest;
239         unsigned int n;
240         int i, found;
241         const char *keywords[] = {"co", "n", NULL};
242
243         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "OI:find_n", (char **)keywords,
244                                          &py_co, &n))
245         {
246                 return NULL;
247         }
248
249         if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1)
250                 return NULL;
251
252         if (UINT_IS_NEG(n)) {
253                 PyErr_SetString(PyExc_RuntimeError, "negative 'n' given");
254                 return NULL;
255         }
256
257         if (self->count != self->count_balance) {
258                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_n()");
259                 return NULL;
260         }
261
262         nearest = MEM_mallocN(sizeof(KDTreeNearest) * n, __func__);
263
264         found = BLI_kdtree_find_nearest_n(self->obj, co, NULL, nearest, n);
265
266         py_list = PyList_New(found);
267
268         for (i = 0; i < found; i++) {
269                 PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
270         }
271
272         MEM_freeN(nearest);
273
274         return py_list;
275 }
276
277 PyDoc_STRVAR(py_kdtree_find_range_doc,
278 ".. method:: find_range(co, radius)\n"
279 "\n"
280 "   Find all points within ``radius`` of ``co``.\n"
281 "\n"
282 "   :arg co: 3d coordinates.\n"
283 "   :type co: float triplet\n"
284 "   :arg radius: Distance to search for points.\n"
285 "   :type radius: float\n"
286 "   :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
287 "   :rtype: :class:`list`\n"
288 );
289 static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
290 {
291         PyObject *py_list;
292         PyObject *py_co;
293         float co[3];
294         KDTreeNearest *nearest = NULL;
295         float radius;
296         int i, found;
297
298         const char *keywords[] = {"co", "radius", NULL};
299
300         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Of:find_range", (char **)keywords,
301                                          &py_co, &radius))
302         {
303                 return NULL;
304         }
305
306         if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1)
307                 return NULL;
308
309         if (radius < 0.0f) {
310                 PyErr_SetString(PyExc_RuntimeError, "negative radius given");
311                 return NULL;
312         }
313
314         if (self->count != self->count_balance) {
315                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_range()");
316                 return NULL;
317         }
318
319         found = BLI_kdtree_range_search(self->obj, co, NULL, &nearest, radius);
320
321         py_list = PyList_New(found);
322
323         for (i = 0; i < found; i++) {
324                 PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
325         }
326
327         if (nearest) {
328                 MEM_freeN(nearest);
329         }
330
331         return py_list;
332 }
333
334
335 static PyMethodDef PyKDTree_methods[] = {
336         {"insert", (PyCFunction)py_kdtree_insert, METH_VARARGS | METH_KEYWORDS, py_kdtree_insert_doc},
337         {"balance", (PyCFunction)py_kdtree_balance, METH_NOARGS, py_kdtree_balance_doc},
338         {"find", (PyCFunction)py_kdtree_find, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_doc},
339         {"find_n", (PyCFunction)py_kdtree_find_n, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_n_doc},
340         {"find_range", (PyCFunction)py_kdtree_find_range, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_range_doc},
341         {NULL, NULL, 0, NULL}
342 };
343
344 PyDoc_STRVAR(py_KDtree_doc,
345 "KdTree(size) -> new kd-tree initialized to hold ``size`` items.\n"
346 "\n"
347 ".. note::\n"
348 "\n"
349 "   :class:`KDTree.balance` must have been called before using any of the ``find`` methods.\n"
350 );
351 PyTypeObject PyKDTree_Type = {
352         PyVarObject_HEAD_INIT(NULL, 0)
353         "KDTree",                                    /* tp_name */
354         sizeof(PyKDTree),                            /* tp_basicsize */
355         0,                                           /* tp_itemsize */
356         /* methods */
357         (destructor)PyKDTree__tp_dealloc,            /* tp_dealloc */
358         NULL,                                        /* tp_print */
359         NULL,                                        /* tp_getattr */
360         NULL,                                        /* tp_setattr */
361         NULL,                                        /* tp_compare */
362         NULL,                                        /* tp_repr */
363         NULL,                                        /* tp_as_number */
364         NULL,                                        /* tp_as_sequence */
365         NULL,                                        /* tp_as_mapping */
366         NULL,                                        /* tp_hash */
367         NULL,                                        /* tp_call */
368         NULL,                                        /* tp_str */
369         NULL,                                        /* tp_getattro */
370         NULL,                                        /* tp_setattro */
371         NULL,                                        /* tp_as_buffer */
372         Py_TPFLAGS_DEFAULT,                          /* tp_flags */
373         py_KDtree_doc,                               /* Documentation string */
374         NULL,                                        /* tp_traverse */
375         NULL,                                        /* tp_clear */
376         NULL,                                        /* tp_richcompare */
377         0,                                           /* tp_weaklistoffset */
378         NULL,                                        /* tp_iter */
379         NULL,                                        /* tp_iternext */
380         (struct PyMethodDef *)PyKDTree_methods,      /* tp_methods */
381         NULL,                                        /* tp_members */
382         NULL,                                        /* tp_getset */
383         NULL,                                        /* tp_base */
384         NULL,                                        /* tp_dict */
385         NULL,                                        /* tp_descr_get */
386         NULL,                                        /* tp_descr_set */
387         0,                                           /* tp_dictoffset */
388         (initproc)PyKDTree__tp_init,                 /* tp_init */
389         (allocfunc)PyType_GenericAlloc,              /* tp_alloc */
390         (newfunc)PyType_GenericNew,                  /* tp_new */
391         (freefunc)0,                                 /* tp_free */
392         NULL,                                        /* tp_is_gc */
393         NULL,                                        /* tp_bases */
394         NULL,                                        /* tp_mro */
395         NULL,                                        /* tp_cache */
396         NULL,                                        /* tp_subclasses */
397         NULL,                                        /* tp_weaklist */
398         (destructor) NULL                            /* tp_del */
399 };
400
401 PyDoc_STRVAR(py_kdtree_doc,
402 "Generic 3-dimentional kd-tree to perform spatial searches."
403 );
404 static struct PyModuleDef kdtree_moduledef = {
405         PyModuleDef_HEAD_INIT,
406         "mathutils.kdtree",                          /* m_name */
407         py_kdtree_doc,                               /* m_doc */
408         0,                                           /* m_size */
409         NULL,                                        /* m_methods */
410         NULL,                                        /* m_reload */
411         NULL,                                        /* m_traverse */
412         NULL,                                        /* m_clear */
413         NULL                                         /* m_free */
414 };
415
416 PyMODINIT_FUNC PyInit_mathutils_kdtree(void)
417 {
418         PyObject *m = PyModule_Create(&kdtree_moduledef);
419
420         if (m == NULL) {
421                 return NULL;
422         }
423
424         /* Register the 'KDTree' class */
425         if (PyType_Ready(&PyKDTree_Type)) {
426                 return NULL;
427         }
428         PyModule_AddObject(m, (char *)"KDTree", (PyObject *) &PyKDTree_Type);
429
430         return m;
431 }