{-# 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, since the constant pool is indexed by integers, and we often need to check if a value is already in the constant pool before inserting it.
Because of the specialised nature, its indexes start at 1, not 0. I would apologise but I'm not sorry.
-}
module H2JVM.Internal.IndexedMap (
    -- * Types
    IndexedMap,
    Index,
    indexValue,

    -- * Construction
    empty,
    singleton,

    -- * Lookup
    lookup,
    lookupIndex,
    lookupIndexWhere,
    isEmpty,

    -- * Insertion
    insert,
    lookupOrInsert,
    lookupOrInsertM,
    lookupOrInsertMOver,

    -- * Conversion
    toVector,
)
where

import Control.Lens (Lens', set, view)
import Control.Monad (forM_)
import Data.Vector (Vector)
import Data.Word (Word32)
import Effectful
import Effectful.State.Static.Local
import GHC.Exts (IsList (..))
import Prelude hiding (lookup)

import Data.IntMap qualified as IM
import Data.Map qualified as M
import Data.Vector qualified as V

-- | An index into the map.
newtype Index = Index Word32 deriving (Index -> Index -> Bool
(Index -> Index -> Bool) -> (Index -> Index -> Bool) -> Eq Index
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Index -> Index -> Bool
== :: Index -> Index -> Bool
$c/= :: Index -> Index -> Bool
/= :: Index -> Index -> Bool
Eq, Integer -> Index
Index -> Index
Index -> Index -> Index
(Index -> Index -> Index)
-> (Index -> Index -> Index)
-> (Index -> Index -> Index)
-> (Index -> Index)
-> (Index -> Index)
-> (Index -> Index)
-> (Integer -> Index)
-> Num Index
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: Index -> Index -> Index
+ :: Index -> Index -> Index
$c- :: Index -> Index -> Index
- :: Index -> Index -> Index
$c* :: Index -> Index -> Index
* :: Index -> Index -> Index
$cnegate :: Index -> Index
negate :: Index -> Index
$cabs :: Index -> Index
abs :: Index -> Index
$csignum :: Index -> Index
signum :: Index -> Index
$cfromInteger :: Integer -> Index
fromInteger :: Integer -> Index
Num, Eq Index
Eq Index =>
(Index -> Index -> Ordering)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Bool)
-> (Index -> Index -> Index)
-> (Index -> Index -> Index)
-> Ord Index
Index -> Index -> Bool
Index -> Index -> Ordering
Index -> Index -> Index
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Index -> Index -> Ordering
compare :: Index -> Index -> Ordering
$c< :: Index -> Index -> Bool
< :: Index -> Index -> Bool
$c<= :: Index -> Index -> Bool
<= :: Index -> Index -> Bool
$c> :: Index -> Index -> Bool
> :: Index -> Index -> Bool
$c>= :: Index -> Index -> Bool
>= :: Index -> Index -> Bool
$cmax :: Index -> Index -> Index
max :: Index -> Index -> Index
$cmin :: Index -> Index -> Index
min :: Index -> Index -> Index
Ord, Int -> Index -> ShowS
[Index] -> ShowS
Index -> String
(Int -> Index -> ShowS)
-> (Index -> String) -> ([Index] -> ShowS) -> Show Index
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Index -> ShowS
showsPrec :: Int -> Index -> ShowS
$cshow :: Index -> String
show :: Index -> String
$cshowList :: [Index] -> ShowS
showList :: [Index] -> ShowS
Show)

-- | Get the raw value of an index.
indexValue :: Index -> Word32
indexValue :: Index -> Word32
indexValue (Index Word32
i) = Word32
i

-- | An indexed map is a map from integer indexes to values, with a reverse map from values to indexes.
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 a value in the map by its index.

>>> lookup @String 1 (singleton "hello")
Just "hello"
-}
lookup :: Index -> IndexedMap a -> Maybe a
lookup :: forall a. Index -> IndexedMap a -> Maybe a
lookup (Index Word32
i) (IndexedMap IntMap a
m Map a Int
_) = Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i) IntMap a
m

