Implement UCS while expressions

Author: LPTK

hkmc2/shared/src/main/scala/hkmc2/bbml/bbML.scala

@@ -287,7 +287,7 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
       val res = freshVar(using ctx)
       constrain(bodyCtx, sk | res)
       (bodyTy, rhsCtx | res, rhsEff | bodyEff)
-    case Term.If(Split.Cons(Branch(cond, Pattern.LitPat(BoolLit(true)), Split.Else(cons)), Split.Else(alts))) =>
+    case Term.IfLike(Keyword.`if`, Split.Cons(Branch(cond, Pattern.LitPat(BoolLit(true)), Split.Else(cons)), Split.Else(alts))) =>
       val (condTy, condCtx, condEff) = typeCode(cond)
       val (consTy, consCtx, consEff) = typeCode(cons)
       val (altsTy, altsCtx, altsEff) = typeCode(alts)
@@ -373,7 +373,7 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
       given Ctx = nextCtx
       constrain(ascribe(term, skolemize(pt))._2, Bot) // * never generalize terms with effects
       (pt, Bot)
-    case (Term.If(branches), ty) => // * propagate
+    case (Term.IfLike(Keyword.`if`, branches), ty) => // * propagate
       typeSplit(branches, S(ty))
     case (Term.Asc(term, ty), rhs) =>
       ascribe(term, typeType(ty))
@@ -536,7 +536,7 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
       case Term.Asc(term, ty) =>
         val res = typeType(ty)(using ctx)
         ascribe(term, res)
-      case Term.If(branches) => typeSplit(branches, N)
+      case Term.IfLike(Keyword.`if`, branches) => typeSplit(branches, N)
       case Term.Region(sym, body) =>
         val nestCtx = ctx.nextLevel
         given Ctx = nestCtx

hkmc2/shared/src/main/scala/hkmc2/codegen/Block.scala

@@ -32,8 +32,10 @@ sealed abstract class Block extends Product with AutoLocated:
     case Match(scrut, arms, dflt, rst) =>
       arms.flatMap(_._2.definedVars).toSet ++ dflt.toList.flatMap(_.definedVars) ++ rst.definedVars
     case End(_) => Set.empty
+    case Break(_, _) => Set.empty
     case Define(defn, rst) => rst.definedVars
     case TryBlock(sub, fin, rst) => sub.definedVars ++ fin.definedVars ++ rst.definedVars
+    case Label(lbl, bod, rst) => bod.definedVars ++ rst.definedVars
 
   // TODO conserve if no changes
   def mapTail(f: BlockTail => BlockTail): Block = this match
@@ -60,6 +62,11 @@ case class Return(res: Result, implct: Bool) extends BlockTail
 
 case class Throw(exc: Result) extends BlockTail
 
+case class Label(label: Local, body: Block, rest: Block) extends BlockTail
+
+case class Break(label: Local, toBeginning: Bool) extends BlockTail
+
+// TODO: remove this form?
 case class Begin(sub: Block, rest: Block) extends Block with ProductWithTail
 
 case class TryBlock(sub: Block, finallyDo: Block, rest: Block) extends Block with ProductWithTail

hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala

@@ -204,52 +204,69 @@ class Lowering(using TL, Raise, Elaborator.State):
             )
     */
 
-    case iftrm: st.If =>
+    case iftrm: st.IfLike =>
 
-      tl.log(s"If $iftrm")
+      tl.log(s"${iftrm.kw} $iftrm")
+      
+      val isIf = iftrm.kw match
+        case syntax.Keyword.`if` => true
+        case syntax.Keyword.`while` => false
+      val isWhile = !isIf
 
       var usesResTmp = false
       lazy val l =
         usesResTmp = true
         new TempSymbol(summon[Elaborator.State].nextUid, S(t))
 
-      def go(split: Split)(using Subst): Block = split match
+      lazy val lbl =
+        new TempSymbol(summon[Elaborator.State].nextUid, S(t))
+      
+      def go(split: Split, topLevel: Bool)(using Subst): Block = split match
         case Split.Let(sym, trm, tl) =>
           term(trm): r =>
