{-# LANGUAGE GADTs, ExistentialQuantification, Rank2Types #-}
module RefMonad (
    TypRep(..),  -- TypRep a, GADT reifying type representation of a
    Typeable,    -- typeclass of types that have a TypRep
        typRep,  -- :: TypRep a
    Dynamic(..), -- Dynamic (TypRep a) a
    castDyn,     -- :: Typeable a => Dynamic -> Maybe a
    
    RefM,    -- RefM s a.  abstract, instance of Monad.
             -- s represents the "heap state"
    runRefM, -- :: (forall s. RefM s a) -> a
             -- "forall" forces us to not care about the initial state of the heap.

    Ref,     -- Ref s a.  abstract reference to objects of type a
    newRef,  -- :: Typeable a => a -> RefM s (Ref a)
    readRef, -- :: Ref s a -> RefM s a
    writeRef -- :: Ref s a -> a -> RefM s ()
)
where
import Data.Map (Map)
import qualified Data.Map as M
import Control.Monad (liftM)
import Control.Monad.State (State, evalState, put, get)
import Data.Maybe (fromMaybe)

---
--- Type representations & type-safe cast
--- using a GADT to represent types
--- (as has been done many times before)
---

data TypRep a where
    TInteger :: TypRep Integer
    TBool :: TypRep Bool
    TList :: TypRep a -> TypRep [a]
    TPair :: TypRep a -> TypRep b -> TypRep (a,b)
    -- etc.

data Teq a b where Refl :: Teq a a

typEq :: TypRep a -> TypRep b -> Maybe (Teq a b)
typEq TInteger TInteger = return Refl
typEq TBool TBool = return Refl
typEq (TList a) (TList b) = do
    Refl <- typEq a b
    return Refl
typEq (TPair a1 a2) (TPair b1 b2) = do
    Refl <- typEq a1 b1
    Refl <- typEq a2 b2
    return Refl
-- etc.
typEq _ _ = fail "types not equal"

castRep :: TypRep a -> a -> TypRep b -> Maybe b
castRep typa a typb = do
    Refl <- typEq typa typb
    return a

---
--- Existence of type representations allows us to use
--- existential quantification to encode dynamic
--- typing.
---
data Dynamic = forall a. Dynamic (TypRep a) a

class Typeable a where
    typRep :: TypRep a

instance Typeable Integer where typRep = TInteger
instance Typeable Bool where typRep = TBool
instance Typeable a => Typeable [a] where typRep = TList typRep
instance (Typeable a, Typeable b) => Typeable (a,b) where typRep = TPair typRep typRep
-- etc.

castDynRep :: Dynamic -> TypRep a -> Maybe a
castDynRep (Dynamic typx x) = castRep typx x

castDyn :: Typeable a => Dynamic -> Maybe a
castDyn dyn = castDynRep dyn typRep

---
--- Now that we have a dynamic type
--- we can use it to hold data in a heap addressed by
--- some arbitrary key type.
---
--- I choose integer because I know we won't run out!
---

type HeapItem = Dynamic
type Heap = Map Integer HeapItem

data Ref s a = Ref Integer (TypRep a)
newtype RefM s a = RefM { unRefM :: State (Integer, Heap) a }

instance Monad (RefM s) where
    return x = RefM $ return x
    m >>= f  = RefM $ unRefM m >>= unRefM . f
    fail s   = RefM $ fail s

-- note constraint on Typeable a here
newRef :: Typeable a => a -> RefM s (Ref s a)
newRef x = RefM $ do
    (top, heap) <- get
    let rep = typRep
    let ref = Ref top rep
    put (top + 1, M.insert top (Dynamic rep x) heap)
    return ref

readRef :: Ref s a -> RefM s a
readRef (Ref n typa) = RefM $ do
    heap <- liftM snd get
    let dyn = fromMaybe (error "can't happen -- elem not in map") $
                  M.lookup n heap
    let result = fromMaybe (error "can't happen -- type mismatch") $
                   castDynRep dyn typa
    return result

writeRef :: Ref s a -> a -> RefM s ()
writeRef (Ref n typa) a = RefM $ do
    (top, heap) <- get
    put (top, M.insert n (Dynamic typa a) heap)

runRefM :: (forall s. RefM s a) -> a
runRefM m = evalState (unRefM m) (0, M.empty)