Move UCS normalization before lowering (#308)

Author: chengluyu

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

@@ -295,19 +295,24 @@ class BBTyper(using elState: Elaborator.State, tl: TL):
       val nestCtx1 = ctx.nest
       val nestCtx2 = ctx.nest
       val patTy = pattern match
-      case Pattern.ClassLike(sym, _, _, _) =>
-        val (clsTy, tv, emptyTy) = sym.asCls.flatMap(_.defn) match
-        case S(cls) =>
-          (ClassLikeType(sym, cls.tparams.map(_ => freshWildcard(sym))), (freshVar(new TempSymbol(S(scrutinee), "scrut"))), ClassLikeType(sym, cls.tparams.map(_ => Wildcard.empty)))
-        case _ =>
-          error(msg"Cannot match ${scrutinee.toString} as ${sym.toString}" -> split.toLoc :: Nil)
-          (Bot, Bot, Bot)
-        scrutinee match // * refine
-          case Ref(sym: LocalSymbol) =>
-            nestCtx1 += sym -> clsTy
-            nestCtx2 += sym -> tv
-          case _ => () // TODO: refine all variables holding this value?
-        clsTy | (tv & Type.mkNegType(emptyTy))
+      case pat: Pattern.ClassLike =>
+        pat.constructor.symbol.flatMap(_.asCls) match
+          case S(sym) =>
+            val (clsTy, tv, emptyTy) = sym.defn.map(sym -> _) match
+            case S((sym, cls)) =>
+              (ClassLikeType(sym, cls.tparams.map(_ => freshWildcard(sym))), (freshVar(new TempSymbol(S(scrutinee), "scrut"))), ClassLikeType(sym, cls.tparams.map(_ => Wildcard.empty)))
+            case _ =>
+              error(msg"Cannot match ${scrutinee.toString} as ${sym.toString}" -> split.toLoc :: Nil)
+              (Bot, Bot, Bot)
+            scrutinee match // * refine
+              case Ref(sym: LocalSymbol) =>
+                nestCtx1 += sym -> clsTy
+                nestCtx2 += sym -> tv
+              case _ => () // TODO: refine all variables holding this value?
+            clsTy | (tv & Type.mkNegType(emptyTy))
+          case N =>
+            error(msg"Not a valid class: ${pat.constructor.describe}" -> pat.constructor.toLoc :: Nil)
+            Bot
       case Pattern.Lit(lit) => lit match
         case _: Tree.BoolLit => BbCtx.boolTy
         case _: Tree.IntLit => BbCtx.intTy

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

@@ -220,7 +220,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
       args(fs)(args => k(Value.Arr(args)))
     case ref @ st.Ref(sym) =>
       sym match
-      case ctx.builtins.source.bms | ctx.builtins.js.bms | ctx.builtins.debug.bms =>
+      case ctx.builtins.source.bms | ctx.builtins.js.bms | ctx.builtins.debug.bms | ctx.builtins.annotations.bms =>
         raise:
           ErrorReport(
             msg"Module '${sym.nme}' is virtual (i.e., \"compiler fiction\"); cannot be used directly" -> t.toLoc ::
@@ -466,29 +466,36 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
               )
             pat match
               case Pattern.Lit(lit) => mkMatch(Case.Lit(lit) -> go(tail, topLevel = false))
-              case Pattern.ClassLike(cls: ClassSymbol, _trm, _args0, _refined)
-                  // Do not elaborate `_trm` when the `cls` is virtual.
-                  if Elaborator.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 make it `Predef.unreachable` here.
-                mkMatch(Case.Cls(cls, unreachableFn) -> go(tail, topLevel = false))
-              case Pattern.ClassLike(cls, trm, args0, _refined) =>
-                subTerm_nonTail(trm): st =>
-                  val args = args0.getOrElse(Nil)
-                  val clsParams = cls match
-                    case cls: ClassSymbol => cls.tree.clsParams
-                    case _: ModuleSymbol => Nil
-                  assert(args0.isEmpty || clsParams.length === args.length)
+              case Pattern.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)
                   def mkArgs(args: Ls[TermSymbol -> BlockLocalSymbol])(using Subst): Case -> Block = args match
                     case Nil =>
