pretty printing of terms + typechecking handle WIP

src/Check.hs

@@ -14,9 +14,10 @@ import Control.Monad.Writer
 import Control.Monad.RWS
 import Control.Lens
 
-import           Data.List((\\))
+import           Data.List((\\), sort)
 import qualified Data.Map as M
 import           Data.Map(Map)
+import           Data.Maybe(catMaybes, isNothing)
 import qualified Data.Set as S
 import           Data.Set(Set)
 import qualified Data.Text as T
@@ -141,33 +142,50 @@ infer (Let v body exp) = do
   env <- ask
   let s = generalize env ty1
   inEnv v s $ infer exp
--- infer (Handle t hs) = do
---   (tyh, eh, lbl) <- inferHandlers hs
---   (ty1, e1) <- infer t
---   fr <- freshRow
---   constr e1 (Row [lbl] (Just fr))
---   constr (TyVar fr) eh
---   constr ty1 tyh
---   return (tyh, eh)
+infer (Handle lbl t hs) = do
+  (ty1, e1) <- infer t
+  fr <- fresh
+  constrRow e1 (Row [lbl] (Just fr))
+  (tyR, eR) <- inferHandlers lbl (ty1, (Row [] (Just fr))) hs
+  return (tyR, eR)
 infer (Lift lbl t) = do
   fr <- fresh
   (ty, eff) <- infer t
   constrRow (Row [lbl] (Just fr)) eff
   return (ty, eff)
 
--- extractLabel :: Handler -> Set Ident
--- extractLabel (Op x _ _ _) = S.singleton x
--- extractLabel (Ret _ _) = S.empty
+inferHandler :: Check m => TyLit -> (Typ, Row) -> Handler -> m (Maybe Ident, Typ)
+inferHandler hLbl (resT, resE) (Op id arg cont exp) = do
+  (lbl, argT, retT) <- lookupOp id
+  if hLbl /= lbl then throw $ 
+    TypeError ("Could not match eff " ++ show lbl ++ " of " ++ show id ++ " with required effect "
+               ++ show hLbl) else return ()
+  (ty, env) <- inEnv arg (Scheme [] argT) 
+               $ inEnv cont (Scheme [] (TyArr retT resE resT)) $ infer exp
+  constrRow env resE
+  return (Just id, ty)
+inferHandler hLbl (resT, resE) (Ret val exp) = do
+  (ty, env) <- inEnv val (Scheme [] resT) $ infer exp
+  constrRow env resE
+  return (Nothing, ty)
+
 
 -- -- TODO We probably should pass type of handled expression,
 -- -- to construct type of continuation.
 
--- inferHandlers :: [Handler] -> Check (Typ, Typ, Ident)
--- inferHandlers hs = do
---   lbls <- mapM lookupEff . S.toList . fold extractLabel S.empty $ hs
---   lbl <- case S.fromList . S.toList $ lbls of
---     [l] -> return l
---     ls  -> throw $ TypeError $ "Handlers had unexpected amount of effects: " ++ show ls
+inferHandlers :: Check m => TyLit -> (Typ, Row) -> [Handler] -> m (Typ, Row)
+inferHandlers lbl (resT, resE) hs = do
+  types <- mapM (inferHandler lbl (resT, resE)) hs
+  desired <- lookupEff lbl
+  -- liftIO $ putDocW 80 $ pretty desired
+  let idents = catMaybes . map fst $ types
+  let ret    = filter isNothing . map fst $ types
+  if sort idents /= desired || ret /= [Nothing]
+    then throw $ TypeError ("Wrong handlers, expected " ++ show desired ++ " got " ++ show idents)
+    else return ()
+  let typs = map snd types
+  zipWithM constrTyp typs (tail typs)
+  return (head typs, resE)
 
 generalize :: Environment -> Typ -> Scheme
 generalize env ty = Scheme (S.toList vars) ty
@@ -213,10 +231,10 @@ unifyR s (Row l1 v1, Row l2 v2) = let
 
       (Just a, Just b) -> do
         fr <- fresh
-        liftIO $ putDocW 80 $ "unifiying" P.<+> pretty (Row l1 v1) P.<+> "and" P.<+> pretty (Row l2 v2) P.<> P.line
+        -- liftIO $ putDocW 80 $ "unifiying" P.<+> pretty (Row l1 v1) P.<+> "and" P.<+> pretty (Row l2 v2) P.<> P.line
         s' <- (extend a (Right $ Row extraInL2 (Just fr)) s >>=
                 extend b (Right $ Row extraInL1 (Just fr)))
-        liftIO $ putDocW 80 $ "resulting subst:" P.<+> pretty s' P.<> P.line
+        -- liftIO $ putDocW 80 $ "resulting subst:" P.<+> pretty s' P.<> P.line
         return $ (s', [])
 
 solve :: Solve m => Subst -> Constraints -> m Subst

src/Environment.hs

@@ -10,6 +10,7 @@ import Control.Monad.Reader
 import Control.Monad.Except
 import Control.Lens
 
+import           Data.List(sort)
 import qualified Data.Map as M
 import           Data.Map (Map)
 
@@ -22,6 +23,7 @@ data Scheme = Scheme [TyVar] Typ
 data Environment = Env
   { _eTypeContext :: Map Ident Scheme
   , _eOperations  :: Map Ident (TyLit, Typ, Typ)
+  , _eEffects     :: Map TyLit [Ident]
   } deriving (Show)
 
 makeLenses ''Environment
@@ -34,13 +36,19 @@ operations = M.fromList
   ]
   where tl = TyLit . TL
 
