{-# LANGUAGE FlexibleContexts, TypeApplications #-}

module Data.IntervalIntMap.Internal.IntervalIntIntMap
    ( IntervalValue(..)
    , Interval(..)
    , IntervalIntMap
    , naiveIntervalMapLookup
    , lookup
    , overlaps
    , overlapsWithKeys
    , naiveOverlaps
    , naiveOverlapsWithKeys
    , NaiveIntervalInt
    , intervalContains
    , partition
    , freeze
#ifdef IS_BUILDING_TEST
    , mkTree
#endif
    ) where
import Prelude hiding (lookup)

import qualified Data.IntervalIntMap.Internal.GrowableVector as GV

import qualified Foreign.Storable as FS
import           Foreign.Ptr (castPtr, plusPtr)
import qualified Data.Set as S
import qualified Data.IntSet as IS
import qualified Data.Vector.Storable as VS
import           Control.Monad.ST (runST)
import           Data.Word (Word32)
import           Data.Ord (comparing)
import           Data.Vector.Algorithms.Tim (sortBy)
import           Control.DeepSeq (NFData(..))


{- DATA STRUCTURE
 -
 - An IntervalValue contains the interval [ivStart, ivPast) and the value. This
 - is a closed-open interval, so represents `x` such that `ivStart <= x <
 - ivPast`.
 -
 - The simplest map is the NaiveIntervalInt, which is just a vector. It is very
 - memory efficient, but needs O(N) to search. However, for small N, it is
 - likely very efficient and we can use this "structure" for testing too.
 -
 - The Tree is very simple: At each node, there is a split value and intervals
 - are completely below it, completely above it, or contain the point.
 -
 -
 - Leafs contain NaiveIntervalInt
 -}


data Interval = Interval !Int !Int
#ifdef IS_BUILDING_TEST
                            deriving (Eq, Show)
#endif

data IntervalValue = IntervalValue
                        { IntervalValue -> Word32
ivStart :: !Word32
                        , IntervalValue -> Word32
ivPast :: !Word32
                        , IntervalValue -> Word32
ivValue :: !Word32
                        }
#ifdef IS_BUILDING_TEST
                              deriving (Show)
#endif

instance Eq IntervalValue where
    (IntervalValue Word32
s0 Word32
e0 Word32
ix0) == :: IntervalValue -> IntervalValue -> Bool
== (IntervalValue Word32
s1 Word32
e1 Word32
ix1) =
            Word32
s0 Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
s1 Bool -> Bool -> Bool
&& Word32
e0 Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
e1 Bool -> Bool -> Bool
&& Word32
ix0 Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
ix1

-- This is necessary to build sets of 'IntervalValue's (e.g., in 'naiveOverlapsWithKeys')
instance Ord IntervalValue where
    (IntervalValue Word32
s0 Word32
e0 Word32
ix0) compare :: IntervalValue -> IntervalValue -> Ordering
`compare` (IntervalValue Word32
s1 Word32
e1 Word32
ix1)
        | Word32
s0 Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word32
s1 = Word32
s0 Word32 -> Word32 -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Word32
s1
        | Word32
e0 Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word32
e1 = Word32
e0 Word32 -> Word32 -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Word32
e1
        | Bool
otherwise = Word32
ix0 Word32 -> Word32 -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Word32
ix1

instance FS.Storable IntervalValue where
    sizeOf :: IntervalValue -> Int
sizeOf IntervalValue
_ = Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 -- aka 12
    alignment :: IntervalValue -> Int
alignment IntervalValue
x = Word32 -> Int
forall a. Storable a => a -> Int
FS.alignment (IntervalValue -> Word32
ivStart IntervalValue
x)
    peek :: Ptr IntervalValue -> IO IntervalValue
peek Ptr IntervalValue
p = Word32 -> Word32 -> Word32 -> IntervalValue
IntervalValue
                    (Word32 -> Word32 -> Word32 -> IntervalValue)
-> IO Word32 -> IO (Word32 -> Word32 -> IntervalValue)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
FS.peek (Ptr IntervalValue -> Ptr Word32
forall a b. Ptr a -> Ptr b
castPtr Ptr IntervalValue
p)
                    IO (Word32 -> Word32 -> IntervalValue)
-> IO Word32 -> IO (Word32 -> IntervalValue)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
FS.peek (Ptr IntervalValue -> Ptr Any
forall a b. Ptr a -> Ptr b
castPtr Ptr IntervalValue
p Ptr Any -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
4)
                    IO (Word32 -> IntervalValue) -> IO Word32 -> IO IntervalValue
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
FS.peek (Ptr IntervalValue -> Ptr Any
forall a b. Ptr a -> Ptr b
castPtr Ptr IntervalValue
p Ptr Any -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
8)
    poke :: Ptr IntervalValue -> IntervalValue -> IO ()
poke Ptr IntervalValue
ptr (IntervalValue Word32
s Word32
p Word32
v) = do
        let ptr' :: Ptr b
ptr' = Ptr IntervalValue -> Ptr b
forall a b. Ptr a -> Ptr b
castPtr Ptr IntervalValue
ptr
        Ptr Word32 -> Int -> Word32 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
FS.pokeElemOff @Word32 Ptr Word32
forall b. Ptr b
ptr' Int
0 Word32
s
        Ptr Word32 -> Int -> Word32 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
FS.pokeElemOff @Word32 Ptr Word32
forall b. Ptr b
ptr' Int
1 Word32
p
        Ptr Word32 -> Int -> Word32 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
FS.pokeElemOff @Word32 Ptr Word32
forall b. Ptr b
ptr' Int
2 Word32
v

intervalContains :: Int -> IntervalValue -> Bool
intervalContains :: Int -> IntervalValue -> Bool
intervalContains Int
p (IntervalValue Word32
s Word32
e Word32
_) =
    let p' :: Word32
p' = Int -> Word32
forall a. Enum a => Int -> a
toEnum Int
p
    in Word32
s Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word32
p' Bool -> Bool -> Bool
&& Word32
p' Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
e

type NaiveIntervalInt = VS.Vector IntervalValue

data IntervalIntMapNode = Leaf NaiveIntervalInt
                        | InnerNode
                            { IntervalIntMapNode -> Int
_nodeSplitValue :: !Int
                            , IntervalIntMapNode -> IntervalIntMapNode
_leftSplit :: !IntervalIntMapNode
                            , IntervalIntMapNode -> IntervalIntMapNode
_centerSplit :: !IntervalIntMapNode
                            , IntervalIntMapNode -> IntervalIntMapNode
_rightSplit :: !IntervalIntMapNode
                            }
#ifdef IS_BUILDING_TEST
                              deriving (Show)
#endif

instance NFData IntervalIntMapNode where
    rnf :: IntervalIntMapNode -> ()
rnf (Leaf NaiveIntervalInt
v) = NaiveIntervalInt -> ()
forall a. NFData a => a -> ()
rnf NaiveIntervalInt
v
    rnf (InnerNode !Int
_ IntervalIntMapNode
left IntervalIntMapNode
center IntervalIntMapNode
right) = IntervalIntMapNode -> ()
forall a. NFData a => a -> ()
rnf IntervalIntMapNode
left () -> () -> ()
`seq` IntervalIntMapNode -> ()
forall a. NFData a => a -> ()
rnf IntervalIntMapNode
center () -> () -> ()
`seq` IntervalIntMapNode -> ()
forall a. NFData a => a -> ()
rnf IntervalIntMapNode
right

newtype IntervalIntMap = IntervalIntMap { IntervalIntMap -> IntervalIntMapNode
_imapRoot :: IntervalIntMapNode }
#ifdef IS_BUILDING_TEST
                              deriving (Show)
#endif

instance NFData IntervalIntMap where
    rnf :: IntervalIntMap -> ()
rnf (IntervalIntMap !IntervalIntMapNode
n) = IntervalIntMapNode -> ()
forall a. NFData a => a -> ()
rnf IntervalIntMapNode
n

partition :: Int -> NaiveIntervalInt -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt)
partition :: Int
-> NaiveIntervalInt
-> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt)
partition Int
p NaiveIntervalInt
vec = (forall s.
 ST s (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
-> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt)
forall a. (forall s. ST s a) -> a
runST ((forall s.
  ST s (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
 -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
-> (forall s.
    ST s (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
-> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt)
forall a b. (a -> b) -> a -> b
$ do
    GrowableVector s IntervalValue
left <- ST s (GrowableVector s IntervalValue)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
m (GrowableVector (PrimState m) a)
GV.new
    GrowableVector s IntervalValue
center <- ST s (GrowableVector s IntervalValue)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
m (GrowableVector (PrimState m) a)
GV.new
    GrowableVector s IntervalValue
right <- ST s (GrowableVector s IntervalValue)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
m (GrowableVector (PrimState m) a)
GV.new
    NaiveIntervalInt -> (IntervalValue -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Storable a) =>
Vector a -> (a -> m b) -> m ()
VS.forM_ NaiveIntervalInt
vec ((IntervalValue -> ST s ()) -> ST s ())
-> (IntervalValue -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \IntervalValue
val ->
        let target :: GrowableVector s IntervalValue
target
                | IntervalValue -> Word32
ivPast IntervalValue
val Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Word32
forall a. Enum a => Int -> a
toEnum Int
p = GrowableVector s IntervalValue
left
                | IntervalValue -> Word32
ivStart IntervalValue
val Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Word32
forall a. Enum a => Int -> a
toEnum Int
p = GrowableVector s IntervalValue
right
                | Bool
otherwise = GrowableVector s IntervalValue
center
        in IntervalValue
-> GrowableVector (PrimState (ST s)) IntervalValue -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
a -> GrowableVector (PrimState m) a -> m ()
GV.pushBack IntervalValue
val GrowableVector s IntervalValue
GrowableVector (PrimState (ST s)) IntervalValue
target
    (,,)
        (NaiveIntervalInt
 -> NaiveIntervalInt
 -> NaiveIntervalInt
 -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
-> ST s NaiveIntervalInt
-> ST
     s
     (NaiveIntervalInt
      -> NaiveIntervalInt
      -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GrowableVector (PrimState (ST s)) IntervalValue
-> ST s NaiveIntervalInt
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
GrowableVector (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowableVector s IntervalValue
GrowableVector (PrimState (ST s)) IntervalValue
left
        ST
  s
  (NaiveIntervalInt
   -> NaiveIntervalInt
   -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
-> ST s NaiveIntervalInt
-> ST
     s
     (NaiveIntervalInt
      -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> GrowableVector (PrimState (ST s)) IntervalValue
-> ST s NaiveIntervalInt
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
GrowableVector (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowableVector s IntervalValue
GrowableVector (PrimState (ST s)) IntervalValue
center
        ST
  s
  (NaiveIntervalInt
   -> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt))
-> ST s NaiveIntervalInt
-> ST s (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> GrowableVector (PrimState (ST s)) IntervalValue
-> ST s NaiveIntervalInt
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
GrowableVector (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowableVector s IntervalValue
GrowableVector (PrimState (ST s)) IntervalValue
right


sortedByEnd :: NaiveIntervalInt -> NaiveIntervalInt
sortedByEnd :: NaiveIntervalInt -> NaiveIntervalInt
sortedByEnd NaiveIntervalInt
vec = (forall s. ST s (MVector s IntervalValue)) -> NaiveIntervalInt
forall a. Storable a => (forall s. ST s (MVector s a)) -> Vector a
VS.create ((forall s. ST s (MVector s IntervalValue)) -> NaiveIntervalInt)
-> (forall s. ST s (MVector s IntervalValue)) -> NaiveIntervalInt
forall a b. (a -> b) -> a -> b
$ do
    MVector s IntervalValue
vec' <- NaiveIntervalInt -> ST s (MVector (PrimState (ST s)) IntervalValue)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VS.thaw NaiveIntervalInt
vec
    (IntervalValue -> IntervalValue -> Ordering)
-> MVector (PrimState (ST s)) IntervalValue -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
sortBy ((IntervalValue -> Word32)
-> IntervalValue -> IntervalValue -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing IntervalValue -> Word32
ivPast) MVector s IntervalValue
MVector (PrimState (ST s)) IntervalValue
vec'
    MVector s IntervalValue -> ST s (MVector s IntervalValue)
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s IntervalValue
vec'

{-|
  Turn a 'NaiveIntervalInt' into an 'IntervalIntMap'
-}
freeze :: NaiveIntervalInt -> IntervalIntMap
freeze :: NaiveIntervalInt -> IntervalIntMap
freeze = Int -> NaiveIntervalInt -> IntervalIntMap
mkTree Int
16

mkTree :: Int -> NaiveIntervalInt -> IntervalIntMap
mkTree :: Int -> NaiveIntervalInt -> IntervalIntMap
mkTree Int
maxSplit NaiveIntervalInt
vec = IntervalIntMapNode -> IntervalIntMap
IntervalIntMap (IntervalIntMapNode -> IntervalIntMap)
-> IntervalIntMapNode -> IntervalIntMap
forall a b. (a -> b) -> a -> b
$ Int -> Int -> NaiveIntervalInt -> IntervalIntMapNode
mkTree' Int
0 Int
maxSplit (NaiveIntervalInt -> NaiveIntervalInt
sortedByEnd NaiveIntervalInt
vec)

maxSplitIters :: Int
maxSplitIters :: Int
maxSplitIters = Int
8

mkTree' :: Int -> Int -> NaiveIntervalInt -> IntervalIntMapNode
mkTree' Int
nIters Int
maxSplit NaiveIntervalInt
vec
    | NaiveIntervalInt -> Int
forall a. Storable a => Vector a -> Int
VS.length NaiveIntervalInt
vec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxSplit = NaiveIntervalInt -> IntervalIntMapNode
Leaf NaiveIntervalInt
vec
    | Int
nIters Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxSplitIters = NaiveIntervalInt -> IntervalIntMapNode
Leaf NaiveIntervalInt
vec
    | Bool
otherwise = Int -> Int -> NaiveIntervalInt -> IntervalIntMapNode
trySplit Int
nIters Int
maxSplit NaiveIntervalInt
vec

trySplit :: Int -> Int -> NaiveIntervalInt -> IntervalIntMapNode
trySplit Int
nIters Int
maxSplit NaiveIntervalInt
vec = Int
-> IntervalIntMapNode
-> IntervalIntMapNode
-> IntervalIntMapNode
-> IntervalIntMapNode
InnerNode (Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
p) (NaiveIntervalInt -> IntervalIntMapNode
r NaiveIntervalInt
left) (NaiveIntervalInt -> IntervalIntMapNode
r NaiveIntervalInt
center) (NaiveIntervalInt -> IntervalIntMapNode
r NaiveIntervalInt
right)
    where
        r :: NaiveIntervalInt -> IntervalIntMapNode
r = Int -> Int -> NaiveIntervalInt -> IntervalIntMapNode
mkTree' Int
nIters' Int
maxSplit
        (NaiveIntervalInt
left, NaiveIntervalInt
center, NaiveIntervalInt
right) = Int
-> NaiveIntervalInt
-> (NaiveIntervalInt, NaiveIntervalInt, NaiveIntervalInt)
partition (Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
p) NaiveIntervalInt
vec
        nIters' :: Int
nIters'
            | Bool
successful = Int
0
            | Bool
otherwise = Int
nIters Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

        -- The criterion for calling it a successful split is a bit random, but seems to work:
        -- If after splitting the largest component is at least maxSplit
        -- smaller than the input, that was a successful split
        successful :: Bool
successful = NaiveIntervalInt -> Int
forall a. Storable a => Vector a -> Int
VS.length NaiveIntervalInt
vec Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((NaiveIntervalInt -> Int) -> [NaiveIntervalInt] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map NaiveIntervalInt -> Int
forall a. Storable a => Vector a -> Int
VS.length [NaiveIntervalInt
left, NaiveIntervalInt
center, NaiveIntervalInt
right]) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxSplit

        -- Choosing a pivot will probably have a big impact on the performance.
        -- We pick the median end-point one, which is probably a decent heuristic
        p :: Word32
p = IntervalValue -> Word32
ivPast (IntervalValue -> Word32) -> IntervalValue -> Word32
forall a b. (a -> b) -> a -> b
$ NaiveIntervalInt -> Int -> IntervalValue
forall a. Storable a => Vector a -> Int -> a
(VS.!) NaiveIntervalInt
vec (NaiveIntervalInt -> Int
forall a. Storable a => Vector a -> Int
VS.length NaiveIntervalInt
vec Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

lookup :: Int -> IntervalIntMap -> IS.IntSet
lookup :: Int -> IntervalIntMap -> IntSet
lookup Int
x (IntervalIntMap IntervalIntMapNode
root) = IntervalIntMapNode -> IntSet
lookup' IntervalIntMapNode
root
    where

        lookup' :: IntervalIntMapNode -> IntSet
lookup' (Leaf NaiveIntervalInt
vec) = Int -> NaiveIntervalInt -> IntSet
naiveIntervalMapLookup Int
x NaiveIntervalInt
vec
        lookup' (InnerNode Int
p IntervalIntMapNode
left IntervalIntMapNode
center IntervalIntMapNode
right)
            | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
p = IntervalIntMapNode -> IntSet
lookup' IntervalIntMapNode
left IntSet -> IntSet -> IntSet
`IS.union` IntervalIntMapNode -> IntSet
lookup' IntervalIntMapNode
center
            | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
p = IntervalIntMapNode -> IntSet
lookup' IntervalIntMapNode
center
            | Bool
otherwise = IntervalIntMapNode -> IntSet
lookup' IntervalIntMapNode
center IntSet -> IntSet -> IntSet
`IS.union` IntervalIntMapNode -> IntSet
lookup' IntervalIntMapNode
right

naiveIntervalMapLookup :: Int -> NaiveIntervalInt -> IS.IntSet
naiveIntervalMapLookup :: Int -> NaiveIntervalInt -> IntSet
naiveIntervalMapLookup Int
x = [Int] -> IntSet
IS.fromList ([Int] -> IntSet)
-> (NaiveIntervalInt -> [Int]) -> NaiveIntervalInt -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> [Int]
forall a. Storable a => Vector a -> [a]
VS.toList (Vector Int -> [Int])
-> (NaiveIntervalInt -> Vector Int) -> NaiveIntervalInt -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntervalValue -> Int) -> NaiveIntervalInt -> Vector Int
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
VS.map (Word32 -> Int
forall a. Enum a => a -> Int
fromEnum (Word32 -> Int)
-> (IntervalValue -> Word32) -> IntervalValue -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntervalValue -> Word32
ivValue) (NaiveIntervalInt -> Vector Int)
-> (NaiveIntervalInt -> NaiveIntervalInt)
-> NaiveIntervalInt
-> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntervalValue -> Bool) -> NaiveIntervalInt -> NaiveIntervalInt
forall a. Storable a => (a -> Bool) -> Vector a -> Vector a
VS.filter (Int -> IntervalValue -> Bool
intervalContains Int
x)

naiveOverlaps :: Interval -> NaiveIntervalInt -> IS.IntSet
naiveOverlaps :: Interval -> NaiveIntervalInt -> IntSet
naiveOverlaps Interval
i = [Int] -> IntSet
IS.fromList ([Int] -> IntSet)
-> (NaiveIntervalInt -> [Int]) -> NaiveIntervalInt -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Interval, Int) -> Int) -> [(Interval, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Interval, Int) -> Int
forall a b. (a, b) -> b
snd ([(Interval, Int)] -> [Int])
-> (NaiveIntervalInt -> [(Interval, Int)])
-> NaiveIntervalInt
-> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Interval -> NaiveIntervalInt -> [(Interval, Int)]
naiveOverlapsWithKeys Interval
i

naiveOverlapsWithKeys :: Interval -> NaiveIntervalInt -> [(Interval, Int)]
naiveOverlapsWithKeys :: Interval -> NaiveIntervalInt -> [(Interval, Int)]
naiveOverlapsWithKeys Interval
i = (IntervalValue -> (Interval, Int))
-> [IntervalValue] -> [(Interval, Int)]
forall a b. (a -> b) -> [a] -> [b]
map IntervalValue -> (Interval, Int)
asPair ([IntervalValue] -> [(Interval, Int)])
-> (NaiveIntervalInt -> [IntervalValue])
-> NaiveIntervalInt
-> [(Interval, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set IntervalValue -> [IntervalValue]
forall a. Set a -> [a]
S.toList (Set IntervalValue -> [IntervalValue])
-> (NaiveIntervalInt -> Set IntervalValue)
-> NaiveIntervalInt
-> [IntervalValue]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Interval -> NaiveIntervalInt -> Set IntervalValue
naiveOverlapsWithKeys' Interval
i

asPair :: IntervalValue -> (Interval, Int)
asPair (IntervalValue Word32
s Word32
e Word32
ix) = (Int -> Int -> Interval
Interval (Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
s) (Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
e), Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
ix)

naiveOverlapsWithKeys' :: Interval -> NaiveIntervalInt -> S.Set IntervalValue
naiveOverlapsWithKeys' :: Interval -> NaiveIntervalInt -> Set IntervalValue
naiveOverlapsWithKeys' (Interval Int
s0 Int
e0) = [IntervalValue] -> Set IntervalValue
forall a. Ord a => [a] -> Set a
S.fromList ([IntervalValue] -> Set IntervalValue)
-> (NaiveIntervalInt -> [IntervalValue])
-> NaiveIntervalInt
-> Set IntervalValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NaiveIntervalInt -> [IntervalValue]
forall a. Storable a => Vector a -> [a]
VS.toList (NaiveIntervalInt -> [IntervalValue])
-> (NaiveIntervalInt -> NaiveIntervalInt)
-> NaiveIntervalInt
-> [IntervalValue]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntervalValue -> Bool) -> NaiveIntervalInt -> NaiveIntervalInt
forall a. Storable a => (a -> Bool) -> Vector a -> Vector a
VS.filter IntervalValue -> Bool
overlap1
    where
        overlap1 :: IntervalValue -> Bool
overlap1 (IntervalValue Word32
s1' Word32
e1' Word32
_)
            | Int
s0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
e0 = Bool
False
            | Word32
s1' Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
e1' = Bool
False
            | Bool
otherwise =
                let
                    s1 :: Int
s1 = Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
s1'
                    e1 :: Int
e1 = Word32 -> Int
forall a. Enum a => a -> Int
fromEnum Word32
e1'
                in (Int
s0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
s1 Bool -> Bool -> Bool
&& Int
s1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
e0) Bool -> Bool -> Bool
|| (Int
s1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
s0 Bool -> Bool -> Bool
&& Int
s0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
e1)

overlaps :: Interval -> IntervalIntMap -> IS.IntSet
overlaps :: Interval -> IntervalIntMap -> IntSet
overlaps Interval
i (IntervalIntMap IntervalIntMapNode
root) = Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
root

overlaps' :: Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i (Leaf NaiveIntervalInt
vec) = Interval -> NaiveIntervalInt -> IntSet
naiveOverlaps Interval
i NaiveIntervalInt
vec
overlaps' Interval
i (InnerNode Int
p IntervalIntMapNode
left IntervalIntMapNode
centre IntervalIntMapNode
right)
    | Interval
i Interval -> Int -> Bool
`intervalAbove` Int
p = Interval -> IntervalIntMapNode -> IntSet
overlaps'  Interval
i IntervalIntMapNode
right IntSet -> IntSet -> IntSet
`IS.union` Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
centre
    | Interval
i Interval -> Int -> Bool
`intervalBelow` Int
p = Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
left IntSet -> IntSet -> IntSet
`IS.union` Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
centre
    | Bool
otherwise = Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
left IntSet -> IntSet -> IntSet
`IS.union` Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
centre IntSet -> IntSet -> IntSet
`IS.union` Interval -> IntervalIntMapNode -> IntSet
overlaps' Interval
i IntervalIntMapNode
right

overlapsWithKeys :: Interval -> IntervalIntMap -> [(Interval, Int)]
overlapsWithKeys :: Interval -> IntervalIntMap -> [(Interval, Int)]
overlapsWithKeys Interval
i (IntervalIntMap IntervalIntMapNode
root) = (IntervalValue -> (Interval, Int))
-> [IntervalValue] -> [(Interval, Int)]
forall a b. (a -> b) -> [a] -> [b]
map IntervalValue -> (Interval, Int)
asPair ([IntervalValue] -> [(Interval, Int)])
-> (Set IntervalValue -> [IntervalValue])
-> Set IntervalValue
-> [(Interval, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set IntervalValue -> [IntervalValue]
forall a. Set a -> [a]
S.toList (Set IntervalValue -> [(Interval, Int)])
-> Set IntervalValue -> [(Interval, Int)]
forall a b. (a -> b) -> a -> b
$ Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
root

overlapsWithKeys' :: Interval -> IntervalIntMapNode -> S.Set IntervalValue
overlapsWithKeys' :: Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i (Leaf NaiveIntervalInt
vec) = Interval -> NaiveIntervalInt -> Set IntervalValue
naiveOverlapsWithKeys' Interval
i NaiveIntervalInt
vec
overlapsWithKeys' Interval
i (InnerNode Int
p IntervalIntMapNode
left IntervalIntMapNode
centre IntervalIntMapNode
right)
    | Interval
i Interval -> Int -> Bool
`intervalAbove` Int
p = Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys'  Interval
i IntervalIntMapNode
right Set IntervalValue -> Set IntervalValue -> Set IntervalValue
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
centre
    | Interval
i Interval -> Int -> Bool
`intervalBelow` Int
p = Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
left Set IntervalValue -> Set IntervalValue -> Set IntervalValue
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
centre
    | Bool
otherwise = Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
left Set IntervalValue -> Set IntervalValue -> Set IntervalValue
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
centre Set IntervalValue -> Set IntervalValue -> Set IntervalValue
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Interval -> IntervalIntMapNode -> Set IntervalValue
overlapsWithKeys' Interval
i IntervalIntMapNode
right

intervalAbove :: Interval -> Int -> Bool
intervalAbove (Interval Int
s Int
_) Int
p = Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
p
intervalBelow :: Interval -> Int -> Bool
intervalBelow (Interval Int
_ Int
e) Int
p = Int
e Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
p