Various Bug Fixes and Tailrec Cleanup

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

@@ -682,6 +682,7 @@ final case class FunDefn(
   lazy val affineInfo: Ls[Int] =
     annotations.collect:
       case Annot.Affine(whichParamList) => whichParamList
+  lazy val allParamSyms: Ls[VarSymbol] = params.flatMap(_.paramSyms)
 
   // * This deliberately is not a lazy val: its initialization would synchronize on the JVM.
   // * Computing the summary is pure, and a reference write is atomic, so concurrent traversals may

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

@@ -13,7 +13,6 @@ import hkmc2.syntax.{Tree, SpreadKind}
 import hkmc2.ScopeData.*
 import hkmc2.Lifter.AccessInfo
 import scala.collection.mutable.{ArrayBuffer, LinkedHashMap}
-import java.lang.instrument.ClassDefinition
 
 /*
 
@@ -207,11 +206,6 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
       val x = f(item)
       if x > l then x else l
 
-  def getParamSyms(f: FunDefn) = f.params.flatMap:
-    case ParamList(_, params, S(rest)) =>
-      params.map(_.sym).appended(rest.sym)
-    case p => p.params.map(_.sym)
-
   def paramsLen(f: FunDefn): Int = f.params.foldLeft(0): (acc, head) =>
     acc + (if head.restParam.isDefined then 1 + head.params.length
       else head.params.length)
@@ -249,7 +243,55 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
       CallArgsResult.Success(hd.appended(rest))
     else
       CallArgsResult.Success(hd)
-
+
+  /**
+    * Applies a parameter list to an argument list, both of which could contain arbitrary spreads.
+    *
+    * @param plist The parameter list to apply.
+    * @param args The argument list.
+    * @return A tuple `(Block => Block, List[Path])`, where:
+    * - The first return value contains the code that breaks up the arguments into the shape expected by the parameter list, and
+    *   should be applied to the remaining code;
+    * - The second return value is a list of paths containing the arguments in order. If there is a spread parameter, it will be
+    *   the last one.
+    */
+  def forceSpread(plist: ParamList, args: List[Arg]): (Block => Block, List[Path]) =
+    // Forcibly spread the args in an array.
+    // Assume the lengths are correct
+    val paramList = plist.params
+    val restParam = plist.restParam
+
+    val tupleSym = TempSymbol(N, "argList")
+
+    val tupleRes = Tuple(false, args)
+
+    // Main args
+    def mainArgs(rest: List[Path]) = (0 until paramList.size).toList.foldRight(rest):
+      case (n, acc) => DynSelect(tupleSym.asSimpleRef, Value.Lit(Tree.IntLit(n)), true) :: acc
+
+    // If the rest param exists, append a slice
+    val (initialBlk: (Block => Block), pathList: List[Path]) =
+      if restParam.isDefined then
+        val sliceResSym = TempSymbol(N, "sliceRes")
+        // runtime.Tuple.slice(tupleSym, paramList.length, 0)
+        val sliceRes = Call(
+          State.runtimeSymbol.asSimpleRef
+            .sel(Tree.Ident("Tuple"), State.tupleSymbol)
+            .sel(Tree.Ident("slice"), State.tupleSliceSymbol),
+          (tupleSym.asSimpleRef.asArg
+            :: Value.Lit(Tree.IntLit(paramList.length)).asArg
+            :: Value.Lit(Tree.IntLit(0)).asArg
+            :: Nil) ne_:: Nil
+        )(CallMetadata.defaultMlsFun)
+        val blk = blockBuilder
+          .assignScoped(tupleSym, tupleRes)
+          .assignScoped(sliceResSym, sliceRes)
+        (blk, mainArgs(sliceResSym.asSimpleRef :: Nil))
+      else
+        (blockBuilder.assignScoped(tupleSym, tupleRes), mainArgs(Nil))
+    end val
+    (initialBlk, pathList)
+
   def optScc(scc: SccOfCalls, owner: Opt[InnerSymbol])(using accessInfo: (ScopeData, AccessMap)): (Opt[FunDefn], List[FunDefn]) =
     // sort the functions so the order is more predictable
     val funs = scc.funs.sortBy(f => f.dSym.uid)
@@ -283,18 +325,28 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
             :: msg"It could self-recurse through this call, which is not a tail call." -> reportLoc
             :: Nil
           ))
