space-time 0.4.0

//
// Copyright 2020, Gobsmacked Labs, LLC.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! SpaceFillingCurve for storing non-point features based on a bounding box.

use crate::index_range::{CoveredRange, IndexRange, OverlappingRange};
use alloc::{boxed::Box, collections::VecDeque, vec, vec::Vec};
use num_integer::div_floor;
#[allow(unused_imports)]
use num_traits::Float;

/// Z-order curve implementation for non-point features.
///
/// Based on [geomesa-z3 scala implementation](https://github.com/locationtech/geomesa/blob/771777d3a9716b04f7dcd27a6b7d1bb822a1b5a7/geomesa-z3/src/main/scala/org/locationtech/geomesa/curve/XZ2SFC.scala)
/// which is based on 'XZ-Ordering: A Space Filling Curve for Objects
/// with Spatial Extension' by Christian Bohm, Gerald Klump, and Hans-Peter Kriegel
pub struct XZ2SFC {
    g: u32,
    x_min: f64,
    x_max: f64,
    y_min: f64,
    y_max: f64,
}

impl XZ2SFC {
    /// Maximum supported resolution. Above this, the `4.pow(g - i)` terms in
    /// `sequence_code`/`sequence_interval` overflow `u64` (`4^32 == 2^64`).
    pub const MAX_G: u32 = 31;

    fn x_size(&self) -> f64 {
        self.x_max - self.x_min
    }

    fn y_size(&self) -> f64 {
        self.y_max - self.y_min
    }

    /// Return an `XZ2SFC`.
    ///
    /// # Panics
    ///
    /// Panics if `g > MAX_G` (31), the resolution above which the curve's
    /// index arithmetic overflows `u64`.
    #[must_use]
    pub fn new(g: u32, x_min: f64, y_min: f64, x_max: f64, y_max: f64) -> Self {
        assert!(g <= Self::MAX_G, "resolution g must be <= {}", Self::MAX_G);
        XZ2SFC {
            g,
            x_min,
            x_max,
            y_min,
            y_max,
        }
    }

    /// An `XZ2SFC` for unprojected coordinates.
    ///
    /// # Panics
    ///
    /// Panics if `g > MAX_G` (31).
    #[must_use]
    pub fn wgs84(g: u32) -> Self {
        assert!(g <= Self::MAX_G, "resolution g must be <= {}", Self::MAX_G);
        XZ2SFC {
            g,
            x_min: -180.0,
            x_max: 180.0,
            y_min: -90.0,
            y_max: 90.0,
        }
    }

    /// Return the index for a bounding box.
    #[must_use]
    pub fn index(&self, xmin: f64, ymin: f64, xmax: f64, ymax: f64) -> u64 {
        let (nxmin, nymin, nxmax, nymax) = self.normalize(xmin, ymin, xmax, ymax);

        let max_dim = (nxmax - nxmin).max(nymax - nymin);

        // This is a slightly different construction but same value as geomesa.
        let el_1 = max_dim.log(0.5).floor() as i32;

        let length: u32 = if el_1 as u32 >= self.g {
            self.g
        } else {
            let w2 = 0.5_f64.powi(el_1 + 1);

            if Self::predicate(nxmin, nxmax, w2) && Self::predicate(nxmin, nxmax, w2) {
                (el_1 + 1) as u32
            } else {
                el_1 as u32
            }
        };

        self.sequence_code(nxmin, nymin, length)
    }

    fn predicate(min: f64, max: f64, w2: f64) -> bool {
        max <= (min / w2).floor() * w2 + 2.0 * w2
    }

