rtree in elm

RTree.elm

module RTree exposing (..)

import BoundingBox2d exposing (BoundingBox2d)
import Point2d exposing (Point2d)
import List.Extra as List


-- bounding boxes can go to infintiy
-- minNode <= maxNode // 2
-- the root node has at least 2 children, unless it is a leaf
-- all leaves appear on the same level


type RTree a
    = Empty
    | NonEmpty (NonEmptyRTree a)


type NonEmptyRTree a
    = Leaf { boundingBox : BoundingBox2d, values : List ( BoundingBox2d, a ) }
    | Node { boundingBox : BoundingBox2d, children : ( NonEmptyRTree a, List (NonEmptyRTree a) ) }


map : (a -> b) -> RTree a -> RTree b
map f rtree =
    case rtree of
        Empty ->
            Empty

        NonEmpty nonEmpty ->
            NonEmpty (mapNonEmpty f nonEmpty)


mapNonEmpty : (a -> b) -> NonEmptyRTree a -> NonEmptyRTree b
mapNonEmpty f rtree =
    case rtree of
        Leaf leaf ->
            Leaf { boundingBox = leaf.boundingBox, values = List.map (Tuple.mapSecond f) leaf.values }

        Node node ->
            let
                ( first, rest ) =
                    node.children
            in
                Node { boundingBox = node.boundingBox, children = ( mapNonEmpty f first, List.map (mapNonEmpty f) rest ) }


depthNonEmpty : NonEmptyRTree a -> Int
depthNonEmpty nonEmpty =
    case nonEmpty of
        Leaf leaf ->
            List.length leaf.values

        Node node ->
            let
                ( first, rest ) =
                    node.children

                folder item accum =
                    accum + depthNonEmpty item
            in
                depthNonEmpty first + List.foldl folder 0 rest


depth : RTree a -> Int
depth rtree =
    case rtree of
        Empty ->
            0

        NonEmpty subtree ->
            depthNonEmpty subtree


empty : RTree a
empty =
    Empty


singleton : BoundingBox2d -> a -> RTree a
singleton boundingBox value =
    NonEmpty (Leaf { boundingBox = boundingBox, values = [ ( boundingBox, value ) ] })


getBoundingBox : NonEmptyRTree a -> BoundingBox2d
getBoundingBox rtree =
    case rtree of
        Leaf { boundingBox } ->
            boundingBox

        Node { boundingBox } ->
            boundingBox


insert : BoundingBox2d -> a -> RTree a -> RTree a
insert boundingBox value rtree =
    case rtree of
        Empty ->
            singleton boundingBox value

        NonEmpty nonEmpty ->
            NonEmpty (insertNonEmpty boundingBox value nonEmpty)


insertNonEmpty : BoundingBox2d -> a -> NonEmptyRTree a -> NonEmptyRTree a
insertNonEmpty boundingBox value rtree =
    case insertNonEmptyHelper boundingBox value rtree of
        Done v ->
            v

        Split x y ->
            Node
                { boundingBox = BoundingBox2d.hull (getBoundingBox x) (getBoundingBox y)
                , children = ( x, [ y ] )
                }


type Split a
    = Done a
    | Split a a


maxNodes : Int
maxNodes =
    4


minNodes : Int
minNodes =
    2


maxByAndRest : (a -> comparable) -> a -> List a -> ( a, List a )
maxByAndRest toComparable default list =
    let
        go items ( smallestSoFar, smallestItemSoFar, accum ) =
            case items of
                item :: remaining ->
                    let
                        current =
                            toComparable item
                    in
                        if current > smallestSoFar then
                            go remaining ( current, item, smallestItemSoFar :: accum )
                        else
                            go remaining ( smallestSoFar, smallestItemSoFar, item :: accum )

                [] ->
                    ( smallestItemSoFar, accum )
    in
        go list ( toComparable default, default, [] )


minByAndRest : (a -> comparable) -> a -> List a -> ( a, List a )
minByAndRest toComparable default list =
    let
        go items ( smallestSoFar, smallestItemSoFar, accum ) =
            case items of
                item :: remaining ->
                    let
                        current =
                            toComparable item
                    in
                        if current < smallestSoFar then
                            go remaining ( current, item, smallestItemSoFar :: accum )
                        else
                            go remaining ( smallestSoFar, smallestItemSoFar, item :: accum )

                [] ->
                    ( smallestItemSoFar, accum )
    in
        go list ( toComparable default, default, [] )


