{-
uNESL: An instrumented interpreter for a NESL subset
Andrzej Filinski, University of Copenhagen

Version 0.10, 12 April 2018

This is a simplistic interpreter for a subset of a NESL-ish language,
instrumented to also compute work and span costs of the expressions
evaluated. For educational purposes only: the absolute performance
is unimpressive. Also, bugs are quite likely, and the current version
of the code is NOT a good example of proper Haskell style.

To start top-level loop, either "runhaskell unesl.hs", or (for
somewhat better performance) "ghc unesl.hs; ./unesl".

Top level commands are of the form:

  exp                         Type-check and evaluate the NESL expression
  function f(x1..xn) = exp    Define a NESL function
  def x = exp                 Define a NESL value
  :load filename.nesl         Read function/value defs from file
  :defs                       Show types of all currently defined identifiers
  :quit                       Exit the interpreter

Changelog:

 0.01   Initial "release"
 0.01a  Indicate (by comments) how to compile w/o Haskeline
 0.01b  Use ";" as separator in comprehensions, as per NESL docs.
        ("," still works, but is deprecated.)
 0.02   Added value definitions ("def x = exp")
 0.03   Added explicit typing of functions [subsequently reverted]
 0.04   Adapted for GHC 7.10 Applicative => Monad brain damage.
 0.05   Minor tweaks
 0.06   Added :defs command
 0.10   Multiple cleanups and fixes; some incompatible syntax tweaks

TODO:
- add builtin scans, reduces
- monadify threading of environments everywhere
- more aliases & compatibility defs in stdlib.nesl
- more general type-class hierarchy
- fix interrupts
-}

{-# LANGUAGE RankNTypes #-}

import Control.Monad
import Text.ParserCombinators.Parsec
import qualified Text.Parsec.Token as PT
import qualified Text.Parsec.Language as PL
import qualified Text.Parsec.Expr as PE
import qualified Data.Array as A
import Data.Char (chr, ord)
import Data.List (transpose)
import Data.Bits ((.&.), (.|.), xor, shift)
import qualified Data.Map as M

-- Comment out the next two imports if you don't have (and can't
-- install) Haskeline on your system. (And uncomment the compatibility
-- defs below)
import Control.Monad.Trans (lift)
import System.Console.Haskeline (InputT, getInputLine, runInputT,
                                 defaultSettings)

-- doesn't quite work right. should probably use lower-level POSIX
-- interface for user interrupts
import qualified Control.Exception as C

---- Syntax

type Id = String

data AVal =
    IVal Int
  | RVal Double
  | BVal Bool
  | CVal Char
  deriving Eq

data Exp =
    Var Id
  | Lit AVal
  | Tup [Exp]
  | Vec [Exp]
  | If Exp Exp Exp
  | Let Pat Exp Exp
  | Over Exp [(Pat, Exp)] 
  | Call Id [Exp]
  | Time Exp
  | Tag SourcePos Exp
  deriving Show

data Pat =
    PVar Id
  | PWild 
  | PTup [Pat]
  deriving Show

data Def = 
    FDef Id [Id] Exp
  | VDef Id Exp

data Top =
    TExp Exp
  | TDef Def
  | TLoad String
  | TExit
  | TDump

---- Evaluation

data Vect a = Vect {vList :: [a], vArray :: A.Array Int a}

mkVect :: [a] -> Vect a
mkVect l = Vect l (A.listArray (0, length l-1) l)

data Val =
    AVal AVal
  | TVal [Val]
  | VVal (Vect Val)
  | FVal ([Val] -> Comp Val)

instance Show AVal where
  show (IVal i) = show i
  show (RVal r) = show r
  show (BVal b) = if b then "T" else "F"
  show (CVal c) = show c

instance Show Val where
  show (AVal a) = show a
  show (TVal vs) = "(" ++ showelts vs ++ ")"
  show (VVal v) = "[" ++ showelts (vList v) ++ "]"
  show (FVal _) = "<function>"

showelts :: [Val] -> String
showelts [] = ""
showelts [x] = show x
showelts (x:xs) = show x ++ ", " ++ showelts xs

size :: Val -> Int
size (AVal _) = 1
size (TVal vs) = sum (map size vs)
size (VVal vv) = 1 + sum (map size (vList vv))
size (FVal _) = 1 -- error?

type RCtx = String

newtype Comp a = Comp {rComp :: RCtx -> Either String (a, Int, Int)}
  
