{- A call-by-need interpreter using a heap to express sharing and support recursive bindings,
   and an environment to avoid substitution.  Extended to support free variables, narrowing,
   non-determinism, and concurrency.  -}

import Map
import Maybe
import List
import Monad
import qualified Forest
import Trace

type Var = String
type Constr = String
type Prim = (String,Concurrent)
type Flex = Bool
type Concurrent = Bool
type Pattern = (Constr,[Var])

data Exp = Var Var
	 | Int Int
	 | Abs Var Exp
         | App Exp Exp
         | Capp Constr [Exp]
	 | Primapp Prim Exp Exp
         | Case Flex Exp [(Pattern,Exp)]
	 | Letrec [(Var,Exp)] Exp
	 | Logic Var Exp
  deriving (Show,Eq)

-- Heap Pointers
type HPtr = Int

-- Values
data Value = 
   VCapp Constr [HPtr]
 | VInt Int
 | VAbs Env Var Exp
 | VLogic HPtr
 deriving (Show,Eq)

data HEntry = 
   HValue Value
 | HThunk Env Exp
 | HBlackhole
 deriving Show

type Env = Map Var HPtr


-- Heap: Free location supply plus bindings
data Heap = Heap {free :: [HPtr], bindings :: Map HPtr HEntry}
instance Show Heap where
  show (Heap free bindings) = show bindings
hfresh :: Heap -> (HPtr,Heap)
hfresh (Heap (v:vs) bindings) = (v,Heap vs bindings)
hempty :: Heap 
hempty = Heap {free = [0..], bindings = mempty}
hget :: Heap -> HPtr -> HEntry
hget (Heap free bindings) v = mget bindings v
hset :: Heap  -> (HPtr,HEntry) -> Heap 
hset (Heap free bindings) (v,e) = Heap free (mset bindings (v,e))


{- Model answers as lists; inherits standard monad defns for lists.
type A a = [a]  
alift = id
sols = id
-}

{- Model answers as forests; gives choice of dfs, bfs, etc. -}
type A a = Forest.Forest a
alift = Forest.lift
dfsols :: A a -> [a]
dfsols = Forest.dfs
bfsols :: A a -> [a]
bfsols = Forest.bfs
sols = bfsols  -- can change to whatever you want

-- Monad carries heap, and returns an answer structure
data State a = Done a 
             | Running (M a)