insertNonEmptyHelper : BoundingBox2d -> a -> NonEmptyRTree a -> Split (NonEmptyRTree a)
insertNonEmptyHelper boundingBox value rtree =
    case rtree of
        Leaf leaf ->
            if List.length leaf.values >= maxNodes then
                let
                    ( a, b ) =
                        quadraticSplit (( boundingBox, value ) :: leaf.values)
                in
                    Split (Leaf a) (Leaf b)
            else
                Done
                    (Leaf
                        { boundingBox = BoundingBox2d.hull leaf.boundingBox boundingBox
                        , values = ( boundingBox, value ) :: leaf.values
                        }
                    )

        Node node ->
            let
                ( first, rest ) =
                    node.children

                ( x, xs ) =
                    minByAndRest (\subtree -> boundingBoxDelta (getBoundingBox subtree) boundingBox) first rest
            in
                case insertNonEmptyHelper boundingBox value x of
                    Done newFirst ->
                        let
                            newBoundingBox =
                                if (getBoundingBox newFirst) == (getBoundingBox x) then
                                    -- bounding box didn't change
                                    node.boundingBox
                                else
                                    List.foldl (\subtree -> BoundingBox2d.hull (getBoundingBox subtree)) (getBoundingBox newFirst) xs
                        in
                            Done
                                (Node
                                    { boundingBox = newBoundingBox
                                    , children = ( newFirst, xs )
                                    }
                                )

                    Split left right ->
                        let
                            size =
                                2 + List.length xs
                        in
                            if size > maxNodes then
                                let
                                    ( a, b ) =
                                        quadraticSplitNode (left :: right :: xs)

                                    newLeft =
                                        case a.values of
                                            y :: ys ->
                                                { boundingBox = a.boundingBox, children = ( y, ys ) }

                                            _ ->
                                                Debug.crash ""

                                    newRight =
                                        case b.values of
                                            y :: ys ->
                                                { boundingBox = b.boundingBox, children = ( y, ys ) }

                                            _ ->
                                                Debug.crash ""
                                in
                                    Split (Node newLeft) (Node newRight)
                            else
                                let
                                    newBoundingBox =
                                        if BoundingBox2d.isContainedIn node.boundingBox boundingBox then
                                            node.boundingBox
                                        else
                                            List.foldl (\subtree -> BoundingBox2d.hull (getBoundingBox subtree))
                                                (BoundingBox2d.hull (getBoundingBox left) (getBoundingBox right))
                                                xs
                                in
                                    Done
                                        (Node
                                            { boundingBox = newBoundingBox
                                            , children = ( left, right :: xs )
                                            }
                                        )


cartesianProduct list =
    case list of
        [] ->
            []

        x :: xs ->
            List.map (\v -> ( x, v )) xs ++ cartesianProduct xs


pickSeedsBy : (a -> BoundingBox2d) -> List a -> Maybe ( a, a )
pickSeedsBy toBoundingBox nodes =
    let
        pairs =
            cartesianProduct nodes

        costFunction ( a, b ) =
            boundingBoxDistance (toBoundingBox a) (toBoundingBox b)
    in
        List.maximumBy costFunction pairs


pickNextBy : (a -> BoundingBox2d) -> BoundingBox2d -> BoundingBox2d -> List a -> Maybe ( a, List a )
pickNextBy toBoundingBox bb bb_ nodes =
    case nodes of
        [] ->
            Nothing

        x :: xs ->
            let
                diff value =
                    let
                        bb__ =
                            toBoundingBox value
                    in
                        abs ((boundingBoxDelta bb bb__) - (boundingBoxDelta bb_ bb__))
            in
                Just (maxByAndRest diff x xs)


split :
    (a -> BoundingBox2d)
    -> BoundingBox2d
    -> List a
    -> BoundingBox2d
    -> List a
    -> List a
    -> ( { boundingBox : BoundingBox2d, values : List a }, { boundingBox : BoundingBox2d, values : List a } )
split toBoundingBox ns_bb ns ms_bb ms nodes =
    case nodes of
        [] ->
            ( { boundingBox = ns_bb, values = ns }
            , { boundingBox = ms_bb, values = ms }
            )

        _ ->
            case pickNextBy toBoundingBox ns_bb ms_bb nodes of
                Nothing ->
                    ( { boundingBox = ns_bb, values = ns }
                    , { boundingBox = ms_bb, values = ms }
                    )

                Just ( first, rest ) ->
                    let
                        bb =
                            toBoundingBox first

                        delta_n =
                            boundingBoxDelta bb ns_bb

                        delta_m =
                            boundingBoxDelta bb ms_bb
                    in
                        if delta_n < delta_m then
                            split toBoundingBox (BoundingBox2d.hull ns_bb bb) (first :: ns) ms_bb ms rest
                        else
                            split toBoundingBox ns_bb ns (BoundingBox2d.hull ms_bb bb) (first :: ms) rest


