content: -- | 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 { unIdMapKey :: Int }
    deriving stock (Eq, Show)
  
  data IdMap a = IdMap
    { idMap :: !(IntMap a)
    , nextId :: !Int
    }
    deriving stock 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 idMapA _) == (IdMap idMapB _) = idMapA == idMapB
  
  instance Functor IdMap where
    fmap f IdMap{idMap, nextId} = IdMap { idMap = fmap f idMap, nextId }
  
  instance Foldable IdMap where
    foldMap f m = Foldable.foldMap f $ idMap m
  
  instance FoldableWithIndex Int IdMap where
    ifoldMap f m = ifoldMap f $ idMap m
  
  instance Traversable IdMap where
    traverse f IdMap{idMap, nextId} =
      fmap (\m -> IdMap { idMap = m, nextId }) (traverse f 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 i) f IdMap{idMap, nextId} =
      case IntMap.lookup i idMap of
        Just v -> fmap update (f v) -- f v <&> \v' -> IntMap.insert k v' m
        Nothing -> pure IdMap{idMap, nextId}
      where
        update :: a -> IdMap a
        update v' =
          IdMap
            { idMap = IntMap.adjust (const v') i idMap
            , nextId
            }
  
  initialId :: Int
  initialId = 0
  
  -- | Get the next available ID.
  --
  -- >>> succId 3
  -- 4
  succId :: Int -> Int
  succId i = i + 1
  
  -- | An initial 'IdMap' with no values.
  --
  -- >>> emptyIdMap
  -- IdMap {idMap = fromList [], nextId = 0}
  emptyIdMap :: IdMap a
  emptyIdMap = IdMap { idMap = mempty, nextId = 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 a IdMap {idMap, nextId} =
    let newMap = IntMap.insert nextId a idMap
        newNextId = nextId + 1
    in (IdMapKey nextId, IdMap { idMap = newMap, nextId = 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 a = insertIdMap 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 (IdMapKey k) IdMap {idMap} = IntMap.lookup k 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 IdMap {idMap} = fmap IdMapKey $ IntMap.keys 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 (IdMapKey k) IdMap {idMap, nextId} =
    IdMap
      { idMap = IntMap.delete k idMap
      , nextId
      }
  