-                      Case.Cls(cls, st) -> go(tail, topLevel = false)
+                      Case.Cls(ctorSym, st) -> go(tail, topLevel = false)
                     case (param, arg) :: args =>
                       val (cse, blk) = mkArgs(args)
                       (cse, Assign(arg, Select(sr, param.id/*FIXME incorrect Ident?*/)(S(param)), blk))
-                  mkMatch(mkArgs(clsParams.iterator.zip(args).collect { case (s1, S(s2)) => (s1, s2) }.toList))
+                  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: ModuleSymbol) => 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 Pattern.Tuple(len, inf) => mkMatch(Case.Tup(len, inf) -> go(tail, topLevel = false))
               case Pattern.Record(entries) =>
                 val objectSym = ctx.builtins.Object
@@ -513,11 +520,17 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
           Throw(Instantiate(Select(Value.Ref(State.globalThisSymbol), Tree.Ident("Error"))(N),
             Value.Lit(syntax.Tree.StrLit("match error")) :: Nil)) // TODO add failed-match scrutinee info
 
-      if k.isInstanceOf[TailOp] && isIf then go(iftrm.normalized, topLevel = true)
+      val normalize = ucs.Normalization()
+      val normalized = tl.scoped("ucs:normalize"):
+        normalize(iftrm.desugared)
+      tl.scoped("ucs:normalized"):
+        tl.log(s"Normalized:\n${Split.display(normalized)}")
+
+      if k.isInstanceOf[TailOp] && isIf then go(normalized, topLevel = true)
       else
         val body = if isWhile
