{-# LANGUAGE FlexibleContexts #-}

module Data.IntervalIntMap
#ifndef IS_BUILDING_TEST
    ( IntervalIntMap
#else
    ( IntervalIntMap(..)
#endif
    , IntervalIntMapAccumulator
    , IM.Interval(..)
    , fromList
    , elems
    , new
    , insert
    , unsafeFreeze
    , lookup
    , map
    , overlaps
    , overlapsWithKeys
    ) where

import Prelude hiding (lookup, map)

import qualified Data.IntervalIntMap.Internal.IntervalIntIntMap as IM
import qualified Data.IntervalIntMap.Internal.GrowableVector as GV
import qualified Data.Vector.Storable as VS
import qualified Data.IntSet as IS
import           Foreign.Storable (Storable(..))
import           Control.Monad.Primitive (PrimMonad, PrimState)
import           Control.Monad (forM_)
import           Control.Monad.ST (runST)
import           Control.Arrow (second)
import           Control.DeepSeq (NFData(..))


{-| The typical interval map structure models a function of the type @ f :: Int
 - -> Maybe a@. That is, each position in the domain is either annotated by an
 - interval or it is not. When you attempt to insert an interval that overlaps
 - with an existing one, the new value may either (1) replace or (2) by
 - combined with the older one.
 -
 - This is **not** the model here. The model here is @f :: Int -> [a]@! An
 - interval map is a bag of intervals which may overlap. When they do overlap
 - and you query at a position where multiple ones could be active, you get all
 - of them (in some reliable, but unspecified, order; currently insertion
 - order, but this is not an API guarantee).
 -
 - The API uses two objects:
 -
 -  'IntervalIntMapAccumulator': allows insertion. This is a Mutable object and
 -  insertions should be in a `PrimMonad`
 -
 -  'IntervalIntMap': allows querying and operations are pure.
 -
 -}


data IntervalIntMap a = IntervalIntMap !IM.IntervalIntMap
                                       !(VS.Vector a)


instance NFData (IntervalIntMap a) where
    rnf :: IntervalIntMap a -> ()
rnf (IntervalIntMap IntervalIntMap
im Vector a
v) = IntervalIntMap -> ()
forall a. NFData a => a -> ()
rnf IntervalIntMap
im () -> () -> ()
`seq` Vector a -> ()
forall a. NFData a => a -> ()
rnf Vector a
v

data IntervalIntMapAccumulator s a = IntervalIntMapAccumulator
                                        !(GV.GrowableVector s (IM.IntervalValue))
                                        !(GV.GrowableVector s a)


-- |Create an 'IntervalIntMap' from a list of (key, value)
fromList :: Storable a => [(IM.Interval, a)] -> IntervalIntMap a
fromList :: [(Interval, a)] -> IntervalIntMap a
fromList [(Interval, a)]
vs = (forall s. ST s (IntervalIntMap a)) -> IntervalIntMap a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (IntervalIntMap a)) -> IntervalIntMap a)
-> (forall s. ST s (IntervalIntMap a)) -> IntervalIntMap a
forall a b. (a -> b) -> a -> b
$ do
    IntervalIntMapAccumulator s a
