Unify patterns in UCS and UPS (#336)

Co-authored-by: Lionel Parreaux <[email protected]>

Author: chengluyu

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

@@ -48,7 +48,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
@@ -60,7 +60,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
@@ -75,7 +75,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) =>
@@ -90,7 +90,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
@@ -110,7 +110,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
@@ -132,7 +132,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
@@ -177,12 +177,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))
@@ -267,7 +267,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
@@ -595,7 +595,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

@@ -45,11 +45,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

@@ -332,7 +332,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)
@@ -820,7 +820,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
           AssignField(runtimePath, stackDepthIdent, depthSym.asPath, mainMatchBlk)(N)
         else mainMatchBlk
 
-      val lbl = blockBuilder.label(loopLbl, withResetDepth).rest(End())
+      val lbl = blockBuilder.label(loopLbl, loop = true, withResetDepth).rest(End())
 
       def createAssignment(sym: Local) = Assign(sym, resumedVal.asPath, End())
 

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

@@ -718,11 +718,11 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
                     (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 = applySubBlockAndReset(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 TryBlock(sub, fin, rst) =>
           val sub2 = applySubBlockAndReset(sub)
           val fin2 = applySubBlockAndReset(fin)