Unify patterns in UCS and UPS

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

@@ -4,6 +4,7 @@ package bbml
 
 import scala.collection.mutable.{HashSet, HashMap, ListBuffer}
 import scala.annotation.tailrec
+import sourcecode.{FileName, Line, Name}
 
 import mlscript.utils.*, shorthands.*
 import utils.*
@@ -73,7 +74,7 @@ object BbCtx:
 end BbCtx
 
 
-class BBTyper(using elState: Elaborator.State, tl: TL):
+class BBTyper(using elState: Elaborator.State, tl: TL)(using Ctx):
   import tl.{trace, log}
 
   private val infVarState = new InfVarUid.State()
@@ -100,7 +101,7 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
     state.lowerBounds = ctx.getRegEnv :: Nil
     InfVar(ctx.lvl, infVarState.nextUid, state, false)(sym, "")
 
-  private def error(msg: Ls[Message -> Opt[Loc]])(using BbCtx) =
+  private def error(using Line, FileName, Name, Raise)(msg: Ls[Message -> Opt[Loc]])(using BbCtx) =
     raise(ErrorReport(msg))
     Bot // TODO: error type?
 
@@ -258,7 +259,7 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
       val res = freshVar(new TempSymbol(S(blk), "ctx"))(using ctx)
       constrain(bodyCtx, sk | res)
       (bodyTy, rhsCtx | res, rhsEff | bodyEff)
-    case Term.IfLike(Keyword.`if`, Split.Let(_, cond, Split.Cons(Branch(_, FlatPattern.Lit(BoolLit(true)), Split.Else(cons)), Split.Else(alts)))) =>
+    case Term.IfLike(Keyword.`if`, SimpleSplit.IfThenElse(cond, cons, alts)) =>
       val (condTy, condCtx, condEff) = typeCode(cond)
       val (consTy, consCtx, consEff) = typeCode(cons)
       val (altsTy, altsCtx, altsEff) = typeCode(alts)
@@ -363,8 +364,8 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
       given BbCtx = nextCtx
       constrain(ascribe(term, skolemize(pt))._2, Bot) // * never generalize terms with effects
       (pt, Bot)
-    case (Term.IfLike(Keyword.`if`, branches), ty) => // * propagate
-      typeSplit(branches, S(ty))
+    case (Term.IfLike(Keyword.`if`, split), ty) => // * propagate
+      typeSplit(split.getExpandedSplit, S(ty))
     case (Term.Asc(term, ty), rhs) =>
       ascribe(term, typeType(ty))
       ascribe(term, rhs)
@@ -550,7 +551,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.IfLike(Keyword.`if`, branches) => typeSplit(branches, N)
+      case Term.IfLike(Keyword.`if`, split) => typeSplit(split.getExpandedSplit, N)
       case reg @ Term.Region(sym, body) =>
         val sk = freshReg(sym)(using ctx)
         val nestCtx = ctx.nestReg(sk)

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

@@ -47,7 +47,7 @@ sealed abstract class Block extends Product:
     // Note that the handler's LHS and body are not part of the current block, so we do not consider them here.
     case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) => rst.definedVars + res
     case TryBlock(sub, fin, rst) => sub.definedVars ++ fin.definedVars ++ rst.definedVars
-    case Label(lbl, bod, rst) => bod.definedVars ++ rst.definedVars
+    case Label(lbl, _, bod, rst) => bod.definedVars ++ rst.definedVars
 
   lazy val size: Int = this match
     case _: Return | _: Throw | _: End | _: Break | _: Continue => 1
@@ -59,7 +59,7 @@ sealed abstract class Block extends Product:
       1 + arms.map(_._2.size).sum + dflt.map(_.size).getOrElse(0) + rst.size
     case Define(_, rst) => 1 + rst.size
     case TryBlock(sub, fin, rst) => 1 + sub.size + fin.size + rst.size
-    case Label(_, bod, rst) => 1 + bod.size + rst.size
+    case Label(_, _, bod, rst) => 1 + bod.size + rst.size
     case HandleBlock(lhs, res, par, args, cls, handlers, bdy, rst) => 1 + handlers.map(_.body.size).sum + bdy.size + rst.size
 
   // TODO conserve if no changes
@@ -74,7 +74,7 @@ sealed abstract class Block extends Product:
       Match(scrut, arms.map(_ -> _.mapTail(f)), dflt.map(_.mapTail(f)), rst)
     case Match(scrut, arms, dflt, rst) =>
       Match(scrut, arms, dflt, rst.mapTail(f))