-            Assign(sym, r, go(tl))
+            Assign(sym, r, go(tl, topLevel))
         case Split.Cons(Branch(scrut, pat, tail), restSplit) =>
           subTerm(scrut): sr =>
             tl.log(s"Binding scrut $scrut to $sr ${summon[Subst].map}")
             val cse = pat match
-              case Pattern.LitPat(lit) => Case.Lit(lit) -> go(tail)
+              case Pattern.LitPat(lit) => Case.Lit(lit) -> go(tail, topLevel = false)
               case Pattern.Class(cls, args0, _refined) =>
                 val args = args0.getOrElse(Nil)
                 val clsDefn = cls.defn.getOrElse(die)
                 val clsParams = clsDefn.paramsOpt.getOrElse(Nil)
                 assert(args0.isEmpty || clsParams.length === args.length)
                 def mkArgs(args: Ls[Param -> BlockLocalSymbol])(using Subst): Case -> Block = args match
-                  case Nil => Case.Cls(cls) -> go(tail)
+                  case Nil => Case.Cls(cls) -> go(tail, topLevel = false)
                   case (param, arg) :: args =>
                     // summon[Subst].+(arg -> Value.Ref(new TempSymbol(summon[Elaborator.State].nextUid, N)))
                     // Assign(arg, Select(sr, Tree.Ident("head")), mkArgs(args))
                     val (cse, blk) = mkArgs(args)
                     (cse, Assign(arg, Select(sr, param.sym.id/*FIXME incorrect Ident?*/), blk))
                 mkArgs(clsParams.zip(args))
             Match(sr, cse :: Nil,
-              // elseBranch,
-              S(go(restSplit)),
+              S(go(restSplit, topLevel = true)),
               End()
             )
         case Split.Else(els) =>
-          if k.isInstanceOf[Ret] then term(els)(k)
-          else term(els)(r => Assign(l, r, End()))
+          if k.isInstanceOf[Ret] && isIf then term(els)(k)
+          else
+            term(els): r =>
+              Assign(l, r,
+                if isWhile && !topLevel then Break(lbl, toBeginning = true)
+                // if isWhile then Break(lbl, toBeginning = !topLevel)
+                else End()
+              )
         case Split.Nil =>
           Throw(Instantiate(Value.Ref(Elaborator.Ctx.errorSymbol),
             Value.Lit(syntax.Tree.StrLit("match error")) :: Nil)) // TODO add failed-match scrutinee info
 
-      if k.isInstanceOf[Ret] then go(iftrm.normalized)
-      else Begin(
-          go(iftrm.normalized),
+      if k.isInstanceOf[Ret] && isIf then go(iftrm.normalized, topLevel = true)
+      else
+        val body = if isWhile
+          then Label(lbl, go(iftrm.normalized, topLevel = true), End())
+          else go(iftrm.normalized, topLevel = true)
+        Begin(
+          body,
           if usesResTmp then k(Value.Ref(l))
           else k(Value.Lit(syntax.Tree.UnitLit(true))) // * it seems this currently never happens
         )

hkmc2/shared/src/main/scala/hkmc2/codegen/js/JSBuilder.scala

@@ -203,6 +203,18 @@ class JSBuilder extends CodeBuilder:
 
     case Throw(res) =>
       doc" # throw ${result(res)}"
+    
+    case Break(lbl, false) =>
+      doc" # break ${getVar(lbl)}"
+    
+    case Break(lbl, true) =>
+      doc" # continue ${getVar(lbl)}"
+    
+    case Label(lbl, bod, rst) =>
+      scope.allocateName(lbl)
+      doc" # ${getVar(lbl)}: while (true) { #{ ${
+        returningTerm(bod)
+      } # break; #}  # }${returningTerm(rst)}"
 
     case TryBlock(sub, fin, rst) =>
       doc" # try { #{ ${returningTerm(sub)

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

