qt.sml

(*  REPRESENTATION CONVENTION: (left, top, right, bottom)
    REPRESENTATION INVARIANT: left < right andalso bottom < top
*)
datatype rectangle = Rect of int * int * int * int;

(*  REPRESENTATION CONVENTION: (extent, horizontal, vertical, 
    topLeft, topRight, bottomLeft, bottomRight)
    REPRESENTATION INVARIANT: 
*)
datatype quadTree = EmptyQuadTree
                  | Qt of rectangle * rectangle list * rectangle list *
                    quadTree * quadTree * quadTree * quadTree;

(*  emptyQtree(e)
    TYPE: rectangle -> quadTree
    PRE: (none)
    POST: empty Qt with extent e
    EXAMPLE:    emptyQtree Rect(~10, 10, 10, ~10) = 
                    Qt(Rect(~10, 10, 10, ~10), [], [], 
                        EmptyQuadTree, EmptyQuadTree, EmptyQuadTree, EmptyQuadTree)
*)
fun emptyQtree e = 
    Qt(e, [], [], EmptyQuadTree, EmptyQuadTree, EmptyQuadTree, EmptyQuadTree)

(*  getOrCreateQtree(q, e)
    TYPE: quadTree * rectangle -> quadTree
    PRE: (none)
    POST: empty Qt with extent e if q = EmptyQuadTree, q otherwise
    EXAMPLE:    getOrCreateQtree(EmptyQuadTree, Rect(~10, 10, 10, ~10)) = 
                    Qt(Rect(~10, 10, 10, ~10), [], [], 
                        EmptyQuadTree, EmptyQuadTree, EmptyQuadTree, EmptyQuadTree)
*)
fun getOrCreateQtree(EmptyQuadTree, e) = emptyQtree e
  | getOrCreateQtree(q, _) = q

(*  TODO: boundaries *)
(*  isPointInRect(r, x, y)
    TYPE: rectangle * int * int -> bool
    PRE: (none)
    POST: true or false depending on if the point (x, y) is within r
    EXAMPLE:    isPointInRect(Rect(~10, 10, 10, ~10), 0, 0) = true
                isPointInRect(Rect(~10, 10, 10, ~10), 10, 10) = false
*)
fun isPointInRect(Rect(l, t, r, b), x, y) =
    l <= x andalso x < r andalso y < t andalso y >= b

(*  TODO: boundaries *)
(*  isRectInRect(r1, r2)
    TYPE: rectangle * rectangle -> bool
    PRE: (none)
    POST: true or false depending on if r1 is within r2
    EXAMPLE:    isRectInRect(Rect(~10, 10, 10, ~10), Rect(~3, 3, 3, ~3)) = true
                isRectInRect(Rect(~10, 10, 10, ~10), Rect(~11, 11, 11, ~11)) = false
*)
fun isRectInRect(Rect(l1, t1, r1, b1), Rect(l2, t2, r2, b2)) =
    l2 >= l1 andalso t2 <= t1 andalso r2 <= r1 andalso b2 >= b1

(*  getRectsOnPoint(rlist, x, y)
    TYPE: rectangle list * int * int -> rectangle list
    PRE: (none)
    POST: rectangles in rlist that contains point (x, y)
    EXAMPLE:    getRectsOnPoint([Rect(~10, 10, 10, ~10), Rect(~3, 3, 3, ~3)], 9, 9) = 
                    [Rect(~10, 10, 10, ~10)]
*)
fun getRectsOnPoint(rectangles, x, y) = 
    List.filter (fn r => isPointInRect(r, x, y)) rectangles

(*  isOnVertical(r1, r2)
    TYPE: rectangle * rectangle -> bool
    PRE: (none)
    POST: true or false depending on if r1 lies on the vertical center line of r2
    EXAMPLE:    isOnVertical(Rect(~10, 10, 10, ~10), Rect(~3, 3, 3, ~3)) = true
*)
fun isOnVertical(Rect(l1, _, r1, _), Rect(l2, _, r2, _)) =
    let
        val cy = (l1 + r1) div 2
    in
        l2 <= cy andalso cy < r2
    end

