From b844d0b67691c4c91893cb2819a1fc2ac36553bc Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Tue, 29 Oct 2024 20:20:05 -0700
Subject: [PATCH] parsing and desugaring of `do` blocks. (Some inference issues
 remain.)

---
 newt/TypeClass.newt        | 26 +++++++++++++++++++++-----
 src/Lib/Elab.idr           | 24 ++++++++++++++++++------
 src/Lib/Parser.idr         | 18 ++++++++++++++++++
 src/Lib/Syntax.idr         | 16 +++++++++++-----
 src/Lib/Tokenizer.idr      |  4 ++--
 tests/black/TypeClass.newt | 26 +++++++++++++++++++++-----
 6 files changed, 91 insertions(+), 23 deletions(-)

diff --git a/newt/TypeClass.newt b/newt/TypeClass.newt
index d83daa0..a0890b3 100644
--- a/newt/TypeClass.newt
+++ b/newt/TypeClass.newt
@@ -3,6 +3,7 @@ module TypeClass
 data Monad : (U -> U) -> U where
   MkMonad : { M : U -> U } ->
             (bind : {A B : U} -> (M A) -> (A -> M B) -> M B) ->
+            (pure : {A : U} -> A -> M A) ->
             Monad M
 
 data Maybe : U -> U where
@@ -18,21 +19,36 @@ bindEither (Left a) amb = Left a
 bindEither (Right b) amb = amb b
 
 EitherMonad : {A : U} -> Monad (Either A)
-EitherMonad = MkMonad {Either A} bindEither
+EitherMonad = MkMonad {Either A} bindEither Right
 
 bindMaybe : {A B : U} -> Maybe A -> (A -> Maybe B) -> Maybe B
 bindMaybe Nothing amb = Nothing
 bindMaybe (Just a) amb = amb a
 
 MaybeMonad : Monad Maybe
-MaybeMonad = MkMonad bindMaybe
+MaybeMonad = MkMonad bindMaybe Just
 
 _>>=_ : {a b : U} -> {m : U -> U} -> {{x : Monad m}} -> (m a) -> (a -> m b) -> m b