{-| the name is a lie! just splitting evenly yields much better results for our input
-}
quadraticSplit :
    List ( BoundingBox2d, a )
    -> ( { boundingBox : BoundingBox2d, values : List ( BoundingBox2d, a ) }, { boundingBox : BoundingBox2d, values : List ( BoundingBox2d, a ) } )
quadraticSplit values =
    case values of
        [] ->
            Debug.crash "Can't split empty list"

        _ ->
            case pickSeedsBy Tuple.first values of
                Nothing ->
                    Debug.crash "Can't split empty list"

                Just ( ( bb1, _ ) as s1, ( bb2, _ ) as s2 ) ->
                    let
                        rest =
                            List.filter (\n -> n /= s1 && n /= s2) values
                    in
                        split Tuple.first bb1 [ s1 ] bb2 [ s2 ] rest


{-| the name is a lie! just splitting evenly yields much better results for our input
-}
quadraticSplitNode :
    List (NonEmptyRTree a)
    -> ( { boundingBox : BoundingBox2d, values : List (NonEmptyRTree a) }, { boundingBox : BoundingBox2d, values : List (NonEmptyRTree a) } )
quadraticSplitNode values =
    case values of
        [] ->
            Debug.crash "Can't split empty list"

        _ ->
            case pickSeedsBy getBoundingBox values of
                Nothing ->
                    Debug.crash "Can't split empty list"

                -- Just ( ( bb1, _ ) as s1, ( bb2, _ ) as s2 ) ->
                Just ( option1, option2 ) ->
                    let
                        rest =
                            List.filter (\n -> n /= option1 && n /= option2) values
                    in
                        split getBoundingBox (getBoundingBox option1) [ option1 ] (getBoundingBox option2) [ option2 ] rest


cheapSplit :
    List ( BoundingBox2d, a )
    -> ( { boundingBox : BoundingBox2d, values : List ( BoundingBox2d, a ) }, { boundingBox : BoundingBox2d, values : List ( BoundingBox2d, a ) } )
cheapSplit values =
    let
        n =
            List.length values // 2

        xs =
            List.take n values

        ys =
            List.drop n values
    in
        case ( List.map Tuple.first xs |> BoundingBox2d.aggregate, List.map Tuple.first ys |> BoundingBox2d.aggregate ) of
            ( Just bb1, Just bb2 ) ->
                ( { boundingBox = bb1, values = xs }
                , { boundingBox = bb2, values = ys }
                )

            _ ->
                Debug.crash "not enough items"


{-| Remove a particular bounding box from the tree
-}
remove : BoundingBox2d -> RTree a -> RTree a
remove boundingBox rtree =
    case rtree of
        Empty ->
            rtree

        NonEmpty nonEmpty ->
            Tuple.second (removeNonEmpty boundingBox nonEmpty)


removeValues : BoundingBox2d -> List ( BoundingBox2d, a ) -> List ( BoundingBox2d, a ) -> Bool -> ( Bool, List ( BoundingBox2d, a ) )
removeValues target values accum dirty =
    case values of
        [] ->
            ( dirty, accum )

        (( firstBoundingBox, _ ) as first) :: rest ->
            if firstBoundingBox == target then
                removeValues target rest accum True
            else
                removeValues target rest (first :: accum) dirty


removeNonEmpty : BoundingBox2d -> NonEmptyRTree a -> ( Bool, RTree a )
removeNonEmpty boundingBox rtree =
    case rtree of
        Leaf leaf ->
            if BoundingBox2d.isContainedIn leaf.boundingBox boundingBox then
                let
                    ( dirty, newValues ) =
                        removeValues boundingBox leaf.values [] False
                in
                    if dirty then
                        case newValues of
                            [] ->
                                ( dirty, Empty )

                            x :: xs ->
                                let
                                    newBoundingBox =
                                        List.foldl (\( bbox, _ ) -> BoundingBox2d.hull bbox) (Tuple.first x) xs
                                in
                                    ( True, NonEmpty (Leaf { boundingBox = newBoundingBox, values = newValues }) )
                    else
                        ( False, NonEmpty (Leaf leaf) )
            else
                ( False, NonEmpty (Leaf leaf) )

        Node node ->
            if BoundingBox2d.isContainedIn node.boundingBox boundingBox then
                let
                    ( first, rest ) =
                        node.children

                    baseCase accum dirty changedBoundingBox =
                        if not dirty then
                            ( False, NonEmpty (Node node) )
                        else
                            case accum of
                                [] ->
                                    ( True, Empty )

                                [ x ] ->
                                    ( True, NonEmpty x )

                                x :: xs ->
                                    let
                                        newBoundingBox =
                                            if changedBoundingBox then
                                                List.foldl (\subtree -> BoundingBox2d.hull (getBoundingBox subtree)) (getBoundingBox x) xs
                                            else
                                                node.boundingBox

                                        newNode =
                                            Node { boundingBox = newBoundingBox, children = ( x, xs ) }
                                    in
                                        ( True, NonEmpty newNode )

                    go children accum dirty changedBoundingBox =
                        case children of
                            [] ->
                                baseCase accum dirty changedBoundingBox

                            firstChild :: remainingChildren ->
                                case removeNonEmpty boundingBox firstChild of
                                    ( _, Empty ) ->
                                        go remainingChildren accum True True

                                    ( True, NonEmpty nonEmpty ) ->
                                        if getBoundingBox nonEmpty == getBoundingBox firstChild then
                                            baseCase (nonEmpty :: remainingChildren ++ accum) True changedBoundingBox
                                        else
                                            baseCase (nonEmpty :: remainingChildren ++ accum) True True

                                    ( False, NonEmpty nonEmpty ) ->
                                        go remainingChildren (nonEmpty :: accum) dirty changedBoundingBox
                in
                    go (first :: rest) [] False False
            else
                ( False, NonEmpty (Node node) )