-
+
+    val (hasWrapper, bms) =
+      if funsLen === 1 then funs.head.params match
+        // Simplest case: One function, one parameter list, no wrapper.
+        case _ :: Nil => (false, funs.head.sym)
+        // When a function has multiple param lists, TailRecOpt flattens them into a single
+        // param list for the internal loop. We need a wrapper function that preserves the
+        // original multi-param-list interface and delegates to the flattened internal loop.
+        case _ => (true, BlockMemberSymbol(funs.head.sym.nme + "$tailrec", Nil, true))
+      // Multiple functions: Create a wrapper.
+      else (true, BlockMemberSymbol(funs.iterator.map(_.sym.nme).mkString("_"), Nil, true))
+    val dSym =
+      if !hasWrapper then funs.head.dSym
+      else TermSymbol(syntax.Fun, owner, Tree.Ident(bms.nme))
+
     val maxParamLen = maxInt(funs, paramsLen)
     val paramSyms =
-        if funs.length === 1 then
-          val syms = getParamSyms(funs.head)
-          if funs.head.params.length === 1 then syms
-          else
-            // Duplicate the params for the internal loop defn (see the doc at the
-            // end of this function), but preserve the names.
-            syms.map(v => VarSymbol(Tree.Ident(v.id.name)))
-        else
-          for i <- 0 until maxParamLen yield VarSymbol(Tree.Ident("param" + i))
+        if hasWrapper then 
+          if funsLen === 1 then
+            funs.head.allParamSyms.map(v => VarSymbol(Tree.Ident(v.id.name)))
+          else for i <- 0 until maxParamLen yield VarSymbol(Tree.Ident("param" + i))
+        else funs.head.allParamSyms
       .toList
     val paramSymsArr = ArrayBuffer.from(paramSyms)
     // Function -> param -> param symbol in the rewritten function
@@ -309,17 +361,29 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
 
     val dSymIds = funs.iterator.map(_.dSym).zipWithIndex.toMap
     val dSymToDefn = funs.iterator.map(f => f.dSym -> f).toMap
-    val bms =
-      if funsLen === 1 then funs.head.sym
-      else BlockMemberSymbol(funs.iterator.map(_.sym.nme).mkString("_"), Nil, true)
-    val dSym =
-      if funsLen === 1 then funs.head.dSym
-      else TermSymbol(syntax.Fun, owner, Tree.Ident(bms.nme))
+
     val loopSym = LabelSymbol(N, "loopLabel")
     val curIdSym = VarSymbol(Tree.Ident("id"))
 
+    val loopDefnPath = owner match
+      case Some(value) => Select(value.asThis, Tree.Ident(bms.nme))(S(dSym))(false)
+      case None => bms.asMemberRef(dSym)
+
+    def rewriteKnownCall(callee: FunDefn, flattenedArgs: List[Path]) =
+      if !hasWrapper then lastWords("Tried to rewrite a call to a non-existent loop definition.")
+      if callee.allParamSyms.length =/= flattenedArgs.length then lastWords("Incorrect function call arity.")
+      // Fill in the argument list with the arguments in the correct position by applying the paramSym -> arg map to
+      // the loop definitions parameter list.
+      val argsMap = (callee.allParamSyms.map(paramSymsMap(callee.dSym)) zip flattenedArgs).toMap
+      val args = paramSyms.map: s =>
+        argsMap.get(s) match
+          case Some(pth) => Arg(N, pth)
+          case None => Arg(N, Value.Lit(Tree.UnitLit(false)))
+      val argsWithId = if funsLen > 1 then Value.Lit(Tree.IntLit(dSymIds(callee.dSym))).asArg :: args else args
+      Call(loopDefnPath, argsWithId ne_:: Nil)(CallMetadata.defaultMlsFun)
+
     class FunRewriter(f: FunDefn) extends BlockTransformerShallow(SymbolSubst.Id):
-      val params = getParamSyms(f)
+      val params = f.allParamSyms
       val paramsSet = f.params.toSet
       val paramsIdxes = params.zipWithIndex.toMap
 
@@ -445,45 +509,12 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
                         case x => x
                       ret
                   case CallArgsResult.ForceSpread =>
-                    // Forcibly spread the args in an array.
-                    // Assume the lengths are correct
-                    val paramList = ogParamList.params
-                    val restParam = ogParamList.restParam
-
-                    val tupleSym = TempSymbol(N, "argList")
-
                     // We can safely remove all of the symbols from this parameter list from `assignedSyms` at this stage,
                     // because the RHS of every parameter will be computed when spreading them in the tuple, which happens
                     // before any of the param symbols are assigned to.
                     assignedSyms --= thisParamSyms
                     paramRewriter.applyArgs(ogArgs): newArgs =>
