Support the use of do as a connective in UCS (hkust-taco#246)

Co-authored-by: Lionel Parreaux <lionel.parreaux@gmail.com>

Author: chengluyu

hkmc2/shared/src/main/scala/hkmc2/semantics/Desugarer.scala

@@ -6,8 +6,8 @@ import mlscript.utils.*, shorthands.*
 import Message.MessageContext
 import utils.TraceLogger
 import hkmc2.syntax.Literal
-import Keyword.{as, and, `else`, is, let, `then`}
-import collection.mutable.HashMap
+import Keyword.{as, and, `do`, `else`, is, let, `then`}
+import collection.mutable.{HashMap, SortedSet}
 import Elaborator.{ctx, Ctxl}
 import ucs.DesugaringBase
 
@@ -16,30 +16,52 @@ object Desugarer:
     infix def unapply(tree: Tree): Opt[(Tree, Tree)] = tree match
       case InfixApp(lhs, `op`, rhs) => S((lhs, rhs))
       case _ => N
-
-  /** An extractor that accepts either `A and B` or `A then B`. */
-  object `~>`:
-    infix def unapply(tree: Tree): Opt[(Tree, Tree \/ Tree)] = tree match
-      case lhs and rhs => S((lhs, L(rhs)))
-      case lhs `then` rhs => S((lhs, R(rhs)))
-      case _ => N
-
+  
   class ScrutineeData:
     val classes: HashMap[ClassSymbol, List[BlockLocalSymbol]] = HashMap.empty
     val tupleLead: HashMap[Int, BlockLocalSymbol] = HashMap.empty
     val tupleLast: HashMap[Int, BlockLocalSymbol] = HashMap.empty
 end Desugarer
 
-class Desugarer(tl: TraceLogger, val elaborator: Elaborator)
+class Desugarer(val elaborator: Elaborator)
     (using raise: Raise, state: Elaborator.State, c: Elaborator.Ctx) extends DesugaringBase:
   import Desugarer.*
   import Elaborator.Ctx
-  import elaborator.term
-  import tl.*
+  import elaborator.term, elaborator.tl.*
+  
+  given Ordering[Loc] = Ordering.by: loc =>
+    (loc.spanStart, loc.spanEnd)
+  
+  /** Keep track of the locations where `do` and `then` are used as connectives. */
+  private val kwLocSets = (SortedSet.empty[Loc], SortedSet.empty[Loc])
+  
+  private def reportInconsistentConnectives(kw: Keyword, kwLoc: Opt[Loc]): Unit =
+    log(kwLocSets)
+    (kwLocSets._1.headOption, kwLocSets._2.headOption) match
+      case (Some(doLoc), Some(thenLoc)) =>
+        raise(ErrorReport(
+          msg"Mixed use of `do` and `then` in the `${kw.name}` expression." -> kwLoc
+            :: msg"Keyword `then` is used here." -> S(thenLoc)
+            :: msg"Keyword `do` is used here." -> S(doLoc) :: Nil
+        ))
+      case _ => ()
+  
+  private def topmostDefault: Split =
+    if kwLocSets._1.nonEmpty then Split.Else(Term.Lit(UnitLit(true))) else Split.End
+  
+  /** An extractor that accepts either `A and B`, `A then B`, and `A do B`. It
+   *  also keeps track of the usage of `then` and `do`.
+   */
+  object `~>`:
+    infix def unapply(tree: Tree): Opt[(Tree, Tree \/ Tree)] = tree match
+      case lhs and rhs => S((lhs, L(rhs)))
+      case lhs `then` rhs => kwLocSets._2 ++= tree.toLoc; S((lhs, R(rhs)))
+      case lhs `do` rhs => kwLocSets._1 ++= tree.toLoc; S((lhs, R(rhs)))
+      case _ => N
 
   // We're working on composing continuations in the UCS translation.
   // The type of continuation is `Split => Ctx => Split`.
-  // The first parameter represents the fallback split, which does not have
+  // The first parameter represents the "backup" split, which does not have
   // access to the bindings in the current match. The second parameter
   // represents the context with bindings in the current match.
 
@@ -82,10 +104,6 @@ class Desugarer(tl: TraceLogger, val elaborator: Elaborator)
 
   def default: Split => Sequel = split => _ => split
 
-  /** Desugar UCS shorthands. */
-  def shorthands(tree: Tree): Sequel = termSplitShorthands(tree, identity):
-    Split.default(Term.Lit(Tree.BoolLit(false)))
-
   private def termSplitShorthands(tree: Tree, finish: Term => Term): Split => Sequel = tree match
     case Block(branches) => branches match
       case Nil => lastWords("encountered empty block")
@@ -166,6 +184,12 @@ class Desugarer(tl: TraceLogger, val elaborator: Elaborator)
           val sym = VarSymbol(ident)
           val fallbackCtx = ctx + (ident.name -> sym)
           Split.Let(sym, term(termTree)(using ctx), elabFallback(fallback)(fallbackCtx)).withLocOf(t)
+        case Modified(Keyword.`do`, doLoc, computation) => fallback => ctx => trace(
+          pre = s"termSplit: do $computation",
+          post = (res: Split) => s"termSplit: else >>> $res"
+        ):
+          val sym = TempSymbol(N, "doTemp")
+          Split.Let(sym, term(computation)(using ctx), elabFallback(fallback)(ctx)).withLocOf(t)
         case Modified(Keyword.`else`, elsLoc, default) => fallback => ctx => trace(
           pre = s"termSplit: else $default",
           post = (res: Split) => s"termSplit: else >>> $res"
@@ -241,6 +265,12 @@ class Desugarer(tl: TraceLogger, val elaborator: Elaborator)
               val sym = VarSymbol(ident)
               val fallbackCtx = ctx + (ident.name -> sym)
               Split.Let(sym, term(termTree)(using ctx), elabFallback(fallbackCtx))
+          case (Tree.Empty(), Modified(Keyword.`do`, doLoc, computation)) => ctx => trace(
+            pre = s"termSplit: do $computation",
+            post = (res: Split) => s"termSplit: else >>> $res"
+          ):
+            val sym = TempSymbol(N, "doTemp")
+            Split.Let(sym, term(computation)(using ctx), elabFallback(ctx))
           case (Tree.Empty(), Modified(Keyword.`else`, elsLoc, default)) => ctx =>
             // TODO: report `rest` as unreachable
             Split.default(term(default)(using ctx))
@@ -322,6 +352,12 @@ class Desugarer(tl: TraceLogger, val elaborator: Elaborator)
           val sym = VarSymbol(ident)
           val fallbackCtx = ctx + (ident.name -> sym)
           Split.Let(sym, term(termTree)(using ctx), elabFallback(backup)(fallbackCtx))
+      case Modified(Keyword.`do`, doLoc, computation) => fallback => ctx => trace(
+        pre = s"patternSplit (do) <<< $computation",
+        post = (res: Split) => s"patternSplit: else >>> $res"
+      ):
+        val sym = TempSymbol(N, "doTemp")
+        Split.Let(sym, term(computation)(using ctx), elabFallback(fallback)(ctx))
       case Modified(Keyword.`else`, elsLoc, body) => backup => ctx => trace(
         pre = s"patternSplit (else) <<< $tree",
         post = (res: Split) => s"patternSplit (else) >>> ${res.showDbg}"
@@ -499,4 +535,20 @@ class Desugarer(tl: TraceLogger, val elaborator: Elaborator)
     ):
       val innermostSplit = subMatches(rest, sequel)(fallback)
       expandMatch(scrutinee, pattern, innermostSplit)(fallback)
+  
+  /** Desugar `case` expressions. */
+  def apply(tree: Case, scrut: VarSymbol)(using Ctx): Split =
+    val topmost = patternSplit(tree.branches, scrut)(Split.End)(ctx)
+    reportInconsistentConnectives(Keyword.`case`, tree.kwLoc)
+    topmost ++ topmostDefault
+  
+  /** Desugar `if` and `while` expressions. */
+  def apply(tree: IfLike)(using Ctx): Split =
+    val topmost = termSplit(tree.split, identity)(Split.End)(ctx)
+    reportInconsistentConnectives(tree.kw, tree.kwLoc)
+    topmost ++ topmostDefault
+  
+  /** Desugar `is` and `and` shorthands. */
+  def apply(tree: InfixApp)(using Ctx): Split =
+    termSplitShorthands(tree, identity)(Split.default(Term.Lit(Tree.BoolLit(false))))(ctx)
 end Desugarer

hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala

@@ -269,8 +269,8 @@ extends Importer:
         Term.Lam(syms, term(rhs)(using nestCtx))
     case InfixApp(lhs, Keyword.`:`, rhs) =>
       Term.Asc(term(lhs), term(rhs))
-    case InfixApp(lhs, Keyword.`is` | Keyword.`and`, rhs) =>
-      val des = new Desugarer(tl, this).shorthands(tree)(ctx)
+    case tree @ InfixApp(lhs, Keyword.`is` | Keyword.`and`, rhs) =>
+      val des = new Desugarer(this)(tree)
       val nor = new ucs.Normalization(tl)(des)
       Term.IfLike(Keyword.`if`, des)(nor)
     case App(Ident("|"), Tree.Tup(lhs :: rhs :: Nil)) =>
@@ -348,8 +348,8 @@ extends Importer:
         Term.New(cls(c, inAppPrefix = false), Nil).withLocOf(tree)
       // case _ =>
       //   raise(ErrorReport(msg"Illegal new expression." -> tree.toLoc :: Nil))
-    case Tree.IfLike(kw, split) =>
-      val desugared = new Desugarer(tl, this).termSplit(split, identity)(Split.End)(ctx)
+    case tree @ Tree.IfLike(kw, _, split) =>
+      val desugared = new Desugarer(this)(tree)
       scoped("ucs:desugared"):
         log(s"Desugared:\n${Split.display(desugared)}")
       val normalized = new ucs.Normalization(tl)(desugared)
@@ -358,10 +358,9 @@ extends Importer:
       Term.IfLike(kw, desugared)(normalized)
     case Tree.Quoted(body) => Term.Quoted(term(body))
     case Tree.Unquoted(body) => Term.Unquoted(term(body))
-    case Tree.Case(branches) =>
+    case tree @ Tree.Case(_, branches) =>
       val scrut = VarSymbol(Ident("caseScrut"))
-      val desugarer = new Desugarer(tl, this)
-      val des = desugarer.patternSplit(branches, scrut)(Split.End)(ctx)
+      val des = new Desugarer(this)(tree, scrut)
       scoped("ucs:desugared"):
         log(s"Desugared:\n${Split.display(des)}")
       val nor = new ucs.Normalization(tl)(des)

hkmc2/shared/src/main/scala/hkmc2/syntax/Keyword.scala

@@ -63,7 +63,11 @@ object Keyword:
 
   val `if` = Keyword("if", N, nextPrec)
   val `while` = Keyword("while", N, curPrec)
-  val `then` = Keyword("then", nextPrec, curPrec)
+  
+  val thenPrec = nextPrec
+  val `then` = Keyword("then", thenPrec, thenPrec)
+  val `do` = Keyword("do", thenPrec, thenPrec)
+  
   val `else` = Keyword("else", nextPrec, curPrec)
   val `case` = Keyword("case", N, N)
   val `fun` = Keyword("fun", N, N)
@@ -81,7 +85,6 @@ object Keyword:
   val `in` = Keyword("in", curPrec, curPrec)
   val `out` = Keyword("out", N, curPrec)
   val `set` = Keyword("set", N, curPrec)
-  val `do` = Keyword("do", N, N)
   val `declare` = Keyword("declare", N, N)
   val `trait` = Keyword("trait", N, N)
   val `mixin` = Keyword("mixin", N, N)
@@ -125,7 +128,7 @@ object Keyword:
     `abstract`, mut, virtual, `override`, declare, public, `private`)
 
   type Infix = `and`.type | `or`.type | `then`.type | `else`.type | `is`.type | `:`.type | `->`.type |
-    `=>`.type | `extends`.type | `restricts`.type | `as`.type
+    `=>`.type | `extends`.type | `restricts`.type | `as`.type | `do`.type
 
   type Ellipsis = `...`.type | `..`.type
 

hkmc2/shared/src/main/scala/hkmc2/syntax/ParseRule.scala

@@ -183,12 +183,12 @@ class ParseRules(using State):
             val items = split match
               case Block(stmts) => stmts.appended(clause)
               case _ => split :: clause :: Nil
-            IfLike(kw, Block(items))
-          case (split, N) => IfLike(kw, split)
+            IfLike(kw, N/* TODO */, Block(items))
+          case (split, N) => IfLike(kw, N/* TODO */, split)
         ,
         Blk(
           ParseRule(s"'${kw.name}' block")(End(()))
-        ) { case (body, _) => IfLike(kw, body) }
+        ) { case (body, _) => IfLike(kw, N/* TODO */, body) }
       )
 
   def typeAliasLike(kw: Keyword, kind: TypeDefKind): Kw[TypeDef] =