@@ -144,7 +144,7 @@ extends Importer:
     case InfixApp(lhs, Keyword.`is` | Keyword.`and`, rhs) =>
       val des = new Desugarer(tl, this).shorthands(tree)(ctx)
       val nor = new ucs.Normalization(tl)(des)
-      Term.If(des)(nor)
+      Term.IfLike(Keyword.`if`, des)(nor)
     case App(Ident("|"), Tree.Tup(lhs :: rhs :: Nil)) =>
       Term.CompType(term(lhs), term(rhs), true)
     case App(Ident("&"), Tree.Tup(lhs :: rhs :: Nil)) =>
@@ -174,14 +174,14 @@ extends Importer:
         Term.New(term(cls), Nil).withLocOf(tree)
       // case _ =>
       //   raise(ErrorReport(msg"Illegal new expression." -> tree.toLoc :: Nil))
-    case Tree.If(split) =>
+    case Tree.IfLike(kw, split) =>
       val desugared = new Desugarer(tl, this).termSplit(split, identity)(Split.Nil)(ctx)
       scoped("ucs:desugared"):
         log(s"Desugared:\n${Split.display(desugared)}")
       val normalized = new ucs.Normalization(tl)(desugared)
       scoped("ucs:normalized"):
         log(s"Normalized:\n${Split.display(normalized)}")
-      Term.If(desugared)(normalized)
+      Term.IfLike(kw, desugared)(normalized)
     case IfElse(InfixApp(InfixApp(scrutinee, Keyword.`is`, ctor @ Ident(cls)), Keyword.`then`, cons), alts) =>
       ctx.get(cls) match
         case S(sym: ClassSymbol) =>
@@ -192,7 +192,7 @@ extends Importer:
             Pattern.Class(sym, N, true)(ctor),
             Split.default(term(cons))
           ) :: Split.default(term(alts)))
-          Term.If(body)(body)
+          Term.IfLike(Keyword.`if`, body)(body)
         case _ =>
           raise(ErrorReport(msg"Illegal pattern $cls." -> tree.toLoc :: Nil))
           Term.Error
@@ -203,7 +203,7 @@ extends Importer:
         scrutVar.ref(),
         Split.default(term(cons))
       ) :: Split.default(term(alts)))
-      Term.If(body)(body)
+      Term.IfLike(Keyword.`if`, body)(body)
     case Tree.Quoted(body) => Term.Quoted(term(body))
     case Tree.Unquoted(body) => Term.Unquoted(term(body))
     case Tree.Case(branches) =>
@@ -215,7 +215,7 @@ extends Importer:
       val nor = new ucs.Normalization(tl)(des)
       scoped("ucs:normalized"):
         log(s"Normalized:\n${Split.display(nor)}")
-      Term.Lam(Param(FldFlags.empty, scrut, N) :: Nil, Term.If(des)(nor))
+      Term.Lam(Param(FldFlags.empty, scrut, N) :: Nil, Term.IfLike(Keyword.`if`, des)(nor))
     case Modified(Keyword.`return`, kwLoc, body) =>
       Term.Ret(term(body))
     case Tree.Region(id: Tree.Ident, body) =>

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

@@ -16,7 +16,7 @@ enum Term extends Statement:
   case TyApp(lhs: Term, targs: Ls[Term])
   case Sel(prefix: Term, nme: Tree.Ident)
   case Tup(fields: Ls[Fld])(val tree: Tree.Tup)
-  case If(desugared: Split)(val normalized: Split)
+  case IfLike(kw: Keyword.`if`.type | Keyword.`while`.type, desugared: Split)(val normalized: Split)
   case Lam(params: Ls[Param], body: Term)
   case FunTy(lhs: Term, rhs: Term, eff: Opt[Term])
   case Forall(tvs: Ls[QuantVar], body: Term)
@@ -63,7 +63,8 @@ enum Term extends Statement:
     case TyApp(lhs, targs) => "type application"
     case Sel(pre, nme) => "selection"
     case Tup(fields) => "tuple literal"