-    case Label(label, body, rest) => Label(label, body, rest.mapTail(f))
+    case Label(label, loop, body, rest) => Label(label, loop, body.mapTail(f), rest.mapTail(f))
     case af @ AssignField(lhs, nme, rhs, rest) =>
       AssignField(lhs, nme, rhs, rest.mapTail(f))(af.symbol)
     case adf @ AssignDynField(lhs, fld, arrayIdx, rhs, rest) =>
@@ -89,7 +89,7 @@ sealed abstract class Block extends Product:
         (pat, arm) => arm.freeVars -- pat.freeVars
     case Return(res, implct) => res.freeVars
     case Throw(exc) => exc.freeVars
-    case Label(label, body, rest) => (body.freeVars - label) ++ rest.freeVars 
+    case Label(label, _, body, rest) => (body.freeVars - label) ++ rest.freeVars 
     case Break(label) => Set(label)
     case Continue(label) => Set(label)
     case Begin(sub, rest) => sub.freeVars ++ rest.freeVars
@@ -109,7 +109,7 @@ sealed abstract class Block extends Product:
         (pat, arm) => arm.freeVarsLLIR -- pat.freeVarsLLIR
     case Return(res, implct) => res.freeVarsLLIR
     case Throw(exc) => exc.freeVarsLLIR
-    case Label(label, body, rest) => (body.freeVarsLLIR - label) ++ rest.freeVarsLLIR 
+    case Label(label, _, body, rest) => (body.freeVarsLLIR - label) ++ rest.freeVarsLLIR 
     case Break(label) => Set.empty
     case Continue(label) => Set.empty
     case Begin(sub, rest) => sub.freeVarsLLIR ++ rest.freeVarsLLIR
@@ -131,7 +131,7 @@ sealed abstract class Block extends Product:
     case AssignDynField(_, _, _, rhs, rest) => rhs.subBlocks ::: rest :: Nil
     case Define(d, rest) => d.subBlocks ::: rest :: Nil
     case HandleBlock(_, _, par, args, _, handlers, body, rest) => par.subBlocks ++ args.flatMap(_.subBlocks) ++ handlers.map(_.body) :+ body :+ rest
-    case Label(_, body, rest) => body :: rest :: Nil
+    case Label(_, _, body, rest) => body :: rest :: Nil
 
     // TODO rm Lam from values and thus the need for these cases
     case Return(r, _) => r.subBlocks
@@ -176,12 +176,12 @@ sealed abstract class Block extends Product:
       then this
       else Match(scrut, newArms, newDflt, newRest)
 
-    case Label(label, body, rest) =>
+    case Label(label, loop, body, rest) =>
       val newBody = body.flattened
       val newRest = rest.flatten(k)
       if (newBody is body) && (newRest is rest)
       then this
-      else Label(label, newBody, newRest)
+      else Label(label, loop, newBody, newRest)
 
     case Begin(sub, rest) =>
       sub.flatten(_ => rest.flatten(k))
@@ -266,7 +266,7 @@ 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 Block
+case class Label(label: Local, loop: Bool, body: Block, rest: Block) extends Block
 
 case class Break(label: Local) extends BlockTail
 case class Continue(label: Local) extends BlockTail
@@ -593,7 +593,7 @@ extension (k: Block => Block)
   def end = k.rest(End())
   def ifthen(scrut: Path, cse: Case, trm: Block, els: Opt[Block] = N): Block => Block =
     k.chain(Match(scrut, cse -> trm :: Nil, els, _))
-  def label(label: Local, body: Block) = k.chain(Label(label, body, _))
+  def label(label: Local, loop: Bool, body: Block) = k.chain(Label(label, loop, body, _))
   def ret(r: Result) = k.rest(Return(r, false))
   def staticif(b: Boolean, f: (Block => Block) => (Block => Block)) = if b then k.transform(f) else k
   def foldLeft[A](xs: Iterable[A])(f: (Block => Block, A) => Block => Block) = xs.foldLeft(k)(f)

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

@@ -44,11 +44,11 @@ class BlockTransformer(subst: SymbolSubst):
                 (arms2 is arms) &&
                 (dflt2 is dflt) && (rst2 is rst)
               then b else Match(scrut2, arms2, dflt2, rst2)
