RTree implementation for 2D spatial data

RTree.fs

(*
  An immutable R-Tree implementation in F#. 
  
  Based on: https://github.com/swimos/swim-rust/tree/main/swim_utilities/swim_rtree (licensed on Apache 2.0)
  
  Author: Bartosz Sypytkowski <b.sypytkowski at gmail.com>
*)

namespace Demos.RTree

open System
open System.Collections.Generic
open System.Runtime.CompilerServices
open System.Text
open Demos

type SplitStrategy =
  | Linear  = 1
  | Quadratic = 2
  
type Config =
  { minCap: int  // minimal number of elements that a node must store - otherwise nodes must be merged
    maxCap: int  // maximal number of elements that a node can store - otherwise it must be split in two
    strategy: SplitStrategy } // strategy used for splitting

/// Interface implemented by an item type that's about to be stored in RTree.
type Spatial =
  abstract Boundary: Rect
  
and [<Struct;NoComparison;CustomEquality>] Point =
  { x: float; y: float }
  static member (+) (a: Point, b: Point) : Point = { x = a.x + b.x ; y = a.y + b.y }
  static member (-) (a: Point, b: Point) : Point = { x = a.x - b.x; y = a.y - b.y }
  member a.Compare(b: Point) =
    if a.x = b.x && a.y = b.y then ValueSome 0
    elif a.x <= b.x && a.y <= b.y then ValueSome -1
    elif a.x >= b.x && a.y >= b.y then ValueSome 1
    else ValueNone
  override a.Equals(o: obj) = match o with :? Point as b -> a.x = b.x && a.y = b.y
  override a.GetHashCode() = HashCode.Combine(a.x.GetHashCode(), a.y.GetHashCode())
  override a.ToString() = sprintf "(%s,%s)" (string a.x) (string a.y)
  interface Spatial with
    member this.Boundary = { low = this; high = this }
  
and [<Struct;NoComparison>] Rect =
  { low: Point
    high: Point }
  override a.ToString() = sprintf "{%s,%s}" (string a.low) (string a.high)
  interface Spatial with
    member this.Boundary = this

[<RequireQualifiedAccess>]
module Point =
  
  let gt (a: Point) (b: Point) : bool = a.Compare b = ValueSome 1
  
  let lt (a: Point) (b: Point) : bool = a.Compare b = ValueSome -1
  
  let eq (a: Point) (b: Point) : bool = a.Compare b = ValueSome 0
  
  let gte (a: Point) (b: Point) : bool = match a.Compare b with ValueSome 1 | ValueSome 0 -> true | _ -> false
  
  let lte (a: Point) (b: Point) : bool = match a.Compare b with ValueSome -1 | ValueSome 0 -> true | _ -> false
  
  let inline min a b = { x = Math.Min(a.x, b.x); y = Math.Min(a.y, b.y) }
  
  let inline max a b = { x = Math.Max(a.x, b.x); y = Math.Max(a.y, b.y) }
  
  let inline mean a b = { x = (a.x + b.x) / 2.0; y = (a.y + b.y) / 2.0 }
  
  let inline anyMatch a b = a.x = b.x || a.y = b.y

[<RequireQualifiedAccess>]
module Rect =
  
  let inline create a b c d : Rect =
    let low = { x = a; y = b }
    let high = { x = c; y = d }
    assert (Point.gte high low)
    { low = low; high = high }
    
  /// Checks if `a` is completely covering `b`.
  let contains (a: Rect) (b: Rect) : bool when 'k :> Spatial =
    Point.lte a.low b.low && Point.gte a.high b.high
    
  /// Checks if twp rectangles intersect with one another.
  let intersects (a: Rect) (b: Rect) : bool =
    not (Point.gt a.low b.high || Point.lt a.high b.low)
    
  /// Returns a `Rect` which contains both rectangles.
  let wrap (a: Rect) (b: Rect) : Rect =
    { low = Point.min a.low b.low ; high = Point.max a.high b.high }
    
  /// Returns an area of a given rectangle. 
  let area (a: Rect) : float =
    let p = (a.high - a.low)
    p.x * p.y    
  
[<Struct>]
type Node<'k,'v when 'k :> Spatial> =
  { level: int; entries: Entry<'k,'v>[] }
  override n.ToString() = Array.string n.entries

and Entry<'k,'v when 'k :> Spatial> =
  | Leaf of key:'k * value:'v
  | Branch of key:Rect * Node<'k,'v>
  override e.ToString() =
    match e with
    | Leaf(key, value) -> sprintf "%s => %s" (string key) (string value)
    | Branch(key, node) -> sprintf "%s => %s" (string key) (string node)
  