-          then Label(lbl, go(iftrm.normalized, topLevel = true), End())
-          else go(iftrm.normalized, topLevel = true)
+          then Label(lbl, go(normalized, topLevel = true), End())
+          else go(normalized, topLevel = true)
         Begin(
           body,
           if usesResTmp then k(Value.Ref(l))

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

@@ -156,6 +156,8 @@ object Elaborator:
       object debug extends VirtualModule(assumeBuiltinMod("debug")):
         val printStack = assumeObject("printStack")
         val getLocals = assumeObject("getLocals")
+      object annotations extends VirtualModule(assumeBuiltinMod("annotations")):
+        val compile = assumeObject("compile")
       def getBuiltinOp(op: Str): Opt[Str] =
         if getBuiltin(op).isDefined then builtinBinOps.get(op) else N
       /** Classes that do not use `instanceof` in pattern matching. */
@@ -210,10 +212,22 @@ object Elaborator:
       BlockMemberSymbol(id.name, Nil, true)
     val matchResultClsSymbol =
       val id = new Ident("MatchResult")
-      ClassSymbol(TypeDef(syntax.Cls, App(id, Tup(Ident("captures") :: Nil)), N, N), id)
+      val td = TypeDef(syntax.Cls, App(id, Tup(Ident("captures") :: Nil)), N, N)
+      val cs = ClassSymbol(td, id)
+      val flag = FldFlags.empty.copy(value = true)
+      val ps = PlainParamList(Param(flag, VarSymbol(Ident("captures")), N, Modulefulness(N)(false)) :: Nil)
+      cs.defn = S(ClassDef.Parameterized(N, syntax.Cls, cs, BlockMemberSymbol(cs.name, Nil),
+        Nil, ps, N, ObjBody(Blk(Nil, Term.Lit(UnitLit(false)))), N, Nil))
+      cs
     val matchFailureClsSymbol =
       val id = new Ident("MatchFailure")
-      ClassSymbol(TypeDef(syntax.Cls, App(id, Tup(Ident("errors") :: Nil)), N, N), id)
+      val td = TypeDef(syntax.Cls, App(id, Tup(Ident("errors") :: Nil)), N, N)
+      val cs = ClassSymbol(td, id)
+      val flag = FldFlags.empty.copy(value = true)
+      val ps = PlainParamList(Param(flag, VarSymbol(Ident("errors")), N, Modulefulness(N)(false)) :: Nil)
+      cs.defn = S(ClassDef.Parameterized(N, syntax.Cls, cs, BlockMemberSymbol(cs.name, Nil),
+        Nil, ps, N, ObjBody(Blk(Nil, Term.Lit(UnitLit(false)))), N, Nil))
+      cs
     val builtinOpsMap =
       val baseBuiltins = builtins.map: op =>
           op -> BuiltinSymbol(op,
@@ -263,14 +277,6 @@ extends Importer:
         N
     case _ => N
 
-  /** To perform a reverse lookup for a term that references a symbol in the current context. */
-  def reference(target: ClassSymbol | ModuleSymbol): Ctxl[Opt[Term]] =
-    def go(ctx: Ctx): Opt[Term] =
-      ctx.env.values.collectFirst:
-        case elem if elem.symbol.flatMap(_.asClsLike).contains(target) => elem.ref(target.id)
-      .orElse(ctx.parent.flatMap(go))
-    go(ctx).map(Term.SynthSel(_, Ident("class"))(S(target)))
-  
   def cls(trm: Term, inAppPrefix: Bool)
       : Ctxl[Term]
       = trace[Term](s"Elab class ${trm}", r => s"~> $r"):
@@ -467,10 +473,7 @@ extends Importer:
       val des = new ucs.Desugarer(this)(tree)
       scoped("ucs:desugared"):
         log(s"Desugared:\n${Split.display(des)}")
-      val nor = new ucs.Normalization(this)(des)
-      scoped("ucs:normalized"):
-        log(s"Normalized:\n${Split.display(nor)}")
-      Term.IfLike(Keyword.`if`, des)(nor)
+      Term.IfLike(Keyword.`if`, des)
     case InfixApp(lhs, Keyword.`then`, rhs) =>
       raise:
         ErrorReport(msg"Unexpected infix use of 'then' keyword here" -> tree.toLoc :: Nil)
@@ -597,23 +600,17 @@ extends Importer:
       val desugared = new ucs.Desugarer(this)(tree)
       scoped("ucs:desugared"):
         log(s"Desugared:\n${Split.display(desugared)}")
-      val normalized = new ucs.Normalization(this)(desugared)
-      scoped("ucs:normalized"):
-        log(s"Normalized:\n${Split.display(normalized)}")
-      Term.IfLike(kw, desugared)(normalized)
+      Term.IfLike(kw, desugared)
     case Tree.Quoted(body) => Term.Quoted(subterm(body))
     case Tree.Unquoted(body) => Term.Unquoted(subterm(body))
     case tree @ Tree.Case(_, branches) =>
       val scrut = VarSymbol(Ident("caseScrut"))
       val des = new ucs.Desugarer(this)(tree, scrut)
       scoped("ucs:desugared"):
         log(s"Desugared:\n${Split.display(des)}")
-      val nor = new ucs.Normalization(this)(des)
-      scoped("ucs:normalized"):
-        log(s"Normalized:\n${Split.display(nor)}")
       Term.Lam(PlainParamList(
           Param(FldFlags.empty, scrut, N, Modulefulness.none) :: Nil
-        ), Term.IfLike(Keyword.`if`, des)(nor))
+        ), Term.IfLike(Keyword.`if`, des))
     case Modified(Keyword.`return`, kwLoc, body) =>
       ctx.getRetHandler match
       case ReturnHandler.Required(sym) =>

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

@@ -3,8 +3,11 @@ package semantics
 
 import mlscript.utils.*, shorthands.*
 import syntax.*, Tree.Ident
+import Elaborator.{Ctx, ctx}
 import ucs.DeBrujinSplit
 
+import Pattern.*
+
 /** Flat patterns for pattern matching */
 enum Pattern extends AutoLocated:
 