containing : Point2d -> RTree a -> List ( BoundingBox2d, a )
containing point rtree =
    case rtree of
        Empty ->
            []

        NonEmpty nonEmpty ->
            containingNonEmptyHelper point [ nonEmpty ] []


containingNonEmptyHelper : Point2d -> List (NonEmptyRTree a) -> List ( BoundingBox2d, a ) -> List ( BoundingBox2d, a )
containingNonEmptyHelper target toVisit accumulator =
    case toVisit of
        [] ->
            accumulator

        (Node node) :: remainingTrees ->
            if BoundingBox2d.contains target node.boundingBox then
                let
                    ( first, rest ) =
                        node.children
                in
                    containingNonEmptyHelper target (first :: rest ++ remainingTrees) accumulator
            else
                containingNonEmptyHelper target remainingTrees accumulator

        (Leaf leaf) :: remainingTrees ->
            if BoundingBox2d.contains target leaf.boundingBox then
                let
                    newlyFound =
                        List.filter (\( valueBoundingBox, _ ) -> BoundingBox2d.contains target valueBoundingBox) leaf.values
                in
                    containingNonEmptyHelper target remainingTrees (newlyFound ++ accumulator)
            else
                containingNonEmptyHelper target remainingTrees accumulator


concatMap : (a -> List b) -> List a -> List b
concatMap toBs =
    List.foldl (\a accum -> toBs a ++ accum) []


search : BoundingBox2d -> RTree a -> List ( BoundingBox2d, a )
search target rtree =
    case rtree of
        Empty ->
            []

        NonEmpty nonEmpty ->
            searchNonEmpty target nonEmpty


searchNonEmpty : BoundingBox2d -> NonEmptyRTree a -> List ( BoundingBox2d, a )
searchNonEmpty target rtree =
    case rtree of
        Node node ->
            if BoundingBox2d.intersects target node.boundingBox then
                let
                    ( first, rest ) =
                        node.children
                in
                    concatMap (searchNonEmpty target) (first :: rest)
            else
                []

        Leaf leaf ->
            if BoundingBox2d.intersects target leaf.boundingBox then
                List.filter (\( valueBoundingBox, _ ) -> BoundingBox2d.intersects target valueBoundingBox) leaf.values
            else
                []


toList : RTree a -> List ( BoundingBox2d, a )
toList rtree =
    case rtree of
        Empty ->
            []

        NonEmpty nonEmpty ->
            toListNonEmpty nonEmpty


toListNonEmpty : NonEmptyRTree a -> List ( BoundingBox2d, a )
toListNonEmpty rtree =
    case rtree of
        Leaf leaf ->
            leaf.values

        Node node ->
            let
                ( first, rest ) =
                    node.children
            in
                concatMap toListNonEmpty (first :: rest)


boundingBoxArea : BoundingBox2d -> Float
boundingBoxArea boundingBox =
    let
        ( width, height ) =
            BoundingBox2d.dimensions boundingBox
    in
        width * height


boundingBoxDelta : BoundingBox2d -> BoundingBox2d -> Float
boundingBoxDelta bb bb_ =
    boundingBoxArea (BoundingBox2d.hull bb bb_) - boundingBoxArea bb


boundingBoxDistance bb bb_ =
    (boundingBoxArea (BoundingBox2d.hull bb bb_)) - (boundingBoxArea bb) - (boundingBoxArea bb_)