--------------------------------------------------------------------------------
-- Most of this code is taken from 'http://github.com/ekmett/stable-maps'.
--------------------------------------------------------------------------------

{-# LANGUAGE Safe #-}

module System.Mem.StableName.Map
    ( Map(..)
    , empty
    , null
    , singleton
    , member
    , notMember
    , insert
    , insertWith
    , insertWith'
    , lookup
    , find
    , findWithDefault
    ) where

import qualified Prelude
import Prelude hiding (lookup, null)
import System.Mem.StableName.Dynamic
import qualified Data.IntMap as IntMap
import Data.IntMap (IntMap)

import Copilot.Core.Error (impossible)

data Map a = Map { Map a -> IntMap [(DynStableName, a)]
getMap  :: IntMap [(DynStableName, a)]
                 , Map a -> Int
getSize :: Int }

empty :: Map a
empty :: Map a
empty = IntMap [(DynStableName, a)] -> Int -> Map a
forall a. IntMap [(DynStableName, a)] -> Int -> Map a
Map IntMap [(DynStableName, a)]
forall a. IntMap a
IntMap.empty Int
0

null :: Map a -> Bool
null :: Map a -> Bool
null (Map IntMap [(DynStableName, a)]
m Int
_) = IntMap [(DynStableName, a)] -> Bool
forall a. IntMap a -> Bool
IntMap.null IntMap [(DynStableName, a)]
m

singleton :: DynStableName -> a -> Map a
singleton :: DynStableName -> a -> Map a
singleton DynStableName
k a
v =
  IntMap [(DynStableName, a)] -> Int -> Map a
forall a. IntMap [(DynStableName, a)] -> Int -> Map a
Map (Int -> [(DynStableName, a)] -> IntMap [(DynStableName, a)]
forall a. Int -> a -> IntMap a
IntMap.singleton (DynStableName -> Int
hashDynStableName DynStableName
k) [(DynStableName
k,a
v)]) Int
1

member :: DynStableName -> Map a -> Bool
member :: DynStableName -> Map a -> Bool
member DynStableName
k Map a
m = case DynStableName -> Map a -> Maybe a
forall v. DynStableName -> Map v -> Maybe v
lookup DynStableName
k Map a
m of
    Maybe a
Nothing -> Bool
False
    Just a
_ -> Bool
True

notMember :: DynStableName -> Map a -> Bool
notMember :: DynStableName -> Map a -> Bool
notMember DynStableName
k Map a
m = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ DynStableName -> Map a -> Bool
forall a. DynStableName -> Map a -> Bool
member DynStableName
k Map a
m

insert :: DynStableName -> a -> Map a -> Map a
insert :: DynStableName -> a -> Map a -> Map a
insert DynStableName
k a
v Map { getMap :: forall a. Map a -> IntMap [(DynStableName, a)]
getMap  = IntMap [(DynStableName, a)]
mp
               , getSize :: forall a. Map a -> Int
getSize = Int
sz }
  = IntMap [(DynStableName, a)] -> Int -> Map a
forall a. IntMap [(DynStableName, a)] -> Int -> Map a
Map (([(DynStableName, a)]
 -> [(DynStableName, a)] -> [(DynStableName, a)])
-> Int
-> [(DynStableName, a)]
-> IntMap [(DynStableName, a)]
-> IntMap [(DynStableName, a)]
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith [(DynStableName, a)]
-> [(DynStableName, a)] -> [(DynStableName, a)]
forall a. [a] -> [a] -> [a]
(++) (DynStableName -> Int
hashDynStableName DynStableName
k) [(DynStableName
k,a
v)] IntMap [(DynStableName, a)]
mp)
        (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | /O(log n)/. Insert with a function for combining the new value and old value.
-- @'insertWith' f key value mp@
-- will insert the pair (key, value) into @mp@ if the key does not exist
-- in the map. If the key does exist, the function will insert the pair
-- @(key, f new_value old_value)@
insertWith :: (a -> a -> a) -> DynStableName -> a -> Map a -> Map a
insertWith :: (a -> a -> a) -> DynStableName -> a -> Map a -> Map a
insertWith a -> a -> a
f DynStableName
k a
v Map { getMap :: forall a. Map a -> IntMap [(DynStableName, a)]
getMap  = IntMap [(DynStableName, a)]
mp
                     , getSize :: forall a. Map a -> Int
getSize = Int
sz }
  = IntMap [(DynStableName, a)] -> Int -> Map a
forall a. IntMap [(DynStableName, a)] -> Int -> Map a
Map (([(DynStableName, a)]
 -> [(DynStableName, a)] -> [(DynStableName, a)])
-> Int
-> [(DynStableName, a)]
-> IntMap [(DynStableName, a)]
-> IntMap [(DynStableName, a)]
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith [(DynStableName, a)]
-> [(DynStableName, a)] -> [(DynStableName, a)]
forall t. t -> [(DynStableName, a)] -> [(DynStableName, a)]
go (DynStableName -> Int
hashDynStableName DynStableName
k) [(DynStableName
k,a
v)] IntMap [(DynStableName, a)]
mp)
        (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    where
        go :: t -> [(DynStableName, a)] -> [(DynStableName, a)]
go t
_ ((DynStableName
k',a
v'):[(DynStableName, a)]
kvs)
            | DynStableName
k DynStableName -> DynStableName -> Bool
forall a. Eq a => a -> a -> Bool
== DynStableName
k' = (DynStableName
k', a -> a -> a
f a
v a
v') (DynStableName, a) -> [(DynStableName, a)] -> [(DynStableName, a)]
forall a. a -> [a] -> [a]
: [(DynStableName, a)]
kvs
            | Bool
otherwise = (DynStableName
k',a
v') (DynStableName, a) -> [(DynStableName, a)] -> [(DynStableName, a)]
forall a. a -> [a] -> [a]
: t -> [(DynStableName, a)] -> [(DynStableName, a)]
go t
forall a. HasCallStack => a
undefined [(DynStableName, a)]
kvs
        go t
_ [] = []

-- | Same as 'insertWith', but with the combining function applied strictly.
insertWith' :: (a -> a -> a) -> DynStableName -> a -> Map a -> Map a
insertWith' :: (a -> a -> a) -> DynStableName -> a -> Map a -> Map a
insertWith' a -> a -> a
f DynStableName
k a
v Map { getMap :: forall a. Map a -> IntMap [(DynStableName, a)]
getMap  = IntMap [(DynStableName, a)]
mp
                      , getSize :: forall a. Map a -> Int
getSize = Int
sz }
  = IntMap [(DynStableName, a)] -> Int -> Map a
forall a. IntMap [(DynStableName, a)] -> Int -> Map a
Map (([(DynStableName, a)]
 -> [(DynStableName, a)] -> [(DynStableName, a)])
-> Int
-> [(DynStableName, a)]
-> IntMap [(DynStableName, a)]
-> IntMap [(DynStableName, a)]
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith [(DynStableName, a)]
-> [(DynStableName, a)] -> [(DynStableName, a)]
forall t. t -> [(DynStableName, a)] -> [(DynStableName, a)]
go (DynStableName -> Int
hashDynStableName DynStableName
k) [(DynStableName
k,a
v)] IntMap [(DynStableName, a)]
mp)
        (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    where
        go :: t -> [(DynStableName, a)] -> [(DynStableName, a)]
go t
_ ((DynStableName
k',a
v'):[(DynStableName, a)]
kvs)
            | DynStableName
k DynStableName -> DynStableName -> Bool
forall a. Eq a => a -> a -> Bool
== DynStableName
k' = let v'' :: a
v'' = a -> a -> a
f a
v a
v' in a
v'' a -> [(DynStableName, a)] -> [(DynStableName, a)]
`seq` (DynStableName
k', a
v'') (DynStableName, a) -> [(DynStableName, a)] -> [(DynStableName, a)]
forall a. a -> [a] -> [a]
: [(DynStableName, a)]
kvs
            | Bool
otherwise = (DynStableName
k', a
v') (DynStableName, a) -> [(DynStableName, a)] -> [(DynStableName, a)]
forall a. a -> [a] -> [a]
: t -> [(DynStableName, a)] -> [(DynStableName, a)]
go t
forall a. HasCallStack => a
undefined [(DynStableName, a)]
kvs
        go t
_ [] = []

-- | /O(log n)/. Lookup the value at a key in the map.
--
-- The function will return the corresponding value as a @('Just' value)@
-- or 'Nothing' if the key isn't in the map.
lookup :: DynStableName -> Map v -> Maybe v
lookup :: DynStableName -> Map v -> Maybe v
lookup DynStableName
k (Map IntMap [(DynStableName, v)]
m Int
_) = do
    [(DynStableName, v)]
pairs <- Int -> IntMap [(DynStableName, v)] -> Maybe [(DynStableName, v)]
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (DynStableName -> Int
hashDynStableName DynStableName
k) IntMap [(DynStableName, v)]
m
    DynStableName -> [(DynStableName, v)] -> Maybe v
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup DynStableName
k [(DynStableName, v)]
pairs

find :: DynStableName -> Map v -> v
find :: DynStableName -> Map v -> v
find DynStableName
k Map v
m = case DynStableName -> Map v -> Maybe v
forall v. DynStableName -> Map v -> Maybe v
lookup DynStableName
k Map v
m of
    Maybe v
Nothing -> String -> String -> v
forall a. String -> String -> a
impossible String
"find" String
"copilot-language"
    Just v
x -> v
x

-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
-- the value at key @k@ or returns the default value @def@
-- when the key is not in the map.
findWithDefault :: v -> DynStableName -> Map v -> v
findWithDefault :: v -> DynStableName -> Map v -> v
findWithDefault v
dflt DynStableName
k Map v
m = v -> (v -> v) -> Maybe v -> v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe v
dflt v -> v
forall a. a -> a
id (Maybe v -> v) -> Maybe v -> v
forall a b. (a -> b) -> a -> b
$ DynStableName -> Map v -> Maybe v
forall v. DynStableName -> Map v -> Maybe v
lookup DynStableName
k Map v
m