@@ -13,37 +16,84 @@ enum Pattern extends AutoLocated:
   /** An individual argument is None when it is not matched, i.e. when an underscore is used there.
     * The whole argument list is None when no argument list is being matched at all, as in `x is Some then ...`. */
   case ClassLike(
-    sym: ClassSymbol | ModuleSymbol,
-    trm: Term,
-    args: Opt[List[Opt[BlockLocalSymbol]]],
-    var refined: Bool,
+      val constructor: Term,
+      val arguments: Opt[Ls[Argument]],
+      val mode: MatchMode,
+      var refined: Bool
   )(val tree: Tree)
 
-  case Synonym(symbol: PatternSymbol, patternArguments: Ls[(split: DeBrujinSplit, tree: Tree)])
-  
   case Tuple(size: Int, inf: Bool)
 
   case Record(entries: List[(Ident -> BlockLocalSymbol)])
+
 
   def subTerms: Ls[Term] = this match
-    case ClassLike(_, t, _, _) => t :: Nil
-    case _: (Lit | Synonym | Tuple | Record) => Nil 
+    case p: ClassLike => p.constructor :: (p.mode match
+      case MatchMode.Default | _: MatchMode.StringPrefix => Nil
+      case MatchMode.Annotated(annotation) => annotation :: Nil)
+    case _: (Lit | Tuple | Record) => Nil 
 
   def children: Ls[Located] = this match
     case Lit(literal) => literal :: Nil
-    case ClassLike(_, t, args, _) =>
-      t :: args.fold(Nil)(_.collect { case S(symbol) => symbol })
-    case Synonym(_, arguments) => arguments.map(_.tree)
+    case ClassLike(ctor, scruts, _, _) => ctor :: scruts.fold(Nil)(_.map(_.scrutinee))
     case Tuple(fields, _) => Nil
     case Record(entries) => entries.flatMap { case (nme, als) => nme :: als :: Nil }
 
   def showDbg: Str = this match
     case Lit(literal) => literal.idStr
-    case ClassLike(sym, t, ps, rfd) => (if rfd then "refined " else "") +
-      sym.nme + ps.fold("")(_.iterator.map(_.fold("_")(_.toString)).mkString("(", ", ", ")"))
-    case Synonym(symbol, arguments) =>
-      symbol.nme + arguments.iterator.map(_.tree.showDbg).mkString("(", ", ", ")")
+    case ClassLike(ctor, args, _, rfd) =>
+      def showCtor(ctor: Term): Str = ctor match
+        // This prints the symbol name without `refNum` and "member:" prefix.
+        case Term.Ref(sym: BlockMemberSymbol) => sym.nme
+        // This prints the symbol without `refNum`.
+        case Term.Ref(sym) => sym.toString
+        case Term.Sel(p, i) => s"${showCtor(p)}.${i.name}"
+        case Term.SynthSel(p, i) => s"${showCtor(p)}.${i.name}"
+        case _ => ctor.showDbg
+      (if rfd then "refined " else "") + showCtor(ctor) +
+        args.fold("")(_.iterator.map(_.scrutinee.nme).mkString("(", ", ", ")"))
     case Tuple(size, inf) => "[]" + (if inf then ">=" else "=") + size
     case Record(Nil) => "{}"
     case Record(entries) =>
       entries.iterator.map(_.name + ": " + _).mkString("{ ", ", ", " }")
