Support the use of do as a connective in UCS

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")
@@ -499,4 +517,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, curPrec)
+  val `do` = Keyword("do", thenPrec, curPrec)
+  
   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 =>

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/ucs/syntax/Do.mls

@@ -0,0 +1,94 @@
+:js
+
+import "../../../mlscript-compile/Option.mls"
+
+open Option
+
+// Using `do` will not cause match errors.
+
+let x = false
+//│ x = false
+
+if x do
+  print("executed")
+  set x = false
+
+x
+//│ = false
+
+if (not of x) do
+  print("executed")
+  set x = false
+//│ > executed
+
+x
+//│ = false
+
+// Completely using `do`
+// =====================
+
+fun f(y) =
+  let x = Some(y)
+  if x is
+    Some(0) do set x = None
+    Some(v) and v % 2 == 0 do set x = Some(v / 2)
+  x
+  
+f(0)
+//│ = None { class: [class None] }
+
+f(42)
+//│ = Some { value: 21 }
+
+f(41)
+//│ = Some { value: 41 }
+
+
+// Mix using `then` and `do`
+// =========================
+
+:e
+fun g(y) =
+  let x = Some(y)
+  if x is
+    Some(0) do set x = None
+    Some(v) and v % 2 == 0 then set x = Some(v / 2)
+  x
+//│ ╔══[ERROR] Mixed use of `do` and `then` in the `if` expression.
+//│ ╟── Keyword `then` is used here.
+//│ ║  l.55: 	    Some(v) and v % 2 == 0 then set x = Some(v / 2)
+//│ ║        	    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+//│ ╟── Keyword `do` is used here.
+//│ ║  l.54: 	    Some(0) do set x = None
+//│ ╙──      	    ^^^^^^^^^^^^^^^^^^^^^^^
+
+g(0)
+//│ = None { class: [class None] }
+
+g(42)
+//│ = Some { value: 21 }
+
+g(41)
+//│ = Some { value: 41 }
+
+
+// Completely using `then`
+// =======================
+
+
+fun h(y) =
+  let x = Some(y)
+  if x is
+    Some(0) then set x = None
+    Some(v) and v % 2 == 0 then set x = Some(v / 2)
+  x
+  
+h(0)
+//│ = None { class: [class None] }
+
+h(42)
+//│ = Some { value: 21 }
+
+:re
+h(41)
+//│ ═══[RUNTIME ERROR] Error: match error