Add MortonUtils

src/Common/MortonUtils.h

@@ -0,0 +1,199 @@
+#pragma once
+
+#include <Core/Types.h>
+#include <Common/BitHelpers.h>
+#include <base/defines.h>
+#include <array>
+#include <iostream>
+
+
+namespace
+{
+    inline UInt64 toMask(UInt64 n)
+    {
+        n |= n >> 1;
+        n |= n >> 2;
+        n |= n >> 4;
+        n |= n >> 8;
+        n |= n >> 16;
+        n |= n >> 32;
+        return n;
+    }
+}
+
+
+/** Splits the interval [first, last] to a set of intervals [first_i, last_i],
+  * each of them determined by a bit prefix: [xxxxxx0000, xxxxxx1111],
+  *
+  * For example, the interval [6, 13] = {5, 7, 8, 9, 10, 11, 12, 13}
+  * will be represented by the set of intervals:
+  * - [6,  7]   0000011*
+  * - [8,  11]  000010**
+  * - [12, 13]  0000110*
+  *
+  * It means that if you have a binary space partition by powers of two,
+  * every of the resulting intervals will fully occupy one of the levels of this partition.
+  */
+template <typename F>
+void intervalBinaryPartition(UInt64 first, UInt64 last, F && callback)
+{
+    /// first = 6:    00000110
+    /// last = 13:    00001101
+    /// first ^ last: 00001011
+    /// mask:         00000111
+    /// split = 7:    00000111
+
+    /// first = 8:    00001000
+    /// last = 13:    00001101
+    /// first ^ last: 00000101
+    /// mask:         00000011
+    /// split = 11:   00001011
+
+    /// first = 8:    00001000
+    /// last = 11:    00001011
+    /// first ^ last: 00000011
+    /// mask:         00000001
+    /// split = 9:    00001001
+
+    /// Another example:
+
+    /// first = 15:   00001111
+    /// last = 31:    00011111
+    /// first ^ last: 00010000
+    /// mask:         00001111
+    /// split = 15:   00001111
+
+    /// Another example:
+
+    /// first = 6:    00000110
+    /// last = 7:     00000111
+    /// first ^ last: 00000001
+    /// mask:         00000000
+    /// split = 11:   00001011
+
+    UInt64 diff = first ^ last;
+    UInt64 mask = toMask(diff) >> 1;
+
+    /// The current interval represents a whole range with fixed prefix.
+    if ((first & mask) == 0 && (last & mask) == mask)
+    {
+        chassert(((last - first + 1) & (last - first)) == 0); /// The interval length is one less than a power of two.
+        callback(first, last);
+        return;
+    }
+
+    UInt64 split = first | mask;
+
+    chassert(split >= first);
+    chassert(split <= last);
+
+    intervalBinaryPartition(first, split, std::forward<F>(callback));
+    if (split < last)
+        intervalBinaryPartition(split + 1, last, std::forward<F>(callback));
+}
+
+
+/** Multidimensional version of binary space partitioning.
+  * It takes a parallelogram - a direct product of intervals (in each dimension),
+  * and splits it into smaller parallelograms - a direct product of partitions across each dimension.
+  */
+template <size_t N, size_t start_idx, typename F>
+void parallelogramBinaryPartitionImpl(
+    std::array<std::pair<UInt64, UInt64>, N> parallelogram,
+    F && callback)
+{
+    intervalBinaryPartition(parallelogram[start_idx].first, parallelogram[start_idx].second,
+        [&](UInt64 a, UInt64 b) mutable
+        {
+            auto new_parallelogram = parallelogram;
+            new_parallelogram[start_idx].first = a;
+            new_parallelogram[start_idx].second = b;
+
+            if constexpr (start_idx + 1 < N)
+                parallelogramBinaryPartitionImpl<N, start_idx + 1>(new_parallelogram, std::forward<F>(callback));
+            else
+                callback(new_parallelogram);
+        });
+}
+
+
+template <size_t N, typename F>
+void parallelogramBinaryPartition(
+    std::array<std::pair<UInt64, UInt64>, N> parallelogram,
+    F && callback)
+{
+    parallelogramBinaryPartitionImpl<N, 0>(parallelogram, std::forward<F>(callback));
+}
+
+
+/** Unpack an interval of Morton curve to parallelograms covered by it across N dimensions.
+  */
+template <size_t N, typename F>
+void mortonIntervalToParallelograms(UInt64 first, UInt64 last, F && callback)
+{
+    intervalBinaryPartition(first, last, [&](UInt64 a, UInt64 b)
+    {
+        std::array<std::pair<UInt64, UInt64>, N> unpacked{};
+
+        for (size_t bit_idx = 0; bit_idx < 64; ++bit_idx)
+        {
+            size_t source_bit = 63 - bit_idx;
+            size_t result_bit = (63 - bit_idx) / N;
+
+            unpacked[source_bit % N].first |= ((a >> source_bit) & 1) << result_bit;
+            unpacked[source_bit % N].second |= ((b >> source_bit) & 1) << result_bit;
+        }
+
+        callback(unpacked);
+    });
+}
+
+
+/** Given a parallelogram, find intervals of Morton curve that cover this parallelogram.
+  * Note: to avoid returning too many intervals, the intervals can be returned larger than exactly needed
+  * (covering some other points, not belonging to the parallelogram).
+  */
+template <size_t N, typename F>
+void parallelogramToPossibleMortonIntervals(
+    std::array<std::pair<UInt64, UInt64>, N> parallelogram,
+    F && callback)
+{
+    parallelogramBinaryPartition<N>(parallelogram, [&](auto part)
+    {
+        size_t suffix_size = 0;
+        for (size_t i = 0; i < N; ++i)
+            if (part[i].second != part[i].first)
+                suffix_size = std::max<size_t>(suffix_size,
+                    bitScanReverse(part[i].second - part[i].first));
+
+        UInt64 first = 0;
+        UInt64 last = 0;
+
+        size_t source_bit_idx = 0;
+        size_t result_bit_idx = 0;
+
+        while (result_bit_idx < 64)
+        {
+            for (size_t i = 0; i < N; ++i)
+            {
+                if (source_bit_idx <= suffix_size)
+                {
+                    last |= (1 << result_bit_idx);
+                }
+                else
+                {
+                    UInt64 bit = (((part[i].first >> source_bit_idx) & 1) << result_bit_idx);
+                    first |= bit;
+                    last |= bit;
+                }
+
+                ++result_bit_idx;
+                if (!(result_bit_idx < 64))
+                    break;
+            }
+            ++source_bit_idx;
+        }
+
+        callback(first, last);
+    });
+}

