Tail Recursion Optimization

compiler/shared/main/scala/mlscript/compiler/ir/Builder.scala

@@ -307,7 +307,10 @@ final class Builder(fresh: Fresh, fnUid: FreshInt, classUid: FreshInt, tag: Fres
 
       import scala.collection.mutable.{ HashSet => MutHSet }
 
-      val cls = grouped.getOrElse(0, Nil).map(buildClassInfo)
+      val cls = ClassInfo(classUid.make, "True", List()) // TODO: add "True" amd "False" at some pointgrouped.getOrElse(0, Nil).map(buildClassInfo)
+        :: ClassInfo(classUid.make, "False", List()) 
+        :: grouped.getOrElse(0, Nil).map(buildClassInfo)
+
       cls.foldLeft(Set.empty)(checkDuplicateField(_, _))
 
       val clsinfo = cls.toSet
@@ -342,10 +345,5 @@ final class Builder(fresh: Fresh, fnUid: FreshInt, classUid: FreshInt, tag: Fres
 
       resolveDefnRef(main, defs, true)
       validate(main, defs)
-
-      // TODO: should properly import built-in types
-      val clsWithBool = clsinfo
-        + ClassInfo(classUid.make, "True", List())
-        + ClassInfo(classUid.make, "False", List())
 
       Program(clsinfo, defs, main)

compiler/shared/main/scala/mlscript/compiler/ir/IR.scala

@@ -112,7 +112,7 @@ enum Expr:
     case Literal(IntLit(lit)) => s"$lit" |> raw
     case Literal(DecLit(lit)) => s"$lit" |> raw
     case Literal(StrLit(lit)) => s"$lit" |> raw
-    case Literal(UnitLit(lit)) => s"$lit" |> raw
+    case Literal(UnitLit(lit)) => (if lit then "undefined" else "null") |> raw
     case CtorApp(ClassInfo(_, name, _), args) =>
       raw(name) <#> raw("(") <#> raw(args |> show_args) <#> raw(")")
     case Select(s, _, fld) =>
@@ -297,7 +297,7 @@ enum LocMarker:
     case MLit(IntLit(lit)) => s"$lit" |> raw
     case MLit(DecLit(lit)) => s"$lit" |> raw
     case MLit(StrLit(lit)) => s"$lit" |> raw
-    case MLit(UnitLit(lit)) => s"$lit" |> raw
+    case MLit(UnitLit(lit)) => (if lit then "undefined" else "null") |> raw
     case _ => raw("...")
 
   def show = s"$tag-" + toDocument.print

compiler/shared/main/scala/mlscript/compiler/optimizer/TailRecOpt.scala

@@ -9,6 +9,62 @@ import Message.MessageContext
 import compiler.ir._
 import compiler.ir.Node._
 
+/*
+
+DOCUMENTATION OF SEMANTICS OF @tailcall and @tailrec
+
+@tailcall: Used to annotate specific function calls. Calls annotated with @tailcall 
+must be tail calls or tail modulo-cons calls. These calls may be optimized to not
+consume additional stack space. If such an optimization is not possible, then the
+compiler will throw an error.
+
+If there are multiple possible candidates for tail modulo-cons calls in a single
+branch of an expression, then @tailcall can be uesd to indicate which one will be
+optimized. For instance in
+
+fun foo() =
+  A(foo(), bar())
+
+we can use @tailcall to annotate the call foo() or bar(). If a call other than the
+last call is annotated with @tailcall, then the remaining functions must be pure
+to ensure that reordering computations does not change the result.
+
+If bar() is impure but you still want to optimize the call foo(), then you can do
+
+fun foo() =
+  let b = bar()
+  let a = @tailcall foo()
+  A(a, b)
+
+because here, you are taking responsibility for the reordering of the computations.
+
+@tailrec: Used to annotate functions. When this annotation is used on a function, say
+@tailrec fun foo(), the compiler will ensure no sequence of recursive calls back to foo()
+consume stack space, i.e. they are all tail calls. Note that a call to foo() may
+consume an arbitrary amount of stack space as long as foo() is only consuming finite
+stack space. For example,
+
+@tailrec fun foo() = bar()
+fun bar() =
+  bar()
+  bar()
+
+is valid. However,
+
+@tailrec fun foo() = bar()
+fun bar() =
+  foo()
+  bar()
+
+is invalid. If we swap the position of foo() and bar() in the body of bar, it is still invalid.
+
+Equivalently, if fun foo() is annotated with @tailrec, let S be the largest strongly
+connected component in the call-graph of the program that contains foo. Then an error
+will be thrown if and only if all edges (calls) connecting the nodes of the strongly
+connected component are tail calls or tail modulo-cons calls.
+
+*/
+
 // fnUid should be the same FreshInt that was used to build the graph being passed into this class
 class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diagnostic => Unit):
   case class LetCtorNodeInfo(node: LetExpr, ctor: Expr.CtorApp, cls: ClassInfo, ctorValName: Name, fieldName: String, idx: Int)
@@ -583,12 +639,14 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
       val trueClass = classes.find(c => c.ident == "True").get
       val falseClass = classes.find(c => c.ident == "False").get
 
-      // CONOTEXT APPLICATION
-
+      // CONTEXT APPLICATION
+      
       val mergedNames = defns.foldLeft("")(_ + "_" + _.name)
 
-      val ctxAppName = mergedNames + "_ctx_app"
-      val ctxCompName = mergedNames + "_ctx_comp"
+      val ctxAppId = fnUid.make
+      val ctxAppName = mergedNames + "_ctx_app$" + ctxAppId
+      val ctxCompId = fnUid.make
+      val ctxCompName = mergedNames + "_ctx_comp$" + ctxCompId
 
       // map integers to classes and fields which will be assigned to
       val classIdMap = classes.map(c => c.id -> c).toMap
@@ -649,7 +707,7 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
         )
       ).attachTag(tag)
 
