Implement Typed IR

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

@@ -12,7 +12,8 @@ import hkmc2.semantics.{Term => st}
 import syntax.{Literal, Tree, SpreadKind, Keyword}
 import semantics.*
 import semantics.Term.*
-import sem.Elaborator.State
+import sem.Elaborator.{Ctx, State, ctx}
+import sourcecode.{FileName, Line}
 
 
 /* Important design notes.
@@ -532,7 +533,7 @@ object HandleBlock:
         N, sym, PlainParamList(Param(FldFlags.empty, handler.resumeSym, N, Modulefulness.none) :: Nil) :: Nil,
         handler.body
         )(N, annotations = Nil)
-      val rSym = TempSymbol(N, "suspendRes")
+      val rSym = TempSymbol(N, erasedType = N, "suspendRes")
       FunDefn.withFreshSymbol(
         S(cls),
         handler.sym,
@@ -678,7 +679,7 @@ final case class FunDefn(
 
 object FunDefn:
   def withFreshSymbol(owner: Opt[InnerSymbol], sym: BlockMemberSymbol, params: Ls[ParamList], body: Block)(configOverride: Opt[Config], annotations: Ls[Annot])(using State) =
-    val tSym = TermSymbol(syntax.Fun, owner, Tree.Ident(sym.nme))
+    val tSym = TermSymbol(syntax.Fun, owner, Tree.Ident(sym.nme), erasedType = N)
     sym.tsym = S(tSym)
     FunDefn(owner, sym, tSym, params, body)(configOverride, annotations)
 
@@ -704,7 +705,7 @@ object ValDefn:
       annotations: Ls[Annot],
     )(using State)
     : ValDefn =
-      ValDefn(tsym = TermSymbol(k, owner, Tree.Ident(sym.nme)), sym = sym, rhs = rhs)(configOverride, annotations)
+      ValDefn(tsym = TermSymbol(k, owner, Tree.Ident(sym.nme), erasedType = rhs.erasedType), sym, rhs)(configOverride, annotations)
 
 
 /*
@@ -836,9 +837,83 @@ enum Case:
     case Tup(_, _) => Set.empty
     case Field(_, _) => Set.empty
 
+/** A primitive type of the block IR. */
+enum PrimitiveType:
+  case Unit, Int, Int31, Num, Str, Bool, Array
+
+  /** The symbol for this primitive type. */
+  def sym(using Ctx, State): ClassLikeSymbol = this match
+    case Unit => summon[State].unitSymbol
+    case Int => ctx.builtins.Int
+    case Int31 => ctx.builtins.Int31
+    case Num => ctx.builtins.Num
+    case Str => ctx.builtins.Str
+    case Bool => ctx.builtins.Bool
+    case Array => ctx.builtins.Array
+
+object ErasedType:
+  def ObjectRef: ErasedType.AnyRef = AnyRef(rsc = false, NoSymbol())
+
+  /** Maps a [[`ClassLikeSymbol`]] into the canonical [[`ErasedType`]]. */
+  def fromClsLikeSymbol(csym: ClassLikeSymbol, rsc: Bool)(using Ctx, State): ErasedType = 
+    PrimitiveType.values.find(_.sym === csym) match 
+      case _ if csym === ctx.builtins.Object => ObjectRef
+      case S(prim) => ErasedType.Primitive(prim)
+      case _ => ErasedType.AnyRef(rsc, csym)
+
+/** A generics-erased type of the Block IR. */
+enum ErasedType:
+  /** 
+    * An reference to a class-like symbol.
+    *
+    * If `csym` is `NoSymbol`, this represents the top type (`Object`).
+    *
+    * - `rsc` is true if this reference is a resource class.
+    */
+  case AnyRef(rsc: Bool, csym: ClassLikeSymbol | NoSymbol)
+
+  /** An primitive type. */
+  case Primitive(prim: PrimitiveType)
+
+  /** The symbol for this erased type. */
+  def sym(using Ctx, State): ClassLikeSymbol = this match
+      case AnyRef(_, csym: ClassLikeSymbol) => csym
+      case AnyRef(_, _: NoSymbol) => ctx.builtins.Object
+      case Primitive(prim) => prim.sym
+
+/** Trait representing a Block IR element that has an [[`ErasedType`]]. */
+trait HasErasedType:
+  /** The [[`ErasedType`]] of this element, or `N` if the erased type is not known. */
+  def erasedType: Opt[ErasedType]
+
+  /** Similar to `erasedType`, but coerces to the top type if the specific erased type is not known. */
+  def erasedType_! : ErasedType = erasedType.getOrElse(ErasedType.ObjectRef)
+
+/** A [[`HasErasedType`]] that can have its erased type refined post-construction. */
+trait HasRefinableErasedType extends HasErasedType:
+  // Implementation Note: Provided for overriding classes to implement `erasedType` directly as an `override var`
+  def erasedType_=(newType: Opt[ErasedType]): Unit
+
+  /** Refines the erased type if it was not previously refined, but allowing for idempotent refinements to the same
+    * type.
+    *
+    * A soft assert will be raised if the erased type was already refined to a different type.
+    */
+  def refineErasedType(newType: ErasedType)(using Line, FileName, Raise): Unit =
+    softAssert(erasedType.forall(_ == newType), s"Cannot refine already-refined erased type $erasedType to $newType")
+    if erasedType.isEmpty then erasedType = S(newType)
+
+extension (lit: Literal) 
+  def erasedType: ErasedType = lit match
+    case Tree.UnitLit(_) => ErasedType.Primitive(PrimitiveType.Unit)
+    case Tree.IntLit(_) => ErasedType.Primitive(PrimitiveType.Int)
+    case Tree.DecLit(_) => ErasedType.Primitive(PrimitiveType.Num)
+    case Tree.StrLit(_) => ErasedType.Primitive(PrimitiveType.Str)
+    case Tree.BoolLit(_) => ErasedType.Primitive(PrimitiveType.Bool)
+
 sealed trait TrivialResult extends Result
 
