geo_index/rtree/sort/
hilbert.rs

1use crate::indices::MutableIndices;
2use crate::r#type::IndexableNum;
3use crate::rtree::sort::{Sort, SortParams};
4
5/// An implementation of hilbert sorting.
6///
7/// The implementation is ported from the original [flatbush](https://github.com/mourner/flatbush)
8/// JavaScript library. The hilbert calculations are originally derived from [a C++
9/// implementation](https://github.com/rawrunprotected/hilbert_curves).
10#[derive(Debug, Clone, Copy, PartialEq)]
11pub struct HilbertSort;
12
13impl<N: IndexableNum> Sort<N> for HilbertSort {
14    fn sort(params: &mut SortParams<N>, boxes: &mut [N], indices: &mut MutableIndices) {
15        let width = params.max_x - params.min_x; // || 1.0;
16        let height = params.max_y - params.min_y; // || 1.0;
17        let mut hilbert_values: Vec<u32> = Vec::with_capacity(params.num_items);
18        let hilbert_max = ((1 << 16) - 1) as f64;
19
20        {
21            // map item centers into Hilbert coordinate space and calculate Hilbert values
22            let mut pos = 0;
23            for _ in 0..params.num_items {
24                let min_x = boxes[pos];
25                pos += 1;
26                let min_y = boxes[pos];
27                pos += 1;
28                let max_x = boxes[pos];
29                pos += 1;
30                let max_y = boxes[pos];
31                pos += 1;
32
33                let x = (hilbert_max
34                    * ((min_x + max_x).to_f64().unwrap() / 2. - params.min_x.to_f64().unwrap())
35                    / width.to_f64().unwrap())
36                .floor() as u32;
37                let y = (hilbert_max
38                    * ((min_y + max_y).to_f64().unwrap() / 2. - params.min_y.to_f64().unwrap())
39                    / height.to_f64().unwrap())
40                .floor() as u32;
41                hilbert_values.push(hilbert(x, y));
42            }
43        }
44
45        // sort items by their Hilbert value (for packing later)
46        sort(
47            &mut hilbert_values,
48            boxes,
49            indices,
50            0,
51            params.num_items - 1,
52            params.node_size,
53        );
54    }
55}
56
57/// Custom quicksort that partially sorts bbox data alongside the hilbert values.
58// Partially taken from static_aabb2d_index under the MIT/Apache license
59fn sort<N: IndexableNum>(
60    values: &mut [u32],
61    boxes: &mut [N],
62    indices: &mut MutableIndices,
63    left: usize,
64    right: usize,
65    node_size: usize,
66) {
67    debug_assert!(left <= right);
68
69    if left / node_size >= right / node_size {
70        return;
71    }
72
73    let midpoint = (left + right) / 2;
74    let pivot = values[midpoint];
75    let mut i = left.wrapping_sub(1);
76    let mut j = right.wrapping_add(1);
77
78    loop {
79        loop {
80            i = i.wrapping_add(1);
81            if values[i] >= pivot {
82                break;
83            }
84        }
85
86        loop {
87            j = j.wrapping_sub(1);
88            if values[j] <= pivot {
89                break;
90            }
91        }
92
93        if i >= j {
94            break;
95        }
96
97        swap(values, boxes, indices, i, j);
98    }
99
100    sort(values, boxes, indices, left, j, node_size);
101    sort(values, boxes, indices, j.wrapping_add(1), right, node_size);
102}
103
104/// Swap two values and two corresponding boxes.
105#[inline]
106fn swap<N: IndexableNum>(
107    values: &mut [u32],
108    boxes: &mut [N],
109    indices: &mut MutableIndices,
110    i: usize,
111    j: usize,
112) {
113    values.swap(i, j);
114
115    let k = 4 * i;
116    let m = 4 * j;
117    boxes.swap(k, m);
118    boxes.swap(k + 1, m + 1);
119    boxes.swap(k + 2, m + 2);
120    boxes.swap(k + 3, m + 3);
121
122    indices.swap(i, j);
123}
124
125// Taken from static_aabb2d_index under the mit/apache license
126// https://github.com/jbuckmccready/static_aabb2d_index/blob/9e6add59d77b74d4de0ac32159db47fbcb3acc28/src/static_aabb2d_index.rs#L486C1-L544C2
127#[inline]
128fn hilbert(x: u32, y: u32) -> u32 {
129    // Fast Hilbert curve algorithm by http://threadlocalmutex.com/
130    // Ported from C++ https://github.com/rawrunprotected/hilbert_curves (public domain)
131    let mut a_1 = x ^ y;
132    let mut b_1 = 0xFFFF ^ a_1;
133    let mut c_1 = 0xFFFF ^ (x | y);
134    let mut d_1 = x & (y ^ 0xFFFF);
135
136    let mut a_2 = a_1 | (b_1 >> 1);
137    let mut b_2 = (a_1 >> 1) ^ a_1;
138    let mut c_2 = ((c_1 >> 1) ^ (b_1 & (d_1 >> 1))) ^ c_1;
139    let mut d_2 = ((a_1 & (c_1 >> 1)) ^ (d_1 >> 1)) ^ d_1;
140
141    a_1 = a_2;
142    b_1 = b_2;
143    c_1 = c_2;
144    d_1 = d_2;
145    a_2 = (a_1 & (a_1 >> 2)) ^ (b_1 & (b_1 >> 2));
146    b_2 = (a_1 & (b_1 >> 2)) ^ (b_1 & ((a_1 ^ b_1) >> 2));
147    c_2 ^= (a_1 & (c_1 >> 2)) ^ (b_1 & (d_1 >> 2));
148    d_2 ^= (b_1 & (c_1 >> 2)) ^ ((a_1 ^ b_1) & (d_1 >> 2));
149
150    a_1 = a_2;
151    b_1 = b_2;
152    c_1 = c_2;
153    d_1 = d_2;
154    a_2 = (a_1 & (a_1 >> 4)) ^ (b_1 & (b_1 >> 4));
155    b_2 = (a_1 & (b_1 >> 4)) ^ (b_1 & ((a_1 ^ b_1) >> 4));
156    c_2 ^= (a_1 & (c_1 >> 4)) ^ (b_1 & (d_1 >> 4));
157    d_2 ^= (b_1 & (c_1 >> 4)) ^ ((a_1 ^ b_1) & (d_1 >> 4));
158
159    a_1 = a_2;
160    b_1 = b_2;
161    c_1 = c_2;
162    d_1 = d_2;
163    c_2 ^= (a_1 & (c_1 >> 8)) ^ (b_1 & (d_1 >> 8));
164    d_2 ^= (b_1 & (c_1 >> 8)) ^ ((a_1 ^ b_1) & (d_1 >> 8));
165
166    a_1 = c_2 ^ (c_2 >> 1);
167    b_1 = d_2 ^ (d_2 >> 1);
168
169    let mut i0 = x ^ y;
170    let mut i1 = b_1 | (0xFFFF ^ (i0 | a_1));
171
172    i0 = (i0 | (i0 << 8)) & 0x00FF00FF;
173    i0 = (i0 | (i0 << 4)) & 0x0F0F0F0F;
174    i0 = (i0 | (i0 << 2)) & 0x33333333;
175    i0 = (i0 | (i0 << 1)) & 0x55555555;
176
177    i1 = (i1 | (i1 << 8)) & 0x00FF00FF;
178    i1 = (i1 | (i1 << 4)) & 0x0F0F0F0F;
179    i1 = (i1 | (i1 << 2)) & 0x33333333;
180    i1 = (i1 | (i1 << 1)) & 0x55555555;
181
182    (i1 << 1) | i0
183}