{- 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,
   and non-determinism.  -}

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

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

data Exp = Var Var
	 | Int Int
	 | Abs Var Exp
         | App Exp Exp
         | Capp Constr [Exp]
	 | Primapp Prim Exp Exp
         | Case 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
newtype M a = M (Heap -> A (a,Heap))
instance Monad M where
  (M m1) >>= k = M (\h -> do (a',h') <- m1 h
			     let M m2 = k a' in m2 h')
  return x = M (\h -> return (x,h))
instance MonadPlus M where
  mzero = M (\ _ -> mzero)
  (M m1) `mplus` (M m2) = M (\h -> m1 h `mplus` m2 h)
fresh :: M HPtr
fresh = M (\ h -> return (hfresh h))
store :: HPtr -> HEntry -> M ()
store p e = M (\ h -> return ((),hset h (p,e)))
fetch :: HPtr -> M HEntry
fetch p = M (\ h -> return (hget h p,h))
run :: M a -> A (a,Heap) 
run (M m) = m hempty
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 e1 e2) = 
     do v1 <- eval env e1
	v2 <- eval env e2
	checkGround v1
	checkGround v2
        return (doPrimapp p v1 v2)
eval env (Case 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 ->
            msum (map f pes)
            where
              f ((c,xs),e') =
                do ps <- mapM (const allocLogic) xs
	           store p0 (HValue (VCapp c ps))                       
                   let env' = foldl mset env (zip xs ps)
		   eval env' e'
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 :: Prim -> 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'

success = Capp "Success" []

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

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

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

amb a b = Logic "dummy" (Case (Var "dummy") [(("Ldummy",[]),a),
					     (("Rdummy",[]),b)])

c = Letrec [("undef", Var "undef")]
	   (amb (Var "undef") (Int 1))

coin = amb (Int 0) (Int 1)

coin1 = Letrec [("x",coin)] (Primapp "add" (Var "x") (Var "x"))

coin2 = Letrec [("x",coin),("y",coin)] (Primapp "add" (Var "x") (Var "y"))

coinf1 = Letrec [("coinf", Abs "dummy" coin),
		 ("x", (App (Var "coinf") (Int 0)))]
	       (Primapp "add" (Var "x") (Var "x"))
coinf2 = Letrec [("coinf", Abs "dummy" coin)]
	       (Primapp "add" (App (Var "coinf") (Int 0)) (App (Var "coinf") (Int 1)))