-    case Label(lbl, bod, rst) =>
+    case Label(lbl, loop, bod, rst) =>
       val lbl2 = applyLocal(lbl)
       val bod2 = applySubBlock(bod)
       val rst2 = applySubBlock(rst)
-      if (lbl2 is lbl) && (bod2 is bod) && (rst2 is rst) then b else Label(lbl2, bod2, rst2)
+      if (lbl2 is lbl) && (bod2 is bod) && (rst2 is rst) then b else Label(lbl2, loop, bod2, rst2)
     case Begin(sub, rst) =>
       val sub2 = applySubBlock(sub)
       val rst2 = applySubBlock(rst)

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

@@ -31,7 +31,7 @@ class BlockTraverser:
         applyCase(arm._1); applySubBlock(arm._2)
       dflt.foreach(applySubBlock)
       applySubBlock(rst)
-    case Label(lbl, bod, rst) => applyLocal(lbl); applySubBlock(bod); applySubBlock(rst)
+    case Label(lbl, loop, bod, rst) => applyLocal(lbl); applySubBlock(bod); applySubBlock(rst)
     case Begin(sub, rst) => applySubBlock(sub); applySubBlock(rst)
     case TryBlock(sub, fin, rst) => applySubBlock(sub); applySubBlock(fin); applySubBlock(rst)
     case Assign(l, r, rst) => applyLocal(l); applyResult(r); applySubBlock(rst)

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

@@ -227,7 +227,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
               Match(scrut, newArms, dfltParts.map(_.head), StateTransition(restId)),
               BlockState(restId, restParts.head, N) :: states
             )
-      case l @ Label(label, body, rest) =>
+      case l @ Label(label, loop, body, rest) =>
         val startId = freshId() // start of body
 
         val PartRet(restNew, restParts) = go(rest)
@@ -586,7 +586,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       End()
     )
 
-    val lbl = blockBuilder.label(loopLbl, mainMatchBlk).rest(End())
+    val lbl = blockBuilder.label(loopLbl, loop = true, mainMatchBlk).rest(End())
 
     val resumedVal = VarSymbol(Tree.Ident("value$"))
 

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

@@ -66,7 +66,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
   val lowerHandlers: Bool = config.effectHandlers.isDefined
   val lift: Bool = config.liftDefns.isDefined
 
-  private lazy val unreachableFn =
+  lazy val unreachableFn =
     Select(Value.Ref(State.runtimeSymbol), Tree.Ident("unreachable"))(N)
 
   def unit: Path =
@@ -519,143 +519,10 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
           lamDef,
           k(lamSym |> Value.Ref.apply))
 
-    /* 
-    case t @ st.If(Split.Let(sym, trm, tail)) =>
-      // term(st.Blk(semantics.LetDecl(sym) :: semantics.DefineVar(sym, trm) :: Nil, st.If(tail)(t.normalized)))(k)
-      term(trm): r =>
-        Assign(sym, r, term(st.If(tail)(t.normalized))(k))
 
-    // TODO rm
-    case st.If(Split.Cons(
-      Branch(scrut, Pattern.LitPat(tru @ Tree.BoolLit(true)), Split.Else(thn)),
-      restSplit
-    )) =>
-
-      val elseBranch = restSplit match
-        case Split.Else(els) => S(els)
-        case Split.Nil => N
-
-      elseBranch match
-      case S(els) if k.isInstanceOf[Ret] =>
-        subTerm(scrut): sr =>
-          // Match(sr, Case.Lit(tru) -> term(thn)(k) :: Nil,
-          //   Some(term(els)(k)), 
-          //   Unreachable
-          // )
-          Match(sr, Case.Lit(tru) -> term(thn)(k) :: Nil,
-            N, 
-            term(els)(k)
-          )
-      case _ =>
-        val l = new TempSymbol(S(t))
-        subTerm(scrut): sr =>
-            Match(sr, Case.Lit(tru) -> subTerm(thn)(r => Assign(l, r, End())) :: Nil,
-              elseBranch.map(els => subTerm(els)(r => Assign(l, r, End()))),
-              k(Value.Ref(l))
-            )
-    */
+    case iftrm: st.IfLike => ucs.Normalization(this)(iftrm)(k)
 