-sealed abstract class Result extends AutoLocated:
+sealed abstract class Result extends AutoLocated, HasErasedType:
 // // * Used for debugging locations:
 // sealed abstract class Result extends AutoLocated with ProductWithExtraInfo:
 //   def extraInfo: Str = toLoc.toString
@@ -923,6 +998,16 @@ sealed abstract class Result extends AutoLocated:
     case Value.Lit(lit) => 0
     case DynSelect(qual, fld, arrayIdx) => qual.size + fld.size
 
+  lazy val erasedType: Opt[ErasedType] = this match
+    case Tuple(_, _) => S(ErasedType.Primitive(PrimitiveType.Array))
+    case Record(_, _) => S(ErasedType.ObjectRef)
+    case Instantiate(_, cls, _) => cls.targetSymbol.flatMap(_.asClsOrMod).map(ErasedType.AnyRef(rsc = false, _))
+    case Value.SimpleRef(sym) => sym.erasedType
+    case Value.MemberRef(_, disamb: (ClassSymbol | ModuleOrObjectSymbol | TypeAliasSymbol)) => disamb.erasedType
+    case Value.This(clsOrMod: (ClassSymbol | ModuleOrObjectSymbol)) => clsOrMod.erasedType
+    case Value.Lit(lit) => S(lit.erasedType)
+    case _ => N
+
 /* mayRaiseEffects indicates whether this call may raise effect (algebraic effect),
  * regardless of whether the check for effect is inserted or not.
  * Note that the check for effect is inserted during HandlerLowering and setting this to true
@@ -992,7 +1077,7 @@ object Value:
       case SimpleRef(l) => l
       case MemberRef(bms, disamb) => bms
       case This(sym) => sym
-
+  
   @deprecated("Use Value.SimpleRef, Value.MemberRef, or Value.This instead.")
   object Ref:
     def apply(l: ValueSymbol | NoSymbol, disamb: Opt[DefinitionSymbol[?]]): Value.RefLike =

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

@@ -1181,7 +1181,7 @@ class BlockSimplifier
                   def go(acc: Block => Block, args: List[(VarSymbol, Result)], mapping: Map[Symbol, Symbol]): Block =
                     args match
                     case Nil =>
-                      val resSym = TempSymbol(N, "inlinedVal")
+                      val resSym = TempSymbol(N, erasedType = N, "inlinedVal")
                       val copier = Copier(resSym, mapping)
                       val newBlk = copier.applyBlock(blk)
                       if extraArgss.isEmpty then
@@ -1190,7 +1190,7 @@ class BlockSimplifier
                         acc(Scoped(Set(resSym), newBlk(
                           k(Call(resSym.asSimpleRef, extraArgss.ne_!)(c.isMlsFun, c.mayRaiseEffects, false)))))
                     case (sym, value) :: argRest =>
-                      val newSym = VarSymbol(sym.id)
+                      val newSym = VarSymbol(sym.id, erasedType = N)
                       go(acc.assignScoped(newSym, value), argRest, mapping + (sym -> newSym))
                   go(blockBuilder, matchedArgs, Map.empty)
           case _ => super.applyResult(r)(k)

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

@@ -22,26 +22,26 @@ class BufferableTransform()(using Ctx, State, Raise):
           cls.bufferable.fold(super.applyDefn(defn)(k)): bufferable =>
             val companionSym = ModuleOrObjectSymbol(DummyTypeDef(syntax.Mod), new Tree.Ident(cls.sym.nme))
             val clsSizeSym = BlockMemberSymbol("size", Nil, false)
-            val clsSizeTermSym = TermSymbol(syntax.ImmutVal, S(companionSym), new Tree.Ident("size"))
+            val clsSizeTermSym = TermSymbol(syntax.ImmutVal, S(companionSym), new Tree.Ident("size"), erasedType = S(ErasedType.Primitive(PrimitiveType.Int)))
             val pubFieldMap: Map[Symbol, Symbol] = cls.publicFields.toMap
             val fields = cls.privateFields ++ cls.publicFields.map(_._2)
             val fieldMap: Map[Symbol, Int] = fields.zipWithIndex.toMap
             def mkSymbolReplacer(params: List[ParamList]): (List[ParamList], Map[SimpleSymbol, SimpleSymbol]) =
               val allVars = params.flatMap(_.allParams).map(_.sym)
               val varMap = allVars
                 .map: sym =>
-                  (sym, VarSymbol(sym.id))
+                  (sym, VarSymbol(sym.id, erasedType = N))
                 .toMap
               def mapParam(p: Param) =
                 Param(p.flags, varMap(p.sym), p.sign, p.modulefulness)
               (params.map(pl => ParamList(pl.flags, pl.params.map(mapParam), pl.restParam.map(mapParam))), varMap.toMap)
             def mkFieldReplacer(buf: VarSymbol, baseIdx: VarSymbol, symMap: Map[SimpleSymbol, SimpleSymbol]) =
               def getOffset(off: Int)(k: Path => Block): Block =
-                val idxSymbol = new TempSymbol(N, "idx")
+                val idxSymbol = new TempSymbol(N, erasedType = S(ErasedType.Primitive(PrimitiveType.Int)), "idx")
                 Scoped(Set.single(idxSymbol), Assign(idxSymbol, Call(State.builtinOpsMap("+").asSimpleRef, (baseIdx.asSimpleRef.asArg :: Value.Lit(Tree.IntLit(off)).asArg :: Nil) ne_:: Nil)(true, false, false),
                   k(DynSelect(buf.asSimpleRef.selSN("buf"), idxSymbol.asSimpleRef, true))))
               def assignToOffset(off: Int, r: Result, rst: Block) =
-                val idxSymbol = new TempSymbol(N, "idx")
+                val idxSymbol = new TempSymbol(N, erasedType = S(ErasedType.Primitive(PrimitiveType.Int)), "idx")
                 Scoped(Set.single(idxSymbol), Assign(idxSymbol, Call(State.builtinOpsMap("+").asSimpleRef, (baseIdx.asSimpleRef.asArg :: Value.Lit(Tree.IntLit(off)).asArg :: Nil) ne_:: Nil)(true, false, false),
                   AssignDynField(buf.asSimpleRef.selSN("buf"), idxSymbol.asSimpleRef, true, r, applyBlock(rst))))
               new BlockTransformer(SymbolSubst.Id):
@@ -73,11 +73,11 @@ class BufferableTransform()(using Ctx, State, Raise):
                         k(res)
                   case _ => super.applyPath(p)(k)
             def transformFunDefn(f: FunDefn, isCtor: Bool): FunDefn =
-              val buf = VarSymbol(new Tree.Ident("buf"))
-              val idx = VarSymbol(new Tree.Ident("idx"))
+              val buf = VarSymbol(new Tree.Ident("buf"), erasedType = N)
+              val idx = VarSymbol(new Tree.Ident("idx"), erasedType = S(ErasedType.Primitive(PrimitiveType.Int)))
               val (newParams, symMap) = mkSymbolReplacer(f.params)
               val blk = mkFieldReplacer(buf, idx, symMap).applyBlock(f.body)
-              FunDefn(f.owner, f.sym, TermSymbol(f.dSym.k, f.dSym.owner, f.dSym.id), PlainParamList(
+              FunDefn(f.owner, f.sym, TermSymbol(f.dSym.k, f.dSym.owner, f.dSym.id, erasedType = N), PlainParamList(
                 Param(FldFlags.empty, buf, N, Modulefulness.none) :: Param(FldFlags.empty, idx, N, Modulefulness.none) :: Nil) :: newParams,
                 if isCtor then Begin(blk, Return(idx.asSimpleRef)) else blk)(configOverride = f.configOverride, annotations = f.annotations)
             val fakeCtor = transformFunDefn(FunDefn.withFreshSymbol(

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

@@ -166,7 +166,7 @@ class Rewrite(val deadParamElimSolver: DeadParamElimSolver)(using Raise):
             val name = instId.mkFunName + s"$$${f.nme}"
             f -> (
               new BlockMemberSymbol(name, Nil, true),
-              new TermSymbol(Fun, N, Tree.Ident(name)))
+              new TermSymbol(Fun, N, Tree.Ident(name), erasedType = N))
           .toMap)
     end mkNewPolyFnSyms
 
@@ -334,12 +334,12 @@ class Rewrite(val deadParamElimSolver: DeadParamElimSolver)(using Raise):
         case ParamList(flags, params, restParam) =>
           val params2 = params.map:
             case p =>
-              val newSym = new VarSymbol(Tree.Ident(p.sym.name))
+              val newSym = new VarSymbol(Tree.Ident(p.sym.name), erasedType = N)
               refreshParamMap(p.sym) = newSym
               Param(p.flags, newSym, p.sign, p.modulefulness)
           val rest2 = restParam.map:
             case p =>
-              val newSym = new VarSymbol(Tree.Ident(p.sym.name))
+              val newSym = new VarSymbol(Tree.Ident(p.sym.name), erasedType = N)
               refreshParamMap(p.sym) = newSym
               Param(p.flags, newSym, p.sign, p.modulefulness)
           ParamList(flags, params2, rest2)

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

@@ -172,7 +172,7 @@ class EtaExpansionRewrite(val etaExpansionSolver: EtaExpansionSolver)(using Rais
           targetShape.drop(existingShape.size).zipWithIndex.map:
             case (count, idx) =>
               val params = (0 until count).toList.map: i =>
-                Param.simple(new VarSymbol(new Tree.Ident(s"eta$$$idx$$$i")))
+                Param.simple(new VarSymbol(new Tree.Ident(s"eta$$$idx$$$i"), erasedType = N))
               EtaParamList(
                 ParamList(ParamListFlags.empty, params, N),
                 params.map(p => Arg(N, p.sym.asSimpleRef)),
@@ -194,7 +194,7 @@ class EtaExpansionRewrite(val etaExpansionSolver: EtaExpansionSolver)(using Rais
             Return(
               Call(fun, (argss ++ activeEtaArgss).ne_!)(c.isMlsFun, c.mayRaiseEffects, c.explicitTailCall))
           case _ =>
-            val tmp = TempSymbol(N, "eta$res")
+            val tmp = TempSymbol(N, erasedType = N, "eta$res")
             Scoped(
               Set.single(tmp),
               Assign(tmp, res2, Return(etaCall(tmp.asPath).withLocOf(res2))))

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

@@ -47,9 +47,19 @@ class FirstClassFunctionTransformer
       case s: TermSymbol if s.k is syntax.Fun =>
         val params = getParamList(l).getOrElse(lastWords(s"Cannot get ${l.nme}'s parameter list."))
         val clsDef = generateFCFunctionClass(ref, params)
-        val tmp = new TempSymbol(None)
+        val tmp = new TempSymbol(None, erasedType = S(ErasedType.AnyRef(rsc = false, clsDef.isym.asClsOrMod.get)))
         val cls = clsDef.sym.asMemberRef(clsDef.isym)
-        Scoped(Set(clsDef.sym, tmp), Define(clsDef, Assign(tmp, Instantiate(false, cls, Nil :: Nil), k(tmp.asSimpleRef))))
+        Scoped(
+          syms = Set(clsDef.sym, tmp),
+          body = Define(
+            clsDef,
+            Assign(
+              tmp,
+              Instantiate(false, cls, Nil :: Nil),
+              k(tmp.asSimpleRef),
+            )
+          )
+        )
       case _ => k(p)
     case sel: Select => sel.symbol match
       case Some(s: TermSymbol) if (s.k is syntax.Fun) =>
@@ -60,9 +70,19 @@ class FirstClassFunctionTransformer
               source = Diagnostic.Source.Compilation)
           PlainParamList(Nil)
         val clsDef = generateFCFunctionClass(sel, params)
-        val tmp = new TempSymbol(None)
+        val tmp = new TempSymbol(None, erasedType = S(ErasedType.AnyRef(rsc = false, clsDef.isym.asClsOrMod.get)))
         val cls = clsDef.sym.asMemberRef(clsDef.isym)
-        Scoped(Set(clsDef.sym, tmp), Define(clsDef, Assign(tmp, Instantiate(false, cls, Nil :: Nil), k(tmp.asSimpleRef))))
+        Scoped(
+          Set(clsDef.sym, tmp),
+          Define(
+            clsDef,
+            Assign(
+              tmp,
+              Instantiate(false, cls, Nil :: Nil),
+              k(tmp.asSimpleRef),
+            )
+          )
+        )
       case Some(_) => k(p)
       case _ =>
         raise(ErrorReport(msg"Cannot determine if ${sel.name.name} is a function." -> sel.toLoc :: Nil,

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

@@ -114,7 +114,7 @@ type StackSafetyMap = collection.Map[FnOrCls, (Int, Block)]
 
 class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise, Elaborator.State, Elaborator.Ctx):
 
-  private def freshTmp(dbgNme: Str = "tmp") = new TempSymbol(N, dbgNme)
+  private def freshTmp(erasedType: Opt[ErasedType], dbgNme: Str = "tmp") = new TempSymbol(N, erasedType, dbgNme)
   private def freshLabel(nme: Str) = new LabelSymbol(N, nme)
 
   private def rtThrowMsg(msg: Str) = Throw(
@@ -134,7 +134,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       case _ => N
 
   object StateTransition:
-    private val transitionSymbol = freshTmp("transition")
+    private val transitionSymbol = freshTmp(erasedType = N, "transition")
     def apply(uid: StateId) =
       Return(PureCall(transitionSymbol.asSimpleRef, List(Value.Lit(Tree.IntLit(uid)))))
     def unapply(blk: Block) = blk match
@@ -143,7 +143,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       case _ => N
 
   object Unwind:
-    private val unwindSymbol = freshTmp("unwind")
+    private val unwindSymbol = freshTmp(erasedType = N, "unwind")
     def apply(uid: StateId, loc: Value) =
       Return(PureCall(unwindSymbol.asSimpleRef, List(Value.Lit(Tree.IntLit(uid)), loc)))
     def unapply(blk: Block) = blk match
@@ -540,7 +540,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
         .flatMap: (sym, idx) =>
           List(intLit(idx), Value.Lit(Tree.StrLit(sym.nme)))
         .map(_.asArg)
-      val debugInfoSym = freshTmp(s"$debugNme$$debugInfo")
+      val debugInfoSym = freshTmp(erasedType = S(ErasedType.Primitive(PrimitiveType.Array)), s"$debugNme$$debugInfo")
       // TODO: properly support spread argument by calculating the correct length.
       val rtArgLists = intLit(fun.params.length) :: fun.params.flatMap: pl =>
         intLit(pl.params.length) :: pl.params.map(p => p.sym.asSimpleRef)
@@ -615,7 +615,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       return b
     val vars = if opt.debug then ctx.resumeInfo.currentLocals else computeRestoreList(parts)
 
-    val pcVar = freshTmp("pc")
+    val pcVar = freshTmp(erasedType = S(ErasedType.Primitive(PrimitiveType.Int)), "pc")
     val mainLoopLbl = freshLabel("main")
 
     val edges = computeEdges(parts)
@@ -690,7 +690,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
           case (acc, f) => f(acc)
         Label(mainLoopLbl, true, matches, End())
 
-    val getSavedTmp = freshTmp("saveOffset")
+    val getSavedTmp = freshTmp(erasedType = S(ErasedType.Primitive(PrimitiveType.Int)), "saveOffset")
     def getSaved(off: BigInt): (Block => Block, Path) =
       if off == 0 then
         return (id, DynSelect(paths.runtimePath.selSN("resumeArr"), paths.runtimePath.selSN("resumeIdx"), true))
@@ -739,7 +739,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       override def applyResult(r: Result)(k: Result => Block) = r match
         case EffectfulResult(r) =>
           // Fallback case, this may lead to unnecessary assignments if it is assign-like
-          val l = freshTmp()
+          val l = freshTmp(erasedType = N)
           Scoped(Set(l), effectCheck(l, r, k(l.asSimpleRef)))
         case _ => super.applyResult(r)(k)
     topLevelTransform.applyBlock(b)