Support Let

4 files changed, 54 insertions, 6 deletions

diff --git a/ast/CC/AST/Source.hs b/ast/CC/AST/Source.hs
index 11b7bc6..e648759 100644
--- a/ast/CC/AST/Source.hs
+++ b/ast/CC/AST/Source.hs
@@ -27,6 +27,7 @@ data Type = TFun Type Type
   deriving (Show, Read)
 
 data Expr = Lam SourceRange [(Name, SourceRange)] Expr
+          | Let SourceRange (Name, SourceRange) Expr Expr
           | Call SourceRange Expr Expr
           | Int SourceRange Int
           | Tup SourceRange [Expr]
@@ -43,6 +44,7 @@ instance Pretty Type where
 
 instance HasRange Expr where
     range (Lam sr _ _) = sr
+    range (Let sr _ _ _) = sr
     range (Call sr _ _) = sr
     range (Int sr _) = sr
     range (Tup sr _) = sr
diff --git a/ast/CC/AST/Typed.hs b/ast/CC/AST/Typed.hs
index 26f7b5c..b12b30a 100644
--- a/ast/CC/AST/Typed.hs
+++ b/ast/CC/AST/Typed.hs
@@ -25,6 +25,7 @@ data TypeScheme = TypeScheme [Int] Type
   deriving (Show, Read)
 
 data Expr = Lam Type Occ Expr
+          | Let Occ Expr Expr
           | Call Type Expr Expr
           | Int Int
           | Tup [Expr]
@@ -36,6 +37,7 @@ data Occ = Occ Name Type
 
 exprType :: Expr -> Type
 exprType (Lam typ _ _) = typ
+exprType (Let _ _ body) = exprType body
 exprType (Call typ _ _) = typ
 exprType (Int _) = TInt
 exprType (Tup es) = TTup (map exprType es)
@@ -60,6 +62,12 @@ instance Pretty Expr where
         precParens p 2 $
             "(\\(" ++ n ++ " :: " ++ pretty t ++ ") -> "
                 ++ prettyPrec 2 e ++ ") :: " ++ pretty ty
+    prettyPrec p (Let (Occ n t) rhs e) =
+        precParens p 2 $
+            "let (" ++ n ++ " :: " ++ pretty t ++ ") = " ++ pretty rhs ++ " " ++
+                (case e of
+                   Let _ _ _ -> pretty e
+                   _ -> "in " ++ pretty e)
     prettyPrec p (Call ty e1 e2) =
         precParens p 2 $
             prettyPrec 3 e1 ++ " " ++ prettyPrec 3 e2 ++ " :: " ++ pretty ty
diff --git a/parser/CC/Parser.hs b/parser/CC/Parser.hs
index bbabe3c..2d2c4b7 100644
--- a/parser/CC/Parser.hs
+++ b/parser/CC/Parser.hs
@@ -64,7 +64,7 @@ pParenType = do
         _ -> return (TTup tys)
 
 pExpr :: Parser Expr
-pExpr = label (pCall <|> pLam) "expression"
+pExpr = label (pLam <|> pLet <|> pCall) "expression"
   where
     pCall = do
         atoms <- many1 pExprAtom
@@ -90,6 +90,31 @@ pExpr = label (pCall <|> pLam) "expression"
         p2 <- getPosition
         return (Lam (SourceRange p1 p2) names body)
 
+    pLet = do
+        p1 <- try $ do
+            whitespace
+            p <- getPosition
+            void (string "let")
+            return p
+        afterKeyword p1
+      where
+        afterKeyword p1 = do
+            whitespace1
+            lhs <- pName0
+            symbol "="
+            rhs <- pExpr
+            let fullRange rest = mergeRange (SourceRange p1 p1) (range rest)
+            choice [ do p1' <- try $ do
+                            whitespace
+                            p1' <- getPosition
+                            void (string "let")
+                            return p1'
+                        rest <- afterKeyword p1'
+                        return (Let (fullRange rest) lhs rhs rest)
+                   , do symbol "in"
+                        body <- pExpr
+                        return (Let (fullRange body) lhs rhs body) ]
+
 pExprAtom :: Parser Expr
 pExprAtom =
     choice [ uncurry (flip Int) <$> pInt
@@ -119,11 +144,15 @@ pInt = try (whitespace >> pInt0)
 pName0 :: Parser (Name, SourceRange)
 pName0 = do
     p1 <- getPosition
-    c <- pWordFirstChar
-    cs <- many pWordMidChar
+    s <- try $ do
+        c <- pWordFirstChar
+        cs <- many pWordMidChar
+        let s = c : cs
+        guard (s `notElem` ["let", "in"])
+        return s
     p2 <- getPosition
     notFollowedBy pWordMidChar
-    return (c : cs, SourceRange p1 p2)
+    return (s, SourceRange p1 p2)
 
 pWordFirstChar :: Parser Char
 pWordFirstChar = letter <|> oneOf "_$#!"
@@ -146,8 +175,9 @@ token s = try (whitespace >> void (string s))
 emptyLines :: Parser ()
 emptyLines = (try (whitespace >> newline) >> emptyLines) <|> try (whitespace >> eof) <|> return ()
 
-whitespace :: Parser ()
-whitespace = void (many (void (char ' ') <|> void (try (string "\n "))))
+whitespace, whitespace1 :: Parser ()
+whitespace  = void (many  (void (char ' ') <|> void (try (string "\n "))))
+whitespace1 = void (many1 (void (char ' ') <|> void (try (string "\n "))))
 
 getPosition :: Parser SourcePos
 getPosition = do
diff --git a/typecheck/CC/Typecheck.hs b/typecheck/CC/Typecheck.hs
index 8678771..8803a62 100644
--- a/typecheck/CC/Typecheck.hs
+++ b/typecheck/CC/Typecheck.hs
@@ -87,6 +87,8 @@ instance Substitute Env where
 instance Substitute T.Expr where
     theta >>! T.Lam ty (T.Occ name ty2) body =
         T.Lam (theta >>! ty) (T.Occ name (theta >>! ty2)) (theta >>! body)
+    theta >>! T.Let (T.Occ name ty) rhs body =
+        T.Let (T.Occ name (theta >>! ty)) (theta >>! rhs) (theta >>! body)
     theta >>! T.Call ty e1 e2 =
         T.Call (theta >>! ty) (theta >>! e1) (theta >>! e2)
     _     >>! expr@(T.Int _) = expr
@@ -152,6 +154,12 @@ infer env expr = case expr of
         let argType = theta >>! argVar
         return (theta, T.Lam (T.TFun argType (T.exprType body'))
                              (T.Occ arg argType) body')
+    S.Let _ (name, _) rhs body -> do
+        (theta1, rhs') <- infer env rhs
+        let varType = T.exprType rhs'
+        let augEnv = envAdd name (T.TypeScheme [] varType) env
+        (theta2, body') <- infer augEnv body
+        return (theta2 <> theta1, T.Let (T.Occ name varType) rhs' body')
     S.Call sr func arg -> do
         (theta1, func') <- infer env func
         (theta2, arg') <- infer (theta1 >>! env) arg