    /// Compute that index ranges that are contained or overlap the bounding box.
    pub fn ranges(
        &self,
        xmin: f64,
        ymin: f64,
        xmax: f64,
        ymax: f64,
        max_ranges: Option<u16>,
    ) -> Vec<Box<dyn IndexRange>> {
        let windows = {
            let (nxmin, nymin, nxmax, nymax) = self.normalize(xmin, ymin, xmax, ymax);
            &[QueryWindow {
                xmin: nxmin,
                ymin: nymin,
                xmax: nxmax,
                ymax: nymax,
            }]
        };

        let range_stop = max_ranges.unwrap_or(u16::MAX);

        self.ranges_impl(windows, range_stop)
    }

    fn ranges_impl(&self, query: &[QueryWindow], range_stop: u16) -> Vec<Box<dyn IndexRange>> {
        let mut ranges: Vec<Box<dyn IndexRange>> = Vec::with_capacity(100);

        let mut remaining: VecDeque<Option<XElement>> = VecDeque::with_capacity(100);

        for el in XElement::level_one_elements() {
            remaining.push_back(Some(el));
        }
        remaining.push_back(LEVEL_TERMINATOR);

        let mut level: u32 = 1;

        while level < self.g && !remaining.is_empty() && ranges.len() < range_stop.into() {
            match remaining.pop_front() {
                Some(LEVEL_TERMINATOR) if !remaining.is_empty() => {
                    level += 1;
                    remaining.push_back(LEVEL_TERMINATOR);
                }
                Some(element) => {
                    self.check_value(element, level, query, &mut ranges, &mut remaining)
                }
                _ => (),
            }
        }

        while let Some(quad) = remaining.pop_front() {
            if let Some(quad) = quad {
                let (min, max) = self.sequence_interval(quad.xmin, quad.ymin, level, false);
                ranges.push(Box::new(OverlappingRange::new(min, max)));
            } else {
                level += 1;
            }
        }

        ranges.sort();

        let mut current: Option<Box<dyn IndexRange>> = None;

        let mut results = vec![];

        for range in ranges {
            if let Some(cur) = current {
                if range.lower() <= cur.upper().saturating_add(1) {
                    let max = cur.upper().max(range.upper());
                    let min = cur.lower();
                    if cur.contained() && range.contained() {
                        current = Some(Box::new(CoveredRange::new(min, max)));
                    } else {
                        current = Some(Box::new(OverlappingRange::new(min, max)));
                    }
                } else {
                    results.push(cur);
                    current = Some(range);
                }
            } else {
                current = Some(range);
            }
        }

        if let Some(current) = current {
            results.push(current);
        }

        results
    }

    /// Curve-code offset contributed by quadrant `q` (`0..=3`) at remaining
    /// resolution `k`: the number of cells skipped by stepping past quadrants
    /// `0..q` of a node whose subtree spans `k` more levels. Shared by
    /// `sequence_code` and `sequence_interval` so the kani overflow proof
    /// exercises the exact arithmetic used in production. Does not overflow for
    /// `k <= MAX_G` (see `code_offset_doesnt_overflow`).
    fn code_offset(q: u64, k: u32) -> u64 {
        div_floor(q * (4_u64.pow(k) - 1), 3)
    }

    fn sequence_code(&self, x: f64, y: f64, length: u32) -> u64 {
        let mut xmin = 0.0;
        let mut ymin = 0.0;
        let mut xmax = 1.0;
        let mut ymax = 1.0;

        let mut cs = 0_u64;

        for i in 0_u32..length {
            let x_center = (xmin + xmax) / 2.0;
            let y_center = (ymin + ymax) / 2.0;

            match (x < x_center, y < y_center) {
                (true, true) => {
                    cs += 1;
                    xmax = x_center;
                    ymax = y_center;
                }
                (false, true) => {
                    cs += 1 + Self::code_offset(1, self.g - i);
                    xmin = x_center;
                    ymax = y_center;
                }
                (true, false) => {
                    cs += 1 + Self::code_offset(2, self.g - i);
                    xmax = x_center;
                    ymin = y_center;
                }
                (false, false) => {
                    cs += 1 + Self::code_offset(3, self.g - i);
                    xmin = x_center;
                    ymin = y_center;
                }
            }
        }
        cs
    }

