module Control.Monad.Atom
( AtomTable
, Atom
, toAtom
, fromAtom
, maybeToAtom
, empty
, evalAtom
, runAtom
, mapping
)
where
import Control.Monad.State
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import GHC.Generics (Generic)
data AtomTable a = T { lastID :: !Int
, to :: !(Map.Map a Int)
, from :: !(IntMap.IntMap a) }
deriving (Generic)
newtype Atom a r = Atom (State (AtomTable a) r)
deriving (Functor, Monad)
toAtom :: Ord a => a -> Atom a Int
toAtom x = Atom $ do
let b = x
t <- get
case Map.lookup b (to t) of
Just j -> return $! j
Nothing -> do
let i = lastID t
i' = i + 1
!t' = t { lastID = i'
, to = Map.insert b i (to t)
, from = IntMap.insert i b (from t) }
put t'
return $! lastID t
maybeToAtom :: Ord a => a -> Atom a (Maybe Int)
maybeToAtom x = Atom $ do
t <- get
return $! Map.lookup x . to $ t
fromAtom :: Int -> Atom a a
fromAtom i = Atom $ do
t <- get
return $! (from t) IntMap.! i
empty :: (Ord a) => AtomTable a
empty = T 0 Map.empty IntMap.empty
runAtom :: (Ord a) => Atom a t -> AtomTable a -> (t, AtomTable a)
runAtom (Atom x) s = runState x s
evalAtom :: (Ord a) => Atom a t -> t
evalAtom = fst . flip runAtom empty
mapping :: (Ord a) => AtomTable a -> Map.Map a Int
mapping = to