{-# LANGUAGE  GeneralizedNewtypeDeriving 
  , NoMonomorphismRestriction 
  , 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
    , 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)

-- | @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)
                            
-- | @Atom@ is a specialized state monad for converting to and from
-- @Int@s.
newtype Atom a r = Atom (State (AtomTable a) r)
   deriving (Functor, Monad)                   


-- | @toAtom x@ converts @x@ to a unique @Int@ in the @Atom@ 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 x@ converts @x@ to a unique @Int@ in the @Atom@
-- monad only if @x@ already has a corresponding @Int@
maybeToAtom :: Ord a => a -> Atom a (Maybe Int)
maybeToAtom x = Atom $ do
  t <- get
  return $! Map.lookup x . to $ t
 
-- | @fromAtom i@ converts the @Int@ @i@ to its corresponding object
-- in the Atom monad.
fromAtom :: Int -> Atom a a            
fromAtom i = Atom $ do
  t <- get
  return $! (from t) IntMap.! i

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

-- | @runAtom c s@ runs computation c in the Atom monad with the
-- initial @AtomTable@ s.
runAtom :: (Ord a) => Atom a t -> AtomTable a -> (t, AtomTable a)
runAtom (Atom x) s = runState x s

-- | @evalAtom c@ runs computation c in the Atom monad with the empty
-- initial @AtomTable@.
evalAtom :: (Ord a) => Atom a t -> t
evalAtom = fst . flip runAtom empty


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