-    case iftrm: st.IfLike =>
-
-      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(S(t))
-
-      lazy val lbl =
-        new TempSymbol(S(t))
-
-      def go(split: Split, topLevel: Bool)(using Subst): Block = split match
-        case Split.Let(sym, trm, tl) =>
-          term_nonTail(trm): r =>
-            Assign(sym, r, go(tl, topLevel))
-        case Split.Cons(Branch(scrut, pat, tail), restSplit) =>
-          subTerm_nonTail(scrut): sr =>
-            tl.log(s"Binding scrut $scrut to $sr (${summon[Subst].map})")
-            // val cse = 
-            def mkMatch(cse: Case -> Block) = Match(sr, cse :: Nil,
-                S(go(restSplit, topLevel = true)),
-                End()
-              )
-            pat match
-              case FlatPattern.Lit(lit) => mkMatch(Case.Lit(lit) -> go(tail, topLevel = false))
-              case FlatPattern.ClassLike(ctor, argsOpt, _mode, _refined) =>
-                /** Make a continuation that creates the match. */
-                def k(ctorSym: ClassLikeSymbol, clsParams: Ls[TermSymbol])(st: Path): Block =
-                  val args = argsOpt.map(_.map(_.scrutinee)).getOrElse(Nil)
-                  // Normalization should reject cases where the user provides
-                  // more sub-patterns than there are actual class parameters.
-                  assert(argsOpt.isEmpty || args.length <= clsParams.length, (argsOpt, clsParams))
-                  def mkArgs(args: Ls[TermSymbol -> BlockLocalSymbol])(using Subst): Case -> Block = args match
-                    case Nil =>
-                      Case.Cls(ctorSym, st) -> go(tail, topLevel = false)
-                    case (param, arg) :: args =>
-                      val (cse, blk) = mkArgs(args)
-                      (cse, Assign(arg, Select(sr, new Tree.Ident(param.id.name).withLocOf(arg))(S(param)), blk))
-                  mkMatch(mkArgs(clsParams.iterator.zip(args).toList))
-                ctor.symbol.flatMap(_.asClsOrMod) match
-                  case S(cls: ClassSymbol) if ctx.builtins.virtualClasses contains cls =>
-                    // [invariant:0] Some classes (e.g., `Int`) from `Prelude` do
-                    // not exist at runtime. If we do lowering on `trm`, backends
-                    // (e.g., `JSBuilder`) will not be able to handle the corresponding selections.
-                    // In this case the second parameter of `Case.Cls` will not be used.
-                    // So we do not elaborate `ctor` when the `cls` is virtual
-                    // and use it `Predef.unreachable` here.
-                    k(cls, Nil)(unreachableFn)
-                  case S(cls: ClassSymbol) => subTerm_nonTail(ctor)(k(cls, cls.tree.clsParams))
-                  case S(mod: ModuleOrObjectSymbol) => subTerm_nonTail(ctor)(k(mod, Nil))
-                  case N =>
-                    // Normalization have already checked the constructor
-                    // resolves to a class or module. Branches with unresolved
-                    // constructors should have been removed.
-                    lastWords("Pattern.ClassLike: constructor is neither a class nor a module")
-              case FlatPattern.Tuple(len, inf) => mkMatch(Case.Tup(len, inf) -> go(tail, topLevel = false))
-              case FlatPattern.Record(entries) =>
-                val objectSym = ctx.builtins.Object
-                mkMatch( // checking that we have an object
-                  Case.Cls(objectSym, Value.Ref(BuiltinSymbol(objectSym.nme, false, false, true, false))),
-                  entries.foldRight(go(tail, topLevel = false)):
-                    case ((fieldName, fieldSymbol), blk) =>
-                      mkMatch(
-                        Case.Field(fieldName, safe = true), // we know we have an object, no need to check again
-                        Assign(fieldSymbol, Select(sr, fieldName)(N), blk)
-                      )
-                )
-        case Split.Else(els) =>
-          if k.isInstanceOf[TailOp] && isIf then term_nonTail(els)(k)
-          else
-            term_nonTail(els): r =>
-              Assign(l, r,
-                if isWhile && !topLevel then Continue(lbl)
-                else End()
-              )
-        case Split.End =>
-          Throw(Instantiate(mut = false, Select(Value.Ref(State.globalThisSymbol), Tree.Ident("Error"))(N),
-            Value.Lit(syntax.Tree.StrLit("match error")).asArg :: Nil)) // TODO add failed-match scrutinee info
-
-      val normalize = ucs.Normalization()
-      val normalized = tl.scoped("ucs:normalize"):
-        normalize(iftrm.desugared)
-      tl.scoped("ucs:normalized"):
-        tl.log(s"Normalized:\n${normalized.prettyPrint}")
-
-      if k.isInstanceOf[TailOp] && isIf then go(normalized, topLevel = true)
-      else
-        val body = if isWhile
-          then Label(lbl, go(normalized, topLevel = true), End())
-          else go(normalized, topLevel = true)
-        Begin(
-          body,
-          if usesResTmp then k(Value.Ref(l))
-          else k(unit) // * it seems this currently never happens
-        )
+    case iftrm: st.SynthIf => ucs.Normalization(this)(iftrm)(k)
 
     case sel @ Sel(prefix, nme) =>
       setupSelection(prefix, nme, sel.sym)(k)
