-- | Module    : Termonad.IdMap
-- Description : A Map that keeps track of the ID of values
-- Copyright   : (c) Dennis Gosnell, 2023
-- License     : BSD3
-- Stability   : experimental
-- Portability : POSIX

module Termonad.IdMap.Internal where

import Termonad.Prelude

import Control.Lens (FoldableWithIndex, ifoldMap, Index, IxValue, Traversal', Ixed (ix))
import qualified Data.Foldable as Foldable
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap

newtype IdMapKey = IdMapKey { IdMapKey -> Int
unIdMapKey :: Int }
  deriving stock (IdMapKey -> IdMapKey -> Bool
(IdMapKey -> IdMapKey -> Bool)
-> (IdMapKey -> IdMapKey -> Bool) -> Eq IdMapKey
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: IdMapKey -> IdMapKey -> Bool
== :: IdMapKey -> IdMapKey -> Bool
$c/= :: IdMapKey -> IdMapKey -> Bool
/= :: IdMapKey -> IdMapKey -> Bool
Eq, Int -> IdMapKey -> ShowS
[IdMapKey] -> ShowS
IdMapKey -> String
(Int -> IdMapKey -> ShowS)
-> (IdMapKey -> String) -> ([IdMapKey] -> ShowS) -> Show IdMapKey
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> IdMapKey -> ShowS
showsPrec :: Int -> IdMapKey -> ShowS
$cshow :: IdMapKey -> String
show :: IdMapKey -> String
$cshowList :: [IdMapKey] -> ShowS
showList :: [IdMapKey] -> ShowS
Show)

data IdMap a = IdMap
  { forall a. IdMap a -> IntMap a
idMap :: !(IntMap a)
  , forall a. IdMap a -> Int
nextId :: !Int
  }
  deriving stock Int -> IdMap a -> ShowS
[IdMap a] -> ShowS
IdMap a -> String
(Int -> IdMap a -> ShowS)
-> (IdMap a -> String) -> ([IdMap a] -> ShowS) -> Show (IdMap a)
forall a. Show a => Int -> IdMap a -> ShowS
forall a. Show a => [IdMap a] -> ShowS
forall a. Show a => IdMap a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> IdMap a -> ShowS
showsPrec :: Int -> IdMap a -> ShowS
$cshow :: forall a. Show a => IdMap a -> String
show :: IdMap a -> String
$cshowList :: forall a. Show a => [IdMap a] -> ShowS
showList :: [IdMap a] -> ShowS
Show

-- | 'IdMap's are equal if they contain the same elements at the same keys.
--
-- >>> let (helloKey, idmapA) = insertIdMap "hello" emptyIdMap
-- >>> let (_, idmapB) = singletonIdMap "hello"
-- >>> idmapA == idmapB
-- True
--
-- Note that if you delete and reinsert a value, it will get a different key,
-- so will no longer be equal.
--
-- >>> let (_, idmapA') = insertIdMap "hello" $ deleteIdMap helloKey idmapA
-- >>> idmapA' == idmapB
-- False
--
-- However, 'IdMap's don't check the 'nextId' field when determining equality.
--
-- >>> let (byeKey, idmapA'') = insertIdMap "bye" idmapA
-- >>> let idmapA''' = deleteIdMap byeKey idmapA''
-- >>> idmapA''' == idmapB
-- True
instance Eq a => Eq (IdMap a) where
  (IdMap IntMap a
idMapA Int
_) == :: IdMap a -> IdMap a -> Bool
== (IdMap IntMap a
idMapB Int
_) = IntMap a
idMapA IntMap a -> IntMap a -> Bool
forall a. Eq a => a -> a -> Bool
== IntMap a
idMapB

instance Functor IdMap where
  fmap :: forall a b. (a -> b) -> IdMap a -> IdMap b
fmap a -> b
f IdMap{IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap, Int
nextId :: forall a. IdMap a -> Int
nextId :: Int
nextId} = IdMap { idMap :: IntMap b
idMap = (a -> b) -> IntMap a -> IntMap b
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f IntMap a
idMap, Int
nextId :: Int
nextId :: Int
nextId }

instance Foldable IdMap where
  foldMap :: forall m a. Monoid m => (a -> m) -> IdMap a -> m
foldMap a -> m
f IdMap a
m = (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
Foldable.foldMap a -> m
f (IntMap a -> m) -> IntMap a -> m
forall a b. (a -> b) -> a -> b
$ IdMap a -> IntMap a
forall a. IdMap a -> IntMap a
idMap IdMap a
m

instance FoldableWithIndex Int IdMap where
  ifoldMap :: forall m a. Monoid m => (Int -> a -> m) -> IdMap a -> m
ifoldMap Int -> a -> m
f IdMap a
m = (Int -> a -> m) -> IntMap a -> m
forall m a. Monoid m => (Int -> a -> m) -> IntMap a -> m
forall i (f :: * -> *) m a.
(FoldableWithIndex i f, Monoid m) =>
(i -> a -> m) -> f a -> m
ifoldMap Int -> a -> m
f (IntMap a -> m) -> IntMap a -> m
forall a b. (a -> b) -> a -> b
$ IdMap a -> IntMap a
forall a. IdMap a -> IntMap a
idMap IdMap a
m

instance Traversable IdMap where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> IdMap a -> f (IdMap b)
traverse a -> f b
f IdMap{IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap, Int
nextId :: forall a. IdMap a -> Int
nextId :: Int
nextId} =
    (IntMap b -> IdMap b) -> f (IntMap b) -> f (IdMap b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\IntMap b
m -> IdMap { idMap :: IntMap b
idMap = IntMap b
m, Int
nextId :: Int
nextId :: Int
nextId }) ((a -> f b) -> IntMap a -> f (IntMap b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> IntMap a -> f (IntMap b)
traverse a -> f b
f IntMap a
idMap)

type instance Index (IdMap a) = IdMapKey
type instance IxValue (IdMap a) = a

instance Ixed (IdMap a) where
  ix :: IdMapKey -> Traversal' (IdMap a) a
  ix :: IdMapKey -> Traversal' (IdMap a) a
ix (IdMapKey Int
i) a -> f a
f IdMap{IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap, Int
nextId :: forall a. IdMap a -> Int
nextId :: Int
nextId} =
    case Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
i IntMap a
idMap of
      Just a
v -> (a -> IdMap a) -> f a -> f (IdMap a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> IdMap a
update (a -> f a
f a
v) -- f v <&> \v' -> IntMap.insert k v' m
      Maybe a
Nothing -> IdMap a -> f (IdMap a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IdMap{IntMap a
idMap :: IntMap a
idMap :: IntMap a
idMap, Int
nextId :: Int
nextId :: Int
nextId}
    where
      update :: a -> IdMap a
      update :: a -> IdMap a
update a
v' =
        IdMap
          { idMap :: IntMap a
idMap = (a -> a) -> Int -> IntMap a -> IntMap a
forall a. (a -> a) -> Int -> IntMap a -> IntMap a
IntMap.adjust (a -> a -> a
forall a b. a -> b -> a
const a
v') Int
i IntMap a
idMap
          , Int
nextId :: Int
nextId :: Int
nextId
          }

initialId :: Int
initialId :: Int
initialId = Int
0

-- | Get the next available ID.
--
-- >>> succId 3
-- 4
succId :: Int -> Int
succId :: Int -> Int
succId Int
i = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- | An initial 'IdMap' with no values.
--
-- >>> emptyIdMap
-- IdMap {idMap = fromList [], nextId = 0}
emptyIdMap :: IdMap a
emptyIdMap :: forall a. IdMap a
emptyIdMap = IdMap { idMap :: IntMap a
idMap = IntMap a
forall a. Monoid a => a
mempty, nextId :: Int
nextId = Int
0 }

-- | Insert a value into an 'IdMap'.  Returns the key for the newly inserted
-- item.
--
-- >>> let (key, idmap) = insertIdMap "hello" emptyIdMap
-- >>> (key, idmap)
-- (IdMapKey {unIdMapKey = 0},IdMap {idMap = fromList [(0,"hello")], nextId = 1})
--
-- >>> insertIdMap "zoom" idmap
-- (IdMapKey {unIdMapKey = 1},IdMap {idMap = fromList [(0,"hello"),(1,"zoom")], nextId = 2})
insertIdMap :: a -> IdMap a -> (IdMapKey, IdMap a)
insertIdMap :: forall a. a -> IdMap a -> (IdMapKey, IdMap a)
insertIdMap a
a IdMap {IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap, Int
nextId :: forall a. IdMap a -> Int
nextId :: Int
nextId} =
  let newMap :: IntMap a
newMap = Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
nextId a
a IntMap a
idMap
      newNextId :: Int
newNextId = Int
nextId Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  in (Int -> IdMapKey
IdMapKey Int
nextId, IdMap { idMap :: IntMap a
idMap = IntMap a
newMap, nextId :: Int
nextId = Int
newNextId })

-- | Create an 'IdMap' with a single value.
--
-- >>> singletonIdMap "hello"
-- (IdMapKey {unIdMapKey = 0},IdMap {idMap = fromList [(0,"hello")], nextId = 1})
--
-- prop> \a -> insertIdMap a emptyIdMap == singletonIdMap a
singletonIdMap :: a -> (IdMapKey, IdMap a)
singletonIdMap :: forall a. a -> (IdMapKey, IdMap a)
singletonIdMap a
a = a -> IdMap a -> (IdMapKey, IdMap a)
forall a. a -> IdMap a -> (IdMapKey, IdMap a)
insertIdMap a
a IdMap a
forall a. IdMap a
emptyIdMap

-- | Lookup the given key in an 'IdMap'.
--
-- >>> let (key, idmap) = insertIdMap "hello" emptyIdMap
-- >>> lookupIdMap key idmap
-- Just "hello"
--
-- Trying to lookup keys that don't exist returns 'Nothing':
--
-- >>> let idmap' = deleteIdMap key idmap
-- >>> lookupIdMap key idmap'
-- Nothing
lookupIdMap :: IdMapKey -> IdMap a -> Maybe a
lookupIdMap :: forall a. IdMapKey -> IdMap a -> Maybe a
lookupIdMap (IdMapKey Int
k) IdMap {IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap} = Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k IntMap a
idMap

-- | List all keys in an 'IdMap'.
--
-- >>> let (_, idmap) = singletonIdMap "hello"
-- >>> let (_, idmap') = insertIdMap "bye" idmap
-- >>> keysIdMap idmap'
-- [IdMapKey {unIdMapKey = 0},IdMapKey {unIdMapKey = 1}]
--
-- Returns the empty list when passed an empty 'IdMap':
--
-- >>> keysIdMap emptyIdMap
-- []
keysIdMap :: IdMap a -> [IdMapKey]
keysIdMap :: forall a. IdMap a -> [IdMapKey]
keysIdMap IdMap {IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap} = (Int -> IdMapKey) -> [Int] -> [IdMapKey]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> IdMapKey
IdMapKey ([Int] -> [IdMapKey]) -> [Int] -> [IdMapKey]
forall a b. (a -> b) -> a -> b
$ IntMap a -> [Int]
forall a. IntMap a -> [Int]
IntMap.keys IntMap a
idMap

-- | Delete a key and its value from the map. When the key is not a member of
-- the map, the original map is returned.
--
-- >>> let (key, idmap) = singletonIdMap "hello"
-- >>> let (_, idmap') = insertIdMap "bye" idmap
-- >>> deleteIdMap key idmap'
-- IdMap {idMap = fromList [(1,"bye")], nextId = 2}
--
-- Deleting a key that does not exist just returns the old map:
--
-- >>> deleteIdMap key idmap'
-- IdMap {idMap = fromList [(1,"bye")], nextId = 2}
deleteIdMap :: IdMapKey -> IdMap a -> IdMap a
deleteIdMap :: forall a. IdMapKey -> IdMap a -> IdMap a
deleteIdMap (IdMapKey Int
k) IdMap {IntMap a
idMap :: forall a. IdMap a -> IntMap a
idMap :: IntMap a
idMap, Int
nextId :: forall a. IdMap a -> Int
nextId :: Int
nextId} =
  IdMap
    { idMap :: IntMap a
idMap = Int -> IntMap a -> IntMap a
forall a. Int -> IntMap a -> IntMap a
IntMap.delete Int
k IntMap a
idMap
    , Int
nextId :: Int
nextId :: Int
nextId
    }