{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
module Data.IMap
    ( IMap
    , Run(..)
    , empty
    , Data.IMap.null
    , singleton
    , insert
    , delete
    , restrict
    , lookup
    , splitLE
    , intersectionWith
    , mapMaybe
    , addToKeys
    , unsafeUnion
    , fromList
    , unsafeRuns
    , unsafeToAscList
    ) where

import Data.List (foldl')
import Data.Monoid
import Data.IntMap.Strict (IntMap)
import GHC.Generics
import Control.DeepSeq
import Prelude hiding (lookup)
import qualified Data.IntMap.Strict as IM

-- | Semantically, 'IMap' and 'IntMap' are identical; but 'IMap' is more
-- efficient when large sequences of contiguous keys are mapped to the same
-- value.
newtype IMap a = IMap { _runs :: IntMap (Run a) } deriving (Show, Functor, Read, Generic, NFData)

{-# INLINE unsafeRuns #-}
-- | This function is unsafe because 'IMap's that compare equal may split their
-- runs into different chunks; consumers must promise that they do not treat
-- run boundaries specially.
unsafeRuns :: IMap a -> IntMap (Run a)
unsafeRuns = _runs

instance Eq a => Eq (IMap a) where
    IMap m == IMap m' = go (IM.toAscList m) (IM.toAscList m') where
        go ((k, Run n a):kvs) ((k', Run n' a'):kvs')
            = k == k' && a == a' && case compare n n' of
                LT -> go kvs ((k'+n, Run (n'-n) a'):kvs')
                EQ -> go kvs kvs'
                GT -> go ((k+n', Run (n-n') a):kvs) kvs'
        go [] [] = True
        go _ _ = False

instance Ord a => Ord (IMap a) where
    compare (IMap m) (IMap m') = go (IM.toAscList m) (IM.toAscList m') where
        go [] [] = EQ
        go [] _  = LT
        go _  [] = GT
        go ((k, Run n a):kvs) ((k', Run n' a'):kvs')
            = compare k k' <> compare a a' <> case compare n n' of
                LT -> go kvs ((k'+n, Run (n'-n) a'):kvs')
                EQ -> go kvs kvs'
                GT -> go ((k+n', Run (n-n') a):kvs) kvs'

-- | Zippy: '(<*>)' combines values at equal keys, discarding any values whose
-- key is in only one of its two arguments.
instance Applicative IMap where
    pure a = IMap . IM.fromDistinctAscList $
        [ (minBound, Run maxBound a)
        , (-1, Run maxBound a)
        , (maxBound-1, Run 2 a)
        ]
    (<*>) = intersectionWith ($)

-- | @Run n a@ represents @n@ copies of the value @a@.
data Run a = Run
    { len :: !Int
    , val :: !a
    } deriving (Eq, Ord, Read, Show, Functor, Generic, NFData)

instance Foldable    Run where foldMap f r = f (val r)
instance Traversable Run where sequenceA (Run n v) = Run n <$> v

empty :: IMap a
empty = IMap IM.empty

null :: IMap a -> Bool
null = IM.null . _runs

singleton :: Int -> Run a -> IMap a
singleton k r
    | len r >= 1 = IMap (IM.singleton k r)
    | otherwise = empty

insert :: Int -> Run a -> IMap a -> IMap a
insert k r m
    | len r < 1 = m
    | otherwise = m { _runs = IM.insert k r (_runs (delete k r m)) }

{-# INLINE delete #-}
delete :: Int -> Run ignored -> IMap a -> IMap a
delete k r m
    | len r < 1 = m
    | otherwise = m { _runs = IM.union (_runs lt) (_runs gt) }
    where
    (lt, ge) = splitLE (k-1) m
    (_ , gt) = splitLE (k+len r-1) ge

-- | Given a range of keys (as specified by a starting key and a length for
-- consistency with other functions in this module), restrict the map to keys
-- in that range. @restrict k r m@ is equivalent to @intersectionWith const m
-- (insert k r empty)@ but potentially more efficient.
restrict :: Int -> Run ignored -> IMap a -> IMap a
restrict k r = id
    . snd
    . splitLE (k-1)
    . fst
    . splitLE (k+len r-1)

lookup :: Int -> IMap a -> Maybe a
lookup k m = case IM.lookupLE k (_runs m) of
    Just (k', Run n a) | k < k'+n -> Just a
    _ -> Nothing

-- | @splitLE n m@ produces a tuple @(le, gt)@ where @le@ has all the
-- associations of @m@ where the keys are @<= n@ and @gt@ has all the
-- associations of @m@ where the keys are @> n@.
splitLE :: Int -> IMap a -> (IMap a, IMap a)
splitLE k m = case IM.lookupLE k (_runs m) of
    Nothing -> (empty, m)
    Just (k', r@(Run n _)) -> case (k' + n - 1 <= k, k' == k) of
        (True , False) -> (m { _runs = lt }, m { _runs = gt })
        (True , True ) -> (m { _runs = IM.insert k r lt }, m { _runs = gt })
        (False, _    ) -> ( m { _runs = IM.insert k'    r { len =     1 + k - k' } lt' }
                          , m { _runs = IM.insert (k+1) r { len = n - 1 - k + k' } gt' }
                          )
        where
        (lt', gt') = IM.split k' (_runs m)
    where
    (lt, gt) = IM.split k (_runs m)

-- | Increment all keys by the given amount. This is like
-- 'IM.mapKeysMonotonic', but restricted to partially-applied addition.
addToKeys :: Int -> IMap a -> IMap a
addToKeys n m = m { _runs = IM.mapKeysMonotonic (n+) (_runs m) }

-- TODO: This is pretty inefficient. IntMap offers some splitting functions
-- that should make it possible to be more efficient here (though the
-- implementation would be significantly messier).
intersectionWith :: (a -> b -> c) -> IMap a -> IMap b -> IMap c
intersectionWith f (IMap runsa) (IMap runsb)
    = IMap . IM.fromDistinctAscList $ merge (IM.toAscList runsa) (IM.toAscList runsb)
    where
    merge as@((ka, ra):at) bs@((kb, rb):bt)
        | ka' < kb = merge at bs
        | kb' < ka = merge as bt
        | otherwise = (kc, Run (kc' - kc + 1) vc) : case compare ka' kb' of
            LT -> merge at bs
            EQ -> merge at bt
            GT -> merge as bt
        where
        ka' = ka + len ra - 1
        kb' = kb + len rb - 1
        kc  = max ka  kb
        kc' = min ka' kb'
        vc  = f (val ra) (val rb)
    merge _ _ = []

mapMaybe :: (a -> Maybe b) -> IMap a -> IMap b
mapMaybe f (IMap runs) = IMap (IM.mapMaybe (traverse f) runs)

fromList :: [(Int, Run a)] -> IMap a
fromList = foldl' (\m (k, r) -> insert k r m) empty

-- | This function is unsafe because 'IMap's that compare equal may split their
-- runs into different chunks; consumers must promise that they do not treat
-- run boundaries specially.
unsafeToAscList :: IMap a -> [(Int, Run a)]
unsafeToAscList = IM.toAscList . _runs

-- | This function is unsafe because it assumes there is no overlap between its
-- arguments. That is, in the call @unsafeUnion a b@, the caller must guarantee
-- that if @lookup k a = Just v@ then @lookup k b = Nothing@ and vice versa.
unsafeUnion :: IMap a -> IMap a -> IMap a
unsafeUnion a b = IMap { _runs = _runs a `IM.union` _runs b }