{- | Lookup a value in the map, returning its index if it exists.

>>> 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 Index
lookupIndex :: forall a. Ord a => a -> IndexedMap a -> Maybe Index
lookupIndex a
a (IndexedMap IntMap a
_ Map a Int
m) = Word32 -> Index
Index (Word32 -> Index) -> (Int -> Word32) -> Int -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Index) -> Maybe Int -> Maybe Index
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 Index
lookupIndexWhere :: forall a. (a -> Bool) -> IndexedMap a -> Maybe Index
lookupIndexWhere a -> Bool
f (IndexedMap IntMap a
m Map a Int
_) = Word32 -> Index
Index (Word32 -> Index) -> ((Int, a) -> Word32) -> (Int, a) -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> ((Int, a) -> Int) -> (Int, a) -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst ((Int, a) -> Index) -> Maybe (Int, a) -> Maybe Index
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.
In other words, this will insert a duplicate value if it already exists in the map, and return a new index for it.

>>> insert "hello" empty
(Index 1,fromList [(1,"hello")])

>>> insert "world" (singleton "hello")
(Index 2,fromList [(1,"hello"),(2,"world")])
>>> insert "hello" (singleton "hello")
(Index 2,fromList [(1,"hello"),(2,"hello")])
-}
insert :: Ord a => a -> IndexedMap a -> (Index, IndexedMap a)
insert :: forall a. Ord a => a -> IndexedMap a -> (Index, IndexedMap a)
insert a
a (IndexedMap IntMap a
m Map a Int
m') = (Word32 -> Index
Index (Word32 -> Index) -> Word32 -> Index
forall a b. (a -> b) -> a -> b
$ Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral 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.
If the value already exists, this will return the existing index and the original map.

>>> lookupOrInsert "hello" (singleton "hello")
(Index 1,fromList [(1,"hello")])
>>> lookupOrInsert "world" (singleton "hello")
(Index 2,fromList [(1,"hello"),(2,"world")])
-}
lookupOrInsert :: Ord a => a -> IndexedMap a -> (Index, IndexedMap a)
lookupOrInsert :: forall a. Ord a => a -> IndexedMap a -> (Index, 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 -> (Word32 -> Index
Index (Word32 -> Index) -> Word32 -> Index
forall a b. (a -> b) -> a -> b
$ Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral 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 -> (Index, IndexedMap a)
forall a. Ord a => a -> IndexedMap a -> (Index, 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')

-- | A monadic version of 'lookupOrInsert' that can be used in a state monad.
lookupOrInsertM :: (State (IndexedMap a) :> r, Ord a) => a -> Eff r Index
lookupOrInsertM :: forall a (r :: [Effect]).
(State (IndexedMap a) :> r, Ord a) =>
a -> Eff r Index
lookupOrInsertM = Lens' (IndexedMap a) (IndexedMap a) -> a -> Eff r Index
forall a (r :: [Effect]) b.
(State a :> r, Ord b) =>
Lens' a (IndexedMap b) -> b -> Eff r Index
lookupOrInsertMOver (IndexedMap a -> f (IndexedMap a))
-> IndexedMap a -> f (IndexedMap a)
forall a. a -> a
Lens' (IndexedMap a) (IndexedMap a)
id

{- | A more general version of 'lookupOrInsertM' that allows you to specify which lens to use for the state.
This is useful if you have an 'IndexedMap' within a larger state, and you want to avoid having to manually get and put the 'IndexedMap' every time you want to lookup or insert a value.
-}
lookupOrInsertMOver :: (State a :> r, Ord b) => Lens' a (IndexedMap b) -> b -> Eff r Index
lookupOrInsertMOver :: forall a (r :: [Effect]) b.
(State a :> r, Ord b) =>
Lens' a (IndexedMap b) -> b -> Eff r Index
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

{- | Check if the map is empty

>>> isEmpty empty
True
>>> isEmpty (singleton "hello")
False
-}
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'. Duplicates are not removed, and the order of the vector is the order of the indexes in the map.
This relies on the fact that 'IndexedMap' is strictly increasing in the key, so we can just generate a vector of the appropriate length and fill it with the values from the map.

>>> 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 -> (Index, IndexedMap a) -> IndexedMap a
forall a b. (a, b) -> b
snd ((Index, IndexedMap a) -> IndexedMap a)
-> (Index, IndexedMap a) -> IndexedMap a
forall a b. (a -> b) -> a -> b
$ a -> IndexedMap a -> (Index, IndexedMap a)
forall a. Ord a => a -> IndexedMap a -> (Index, 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

-- | 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)] Index)
-> 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)] Index
forall a (r :: [Effect]).
(State (IndexedMap a) :> r, Ord a) =>
a -> Eff r Index
lookupOrInsertM
            Vector a
-> (a -> Eff '[State (IndexedMap a)] Index)
-> 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)] Index
forall a (r :: [Effect]).
(State (IndexedMap a) :> r, Ord a) =>
a -> Eff r Index
lookupOrInsertM

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