(*  isOnHorizontal(r1, r2)
    TYPE: rectangle * rectangle -> bool
    PRE: (none)
    POST: true or false depending on if r1 lies on the horizontal center line of r2
    EXAMPLE:    isOnHorizontal(Rect(~10, 10, 10, ~10), Rect(~3, 3, 3, ~3)) = true
*)
fun isOnHorizontal(Rect(_, t1, _, b1), Rect(_, t2, _, b2)) =
    let
        val cx = (t1 + b1) div 2    
    in
        b2 <= cx andalso cx < t2
    end

(*  calculateQuadrants(q)
    TYPE: quadTree -> (rectangle * rectangle * rectangle * rectangle)
    PRE: q <> EmptyQuadTree
    POST: (topLeftRect, topRightRect, bottomLeftRect, bottomRightRect)
*)
fun calculateQuadrants(Qt(Rect(l, t, r, b), _, _, _, _, _, _)) = 
    let
        val cx = (l + r) div 2
        val cy = (t + b) div 2

        val tlRect = Rect(l, t, cx, cy + 1)
        val trRect = Rect(cx + 1, t, r, cy + 1)
        val blRect = Rect(l, cy, cx, b)
        val brRect = Rect(cx + 1, cy, r, b)
    in
        (tlRect, trRect, blRect, brRect)
    end

(*  calculateQuadrants(q, x, y)
    TYPE: quadTree -> quadTree
    PRE: q <> EmptyQuadTree
    POST: the quadrant in q that contains point (x, y)
*)
fun getQuadrantByPosition(qt as Qt(_, _, _, tl, tr, bl, br), x, y) =
    let
        val (tlRect, trRect, blRect, brRect) = calculateQuadrants(qt)
    in
        if isPointInRect(blRect, x, y) then
            bl
        else if isPointInRect(tlRect, x, y) then
            tl
        else if isPointInRect(trRect, x, y) then
            tr
        else
            br
    end

(*  insert(q, r)
    TYPE: quadTree * rectangle -> quadTree
    PRE: q <> EmptyQuadTree
    POST: q updated with r
*)
fun insert(qt as Qt(e, h, v, tl, tr, bl, br), rt) =
    if isOnVertical(e, rt) then
        Qt(e, h, rt::v, tl, tr, bl, br)
    else if isOnHorizontal(e, rt) then
        Qt(e, rt::h, v, tl, tr, bl, br)
    else
        let
            val (tlRect, trRect, blRect, brRect) = calculateQuadrants(qt)
        in
            if isRectInRect(tlRect, rt) then
                Qt(e, h, v, insert(getOrCreateQtree(tl, tlRect), rt), tr, bl, br)
            else if isRectInRect(trRect, rt) then
                Qt(e, h, v, tl, insert(getOrCreateQtree(tr, trRect), rt), bl, br)
            else if isRectInRect(blRect, rt) then
                Qt(e, h, v, tl, tr, insert(getOrCreateQtree(bl, blRect), rt), br)
            else
                Qt(e, h, v, tl, tr, bl, insert(getOrCreateQtree(br, brRect), rt))
        end

(*  query(q, x, y)
    TYPE: quadTree * int * int -> rectangle list
    PRE: (none)
    POST: rectangles in q that contain point (x, y)
*)
fun query(EmptyQuadTree, _, _) = []
  | query(qt as Qt(e as Rect(l, t, r, b), h, v, _, _, _, _), x, y) =
    let
        val rects = getRectsOnPoint(v, x, y)@getRectsOnPoint(h, x, y)

        val cx = (l + r) div 2
        val cy = (t + b) div 2
    in
        (* Abort query as soon we find out that the point lies on
           a center line *)
        if x = cx orelse y = cy then
            rects
        else
            rects@query(getQuadrantByPosition(qt, x, y), x, y)
    end
;