{-# LANGUAGE  GeneralizedNewtypeDeriving
  , NoMonomorphismRestriction
  , MultiParamTypeClasses
  , TypeFamilies
  , FlexibleInstances
  , FlexibleContexts
  , DeriveGeneric
  , BangPatterns #-}
-- | The Atom monad provides functions which convert objects to unique
-- atoms (represented as Ints). Example:
--
-- > example = evalAtom $ do
-- >   xs <- mapM toAtom "abcabd"
-- >   zs <- mapM fromAtom xs
-- >   return $ zip zs xs
--
-- >>> example
-- >>> [('a',0),('b',1),('c',2),('a',0),('b',1),('d',3)]

module Control.Monad.Atom
    ( AtomTable
    , Atom
    , AtomT
    , toAtom
    , fromAtom
    , maybeToAtom
    , empty
    , evalAtom
    , evalAtomT
    , runAtom
    , runAtomT
    , mapping
    )
where
import Control.Monad.State
import Control.Monad.Identity
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import GHC.Generics (Generic)

-- | @AtomTable@ holds the state necessary for converting to and from
-- @Int@s.
data AtomTable a = T { lastID :: {-# UNPACK #-} !Int
                     , to     :: !(Map.Map a Int)
                     , from   :: !(IntMap.IntMap a) }
                   deriving (Generic)


-- | @AtomT@ is a specialized state monad transformer for converting
-- to and from @Int@s.
newtype AtomT a m r = AtomT (StateT (AtomTable a) m r)
   deriving (Functor, Monad, MonadTrans, MonadIO)
-- | @Atom@ is a specialized state monad for converting to and from
-- @Int@s.
newtype Atom a r = Atom (AtomT a Identity r)
   deriving (Functor, Monad)

class (Monad m) => MonadAtom m where
  type Key m
  -- | @toAtom x@ converts @x@ to a unique @Int@ in the @Atom@ monad
  toAtom :: Key m -> m Int
  -- | @maybeToAtom x@ converts @x@ to a unique @Int@ in the @Atom@
  -- monad only if @x@ already has a corresponding @Int@
  maybeToAtom :: Key m -> m (Maybe Int)
  -- | @fromAtom i@ converts the @Int@ @i@ to its corresponding object
  -- in the Atom monad.
  fromAtom :: Int -> m (Key m)
instance (Ord a, Monad m) => MonadAtom (AtomT a m) where
  type Key (AtomT a m) = a
  toAtom = AtomT . toAtom'
  {-# INLINE toAtom #-}
  maybeToAtom = AtomT . maybeToAtom'
  {-# INLINE maybeToAtom #-}
  fromAtom = AtomT . fromAtom'
  {-# INLINE fromAtom #-}
instance (Ord a) => MonadAtom (Atom a) where
  type Key (Atom a) = a
  toAtom = Atom . toAtom
  {-# INLINE toAtom #-}
  maybeToAtom = Atom . maybeToAtom
  {-# INLINE maybeToAtom #-}
  fromAtom = Atom . fromAtom
  {-# INLINE fromAtom #-}

-- | @runAtomT c s@ runs computation c in the AtomT monad transformer
-- with the initial @AtomTable@ s.
runAtomT :: (Ord a, Monad m) => AtomT a m r -> AtomTable a -> m (r, AtomTable a)
runAtomT (AtomT x) = runStateT x 
{-# INLINE runAtomT #-}

-- | @runAtom c s@ runs computation c in the Atom monad with the
-- initial @AtomTable@ s.
runAtom :: (Ord a) => Atom a r -> AtomTable a -> (r, AtomTable a)
runAtom (Atom x) s = runIdentity (runAtomT x s)
{-# INLINE runAtom #-}

-- | @evalAtomT c@ runs computation c in the AtomT monad transformer
-- with the empty initial @AtomTable@.
evalAtomT :: (Ord a, Monad m) => AtomT a m r -> m r
evalAtomT (AtomT x) = evalStateT x empty
{-# INLINE evalAtomT #-}

-- | @evalAtom c@ runs computation c in the Atom monad with the empty
-- initial @AtomTable@.
evalAtom :: (Ord a) => Atom a r -> r
evalAtom (Atom x) = runIdentity (evalAtomT x)
{-# INLINE evalAtom #-}

toAtom' :: (Monad m, Ord a) => a -> StateT (AtomTable a) m Int
toAtom' x = 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
{-# SPECIALIZE toAtom' :: (Ord a) => a -> StateT (AtomTable a) Identity Int #-}

maybeToAtom' :: (Ord a, Monad m) =>
                a -> StateT (AtomTable a) m (Maybe Int)
maybeToAtom' x = do
  t <- get
  return $! Map.lookup x . to $ t
{-# SPECIALIZE  maybeToAtom' :: (Ord a) =>
                a -> StateT (AtomTable a) Identity (Maybe Int) #-}

fromAtom' ::  Monad m => Int -> StateT (AtomTable a) m a
fromAtom' i = do
  t <- get
  return $! from t IntMap.! i
{-# SPECIALIZE fromAtom' :: Int -> StateT (AtomTable a) Identity a #-}

-- | @empty@ is the initial empty @AtomTable@
empty :: (Ord a) => AtomTable a
empty = T 0 Map.empty IntMap.empty

-- | The mapping stored in the atom table
mapping :: (Ord a) => AtomTable a -> Map.Map a Int
mapping = to
{-# INLINE mapping #-}