@@ -889,7 +756,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
           .assign(arrSym, Tuple(mut = false, (l4 :: l5 :: Nil).map(s => Value.Ref(s).asArg)))
           .rest(setupTerm("Blk", Value.Ref(arrSym) :: Value.Ref(l3) :: Nil)(k))
       }
-    case IfLike(syntax.Keyword.`if`, split) => quoteSplit(split): r =>
+    case IfLike(syntax.Keyword.`if`, split) => quoteSplit(split.getExpandedSplit): r =>
       val l = new TempSymbol(N)
       Assign(l, r, setupTerm("IfLike", setupQuotedKeyword("If") :: Value.Ref(l) :: Nil)(k))
     case Unquoted(body) => term(body)(k)
@@ -921,18 +788,18 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
     (mtds, publicFlds, privateFlds, ctor)
 
   /** Compile the pattern definition into `unapply` and `unapplyStringPrefix`
-   *  methods using the `NaiveCompiler`, which transliterate the pattern into
+   *  methods using the `SplitCompiler`, which transliterate the pattern into
    *  UCS splits that backtrack without any optimizations. */
   def compilePatternMethods(defn: PatternDef)(using Subst):
       // The return type is intended to be consistent with `gatherMembers`
       (Ls[FunDefn], Ls[BlockMemberSymbol -> TermSymbol], Ls[TermSymbol], Block) =
-    val compiler = new ups.NaiveCompiler
+    val compiler = new ups.SplitCompiler
     val methods = compiler.compilePattern(defn)
-    val mtds = methods
-      .flatMap: td =>
-        td.body.map: bod =>
-          val (paramLists, bodyBlock) = setupFunctionDef(td.params, bod, S(td.sym.nme))
-          FunDefn(td.owner, td.sym, paramLists, bodyBlock)
+    // We only need `owner`, `sym`, `params` and `body`
+    val mtds = methods.map:
+      case (sym, params, split) =>
+        val (paramLists, bodyBlock) = setupFunctionDef(params :: Nil, split, S(sym.nme))
+        FunDefn(N, sym, paramLists, bodyBlock)
     (mtds, Nil, Nil, End())
 
   def args(elems: Ls[Elem])(k: Ls[Arg] => Block)(using Subst): Block =
@@ -1050,6 +917,11 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
       (using Subst): (List[ParamList], Block) =
     (paramLists, returnedTerm(bodyTerm))
 
+  /** Same as the other overload, but expect the body to be a `Split`. */
+  def setupFunctionDef(paramLists: List[ParamList], bodySplit: Split, name: Option[Str])
+      (using Subst): (List[ParamList], Block) =
+    (paramLists, ucs.Normalization(this)(bodySplit)(Ret))
+
   def reportAnnotations(target: Statement, annotations: Ls[Annot]): Unit =
     annotations.foreach:
       case Annot.Untyped => ()

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

@@ -37,9 +37,9 @@ object Printer:
       doc"match ${mkDocument(scrut)} #{  # ${docCases} # else #{  # ${docDefault} #}  #}  # in # ${mkDocument(rest)}"
     case Return(res, implct) => doc"return ${mkDocument(res)}"
     case Throw(exc) => doc"throw ${mkDocument(exc)}"
-    case Label(label, body, rest) =>
+    case Label(label, loop, body, rest) =>
       val l2 = summon[Scope].allocateName(label)
-      doc"label $l2 = ${mkDocument(body)} in # ${mkDocument(rest)}"
+      doc"labelled ${if loop then "loop" else "block"} $l2 = ${mkDocument(body)} in # ${mkDocument(rest)}"
     case Break(label) =>
       doc"break ${getVar(label)}"
     case Continue(label) =>