content: {-
  Copyright   : (C) 2016-2021 QBayLogic B.V.
                         2022 Alexander McKenna
  License     : BSD2 (see the file LICENSE)
  Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>
  -}
  
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Clash.Data.UniqMap
    ( UniqMap(..)
    , empty
    , singleton
    , singletonUnique
    , null
    , insert
    , insertUnique
    , insertWith
    , insertMany
    , lookup
    , find
    , elem
    , notElem
    , filter
    , mapMaybe
    , foldrWithUnique
    , foldlWithUnique'
    , delete
    , deleteMany
    , unionWith
    , difference
    , disjoint
    , submap
    , fromList
    , toList
    , keys
    , elems
    ) where
  
  #if MIN_VERSION_ghc(9,8,4) || (MIN_VERSION_ghc(9,6,7) && !MIN_VERSION_ghc(9,8,0))
  #define UNIQUE_IS_WORD64
  #endif
  
  import           Prelude hiding (elem, filter, lookup, notElem, null)
  
  import           Control.DeepSeq (NFData)
  import           Data.Binary (Binary (..))
  import           Data.Bifunctor (first)
  import           Data.Function (on)
  #ifdef UNIQUE_IS_WORD64
  import           GHC.Data.Word64Map.Strict (Word64Map)
  import qualified GHC.Data.Word64Map.Strict as IntMap
  #else
  import           Data.IntMap.Strict (IntMap)
  import qualified Data.IntMap.Strict as IntMap
  #endif
  import qualified Data.List as List (foldl')
  
  #if !MIN_VERSION_containers(0,6,2)
  import qualified Data.IntMap.Extra as IntMap
  #endif
  
  #if MIN_VERSION_prettyprinter(1,7,0)
  import           Prettyprinter
  #else
  import           Data.Text.Prettyprint.Doc
  #endif
  
  import           Clash.Pretty
  import           Clash.Unique (Unique, Uniquable(getUnique))
  
  -- | A map indexed by a 'Unique'. Typically the elements of this map are also
  -- uniqueable and provide their own key, however a unique can be associated
  -- with any value.
  newtype UniqMap a
  #ifdef UNIQUE_IS_WORD64
    = UniqMap { uniqMapToIntMap :: Word64Map a }
  #else
    = UniqMap { uniqMapToIntMap :: IntMap a }
  #endif
    deriving stock Traversable
    deriving newtype
      ( Foldable
      , Functor
      , Monoid
      , NFData
      , Semigroup
      , Show
      )
  #ifdef UNIQUE_IS_WORD64
  instance Binary a => Binary (UniqMap a) where
    put (UniqMap m) = put (IntMap.size m) <> mapM_ put (IntMap.toAscList m)
    get             = fmap (UniqMap . IntMap.fromDistinctAscList) get
  #else
    deriving newtype Binary
  #endif
  
  instance ClashPretty a => ClashPretty (UniqMap a) where
    clashPretty xs =
      brackets $ fillSep $ punctuate comma $
        [ fromPretty k <+> ":->" <+> clashPretty v
        | (k, v) <- toList xs
        ]
  
  -- | An empty map.
  empty :: UniqMap a
  empty =
    UniqMap IntMap.empty
  
  {-# SPECIALIZE singleton :: Unique -> b -> UniqMap b #-}
  -- | A map containing a single value indexed by the given key's unique.
  singleton :: Uniquable a => a -> b -> UniqMap b
  singleton k v =
    UniqMap (IntMap.singleton (getUnique k) v)
  
  {-# SPECIALIZE singletonUnique :: Unique -> UniqMap Unique #-}
  -- | A map containing a single value indexed by the value's unique.
  singletonUnique :: Uniquable a => a -> UniqMap a
  singletonUnique v =
    singleton (getUnique v) v
  
  -- | Check if the map is empty.
  null :: UniqMap a -> Bool
  null =
    IntMap.null . uniqMapToIntMap
  
  {-# SPECIALIZE insert :: Unique -> b -> UniqMap b -> UniqMap b #-}
  -- | Insert a new key-value pair into the map.
  insert :: Uniquable a => a -> b -> UniqMap b -> UniqMap b
  insert k v =
    UniqMap . IntMap.insert (getUnique k) v . uniqMapToIntMap
  
  {-# SPECIALIZE insertUnique :: Unique -> UniqMap Unique -> UniqMap Unique #-}
  -- | Insert a new value into the map, using the unique of the value as the key.
  insertUnique :: Uniquable a => a -> UniqMap a -> UniqMap a
  insertUnique v =
    insert (getUnique v) v
  
  -- | Insert a new key-value pair into the map, using the given combining
  -- function if there is already an entry with the same unique in the map.
  insertWith :: Uniquable a => (b -> b -> b) -> a -> b -> UniqMap b -> UniqMap b
  insertWith f k v =
    UniqMap . IntMap.insertWith f (getUnique k) v . uniqMapToIntMap
  
  -- | Insert a list of key-value pairs into the map.
  insertMany :: Uniquable a => [(a, b)] -> UniqMap b -> UniqMap b
  insertMany kvs xs =
    List.foldl' (\acc (k, v) -> insert k v acc) xs kvs
  
  {-# SPECIALIZE lookup :: Unique -> UniqMap b -> Maybe b #-}
  -- | Lookup an item in the map, using the unique of the given key.
  lookup :: Uniquable a => a -> UniqMap b -> Maybe b
  lookup k =
    IntMap.lookup (getUnique k) . uniqMapToIntMap
  
  {-# SPECIALIZE find :: Unique -> UniqMap b -> b #-}
  -- | Lookup and item in the map, using the unique of the given key. If the item
  -- is not found in the map an error is raised.
  find :: Uniquable a => a -> UniqMap b -> b
  find k =
    let notFound =
          error ("find: Key " <> show (getUnique k) <> " is not in the UniqMap")
     in IntMap.findWithDefault notFound (getUnique k) . uniqMapToIntMap
  
  {-# SPECIALIZE elem :: Unique -> UniqMap b -> Bool #-}
  -- | Check if there is an entry in the map for the unique of the given value.
  elem :: Uniquable a => a -> UniqMap b -> Bool
  elem k =
    IntMap.member (getUnique k) . uniqMapToIntMap
  
  {-# SPECIALIZE notElem :: Unique -> UniqMap b -> Bool #-}
  -- | Check if there is not an entry in the map for the unique of the given
  -- value.
  notElem :: Uniquable a => a -> UniqMap b -> Bool
  notElem k =
    IntMap.notMember (getUnique k) . uniqMapToIntMap
  
  -- | Filter all elements in the map according to some predicate.
  filter :: (b -> Bool) -> UniqMap b -> UniqMap b
  filter p =
    UniqMap . IntMap.filter p . uniqMapToIntMap
  
  -- | Apply a function to all elements in the map, keeping those where the
  -- result is not @Nothing@.
  mapMaybe :: (a -> Maybe b) -> UniqMap a -> UniqMap b
  mapMaybe f =
    UniqMap . IntMap.mapMaybe f . uniqMapToIntMap
  
  -- | Lazily right-fold over the map using the given function.
  foldrWithUnique :: (Unique -> a -> b -> b) -> b -> UniqMap a -> b
  foldrWithUnique f x =
    IntMap.foldrWithKey f x . uniqMapToIntMap
  
  -- | Strictly left-fold over the map using the given function.
  foldlWithUnique' :: (b -> Unique -> a -> b) -> b -> UniqMap a -> b
  foldlWithUnique' f x =
    IntMap.foldlWithKey' f x . uniqMapToIntMap
  
  {-# SPECIALIZE delete :: Unique -> UniqMap b -> UniqMap b #-}
  -- | Delete the entry in the map indexed by the unique of the given value.
  delete :: Uniquable a => a -> UniqMap b -> UniqMap b
  delete k =
    UniqMap . IntMap.delete (getUnique k) . uniqMapToIntMap
  
  -- | Delete all entries in the map indexed by the uniques of the given values.
  deleteMany :: Uniquable a => [a] -> UniqMap b -> UniqMap b
  deleteMany ks xs =
    List.foldl' (\acc k -> delete k acc) xs ks
  
  -- | Merge two unique maps, using the given combining funcion if a value with
  -- the same unique key exists in both maps.
  unionWith :: (b -> b -> b) -> UniqMap b -> UniqMap b -> UniqMap b
  unionWith f xs ys =
    UniqMap ((IntMap.unionWith f `on` uniqMapToIntMap) xs ys)
  
  -- | Filter the first map to only contain keys which are not in the second map.
  difference :: UniqMap b -> UniqMap b -> UniqMap b
  difference xs ys =
    UniqMap ((IntMap.difference `on` uniqMapToIntMap) xs ys)
  
  -- | Check if there are no common keys between two maps.
  disjoint :: UniqMap b -> UniqMap b -> Bool
  disjoint =
    IntMap.disjoint `on` uniqMapToIntMap
  
  -- | Check if one map is a submap of another.
  submap :: UniqMap b -> UniqMap b -> Bool
  submap =
    -- We only check that the keys of the map make it a submap, the elements do
    -- not need to be equal. Maybe this should be changed?
    IntMap.isSubmapOfBy (\_ _ -> True) `on` uniqMapToIntMap
  
  {-# SPECIALIZE fromList :: [(Unique, b)] -> UniqMap b #-}
  -- | Convert a list of key-value pairs to a map.
  fromList :: Uniquable a => [(a, b)] -> UniqMap b
  fromList =
    UniqMap . IntMap.fromList . fmap (first getUnique)
  
  -- | Convert a map to a list of unique-value pairs.
  toList :: UniqMap b -> [(Unique, b)]
  toList =
    IntMap.toList . uniqMapToIntMap
  
  -- | Get the unique keys of a map.
  keys :: UniqMap b -> [Unique]
  keys =
    IntMap.keys . uniqMapToIntMap
  
  -- | Get the values of a map.
  elems :: UniqMap b -> [b]
  elems =
    IntMap.elems . uniqMapToIntMap
  