Make tail-rec optimization improvements

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

@@ -969,6 +969,16 @@ case class Call(fun: Path, argss: NELs[Ls[Arg]])(val metadata: CallMetadata) ext
         case fd: FunDefn => argss.lengthCompare(fd.params.length) < 0
         case _ => false
     case _ => false
+  // `metadata` lives in a secondary constructor list, so case-class equality
+  // would otherwise ignore annotations such as @tailcall.
+  override def equals(obj: Any): Bool = obj match
+    case that: Call =>
+      fun == that.fun &&
+        argss == that.argss &&
+        metadata == that.metadata
+    case _ => false
+  override def hashCode: Int =
+    (fun, argss, metadata).hashCode
 
 object Call:
 

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

@@ -55,10 +55,19 @@ class CompilationPipeline(using Config, Raise, State, Ctx, SymbolPrinter):
       else prog
     runPass("ClassParamFlattener")(ClassParamFlattener.apply)
     runPass("ReflectionInstrumenter")(ReflectionInstrumenter(using summon).apply)
-    runPass("TailRecOpt")(TailRecOpt().transform)
     preOptimizeHook(result)
+
+    // * We run this pass here first, before inlining so that the @tailrec/@tailcall annotations
+    // * can be properly checked.
+    runPass("TailRecOpt")(TailRecOpt(true).transform)
+
     runPass("WorkerWrapper")(WorkerWrapper(symbolsToPreserve, otl, printer))
     runPass("BlockSimplifier")(BlockSimplifier(symbolsToPreserve, otl, printer).apply)
     runPass("DeadParamElim")(otl.givenIn(DeadParamElim.apply))
 
+    // * More tailrec opportunities might be revealed after WorkerWrapper + BlockSimplifier,
+    // * which might bring split curried recursive calls (such as those coming out of Deforest + EtaExpansion)
+    // * into proper tail positions.
+    runPass("TailRecOpt")(TailRecOpt(false).transform)
+
     result

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

@@ -68,10 +68,24 @@ connected component are tail calls.
 */
 
 // This optimization assumes the lifter has been run.
-class TailRecOpt(using State, TL, Raise):
+class TailRecOpt(checkAnnotations: Bool)(using State, TL, Raise):
 
   type AccessMap = Map[ScopedInfo, AccessInfo]
 
+  private def isExactlySaturatedCall(c: Call, f: FunDefn): Bool =
+    val cmp = c.argss.sizeCompare(f.params.size)
+    // A zero-argument-list definition may still be invoked by a `Call` node:
+    // the callee body is evaluated as a nullary thunk, and the call's argument
+    // lists are then applied to the returned value. For non-nullary callees,
+    // passing more argument lists than the callee can receive violates the IR
+    // invariant that a `Call` node does not apply the result of a non-nullary
+    // callee.
+    softAssert(
+      cmp <= 0 || f.params.isEmpty,
+      s"Call node passes ${c.argss.size} argument lists to ${f.dSym.showDbg}, which can receive ${f.params.size}.",
+    )
+    cmp === 0
+
   object CallToFun:
     def unapply(c: Call): Opt[TermSymbol] = c match
       case Call(fun = Value.MemberRef(_, r: TermSymbol)) => S(r)
@@ -115,24 +129,34 @@ class TailRecOpt(using State, TL, Raise):
     override def applyBlock(b: Block): Unit = b match
       case TailCallShape(r, c) => getFun(r) match
         case Some(value) =>
-          if c.argss.size != value.params.size then
-            if c.metadata.explicitTailCall then
+          // Only exactly saturated calls can be rewritten as direct loop jumps.
+          // Under-applied calls only build closures for later argument lists;
+          // over-applied calls cannot be rewritten as direct jumps either.
+          if isExactlySaturatedCall(c, value) then
+            edges ::= CallEdge.TailCall(f.dSym, r)(c)
+          else
+            if checkAnnotations && c.metadata.explicitTailCall then
               raise(ErrorReport(msg"Only fully applied calls may be marked @tailcall." -> c.toLoc :: Nil))
-          else edges ::= CallEdge.TailCall(f.dSym, r)(c)
         case None =>
-          if c.metadata.explicitTailCall then
+          if checkAnnotations && c.metadata.explicitTailCall then
             raise(ErrorReport(msg"Only functions in this compilation unit may be marked @tailcall." -> c.toLoc :: Nil))
       case Return(c: Call) =>
-        if c.metadata.explicitTailCall then
+        if checkAnnotations && c.metadata.explicitTailCall then
           raise(ErrorReport(msg"Only direct calls in tail position may be marked @tailcall." -> c.toLoc :: Nil))
       case _ => super.applyBlock(b)
 
     override def applyResult(r: Result): Unit = r match
       case c: Call =>
-        if c.metadata.explicitTailCall then
+        if checkAnnotations && c.metadata.explicitTailCall then
           raise(ErrorReport(msg"This call is not in tail position." -> c.toLoc :: Nil))
         c match
-          case CallToFun(r) => edges ::= CallEdge.NormalCall(f.dSym, r)(c)
+          case CallToFun(r) => getFun(r) match
+            // Under-applied curried calls only build closures for later argument
+            // lists; they do not execute the callee body and therefore do not
+            // form recursive call-graph edges.
+            case Some(value) if isExactlySaturatedCall(c, value) =>
+              edges ::= CallEdge.NormalCall(f.dSym, r)(c)
+            case _ =>
           case _ =>
       case _ => super.applyResult(r)
 