-                      val tupleRes = Tuple(false, newArgs)
-
-                      // Main args
-                      def mainArgs(rest: List[Path]) = (0 until paramList.size).toList.foldRight(rest):
-                        case (n, acc) => DynSelect(tupleSym.asSimpleRef, Value.Lit(Tree.IntLit(n)), true) :: acc
-
-                      // If the rest param exists, append a slice
-                      val (initialBlk: (Block => Block), pathList: List[Path]) =
-                        if restParam.isDefined then
-                          val sliceResSym = TempSymbol(N, "sliceRes")
-                          // runtime.Tuple.slice(tupleSym, paramList.length, 0)
-                          val sliceRes = Call(
-                            State.runtimeSymbol.asSimpleRef
-                              .sel(Tree.Ident("Tuple"), State.tupleSymbol)
-                              .sel(Tree.Ident("slice"), State.tupleSliceSymbol),
-                            (tupleSym.asSimpleRef.asArg
-                              :: Value.Lit(Tree.IntLit(paramList.length)).asArg
-                              :: Value.Lit(Tree.IntLit(0)).asArg
-                              :: Nil) ne_:: Nil
-                          )(CallMetadata.defaultMlsFun)
-                          val blk = blockBuilder
-                            .assignScoped(tupleSym, tupleRes)
-                            .assignScoped(sliceResSym, sliceRes)
-                          (blk, mainArgs(sliceResSym.asSimpleRef :: Nil))
-                        else
-                          (blockBuilder.assignScoped(tupleSym, tupleRes), mainArgs(Nil))
-                      end val
+                      val (initialBlk, pathList) = forceSpread(ogParamList, newArgs)
                       val paramAssignments = (thisParamSyms zip pathList).foldRight[Block](rest):
                         case ((sym, path), restBlk) => Assign(sym, path, restBlk)
 
@@ -498,6 +529,41 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
         // Not a tail call
         case _ => super.applyBlock(b)
 
+      // We explicitly rewrite *all* fully applied calls to functions within the SCC. Tail calls are rewritten using a
+      // `continue`, while other calls are rewritten as a call to the merged function. This is to make the tailrec
+      // optimizer pass idempotent. Without this, the subsequent tailrec optimizer passes could see the merged and wrapper
+      // functions as an SCC and try to rewrite the wrapper's tail call. Note that selections (and thus calls to
+      // parameter-less module methods) are not yet supported.
+      override def applyResult(r: Result)(k: Result => Block): Block = r match
+        case c @ CallToFun(calleeSym) if hasWrapper => funsMap.get(calleeSym) match
+          case Some(callee) if isExactlySaturatedCall(c, callee) =>
+            val calleeParamsMap = paramSymsMap(callee.dSym)
+            val argListResults = callee.params.zip(c.argss).map:
+              case (params, args) => (params, params.paramSyms.map(calleeParamsMap), args, rewriteArgsList(callee, params, args))
+            // Note: even though the convention right now is to just have the first k_f arguments of the merged function be the arguments
+            // of each function f, where f has k_f total parameters, we still explicitly find the correct order of arguments without
+            // using this convention. This is because this assumption may not hold in the future, i.e. when arguments could be assigned
+            // types for the typed IR.
+            var pre = blockBuilder
+            // Create the list of flattened args.
+            val argsInOrder: List[Path] = argListResults.flatMap: (ogParams, thisParamSyms, args, argsRes) =>
+              argsRes match
+                case CallArgsResult.Success(res) => res.map:
+                  case r: Path => r
+                  case r: Result =>
+                    val newSym = TempSymbol(N)
+                    pre = pre.assignScoped(newSym, r)
+                    newSym.asPath
+                case CallArgsResult.ForceSpread =>
+                  val (initialBlk, paths) = forceSpread(ogParams, args)
+                  pre = pre.chain(initialBlk)
+                  paths
+            pre.rest(k(rewriteKnownCall(callee, argsInOrder)))
+
+          case _ => super.applyResult(r)(k)
+        case _ => super.applyResult(r)(k)
+
+
       def rewrite(b: Block): Block =
         // Rewrite the result with symbols pointing to the merged function parameters and possibly the copied parameters (see `copiedParams`).
         val blk = applyBlock(symRewriter.applyBlock(b))
@@ -514,67 +580,35 @@ class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
 
     val loop = Label(loopSym, true, switch, End())
 
