{-# LANGUAGE PartialTypeSignatures #-}

{- | An indexed map is an efficient map with integer keys, that can efficiently retrieve the key from a value.
This is used to efficiently build up a constant pool without duplicating entries.
Because of the specialised nature, its indexes start at 1, not 0. I would apologise but I'm not sorry.
-}
module Data.IndexedMap where

import Control.Lens (Lens', set, view)
import Control.Monad (forM_)

import Data.IntMap qualified as IM
import Data.Map qualified as M
import Data.Vector (Vector)
import Data.Vector qualified as V
import Effectful
import Effectful.State.Static.Local
import GHC.Exts (IsList (..))
import Prelude hiding (lookup)

data IndexedMap a = IndexedMap !(IM.IntMap a) !(M.Map a Int)

{- | An empty indexed map
>>> lookup @String 1 empty
Nothing
-}
empty :: IndexedMap a
empty :: forall a. IndexedMap a
empty = IntMap a -> Map a Int -> IndexedMap a
forall a. IntMap a -> Map a Int -> IndexedMap a
IndexedMap IntMap a
forall a. IntMap a
IM.empty Map a Int
forall k a. Map k a
M.empty

{- | Create an indexed map with a single element
>>> lookup @String 1 (singleton "hello")
Just "hello"
>>> lookup @String 2 (singleton "hello")
Nothing
-}
singleton :: (Ord a) => a -> IndexedMap a
singleton :: forall a. Ord a => a -> IndexedMap a
singleton a
a = IntMap a -> Map a Int -> IndexedMap a
forall a. IntMap a -> Map a Int -> IndexedMap a
IndexedMap (Int -> a -> IntMap a
forall a. Int -> a -> IntMap a
IM.singleton Int
1 a
a) (a -> Int -> Map a Int
forall k a. k -> a -> Map k a
M.singleton a
a Int
1)

lookup :: Int -> IndexedMap a -> Maybe a
lookup :: forall a. Int -> IndexedMap a -> Maybe a
lookup Int
i (IndexedMap IntMap a
m Map a Int
_) = Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
i IntMap a
m

{- | Lookup a value in the map
>>> lookupIndex @String "hello" (singleton "hello")
Just 1
>>> lookupIndex @String "hello" (singleton "world")
Nothing
>>> lookupIndex @String "hello" (singleton "world" <> singleton "hello")
Just 2
-}
lookupIndex :: (Ord a) => a -> IndexedMap a -> Maybe Int
lookupIndex :: forall a. Ord a => a -> IndexedMap a -> Maybe Int
lookupIndex a
a (IndexedMap IntMap a
_ Map a Int
m) = a -> Map a Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a Map a Int
m

{- | Find the index of the first element that satisfies the predicate, if any
>>> lookupIndexWhere (== "hello") (singleton "hello")
Just 1

>>> lookupIndexWhere (== "hello") (singleton "world")
Nothing
-}
lookupIndexWhere :: (a -> Bool) -> IndexedMap a -> Maybe Int
lookupIndexWhere :: forall a. (a -> Bool) -> IndexedMap a -> Maybe Int
lookupIndexWhere a -> Bool
f (IndexedMap IntMap a
m Map a Int
_) = (Int, a) -> Int
forall a b. (a, b) -> a
fst ((Int, a) -> Int) -> Maybe (Int, a) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap a -> Maybe (Int, a)
forall a. IntMap a -> Maybe (Int, a)
IM.lookupMin ((a -> Bool) -> IntMap a -> IntMap a
forall a. (a -> Bool) -> IntMap a -> IntMap a
IM.filter a -> Bool
f IntMap a
m)

-- | Insert a value into the map without checking if it already exists
insert :: (Ord a) => a -> IndexedMap a -> (Int, IndexedMap a)
insert :: forall a. Ord a => a -> IndexedMap a -> (Int, IndexedMap a)
insert a
a (IndexedMap IntMap a
m Map a Int
m') = (Int
i, IntMap a -> Map a Int -> IndexedMap a
forall a. IntMap a -> Map a Int -> IndexedMap a
IndexedMap (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
i a
a IntMap a
m) (a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
a Int
i Map a Int
m'))
  where
    i :: Int
i = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IntMap a -> Int
forall a. IntMap a -> Int
IM.size IntMap a
m

-- | Lookup a value in the map, or insert it if it doesn't exist
lookupOrInsert :: (Ord a) => a -> IndexedMap a -> (Int, IndexedMap a)
lookupOrInsert :: forall a. Ord a => a -> IndexedMap a -> (Int, IndexedMap a)
lookupOrInsert a
a (IndexedMap IntMap a
m Map a Int
m') = case a -> Map a Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a Map a Int
m' of
    Just Int
i -> (Int
i, IntMap a -> Map a Int -> IndexedMap a
forall a. IntMap a -> Map a Int -> IndexedMap a
IndexedMap IntMap a
m Map a Int
m')
    Maybe Int
Nothing -> a -> IndexedMap a -> (Int, IndexedMap a)
forall a. Ord a => a -> IndexedMap a -> (Int, IndexedMap a)
insert a
a (IntMap a -> Map a Int -> IndexedMap a
forall a. IntMap a -> Map a Int -> IndexedMap a
IndexedMap IntMap a
m Map a Int
m')

lookupOrInsertM :: (State (IndexedMap a) :> r, Ord a) => a -> Eff r Int
lookupOrInsertM :: forall a (r :: [Effect]).
(State (IndexedMap a) :> r, Ord a) =>
a -> Eff r Int
lookupOrInsertM = Lens' (IndexedMap a) (IndexedMap a) -> a -> Eff r Int
forall a (r :: [Effect]) b.
(State a :> r, Ord b) =>
Lens' a (IndexedMap b) -> b -> Eff r Int
lookupOrInsertMOver (IndexedMap a -> f (IndexedMap a))
-> IndexedMap a -> f (IndexedMap a)
forall a. a -> a
Lens' (IndexedMap a) (IndexedMap a)
id

lookupOrInsertMOver :: (State a :> r, Ord b) => Lens' a (IndexedMap b) -> b -> Eff r Int
lookupOrInsertMOver :: forall a (r :: [Effect]) b.
(State a :> r, Ord b) =>
Lens' a (IndexedMap b) -> b -> Eff r Int
lookupOrInsertMOver Lens' a (IndexedMap b)
lens b
a = do
    i <- (a -> IndexedMap b) -> Eff r (IndexedMap b)
forall s (es :: [Effect]) a.
(HasCallStack, State s :> es) =>
(s -> a) -> Eff es a
gets (Getting (IndexedMap b) a (IndexedMap b) -> a -> IndexedMap b
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (IndexedMap b) a (IndexedMap b)
Lens' a (IndexedMap b)
lens)
    let (idx, new) = lookupOrInsert a i
    modify (set lens new)
    pure idx

isEmpty :: IndexedMap a -> Bool
isEmpty :: forall a. IndexedMap a -> Bool
isEmpty (IndexedMap IntMap a
m Map a Int
_) = IntMap a -> Bool
forall a. IntMap a -> Bool
IM.null IntMap a
m

{- | O(n) conversion to a vector
This relies on the fact that IndexedMap is strictly increasing in the key

>>> toVector (singleton @Int 1)
[1]

>>> toVector (singleton @Int 1 <> singleton 2)
[1,2]

>>> toVector (singleton @Int 1 <> singleton 2 <> singleton 1)
[1,2]
-}
toVector :: IndexedMap a -> Vector a
toVector :: forall a. IndexedMap a -> Vector a
toVector IndexedMap a
i | IndexedMap a -> Bool
forall a. IndexedMap a -> Bool
isEmpty IndexedMap a
i = Vector a
forall a. Vector a
V.empty
toVector (IndexedMap IntMap a
im Map a Int
_) = do
    let (Int
maxIndex, a
_) = IntMap a -> (Int, a)
forall a. IntMap a -> (Int, a)
IM.findMax IntMap a
im
    Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
maxIndex ((IntMap a
im IntMap a -> Int -> a
forall a. IntMap a -> Int -> a
IM.!) (Int -> a) -> (Int -> Int) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+))

--
-- Instances
--

instance (Show a) => Show (IndexedMap a) where
    show :: IndexedMap a -> String
show (IndexedMap IntMap a
im Map a Int
_) = IntMap a -> String
forall a. Show a => a -> String
show IntMap a
im

instance (Eq a) => Eq (IndexedMap a) where
    (IndexedMap IntMap a
im Map a Int
_) == :: IndexedMap a -> IndexedMap a -> Bool
== (IndexedMap IntMap a
im' Map a Int
_) = IntMap a
im IntMap a -> IntMap a -> Bool
forall a. Eq a => a -> a -> Bool
== IntMap a
im'

instance (Ord a) => Ord (IndexedMap a) where
    compare :: IndexedMap a -> IndexedMap a -> Ordering
compare (IndexedMap IntMap a
im Map a Int
_) (IndexedMap IntMap a
im' Map a Int
_) = IntMap a -> IntMap a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare IntMap a
im IntMap a
im'

instance Foldable IndexedMap where
    foldMap :: forall m a. Monoid m => (a -> m) -> IndexedMap a -> m
foldMap a -> m
f (IndexedMap IntMap a
im Map a Int
_) = (a -> m) -> IntMap a -> m
forall m a. Monoid m => (a -> m) -> IntMap a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f IntMap a
im

instance (Ord a) => IsList (IndexedMap a) where
    type Item (IndexedMap a) = a
    fromList :: [Item (IndexedMap a)] -> IndexedMap a
fromList = (a -> IndexedMap a -> IndexedMap a)
-> IndexedMap a -> [a] -> IndexedMap a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
a IndexedMap a
b -> (Int, IndexedMap a) -> IndexedMap a
forall a b. (a, b) -> b
snd ((Int, IndexedMap a) -> IndexedMap a)
-> (Int, IndexedMap a) -> IndexedMap a
forall a b. (a -> b) -> a -> b
$ a -> IndexedMap a -> (Int, IndexedMap a)
forall a. Ord a => a -> IndexedMap a -> (Int, IndexedMap a)
insert a
a IndexedMap a
b) IndexedMap a
forall a. IndexedMap a
empty
    toList :: IndexedMap a -> [Item (IndexedMap a)]
toList = Vector a -> [a]
Vector a -> [Item (Vector a)]
forall l. IsList l => l -> [Item l]
toList (Vector a -> [a])
-> (IndexedMap a -> Vector a) -> IndexedMap a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IndexedMap a -> Vector a
forall a. IndexedMap a -> Vector a
toVector

{- | Semigroup instance for IndexedMap
 | This is a left-biased union of the two maps
-}
instance (Ord a) => Semigroup (IndexedMap a) where
    IndexedMap a
l <> :: IndexedMap a -> IndexedMap a -> IndexedMap a
<> IndexedMap a
r =
        Eff '[] (IndexedMap a) -> IndexedMap a
forall a. HasCallStack => Eff '[] a -> a
runPureEff (Eff '[] (IndexedMap a) -> IndexedMap a)
-> Eff '[] (IndexedMap a) -> IndexedMap a
forall a b. (a -> b) -> a -> b
$ IndexedMap a
-> Eff '[State (IndexedMap a)] () -> Eff '[] (IndexedMap a)
forall s (es :: [Effect]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es s
execState IndexedMap a
forall a. IndexedMap a
empty (Eff '[State (IndexedMap a)] () -> Eff '[] (IndexedMap a))
-> Eff '[State (IndexedMap a)] () -> Eff '[] (IndexedMap a)
forall a b. (a -> b) -> a -> b
$ do
            Vector a
-> (a -> Eff '[State (IndexedMap a)] Int)
-> Eff '[State (IndexedMap a)] ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (IndexedMap a -> Vector a
forall a. IndexedMap a -> Vector a
toVector IndexedMap a
l) a -> Eff '[State (IndexedMap a)] Int
forall a (r :: [Effect]).
(State (IndexedMap a) :> r, Ord a) =>
a -> Eff r Int
lookupOrInsertM
            Vector a
-> (a -> Eff '[State (IndexedMap a)] Int)
-> Eff '[State (IndexedMap a)] ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (IndexedMap a -> Vector a
forall a. IndexedMap a -> Vector a
toVector IndexedMap a
r) a -> Eff '[State (IndexedMap a)] Int
forall a (r :: [Effect]).
(State (IndexedMap a) :> r, Ord a) =>
a -> Eff r Int
lookupOrInsertM

-- | Monoid instance for IndexedMap
instance (Ord a) => Monoid (IndexedMap a) where
    mempty :: IndexedMap a
mempty = IndexedMap a
forall a. IndexedMap a
empty