+effToOps :: Map TyLit [Ident]
+effToOps = fmap sort . 
+           M.fromListWith (++) .
+           map (\(id, l) -> (l, [id])) . M.toList .
+           fmap (\(l, _, _) -> l) $ operations
+
 effects :: Map Ident Scheme
 effects = fmap (\(eff, a, b) -> 
   let v = TV "'a" in
     Scheme [v] (TyArr a (Row [eff] (Just v)) b)) operations
 
 initEnv :: Environment
-initEnv = Env effects operations
+initEnv = Env effects operations effToOps
 
 lookup :: (MonadReader Environment m, MonadError Error m) => Ident -> m Scheme
 lookup v = do
@@ -49,12 +57,19 @@ lookup v = do
     Just t -> return t
     Nothing -> throwError $ UnboundVariable (show v)
 
-lookupEff :: (MonadReader Environment m, MonadError Error m) => Ident -> m (TyLit, Typ, Typ)
-lookupEff v = do
+lookupOp :: (MonadReader Environment m, MonadError Error m) => Ident -> m (TyLit, Typ, Typ)
+lookupOp v = do
   ml <- asks (\e -> e ^. eOperations . to (M.lookup v))
   case ml of
     Just l  -> return l
     Nothing -> throwError $ UnknownOperation v
 
+lookupEff :: (MonadReader Environment m, MonadError Error m) => TyLit -> m [Ident]
+lookupEff v = do
+  ml <- asks (\e -> e ^. eEffects . to (M.lookup v))
+  case ml of
+    Just l  -> return l
+    Nothing -> throwError $ UnknownEffect v
+
 inEnv :: (MonadReader Environment m) => Ident -> Scheme -> m a -> m a
 inEnv v s = local (eTypeContext %~ M.insert v s)

src/Error.hs

@@ -10,6 +10,7 @@ data Error
   = ParseError String
   | UnboundVariable String
   | UnknownOperation Ident
+  | UnknownEffect TyLit
   | TypeError String
   | KindError String
   | UnificationError String

src/Lex.hs

@@ -39,7 +39,7 @@ lexer = P.makeTokenParser $ P.LanguageDef
   , P.opStart         = P.oneOf "-<>"
   , P.opLetter        = P.oneOf "-<>"
   , P.reservedNames   = map T.unpack $ S.toList reservedTokens
-  , P.reservedOpNames = [ "=", "->", ":" ]
+  , P.reservedOpNames = [ "=", "->", ":", "," ]
   , P.caseSensitive   = True
   }
 
@@ -52,6 +52,10 @@ reservedTokens = S.fromList
   [ "let"
   , "run"
   , "in"
+  , "fn"
+  , "handle"
+  , "with"
+  , "return"
   ]
 
 isReserved txt = S.member txt reservedTokens
@@ -62,10 +66,10 @@ resOp = P.reservedOp lexer
 
 parens = P.parens lexer
 
-str s = lexeme $ P.string s 
+str s = lexeme $ P.string s
 
 val =  VInt <$> P.integer lexer
