{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
module Data.RTree.Base
(
RTree (..)
, empty
, singleton
, insert
, insertWith
, delete
, mapMaybe
, union
, unionWith
, lookup
, intersectWithKey
, intersect
, lookupRange
, lookupRangeWithKey
, lookupContainsRangeWithKey
, lookupContainsRange
, length
, null
, keys
, values
, fromList
, toList
, foldWithMBB
, pp
, isValid
, unionDistinct
, unionDistinctWith
, fromList'
, unionDistinctSplit
, depth
, areaIncreasesWith
, partition
, getChildren
, unionMBB'
, createNodeWithChildren
, n
, splitNode
, node
)
where
import Control.DeepSeq (NFData, rnf)
import Data.Binary
import Data.Function (on)
import Data.List (maximumBy, minimumBy, partition)
import qualified Data.List as L (length, map)
import Data.Maybe (catMaybes, isJust)
import qualified Data.Maybe as Maybe (mapMaybe)
import Data.Typeable (Typeable)
import Data.Semigroup
import GHC.Generics (Generic)
import Prelude hiding (length, lookup, map, null)
import Data.RTree.MBB hiding (mbb)
data RTree a =
Node4 {getMBB :: {-# UNPACK #-} ! MBB, getC1 :: ! (RTree a), getC2 :: ! (RTree a), getC3 :: ! (RTree a), getC4 :: ! (RTree a) }
| Node3 {getMBB :: {-# UNPACK #-} ! MBB, getC1 :: ! (RTree a), getC2 :: ! (RTree a), getC3 :: ! (RTree a) }
| Node2 {getMBB :: {-# UNPACK #-} ! MBB, getC1 :: ! (RTree a), getC2 :: ! (RTree a) }
| Node {getMBB :: MBB, getChildren' :: [RTree a] }
| Leaf {getMBB :: {-# UNPACK #-} ! MBB, getElem :: a}
| Empty
deriving (Show, Eq, Typeable, Generic, Functor)
m, n :: Int
m = 2
n = 4
unionMBB' :: RTree a -> RTree a -> MBB
unionMBB' = unionMBB `on` getMBB
empty :: RTree a
empty = Empty
null :: RTree a -> Bool
null Empty = True
null _ = False
singleton :: MBB -> a -> RTree a
singleton = Leaf
node :: MBB -> [RTree a] -> RTree a
node mbb [x,y] = Node2 mbb x y
node mbb [x,y,z] = Node3 mbb x y z
node mbb [x,y,z,w] = Node4 mbb x y z w
node _ [] = error "node: empty"
node mbb xs = Node mbb xs
createNodeWithChildren :: [RTree a] -> RTree a
createNodeWithChildren c = node (unionsMBB $ getMBB <$> c) c
norm :: RTree a -> RTree a
norm (Node4 mbb x y z w) = Node mbb [x,y,z,w]
norm (Node3 mbb x y z) = Node mbb [x,y,z]
norm (Node2 mbb x y) = Node mbb [x,y]
norm x = x
getChildren :: RTree a -> [RTree a]
getChildren Empty = error "getChildren: Empty"
getChildren Leaf{} = error "getChildren: Leaf"
getChildren t = getChildren' $ norm t
fromList :: [(MBB, a)] -> RTree a
fromList l = fromList' $ uncurry singleton <$> l
fromList' :: [RTree a] -> RTree a
fromList' [] = empty
fromList' ts = foldr1 unionDistinct ts
toList :: RTree a -> [(MBB, a)]
toList Empty = []
toList (Leaf mbb x) = [(mbb, x)]
toList t = concatMap toList $ getChildren t
keys :: RTree a -> [MBB]
keys = foldWithMBB (\mbb _ -> [mbb]) (const concat) []
values :: RTree a -> [a]
values = foldWithMBB (const pure) (const concat) []
insertWith :: (a -> a -> a) -> MBB -> a -> RTree a -> RTree a
insertWith f mbb e = unionDistinctWith f (singleton mbb e)
insert :: MBB -> a -> RTree a -> RTree a
insert = insertWith const
simpleMergeEqNode :: (a -> a -> a) -> RTree a -> RTree a -> RTree a
simpleMergeEqNode f l@Leaf{} r = Leaf (getMBB l) (on f getElem l r)
simpleMergeEqNode _ l _ = l
unionDistinctWith :: (a -> a -> a) -> RTree a -> RTree a -> RTree a
unionDistinctWith _ Empty{} t = t
unionDistinctWith _ t Empty{} = t
unionDistinctWith f t1@Leaf{} t2@Leaf{}
| on (==) getMBB t1 t2 = simpleMergeEqNode f t1 t2
| otherwise = createNodeWithChildren [t1, t2]
unionDistinctWith f left right
| depth left > depth right = unionDistinctWith f right left
| depth left == depth right = fromList' $ getChildren left ++ [right]
| L.length (getChildren newNode) > n = createNodeWithChildren $ splitNode newNode
| otherwise = newNode
where
newNode = addLeaf f left right
unionDistinct :: RTree a -> RTree a -> RTree a
unionDistinct = unionDistinctWith const
addLeaf :: (a -> a -> a) -> RTree a -> RTree a -> RTree a
addLeaf f left right
| depth left + 1 == depth right = node (newNode `unionMBB'` right) (newNode : nonEq)
| otherwise = node (left `unionMBB'` right) newChildren
where
newChildren = findNodeWithMinimalAreaIncrease f left (getChildren right)
(eq, nonEq) = partition (on (==) getMBB left) $ getChildren right
newNode = case eq of
[] -> left
[x] -> simpleMergeEqNode f left x
_ -> error "addLeaf: invalid RTree"
findNodeWithMinimalAreaIncrease :: (a -> a -> a) -> RTree a -> [RTree a] -> [RTree a]
findNodeWithMinimalAreaIncrease f leaf children = splitMinimal xsAndIncrease
where
xsAndIncrease = zip children $ areaIncreasesWith leaf <$> children
minimalIncrease = minimum $ snd <$> xsAndIncrease
splitMinimal [] = []
splitMinimal ((t,mbb):xs)
| mbb == minimalIncrease = unionDistinctSplit f leaf t ++ (fst <$> xs)
| otherwise = t : splitMinimal xs
unionDistinctSplit :: (a -> a -> a) -> RTree a -> RTree a -> [RTree a]
unionDistinctSplit f leaf e
| L.length (getChildren newLeaf) > n = splitNode newLeaf
| otherwise = [newLeaf]
where
newLeaf = addLeaf f leaf e
splitNode :: RTree a -> [RTree a]
splitNode Leaf{} = error "splitNode: Leaf"
splitNode e = [createNodeWithChildren x1, createNodeWithChildren x2]
where
(l, r) = findGreatestArea $ getChildren e
(x1, x2) = quadSplit [l] [r] unfinished
unfinished = filter (on (/=) getMBB l) $ filter (on (/=) getMBB r) $ getChildren e
findGreatestArea :: [RTree a] -> (RTree a, RTree a)
findGreatestArea xs = (x', y')
where
xs' = zip xs [(1::Int)..]
listOfTripels = [(fst x, fst y, on unionMBB' fst x y) | x <- xs', y <- xs', ((<) `on` snd) x y]
(x', y', _) = maximumBy (compare `on` (\(_,_,x) -> area x)) listOfTripels
quadSplit :: [RTree a] -> [RTree a] -> [RTree a] -> ([RTree a], [RTree a])
quadSplit left right [] = (left, right)
quadSplit left right unfinished
| L.length left + L.length unfinished <= m = (left ++ unfinished, right)
| L.length right + L.length unfinished <= m = (left, right ++ unfinished)
| isLeft'' = quadSplit (minimumElem : left) right newRest
| otherwise = quadSplit left (minimumElem : right) newRest
where
makeTripel x = (x, isLeft, growth)
where
isLeft = areaIncreasesWithLeft < areaIncreasesWithRight
growth = if isLeft
then areaIncreasesWithLeft
else areaIncreasesWithRight
areaIncreasesWithLeft = areaIncreasesWith x $ createNodeWithChildren left
areaIncreasesWithRight = areaIncreasesWith x $ createNodeWithChildren right
(minimumElem, isLeft'', _) = minimumBy (compare `on` (\(_,_,g) -> g)) $ makeTripel <$> unfinished
newRest = filter (on (/=) getMBB minimumElem) unfinished
areaIncreasesWith :: RTree a -> (RTree a) -> Double
areaIncreasesWith newElem current = newArea - currentArea
where
currentArea = area $ getMBB current
newArea = area $ unionMBB' newElem current
lookup :: MBB -> RTree a -> Maybe a
lookup _ Empty = Nothing
lookup mbb t@Leaf{}
| mbb == getMBB t = Just $ getElem t
| otherwise = Nothing
lookup mbb t = case founds of
[] -> Nothing
x:_ -> Just x
where
matches = filter (\x -> (getMBB x) `containsMBB` mbb) $ getChildren t
founds = catMaybes $ L.map (lookup mbb) matches
intersectWithKey :: MBB -> RTree a -> [(MBB, a)]
intersectWithKey _ Empty = []
intersectWithKey mbb t@Leaf{}
| isJust $ intersectMBB mbb (getMBB t) = [(getMBB t, getElem t)]
| otherwise = []
intersectWithKey mbb t = founds
where matches = filter intersectRTree $ getChildren t
founds = concatMap (intersectWithKey mbb) matches
intersectRTree x = isJust $ mbb `intersectMBB` (getMBB x)
intersect :: MBB -> RTree a -> [a]
intersect mbb t = snd <$> intersectWithKey mbb t
lookupRangeWithKey :: MBB -> RTree a -> [(MBB, a)]
lookupRangeWithKey _ Empty = []
lookupRangeWithKey mbb t@Leaf{}
| mbb `containsMBB` (getMBB t) = [(getMBB t, getElem t)]
| otherwise = []
lookupRangeWithKey mbb t = founds
where
matches = filter intersectRTree $ getChildren t
founds = concatMap (lookupRangeWithKey mbb) matches
intersectRTree x = isJust $ mbb `intersectMBB` (getMBB x)
lookupRange :: MBB -> RTree a -> [a]
lookupRange mbb t = snd <$> (lookupRangeWithKey mbb t)
lookupContainsRangeWithKey :: MBB -> RTree a -> [(MBB, a)]
lookupContainsRangeWithKey _ Empty = []
lookupContainsRangeWithKey mbb t@Leaf{}
| (getMBB t) `containsMBB` mbb = [(getMBB t, getElem t)]
| otherwise = []
lookupContainsRangeWithKey mbb t = founds
where
matches = filter intersectRTree $ getChildren t
founds = concatMap (lookupContainsRangeWithKey mbb) matches
intersectRTree x = (getMBB x) `containsMBB` mbb
lookupContainsRange :: MBB -> RTree a -> [a]
lookupContainsRange mbb t = snd <$> (lookupContainsRangeWithKey mbb t)
delete :: MBB -> RTree a -> RTree a
delete _ Empty = Empty
delete mbb t@Leaf{}
| mbb == getMBB t = Empty
| otherwise = t
delete mbb root
| L.length (getChildren newRoot) == 1 = head $ getChildren newRoot
| otherwise = newRoot
where
newRoot = delete' mbb root
delete' :: MBB -> RTree a -> RTree a
delete' mbb t@Leaf{}
| mbb == getMBB t = Empty
| otherwise = t
delete' mbb t = fromList' $ orphans ++ [newValidNode]
where
(matches, noMatches) = partition (\x -> (getMBB x) `containsMBB` mbb) $ getChildren t
matches' = filter (not . null) $ L.map (delete' mbb) matches
(orphans, validMatches) = foldr handleInvalid ([], []) matches'
handleInvalid l@Leaf{} (orphans', validMatches') = (orphans', l:validMatches')
handleInvalid invalidNode (orphans', validMatches')
| L.length children < m = (children ++ orphans', validMatches')
| otherwise = (orphans', invalidNode:validMatches')
where
children = getChildren invalidNode
newValidNode = createNodeWithChildren $ validMatches ++ noMatches
foldWithMBB :: (MBB -> a -> b) -> (MBB -> [b] -> b) -> b -> RTree a -> b
foldWithMBB _ _ n' Empty = n'
foldWithMBB f _ _ t@Leaf{} = f (getMBB t) (getElem t)
foldWithMBB f g n' t = g (getMBB t) $ foldWithMBB f g n' <$> (getChildren t)
unionWith :: (a -> a -> a) -> RTree a -> RTree a -> RTree a
unionWith _ l Empty = l
unionWith _ Empty r = r
unionWith f t1 t2
| depth t1 <= depth t2 = foldr (uncurry (insertWith f)) t2 (toList t1)
| otherwise = unionWith f t2 t1
union :: RTree a -> RTree a -> RTree a
union = unionWith const
mapMaybe :: (a -> Maybe b) -> RTree a -> RTree b
mapMaybe f t = fromList $ Maybe.mapMaybe func $ toList t
where
func (mbb,x) = case f x of
Nothing -> Nothing
Just x' -> Just (mbb, x')
isValid :: Show b => b -> RTree a -> Bool
isValid _ Empty = True
isValid _ Leaf{} = True
isValid context x = case L.length c >= m && L.length c <= n && (and $ (isValid context) <$> c) && (isBalanced x) of
True -> True
False -> error ( "invalid " ++ show (L.length c) ++ " " ++ show context )
where
isBalanced :: RTree a -> Bool
isBalanced (Leaf _ _ ) = True
isBalanced x' = (and $ isBalanced <$> c') && (and $ (== depth (head c')) <$> (depth <$> c'))
where
c' = getChildren x'
c = getChildren x
i_ :: String
i_ = " "
pp :: (Show a) => RTree a -> IO ()
pp = pp' ""
pp' :: (Show a) => String -> RTree a -> IO ()
pp' i Empty = putStrLn $ i ++ "Empty"
pp' i (Leaf mbb x) = putStrLn $ i ++ "Leaf " ++ (show mbb) ++ " " ++ (show x)
pp' i (Node mbb cs) = do
putStrLn $ i ++ "Node " ++ (show mbb)
mapM_ (pp' (i ++ i_)) cs
pp' i (Node2 mbb c1 c2) = do
putStrLn $ i ++ "Node2 " ++ (show mbb)
mapM_ (pp' (i ++ i_)) [c1, c2]
pp' i (Node3 mbb c1 c2 c3) = do
putStrLn $ i ++ "Node3 " ++ (show mbb)
mapM_ (pp' (i ++ i_)) [c1, c2, c3]
pp' i (Node4 mbb c1 c2 c3 c4) = do
putStrLn $ i ++ "Node4 " ++ (show mbb)
mapM_ (pp' (i ++ i_)) [c1, c2, c3, c4]
depth :: RTree a -> Int
depth Empty = 0
depth (Leaf _ _ ) = 1
depth t = 1 + (depth $ head $ getChildren t)
length :: RTree a -> Int
length Empty = 0
length (Leaf {}) = 1
length t = sum $ length <$> (getChildren t)
instance NFData a => NFData (RTree a) where
rnf (Empty) = ()
rnf (Leaf _ e) = rnf e
rnf (Node _ cs) = rnf cs
rnf (Node2 _ c1 c2) = rnf c1 `seq` rnf c2
rnf (Node3 _ c1 c2 c3) = rnf c1 `seq` rnf c2 `seq` rnf c3
rnf (Node4 _ c1 c2 c3 c4) = rnf c1 `seq` rnf c2 `seq` rnf c3 `seq` rnf c4
instance (Binary a) => Binary (RTree a) where
put (Empty) = put (0::Word8)
put (Leaf mbb e) = put (1::Word8) >> put mbb >> put e
put t = put (2::Word8) >> put (getMBB t) >> put (getChildren t)
get = do
!tag <- getWord8
case tag of
0 -> return Empty
1 -> do
!mbb <- get
!e <- get
return $! Leaf mbb e
2 -> do
!mbb <- get
!c <- get
return $! node mbb c
_ -> fail "RTree.get: error while decoding RTree"
instance (Semigroup a) => Semigroup (RTree a) where
(<>) = unionWith (<>)
instance Monoid a => Monoid (RTree a) where
mempty = empty
mappend = (<>)