Delay conversion of Truffle function body nodes until the function is invoked

CHANGELOG.md

@@ -188,6 +188,7 @@
 - [Fixed execution of defaulted arguments of Atom Constructors][3358]
 - [Converting Enso Date to java.time.LocalDate and back][3374]
 - [Functions with all-defaulted arguments now execute automatically][3414]
+- [Delay construction of Truffle nodes to speed initialization][3429]
 - [Frgaal compiler integration to allow for latest Java constructs][3421]
 
 [3227]: https://github.com/enso-org/enso/pull/3227
@@ -201,6 +202,7 @@
 [3412]: https://github.com/enso-org/enso/pull/3412
 [3414]: https://github.com/enso-org/enso/pull/3414
 [3417]: https://github.com/enso-org/enso/pull/3417
+[3429]: https://github.com/enso-org/enso/pull/3429
 [3421]: https://github.com/enso-org/enso/pull/3421
 
 # Enso 2.0.0-alpha.18 (2021-10-12)

engine/runtime/src/main/java/org/enso/interpreter/node/MethodRootNode.java

@@ -1,8 +1,10 @@
 package org.enso.interpreter.node;
 
 import com.oracle.truffle.api.dsl.ReportPolymorphism;
+import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.NodeInfo;
 import com.oracle.truffle.api.source.SourceSection;
+import java.util.function.Supplier;
 import org.enso.interpreter.Language;
 import org.enso.interpreter.runtime.callable.atom.AtomConstructor;
 import org.enso.interpreter.runtime.scope.LocalScope;
