-- | A monad for binding values to tags to ensure sharing, 
-- with the added twist that the value can be polymorphic
-- and each monomorphic instance is bound separately.
module Control.Monad.TagShare(
  -- ** Dynamic map
  DynMap,
  dynEmpty,
  dynInsert,
  dynLookup,
  -- ** Sharing monad
  Sharing,
  runSharing,
  share
  ) where
import Control.Monad.State
import Data.Typeable
import Data.Dynamic(Dynamic, fromDynamic, toDyn)
import Data.Map as M

-- |  A dynamic map with type safe
-- insertion and lookup.
newtype DynMap tag = 
  DynMap (M.Map (tag, TypeRep) Dynamic) 
  deriving Show

dynEmpty :: DynMap tag
dynEmpty = DynMap M.empty  
  
dynInsert :: (Typeable a, Ord tag) => 
                tag -> a -> DynMap tag -> DynMap tag
dynInsert u a (DynMap m) = 
          DynMap (M.insert (u,typeOf a) (toDyn a) m)

dynLookup :: (Typeable a, Ord tag) => 
                tag -> DynMap tag -> Maybe a
dynLookup u (DynMap m) = hlp fun undefined where 
    hlp :: Typeable a => 
      (TypeRep -> Maybe a) -> a -> Maybe a
    hlp f a = f (typeOf a)
    fun tr = M.lookup (u,tr) m >>= fromDynamic

 
-- | A sharing monad
-- with a function that binds a tag to a value.
type Sharing tag a = State (DynMap tag) a

runSharing :: Sharing tag a -> a
runSharing m = evalState m dynEmpty

share :: (Typeable a, Ord tag) => 
  tag -> Sharing tag a -> Sharing tag a
share t m = do
    mx <- gets $ (dynLookup t)
    case mx of
      Just e      ->  return e
      Nothing     ->  mfix $ \e -> do
        modify (dynInsert t e)
        m