content: {-# 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 }
  