Files
sml-projects/fcore/quad-tree.sml

317 lines
9.5 KiB
Standard ML

signature QUAD_TREE =
sig
type t
val insert: int * int * int * int * int * t -> t
val getCollisions: int * int * int * int * int * t -> int list
val helpGetCollisions: int * int * int * int * int * int list * t -> int list
val hasCollisionAt: int * int * int * int * int * t -> bool
val getItemID: int * int * int * int * t -> int
val create: int * int -> t
end
structure QuadTree: QUAD_TREE =
struct
open QuadTreeType
type item = QuadTreeType.item
fun create (width, height) =
LEAF {items = Vector.fromList [], x = 0, y = 0, w = width, h = height}
fun mkItem (id, startX, startY, width, height) : item =
{ itemID = id
, startX = startX
, startY = startY
, width = width
, height = height
}
type t = QuadTreeType.t
(* max size of vector before we split it further *)
val maxSize = 3
fun mkTopLeft (x, y, w, h, items) =
let
val items = Vector.fromList items
val hw = w div 2
val hh = h div 2
in
LEAF {items = items, x = x, y = y, w = hw, h = hh}
end
fun mkTopRight (x, y, w, h, items) =
let
val items = Vector.fromList items
val hw = w div 2
val hh = h div 2
val x = x + hw
in
LEAF {items = items, x = x, y = y, w = hw, h = hh}
end
fun mkBottomLeft (x, y, w, h, items) =
let
val items = Vector.fromList items
val hw = w div 2
val hh = h div 2
val y = y + hh
in
LEAF {items = items, x = x, y = y, w = hw, h = hh}
end
fun mkBottomRight (x, y, w, h, items) =
let
val items = Vector.fromList items
val hw = w div 2
val hh = h div 2
val x = x + hw
val y = y + hh
in
LEAF {items = items, x = x, y = y, w = hw, h = hh}
end
fun splitLeaf
( x
, y
, w
, h
, tl: item list
, tr: item list
, bl: item list
, br: item list
, elements
, pos
) =
if pos < 0 then
let
val tl = mkTopLeft (x, y, w, h, tl)
val tr = mkTopRight (x, y, w, h, tr)
val bl = mkBottomLeft (x, y, w, h, bl)
val br = mkBottomRight (x, y, w, h, br)
in
NODE
{ topLeft = tl
, topRight = tr
, bottomLeft = bl
, bottomRight = br
, x = x
, y = y
, w = w
, h = h
}
end
else
let
val item = Vector.sub (elements, pos)
val {startX = iX, startY = iY, width = iW, height = iH, ...} = item
val vtl = visitTopLeft (iX, iY, iW, iH, x, y, w, h)
val vtr = visitTopRight (iX, iY, iW, iH, x, y, w, h)
val vbl = visitBottomLeft (iX, iY, iW, iH, x, y, w, h)
val vbr = visitBottomRight (iX, iY, iW, iH, x, y, w, h)
val tl = if vtl then item :: tl else tl
val tr = if vtr then item :: tr else tr
val bl = if vbl then item :: bl else bl
val br = if vbr then item :: br else br
in
splitLeaf (x, y, w, h, tl, tr, bl, br, elements, pos - 1)
end
fun insert (iX, iY, iW, iH, itemID, tree: t) =
case tree of
NODE {topLeft, topRight, bottomLeft, bottomRight, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
let
(* we are not necessarily inserting into all nodes.
* If isCollidingPlus returns false recursively,
* we return the same node back. *)
val tl = insert (iX, iY, iW, iH, itemID, topLeft)
val tr = insert (iX, iY, iW, iH, itemID, topRight)
val bl = insert (iX, iY, iW, iH, itemID, bottomLeft)
val br = insert (iX, iY, iW, iH, itemID, bottomRight)
in
NODE
{ topLeft = tl
, topRight = tr
, bottomLeft = bl
, bottomRight = br
, x = x
, y = y
, w = w
, h = h
}
end
else
tree
| LEAF {items, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
if Vector.length items + 1 > maxSize then
(* have to calculate quadrants and split *)
let
val pos = Vector.length items - 1
val item = mkItem (itemID, iX, iY, iW, iH)
val vtl = visitTopLeft (iX, iY, iW, iH, x, y, w, h)
val vtr = visitTopRight (iX, iY, iW, iH, x, y, w, h)
val vbl = visitBottomLeft (iX, iY, iW, iH, x, y, w, h)
val vbr = visitBottomRight (iX, iY, iW, iH, x, y, w, h)
val tl = if vtl then [item] else []
val tr = if vtr then [item] else []
val bl = if vbl then [item] else []
val br = if vbr then [item] else []
in
splitLeaf (x, y, w, h, tl, tr, bl, br, items, pos)
end
else
(* can insert itemID in items vector *)
let
val item = mkItem (itemID, iX, iY, iW, iH)
val items = Vector.concat [items, Vector.fromList [item]]
in
LEAF {items = items, x = x, y = y, w = w, h = h}
end
else
(* bounds of new item don't fit inside leaf so return old tree *)
tree
fun getCollisionsVec (iX, iY, iW, iH, itemID, pos, elements, acc) =
if pos = Vector.length elements then
acc
else
let
val item = Vector.sub (elements, pos)
val acc =
if isCollidingItem (iX, iY, iW, iH, itemID, item) then
#itemID item :: acc
else
acc
in
getCollisionsVec (iX, iY, iW, iH, itemID, pos + 1, elements, acc)
end
fun getCollisionsAll (iX, iY, iW, iH, itemID, acc, tree) =
case tree of
NODE {topLeft, topRight, bottomLeft, bottomRight, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
let
val acc = getCollisionsAll (iX, iY, iW, iH, itemID, acc, topLeft)
val acc = getCollisionsAll (iX, iY, iW, iH, itemID, acc, topRight)
val acc = getCollisionsAll (iX, iY, iW, iH, itemID, acc, bottomLeft)
in
getCollisionsAll (iX, iY, iW, iH, itemID, acc, bottomRight)
end
else
acc
| LEAF {items, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
getCollisionsVec (iX, iY, iW, iH, itemID, 0, items, acc)
else
acc
fun helpGetCollisions (iX, iY, iW, iH, itemID, acc, tree: t) =
case tree of
NODE {topLeft, topRight, bottomLeft, bottomRight, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
let
val acc = helpGetCollisions (iX, iY, iW, iH, itemID, acc, topLeft)
val acc = helpGetCollisions (iX, iY, iW, iH, itemID, acc, topRight)
val acc = helpGetCollisions
(iX, iY, iW, iH, itemID, acc, bottomLeft)
in
helpGetCollisions (iX, iY, iW, iH, itemID, acc, bottomRight)
end
else
acc
| LEAF {items, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
getCollisionsVec (iX, iY, iW, iH, itemID, 0, items, acc)
else
acc
fun getCollisions (itemX, itemY, itemWidth, itemHeight, itemID, tree) =
helpGetCollisions (itemX, itemY, itemWidth, itemHeight, itemID, [], tree)
fun hasCollisionAtVec (iX, iY, iW, iH, itemID, pos, elements) =
if pos = Vector.length elements then
false
else
let
val item = Vector.sub (elements, pos)
in
isCollidingItem (iX, iY, iW, iH, itemID, item)
orelse hasCollisionAtVec (iX, iY, iW, iH, itemID, pos + 1, elements)
end
fun hasCollisionAt (iX, iY, iW, iH, itemID, tree) =
case tree of
NODE {topLeft, topRight, bottomLeft, bottomRight, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
hasCollisionAt (iX, iY, iW, iH, itemID, topLeft)
orelse hasCollisionAt (iX, iY, iW, iH, itemID, topRight)
orelse hasCollisionAt (iX, iY, iW, iH, itemID, bottomLeft)
orelse hasCollisionAt (iX, iY, iW, iH, itemID, bottomRight)
else
false
| LEAF {items, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
hasCollisionAtVec (iX, iY, iW, iH, itemID, 0, items)
else
false
fun getItemIDVec (iX, iY, iW, iH, pos, elements) =
if pos = Vector.length elements then
~1
else
let
val item = Vector.sub (elements, pos)
in
if isCollidingItem (iX, iY, iW, iH, ~1, item) then #itemID item
else getItemIDVec (iX, iY, iW, iH, pos + 1, elements)
end
fun getItemID (iX, iY, iW, iH, tree) =
case tree of
NODE {topLeft, topRight, bottomLeft, bottomRight, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
let
val try1 = getItemID (iX, iY, iW, iH, topLeft)
val try2 = getItemID (iX, iY, iW, iH, topRight)
val try3 = getItemID (iX, iY, iW, iH, bottomLeft)
val try4 = getItemID (iX, iY, iW, iH, bottomRight)
(* get max: we assume query was narrow enough
* that only one ID is valid *)
val a = Int.max (try1, try2)
val a = Int.max (a, try3)
val a = Int.max (a, try4)
in
a
end
else
~1
| LEAF {items, x, y, w, h} =>
if isCollidingPlus (iX, iY, iW, iH, x, y, w, h) then
getItemIDVec (iX, iY, iW, iH, 0, items)
else
~1
end