-      val appDefn = Defn(fnUid.make, ctxAppName, List(appCtxName, appValName), 1, appNode, false)
+      val appDefn = Defn(ctxAppId, ctxAppName, List(appCtxName, appValName), 1, appNode, false)
 
       // CONTEXT COMPOSITION
       val cmpCtx1Name = Name("ctx1")
@@ -690,7 +748,7 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
         ).attachTag(tag)
       ).attachTag(tag)
 
-      val cmpDefn = Defn(fnUid.make, ctxCompName, List(cmpCtx1Name, cmpCtx2Name), 1, cmpNode, false)
+      val cmpDefn = Defn(ctxCompId, ctxCompName, List(cmpCtx1Name, cmpCtx2Name), 1, cmpNode, false)
 
       // We use tags to identify nodes
       // a bit hacky but it's the most elegant way
@@ -790,7 +848,8 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
 
       def rewriteDefn(d: Defn): Defn =
         val transformed = transformNode(d.body)
-        Defn(fnUid.make, d.name + "_modcons", Name("ctx") :: d.params, d.resultNum, transformed, d.isTailRec)
+        val id = fnUid.make
+        Defn(id, d.name + "_modcons$" + id, Name("ctx") :: d.params, d.resultNum, transformed, d.isTailRec)
 
       // returns (new defn, mod cons defn)
       // where new defn has the same signature and ids as the original, but immediately calls the mod cons defn
@@ -843,6 +902,8 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
     // To build the case block, we need to compare integers and check if the result is "True"
     val trueClass = classes.find(c => c.ident == "True").get
     val falseClass = classes.find(c => c.ident == "False").get
+    // undefined for dummy values
+    val dummyVal = Expr.Literal(UnitLit(true))
 
     // join points need to be rewritten. For now, just combine them into the rest of the function. They will be inlined anyways
     val defns = component.nodes ++ component.joinPoints
@@ -913,10 +974,10 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
 
       val stackFrame = trName :: defnsList.flatMap(d => d.params.map(n => nameMaps(d.id)(n))) // take union of stack frames
 
