Initialize builtin methods lazily in non-AOT mode (#8654)

Registration of builtin methods is done upfront, similarly to AOT mode, but the initialization is delayed until the first use. One still has to read the full list of builtin methods and register them in the builtins scope or we would be getting the `NoSuchMethod` errors otherwise. But the most expensive operation, initialization of classes and getting a method via a reflection, is all done on the first usage. The result is an improved startup and negligible performance impact in regular usage. Closes #8423. # Important Notes Visible speedup. 1. (old) Program returning plain "Hello World" (no stdlib loading) ![Screenshot from 2024-01-02 17-34-35](https://github.com/enso-org/enso/assets/292128/b8348cbb-baf6-4292-8c0e-c8a3492c6583) 2. (new) Program returning plain "Hello World" (no stdlib loading) ![Screenshot from 2024-01-02 17-34-54](https://github.com/enso-org/enso/assets/292128/63dd6c4f-49e2-4a5e-ab95-2460c155b91c) About 10% improvement for the overall execution. 3. (old) Program printing "Hello World" (some stdlib loading) ![Screenshot from 2024-01-02 17-40-21](https://github.com/enso-org/enso/assets/292128/cf3f464c-8f00-4adf-9499-a6ac86430da0) 4. (new) Program printing "Hello World" (some stdlib loading) ![Screenshot from 2024-01-02 17-40-41](https://github.com/enso-org/enso/assets/292128/3ef99893-ccaf-4ea7-90ba-aa090f904a50) About 2% improvement for the overall execution. Minor price to pay for lazy initialization: ![Screenshot from 2024-01-02 17-56-12](https://github.com/enso-org/enso/assets/292128/faa682d8-5370-43d2-9fc9-3454af5749a4)
2024-12-23 10:05:06 +03:00 · 2024-01-04 09:55:21 +01:00 · 2024-01-04 09:55:21 +01:00 · 20531d51df
commit 20531d51df
parent a215521c10
3 changed files with 204 additions and 80 deletions
--- a/engine/runtime/src/main/java/org/enso/interpreter/runtime/builtin/Builtins.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/runtime/builtin/Builtins.java
@ -1,6 +1,7 @@
 package org.enso.interpreter.runtime.builtin;

 import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.TruffleOptions;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStream;
@ -16,6 +17,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.enso.compiler.Passes;
 import org.enso.compiler.context.CompilerContext;
@ -50,17 +52,22 @@ import org.enso.interpreter.runtime.Module;
 import org.enso.interpreter.runtime.callable.function.Function;
 import org.enso.interpreter.runtime.data.Type;
 import org.enso.interpreter.runtime.scope.ModuleScope;
+import org.enso.interpreter.runtime.util.CachingSupplier;
 import org.enso.pkg.QualifiedName;

 /** Container class for static predefined atoms, methods, and their containing scope. */
 public final class Builtins {

-  private static final List<Constructor<? extends Builtin>> loadedBuiltinConstructors;
-  private static final Map<String, LoadedBuiltinMethod> loadedBuiltinMethods;
+  private static List<Constructor<? extends Builtin>> loadedBuiltinConstructors;
+  private static Map<String, LoadedBuiltinMetaMethod> loadedBuiltinMethodsMeta;
+  private static Map<String, LoadedBuiltinMethod> loadedBuiltinMethods;

  static {
    loadedBuiltinConstructors = readBuiltinTypes();
-    loadedBuiltinMethods = readBuiltinMethodsMethods();
+    loadedBuiltinMethodsMeta = readBuiltinMethodsMeta();
+    if (TruffleOptions.AOT) {
+      loadedBuiltinMethods = loadBuiltinMethodClassesEarly(loadedBuiltinMethodsMeta);
+    }
  }

  public static final String PACKAGE_NAME = "Builtins";
@ -75,7 +82,7 @@ public final class Builtins {
  }

  private final Map<Class<? extends Builtin>, Builtin> builtins;
-  private final Map<String, Map<String, LoadedBuiltinMethod>> builtinMethodNodes;
+  private final Map<String, Map<String, Supplier<LoadedBuiltinMethod>>> builtinMethodNodes;
  private final Map<String, Builtin> builtinsByName;

  private final Error error;
@ -129,8 +136,12 @@ public final class Builtins {
            .collect(
                Collectors.toMap(
                    v -> v.getType().getName(), java.util.function.Function.identity()));
-    builtinMethodNodes = readBuiltinMethodsMetadata(loadedBuiltinMethods, scope);
-    registerBuiltinMethods(scope, language);
+    if (TruffleOptions.AOT) {
+      builtinMethodNodes = readBuiltinMethodsMetadata(loadedBuiltinMethods, scope);
+      registerBuiltinMethods(scope, language);
+    } else {
+      builtinMethodNodes = registerBuiltinMethodsLazily(scope, language);
+    }

    ordering = getBuiltinType(Ordering.class);
    comparable = getBuiltinType(Comparable.class);
@ -165,10 +176,20 @@ public final class Builtins {
    special = new Special(language);
  }

+  private static Map<String, LoadedBuiltinMethod> loadBuiltinMethodClassesEarly(
+      Map<String, LoadedBuiltinMetaMethod> map) {
+    Map<String, LoadedBuiltinMethod> methods = new HashMap<>();
+    map.forEach(
+        (key, value) -> {
+          methods.put(key, value.toMethod());
+        });
+    return methods;
+  }
+
  /**
   * Registers builtin methods with their corresponding Atom Constructor's owners. That way
   * "special" builtin types have builtin methods in the scope without requiring everyone to always
-   * import full stdlib
+   * import full stdlib.
   *
   * @param scope Builtins scope
   * @param language The language the resulting function nodes should be associated with
@ -176,8 +197,73 @@ public final class Builtins {
  private void registerBuiltinMethods(ModuleScope scope, EnsoLanguage language) {
    for (Builtin builtin : builtins.values()) {
      var type = builtin.getType();
-      String tpeName = type.getName();
-      Map<String, LoadedBuiltinMethod> methods = builtinMethodNodes.get(tpeName);
+      Map<String, Supplier<LoadedBuiltinMethod>> methods = builtinMethodNodes.get(type.getName());
+      if (methods != null) {
+        // Register a builtin method iff it is marked as auto-register.
+        // Methods can only register under a type or, if we deal with a static method, it's
+        // eigen-type.
+        // Such builtins are available on certain types without importing the whole stdlib, e.g. Any
+        // or Number.
+        methods.forEach(
+            (key, value) -> {
+              LoadedBuiltinMethod meth = value.get();
+              Type tpe =
+                  meth.isAutoRegister ? (!meth.isStatic() ? type : type.getEigentype()) : null;
+              if (tpe != null) {
+                Optional<BuiltinFunction> fun = meth.toFunction(language, false);
+                fun.ifPresent(f -> scope.registerMethod(tpe, key, f.getFunction()));
+              }
+            });
+      }
+    }
+  }
+
+  /**
+   * Register builtin methods and initialize them lazily in the provided scope. This method differs
+   * from `registerBuiltinMethods` where all methods are initialized by the time they are
+   * registered..
+   *
+   * @param scope Builtins scope
+   * @param language The language the resulting function nodes should be associated with
+   * @return map from types to builtin methods
+   */
+  private Map<String, Map<String, Supplier<LoadedBuiltinMethod>>> registerBuiltinMethodsLazily(
+      ModuleScope scope, EnsoLanguage language) {
+    Map<String, Map<String, Supplier<LoadedBuiltinMethod>>> builtinMethodNodes = new HashMap<>();
+    Map<String, Map<String, LoadedBuiltinMetaMethod>> builtinMetaMethods = new HashMap<>();
+    loadedBuiltinMethodsMeta.forEach(
+        (fullName, meta) -> {
+          String[] builtinName = fullName.split("\\.");
+          if (builtinName.length != 2) {
+            throw new CompilerError("Invalid builtin metadata for " + fullName);
+          }
+          String builtinMethodOwner = builtinName[0];
+          String builtinMethodName = builtinName[1];
+          Optional.ofNullable(scope.getTypes().get(builtinMethodOwner))
+              .ifPresentOrElse(
+                  constr -> {
+                    Map<String, Supplier<LoadedBuiltinMethod>> atomNodes =
+                        getOrUpdate(builtinMethodNodes, constr.getName());
+                    atomNodes.put(builtinMethodName, new CachingSupplier<>(() -> meta.toMethod()));
+
+                    Map<String, LoadedBuiltinMetaMethod> atomNodesMeta =
+                        getOrUpdate(builtinMetaMethods, constr.getName());
+                    atomNodesMeta.put(builtinMethodName, meta);
+                  },
+                  () -> {
+                    Map<String, Supplier<LoadedBuiltinMethod>> atomNodes =
+                        getOrUpdate(builtinMethodNodes, builtinMethodOwner);
+                    atomNodes.put(builtinMethodName, new CachingSupplier<>(() -> meta.toMethod()));
+
+                    Map<String, LoadedBuiltinMetaMethod> atomNodesMeta =
+                        getOrUpdate(builtinMetaMethods, builtinMethodOwner);
+                    atomNodesMeta.put(builtinMethodName, meta);
+                  });
+        });
+
+    for (Builtin builtin : builtins.values()) {
+      var type = builtin.getType();
+      Map<String, LoadedBuiltinMetaMethod> methods = builtinMetaMethods.get(type.getName());
      if (methods != null) {
        // Register a builtin method iff it is marked as auto-register.
        // Methods can only register under a type or, if we deal with a static method, it's
@ -187,14 +273,25 @@ public final class Builtins {
        methods.forEach(
            (key, value) -> {
              Type tpe =
-                  value.isAutoRegister ? (!value.isStatic() ? type : type.getEigentype()) : null;
+                  value.isAutoRegister() ? (!value.isStatic() ? type : type.getEigentype()) : null;
              if (tpe != null) {
-                Optional<BuiltinFunction> fun = value.toFunction(language, false);
-                fun.ifPresent(f -> scope.registerMethod(tpe, key, f.getFunction()));
+                Supplier<Function> supplier =
+                    () -> value.toMethod().toFunction(language, false).get().getFunction();
+                scope.registerMethod(tpe, key, supplier);
              }
            });
      }
    }
+    return builtinMethodNodes;
+  }
+
+  private <T> Map<String, T> getOrUpdate(Map<String, Map<String, T>> map, String key) {
+    Map<String, T> entry = map.get(key);
+    if (entry == null) {
+      entry = new HashMap<>();
+      map.put(key, entry);
+    }
+    return entry;
  }

  /**
@ -279,17 +376,18 @@ public final class Builtins {
  private Map<Class<? extends Builtin>, Builtin> initializeBuiltinTypes(
      List<Constructor<? extends Builtin>> constrs, EnsoLanguage language, ModuleScope scope) {
    Map<Class<? extends Builtin>, Builtin> builtins = new HashMap<>();
-    constrs.forEach(
-        constr -> {
-          try {
-            Builtin builtin = constr.newInstance();
-            builtins.put(builtin.getClass(), builtin);
-          } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
-            e.printStackTrace();
-            throw new CompilerError("Invalid builtin type entry: " + constr);
-          }
-        });
-    builtins.values().forEach(b -> b.initialize(language, scope, builtins));
+
+    for (var constr : constrs) {
+      try {
+        Builtin builtin = constr.newInstance();
+        builtins.put(builtin.getClass(), builtin);
+      } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
+        throw new CompilerError("Invalid builtin type entry: " + constr, e);
+      }
+    }
+    for (var b : builtins.values()) {
+      b.initialize(language, scope, builtins);
+    }
    return builtins;
  }

@ -300,10 +398,10 @@ public final class Builtins {
   * @param scope Builtins scope
   * @return A map of builtin method nodes per builtin type name
   */
-  private Map<String, Map<String, LoadedBuiltinMethod>> readBuiltinMethodsMetadata(
+  private Map<String, Map<String, Supplier<LoadedBuiltinMethod>>> readBuiltinMethodsMetadata(
      Map<String, LoadedBuiltinMethod> classes, ModuleScope scope) {

-    Map<String, Map<String, LoadedBuiltinMethod>> methodNodes = new HashMap<>();
+    Map<String, Map<String, Supplier<LoadedBuiltinMethod>>> methodNodes = new HashMap<>();
    classes.forEach(
        (fullBuiltinName, builtin) -> {
          String[] builtinName = fullBuiltinName.split("\\.");
@ -315,24 +413,14 @@ public final class Builtins {
          Optional.ofNullable(scope.getTypes().get(builtinMethodOwner))
              .ifPresentOrElse(
                  constr -> {
-                    Map<String, LoadedBuiltinMethod> atomNodes =
-                        methodNodes.get(builtinMethodOwner);
-                    if (atomNodes == null) {
-                      atomNodes = new HashMap<>();
-                      // TODO: move away from String Map once Builtins are gone
-                      methodNodes.put(constr.getName(), atomNodes);
-                    }
-                    atomNodes.put(builtinMethodName, builtin);
+                    Map<String, Supplier<LoadedBuiltinMethod>> atomNodes =
+                        getOrUpdate(methodNodes, constr.getName());
+                    atomNodes.put(builtinMethodName, new CachingSupplier<>(builtin));
                  },
                  () -> {
-                    Map<String, LoadedBuiltinMethod> atomNodes =
-                        methodNodes.get(builtinMethodOwner);
-                    if (atomNodes == null) {
-                      atomNodes = new HashMap<>();
-                      // TODO: move away from String Map once Builtins are gone
-                      methodNodes.put(builtinMethodOwner, atomNodes);
-                    }
-                    atomNodes.put(builtinMethodName, builtin);
+                    Map<String, Supplier<LoadedBuiltinMethod>> atomNodes =
+                        getOrUpdate(methodNodes, builtinMethodOwner);
+                    atomNodes.put(builtinMethodName, new CachingSupplier<>(builtin));
                  });
        });
    return methodNodes;
@ -347,7 +435,7 @@ public final class Builtins {
   *
   * @return A map of builtin method nodes per builtin type name
   */
-  private static Map<String, LoadedBuiltinMethod> readBuiltinMethodsMethods() {
+  private static Map<String, LoadedBuiltinMetaMethod> readBuiltinMethodsMeta() {
    ClassLoader classLoader = Builtins.class.getClassLoader();
    List<String> lines;
    try (InputStream resource = classLoader.getResourceAsStream(MethodDefinition.META_PATH)) {
@ -374,18 +462,9 @@ public final class Builtins {
              boolean isStatic = java.lang.Boolean.valueOf(builtinMeta[2]);
              boolean isAutoRegister = java.lang.Boolean.valueOf(builtinMeta[3]);

-              try {
-                @SuppressWarnings("unchecked")
-                Class<BuiltinRootNode> clazz =
-                    (Class<BuiltinRootNode>) Class.forName(builtinMeta[1]);
-                Method meth = clazz.getMethod("makeFunction", EnsoLanguage.class, boolean.class);
-                LoadedBuiltinMethod meta = new LoadedBuiltinMethod(meth, isStatic, isAutoRegister);
-                return new AbstractMap.SimpleEntry<String, LoadedBuiltinMethod>(
-                    builtinMeta[0], meta);
-              } catch (ClassNotFoundException | NoSuchMethodException e) {
-                e.printStackTrace();
-                throw new CompilerError("Invalid builtin method " + line);
-              }
+              return new AbstractMap.SimpleEntry<>(
+                  builtinMeta[0],
+                  new LoadedBuiltinMetaMethod(builtinMeta[1], isStatic, isAutoRegister));
            })
        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
  }
@ -402,10 +481,14 @@ public final class Builtins {
  public Optional<BuiltinFunction> getBuiltinFunction(
      String type, String methodName, EnsoLanguage language, boolean isStaticInstance) {
    // TODO: move away from String mapping once Builtins is gone
-    Map<String, LoadedBuiltinMethod> atomNodes = builtinMethodNodes.get(type);
-    if (atomNodes == null) return Optional.empty();
-    LoadedBuiltinMethod builtin = atomNodes.get(methodName);
-    if (builtin == null) return Optional.empty();
+    Map<String, Supplier<LoadedBuiltinMethod>> atomNodes = builtinMethodNodes.get(type);
+    if (atomNodes == null) {
+      return Optional.empty();
+    }
+    LoadedBuiltinMethod builtin = atomNodes.get(methodName).get();
+    if (builtin == null) {
+      return Optional.empty();
+    }
    return builtin.toFunction(language, isStaticInstance);
  }

@ -666,6 +749,43 @@ public final class Builtins {
    return module;
  }

+  private static class LoadedBuiltinMetaMethod {
+
+    private LoadedBuiltinMethod method;
+    private final String className;
+    private final boolean staticMethod;
+    private final boolean autoRegister;
+
+    private LoadedBuiltinMetaMethod(String className, boolean staticMethod, boolean autoRegister) {
+      this.className = className;
+      this.staticMethod = staticMethod;
+      this.autoRegister = autoRegister;
+      this.method = null;
+    }
+
+    boolean isStatic() {
+      return staticMethod;
+    }
+
+    boolean isAutoRegister() {
+      return autoRegister;
+    }
+
+    LoadedBuiltinMethod toMethod() {
+      if (method == null) {
+        try {
+          @SuppressWarnings("unchecked")
+          Class<BuiltinRootNode> clazz = (Class<BuiltinRootNode>) Class.forName(className);
+          Method meth = clazz.getMethod("makeFunction", EnsoLanguage.class, boolean.class);
+          method = new LoadedBuiltinMethod(meth, staticMethod, autoRegister);
+        } catch (ClassNotFoundException | NoSuchMethodException e) {
+          throw new CompilerError("Invalid builtin method " + className, e);
+        }
+      }
+      return method;
+    }
+  }
+
  private record LoadedBuiltinMethod(Method meth, boolean isStatic, boolean isAutoRegister) {
    Optional<BuiltinFunction> toFunction(EnsoLanguage language, boolean isStaticInstance) {
      try {
--- a/engine/runtime/src/main/java/org/enso/interpreter/runtime/scope/ModuleScope.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/runtime/scope/ModuleScope.java
@ -18,6 +18,7 @@ import org.enso.interpreter.runtime.data.EnsoObject;
 import org.enso.interpreter.runtime.data.Type;
 import org.enso.interpreter.runtime.error.RedefinedConversionException;
 import org.enso.interpreter.runtime.error.RedefinedMethodException;
+import org.enso.interpreter.runtime.util.CachingSupplier;

 /** A representation of Enso's per-file top-level scope. */
 public final class ModuleScope implements EnsoObject {
@ -155,28 +156,6 @@ public final class ModuleScope implements EnsoObject {
    }
  }

-  private static final class CachingSupplier<T> implements Supplier<T> {
-    private final Supplier<T> supply;
-    private T memo;
-
-    CachingSupplier(Supplier<T> supply) {
-      this.supply = supply;
-    }
-
-    CachingSupplier(T memo) {
-      this.supply = null;
-      this.memo = memo;
-    }
-
-    @Override
-    public T get() {
-      if (memo == null) {
-        memo = supply.get();
-      }
-      return memo;
-    }
-  }
-
  /**
   * Returns a list of the conversion methods defined in this module for a given constructor.
   *
--- a/engine/runtime/src/main/java/org/enso/interpreter/runtime/util/CachingSupplier.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/runtime/util/CachingSupplier.java
@ -0,0 +1,25 @@
+package org.enso.interpreter.runtime.util;
+
+import java.util.function.Supplier;
+
+public final class CachingSupplier<T> implements Supplier<T> {
+  private final Supplier<T> supply;
+  private T memo;
+
+  public CachingSupplier(Supplier<T> supply) {
+    this.supply = supply;
+  }
+
+  public CachingSupplier(T memo) {
+    this.supply = null;
+    this.memo = memo;
+  }
+
+  @Override
+  public T get() {
+    if (memo == null) {
+      memo = supply.get();
+    }
+    return memo;
+  }
+}