KDTree2d.hs

{-# LANGUAGE BangPatterns #-}

-- KDTree code
--   by Matthew Sottile <matt@galois.com> <mjsottile@computer.org>
--
module KDTree2d (
  KDTreeNode (..),
  kdtAddPoints,
  kdtAddPoint,
  kdtRangeSearch,
  kdtCollisionDetect,
  kdtInBounds,
  dumpKDTree,
  mapKDTree,
  kdtreeToList,
) where

import Data.Maybe
import System.IO
import Vec2


data KDTreeNode a
  = Empty
  | Node !(KDTreeNode a) !Vec2 !a !(KDTreeNode a)
  deriving (Show)


-- | Flatten out a KDTree to a list.
kdtreeToList :: KDTreeNode a -> [a]
kdtreeToList Empty = []
kdtreeToList (Node l _ x r) = [x] ++ kdtreeToList l ++ kdtreeToList r


-- | Apply a worker function to all elements of a KDTree.
mapKDTree :: KDTreeNode a -> (a -> b) -> [b]
mapKDTree Empty _ = []
mapKDTree (Node l _p n r) f = f n : (mapKDTree l f ++ mapKDTree r f)


kdtAddWithDepth :: KDTreeNode a -> Vec2 -> a -> Int -> KDTreeNode a
kdtAddWithDepth Empty pos dat _ =
  Node Empty pos dat Empty
kdtAddWithDepth (Node left npos ndata right) pos dat d
  | vecDimSelect pos d < vecDimSelect npos d =
      Node (kdtAddWithDepth left pos dat d') npos ndata right
  | otherwise =
      Node left npos ndata (kdtAddWithDepth right pos dat d')
  where
    d' = if (d == 1) then 0 else 1


kdtAddPoint :: KDTreeNode a -> Vec2 -> a -> KDTreeNode a
kdtAddPoint t p d =
  kdtAddWithDepth t p d 0


kdtInBounds :: Vec2 -> Vec2 -> Vec2 -> Bool
kdtInBounds p bMin bMax =
  vecLessThan p bMax && vecGreaterThan p bMin


-- X dimension
kdtRangeSearchRecX :: KDTreeNode a -> Vec2 -> Vec2 -> [(Vec2, a)]
kdtRangeSearchRecX Empty _ _ = []
kdtRangeSearchRecX (Node left npos ndata right) bMin bMax
  | nc < mnc =
      nextfun right bMin bMax
  | nc > mxc =
      nextfun left bMin bMax
  | kdtInBounds npos bMin bMax =
      (npos, ndata)
        : (nextfun right bMin bMax ++ nextfun left bMin bMax)
  | otherwise =
      nextfun right bMin bMax ++ nextfun left bMin bMax
  where
    Vec2 nc _ = npos
    Vec2 mnc _ = bMin
    Vec2 mxc _ = bMax
    nextfun = kdtRangeSearchRecY


-- Y dimension
kdtRangeSearchRecY :: (KDTreeNode a) -> Vec2 -> Vec2 -> [(Vec2, a)]
kdtRangeSearchRecY Empty _ _ = []
kdtRangeSearchRecY (Node left npos ndata right) bMin bMax
  | nc < mnc =
      nextfun right bMin bMax
  | nc > mxc =
      nextfun left bMin bMax
  | (kdtInBounds npos bMin bMax) =
      (npos, ndata)
        : (nextfun right bMin bMax ++ nextfun left bMin bMax)
  | otherwise =
      nextfun right bMin bMax ++ nextfun left bMin bMax
  where
    Vec2 _ nc = npos
    Vec2 _ mnc = bMin
    Vec2 _ mxc = bMax
    nextfun = kdtRangeSearchRecX


kdtRangeSearch :: (KDTreeNode a) -> Vec2 -> Vec2 -> [(Vec2, a)]
kdtRangeSearch t bMin bMax =
  kdtRangeSearchRecX t bMin bMax


kdtAddPoints :: [(Vec2, a)] -> (KDTreeNode a) -> (KDTreeNode a)
kdtAddPoints [] t = t
kdtAddPoints ((pt, dat) : ps) t =
  kdtAddPoints ps $ kdtAddPoint t pt dat


singleCollision :: Vec2 -> Vec2 -> Vec2 -> Double -> a -> Maybe (Vec2, a)
singleCollision pt start a eps dat
  | sqrd_dist < eps * eps =
      Just (vecAdd start p, dat)
  | otherwise =
      Nothing
  where
    b = vecSub pt start
    xhat = (vecDot a b) / (vecDot a a)
    p = vecScale a xhat
    e = vecSub p b
    sqrd_dist = vecDot e e


kdtCollisionDetect :: KDTreeNode a -> Vec2 -> Vec2 -> Double -> [(Vec2, a)]
kdtCollisionDetect root !start !end !eps =
  colls
  where
    Vec2 sx sy = start
    Vec2 ex ey = end
    rmin = Vec2 (min sx ex - eps) (min sy ey - eps)
    rmax = Vec2 (max sx ex + eps) (max sy ey + eps)
    pts = kdtRangeSearch root rmin rmax
    a = vecSub end start
    colls = mapMaybe (\(pt, dat) -> singleCollision pt start a eps dat) pts


-- Dumping --------------------------------------------------------------------

-- | Dump a KDTree to a file
dumpKDTree :: KDTreeNode Int -> FilePath -> IO ()
dumpKDTree kdt name =
  do
    h <- openFile name WriteMode
    hPutStrLn h "n x y z"
    dumpKDTreeInner kdt h
    hClose h


-- | Dump a KDTree to a handle.
dumpKDTreeInner :: KDTreeNode Int -> Handle -> IO ()
dumpKDTreeInner kdt h =
  case kdt of
    Empty -> return ()
    Node l v d r ->
      do
        printVec v h d
        dumpKDTreeInner l h
        dumpKDTreeInner r h


-- | Print a vector to a handle.
printVec :: Vec2 -> Handle -> Int -> IO ()
printVec (Vec2 x y) h i =
  hPutStrLn h $ show i ++ " " ++ show x ++ " " ++ show y