-    case If(body) => "`if` expression"
+    case IfLike(Keyword.`if`, body) => "`if` expression"
+    case IfLike(Keyword.`while`, body) => "`while` expression"
     case Lam(params, body) => "function literal"
     case FunTy(lhs, rhs, eff) => "function type"
     case Forall(tvs, body) => "universal quantification"
@@ -95,7 +96,7 @@ sealed trait Statement extends AutoLocated:
     case TyApp(pre, tarsg) => pre :: tarsg
     case Sel(pre, _) => pre :: Nil
     case Tup(fields) => fields.map(_.value)
-    case If(body) => Nil // TODO
+    case IfLike(_, body) => Nil // TODO
     case Lam(params, body) => body :: Nil
     case Blk(stats, res) => stats.flatMap(_.subTerms) ::: res :: Nil
     case Quoted(term) => term :: Nil
@@ -158,7 +159,7 @@ sealed trait Statement extends AutoLocated:
     case Forall(tvs, body) => s"forall ${tvs.mkString(", ")}: ${body.toString}"
     case WildcardTy(in, out) => s"in ${in.map(_.toString).getOrElse("⊥")} out ${out.map(_.toString).getOrElse("⊤")}"
     case Sel(pre, nme) => s"${pre.showDbg}.${nme.name}"
-    case If(body) => s"if { ${body.showDbg} }"
+    case IfLike(kw, body) => s"${kw.name} { ${body.showDbg} }"
     case Lam(params, body) => s"λ${params.map(_.showDbg).mkString(", ")}. ${body.showDbg}"
     case Blk(stats, res) =>
       (stats.map(_.showDbg + "; ") :+ (res match { case Lit(Tree.UnitLit(true)) => "" case x => x.showDbg + " " }))

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

@@ -59,6 +59,7 @@ object Keyword:
   // val `;` = Keyword(";", ascPrec, eqPrec)
 
   val `if` = Keyword("if", N, nextPrec)
+  val `while` = Keyword("while", N, curPrec)
   val `then` = Keyword("then", nextPrec, curPrec)
   val `else` = Keyword("else", nextPrec, curPrec)
   val `case` = Keyword("case", N, N)
@@ -75,9 +76,8 @@ object Keyword:
   val `rec` = Keyword("rec", N, N)
   val `in` = Keyword("in", curPrec, curPrec)
   val `out` = Keyword("out", N, curPrec)
-  val `set` = Keyword("set", N, N)
+  val `set` = Keyword("set", N, curPrec)
   val `do` = Keyword("do", N, N)
-  val `while` = Keyword("while", N, N)
   val `declare` = Keyword("declare", N, N)
   val `trait` = Keyword("trait", N, N)
   val `mixin` = Keyword("mixin", N, N)

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

@@ -159,29 +159,11 @@ object ParseRule:
         // ) { case (lhs, body) => Let(lhs, lhs, body) }
       )
 