-   <|> pure VUnit <* res "()"
+   <|> pure VUnit <* P.try (str "()")
 
 tyVar = lexeme $ P.try $ do
   t <- (do

src/Parse.hs

@@ -34,7 +34,7 @@ term =  Let    <$> (res "let" *> ident) <*> (resOp "=" *>  term) <*> (res "in" *
     <|> Abs    <$> (res "fn"  *> ident) <*> (resOp "->" *> term)
     <|> Handle <$> (res "handle" *> tyLit) <*> (res "in" *> term) <*> (res "with" *> handlers)
     <|> Lift   <$> (res "lift"   *> tyLit) <*> (res "in" *> term2)
-    <|> (P.try $ Bind   <$> (ident <* resOp "<-") <*> term <*> (resOp ";" *> term))
+    <|> (P.try $ Bind   <$> (ident <* resOp "<-") <*> term <*> (resOp "," *> term))
     <|> term1
 
 term1 =  foldl1 App <$> P.many1 term2
@@ -43,9 +43,9 @@ term2 =  Var <$> ident
      <|> Lit <$> val
      <|> parens term
 
-handlers = many  (  Op  <$> ident <*> ident <*> (resOp "," *> ident) <*> (resOp "->" *> term)
-                <|> Ret <$> (res "return" *> ident) <*> (resOp "->" *> term)
-                 )
+handlers = P.endBy (  Op  <$> ident <*> ident <*> (resOp "," *> ident) <*> (resOp "->" *> term)
+                  <|> Ret <$> (res "return" *> ident) <*> (resOp "->" *> term)
+                   ) (resOp ";")
 
 typ =  (P.try $ TyArr <$> typ1 <*> (resOp "->" *> row) <*> typ)
    <|> typ1

src/Print.hs

@@ -24,20 +24,53 @@ instance Pretty TyLit where
 instance Pretty Ident where
   pretty (ID i) = pretty i
 
+instance Pretty Val where
+  pretty (VInt i) = pretty i
+  pretty VUnit    = "()"
+
 instance Pretty Typ where
   pretty = prettyType 0
 
 instance Pretty Row where
   pretty = prettyRow
 
+instance Pretty Handler where
+  pretty (Op id arg cont exp) = 
+    pretty id <+> pretty arg <> "," <+> pretty cont <+> "->" <+> pretty exp
+  pretty (Ret id exp) = "return" <+> pretty id <+> "->" <+> pretty exp
+
+instance Pretty Term where
+  pretty = prettyTerm 0
+
+priority :: Int -> Int -> Doc ann -> Doc ann
+priority x y = if x <= y then id else parens 
+
 prettyType :: Int -> Typ -> Doc ann
 prettyType x (TyArr a r b) = 
-  (if x < 1 then id else parens) $ (prettyType 1 a <+> "->" <+> pretty r <+> prettyType 0 b)
+  priority x 0 $ prettyType 1 a <+> "->" <+> pretty r <+> prettyType 0 b
 prettyType _ (TyVar v) = pretty v
 prettyType _ (TyLit l) = pretty l
 
 prettyRow :: Row -> Doc ann
 prettyRow (Row [] (Just v)) = pretty v
 prettyRow (Row ls (Just v)) = 
   lbracket <> sep (punctuate "," $ map pretty ls) <+> pipe <+> pretty v <> rbracket
-prettyRow (Row ls Nothing)  = brackets $ sep (punctuate "," $ map pretty ls)
+prettyRow (Row ls Nothing)  = brackets $ sep (punctuate "," $ map pretty ls)
+
+prettyTerm :: Int -> Term -> Doc ann
+prettyTerm x (App a b) = 
+  priority x 1 $ prettyTerm 2 a <+> prettyTerm 2 b
+prettyTerm x (Let id b e) = 
+  priority x 0 $ "let" <+> pretty id <+> "=" <+> prettyTerm 0 b <+> "in" <+> prettyTerm 0 e
+prettyTerm x (Abs id b) = 
+  priority x 0 $ "fn" <+> pretty id <+> "->" <+> prettyTerm 0 b
+prettyTerm x (Var id) = pretty id
+prettyTerm x (Lit l) = pretty l
+prettyTerm x (Handle tl t hs) =
+  priority x 0 $ "handle" <+> pretty tl <+> "in" <+> prettyTerm 0 t <+> "with" 
+                          <+> align (vsep $ punctuate ";" $ map pretty hs)
+prettyTerm x (Lift tl t) =
+  priority x 0 $ "lift" <+> pretty tl <+> "in" <+> prettyTerm 2 t
+prettyTerm x (Bind id a b) =
+  priority x 0 $ pretty id <+> "<-" <+> prettyTerm 0 a <+> ";" <+> prettyTerm 0 b
+

test/Spec.hs

@@ -32,23 +32,31 @@ unEither x = case x of
 test :: (MonadError Error m, MonadIO m) => Text -> m ()
 test s = do
   t <- parse term s
+  liftIO $ putDocW 80 $ pretty t <> line
   (p, st, cs) <- liftIO $ runCheck $ process t
   (t, e) <- unEither p
-  liftIO $ putDocW 80 $ pretty cs <> line <> pretty t <+> "|" <+> pretty e <> line
+  liftIO $ putDocW 80 $ -- pretty cs <> line <> 
+                        pretty t <+> "|" <+> pretty e <> line
+
+reportError :: (MonadIO m) => ExceptT Error m a -> m ()
+reportError x = do
+  e <- runExceptT x
+  case e of
+    Left err -> liftIO $ print err
+    Right x  -> return ()
 
 main :: IO ()
 main = do
   putStrLn ""
-  mapM_ (\s -> putStrLn (T.unpack s) >> (runExceptT (test s)) >> putStrLn "")
+  mapM_ (\s -> reportError (test s) >> putStrLn "")
     [ "5"
-    -- , "fn x -> 5"
-    -- , "(fn x -> 5) ()"
-    -- , "print"
+    , "fn x -> 5"
+    , "(fn x -> 5) ()"
+    , "print"
     , "print 5"
-    -- , "fn x -> print x"
-    -- , "fn x -> fn y -> (fn a -> print 4) (y x)"
-    -- , "fn x -> fn y -> z <- y x; print 5"
-    -- , "let id = fn x -> x in id id"
-    -- , "fn x -> x (x 5)"
-    -- , "fn x -> x 5 (x ())"
+    , "fn x -> print x"
+    , "handle ST in put 5 with \
+       \ put x, r -> fn s -> (r ()) x; \
+       \ get u, r -> fn s -> (r s) s;  \
+       \ return x -> fn s -> s;"
     ]