-    val sel = owner match
-      case Some(value) => Select(value.asThis, Tree.Ident(bms.nme))(S(dSym))(false)
-      case None => bms.asMemberRef(dSym)
-
-    val rewrittenFuns =
-      if funsLen === 1 then Nil
-      else funs.map: f =>
-        val paramArgs = getParamSyms(f).map(s => s.asSimpleRef.asArg)
-        val args =
-          Value.Lit(Tree.IntLit(dSymIds(f.dSym))).asArg
-            :: paramArgs
-            ::: List.fill(maxParamLen - paramArgs.length)(Value.Lit(Tree.UnitLit(false)).asArg)
+    if !hasWrapper then
+      val f = funs.head
+      val defn = FunDefn(
+        owner, bms, dSym,
+        f.params,
+        loop)(N, annotations = f.annotations)
+      (N, defn :: Nil)
+    else
+      val wrappers = funs.map: f =>
+        val paramArgs = f.allParamSyms.map(s => s.asSimpleRef)
         val newBod = Return(
-          Call(sel, args ne_:: Nil)(CallMetadata.defaultMlsFun),
+          rewriteKnownCall(f, paramArgs),
         )
-        FunDefn(f.owner, f.sym, f.dSym, f.params, newBod)(N, f.annotations)
-
-    funs match
-      case (f @ FunDefn(params = _ :: Nil | Nil)) :: Nil =>
-        val defn = FunDefn(
-          owner, bms, dSym,
-          f.params,
-          loop)(N, annotations = f.annotations)
-        (N, defn :: Nil)
-      case f :: Nil =>
-        // When a function has multiple param lists, TailRecOpt flattens them into a single
-        // param list for the internal loop. We need a wrapper function that preserves the
-        // original multi-param-list interface and delegates to the flattened internal loop.
-        val params = paramSyms.map(Param.simple(_))
-
-        val loopBms = BlockMemberSymbol(bms.nme + "$tailrec", Nil, true)
-        val loopDSym = TermSymbol(syntax.Fun, owner, Tree.Ident(loopBms.nme))
-        val loopAnnots =
-          if f.inline then Annot.Inline :: Annot.Private :: Nil
-          else Annot.Private :: Nil
-        val internalLoopDefn = FunDefn(
-          owner, loopBms, loopDSym,
-          PlainParamList(params) :: Nil,
-          loop)(N, annotations = loopAnnots)
-        val paramArgs = getParamSyms(f).map(s => s.asSimpleRef.asArg)
-        val internalSel = owner match
-          case Some(value) => Select(value.asThis, Tree.Ident(loopBms.nme))(S(loopDSym))(false)
-          case None => loopBms.asMemberRef(loopDSym)
-        val wrapperBod = Return(
-          Call(internalSel, paramArgs ne_:: Nil)(CallMetadata.defaultMlsFun),
-        )
-        val newAnnots = if f.inline then f.annotations else Annot.Inline :: f.annotations
-        val wrapperDefn = FunDefn(f.owner, f.sym, f.dSym, f.params, wrapperBod)(
-          f.configOverride, annotations = newAnnots)
-        (S(internalLoopDefn), wrapperDefn :: Nil)
-      case _ =>
-        val newParamLists =
-          val initial = paramSyms.map(Param.simple(_))
+        val annots = if f.inline then f.annotations else Annot.Inline :: f.annotations 
+        FunDefn(f.owner, f.sym, f.dSym, f.params, newBod)(N, annots)
+      val newParamLists =
+        val initial = paramSyms.map(Param.simple(_))
+        if funsLen > 1 then
           PlainParamList(Param.simple(curIdSym) :: initial) :: Nil
-
-        val loopDefn = FunDefn(
-          owner, bms, dSym,
-          newParamLists,
-          loop)(N, annotations = Nil) // Q: maybe should be Private?
-
-        (S(loopDefn), rewrittenFuns)
+        else
+          PlainParamList(initial) :: Nil
+      val annotations = 
+        if funsLen == 1 && funs.head.inline then Annot.Inline :: Annot.Private :: Nil
+        else Annot.Private :: Nil
+      val loopDefn = FunDefn(
+        owner, bms, dSym,
+        newParamLists,
+        loop)(N, annotations = annotations)
+      (S(loopDefn), wrappers)
 
   def optFunctions(fs: List[FunDefn], owner: Opt[InnerSymbol])(using (ScopeData, AccessMap)) =
     val (newFsOpt, fsOpt) = partFns(fs).map(optScc(_, owner)).foldLeft[(List[FunDefn], List[FunDefn])](Nil, Nil):