-  val prefixRules: ParseRule[Tree] = ParseRule("start of statement", omitAltsStr = true)(
-    letLike(`let`),
-    letLike(`set`),
-    Kw(`new`):
-      ParseRule("`new` keyword"):
-        Expr(ParseRule("`new` expression")(End(())))((body, _: Unit) => New(body))
-    ,
-    Kw(`in`):
-      ParseRule("modifier keyword `in`"):
-        Expr(
-          ParseRule("`in` expression")(
-            Kw(`out`)(ParseRule(s"modifier keyword `out`")(standaloneExpr)).map(s => S(Tree.Modified(`out`, N/* TODO */, s))),
-            End(N),
-          )
-        ) {
-          case (lhs, N) => Tree.Modified(`in`, N/* TODO */, lhs)
-          case (lhs, S(rhs)) => Tup(Tree.Modified(`in`, N/* TODO */, lhs) :: rhs :: Nil)
-        }
-    ,
-    Kw(`if`):
-      ParseRule("`if` keyword")(
+  def ifLike(kw: `if`.type | `while`.type): Alt[Tree] =
+    Kw(kw):
+      ParseRule(s"'${kw.name}' keyword")(
         Expr(
-          ParseRule("`if` expression")(
+          ParseRule(s"'${kw.name}' expression")(
             End(N),
             Kw(`else`):
               ParseRule(s"`else` keyword")(
@@ -196,14 +178,35 @@ object ParseRule:
             val items = split match
               case Block(stmts) => stmts.appended(clause)
               case _ => split :: clause :: Nil
-            If(Block(items))
-          case (split, N) => If(split)
+            IfLike(kw, Block(items))
+          case (split, N) => IfLike(kw, split)
         ,
         Blk(
-          ParseRule("`if` block")(End(()))
-        ) { case (body, _) => If(body) }
+          ParseRule(s"'${kw.name}' block")(End(()))
+        ) { case (body, _) => IfLike(kw, body) }
       )
+  
+  val prefixRules: ParseRule[Tree] = ParseRule("start of statement", omitAltsStr = true)(
+    letLike(`let`),
+    letLike(`set`),
+    Kw(`new`):
+      ParseRule("`new` keyword"):
+        Expr(ParseRule("`new` expression")(End(())))((body, _: Unit) => New(body))
+    ,
+    Kw(`in`):
+      ParseRule("modifier keyword `in`"):
+        Expr(
+          ParseRule("`in` expression")(
+            Kw(`out`)(ParseRule(s"modifier keyword `out`")(standaloneExpr)).map(s => S(Tree.Modified(`out`, N/* TODO */, s))),
+            End(N),
+          )
+        ) {
+          case (lhs, N) => Tree.Modified(`in`, N/* TODO */, lhs)
+          case (lhs, S(rhs)) => Tup(Tree.Modified(`in`, N/* TODO */, lhs) :: rhs :: Nil)
+        }
     ,
+    ifLike(`if`),
+    ifLike(`while`),
     Kw(`else`):
       ParseRule("`else` clause")(
         Expr(ParseRule("`else` expression")(End(()))):

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

@@ -63,7 +63,7 @@ enum Tree extends AutoLocated:
   case Sel(prefix: Tree, name: Ident)
   case InfixApp(lhs: Tree, kw: Keyword.Infix, rhs: Tree)
   case New(body: Tree)
-  case If(split: Tree)
+  case IfLike(kw: Keyword.`if`.type | Keyword.`while`.type, split: Tree)
   @deprecated("Use If instead", "hkmc2-ucs")
   case IfElse(cond: Tree, alt: Tree)
   case Case(branches: Tree)
@@ -88,7 +88,7 @@ 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 If(split) => split :: Nil
+    case IfLike(_, split) => split :: Nil
     case IfElse(cond, alt) => cond :: alt :: Nil
     case Case(bs) => Ls(bs)
     case Region(name, body) => name :: body :: Nil
@@ -122,7 +122,8 @@ enum Tree extends AutoLocated:
     case Sel(prefix, name) => "selection"
     case InfixApp(lhs, kw, rhs) => "infix application"
     case New(body) => "new"
-    case If(split) => "if expression"
+    case IfLike(Keyword.`if`, split) => "if expression"
+    case IfLike(Keyword.`while`, split) => "while expression"
     case IfElse(cond, alt) => "if-then-else"
     case Case(branches) => "case"
     case Region(name, body) => "region"

hkmc2/shared/src/test/mlscript-compile/Stack.mjs

@@ -0,0 +1,28 @@
+const Stack = new class Stack {
+  constructor() {
+    this.Cons = class Cons {
+      constructor(head, tail) {
+        this.head = head;
+        this.tail = tail;
+        
+      }
+      toString() { return "Cons(" + this.head + ", " + this.tail + ")"; }
+    };
+    this.Nil = new class Nil {
+      constructor() {
+        
+      }
+      toString() { return "Nil"; }
+    };
+  }
+  isEmpty(xs) {
+    if (xs === this.Nil) {
+      return true
+    } else {
+      return false
+    }
+  }
+  toString() { return "Stack"; }
+};
+undefined
+export default Stack;