    fn check_value(
        &self,
        quad: Option<XElement>,
        level: u32,
        query: &[QueryWindow],
        ranges: &mut Vec<Box<dyn IndexRange>>,
        remaining: &mut VecDeque<Option<XElement>>,
    ) {
        if let Some(quad) = quad {
            if Self::is_contained(quad, query) {
                let (min, max) = self.sequence_interval(quad.xmin, quad.ymin, level, false);
                ranges.push(Box::new(CoveredRange::new(min, max)));
            } else if Self::is_overlapped(quad, query) {
                let (min, max) = self.sequence_interval(quad.xmin, quad.ymin, level, true);
                ranges.push(Box::new(OverlappingRange::new(min, max)));
                for el in quad.children() {
                    remaining.push_back(Some(el));
                }
            }
        }
    }

    fn is_contained(quad: XElement, query: &[QueryWindow]) -> bool {
        for q in query {
            if quad.is_contained(q) {
                return true;
            }
        }
        false
    }

    fn is_overlapped(quad: XElement, query: &[QueryWindow]) -> bool {
        for q in query {
            if quad.overlaps(q) {
                return true;
            }
        }
        false
    }

    fn sequence_interval(&self, x: f64, y: f64, length: u32, partial: bool) -> (u64, u64) {
        let min = self.sequence_code(x, y, length);

        let max = if partial {
            min
        } else {
            min + Self::code_offset(1, self.g - length + 1)
        };

        (min, max)
    }

    fn normalize(&self, x_min: f64, y_min: f64, x_max: f64, y_max: f64) -> (f64, f64, f64, f64) {
        assert!(x_min <= x_max && y_min <= y_max);
        assert!(
            x_min >= self.x_min
                && x_max <= self.x_max
                && y_min >= self.y_min
                && y_max <= self.y_max
        );

        (
            (x_min - self.x_min) / self.x_size(),
            (y_min - self.y_min) / self.y_size(),
            (x_max - self.x_min) / self.x_size(),
            (y_max - self.y_min) / self.y_size(),
        )
    }
}

const LEVEL_TERMINATOR: Option<XElement> = None;

#[derive(Debug, Clone, Copy)]
struct QueryWindow {
    pub xmin: f64,
    pub ymin: f64,
    pub xmax: f64,
    pub ymax: f64,
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct XElement {
    xmin: f64,
    ymin: f64,
    xmax: f64,
    ymax: f64,
    length: f64,
}

impl XElement {
    const fn new(xmin: f64, ymin: f64, xmax: f64, ymax: f64, length: f64) -> Self {
        XElement {
            xmin,
            ymin,
            xmax,
            ymax,
            length,
        }
    }

    fn xext(&self) -> f64 {
        self.xmax + self.length
    }

    fn yext(&self) -> f64 {
        self.ymax + self.length
    }

    fn is_contained(&self, window: &QueryWindow) -> bool {
        window.xmin <= self.xmin
            && window.ymin <= self.ymin
            && window.xmax >= self.xext()
            && window.ymax >= self.yext()
    }

    fn overlaps(&self, window: &QueryWindow) -> bool {
        window.xmax >= self.xmin
            && window.ymax >= self.ymin
            && window.xmin <= self.xext()
            && window.ymin <= self.yext()
    }

    fn level_one_elements() -> Vec<XElement> {
        XElement::new(0.0, 0.0, 1.0, 1.0, 1.0).children()
    }

    fn children(&self) -> Vec<XElement> {
        let x_center = (self.xmin + self.xmax) / 2.0;
        let y_center = (self.ymin + self.ymax) / 2.0;
        let len = self.length / 2.0;

        vec![
            XElement::new(self.xmin, self.ymin, x_center, y_center, len),
            XElement::new(x_center, self.ymin, self.xmax, y_center, len),
            XElement::new(self.xmin, y_center, x_center, self.ymax, len),
            XElement::new(x_center, y_center, self.xmax, self.ymax, len),
        ]
    }
}

#[cfg(kani)]
mod kani_proofs {
    use super::*;