instance Monad Comp where  -- exploiting that (Either String) is a monad
  return a = Comp (\c -> return (a, 0, 0))
  m >>= f = Comp (\c -> do (a, w1, s1) <- rComp m c
                           (b, w2, s2) <- rComp (f a) c
                           return (b, w1+w2, s1+s2))
  fail s = Comp (\c -> Left $ "In " ++ c ++ ": " ++ s)

-- never explicitly used
instance Functor Comp where fmap = liftM
instance Applicative Comp where pure = return; (<*>) = ap

returnc :: (Int, Int) -> a -> Comp a
returnc (w,s) a = Comp (\c -> return (a, w, s))

par :: [Comp a] -> Comp [a]
par [] = returnc (1,1) []
par (t : ts) = Comp (\c -> do (a, w1, s1) <- rComp t c
                              (as, w2, s2)  <- rComp (par ts) c 
                              return (a:as, w1+w2, s1 `max` s2))

type VEnv = M.Map Id Val

-- assumes Exp has been type-checked
eval :: Exp -> VEnv -> Comp Val
eval (Var s) r = 
  case M.lookup s r of
    Just a -> return a
    Nothing -> error ("bad variable: " ++ s)
eval (Lit l) r = 
  return (AVal l)
eval (Tup es) r = 
  do vs <- mapM (\e -> eval e r) es
     return $ TVal vs
eval (Vec es) r = 
  do vs <- mapM (\e -> eval e r) es
     returnc (1 + sum (map size vs), 1) $ VVal (mkVect vs)
eval (If e0 e1 e2) r =
  do AVal (BVal b) <- eval e0 r
     if b then eval e1 r else eval e2 r 
eval (Let p e1 e2) r =
  do v1 <- eval e1 r
     eval e2 (bind p v1 `M.union` r)
eval (Over e0 qs) r =
  do vs <- mapM (\(p, e) -> do VVal v <- eval e r; return v) qs
     let ls = map (length . vList) vs
     let ps = map (\ (p, e) -> p) qs
     (if all (\ l -> l == head ls) (tail ls) then
       do vs0 <- par [eval e0 (M.unions (zipWith bind ps w) `M.union` r) 
                      | w <- transpose (map vList vs)]
          return $ VVal (mkVect vs0)
      else fail $ "Over: length mismatch: " ++ show ls)
eval (Call i es) r =
  do vs <- mapM (\e -> eval e r) es
     case M.lookup i r of
       Just (FVal f) -> f vs
       Nothing -> error ("bad function: " ++ i)
eval (Time e) r =
  Comp (\c ->
    do (v, nw, ns) <- rComp (eval e r) "Time"
       return (TVal [v, AVal (IVal nw), AVal (IVal ns)], nw, ns))
eval (Tag t e) r = 
  Comp (\c -> rComp (eval e r) (show t))

bind :: Pat -> Val -> VEnv
bind (PVar x) v = M.singleton x v
bind PWild v = M.empty
bind (PTup ps) (TVal vs) = M.unions (zipWith bind ps vs)

---- Typing

newtype TyVar = TV Int
  deriving Eq

data Kind = 
    KAny
  | KNum

data Type = 
    TCon String [Type]
  | TPar Int
  | TVar TyVar Kind

tBase s = TCon s []
tInt = tBase "int"
tReal = tBase "real"
tBool = tBase "bool"
tChar = tBase "char"

tTup ts = TCon "()" ts
tVec t = TCon "[]" [t]
tFun ts t = TCon "->" [tTup ts, t]

instance Show Type where
  show (TCon "()" ts) = 
   "(" ++ showl ts ++ ")" where
     showl [] = ""
     showl [x] = show x
     showl (x:xs) = show x ++ ", " ++ showl xs
  show (TCon "[]" [t]) = "[" ++ show t ++ "]"
  show (TCon "->" [t1,t2]) = show t1 ++ " -> " ++ show t2
  show (TCon s []) = s
  show (TPar n) = [chr (ord 'a' + n)]

data PType =
    PTBody Type
  | PTAbs Int Kind PType

instance Show PType where
  show (PTBody t) = show t
  show (PTAbs i KAny pt) = "@" ++ show (TPar i) ++ "." ++ show pt
  show (PTAbs i KNum pt) = "@#" ++ show (TPar i) ++ "." ++ show pt