newtype M a = M (Heap -> A (State a,Heap))
instance Monad M where
  (M m1) >>= k = M (\ h -> 
      do (s,h') <- m1 h
         case s of
           Done a' -> let M m2 = k a' in m2 h'
           Running m' -> return (Running (m' >>= k),h'))
  return x = M (\ h -> return (Done x, h))
instance MonadPlus M where
  mzero = M (\ h -> mzero)
  (M m1) `mplus` (M m2) = M (\ h -> (m1 h `mplus` m2 h))
fresh :: M HPtr
fresh = M (\ h -> let (v,h') = hfresh h in return (Done v,h'))
store :: HPtr -> HEntry -> M ()
store p e = M (\ h -> return (Done (),hset h (p,e)))
fetch :: HPtr -> M HEntry
fetch p = M (\ h -> return (Done (hget h p), h))
yield :: M a -> M a
yield m = M (\ h -> return (Running m,h))
run :: M a -> A (a,Heap)
run = run' hempty
 where 
   run' :: Heap -> M a -> A (a, Heap)
   run' h (M m) =
     do (s,h') <- m h
        case s of
          Done a -> return (a, h')
          Running m' -> run' h' m' 
conc :: M a -> M b -> M (a,b) 
conc m1 m2 = (m1 `thn` m2) `mplus` (liftM swap (m2 `thn` m1))
    where thn :: M a -> M b -> M (a,b)
          (M m1') `thn` m2 =  M (\ h -> 
             do (s,h') <- m1' h
                case s of
                  Done a -> return (Running  (do b <- m2 
						 return (a,b)),h')
                  Running m' -> let M m'' = conc m' m2 in m'' h')
	  swap :: (a,b) -> (b,a)
          swap (a,b) = (b,a)
lift :: M a -> M a
lift (M m) = M (\ h -> alift (m h))

eval :: Env -> Exp -> M Value
eval env (Var x) =
     lift $
     do let p = mget env x
        h <- fetch p
        case h of
	     HThunk env' e' -> 
	       do store p HBlackhole
		  v' <- eval env' e'
                  store p (HValue v')
                  return v'
             HValue v -> return v
	     HBlackhole -> mzero
eval env (Int i) = return (VInt i)
eval env (Abs x b) = return (VAbs env x b)
eval env (App e0 e1) = 
     do VAbs env' x b <- eval env e0
        p1 <- fresh
        store p1 (HThunk env e1)
	let env'' = mset env' (x,p1)
	eval env'' b
eval env (Capp c es) = 
     do ps <- mapM (const fresh) es
	zipWithM_ store ps (map (HThunk env) es)
	return (VCapp c ps)
eval env (Primapp (p,concurrent) e1 e2) = 
     do (v1,v2) <-
	  if concurrent then
            conc (eval env e1) (eval env e2)
          else
            do v1 <- eval env e1
               v2 <- eval env e2
               return (v1,v2)
        checkGround v1
	checkGround v2
	return (doPrimapp p v1 v2)
eval env (Case flex e pes) = 
     do v <- eval env e
        case v of
	  VCapp c0 ps ->
    	    do let plookup [] = mzero
                   plookup (((c,xs),b):pes) | c == c0   = return (xs,b)
	                                    | otherwise = plookup pes
	       (xs,b) <- plookup pes
	       let env' = foldl mset env (zip xs ps)
	       eval env' b
          VLogic p0 | flex ->
            msum (map f pes)
            where
              f ((c,xs),e') =
                do ps <- mapM (const allocLogic) xs
                   let env' = foldl mset env (zip xs ps) 
	           store p0 (HValue (VCapp c ps))                       
		   yield $
		     eval env' e'
	  _ -> mzero
eval env (Letrec xes e) = 
     do let (xs,es) = unzip xes
	ps <- mapM (const fresh) xes
	let env' = foldl mset env (zip xs ps)
	zipWithM_ store ps (map (HThunk env') es)
        eval env' e
eval env (Logic x e) = 
     do p <- allocLogic 
        let env' = mset env (x,p)
        eval env' e

allocLogic :: M HPtr
allocLogic = 
     do p <- fresh
        store p (HValue (VLogic p))
        return p

checkGround :: Value -> M ()
checkGround (VLogic _) = mzero
checkGround _ = return ()

doPrimapp :: String -> Value -> Value -> Value
doPrimapp "eq" (VInt i1) (VInt i2) | i1 == i2  = VCapp "True" []
	       		           | otherwise = VCapp "False" []
doPrimapp "add" (VInt i) (VInt j) = VInt (i+j)
doPrimapp "sub" (VInt i) (VInt j) = VInt (i-j)
doPrimapp "mul" (VInt i) (VInt j) = VInt (i*j)
doPrimapp "and" (VCapp "True" []) (VCapp "True" []) = VCapp "True" []
doPrimapp "and" (VCapp "True" []) (VCapp "False" []) = VCapp "False" []
doPrimapp "and" (VCapp "False" []) (VCapp "True" []) = VCapp "False" []
doPrimapp "and" (VCapp "False" []) (VCapp "False" []) = VCapp "False" []

interp :: Exp -> A (Value,Heap)
interp e = run (eval mempty e)

interp' :: Exp -> [(Value,Heap)]
interp'=  sols . interp

interp'' :: Exp -> [Value]
interp'' = nub . map fst . interp'

true = Capp "True" []
false = Capp "False" []

ifthenelse c t f = Case False c [(("True",[]), t),
				 (("False",[]), f)] 

eq0 e = Primapp ("eq",False) e (Int 0)

pand = Primapp ("and",True)

rd x k = Case False (Var x) [(("True",[]), k)]
wr x k = Case True (Var x) [(("True",[]), k)]

f = Logic "x" (pand (wr "x" true)
                    (rd "x" true))

g = Logic "x" (pand (rd "x" true)
		    (wr "x" true))

q = Logic "x" (Logic "y" 
	(pand (rd "x" (wr "y" true))
	      (pand (wr "x" true)
                    (rd "y" true))))


r = Logic "x" (Logic "y" (pand (rd "x" (wr "y" (Var "x")))
	  	    	       (wr "x" (rd "y" (Var "y")))))