    /// `XZ2SFC::code_offset` — the single arithmetic helper behind every
    /// `sequence_code`/`sequence_interval` curve-code term — never overflows
    /// `u64` under the `g <= MAX_G` bound enforced by the constructors. The
    /// resolution argument `k` (`g - i` with `i` in `0..length <= g`, or
    /// `g - length + 1`) lies in `1..=g`, and the quadrant `q` in `0..=3`, so
    /// the proof ranges over every value reachable in production.
    #[kani::proof]
    #[kani::unwind(33)]
    fn code_offset_doesnt_overflow() {
        let g: u32 = kani::any();
        kani::assume(g <= XZ2SFC::MAX_G);

        let k: u32 = kani::any();
        kani::assume(k >= 1 && k <= g);

        let q: u64 = kani::any();
        kani::assume(q <= 3);

        let _ = XZ2SFC::code_offset(q, k);
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_query_bounding_boxes() {
        let sfc = XZ2SFC::wgs84(12);
        let polygon = sfc.index(10.0, 10.0, 12.0, 12.0);

        let containing = [
            (9.0, 9.0, 13.0, 13.0),
            (-180.0, -90.0, 180.0, 90.0),
            (0.0, 0.0, 180.0, 90.0),
            (0.0, 0.0, 20.0, 20.0),
        ];
        let overlapping = [
            (11.0, 11.0, 13.0, 13.0),
            (9.0, 9.0, 11.0, 11.0),
            (10.5, 10.5, 11.5, 11.5),
            (11.0, 11.0, 11.0, 11.0),
        ];
        let disjoint = [
            (-180.0, -90.0, 8.0, 8.0),
            (0.0, 0.0, 8.0, 8.0),
            (9.0, 9.0, 9.5, 9.5),
            (20.0, 20.0, 180.0, 90.0),
        ];

        for bbox in &[containing, overlapping].concat() {
            let ranges = sfc.ranges(bbox.0, bbox.1, bbox.2, bbox.3, None);
            assert!(
                ranges
                    .iter()
                    .any(|r| r.lower() <= polygon && polygon <= r.upper())
            );
        }

        for bbox in &disjoint {
            let ranges = sfc.ranges(bbox.0, bbox.1, bbox.2, bbox.3, None);
            assert!(
                !ranges
                    .iter()
                    .any(|r| r.lower() <= polygon && polygon <= r.upper())
            );
        }
    }

    #[test]
    fn test_index() {
        let sfc = XZ2SFC::wgs84(12);
        assert_eq!(sfc.index(10.0, 10.0, 12.0, 12.0), 16841390);
        assert_eq!(sfc.index(-180.0, -90.0, -180.0, -90.0), 12);
        assert_eq!(sfc.index(-180.0, -90.0, 0.0, 0.0), 2);
        assert_eq!(sfc.index(10.0, -90.0, 12.0, -89.0), 5599580);
        assert_eq!(sfc.index(79.9, 0.5, 79.9, 0.5), 17236267);
    }

    #[test]
    fn test_ranges() {
        let sfc = XZ2SFC::wgs84(20);

        assert_eq!(sfc.ranges(-0.5, -0.5, 0.5, 0.5, None).len(), 8077);
        assert!(sfc.ranges(-0.5, -0.5, 0.5, 0.5, Some(1000)).len() < 1000);

        assert_eq!(sfc.ranges(55.758, 20.5, 55.759, 21.5, None).len(), 5883);
        assert_eq!(sfc.ranges(-55.758, 20.5, -55.755, 21.5, None).len(), 8070);

        let ranges = sfc.ranges(-55.758, 20.5, -55.755, 21.5, None);

        assert_eq!(ranges.first().map(|r| r.lower()), Some(1));
        assert_eq!(ranges.last().map(|r| r.upper()), Some(847016214083));
    }
}

#[cfg(test)]
mod range_query_tests {
    use super::*;
    use quickcheck_macros::quickcheck;