-      // TODO: This works fine for now, but ideally should find a way to guarantee the new
-      // name is unique
-      val newName = defns.foldLeft("")(_ + "_" + _.name) + "_opt"
-      val jpName = defns.foldLeft("")(_ + "_" + _.name) + "_opt_jp"
+      val newId = fnUid.make
+      val newName = defns.foldLeft("")(_ + "_" + _.name) + "_opt$" + newId
+      val jpId = fnUid.make
+      val jpName = defns.foldLeft("")(_ + "_" + _.name) + "_opt_jp$" + jpId
 
       val newDefnRef = DefnRef(Right(newName))
       val jpDefnRef = DefnRef(Right(jpName))
@@ -941,15 +1002,16 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
             Jump(jpDefnRef, transformStackFrame(args, defnInfoMap(defn.expectDefn.id))).attachTag(tag)
           else LetCall(names, defn, args, isTailRec, transformNode(body))().attachTag(tag)
 
-      // Tail calls to another function in the component will be replaced with a tail call
+      // Tail calls to another function in the component will be replaced with a call
       // to the merged function
+      // i.e. for mutually tailrec functions f(a, b) and g(c, d),
+      // f's body will be replaced with a call f_g(a, b, *, *), where * is a dummy value
       def transformDefn(defn: Defn): Defn =
-        // TODO: Figure out how to substitute variables with dummy variables.
         val info = defnInfoMap(defn.id)
 
         val start =
-          stackFrame.take(info.stackFrameIdx).drop(1).map { _ => Expr.Literal(IntLit(0)) } // we drop tailrecBranch and replace it with the defn id
-        val end = stackFrame.drop(info.stackFrameIdx + defn.params.size).map { _ => Expr.Literal(IntLit(0)) }
+          stackFrame.take(info.stackFrameIdx).drop(1).map { _ => dummyVal } // we drop tailrecBranch and replace it with the defn id
+        val end = stackFrame.drop(info.stackFrameIdx + defn.params.size).map { _ => dummyVal }
         val args = asLit(info.defn.id) :: start ::: defn.params.map(Expr.Ref(_)) ::: end
 
         // We use a let call instead of a jump to avoid newDefn from being turned into a join point,
