z_order.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
#pragma once
 
 
#include "signet/types.hpp"
 
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <numeric>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>
 
namespace signet::forge::z_order {
 
// ===========================================================================
// Value normalization -- preserve sort order as unsigned integers
// ===========================================================================
 
inline uint32_t normalize_int32(int32_t v) {
    return static_cast<uint32_t>(v) ^ 0x80000000u;
}
 
inline uint64_t normalize_int64(int64_t v) {
    return static_cast<uint64_t>(v) ^ 0x8000000000000000ULL;
}
 
inline uint32_t normalize_float(float v) {
    uint32_t bits;
    std::memcpy(&bits, &v, 4);
    uint32_t mask = -static_cast<int32_t>(bits >> 31) | 0x80000000u;
    return bits ^ mask;
}
 
inline uint64_t normalize_double(double v) {
    uint64_t bits;
    std::memcpy(&bits, &v, 8);
    uint64_t mask = -static_cast<int64_t>(bits >> 63) | 0x8000000000000000ULL;
    return bits ^ mask;
}
 
inline uint32_t normalize_string(std::string_view v) {
    uint32_t result = 0;
    for (size_t i = 0; i < 4 && i < v.size(); ++i) {
        result |= static_cast<uint32_t>(static_cast<uint8_t>(v[i])) << (24 - 8 * i);
    }
    return result;
}
 
inline uint32_t truncate_to_32(uint64_t v) {
    return static_cast<uint32_t>(v >> 32);
}
 
// ===========================================================================
// Morton code — bit interleaving
// ===========================================================================
 
inline uint64_t morton_2d(uint32_t x, uint32_t y) {
    auto spread = [](uint32_t v) -> uint64_t {
        uint64_t r = v;
        r = (r | (r << 16)) & 0x0000FFFF0000FFFFULL;
        r = (r | (r <<  8)) & 0x00FF00FF00FF00FFULL;
        r = (r | (r <<  4)) & 0x0F0F0F0F0F0F0F0FULL;
        r = (r | (r <<  2)) & 0x3333333333333333ULL;
        r = (r | (r <<  1)) & 0x5555555555555555ULL;
        return r;
    };
    return spread(x) | (spread(y) << 1);
}
 
inline void deinterleave_2d(uint64_t code, uint32_t& x, uint32_t& y) {
    auto compact = [](uint64_t v) -> uint32_t {
        v &= 0x5555555555555555ULL;
        v = (v | (v >> 1))  & 0x3333333333333333ULL;
        v = (v | (v >> 2))  & 0x0F0F0F0F0F0F0F0FULL;
        v = (v | (v >> 4))  & 0x00FF00FF00FF00FFULL;
        v = (v | (v >> 8))  & 0x0000FFFF0000FFFFULL;
        v = (v | (v >> 16)) & 0x00000000FFFFFFFFULL;
        return static_cast<uint32_t>(v);
    };
    x = compact(code);
    y = compact(code >> 1);
}
 
inline std::vector<uint8_t> morton_nd(const std::vector<uint32_t>& normalized,
                                       size_t bits_per_col = 32) {
    size_t n = normalized.size();
    if (n == 0) return {};
 
    size_t total_bits = n * bits_per_col;
    size_t total_bytes = (total_bits + 7) / 8;
    std::vector<uint8_t> result(total_bytes, 0);
 
    // Round-robin interleave: bit position `b` of column `c` maps to
    // output bit position `b * n + c`, with MSB first.
    size_t out_bit = 0;
    for (size_t b = 0; b < bits_per_col; ++b) {
        size_t src_bit = bits_per_col - 1 - b;  // MSB first
        for (size_t c = 0; c < n; ++c) {
            if (normalized[c] & (1u << src_bit)) {
                size_t byte_idx = out_bit / 8;
                size_t bit_in_byte = 7 - (out_bit % 8);
                result[byte_idx] |= static_cast<uint8_t>(1u << bit_in_byte);
            }
            ++out_bit;
        }
    }
 
    return result;
}
 
// ===========================================================================
// ZOrderColumn -- describes a column's data for Z-order sorting
// ===========================================================================
 
struct ZOrderColumn {
    PhysicalType type;   
    const void* data;    
    size_t count;        
};
 
// ===========================================================================
// ZOrderSorter -- sorts row indices by Morton key
// ===========================================================================
 
struct ZOrderSorter {
    [[nodiscard]] static std::vector<size_t> sort(
            size_t num_rows,
            const std::vector<ZOrderColumn>& columns) {
 
        if (num_rows == 0 || columns.empty()) return {};
 
        size_t n_cols = columns.size();
 
        // Validate that each column's count matches num_rows (CWE-787: OOB write)
        for (size_t c = 0; c < n_cols; ++c) {
            if (columns[c].count != num_rows) {
                throw std::out_of_range(
                    "ZOrderSorter::sort: column " + std::to_string(c) +
                    " count (" + std::to_string(columns[c].count) +
                    ") != num_rows (" + std::to_string(num_rows) + ")");
            }
        }
 
        // Normalize all values to uint32_t
        std::vector<std::vector<uint32_t>> normalized(n_cols);
        for (size_t c = 0; c < n_cols; ++c) {
            normalized[c].resize(num_rows);
            normalize_column(columns[c], normalized[c]);
        }
 
        // Fast path: 2 columns — use morton_2d with uint64_t sort keys
        if (n_cols == 2) {
            std::vector<uint64_t> keys(num_rows);
            for (size_t r = 0; r < num_rows; ++r) {
                keys[r] = morton_2d(normalized[0][r], normalized[1][r]);
            }
 
            std::vector<size_t> perm(num_rows);
            std::iota(perm.begin(), perm.end(), 0);
            std::sort(perm.begin(), perm.end(),
                      [&keys](size_t a, size_t b) { return keys[a] < keys[b]; });
            return perm;
        }
 
        // General path: N columns — use morton_nd byte-array sort keys
        std::vector<std::vector<uint8_t>> keys(num_rows);
        std::vector<uint32_t> row_vals(n_cols);
        for (size_t r = 0; r < num_rows; ++r) {
            for (size_t c = 0; c < n_cols; ++c) {
                row_vals[c] = normalized[c][r];
            }
            keys[r] = morton_nd(row_vals);
        }
 
        std::vector<size_t> perm(num_rows);
        std::iota(perm.begin(), perm.end(), 0);
        std::sort(perm.begin(), perm.end(),
                  [&keys](size_t a, size_t b) { return keys[a] < keys[b]; });
        return perm;
    }
 
private:
    static void normalize_column(const ZOrderColumn& col,
                                  std::vector<uint32_t>& out) {
        if (col.data == nullptr) {
            throw std::invalid_argument(
                "ZOrderSorter: column data pointer is null");
        }
        switch (col.type) {
            case PhysicalType::INT32: {
                // CWE-704, C++ [basic.align] — alignment check before pointer cast
                if (reinterpret_cast<uintptr_t>(col.data) % alignof(int32_t) != 0) {
                    throw std::invalid_argument(
                        "ZOrderSorter: INT32 data pointer not aligned for int32_t");
                }
                const auto* vals = static_cast<const int32_t*>(col.data);
                for (size_t i = 0; i < col.count; ++i) {
                    out[i] = normalize_int32(vals[i]);
                }
                break;
            }
            case PhysicalType::INT64: {
                // CWE-704, C++ [basic.align] — alignment check before pointer cast
                if (reinterpret_cast<uintptr_t>(col.data) % alignof(int64_t) != 0) {
                    throw std::invalid_argument(
                        "ZOrderSorter: INT64 data pointer not aligned for int64_t");
                }
                const auto* vals = static_cast<const int64_t*>(col.data);
                for (size_t i = 0; i < col.count; ++i) {
                    out[i] = truncate_to_32(normalize_int64(vals[i]));
                }
                break;
            }
            case PhysicalType::FLOAT: {
                // CWE-704, C++ [basic.align] — alignment check before pointer cast
                if (reinterpret_cast<uintptr_t>(col.data) % alignof(float) != 0) {
                    throw std::invalid_argument(
                        "ZOrderSorter: FLOAT data pointer not aligned for float");
                }
                const auto* vals = static_cast<const float*>(col.data);
                for (size_t i = 0; i < col.count; ++i) {
                    out[i] = normalize_float(vals[i]);
                }
                break;
            }
            case PhysicalType::DOUBLE: {
                // CWE-704, C++ [basic.align] — alignment check before pointer cast
                if (reinterpret_cast<uintptr_t>(col.data) % alignof(double) != 0) {
                    throw std::invalid_argument(
                        "ZOrderSorter: DOUBLE data pointer not aligned for double");
                }
                const auto* vals = static_cast<const double*>(col.data);
                for (size_t i = 0; i < col.count; ++i) {
                    out[i] = truncate_to_32(normalize_double(vals[i]));
                }
                break;
            }
            case PhysicalType::BYTE_ARRAY: {
                // Interpret data as array of std::string pointers
                const auto* vals = static_cast<const std::string*>(col.data);
                for (size_t i = 0; i < col.count; ++i) {
                    out[i] = normalize_string(vals[i]);
                }
                break;
            }
            default: {
                // For unsupported types, use zero (all rows sort equally)
                for (size_t i = 0; i < col.count; ++i) {
                    out[i] = 0;
                }
                break;
            }
        }
    }
};
 
} // namespace signet::forge::z_order