type UEnv = [(TyVar,Type)] -- should probably turn into a map as well

type TCtx = String

newtype Constr a = 
  Constr {rConstr :: UEnv -> Int -> TCtx -> Either String (a, UEnv, Int)}

instance Monad Constr where
  return a = Constr (\u n c -> return (a, [], n))
  m >>= f = Constr (\u n c -> do (a, u1, n') <- rConstr m u n c
                                 (b, u2, n'') <- rConstr (f a) (u++u1) n' c
                                 return (b, u1++u2, n''))
  fail s = Constr (\u n c -> Left ("In " ++ c ++ ": " ++ s))

-- never explicitly used
instance Functor Constr where fmap = liftM
instance Applicative Constr where pure = return; (<*>) = ap

newkvar :: Kind -> Constr Type
newkvar k = Constr (\u n c -> return (TVar (TV n) k, [], n+1))

newvar = newkvar KAny

chase :: TyVar -> Kind -> UEnv -> Type
chase x k u = case lookup x u of
                Nothing -> TVar x k
                Just (TVar y k') -> chase y k' u
                Just t -> t

fetch :: Type -> Constr Type
fetch (TVar x k) = Constr (\u n c -> return (chase x k u, [], n))
fetch t = return t

-- clean up!
addb :: TyVar -> Kind -> Type -> Constr ()
addb x KAny t = Constr (\u n c -> return ((), [(x,t)], n))
addb x KNum t@(TPar n) = Constr (\u n c -> return ((), [(x,t)], n))
addb x KNum t@(TCon "int" []) = Constr (\u n c -> return ((), [(x,t)], n))
addb x KNum t@(TCon "real" []) = Constr (\u n c -> return ((), [(x,t)], n))
addb x KNum t@(TCon s _) = fail $ "non-numeric type: " ++ s
addb x KNum t@(TVar _ KNum) = Constr (\u n c -> return ((), [(x,t)], n))
addb x KNum (TVar y KAny) = Constr (\u n c -> return ((), [(y,TVar x KNum)], n))

nooccur :: TyVar -> Type -> Constr ()
nooccur x t =
  do t' <- fetch t
     case t' of
       TVar y _ -> if x == y then fail "circularity" else return ()
       TPar _ -> return ()
       TCon _ ts -> mapM_ (nooccur x) ts

unify :: Type -> Type -> Constr ()
unify t1 t2 = do t1' <- fetch t1
                 t2' <- fetch t2
                 unifyc (t1', t2')

unifyc (TVar x _, TVar y _) | x == y = return ()
unifyc (TVar x k, t) = do nooccur x t; addb x k t
unifyc (t, TVar x k) = do nooccur x t; addb x k t
unifyc (TCon s1 ts1, TCon s2 ts2) =
  if s1 == s2 then
    if length ts1 == length ts2
    then zipWithM_ unify ts1 ts2
    else fail $ "clash: " ++ s1 ++ " found arity " ++
                show (length ts1) ++ " , expected " ++ show (length ts2)
  else fail $ "clash: found " ++ s1 ++ ", expected " ++ s2

type TEnv = M.Map Id PType

chkExp :: Exp -> TEnv -> Constr Type
chkExp (Lit (IVal _)) g = return tInt
chkExp (Lit (RVal _)) g = return tReal
chkExp (Lit (BVal _)) g = return tBool
chkExp (Lit (CVal _)) g = return tChar
chkExp (Var i) g = 
  case M.lookup i g of
    Just pt -> inst pt
    Nothing -> fail ("Unbound var: " ++ i)
chkExp (Tup es) g = 
  do ts <- mapM (\e -> chkExp e g) es
     return $ tTup ts
chkExp (Vec es) g =
  do t0 <- newvar
     mapM_ (\e -> do t <- chkExp e g; unify t t0) es
     return $ tVec t0
chkExp (If e0 e1 e2) g =
  do t0 <- chkExp e0 g
     unify t0 tBool
     t1 <- chkExp e1 g
     t2 <- chkExp e2 g
     unify t2 t1
     return t1
chkExp (Let p e0 e1) g =
 do t0 <- chkExp e0 g
    (tp, gp) <- chkPat p 
    unify t0 tp
    chkExp e1 (gp `M.union` g)
chkExp (Over e0 qs) g =
  do gs <- mapM (\ (p,e) -> do t <- chkExp e g
                               (tp, gp) <- chkPat p
                               unify t (tVec tp)
                               return gp) qs
     t0 <- chkExp e0 (M.unions gs `M.union` g)  -- should check disjointness
     return $ tVec t0
chkExp (Call f es) g =
  case M.lookup f g of
    Nothing -> fail ("Unbound function: " ++ f)
    Just pt -> do tf <- inst pt
                  ts <- mapM (\e -> chkExp e g) es
                  t <- newvar
                  unify (tFun ts t) tf
                  return t
chkExp (Time e) g =
  do t <- chkExp e g
     return $ tTup [t, tInt, tInt]
chkExp (Tag t e) g = 
  Constr (\u n c -> rConstr (chkExp e g) u n (show t))

chkPat :: Pat -> Constr (Type, TEnv)
chkPat (PVar x) = 
  do t <- newvar
     return (t, M.singleton x (PTBody t))
chkPat PWild = 
  do t <- newvar
     return (t, M.empty)
chkPat (PTup ps) =
  do trs <- mapM chkPat ps
     return (tTup (map fst trs), M.unions (map snd trs)) 
       -- should check disjointness

gener :: Type -> Constr PType
gener t =
  let abs [] n t = PTBody t
      abs (k:ks) n t = PTAbs n k (abs ks (n+1) t)
  in do (t', l) <- gener1 t []
        return $ abs l 0 t'

gener1 :: Type -> [Kind] -> Constr (Type, [Kind])
gener1 t l =
  do t0 <- fetch t
     case t0 of
       TVar x k -> do unify (TVar x k) (TPar (length l))
                      return $ (TPar (length l), l++[k])
       TPar i -> return $ (TPar i, l)
       TCon s ts -> let gl [] m = return ([], m)
                        gl (t0:tr) m = do (t0',m') <- gener1 t0 m
                                          (tr',m'') <- gl tr m'
                                          return (t0':tr', m'')
                    in do (ts',l') <- gl ts l
                          return $ (TCon s ts', l')

chkDef :: Def -> TEnv -> Constr (Id, PType)
chkDef (FDef f xs e) g =
  do ts <- mapM (\x -> newvar) xs
     t0 <- newvar
     let tf = tFun ts t0
     let g1 = M.insert f (PTBody tf)
                         (M.fromList (zipWith (\x t -> (x, PTBody t)) xs ts)
                            `M.union` g)
     te <- chkExp e g1
     unify te t0
     ptf <- gener tf
     return (f, ptf)
chkDef (VDef x e) g =
  do te <- chkExp e g
     pte <- gener te
     return (x, pte)

subst :: Type -> Int -> Type -> Type
subst t0@(TPar x) y t1 = if x == y then t1 else t0
subst t0@(TVar v k) y t1 = t0
subst (TCon s ts) y t1 = TCon s (map (\t -> subst t y t1) ts)

inst :: PType -> Constr Type
inst (PTBody t) = return t
inst (PTAbs i k pt) = do tx <- newkvar k
                         t <- inst pt
                         return $ subst t i tx

---- Initial environment

tP0 = TPar 0


ptFun :: [Type] -> Type -> PType
ptFun ts t = PTBody $ tFun ts t

primop :: ([AVal] -> AVal) -> Val
primop f = FVal (\as -> returnc (1,1) $ AVal (f [v | AVal v <- as]))

numop1 :: (forall a. Num a => a -> a) -> [AVal] -> AVal
numop1 o [IVal n1] = IVal (o n1)
numop1 o [RVal n1] = RVal (o n1)

numop2 :: (forall a. (Num a, Ord a) => a -> a -> a) -> [AVal] -> AVal
numop2 o [IVal n1, IVal n2] = IVal (n1 `o` n2)
numop2 o [RVal n1, RVal n2] = RVal (n1 `o` n2)

compop :: (forall a. (Ord a) => a -> a -> Bool) -> [AVal] -> AVal
compop o [IVal n1, IVal n2] = BVal (n1 `o` n2)
compop o [RVal n1, RVal n2] = BVal (n1 `o` n2)

seglist :: [Int] -> [a] -> [[a]]
seglist [] [] = []
seglist (n:ns) l = take n l : seglist ns (drop n l)

scatter :: [(Int, a)] -> Int -> [[a]]
scatter ps n = 
  let ins ls (i,x) = let (h:t) = drop i ls in take i ls ++ (x:h):t
  in map reverse (foldl ins (replicate n []) ps)

initEnv :: [(Id, PType, Val)]
initEnv =
  [("T", PTBody tBool, AVal (BVal True)),
   ("true", PTBody tBool, AVal (BVal True)),
   ("F", PTBody tBool, AVal (BVal False)),
   ("false", PTBody tBool, AVal (BVal False)),
   ("pi", PTBody tReal, AVal (RVal pi)),

   ("_plus", PTAbs 0 KNum $ ptFun [tP0, tP0] tP0,
    primop (numop2 (+))),
   ("_minus", PTAbs 0 KNum $ ptFun [tP0, tP0] tP0,
    primop (numop2 (-))),
   ("_uminus", PTAbs 0 KNum $ ptFun [tP0] tP0,
    primop (numop1 negate)),
   ("_times", PTAbs 0 KNum $ ptFun [tP0, tP0] tP0,
    primop (numop2 (*))),
   ("_times", PTAbs 0 KNum $ ptFun [tP0, tP0] tP0,
    primop (numop2 (*))),
   ("_div", PTAbs 0 KNum $ ptFun [tP0, tP0] tP0,
    primop (\vs -> case vs of [IVal n1, IVal n2] -> IVal (n1 `div` n2)
                              [RVal n1, RVal n2] -> RVal (n1 / n2))),
   ("_mod", ptFun [tInt, tInt] tInt,
    primop (\ [IVal n1, IVal n2] -> IVal (n1 `mod` n2))),
   ("pow", PTAbs 0 KNum (ptFun [tP0, tInt] tP0),
    primop (\ [v, IVal n2] -> numop1 (^ n2) [v])),
   ("min", PTAbs 0 KNum (ptFun [tP0, tP0] tP0),
    primop (numop2 min)),
   ("max", PTAbs 0 KNum (ptFun [tP0, tP0] tP0),
    primop (numop2 max)),

   ("_and", ptFun [tInt, tInt] tInt,
    primop (\ [IVal n1, IVal n2] -> IVal (n1 .&. n2))),
   ("_or", ptFun [tInt, tInt] tInt,
    primop (\ [IVal n1, IVal n2] -> IVal (n1 .|. n2))),
   ("_xor", ptFun [tInt, tInt] tInt,
    primop (\ [IVal n1, IVal n2] -> IVal (n1 `xor` n2))),
   ("lshift", ptFun [tInt, tInt] tInt,
    primop (\ [IVal n1, IVal n2] -> IVal (shift n1 n2))),
   ("rshift", ptFun [tInt, tInt] tInt,
    primop (\ [IVal n1, IVal n2] -> IVal (shift n1 (-n2)))),

   ("_eq", PTAbs 0 KNum $ ptFun [tP0, tP0] tBool,
    primop (compop (==))),
   ("_leq", PTAbs 0 KNum $ ptFun [tP0, tP0] tBool,
    primop (compop (<=))),

   ("not", ptFun [tBool] tBool, primop (\ [BVal b] -> BVal (not b))),
   ("chr", ptFun [tInt] tChar, primop (\ [IVal n] -> CVal (chr n))),
   ("ord", ptFun [tChar] tInt, primop (\ [CVal c] -> IVal (ord c))),
   ("real", ptFun [tInt] tReal, primop (\ [IVal n] -> RVal (fromIntegral n))),
   ("floor", ptFun [tReal] tInt, primop (\ [RVal n] -> IVal (floor n))),
   ("sqrt", ptFun [tReal] tReal, primop (\ [RVal n] -> RVal (sqrt n))),
   ("sqrt", ptFun [tReal] tReal, primop (\ [RVal n] -> RVal (sqrt n))),
   ("sin", ptFun [tReal] tReal, primop (\ [RVal n] -> RVal (sqrt n))),
   ("cos", ptFun [tReal] tReal, primop (\ [RVal n] -> RVal (sqrt n))),
   ("ln", ptFun [tReal] tReal, primop (\ [RVal n] -> RVal (sqrt n))),
   ("exp", ptFun [tReal] tReal, primop (\ [RVal n] -> RVal (sqrt n))),

   ("_length", PTAbs 0 KAny $ ptFun [tVec tP0] tInt,
    FVal (\ [VVal vs] -> let (lo,hi) = A.bounds (vArray vs)
                         in return $ AVal (IVal (hi-lo+1)))),
   ("_sub", PTAbs 0 KAny $ ptFun [tVec tP0, tInt] tP0,
    FVal (\ [VVal vs, AVal (IVal n)] ->
               let va = vArray vs
                   (lo,hi) = A.bounds va
               in if lo <= n && n <= hi
                  then returnc (1,1) $ va A.! n
                  else fail $ "bad subscript: " ++ show n ++ 
                              ", vector length: " ++ show (hi-lo+1))),
   ("index", ptFun [tInt] (tVec tInt),
    FVal (\ [AVal (IVal n)] -> 
           returnc (n,1) $ VVal (mkVect [AVal (IVal i) | i <- [0..n-1]]))),

   ("_append", PTAbs 0 KAny $ ptFun [tVec tP0, tVec tP0] (tVec tP0),
    FVal (\ [VVal v1, VVal v2] -> 
           let v = vList v1 ++ vList v2
           in returnc (length v,1) (VVal (mkVect v)))),
   ("concat", PTAbs 0 KAny $ ptFun [tVec (tVec tP0)] (tVec tP0),
    FVal (\ [VVal vs] -> 
           let v = concat [vList v | VVal v <- vList vs]
           in returnc (1,1) (VVal (mkVect v)))),
   ("partition", PTAbs 0 KAny $ ptFun [tVec tP0, tVec tInt] (tVec (tVec tP0)),
    FVal (\ [VVal vs, VVal ls] ->
           let ns = [n | AVal (IVal n) <- vList ls]
               l = length (vList vs)
           in if sum ns == l 
              then returnc (length ns,1)
                    $ VVal (mkVect [VVal (mkVect v) | v <- seglist ns (vList vs)])
              else fail "partition: length mismatch")),
   ("scatter", PTAbs 0 KAny $ ptFun [tVec (tTup [tInt, tP0]), tInt]
                                     (tVec (tVec tP0)),
    FVal (\ [VVal v, AVal (IVal n)] ->
           let ps = [(i,a) | TVal [AVal (IVal i), a] <- vList v]
           in if all (\(i, _) -> 0 <= i && i < n) ps
              then returnc (n+length ps,1) $
                     VVal (mkVect (map (VVal . mkVect) (scatter ps n)))
              else fail "scatter: bad index")),

   ("error", PTAbs 0 KAny $ ptFun [tVec tChar] tP0,
    FVal (\ [VVal s] -> fail [c | AVal (CVal c) <- vList s]))]

---- REPL

type CEnv = (VEnv, TEnv)

doExp :: Exp -> CEnv -> IO ()
doExp e (r,g) =
  case rConstr (do t <- chkExp e g; gener t) [] 0 "Top" of
    Right (pt, _, _) -> 
      C.catch
        (case rComp (eval e r) "TopExp" of
           Right (v, nw, ns) -> 
              do putStrLn $ show v ++ " : " ++ show pt
                 putStrLn $ "[Work: " ++ show nw ++ ", span: " ++ show ns ++ "]"
           Left s -> putStrLn $ "Runtime error: " ++ s)
        (\e -> putStrLn $ "Low-level error: " ++ show (e::C.SomeException))
    Left s ->
      putStrLn $ "Type error: " ++ s

doDef :: Def -> CEnv -> IO (Maybe CEnv)
doDef d (r,g) =
  case rConstr (chkDef d g) [] 0 "TopDef" of
    Right ((i, pt), _, _) -> 
      case d of
        FDef i xs e -> 
          let r1 = M.insert i (FVal f) r -- recursive!
              f vs = Comp(\c -> rComp (eval e (M.fromList (zip xs vs) `M.union` r1)) i)
          in do putStrLn $ "Defined " ++ i ++ " : " ++ show pt
                return $ Just (r1, M.insert i pt g)
        VDef x e ->
          case rComp (eval e r) ("Def " ++ x) of
            Right (v, _, _) -> 
              do putStrLn $ "Defined " ++ x ++ " : " ++ show pt
                 return $ Just (M.insert i v r, M.insert x pt g)
            Left s -> 
              do putStrLn $ "Runtime error in def " ++ x ++ ": " ++ s
                 return Nothing
    Left s ->
      do putStrLn $ "Type error in definition: " ++ s
         return Nothing

doDefs :: [Def] -> CEnv -> IO CEnv
doDefs [] c = return c
doDefs (d:ds) c = do mc' <- doDef d c
                     case mc' of
                       Nothing -> return c
                       Just c' -> doDefs ds c'

doTop :: Top -> CEnv -> IO CEnv
doTop (TExp e) c  =
  do doExp e c
     return c
doTop (TDef d) c =
  do mc' <- doDef d c
     case mc' of
       Nothing -> return c
       Just c' -> return c'
doTop (TLoad fn) c =
  do s <- C.catch (readFile fn)
                  (\e -> do print (e::IOError); return "")
     case parse parseDefs fn s of
       Right ds -> doDefs ds c
       Left err -> do print err; return c
doTop TDump c =
  do mapM_ (\(i,t) -> putStrLn $ i ++ " : " ++ show t) (M.toList (snd c))
     return c

-- Haskeline compatibility defs
{-
type InputT m a = m a
lift = id
getInputLine p = do putStr p; s <- getLine; return (Just s)
defaultSettings = ()
runInputT () m = m
-}

repl :: CEnv -> InputT IO ()
repl r = 
  do ms <- getInputLine "$> "
     case ms of
       Nothing -> return ()
       Just "" -> repl r
       Just s ->
         case parse (do {whitespace; e <- parseTop; eof; return e})
                    "<interactive>" s of
           Right TExit -> lift $ putStrLn "Bye!"
           Right tc -> do r' <- lift (doTop tc r); repl r'
           Left s -> do lift (print s); repl r

main :: IO ()
main = runInputT defaultSettings $
         repl (M.fromList (map (\(i, _, v) -> (i,v)) initEnv),
               M.fromList (map (\(i, pt, _) -> (i,pt)) initEnv))

---- Parsing

langDef :: PT.LanguageDef st
langDef =
   PL.emptyDef
    { PT.commentLine = "%",
      PT.opLetter = oneOf "=+", -- bit of a hack, but ok for now
      PT.reservedOpNames = ["+", "-", "++", "=", "/", "/=", "*", "^",
                            "<", "<=", ">", ">=", "#", "&"],
      PT.caseSensitive = False
    }

lexer = PT.makeTokenParser langDef

whitespace = PT.whiteSpace lexer
symbol = PT.symbol lexer
oper = PT.reservedOp lexer
keyword = PT.reserved lexer

parseVar :: Parser Id
parseVar = PT.identifier lexer

parsePat :: Parser Pat
parsePat =
  do x <- parseVar
     return $ PVar x
  <|>
  do keyword "_"
     return PWild
  <|>
  do symbol "("
     l <- parsePat `sepBy` (symbol ",")
     symbol ")"
     return $ PTup l

binop :: String -> Exp -> Exp -> Exp
binop s x y = Call s [x,y]

unop :: String -> Exp -> Exp
unop s x = Call s [x]

parseExp :: Parser Exp
parseExp =
  do keyword "let"
     binds <- (do p <- parsePat
                  oper "="
                  e1 <- parseExp
                  return (p,e1))
              `sepBy1` (symbol ";")
     keyword "in"
     e2 <- parseExp
     return $ foldr (\(p,e1) e2 -> Let p e1 e2) e2 binds
  <|>
  do sp <- getPosition
     keyword "if"
     e1 <- parseExp
     keyword "then"
     e2 <- parseExp
     keyword "else"
     e3 <- parseExp
     return $ Tag sp (If e1 e2 e3)
  <|>
  do keyword "time"
     e <- parseExp
     return $ Time e
  <|>
  PE.buildExpressionParser opTable parsePrefix

mkop1 ps f =
  do sp <- getPosition; ps; return $ \x -> Tag sp (f x)

mkop2 ps f =
  do sp <- getPosition; ps; return $ \x y -> Tag sp (f x y)

opTable =
  [[PE.Infix (mkop2 (oper "*") (binop "_times")) PE.AssocLeft,
    PE.Infix (mkop2 (oper "/") (binop "_div")) PE.AssocLeft,
--  PE.Infix (mkop2 (oper "%") $ binop "_mod") PE.AssocLeft],
    PE.Infix (mkop2 (keyword "mod") (binop "_mod")) PE.AssocLeft],
   [PE.Infix (mkop2 (oper "++") (binop "_append")) PE.AssocLeft,
    PE.Infix (mkop2 (oper "+") (binop "_plus")) PE.AssocLeft,
    PE.Infix (mkop2 (oper "-") (binop "_minus")) PE.AssocLeft],
   [PE.Infix (mkop2 (oper "==") (binop "_eq")) PE.AssocNone,
    PE.Infix (mkop2 (oper "/=") (\x y -> unop "not" (binop "_eq" x y)))
      PE.AssocNone,
    PE.Infix (mkop2 (oper "<=") (binop "_leq")) PE.AssocNone,
    PE.Infix (mkop2 (oper ">=") (\x y -> binop "_leq" y x)) PE.AssocNone,
    PE.Infix (mkop2 (oper "<") (\x y -> unop "not" (binop "_leq" y x)))
      PE.AssocNone,
    PE.Infix (mkop2 (oper ">") (\x y -> unop "not" (binop "_leq" x y)))
      PE.AssocNone],
   [PE.Infix (mkop2 (keyword "and") (binop "_and")) PE.AssocLeft,
    PE.Infix (mkop2 (keyword "or") (binop "_or")) PE.AssocLeft,
    PE.Infix (mkop2 (keyword "xor") (binop "_xor")) PE.AssocLeft]]

-- Parsec's native Prefix ops don't nest
parsePrefix :: Parser Exp
parsePrefix =
  do f <- choice [mkop1 (oper "#") (unop "_length"),
                  mkop1 (oper "&") (unop "index"),
                  mkop1 (oper "-") (unop "_uminus")]
     e <- parsePrefix
     return $ f e
  <|>
  parseSub

parseSub :: Parser Exp
parseSub = do e <- parseAtom
              ss <- many sub
              return $ foldl (\e1 (e2,p) -> Tag p (Call "_sub" [e1,e2])) e ss
            where sub = do sp <- getPosition 
                           symbol "["
                           e1 <- parseExp
                           symbol "]"
                           return (e1, sp)

parseAtom :: Parser Exp
parseAtom = do sp <- getPosition
               x <- parseVar
               (do symbol "("
                   es <- parseExp `sepBy` (symbol ",")
                   symbol ")"
                   return $ Tag sp (Call x es)
                <|>
                (return $ Tag sp (Var x)))
            <|>
            do s <- many1 digit
               (do char '.'
                   s2 <- many1 digit
                   whitespace
                   return $ Lit (RVal (read (s ++ "." ++ s2)))
                <|>
                do whitespace
                   return $ Lit (IVal (read s)))
            <|>
            do char '\''
               c <- anyChar
               char '\''
               whitespace
               return $ Lit (CVal c)
            <|>
            do char '"'
               s <- manyTill anyChar (char '"')
               whitespace
               return $ Vec (map (Lit . CVal) s)
            <|>
            do symbol "("
               es <- parseExp `sepBy` (symbol ",")
               symbol ")"
               case es of
                 [e] -> return e
                 _ -> return $ Tup es
            <|>
            do symbol "{"
               e <- parseExp
               symbol ":"
               qs <- parseQual `sepBy1` (symbol ";" {- <|> symbol "," -})
               me <- option Nothing
                      (do symbol "|"
                          e <- parseExp
                          return $ Just e)
               symbol "}"
               case me of
                 Nothing -> return $ Over e qs 
                 Just ef -> return $ Call "concat" [Over (If ef (Vec [e]) (Vec [])) qs]
            <|>
            do sp <- getPosition
               symbol "["
               es <- parseExp `sepBy` (symbol ",")
               symbol "]"
               return $ Tag sp (Vec es)

parseQual :: Parser (Pat, Exp)
parseQual = do p <- parsePat
               symbol "in" {- <|> symbol "<-" -}
               e <- parseExp
               return (p, e)

parseDef :: Parser Def
parseDef = do keyword "function"
              f <- parseVar
              symbol "("
              as <- parseVar `sepBy` (symbol ",")
              symbol ")"
              oper "="
              e <- parseExp         
              return $ FDef f as e
           <|> 
           do keyword "def"
              x <- parseVar
              oper "="
              e <- parseExp         
              return $ VDef x e
            
parseDefs :: Parser [Def]
parseDefs = do whitespace
               ds <- many parseDef
               eof
               return ds

parseTop :: Parser Top
parseTop = do d <- parseDef
              return $ TDef d
           <|>
           do d <- keyword ":load"
              s <- many1 anyChar
              return $ TLoad s
           <|>              
           do d <- keyword ":quit"
              return $ TExit
           <|>
           do d <- keyword ":defs"
              return $ TDump
           <|>
           do e <- parseExp
              return $ TExp e