96 lines
2.6 KiB
Haskell
96 lines
2.6 KiB
Haskell
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
|
|
{-# LANGUAGE DeriveTraversable #-}
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
module Data.Id (
|
|
Id(..)
|
|
, IdMap
|
|
, Indexed
|
|
, at
|
|
, delete
|
|
, empty
|
|
, filterWithKey
|
|
, fromList
|
|
, insert
|
|
, keysSet
|
|
, lookup
|
|
, mapWithKey
|
|
, member
|
|
, minViewWithKey
|
|
, register
|
|
, singleton
|
|
, size
|
|
, toList
|
|
, union
|
|
) where
|
|
|
|
import Control.Monad.State.Strict (MonadState, modify, gets)
|
|
import Data.IntMap (IntMap, (!))
|
|
import qualified Data.IntMap as IntMap (
|
|
delete, empty, filterWithKey, fromList, keysSet, insert, lookup, mapWithKey
|
|
, maxViewWithKey, member, minViewWithKey, size, toList, union
|
|
)
|
|
import Data.IntSet (IntSet)
|
|
import Prelude hiding (lookup)
|
|
|
|
newtype Id a = Id {
|
|
getId :: Int
|
|
} deriving (Eq, Enum, Ord, Show)
|
|
newtype IdMap a b = IdMap {
|
|
intMap :: IntMap b
|
|
} deriving (Show, Functor, Semigroup, Monoid, Foldable, Traversable)
|
|
type Indexed a = IdMap a a
|
|
|
|
at :: IdMap a b -> Id a -> b
|
|
at (IdMap idMap) = (idMap !) . getId
|
|
|
|
lookup :: Id a -> IdMap a b -> Maybe b
|
|
lookup (Id a) (IdMap idMap) = IntMap.lookup a idMap
|
|
|
|
size :: IdMap a b -> Int
|
|
size = IntMap.size . intMap
|
|
|
|
member :: Id a -> IdMap a b -> Bool
|
|
member (Id a) (IdMap idMap) = IntMap.member a idMap
|
|
|
|
empty :: IdMap a b
|
|
empty = IdMap {intMap = IntMap.empty}
|
|
|
|
singleton :: Id a -> b -> IdMap a b
|
|
singleton a b = fromList [(a, b)]
|
|
|
|
insert :: Id a -> b -> IdMap a b -> IdMap a b
|
|
insert (Id a) b (IdMap idMap) = IdMap {intMap = IntMap.insert a b idMap}
|
|
|
|
delete :: Id a -> IdMap a b -> IdMap a b
|
|
delete (Id a) (IdMap idMap) = IdMap {intMap = IntMap.delete a idMap}
|
|
|
|
minViewWithKey :: IdMap a b -> Maybe ((Id a, b), IdMap a b)
|
|
minViewWithKey = fmap wrap . IntMap.minViewWithKey . intMap
|
|
where
|
|
wrap ((key, b), idMap) = ((Id key, b), IdMap idMap)
|
|
|
|
union :: IdMap a b -> IdMap a b -> IdMap a b
|
|
union (IdMap intMap1) (IdMap intMap2) =
|
|
IdMap {intMap = IntMap.union intMap1 intMap2}
|
|
|
|
mapWithKey :: (Id a -> b -> c) -> IdMap a b -> IdMap a c
|
|
mapWithKey f (IdMap idMap) = IdMap {intMap = IntMap.mapWithKey (f . Id) idMap}
|
|
|
|
filterWithKey :: (Id a -> b -> Bool) -> IdMap a b -> IdMap a b
|
|
filterWithKey f = IdMap . IntMap.filterWithKey (f . Id) . intMap
|
|
|
|
fromList :: [(Id a, b)] -> IdMap a b
|
|
fromList = IdMap . IntMap.fromList . fmap (\(key, b) -> (getId key, b))
|
|
|
|
toList :: IdMap a b -> [(Id a, b)]
|
|
toList = fmap (\(key, b) -> (Id key, b)) . IntMap.toList . intMap
|
|
|
|
keysSet :: IdMap a b -> IntSet
|
|
keysSet = IntMap.keysSet . intMap
|
|
|
|
register :: MonadState (IdMap a b) m => b -> m (Id a)
|
|
register b = do
|
|
newId <- gets (Id . maybe 0 ((+1) . fst . fst) . IntMap.maxViewWithKey . intMap)
|
|
modify (insert newId b)
|
|
return newId
|