type RTree<'k,'v when 'k :> Spatial> =
  { root: Node<'k,'v>
    config: Config }

type Group<'k,'v  when 'k :> Spatial>  = ValueTuple<Entry<'k,'v>[], Rect>
  
/// Struct storing statistics for a single dimension 
[<Struct;IsByRefLike>]
type private DimStats =
  { mutable minLow: float
    mutable maxHigh: float
    
    mutable maxLow: float
    mutable lowIndex: int
    
    mutable minHigh: float
    mutable highIndex: int }

[<RequireQualifiedAccess>]
module private DimStats =
  let create () =
    { minLow = Double.MaxValue
      maxHigh = Double.MinValue
      maxLow = Double.MinValue
      minHigh = Double.MaxValue
      lowIndex = 0
      highIndex = 0 }
    
  let inline vertFarest (s: DimStats inref) = Math.Abs(s.maxHigh - s.minLow)
  
  let inline vertNearest (s: DimStats inref) = Math.Abs(s.minHigh - s.maxLow)
  
[<RequireQualifiedAccess>]
module private Entry =
  
  let inline boundary e = 
    match e with
    | Leaf(key, _) -> key.Boundary
    | Branch(rect, _) -> rect
    
[<RequireQualifiedAccess>]
module private SplitStrategy =
  
  let inline private computeDim (s: DimStats byref) (lo: float) (hi: float) (i: int) =
    s.minLow <- Math.Min(s.minLow, lo)
    s.maxHigh <- Math.Max(s.maxHigh, hi)
    if lo > s.maxLow then
      s.maxLow <- lo
      s.lowIndex <- i
    if hi < s.minHigh then
      s.minHigh <- hi
      s.highIndex <- i
  
  let linear (entries: #IReadOnlyList<Entry<'k,'v>>) =
    let mutable x = 0
    let mutable y = 1
    /// Entry statistics on X dimension
    let mutable dx = DimStats.create ()
    /// Entry statistics on Y dimension
    let mutable dy = DimStats.create ()
    
    if entries.Count > 2 then
      for i in 0..(entries.Count-1) do
        let e = entries.[i]
        let rect = Entry.boundary e
        computeDim &dx rect.low.x rect.high.x i
        computeDim &dy rect.low.y rect.high.y i
    
    let lenX, lenY = DimStats.vertFarest &dx, DimStats.vertFarest &dy
    let sepX, sepY = DimStats.vertNearest &dx, DimStats.vertNearest &dy
    let normX, normY = (sepX / lenX), (sepY / lenY)
    let lowIdx, highIdx = if normX > normY then dx.lowIndex, dx.highIndex else dy.lowIndex, dy.highIndex
    x <- lowIdx
    y <- highIdx
    
    let cmp = x.CompareTo y
    if cmp < 0 then struct(x, y)
    elif cmp > 0 then struct(y, x)
    elif x = 0 then struct(0, 1)
    else struct(0, y)
  
  let quadratic (entries: #IReadOnlyList<Entry<'k,'v>>) =
    let mutable x = 0
    let mutable y = 1
    let mutable maxDiff = Double.MinValue
    if entries.Count > 2 then
      for i in 0..(entries.Count-1) do
        for j in 1..(entries.Count-1) do
          let first = Entry.boundary entries.[i]
          let second = Entry.boundary entries.[j]
          let combined = Rect.wrap first second
          let diff = (Rect.area combined) - (Rect.area first) - (Rect.area second)
          if diff > maxDiff then
            maxDiff <- diff
            x <- i
            y <- j
    struct(x, y)
  
  let inline split (entries: #IReadOnlyList<Entry<'k,'v>>) (strategy: SplitStrategy) =
    match strategy with
    | SplitStrategy.Linear  -> linear entries
    | SplitStrategy.Quadratic -> quadratic entries
    
  let inline nextLinear (entries: #IReadOnlyList<Entry<'k,'v>>) boundary =
    let rect = Rect.wrap boundary (Entry.boundary entries.[0])
    struct(0, rect, 0)
  
  let private preference (e: Entry<'k,'v>) rect1 rect2 =
    let bound = (Entry.boundary e)
    let wrap1 = Rect.wrap rect1 bound
    let diff1 = Rect.area wrap1 - Rect.area rect1
    let wrap2 = Rect.wrap rect2 bound
    let diff2 = Rect.area wrap2 - Rect.area rect2
    struct(diff1, wrap1, diff2, wrap2)
  
  let private selectGroup rect1 rect2 len1 len2 diff1 diff2 =
    if diff1 < diff2 then 0
    elif diff2 < diff1 then 1
    elif Rect.area rect1 < Rect.area rect2 then 0
    elif Rect.area rect2 < Rect.area rect1 then 1
    elif len1 < len2 then 0
    elif len2 < len1 then 1
    else 0
  
  let inline nextQuadratic (entries: #IReadOnlyList<Entry<'k,'v>>) rect1 rect2 len1 len2 =
    let mutable idx = 0
    let mutable e = entries.[idx]
    let struct(pref1, exp1, pref2, exp2) = preference e rect1 rect2
    let mutable maxPreferenceDiff = Math.Abs(pref1 - pref2)
    let mutable group = selectGroup rect1 rect2 len1 len2 pref1 pref2
    let mutable expanded = if group = 0 then exp1 else exp2
    
    for i in 1..(entries.Count - 1) do
      let e = entries.[i]
      let struct(pref1, exp1, pref2, exp2) = preference e rect1 rect2
      let preferenceDiff = Math.Abs(pref1 - pref2)
      if maxPreferenceDiff <= preferenceDiff then
        maxPreferenceDiff <- preferenceDiff
        idx <- i
        group <- selectGroup rect1 rect2 len1 len2 pref1 pref2
        expanded <- if group = 0 then exp1 else exp2
    
    struct(idx, expanded, group)
  
[<Struct>]
type private Split<'k,'v when 'k :> Spatial> =
  // Node was updated but no need for splitting it up
  | NoSplit of node:Node<'k,'v>
  // Node was updated, but number of children surpassed the threshold, therefore it needs to be split
  | Split of left:Entry<'k,'v> * right:Entry<'k,'v>
  
[<RequireQualifiedAccess>]
module private Node =
  
  let root =
    { entries = [||]
      level = 0 }
    
  let inline isLeaf n = n.level = 0
  
  let length n = n.entries.Length
      
  let find (area: Rect) (n: Node<'k,'v>) =
    let rec loop (acc: ResizeArray<'v>) (area: Rect) (n: Node<'k,'v>) =
      for e in n.entries do
        match e with
        | Leaf(key, value) when Rect.contains area key.Boundary -> acc.Add value
        | Branch(r, child) when Rect.intersects area r -> loop acc area child
        | _ -> ()

    let found = ResizeArray()
    loop found area n
    found.ToArray()
    
  let split (config: Config) (n: Node<'k,'v>) =
    let split (config: Config) (entries: Entry<'k,'v>[]) : ValueTuple<Group<'k,'v>,Group<'k,'v>> =
      let struct(i1, i2) = SplitStrategy.split entries config.strategy
      
      let entries = ResizeArray(entries)
      // `i2` is always greater than `i1` 
      let seed2 = entries.[i2]
      let seed1 = entries.[i1]
      entries.RemoveAt i2
      entries.RemoveAt i1
      
      let mutable bound1 = Entry.boundary seed1
      let group1 = ResizeArray<_>(config.minCap)
      group1.Add seed1

      let mutable bound2 = Entry.boundary seed2
      let group2 = ResizeArray<_>(config.minCap)
      group2.Add seed2
      
      while entries.Count > 0 do
        if entries.Count + group1.Count = config.minCap then
          for e in entries do
            bound1 <- Rect.wrap bound1 (Entry.boundary e)
            group1.Add e
          entries.Clear ()
        elif entries.Count + group2.Count = config.minCap then
          for e in entries do
            bound2 <- Rect.wrap bound2 (Entry.boundary e)
            group2.Add e
          entries.Clear ()
        else
          let struct(idx, expanded, group) =
            match config.strategy with
            | SplitStrategy.Linear -> SplitStrategy.nextLinear entries bound1
            | SplitStrategy.Quadratic -> SplitStrategy.nextQuadratic entries bound1 bound2 group1.Count group2.Count
            
          let e = entries.[idx]
          entries.RemoveAt idx
          if group = 0 then
            bound1 <- expanded
            group1.Add e
          else
            bound2 <- expanded
            group2.Add e

      let first: Group<'k,'v> = struct(group1.ToArray(), bound1)
      let second: Group<'k,'v> = struct(group2.ToArray(), bound2)
      printfn "split group in two\n\t%s=%s\n\t%s=%s" (string bound1) (Array.string (group1.ToArray())) (string bound1) (Array.string (group2.ToArray()))
      struct(first, second)
      
    let struct(struct(group1, bound1), struct(group2, bound2)) = split config n.entries
    let first  = Entry.Branch(bound1, { level = n.level; entries = group1 })
    let second = Entry.Branch(bound2, { level = n.level; entries = group2 })
    struct(first, second)
    
  let private addLeaf config item n =
    let rect = Entry.boundary item
    let idx = n.entries |> Array.tryFindIndex (function Leaf(k, _) -> k.Boundary = rect | _ -> false)
    match idx with
    | None ->
      printfn "add %s at %s" (string item) (string n)
      let n = { n with entries = Array.insertAt (Array.length n.entries) item n.entries }
      if n.entries.Length > config.maxCap then
        let struct(l, r) = split config n
        Split(l,r)
      else 
        NoSplit n
    | Some idx ->
      printfn "replace %s at index %i at %s" (string item) idx (string n)
      let n = { n with entries = Array.replace idx item n.entries }
      NoSplit n
          
  let private addBranch item n =
    let mutable minEntry = n.entries.[0]
    let mutable minEntryIdx = 0
    let itemBoundary = (Entry.boundary item)
    let mutable minRect = Rect.wrap (Entry.boundary minEntry) itemBoundary
    let mutable minDiff = Rect.area minRect - Rect.area (Entry.boundary minEntry)
    for i in 1..(n.entries.Length-1) do
      let e = n.entries.[i]
      let bound = (Entry.boundary e)
      let expanded = Rect.wrap (Entry.boundary item) bound
      let diff = Rect.area expanded - Rect.area bound
      if diff < minDiff then
        minDiff <- diff
        minRect <- expanded
        minEntry <- e
        minEntryIdx <- i
    let (Branch(r, child)) = minEntry
    printfn "add branch %s at %s" (string item) (string r)
    (child, minRect, minEntryIdx)
    
    
  let rec add (config: Config) level (item: Entry<'k,'v>) (n: Node<'k,'v>) : Split<'k,'v> =
    match item with
    | Branch _ when level = n.level ->
      let n = { n with entries = Array.insertAt (Array.length n.entries) item n.entries }
      if n.entries.Length > config.maxCap then
        let struct(l, r) = split config n
        Split(l,r)
      else 
        NoSplit n
    | _ ->
      if isLeaf n then
        addLeaf config item n
      else
        let (child, minRect, minEntryIdx) = addBranch item n
        match add config level item child with
        | Split(first, second) ->
          let n = { n with entries = Array.append (Array.removeAt minEntryIdx n.entries) [|first;second|] }
          if n.entries.Length > config.maxCap then
            let struct(l, r) = split config n
            Split(l,r)
          else 
            NoSplit n
        | NoSplit node ->
          let e = Branch(minRect, node)
          let n = { n with entries = Array.replace minEntryIdx e n.entries }
          NoSplit n
  
  let rec remove (config: Config) rect (n: Node<'k,'v>) =
    if isLeaf n then
      let idx =
        n.entries
        |> Array.findIndex (function Leaf(k,_) -> rect = k.Boundary | _ -> false)
      
      if idx = -1 then
        (n, None, None)
      else
        let removed = n.entries.[idx]
        let n = { n with entries = Array.removeAt idx n.entries }
        (n, Some removed, None)
    else
      let mutable entryIdx = -1
      let mutable removedOption = None
      let mutable entries = Array.copy n.entries
      
      let mutable i = 0
      while i < entries.Length do
        let e = entries.[i]
        let r = Entry.boundary e
        if Rect.contains r rect then
          let (Branch(r, child)) = e
          let (child, removed, orphans) = remove config rect child
          match removed with
          | None -> ()
          | Some key ->
            removedOption <- Some((removed, orphans))
            let removedRect: Rect = Entry.boundary key
            let r =
              if Point.anyMatch removedRect.low r.low || Point.anyMatch removedRect.high r.high then
                child.entries
                |> Array.map Entry.boundary
                |> Array.reduce Rect.wrap
              else r
            entries.[i] <- Branch(r, child)
            if child.entries.Length < config.minCap then
              entryIdx <- i
              i <- entries.Length // break;
        i <- i + 1
             
      let n = { n with entries = entries } 
      match removedOption with
      | None -> (n, None, None)
      | Some(removed, orphans) when entryIdx < 0 -> (n, removed, orphans)
      | Some(removed, orphans) ->
        let orphan = n.entries.[entryIdx]
        let n = { n with entries = Array.removeAt entryIdx n.entries }        
        let orphans = Array.append (Option.defaultValue [||] orphans) [|orphan|]
        (n, removed, Some orphans)
  
[<RequireQualifiedAccess>]
module RTree =
  
  let create (minCap: int) (maxCap: int) (strategy: SplitStrategy) =
    assert (minCap <= maxCap)
    let config =
      { minCap = minCap
        maxCap = maxCap
        strategy = strategy }
    { root = Node.root
      config = config }
    
  /// Returns a list of spatial items stored in current RTree, that exists within
  /// the bounds of provided `area`. 
  let within (area: Rect) (tree: RTree<'k,'v>) : 'v[] =
    Node.find area tree.root
    
  let private insert config level item (node: Node<'k,'v>) =
    match Node.add config level item node with
    | NoSplit n   -> n
    | Split(l, r) -> { entries = [|l;r|]; level = node.level + 1 }
  
  let remove (key: 'k) (tree: RTree<'k,'v>) =
    let rect = key.Boundary
    let (root, removed, orphans) = Node.remove tree.config rect tree.root
    let mutable root =
      if root.entries.Length = 1 && not (Node.isLeaf root) then
        match root.entries.[0] with
        | Branch(_, n) -> n
        | _ -> root
      else root
    match orphans with
    | None -> ()
    | Some orphans ->
      for orphan in orphans do
        match orphan with
        | Leaf _ ->
          root <- insert tree.config 0 orphan root
        | Branch(_, n) ->
          for e in n.entries do
            root <- insert tree.config n.level e root
    { tree with root = root }
  
  let add key item (tree: RTree<'k,'v>) =
    let e = Leaf(key, item)
    let root = insert tree.config 0 e tree.root
    { tree with root = root }

RTreeTests.fs

module Demos.RTree.Tests

open Expecto
open Demos.RTree
let private tree =
    let cities =
        [ (-3.11, 55.57), "Edinburgh"
          (12.20, 45.26 ), "Venice"
          (-6.16, 53.21), "Dublin"
          (14.25, 50.05), "Prague"
          (2.21, 48.51), "Paris"
          (-2.59, 53.24), "Liverpool"
          (-2.14, 53.28), "Manchester"
          (4.54, 52.22), "Amsterdam"
          (4.29, 51.56), "Rotterdam"
          (-1.15, 51.45), "Oxford"
          (0.07, 52.12), "Cambridge"
          (0.0, 51.3), "London" ]
    let tree = RTree.create 2 4 SplitStrategy.Quadratic
    cities |> List.fold (fun acc ((x,y), city) -> RTree.add {x=x;y=y} city acc) tree

[<Tests>]
let tests = testList "RTree" [
    test "find (not found)" {
        // Rect from Cracow to Vilnus
        let rect = Rect.create 19.56 50.04 25.17 54.41
        let actual = tree |> RTree.within rect
        Expect.isEmpty actual $"RTree.find within bounds ${rect} should not be found"
    }
    
    test "find 1 element" {
        // Rect from Munich to Varsaw
        let rect = Rect.create 14.25 50.05 14.25 50.05
        let actual =
            tree
            |> RTree.within rect
        Expect.equal [|"Prague"|] actual $"RTree.find within bounds ${rect} should find 1 element"
    }
    
    test "find 5 element" {
        let rect = Rect.create -4.09 50.22 1.18 54.58
        let actual =
            tree
            |> RTree.within rect
            |> Set.ofArray
        let expected = Set.ofList [ "London"; "Oxford"; "Cambridge"; "Liverpool"; "Manchester" ]
        Expect.equal expected actual $"RTree within bounds ${rect} should find 5 elements"
    }
    
    test "remove" {
        let tree = tree |> RTree.remove { x = 14.25; y = 50.05 }
        let rect = Rect.create 14.25 50.05 14.25 50.05
        let actual = tree |> RTree.within rect
        Expect.isEmpty actual $"RTree should not found any items within removed bounds, but {actual.Length} found"
    }
    
    test "replace" {
        let tree = tree |> RTree.add { x = 14.25; y = 50.05 } "Prague (2)"
        let rect = Rect.create 14.25 50.05 14.25 50.05
        let actual = tree |> RTree.within rect
        Expect.equal [|"Prague (2)"|] actual $"RTree.find within bounds ${rect} should find 1 element"
    }
]