@@ -44,12 +46,30 @@ private static String shortName(String atomName, String methodName) {
    * @param language the language identifier
    * @param localScope a description of the local scope
    * @param moduleScope a description of the module scope
-   * @param body the program body to be executed
-   * @param section a mapping from {@code body} to the program source
+   * @param body the program provider to be executed
+   * @param section a mapping from {@code provider} to the program source
    * @param atomConstructor the constructor this method is defined for
    * @param methodName the name of this method
    * @return a node representing the specified closure
    */
+  public static MethodRootNode build(
+      Language language,
+      LocalScope localScope,
+      ModuleScope moduleScope,
+      Supplier<ExpressionNode> body,
+      SourceSection section,
+      AtomConstructor atomConstructor,
+      String methodName) {
+    return build(
+        language,
+        localScope,
+        moduleScope,
+        new LazyBodyNode(body),
+        section,
+        atomConstructor,
+        methodName);
+  }
+
   public static MethodRootNode build(
       Language language,
       LocalScope localScope,
@@ -87,4 +107,19 @@ public AtomConstructor getAtomConstructor() {
   public String getMethodName() {
     return methodName;
   }
+
+  private static class LazyBodyNode extends ExpressionNode {
+    private final Supplier<ExpressionNode> provider;
+
+    LazyBodyNode(Supplier<ExpressionNode> body) {
+      this.provider = body;
+    }
+
+    @Override
+    public Object executeGeneric(VirtualFrame frame) {
+      ExpressionNode newNode = replace(provider.get());
+      notifyInserted(newNode);
+      return newNode.executeGeneric(frame);
+    }
+  }
 }

engine/runtime/src/main/java/org/enso/interpreter/node/callable/function/BlockNode.java

@@ -30,6 +30,10 @@ public static BlockNode build(ExpressionNode[] expressions, ExpressionNode retur
     return new BlockNode(expressions, returnExpr);
   }
 
+  public static BlockNode buildSilent(ExpressionNode[] expressions, ExpressionNode returnExpr) {
+    return new BlockNode(expressions, returnExpr);
+  }
+
   /**
    * Executes the body of the function.
    *

engine/runtime/src/main/java/org/enso/interpreter/runtime/callable/atom/AtomConstructor.java

@@ -111,7 +111,7 @@ private Function buildConstructorFunction(
       ArgumentDefinition[] args) {
 
     ExpressionNode instantiateNode = InstantiateNode.build(this, varReads);
-    BlockNode instantiateBlock = BlockNode.build(assignments, instantiateNode);
+    BlockNode instantiateBlock = BlockNode.buildSilent(assignments, instantiateNode);
     RootNode rootNode =
         ClosureRootNode.build(
             null,

engine/runtime/src/main/scala/org/enso/compiler/codegen/IrToTruffle.scala

@@ -290,18 +290,19 @@ class IrToTruffle(
 
         val function = methodDef.body match {
           case fn: IR.Function =>
-            val (body, arguments) =
-              expressionProcessor.buildFunctionBody(fn.arguments, fn.body)
+            val bodyBuilder =
+              new expressionProcessor.BuildFunctionBody(fn.arguments, fn.body)
             val rootNode = MethodRootNode.build(
               language,
               expressionProcessor.scope,
               moduleScope,
-              body,
+              () => bodyBuilder.bodyNode(),
               makeSection(methodDef.location),
               cons,
               methodDef.methodName.name
             )
             val callTarget = Truffle.getRuntime.createCallTarget(rootNode)
+            val arguments  = bodyBuilder.args()
             new RuntimeFunction(
               callTarget,
               null,
@@ -348,18 +349,19 @@ class IrToTruffle(
 
         val function = methodDef.body match {
           case fn: IR.Function =>
-            val (body, arguments) =
-              expressionProcessor.buildFunctionBody(fn.arguments, fn.body)
+            val bodyBuilder =
+              new expressionProcessor.BuildFunctionBody(fn.arguments, fn.body)
             val rootNode = MethodRootNode.build(
               language,
               expressionProcessor.scope,
               moduleScope,
-              body,
+              () => bodyBuilder.bodyNode(),
               makeSection(methodDef.location),
               toType,
               methodDef.methodName.name
             )
             val callTarget = Truffle.getRuntime.createCallTarget(rootNode)
+            val arguments  = bodyBuilder.args()
             new RuntimeFunction(
               callTarget,
               null,
@@ -1187,59 +1189,72 @@ class IrToTruffle(
       * @return a node for the final shape of function body and pre-processed
       *         argument definitions.
       */
-    def buildFunctionBody(
-      arguments: List[IR.DefinitionArgument],
-      body: IR.Expression
-    ): (BlockNode, Array[ArgumentDefinition]) = {
-      val argFactory = new DefinitionArgumentProcessor(scopeName, scope)
-
-      val argDefinitions = new Array[ArgumentDefinition](arguments.size)
-      val argExpressions = new ArrayBuffer[RuntimeExpression]
-      val seenArgNames   = mutable.Set[String]()
-
-      // Note [Rewriting Arguments]
-      val argSlots =
-        arguments.zipWithIndex.map { case (unprocessedArg, idx) =>
-          val arg = argFactory.run(unprocessedArg, idx)
-          argDefinitions(idx) = arg
-
-          val occInfo = unprocessedArg
-            .unsafeGetMetadata(
-              AliasAnalysis,
-              "No occurrence on an argument definition."
+    class BuildFunctionBody(
+      val arguments: List[IR.DefinitionArgument],
+      val body: IR.Expression
+    ) {
+      private val argFactory = new DefinitionArgumentProcessor(scopeName, scope)
+      private lazy val slots = computeSlots()
+      private lazy val bodyN = computeBodyNode()
+
+      def args(): Array[ArgumentDefinition] = slots._2
+      def bodyNode(): BlockNode             = bodyN
+
+      private def computeBodyNode(): BlockNode = {
+        val (argSlots, _, argExpressions) = slots
+
+        val bodyExpr = body match {
+          case IR.Foreign.Definition(lang, code, _, _, _) =>
+            buildForeignBody(
+              lang,
+              code,
+              arguments.map(_.name.name),
+              argSlots
             )
-            .unsafeAs[AliasAnalysis.Info.Occurrence]
+          case _ => ExpressionProcessor.this.run(body)
+        }
+        BlockNode.build(argExpressions.toArray, bodyExpr)
+      }
 
-          val slot = scope.createVarSlot(occInfo.id)
-          val readArg =
-            ReadArgumentNode.build(idx, arg.getDefaultValue.orElse(null))
-          val assignArg = AssignmentNode.build(readArg, slot)
+      private def computeSlots(): (
+        List[FrameSlot],
+        Array[ArgumentDefinition],
+        ArrayBuffer[RuntimeExpression]
+      ) = {
+        val seenArgNames   = mutable.Set[String]()
+        val argDefinitions = new Array[ArgumentDefinition](arguments.size)
+        val argExpressions = new ArrayBuffer[RuntimeExpression]
+        // Note [Rewriting Arguments]
+        val argSlots = arguments.zipWithIndex.map {
+          case (unprocessedArg, idx) =>
+            val arg = argFactory.run(unprocessedArg, idx)
+            argDefinitions(idx) = arg
+
+            val occInfo = unprocessedArg
+              .unsafeGetMetadata(
+                AliasAnalysis,
+                "No occurrence on an argument definition."
+              )
+              .unsafeAs[AliasAnalysis.Info.Occurrence]
 
-          argExpressions.append(assignArg)
+            val slot = scope.createVarSlot(occInfo.id)
+            val readArg =
+              ReadArgumentNode.build(idx, arg.getDefaultValue.orElse(null))
+            val assignArg = AssignmentNode.build(readArg, slot)
 
-          val argName = arg.getName
+            argExpressions.append(assignArg)
 
-          if (seenArgNames contains argName) {
-            throw new IllegalStateException(
-              s"A duplicate argument name, $argName, was found during codegen."
-            )
-          } else seenArgNames.add(argName)
-          slot
-        }
+            val argName = arg.getName
 
-      val bodyExpr = body match {
-        case IR.Foreign.Definition(lang, code, _, _, _) =>
-          buildForeignBody(
-            lang,
-            code,
-            arguments.map(_.name.name),
-            argSlots
-          )
-        case _ => this.run(body)
+            if (seenArgNames contains argName) {
+              throw new IllegalStateException(
+                s"A duplicate argument name, $argName, was found during codegen."
+              )
+            } else seenArgNames.add(argName)
+            slot
+        }
+        (argSlots, argDefinitions, argExpressions)
       }
-
-      val fnBodyNode = BlockNode.build(argExpressions.toArray, bodyExpr)
-      (fnBodyNode, argDefinitions)
     }
 
     private def buildForeignBody(
@@ -1269,18 +1284,18 @@ class IrToTruffle(
       body: IR.Expression,
       location: Option[IdentifiedLocation]
     ): CreateFunctionNode = {
-      val (fnBodyNode, argDefinitions) = buildFunctionBody(arguments, body)
+      val bodyBuilder = new BuildFunctionBody(arguments, body)
       val fnRootNode = ClosureRootNode.build(
         language,
         scope,
         moduleScope,
-        fnBodyNode,
+        bodyBuilder.bodyNode(),
         makeSection(location),
         scopeName
       )
       val callTarget = Truffle.getRuntime.createCallTarget(fnRootNode)
 
-      val expr = CreateFunctionNode.build(callTarget, argDefinitions)
+      val expr = CreateFunctionNode.build(callTarget, bodyBuilder.args())
 
       setLocation(expr, location)
     }