{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Data.RTree.Strict
(
MBB
, MBB.mbb
, RTree ()
, toLazy
, toStrict
, empty
, singleton
, insert
, insertWith
, delete
, mapMaybe
, union
, unionWith
, lookup
, intersectWithKey
, intersect
, lookupRange
, lookupRangeWithKey
, lookupContainsRange
, lookupContainsRangeWithKey
, length
, null
, keys
, values
, fromList
, toList
) where
import Prelude hiding (lookup, length, null, map)
import Data.Binary
import Data.Function (on)
import qualified Data.List as L (length)
import qualified Data.Maybe as Maybe (mapMaybe)
import Data.Semigroup
import Data.Typeable (Typeable)
import Control.DeepSeq (NFData)
import GHC.Generics (Generic)
import qualified Data.RTree.Base as Lazy
import Data.RTree.MBB hiding (mbb)
import qualified Data.RTree.MBB as MBB
newtype RTree a = RTree {toLazy' :: Lazy.RTree a}
deriving (Show, Eq, Typeable, Generic, NFData, Binary, Monoid, Semigroup)
toStrict :: Lazy.RTree a -> RTree a
toStrict t = map id (RTree t)
toLazy :: RTree a -> Lazy.RTree a
toLazy = toLazy'
empty :: RTree a
empty = RTree Lazy.Empty
null :: RTree a -> Bool
null = Lazy.null . toLazy
singleton :: MBB -> a -> RTree a
singleton mbb !x = RTree $ Lazy.Leaf mbb x
fromList :: [(MBB, a)] -> RTree a
fromList l = RTree $ fromList' $ (toLazy . (uncurry singleton)) <$> l
fromList' :: [Lazy.RTree a] -> Lazy.RTree a
fromList' [] = Lazy.empty
fromList' ts = foldr1 unionDistinct ts
toList :: RTree a -> [(MBB, a)]
toList = Lazy.toList . toLazy
keys :: RTree a -> [MBB]
keys = Lazy.keys . toLazy
values :: RTree a -> [a]
values = Lazy.values . toLazy
insertWith :: (a -> a -> a) -> MBB -> a -> RTree a -> RTree a
insertWith f mbb e oldRoot = RTree $ insertWithStrictLazy f mbb e (toLazy oldRoot)
insertWithStrictLazy :: (a -> a -> a) -> MBB -> a -> Lazy.RTree a -> Lazy.RTree a
insertWithStrictLazy f mbb e oldRoot = unionDistinctWith f (toLazy $ singleton mbb e) oldRoot
insert :: MBB -> a -> RTree a -> RTree a
insert = insertWith const
simpleMergeEqNode :: (a -> a -> a) -> Lazy.RTree a -> Lazy.RTree a -> Lazy.RTree a
simpleMergeEqNode f l@Lazy.Leaf{} r = Lazy.Leaf (Lazy.getMBB l) $! (on f Lazy.getElem l r)
simpleMergeEqNode _ l _ = l
unionDistinctWith :: (a -> a -> a) -> Lazy.RTree a -> Lazy.RTree a -> Lazy.RTree a
unionDistinctWith _ Lazy.Empty{} t = t
unionDistinctWith _ t Lazy.Empty{} = t
unionDistinctWith f t1@Lazy.Leaf{} t2@Lazy.Leaf{}
| on (==) Lazy.getMBB t1 t2 = simpleMergeEqNode f t1 t2
| otherwise = Lazy.createNodeWithChildren [t1, t2]
unionDistinctWith f left right
| Lazy.depth left > Lazy.depth right = unionDistinctWith f right left
| Lazy.depth left == Lazy.depth right = fromList' $ (Lazy.getChildren left) ++ [right]
| (L.length $ Lazy.getChildren newNode) > Lazy.n = Lazy.createNodeWithChildren $ Lazy.splitNode newNode
| otherwise = newNode
where
newNode = addLeaf f left right
unionDistinct :: Lazy.RTree a -> Lazy.RTree a -> Lazy.RTree a
unionDistinct = unionDistinctWith const
addLeaf :: (a -> a -> a) -> Lazy.RTree a -> Lazy.RTree a -> Lazy.RTree a
addLeaf f left right
| Lazy.depth left + 1 == Lazy.depth right = Lazy.node (newNode `Lazy.unionMBB'` right) (newNode : nonEq)
| otherwise = Lazy.node (left `Lazy.unionMBB'` right) newChildren
where
newChildren = findNodeWithMinimalAreaIncrease f left (Lazy.getChildren right)
(eq, nonEq) = Lazy.partition (on (==) Lazy.getMBB left) $ Lazy.getChildren right
newNode = case eq of
[] -> left
[x] -> simpleMergeEqNode f left x
_ -> error "addLeaf: invalid RTree"
findNodeWithMinimalAreaIncrease :: (a -> a -> a) -> Lazy.RTree a -> [Lazy.RTree a] -> [Lazy.RTree a]
findNodeWithMinimalAreaIncrease f leaf children = splitMinimal xsAndIncrease
where
xsAndIncrease = zip children ((Lazy.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) -> Lazy.RTree a -> Lazy.RTree a -> [Lazy.RTree a]
unionDistinctSplit f leaf e
| (L.length $ Lazy.getChildren newLeaf) > Lazy.n = Lazy.splitNode newLeaf
| otherwise = [newLeaf]
where
newLeaf = addLeaf f leaf e
lookup :: MBB -> RTree a -> Maybe a
lookup mbb = Lazy.lookup mbb . toLazy
intersectWithKey :: MBB -> RTree a -> [(MBB, a)]
intersectWithKey mbb = Lazy.intersectWithKey mbb . toLazy
intersect :: MBB -> RTree a -> [a]
intersect mbb = Lazy.intersect mbb . toLazy
lookupRangeWithKey :: MBB -> RTree a -> [(MBB, a)]
lookupRangeWithKey mbb = Lazy.lookupRangeWithKey mbb . toLazy
lookupRange :: MBB -> RTree a -> [a]
lookupRange mbb = Lazy.lookupRange mbb . toLazy
lookupContainsRangeWithKey :: MBB -> RTree a -> [(MBB, a)]
lookupContainsRangeWithKey mbb = Lazy.lookupContainsRangeWithKey mbb . toLazy
lookupContainsRange :: MBB -> RTree a -> [a]
lookupContainsRange mbb = Lazy.lookupContainsRange mbb .toLazy
delete :: MBB -> RTree a -> RTree a
delete mbb = RTree . Lazy.delete mbb . toLazy
unionWith :: (a -> a -> a) -> RTree a -> RTree a -> RTree a
unionWith f' l' r' = RTree $ unionWith' f' (toLazy l') (toLazy r')
where
unionWith' _ l Lazy.Empty = l
unionWith' _ Lazy.Empty r = r
unionWith' f t1 t2
| Lazy.depth t1 <= Lazy.depth t2 = foldr (uncurry (insertWithStrictLazy f)) t2 (Lazy.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')
map :: (a -> b) -> RTree a -> RTree b
map f' = RTree . map' f' . toLazy
where
map' f (Lazy.Node4 mbb x y z w) = Lazy.Node4 mbb (map' f x) (map' f y) (map' f z) (map' f w)
map' f (Lazy.Node3 mbb x y z) = Lazy.Node3 mbb (map' f x) (map' f y) (map' f z)
map' f (Lazy.Node2 mbb x y) = Lazy.Node2 mbb (map' f x) (map' f y)
map' f (Lazy.Node mbb xs) = Lazy.Node mbb (map' f <$> xs)
map' f (Lazy.Leaf mbb e) = toLazy $ singleton mbb (f e)
map' _ Lazy.Empty = Lazy.Empty
length :: RTree a -> Int
length = Lazy.length . toLazy
instance Functor RTree where
fmap = map