    /// `index` must not panic on the float→int cast edge cases: a degenerate
    /// point box (`max_dim == 0`, where `log(0.5)` is `+inf` and the `el_1`
    /// cast saturates to `i32::MAX`), the full-extent box (`max_dim == 1`), and
    /// a sub-cell box.
    #[test]
    fn index_handles_cast_edge_cases() {
        let sfc = XZ2SFC::wgs84(12);
        let _ = sfc.index(10.0, 20.0, 10.0, 20.0); // point box -> length = g
        let _ = sfc.index(-180.0, -90.0, 180.0, 90.0); // full extent -> el_1 == 0
        let _ = sfc.index(-179.999_9, -89.999_9, -179.999_8, -89.999_8);
    }

    /// Every `(code, element)` in the quadtree down to depth `g`.
    fn enumerate(sfc: &XZ2SFC) -> Vec<(u64, XElement)> {
        let mut out = Vec::new();
        let mut stack: Vec<(XElement, u32)> = XElement::level_one_elements()
            .into_iter()
            .map(|e| (e, 1))
            .collect();
        while let Some((e, level)) = stack.pop() {
            let code = sfc.sequence_code(e.xmin, e.ymin, level);
            if level < sfc.g {
                for c in e.children() {
                    stack.push((c, level + 1));
                }
            }
            out.push((code, e));
        }
        out
    }

    /// End-to-end property for the XZ2 range engine. For any query window and
    /// any `range_stop` (including the small values that force early
    /// `bottom_out`), `ranges_impl` must be
    ///
    ///  * COMPLETE — every element whose (enlarged) extent overlaps the query is covered
    ///    by some returned range, so no candidate object is missed, and
    ///  * SOUND — every element whose code lands in a `CoveredRange` is fully contained
    ///    in the query (overlapping ranges may spill, covered ones may not).
    ///
    /// Driven in normalized `[0, 1]` space with dyadic query coordinates so all
    /// element-boundary comparisons are exact in `f64`.
    #[quickcheck]
    fn xz2_ranges_complete_and_sound(g_seed: u8, coords: (u8, u8, u8, u8), range_stop: u8) -> bool {
        let g = u32::from(g_seed % 4) + 1; // 1..=4
        let sfc = XZ2SFC::wgs84(g);

        let denom = f64::from(1u32 << g); // 2^g
        let to_unit = |v: u8| f64::from(u32::from(v) % ((1u32 << g) + 1)) / denom;
        let (a, b, c, d) = coords;
        let (mut xmin, mut xmax) = (to_unit(a), to_unit(c));
        if xmin > xmax {
            core::mem::swap(&mut xmin, &mut xmax);
        }
        let (mut ymin, mut ymax) = (to_unit(b), to_unit(d));
        if ymin > ymax {
            core::mem::swap(&mut ymin, &mut ymax);
        }
        let window = QueryWindow {
            xmin,
            ymin,
            xmax,
            ymax,
        };

        // `0` means "no limit"; otherwise an aggressive cap that forces bottom-out.
        let range_stop = if range_stop == 0 {
            u16::MAX
        } else {
            u16::from(range_stop)
        };
        let ranges = sfc.ranges_impl(&[window], range_stop);

        let elements = enumerate(&sfc);

        // Completeness.
        for (code, e) in &elements {
            if e.overlaps(&window)
                && !ranges
                    .iter()
                    .any(|r| r.lower() <= *code && *code <= r.upper())
            {
                return false;
            }
        }

        // Soundness of covered ranges.
        for (code, e) in &elements {
            for r in &ranges {
                if r.contained()
                    && r.lower() <= *code
                    && *code <= r.upper()
                    && !e.is_contained(&window)
                {
                    return false;
                }
            }
        }

        true
    }
}