@@ -150,7 +174,7 @@ class TailRecOpt(using State, TL, Raise):
     val cg = buildCallGraph(fs).filter: c =>
       val cond = defnSyms.contains(c.f1) && defnSyms.contains(c.f2)
       c.match
-        case c: CallEdge.TailCall if c.call.metadata.explicitTailCall && !cond =>
+        case c: CallEdge.TailCall if checkAnnotations && c.call.metadata.explicitTailCall && !cond =>
           raise(ErrorReport(
             msg"This tail call exits the current scope and is not optimized." -> c.call.toLoc :: Nil))
         case _ =>
@@ -167,7 +191,7 @@ class TailRecOpt(using State, TL, Raise):
       .groupBy: c =>
         val s1 = sccMap(c.f1)
         val s2 = sccMap(c.f2)
-        if s1 =/= s2 && c.call.metadata.explicitTailCall then
+        if checkAnnotations && s1 =/= s2 && c.call.metadata.explicitTailCall then
           raise(ErrorReport(
             msg"This call is not optimized as it does not directly recurse through its parent function." -> c.call.toLoc :: Nil))
           -1
@@ -247,13 +271,13 @@ class TailRecOpt(using State, TL, Raise):
       case c: CallEdge.NormalCall => c.f2 -> c.call
     val nonTailCalls = nonTailCallsLs.toMap
 
-    if nonTailCallsLs.sizeCompare(calls) === 0 then
-      for f <- funs if f.tailRec do
-        raise(WarningReport(msg"This function does not directly self-recurse, but is marked @tailrec." -> f.dSym.toLoc :: Nil))
+    if calls.isEmpty then
+      for f <- funs if checkAnnotations && f.tailRec do
+        raise(WarningReport(msg"This function is marked @tailrec but has no apparent tail calls." -> f.dSym.toLoc :: Nil))
       return (N, funs)
 
     if !nonTailCalls.isEmpty then
-      for f <- funs if f.tailRec do
+      for f <- funs if checkAnnotations && f.tailRec do
         val reportLoc = nonTailCalls.get(f.dSym) match
           // always display a call to f, if possible
           case Some(value) => value.toLoc 
@@ -263,6 +287,7 @@ class TailRecOpt(using State, TL, Raise):
             :: msg"It could self-recurse through this call, which is not a tail call." -> reportLoc
             :: Nil
           ))
+      return (N, funs)
 
     val maxParamLen = maxInt(funs, paramsLen)
     val paramSyms =
@@ -350,11 +375,11 @@ class TailRecOpt(using State, TL, Raise):
           case None => super.applyBlock(b)
           case Some(id) =>
             val callee = dSymToDefn(calleeSym)
-            val calleeParamsMap = paramSymsMap(callee.dSym)
             // We require the call to be fully applied.
-            if c.argss.size != callee.params.size then
+            if !isExactlySaturatedCall(c, callee) then
               super.applyBlock(b)
             else
+              val calleeParamsMap = paramSymsMap(callee.dSym)
               // The code used to continute the loop.
               val cont =
                 if funsLen === 1 then Continue(loopSym)
@@ -585,7 +610,7 @@ class TailRecOpt(using State, TL, Raise):
     // Class methods cannot yet be optimized as they cannot yet be marked final.
 
     if c.k is syntax.Cls then
-      reportClassesTailrec(c)
+      if checkAnnotations then reportClassesTailrec(c)
       val companion = c.companion.mapConserve: comp =>
         val cMtds = optFunctionsFlat(comp.methods, S(comp.isym))
         if cMtds is comp.methods
@@ -659,20 +684,21 @@ class TailRecOpt(using State, TL, Raise):
       optFNew.foldLeft(transformer.applyBlock(b)):
         case (acc, f) => Define(f, acc))
 
-    // Report @tailrec on functions that weren't processed by the optimization above,
-    // e.g. nested functions or functions with @config(tailRecOpt: false).
-    // Class/module methods are handled separately by optClasses and are skipped here.
-    val tailRecFunSyms = tailRecFuns.map(_.dSym).toSet
-    new BlockTraverser:
-      override def applyFunDefn(fun: FunDefn): Unit =
-        if fun.tailRec && !tailRecFunSyms.contains(fun.dSym) then
-          raise(ErrorReport(
-            msg"This @tailrec function was not processed by the tail-call optimizer." -> fun.dSym.toLoc :: Nil))
-        super.applyFunDefn(fun)
-      override def applyDefn(defn: Defn): Unit = defn match
-        case _: ClsLikeDefn => ()
-        case _ => super.applyDefn(defn)
-    .applyBlock(result)
+    if checkAnnotations then
+      // Report @tailrec on functions that weren't processed by the optimization above,
+      // e.g. nested functions or functions with @config(tailRecOpt: false).
+      // Class/module methods are handled separately by optClasses and are skipped here.
+      val tailRecFunSyms = tailRecFuns.map(_.dSym).toSet
+      new BlockTraverser:
+        override def applyFunDefn(fun: FunDefn): Unit =
+          if fun.tailRec && !tailRecFunSyms.contains(fun.dSym) then
+            raise(ErrorReport(
+              msg"This @tailrec function was not processed by the tail-call optimizer." -> fun.dSym.toLoc :: Nil))
+          super.applyFunDefn(fun)
+        override def applyDefn(defn: Defn): Unit = defn match
+          case _: ClsLikeDefn => ()
+          case _ => super.applyDefn(defn)
+      .applyBlock(result)
 
     if result is b
     then prog