acc <- ST s (IntervalIntMapAccumulator s a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
m (IntervalIntMapAccumulator (PrimState m) a)
new
    [(Interval, a)] -> ((Interval, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Interval, a)]
vs (((Interval, a) -> ST s ()) -> ST s ())
-> ((Interval, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Interval
i,a
v) -> Interval
-> a -> IntervalIntMapAccumulator (PrimState (ST s)) a -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Interval -> a -> IntervalIntMapAccumulator (PrimState m) a -> m ()
insert Interval
i a
v IntervalIntMapAccumulator s a
IntervalIntMapAccumulator (PrimState (ST s)) a
acc
    IntervalIntMapAccumulator (PrimState (ST s)) a
-> ST s (IntervalIntMap a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
IntervalIntMapAccumulator (PrimState m) a -> m (IntervalIntMap a)
unsafeFreeze IntervalIntMapAccumulator s a
IntervalIntMapAccumulator (PrimState (ST s)) a
acc

elems :: Storable a => IntervalIntMap a -> [a]
elems :: IntervalIntMap a -> [a]
elems (IntervalIntMap IntervalIntMap
_ Vector a
vals) = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
VS.toList Vector a
vals

-- |New (empty) accumulator
new :: (PrimMonad m, Storable a) => m (IntervalIntMapAccumulator (PrimState m) a)
new :: m (IntervalIntMapAccumulator (PrimState m) a)
new = GrowableVector (PrimState m) IntervalValue
-> GrowableVector (PrimState m) a
-> IntervalIntMapAccumulator (PrimState m) a
forall s a.
GrowableVector s IntervalValue
-> GrowableVector s a -> IntervalIntMapAccumulator s a
IntervalIntMapAccumulator (GrowableVector (PrimState m) IntervalValue
 -> GrowableVector (PrimState m) a
 -> IntervalIntMapAccumulator (PrimState m) a)
-> m (GrowableVector (PrimState m) IntervalValue)
-> m (GrowableVector (PrimState m) a
      -> IntervalIntMapAccumulator (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (GrowableVector (PrimState m) IntervalValue)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
m (GrowableVector (PrimState m) a)
GV.new m (GrowableVector (PrimState m) a
   -> IntervalIntMapAccumulator (PrimState m) a)
-> m (GrowableVector (PrimState m) a)
-> m (IntervalIntMapAccumulator (PrimState m) a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> m (GrowableVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
m (GrowableVector (PrimState m) a)
GV.new


-- |Insert a value into an accumulator
insert :: (PrimMonad m, Storable a) => IM.Interval -> a -> IntervalIntMapAccumulator (PrimState m) a -> m ()
insert :: Interval -> a -> IntervalIntMapAccumulator (PrimState m) a -> m ()
insert (IM.Interval Int
s Int
e) a
v (IntervalIntMapAccumulator GrowableVector (PrimState m) IntervalValue
ivs GrowableVector (PrimState m) a
dat) = do
    Int
ix <- GrowableVector (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
GrowableVector (PrimState m) a -> m Int
GV.length GrowableVector (PrimState m) a
dat
    a -> GrowableVector (PrimState m) a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
a -> GrowableVector (PrimState m) a -> m ()
GV.pushBack a
v GrowableVector (PrimState m) a
dat
    IntervalValue -> GrowableVector (PrimState m) IntervalValue -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
a -> GrowableVector (PrimState m) a -> m ()
GV.pushBack (Word32 -> Word32 -> Word32 -> IntervalValue
IM.IntervalValue (Int -> Word32
forall a. Enum a => Int -> a
toEnum Int
s) (Int -> Word32
forall a. Enum a => Int -> a
toEnum Int
e) (Int -> Word32
forall a. Enum a => Int -> a
toEnum Int
ix)) GrowableVector (PrimState m) IntervalValue
ivs


-- |Transform an 'IntervalIntMapAccumulator' into an 'IntervalIntMap'. This is
--unsafe as the accumulator should **not** be used after this operation is
--performed.
unsafeFreeze :: (PrimMonad m, Storable a) => IntervalIntMapAccumulator (PrimState m) a -> m (IntervalIntMap a)
unsafeFreeze :: IntervalIntMapAccumulator (PrimState m) a -> m (IntervalIntMap a)
unsafeFreeze (IntervalIntMapAccumulator GrowableVector (PrimState m) IntervalValue
ivs GrowableVector (PrimState m) a
values) =
    IntervalIntMap -> Vector a -> IntervalIntMap a
forall a. IntervalIntMap -> Vector a -> IntervalIntMap a
IntervalIntMap
        (IntervalIntMap -> Vector a -> IntervalIntMap a)
-> m IntervalIntMap -> m (Vector a -> IntervalIntMap a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (NaiveIntervalInt -> IntervalIntMap
IM.freeze (NaiveIntervalInt -> IntervalIntMap)
-> m NaiveIntervalInt -> m IntervalIntMap
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GrowableVector (PrimState m) IntervalValue -> m NaiveIntervalInt
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
GrowableVector (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowableVector (PrimState m) IntervalValue
ivs)
        m (Vector a -> IntervalIntMap a)
-> m (Vector a) -> m (IntervalIntMap a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> GrowableVector (PrimState m) a -> m (Vector a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
GrowableVector (PrimState m) a -> m (Vector a)
GV.unsafeFreeze GrowableVector (PrimState m) a
values

indexAll :: Storable a => VS.Vector a -> IS.IntSet -> [a]
indexAll :: Vector a -> IntSet -> [a]
indexAll Vector a
values = ((Int -> a) -> [Int] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> a) -> [Int] -> [a]) -> (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Storable a => Vector a -> Int -> a
(VS.!) Vector a
values) ([Int] -> [a]) -> (IntSet -> [Int]) -> IntSet -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> [Int]
IS.toList

-- |Lookup all values whose keys intersect the given position
lookup ::  Storable a => Int -> IntervalIntMap a -> [a]
lookup :: Int -> IntervalIntMap a -> [a]
lookup Int
p (IntervalIntMap IntervalIntMap
imap Vector a
values) = Vector a -> IntSet -> [a]
forall a. Storable a => Vector a -> IntSet -> [a]
indexAll Vector a
values (IntSet -> [a]) -> IntSet -> [a]
forall a b. (a -> b) -> a -> b
$ Int -> IntervalIntMap -> IntSet
IM.lookup Int
p IntervalIntMap
imap

-- |Map: note that both the input and output types must be instances of
-- Storable, so this is not a functor.
map :: (Storable a, Storable b) => (a -> b) -> IntervalIntMap a -> IntervalIntMap b
map :: (a -> b) -> IntervalIntMap a -> IntervalIntMap b
map a -> b
f (IntervalIntMap IntervalIntMap
im Vector a
vs) = IntervalIntMap -> Vector b -> IntervalIntMap b
forall a. IntervalIntMap -> Vector a -> IntervalIntMap a
IntervalIntMap IntervalIntMap
im ((a -> b) -> Vector a -> Vector b
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
VS.map a -> b
f Vector a
vs)

-- |Lookup all values that overlap with the given input
overlaps :: Storable a => IM.Interval -> IntervalIntMap a -> [a]
overlaps :: Interval -> IntervalIntMap a -> [a]
overlaps Interval
i = ((Interval, a) -> a) -> [(Interval, a)] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Interval, a) -> a
forall a b. (a, b) -> b
snd ([(Interval, a)] -> [a])
-> (IntervalIntMap a -> [(Interval, a)]) -> IntervalIntMap a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Interval -> IntervalIntMap a -> [(Interval, a)]
forall a.
Storable a =>
Interval -> IntervalIntMap a -> [(Interval, a)]
overlapsWithKeys Interval
i

-- |Lookup all values that overlap with the given input
overlapsWithKeys :: Storable a => IM.Interval -> IntervalIntMap a -> [(IM.Interval,a)]
overlapsWithKeys :: Interval -> IntervalIntMap a -> [(Interval, a)]
overlapsWithKeys Interval
i (IntervalIntMap IntervalIntMap
imap Vector a
values) = ((Interval, Int) -> (Interval, a))
-> [(Interval, Int)] -> [(Interval, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> a) -> (Interval, Int) -> (Interval, a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((Int -> a) -> (Interval, Int) -> (Interval, a))
-> (Int -> a) -> (Interval, Int) -> (Interval, a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Storable a => Vector a -> Int -> a
(VS.!) Vector a
values) ([(Interval, Int)] -> [(Interval, a)])
-> [(Interval, Int)] -> [(Interval, a)]
forall a b. (a -> b) -> a -> b
$ Interval -> IntervalIntMap -> [(Interval, Int)]
IM.overlapsWithKeys Interval
i IntervalIntMap
imap