+
+object Pattern:
+  /** Represent the type of arguments in `ClassLike` patterns. This type alias
+   *  is used to reduce repetition in the code.
+   * 
+   *  - Field `pattern` is for error messages.
+   *  - Field `split` is for pattern compilation.
+   *    **TODO(ucs/rp)**: Replace with suitable representation when implement
+   *    the new pattern compilation.
+   */
+  type Argument = (scrutinee: BlockLocalSymbol, pattern: Tree, split: Opt[DeBrujinSplit])
+  
+  /** A class-like pattern whose symbol is resolved to a class. */
+  object Class:
+    def unapply(p: Pattern): Opt[ClassSymbol] = p match
+      case p: Pattern.ClassLike => p.constructor.symbol.flatMap(_.asCls)
+      case _ => N
+  
+  /** A class-like pattern whose symbol is resolved to a module. */
+  object Module:
+    def unapply(p: Pattern): Opt[ModuleSymbol] = p match
+      case p: Pattern.ClassLike => p.constructor.symbol.flatMap(_.asModOrObj)
+      case _ => N
+  
+  enum MatchMode:
+    /** The default mode. If the constructor resolves to:
+     *  - a `ClassSymbol`, then check if the scrutinee is an instance;
+     *  - a `ModuleSymbol`, then check if the scrutinee is the object;
+     *  - a `PatternSymbol`, then call `unapply` on the pattern.
+     */
+    case Default
+    /** Call `unapplyStringPrefix` instead of `unapply`. */
+    case StringPrefix(prefix: TempSymbol, postfix: TempSymbol)
+    /** The pattern is annotated. The normalization will intepret the pattern
+     *  matching behavior based on the resolved symbol
+     */
+    case Annotated(annotation: Term)
+    
+  object ClassLike:
+    def apply(constructor: Term, arguments: Opt[Ls[BlockLocalSymbol]]): ClassLike =
+      ClassLike(constructor, arguments.map(_.map(s => (s, Tree.Dummy, N))), MatchMode.Default, false)(Tree.Dummy)

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

@@ -289,7 +289,7 @@ class Resolver(tl: TraceLogger)
             case Split.Else(default) =>
               traverse(default, expect = Any)
             case Split.End =>
-          split(t.normalized)
+          split(t.desugared)
 
         case Term.New(cls, args, rft) =>
           traverse(cls, expect = Any)

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

@@ -40,6 +40,7 @@ abstract class Symbol(using State) extends Located:
     case _ => N
   def asMod: Opt[ModuleSymbol] = asModOrObj.filter(_.tree.k is Mod)
   def asObj: Opt[ModuleSymbol] = asModOrObj.filter(_.tree.k is Obj)
+  def asClsOrMod: Opt[ClassSymbol | ModuleSymbol] = asCls orElse asModOrObj
   /* 
   def asTrm: Opt[TermSymbol] = this match
     case trm: TermSymbol => S(trm)
@@ -224,6 +225,7 @@ case class ErrorSymbol(val nme: Str, tree: Tree)(using State) extends MemberSymb
 
 sealed trait ClassLikeSymbol extends Symbol:
   self: MemberSymbol[? <: ClassDef | ModuleDef] =>
+  val tree: Tree.TypeDef
   def subst(using sub: SymbolSubst): ClassLikeSymbol
 
 

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

@@ -64,7 +64,7 @@ sealed trait ResolvableImpl:
     case S(td: ClassLikeDef) => S(td)
     case _ => N
 
-  def withIArgs(iargsLs: Ls[Term.Tup]): Term = 
+  def withIArgs(iargsLs: Ls[Term.Tup]): this.type = 
     if !(this.iargsLs.isEmpty || this.iargsLs.get == iargsLs) then
       lastWords:
         s"the implicit arguments for term ${t.showDbg} " +
@@ -89,7 +89,7 @@ enum Term extends Statement:
   case SynthSel(prefix: Term, nme: Tree.Ident)(var sym: Opt[FieldSymbol]) extends Term with ResolvableImpl
   case DynSel(prefix: Term, fld: Term, arrayIdx: Bool)
   case Tup(fields: Ls[Elem])(val tree: Tree.Tup)
-  case IfLike(kw: Keyword.`if`.type | Keyword.`while`.type, desugared: Split)(val normalized: Split)
+  case IfLike(kw: Keyword.`if`.type | Keyword.`while`.type, desugared: Split)
   case Lam(params: ParamList, body: Term)
   case FunTy(lhs: Term, rhs: Term, eff: Opt[Term])
   case Forall(tvs: Ls[QuantVar], outer: Opt[VarSymbol], body: Term)