src/Common/tests/gtest_morton_utils.cpp

@@ -0,0 +1,107 @@
+#include <gtest/gtest.h>
+#include <iostream>
+#include <Common/MortonUtils.h>
+
+
+GTEST_TEST(MortonUtils, Intervals)
+{
+    {
+        std::stringstream res;
+        intervalBinaryPartition(6, 13, [&](UInt64 first, UInt64 last)
+        {
+            res << first << ", " << last << "; ";
+        });
+        ASSERT_EQ(res.str(), "6, 7; 8, 11; 12, 13; ");
+    }
+
+    {
+        std::stringstream res;
+        intervalBinaryPartition(15, 31, [&](UInt64 first, UInt64 last)
+        {
+            res << first << ", " << last << "; ";
+        });
+        ASSERT_EQ(res.str(), "15, 15; 16, 31; ");
+    }
+
+    {
+        std::stringstream res;
+        intervalBinaryPartition(15, 16, [&](UInt64 first, UInt64 last)
+        {
+            res << first << ", " << last << "; ";
+        });
+        ASSERT_EQ(res.str(), "15, 15; 16, 16; ");
+    }
+
+    {
+        std::stringstream res;
+        intervalBinaryPartition(191, 769, [&](UInt64 first, UInt64 last)
+        {
+            res << first << ", " << last << "; ";
+        });
+        ASSERT_EQ(res.str(), "191, 191; 192, 255; 256, 511; 512, 767; 768, 769; ");
+    }
+
+    {
+        std::array<std::pair<UInt64, UInt64>, 2> input = {std::pair{6, 13}, std::pair{15, 31}};
+
+        std::stringstream res;
+        parallelogramBinaryPartition<2>(input, [&](auto parallelogram)
+        {
+            res << "[" << parallelogram[0].first << ", " << parallelogram[0].second
+                << "] x [" << parallelogram[1].first << ", " << parallelogram[1].second
+                << "]; ";
+        });
+
+        ASSERT_EQ(res.str(), "[6, 7] x [15, 15]; [6, 7] x [16, 31]; [8, 11] x [15, 15]; [8, 11] x [16, 31]; [12, 13] x [15, 15]; [12, 13] x [16, 31]; ");
+    }
+
+    {
+        std::array<std::pair<UInt64, UInt64>, 2> input = {std::pair{23, 24}, std::pair{15, 16}};
+
+        std::stringstream res;
+        parallelogramBinaryPartition<2>(input, [&](auto parallelogram)
+        {
+            res << "[" << parallelogram[0].first << ", " << parallelogram[0].second
+                << "] x [" << parallelogram[1].first << ", " << parallelogram[1].second
+                << "]; ";
+        });
+
+        ASSERT_EQ(res.str(), "[23, 23] x [15, 15]; [23, 23] x [16, 16]; [24, 24] x [15, 15]; [24, 24] x [16, 16]; ");
+    }
+
+    {
+        std::stringstream res;
+        mortonIntervalToParallelograms<2>(191, 769, [&](auto parallelogram)
+        {
+            res << "[" << parallelogram[0].first << ", " << parallelogram[0].second
+                << "] x [" << parallelogram[1].first << ", " << parallelogram[1].second
+                << "]; ";
+        });
+
+        ASSERT_EQ(res.str(), "[7, 7] x [15, 15]; [8, 15] x [8, 15]; [16, 31] x [0, 15]; [0, 15] x [16, 31]; [16, 17] x [16, 16]; ");
+    }
+
+    {
+        std::array<std::pair<UInt64, UInt64>, 2> input = {std::pair{23, 24}, std::pair{15, 16}};
+
+        std::stringstream res;
+        parallelogramToPossibleMortonIntervals<2>(input, [&](UInt64 first, UInt64 last)
+        {
+            res << first << ", " << last << "; ";
+        });
+
+        std::cerr << res.str() << "\n";
+    }
+
+    {
+        std::array<std::pair<UInt64, UInt64>, 2> input = {std::pair{6, 13}, std::pair{15, 31}};
+
+        std::stringstream res;
+        parallelogramToPossibleMortonIntervals<2>(input, [&](UInt64 first, UInt64 last)
+        {
+            res << first << ", " << last << "; ";
+        });
+
+        std::cerr << res.str() << "\n";
+    }
+}