@@ -258,7 +258,7 @@ class ParseRules(using State):
     ,
     Kw(`case`):
       ParseRule("`case` keyword")(
-        Blk(ParseRule("`case` branches")(End(())))((body, _: Unit) => Case(body))
+        Blk(ParseRule("`case` branches")(End(())))((body, _: Unit) => Case(N/* TODO */, body))
       )
     ,
     Kw(`region`):
@@ -358,6 +358,7 @@ class ParseRules(using State):
     genInfixRule(`:`, (rhs, _: Unit) => lhs => InfixApp(lhs, `:`, rhs)),
     genInfixRule(`extends`, (rhs, _: Unit) => lhs => InfixApp(lhs, `extends`, rhs)),
     genInfixRule(`restricts`, (rhs, _: Unit) => lhs => InfixApp(lhs, `restricts`, rhs)),
+    genInfixRule(`do`, (rhs, _: Unit) => lhs => InfixApp(lhs, `do`, rhs)),
   )
 
 end ParseRules

hkmc2/shared/src/main/scala/hkmc2/syntax/Parser.scala

@@ -547,7 +547,7 @@ abstract class Parser(
               val ele = simpleExprImpl(prec)
               term match
                 case InfixApp(lhs, Keyword.`then`, rhs) =>
-                  Quoted(IfLike(Keyword.`if`, Block(
+                  Quoted(IfLike(Keyword.`if`, S(l0), Block(
                     InfixApp(Unquoted(lhs), Keyword.`then`, Unquoted(rhs)) :: Modified(Keyword.`else`, N, Unquoted(ele)) :: Nil
                   )))
                 case tk =>
@@ -900,6 +900,12 @@ abstract class Parser(
       case (NEWLINE, _) :: (KEYWORD(kw), _) :: _
       if kw.canStartInfixOnNewLine && kw.leftPrecOrMin > prec
       && infixRules.kwAlts.contains(kw.name)
+      && (kw isnt Keyword.`do`) // This is to avoid the following case:
+        //  ```
+        //  0 then "null"
+        //  do console.log("non-null")
+        //  ```
+        // Otherwise, `do` will be parsed as an infix operator
       =>
         consume
         exprCont(acc, prec, allowNewlines = false)

hkmc2/shared/src/main/scala/hkmc2/syntax/Tree.scala

@@ -66,10 +66,10 @@ enum Tree extends AutoLocated:
   case Sel(prefix: Tree, name: Ident)
   case InfixApp(lhs: Tree, kw: Keyword.Infix, rhs: Tree)
   case New(body: Tree)
-  case IfLike(kw: Keyword.`if`.type | Keyword.`while`.type, split: Tree)
+  case IfLike(kw: Keyword.`if`.type | Keyword.`while`.type, kwLoc: Opt[Loc], split: Tree)
   @deprecated("Use If instead", "hkmc2-ucs")
   case IfElse(cond: Tree, alt: Tree)
-  case Case(branches: Tree)
+  case Case(kwLoc: Opt[Loc], branches: Tree)
   case Region(name: Tree, body: Tree)
   case RegRef(reg: Tree, value: Tree)
   case Effectful(eff: Tree, body: Tree)
@@ -95,9 +95,9 @@ enum Tree extends AutoLocated:
     case InfixApp(lhs, _, rhs) => Ls(lhs, rhs)
     case TermDef(k, head, rhs) => head :: rhs.toList
     case New(body) => body :: Nil
-    case IfLike(_, split) => split :: Nil
+    case IfLike(_, _, split) => split :: Nil
     case IfElse(cond, alt) => cond :: alt :: Nil
-    case Case(bs) => Ls(bs)
+    case Case(_, bs) => Ls(bs)
     case Region(name, body) => name :: body :: Nil
     case RegRef(reg, value) => reg :: value :: Nil
     case Effectful(eff, body) => eff :: body :: Nil
@@ -133,9 +133,9 @@ enum Tree extends AutoLocated:
     case Sel(prefix, name) => "selection"
     case InfixApp(lhs, kw, rhs) => "infix operation"
     case New(body) => "new"
-    case IfLike(Keyword.`if`, split) => "if expression"
-    case IfLike(Keyword.`while`, split) => "while expression"
-    case Case(branches) => "case"
+    case IfLike(Keyword.`if`, _, split) => "if expression"
+    case IfLike(Keyword.`while`, _, split) => "while expression"
+    case Case(_, branches) => "case"
     case Region(name, body) => "region"
     case RegRef(reg, value) => "region reference"
     case Effectful(eff, body) => "effectful"

hkmc2/shared/src/test/mlscript-compile/Predef.mls

@@ -30,15 +30,14 @@ class MatchResult(captures)
 class MatchFailure(errors)
 
 fun checkArgs(functionName, expected, isUB, got) =
-  if got < expected || isUB && got > expected then
+  if got < expected || isUB && got > expected do
     let name = if functionName.length > 0 then " '" + functionName + "'" else ""
     throw globalThis.Error("Function" + name + " expected " + expected + " arguments but got " + got)
     // TODO
     // throw globalThis.Error("Function" + name + " expected "
     //   + expected
     //   + (if isUB then "" else " at least")
     //   + " arguments but got " + got)
-  else ()
 
 module TraceLogger with
   mut val enabled = false

hkmc2/shared/src/test/mlscript-compile/apps/Accounting.mls

@@ -30,9 +30,8 @@ class Line(val name: Str, val proj: Project, val starting_balance: Num, val isMa
   fun expense(amt) =
     set balance = balance -. amt
   fun mustBeEmpty() =
-    if balance > 10_000 then
+    if balance > 10_000 do
       warnings.push of "> **❗️** Unspent balance of " ~ name ~ ": `" ~ display(balance) ~ "`"
-    else () // TODO allow omitting else branch
 
 val lines = []
 

hkmc2/shared/src/test/mlscript/basics/OpBlocks.mls

@@ -86,6 +86,7 @@ fun f(x) = if x
 //│       Ident of "x"
 //│   rhs = S of IfLike:
 //│     kw = keyword 'if'
+//│     kwLoc = N
 //│     split = App:
 //│       lhs = Ident of "x"
 //│       rhs = OpBlock of Ls of 
@@ -120,10 +121,10 @@ fun f(x) = if x
   > 0 then "a"
   is 0 then "b"
 //│ ╔══[PARSE ERROR] Expect an operator instead of 'is' keyword
-//│ ║  l.121: 	  is 0 then "b"
+//│ ║  l.122: 	  is 0 then "b"
 //│ ╙──       	  ^^
 //│ ╔══[PARSE ERROR] Unexpected 'is' keyword here
-//│ ║  l.121: 	  is 0 then "b"
+//│ ║  l.122: 	  is 0 then "b"
 //│ ╙──       	  ^^
 //│ ═══[ERROR] Unrecognized operator branch.
 
@@ -135,11 +136,11 @@ fun f(x) = if x
   foo(A) then a
   bar(B) then b
 //│ ╔══[ERROR] Unrecognized term split (juxtaposition).
-//│ ║  l.134: 	fun f(x) = if x
+//│ ║  l.135: 	fun f(x) = if x
 //│ ║         	              ^
-//│ ║  l.135: 	  foo(A) then a
+//│ ║  l.136: 	  foo(A) then a
 //│ ║         	^^^^^^^^^^^^^^^
-//│ ║  l.136: 	  bar(B) then b
+//│ ║  l.137: 	  bar(B) then b
 //│ ╙──       	^^^^^^^^^^^^^^^
 
 
@@ -148,10 +149,10 @@ fun f(x) = if x
   is 0 then "a"
   is 1 then "b"
 //│ ╔══[PARSE ERROR] Expected start of statement in this position; found 'is' keyword instead
-//│ ║  l.149: 	  is 1 then "b"
+//│ ║  l.150: 	  is 1 then "b"
 //│ ╙──       	  ^^
 //│ ╔══[PARSE ERROR] Expected end of input; found literal instead
-//│ ║  l.149: 	  is 1 then "b"
+//│ ║  l.150: 	  is 1 then "b"
 //│ ╙──       	     ^
 //│ = [Function: f]