@@ -978,10 +1040,10 @@ class TailRecOpt(fnUid: FreshInt, classUid: FreshInt, tag: FreshInt, raise: Diag
 
       val newNode = makeSwitch(trName, valsAndNodes.tail, valsAndNodes.head._2)(trueClass, falseClass)
 
-      val jpDefn = Defn(fnUid.make, jpName, stackFrame, resultNum, newNode, false)
+      val jpDefn = Defn(jpId, jpName, stackFrame, resultNum, newNode, false)
 
       val jmp = Jump(jpDefnRef, stackFrame.map(Expr.Ref(_))).attachTag(tag)
-      val newDefn = Defn(fnUid.make, newName, stackFrame, resultNum, jmp, defnsNoJp.find { _.isTailRec }.isDefined )
+      val newDefn = Defn(newId, newName, stackFrame, resultNum, jmp, defnsNoJp.find { _.isTailRec }.isDefined )
 
       jpDefnRef.defn = Left(jpDefn)
       newDefnRef.defn = Left(newDefn)

compiler/shared/test/diff-ir/IR.mls

@@ -17,7 +17,7 @@ foo()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, mktup2, [x$0,y$0],
 //│ 1, 
 //│ let* (x$1) = mktup(x$0,y$0) in -- #7
@@ -55,7 +55,7 @@ bar()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, foo, [pair$0],
 //│ 1, 
 //│ case pair$0 of -- #16
@@ -105,7 +105,7 @@ foo()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, silly, [pair$0],
 //│ 1, 
 //│ let x$0 = 0 in -- #29
@@ -163,7 +163,7 @@ foo()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, inc_fst, [pair$0],
 //│ 1, 
 //│ let x$0 = 2 in -- #15
@@ -208,7 +208,7 @@ foo()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, inc_fst, [pair$0],
 //│ 1, 
 //│ let x$0 = 0 in -- #15
@@ -256,7 +256,7 @@ bar()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Left, [x]),ClassInfo(1, Right, [y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Left, [x]),ClassInfo(3, Right, [y])}, {
 //│ Def(0, foo, [a$0,b$0],
 //│ 1, 
 //│ case a$0 of -- #36
@@ -298,15 +298,13 @@ bar()
 //│ 2
 
 :interpIR
-class True
-class False
 class Pair(x, y)
 fun foo(a) = a.x + a.y
 fun bar() =
   foo(Pair(1, 0))
 bar()
-//│ |#class| |True|↵|#class| |False|↵|#class| |Pair|(|x|,| |y|)|↵|#fun| |foo|(|a|)| |#=| |a|.x| |+| |a|.y|↵|#fun| |bar|(||)| |#=|→|foo|(|Pair|(|1|,| |0|)|)|←|↵|bar|(||)|
-//│ Parsed: {class True {}; class False {}; class Pair(x, y,) {}; fun foo = (a,) => +((a).x,)((a).y,); fun bar = () => {foo(Pair(1, 0,),)}; bar()}
+//│ |#class| |Pair|(|x|,| |y|)|↵|#fun| |foo|(|a|)| |#=| |a|.x| |+| |a|.y|↵|#fun| |bar|(||)| |#=|→|foo|(|Pair|(|1|,| |0|)|)|←|↵|bar|(||)|
+//│ Parsed: {class Pair(x, y,) {}; fun foo = (a,) => +((a).x,)((a).y,); fun bar = () => {foo(Pair(1, 0,),)}; bar()}
 //│ 
 //│ 
 //│ IR:
@@ -350,7 +348,7 @@ bar()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, C1, [x,y]),ClassInfo(1, C2, [z])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, C1, [x,y]),ClassInfo(3, C2, [z])}, {
 //│ Def(0, foo, [a$0],
 //│ 1, 
 //│ case a$0 of -- #15
@@ -401,7 +399,7 @@ baz()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, foo, [a$0,b$0],
 //│ 1, 
 //│ let x$0 = a$0.x in -- #21
@@ -451,7 +449,7 @@ foo()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, Pair, [x,y])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, Pair, [x,y])}, {
 //│ Def(0, foo, [],
 //│ 1, 
 //│ let x$0 = Pair(0,1) in -- #10
@@ -484,7 +482,7 @@ foo()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, S, [s]),ClassInfo(1, O, [])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, S, [s]),ClassInfo(3, O, [])}, {
 //│ Def(0, foo, [],
 //│ 1, 
 //│ let x$0 = O() in -- #10

compiler/shared/test/diff-ir/IRComplex.mls

@@ -26,7 +26,7 @@ bar()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, A, [x,y,z]),ClassInfo(1, B, [m,n])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, A, [x,y,z]),ClassInfo(3, B, [m,n])}, {
 //│ Def(0, complex_foo, [t$0],
 //│ 1, 
 //│ case t$0 of -- #63
@@ -117,7 +117,7 @@ bar()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, A, [w,x]),ClassInfo(1, B, [y]),ClassInfo(2, C, [z])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, A, [w,x]),ClassInfo(3, B, [y]),ClassInfo(4, C, [z])}, {
 //│ Def(0, complex_foo, [t$0],
 //│ 1, 
 //│ let x$0 = +(1,2) in -- #140
@@ -256,7 +256,7 @@ bar()
 //│ IR:
 //│ 
 //│ Promoted:
-//│ Program({ClassInfo(0, A, [w,x]),ClassInfo(1, B, [y]),ClassInfo(2, C, [z])}, {
+//│ Program({ClassInfo(0, True, []),ClassInfo(1, False, []),ClassInfo(2, A, [w,x]),ClassInfo(3, B, [y]),ClassInfo(4, C, [z])}, {
 //│ Def(0, complex_foo, [t$0],
 //│ 1, 
 //│ let x$0 = +(1,2) in -- #150