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