-_>>=_ {a} {b} {m} {{MkMonad bind'}}  ma amb = bind' {a} {b} ma amb
+_>>=_ {a} {b} {m} {{MkMonad bind' _}}  ma amb = bind' {a} {b} ma amb
 
-infixl 1 _>>=_
+pure : {a : U} {m : U -> U} {{_ : Monad m}} -> a -> m a
+pure {_} {_} {{MkMonad _ pure'}} a = pure' a
+
+infixl 1 _>>=_ _>>_
+
+_>>_ : {a b : U} -> {m : U -> U} -> {{x : Monad m}} -> m a -> m b -> m b
 
 ptype Int
 
+test : Maybe Int
+test = pure 10
+
 foo : Int -> Maybe Int
-foo x = (Just x) >>= (\ x => Just 10)
+foo x = Just 42 >> Just x >>= (\ x => pure {_} {Maybe} 10)
+
+bar : Int -> Maybe Int
+bar x = do
+  let y = x
+  z <- Just x
+  -- This is not sorting out the Maybe...
+  pure {_} {_} z
diff --git a/src/Lib/Elab.idr b/src/Lib/Elab.idr
index 26118a9..61298eb 100644
--- a/src/Lib/Elab.idr
+++ b/src/Lib/Elab.idr
@@ -257,10 +257,14 @@ parameters (ctx: Context)
     debug "env \{show ctx.env}"
     debug "types \{show $ ctx.types}"
     case (t',u') of
+
+      -- flex/flex
       (VMeta fc k sp,  VMeta fc' k' sp' ) =>
         if k == k' then unifySpine l (k == k') sp sp'
         -- TODO, might want to try the other way, too.
-        else solve l k sp (VMeta fc' k' sp') >> pure neutral
+        else if length sp < length sp'
+          then solve l k' sp' (VMeta fc k sp) >> pure neutral
+          else solve l k sp (VMeta fc' k' sp') >> pure neutral
       (t,           VMeta fc' i' sp') => solve l i' sp' t >> pure neutral
       (VMeta fc i sp, t'           ) => solve l i sp t' >> pure neutral
       (VPi fc _ _ a b, VPi fc' _ _ a' b') => [| unify l a a' <+> unify (S l) !(b $$ VVar emptyFC l [<]) !(b' $$ VVar emptyFC l [<]) |]
@@ -612,7 +616,8 @@ buildCase ctx prob scnm scty (dcName, arity, ty) = do
     makeConstr [] (pat :: pats) = ?extra_patterns
     makeConstr ((MkBind nm Implicit x) :: xs) [] = (nm, PatWild emptyFC Implicit) :: makeConstr xs []
     makeConstr ((MkBind nm Auto x) :: xs) [] = (nm, PatWild emptyFC Auto) :: makeConstr xs []
-    makeConstr ((MkBind nm Explicit x) :: xs) [] = ?extra_binders_2
+    -- FIXME need a proper error, but requires wiring M three levels down
+    makeConstr ((MkBind nm Explicit x) :: xs) [] = ?insufficient_patterns
     makeConstr ((MkBind nm Explicit x) :: xs) (pat :: pats) =
       if getIcit pat == Explicit
         then (nm, pat) :: makeConstr xs pats
@@ -822,7 +827,18 @@ showDef : Context -> List String -> Nat -> Val -> M String
 showDef ctx names n v@(VVar _ n' [<]) =  if n == n' then pure "" else pure "= \{pprint names !(quote ctx.lvl v)}"
 showDef ctx names n v = pure "= \{pprint names !(quote ctx.lvl v)}"
 
+
+undo : List DoStmt -> M Raw
+undo [] = error emptyFC "do block must end in expression"
+undo ((DoExpr fc tm) :: Nil) = pure tm
+-- TODO decide if we want to use >> or just >>=
+undo ((DoExpr fc tm) :: xs) = pure $ RApp fc (RApp fc (RVar fc "_>>=_") tm Explicit) (RLam fc "_" Explicit !(undo xs)) Explicit
+-- undo ((DoExpr fc tm) :: xs) = pure $ RApp fc (RApp fc (RVar fc "_>>_") tm Explicit) !(undo xs) Explicit
+undo ((DoLet fc nm tm) :: xs) = RLet fc nm (RImplicit fc) tm <$> undo xs
+undo ((DoArrow fc nm tm) :: xs) = pure $ RApp fc (RApp fc (RVar fc "_>>=_") tm Explicit) (RLam fc nm Explicit !(undo xs)) Explicit
+
 check ctx tm ty = case (tm, !(forceType ty)) of
+  (RDo fc stmts, ty) => check ctx !(undo stmts) ty
   (RCase fc rsc alts, ty) => do
     (sc, scty) <- infer ctx rsc
     scty <- forceMeta scty
@@ -989,7 +1005,3 @@ infer ctx (RLit fc (LString str)) = pure (Lit fc (LString str), !(primType fc "S
 infer ctx (RLit fc (LInt i)) = pure (Lit fc (LInt i), !(primType fc "Int"))
 
 infer ctx tm = error (getFC tm) "Implement infer \{show tm}"
-
--- The idea here is to insert a hole for a parse error
--- but the parser doesn't emit this yet.
--- infer ctx (RParseError str) = ?todo_insert_meta
diff --git a/src/Lib/Parser.idr b/src/Lib/Parser.idr
index 1a27f42..322393d 100644
--- a/src/Lib/Parser.idr
+++ b/src/Lib/Parser.idr
@@ -203,10 +203,28 @@ caseExpr = do
   alts <- startBlock $ someSame $ caseAlt
   pure $ RCase fc sc alts
 
+doArrow : Parser DoStmt
+doArrow = do
+  fc <- getPos
+  name <- try $ ident <* keyword "<-"
+  expr <- term
+  pure $ DoArrow fc name expr
+
+doStmt : Parser DoStmt
+doStmt
+  = DoArrow <$> getPos <*> (try $ ident <* keyword "<-") <*> term
+  <|> DoLet <$> getPos <* keyword "let" <*> ident <* keyword "=" <*> term
+  <|> DoExpr <$> getPos <*> term
+
+doExpr : Parser Raw
+doExpr = RDo <$> getPos <* keyword "do" <*> (startBlock $ someSame doStmt)
+
+
 -- This hits an idris codegen bug if parseOp is last and Lazy
 term =  caseExpr
     <|> letExpr
     <|> lamExpr
+    <|> doExpr
     <|> parseOp
 
 varname : Parser String
diff --git a/src/Lib/Syntax.idr b/src/Lib/Syntax.idr
index df95d2c..556c488 100644
--- a/src/Lib/Syntax.idr
+++ b/src/Lib/Syntax.idr
@@ -59,6 +59,12 @@ record Clause where
 public export
 data RCaseAlt = MkAlt Raw Raw
 
+public export
+data DoStmt : Type where
+  DoExpr : (fc : FC) -> Raw -> DoStmt
+  DoLet : (fc : FC) -> String -> Raw -> DoStmt
+  DoArrow : (fc: FC) -> String -> Raw -> DoStmt
+
 data Raw : Type where
   RVar : (fc : FC) -> (nm : Name) -> Raw
   RLam : (fc : FC) -> (nm : String) -> (icit : Icit) -> (ty : Raw) -> Raw
@@ -71,8 +77,8 @@ data Raw : Type where
   RCase : (fc : FC) -> (scrut : Raw) -> (alts : List RCaseAlt) -> Raw
   RImplicit : (fc : FC) -> Raw
   RHole : (fc : FC) -> Raw
-  -- not used, but intended to allow error recovery
-  RParseError : (fc : FC) -> String -> Raw
+  RDo : (fc : FC) -> List DoStmt -> Raw
+
 
 %name Raw tm
 
@@ -89,7 +95,7 @@ HasFC Raw where
   getFC (RCase fc scrut alts) = fc
   getFC (RImplicit fc) = fc
   getFC (RHole fc) = fc
-  getFC (RParseError fc str) = fc
+  getFC (RDo fc stmts) = fc
 -- derive some stuff - I'd like json, eq, show, ...
 
 
@@ -181,7 +187,7 @@ Show Raw where
   show (RApp _ x y z) = foo [ "App", show x, show y, show z]
   show (RLam _ x i y) = foo [ "Lam", show x, show i, show y]
   show (RCase _ x xs) = foo [ "Case", show x, show xs]
-  show (RParseError _ str) = foo [ "ParseError", "str"]
+  show (RDo _ stmts) = foo [ "DO", "FIXME"]
   show (RU _) = "U"
 
 export
@@ -233,7 +239,7 @@ Pretty Raw where
     asDoc p (RCase _ x xs) = text "TODO - RCase"
     asDoc p (RImplicit _) = text "_"
     asDoc p (RHole _) = text "?"
-    asDoc p (RParseError _ str) = text "ParseError \{str}"
+    asDoc p (RDo _ stmts) = text "TODO - RDo"
 
 export
 Pretty Module where
diff --git a/src/Lib/Tokenizer.idr b/src/Lib/Tokenizer.idr
index 18654df..d4ecfee 100644
--- a/src/Lib/Tokenizer.idr
+++ b/src/Lib/Tokenizer.idr
@@ -5,11 +5,11 @@ import Text.Lexer.Tokenizer
 import Lib.Token
 
 keywords : List String
-keywords = ["let", "in", "where", "case", "of", "data", "U",
+keywords = ["let", "in", "where", "case", "of", "data", "U", "do",
             "ptype", "pfunc", "module", "infixl", "infixr", "infix"]
 
 specialOps : List String
-specialOps = ["->", ":", "=>", ":=", "="]
+specialOps = ["->", ":", "=>", ":=", "=", "<-"]
 
 checkKW : String -> Token Kind
 checkKW s = if elem s keywords then Tok Keyword s else Tok Ident s
diff --git a/tests/black/TypeClass.newt b/tests/black/TypeClass.newt
index d83daa0..21a261b 100644
--- a/tests/black/TypeClass.newt
+++ b/tests/black/TypeClass.newt
@@ -3,6 +3,7 @@ module TypeClass
 data Monad : (U -> U) -> U where
   MkMonad : { M : U -> U } ->
             (bind : {A B : U} -> (M A) -> (A -> M B) -> M B) ->
+            (pure : {A : U} -> A -> M A) ->
             Monad M
 
 data Maybe : U -> U where
@@ -18,21 +19,36 @@ bindEither (Left a) amb = Left a
 bindEither (Right b) amb = amb b
 
 EitherMonad : {A : U} -> Monad (Either A)
-EitherMonad = MkMonad {Either A} bindEither
+EitherMonad = MkMonad {Either A} bindEither Right
 
 bindMaybe : {A B : U} -> Maybe A -> (A -> Maybe B) -> Maybe B
 bindMaybe Nothing amb = Nothing
 bindMaybe (Just a) amb = amb a
 
 MaybeMonad : Monad Maybe
-MaybeMonad = MkMonad bindMaybe
+MaybeMonad = MkMonad bindMaybe Just
 
 _>>=_ : {a b : U} -> {m : U -> U} -> {{x : Monad m}} -> (m a) -> (a -> m b) -> m b
-_>>=_ {a} {b} {m} {{MkMonad bind'}}  ma amb = bind' {a} {b} ma amb
+_>>=_ {a} {b} {m} {{MkMonad bind' _}}  ma amb = bind' {a} {b} ma amb
 
-infixl 1 _>>=_
+pure : {a : U} {m : U -> U} {{_ : Monad m}} -> a -> m a
+pure {_} {_} {{MkMonad _ pure'}} a = pure' a
+
+infixl 1 _>>=_ _>>_
+
+_>>_ : {a b : U} -> {m : U -> U} -> {{x : Monad m}} -> m a -> m b -> m b
 
 ptype Int
 
+test : Maybe Int
+test = pure 10
+
 foo : Int -> Maybe Int
-foo x = (Just x) >>= (\ x => Just 10)
+foo x = Just 42 >> Just x >>= (\ x => pure {_} {Maybe} 10)
+
+bar : Int -> Maybe Int
+bar x = do
+  let y = x
+  z <- Just x
+  -- NOW each of these implicits is a different error, fix them
+  pure {_} {Maybe} z