Optimize new inference in method args (#1845)

pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java

@@ -878,7 +878,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
     if (resolution instanceof LexicalMethod method) {
       var levelsUp = method.levelsUp();
       var identifier = org.pkl.core.runtime.Identifier.method(name, method.isLocal());
-      var args = visitArgumentList(argList);
+      var argInfo = visitArgumentList(argList);
       var needsConst =
           switch (constLevel) {
             case NONE -> false;
@@ -891,35 +891,57 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
         var getModuleNode = new GetTypeAliasModuleNode(sourceSection);
         if (method.isObjectMethod()) {
           return new InvokeQualifiedObjectMethodNode(
-              sourceSection, identifier, args, needsConst, getModuleNode);
+              sourceSection,
+              identifier,
+              argInfo.getFirst(),
+              needsConst,
+              getModuleNode,
+              argInfo.getSecond());
         }
         if (method.isOnClosedClass() || method.isLocal() || method.isExternal()) {
           return new InvokeQualifiedClassMethodNode(
-              sourceSection, identifier, args, needsConst, getModuleNode);
+              sourceSection,
+              identifier,
+              argInfo.getFirst(),
+              needsConst,
+              getModuleNode,
+              argInfo.getSecond());
         }
         return InvokeMethodVirtualNodeGen.create(
             sourceSection,
             identifier,
-            args,
+            argInfo.getFirst(),
             MemberLookupMode.IMPLICIT_LEXICAL,
             needsConst,
+            argInfo.getSecond(),
             getModuleNode,
             GetClassNodeGen.create(null));
       }
       if (method.isObjectMethod()) {
         return new InvokeLexicalObjectMethodNode(
-            sourceSection, identifier, levelsUp, args, needsConst);
+            sourceSection,
+            identifier,
+            levelsUp,
+            argInfo.getFirst(),
+            needsConst,
+            argInfo.getSecond());
       }
       if (method.isOnClosedClass() || method.isLocal() || method.isExternal()) {
         return new InvokeLexicalClassMethodNode(
-            sourceSection, identifier, levelsUp, args, needsConst);
+            sourceSection,
+            identifier,
+            levelsUp,
+            argInfo.getFirst(),
+            needsConst,
+            argInfo.getSecond());
       }
       return InvokeMethodVirtualNodeGen.create(
           sourceSection,
           identifier,
-          args,
+          argInfo.getFirst(),
           MemberLookupMode.IMPLICIT_LEXICAL,
           needsConst,
+          argInfo.getSecond(),
           levelsUp == 0 ? new GetReceiverNode() : new GetEnclosingReceiverNode(levelsUp),
           GetClassNodeGen.create(null));
     } else if (resolution instanceof ImplicitBaseMethod) {
@@ -941,21 +963,25 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
         var baseModule = BaseModule.getModule();
         var method = baseModule.getVmClass().getDeclaredMethod(identifier);
         assert method != null;
+        var argInfo = visitArgumentList(argList);
         return new InvokeMethodDirectNode(
             createSourceSection(expr),
             method,
             new ConstantValueNode(baseModule),
-            visitArgumentList(argList));
+            argInfo.getFirst(),
+            argInfo.getSecond());
       }
     } else if (resolution instanceof ImplicitThisMethod) {
       var isCustomThis = scope.isCustomThisScope();
       var needsConst = constLevel == ConstLevel.ALL && constDepth == -1 && !isCustomThis;
+      var argInfo = visitArgumentList(argList);
       return InvokeMethodVirtualNodeGen.create(
           sourceSection,
           org.pkl.core.runtime.Identifier.get(name),
-          visitArgumentList(argList),
+          argInfo.getFirst(),
           MemberLookupMode.IMPLICIT_THIS,
           needsConst,
+          argInfo.getSecond(),
           VmUtils.createThisNode(VmUtils.unavailableSourceSection(), isCustomThis),
           GetClassNodeGen.create(null));
     } else {
@@ -1053,7 +1079,9 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
     var parent = expr.parent();
     var scope = symbolTable.getCurrentScope();
 
-    while (parent instanceof IfExpr
+    // keep in sync with isImplicitNewExpr
+    while (parent instanceof IfExpr ifExpr
+            && (ifExpr.getThen() == child || ifExpr.getEls() == child)
         || parent instanceof TraceExpr
         || parent instanceof LetExpr letExpr && letExpr.getExpr() == child) {
 
@@ -1135,8 +1163,9 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
             .build();
       }
 
+      var argInfo = visitArgumentList(argCtx);
       return InvokeSuperMethodNodeGen.create(
-          sourceSection, memberName, visitArgumentList(argCtx), needsConst);
+          sourceSection, memberName, argInfo.getFirst(), needsConst, argInfo.getSecond());
     }
 
     // superproperty call
@@ -1150,11 +1179,10 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
 
   @Override
   public ExpressionNode visitQualifiedAccessExpr(QualifiedAccessExpr expr) {
-    if (expr.getArgumentList() != null) {
-      return doVisitMethodAccessExpr(expr);
-    }
-
-    return doVisitPropertyInvocationExpr(expr);
+    var argList = expr.getArgumentList();
+    return argList != null
+        ? doVisitMethodAccessExpr(expr, argList)
+        : doVisitPropertyInvocationExpr(expr);
   }
 
   @Override
@@ -2306,13 +2334,32 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
   }
 
   @Override
-  public ExpressionNode[] visitArgumentList(ArgumentList argumentList) {
+  public Pair<ExpressionNode[], Boolean> visitArgumentList(ArgumentList argumentList) {
     var args = argumentList.getArguments();
     var res = new ExpressionNode[args.size()];
+    var argsRequireInference = false;
     for (int i = 0; i < res.length; i++) {
-      res[i] = visitExpr(args.get(i));
+      var expr = args.get(i);
+      res[i] = visitExpr(expr);
+      argsRequireInference = argsRequireInference || isImplicitNewExpr(expr);
     }
-    return res;
+    return Pair.of(res, argsRequireInference);
+  }
+
+  private static boolean isImplicitNewExpr(Expr expr) {
+    // keep in sync with doVisitNewExprWithInferredParent
+    if (expr instanceof NewExpr newExpr && newExpr.getType() == null) {
+      return true;
+    } else if (expr instanceof IfExpr ifExpr) {
+      return isImplicitNewExpr(ifExpr.getThen()) || isImplicitNewExpr(ifExpr.getEls());
+    } else if (expr instanceof TraceExpr traceExpr) {
+      return isImplicitNewExpr(traceExpr.getExpr());
+    } else if (expr instanceof ParenthesizedExpr parenthesizedExpr) {
+      return isImplicitNewExpr(parenthesizedExpr.getExpr());
+    } else if (expr instanceof LetExpr letExpr) {
+      return isImplicitNewExpr(letExpr.getExpr());
+    }
+    return false;
   }
 
   @Override
@@ -2888,12 +2935,12 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
     return ReadPropertyNodeGen.create(sourceSection, propertyName, needsConst, receiver);
   }
 
-  private ExpressionNode doVisitMethodAccessExpr(QualifiedAccessExpr expr) {
+  private ExpressionNode doVisitMethodAccessExpr(QualifiedAccessExpr expr, ArgumentList argList) {
     var sourceSection = createSourceSection(expr);
     var functionName = toIdentifier(expr.getIdentifier().getValue());
-    var argCtx = expr.getArgumentList();
     var receiver = visitExpr(expr.getExpr());
     var needsConst = needsConst(receiver);
+    var argInfo = visitArgumentList(argList);
 
     if (expr.isNullable()) {
       //noinspection ConstantConditions
@@ -2902,9 +2949,10 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
           InvokeMethodVirtualNodeGen.create(
               sourceSection,
               functionName,
-              visitArgumentList(argCtx),
+              argInfo.getFirst(),
               MemberLookupMode.EXPLICIT_RECEIVER,
               needsConst,
+              argInfo.getSecond(),
               PropagateNullReceiverNodeGen.create(unavailableSourceSection(), receiver),
               GetClassNodeGen.create(null)));
     }
@@ -2913,9 +2961,10 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
     return InvokeMethodVirtualNodeGen.create(
         sourceSection,
         functionName,
-        visitArgumentList(argCtx),
+        argInfo.getFirst(),
         MemberLookupMode.EXPLICIT_RECEIVER,
         needsConst,
+        argInfo.getSecond(),
         receiver,
         GetClassNodeGen.create(null));
   }

pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeLexicalMethodNode.java

@@ -31,8 +31,9 @@ public abstract sealed class AbstractInvokeLexicalMethodNode
       Identifier methodName,
       int levelsUp,
       ExpressionNode[] argumentNodes,
-      boolean needsConst) {
-    super(sourceSection, methodName, argumentNodes, needsConst);
+      boolean needsConst,
+      boolean argsRequireInference) {
+    super(sourceSection, methodName, argumentNodes, needsConst, argsRequireInference);
     this.levelsUp = levelsUp;
   }
 

pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeLexicalOrQualifiedMethodNode.java

@@ -47,8 +47,9 @@ public abstract sealed class AbstractInvokeLexicalOrQualifiedMethodNode
       SourceSection sourceSection,
       Identifier methodName,
       ExpressionNode[] argumentNodes,
-      boolean needsConst) {
-    super(sourceSection, argumentNodes);
+      boolean needsConst,
+      boolean argsRequireInference) {
+    super(sourceSection, argumentNodes, argsRequireInference);
     this.methodName = methodName;
     this.needsConst = needsConst;
     this.isConstChecked = false;

pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeMethodNode.java

@@ -28,10 +28,13 @@ import org.pkl.core.runtime.VmUtils;
 public abstract class AbstractInvokeMethodNode extends ExpressionNode {
 
   @Children protected final ExpressionNode[] argumentNodes;
+  protected final boolean argsRequireInference;
 
-  public AbstractInvokeMethodNode(SourceSection sourceSection, ExpressionNode[] argumentNodes) {
+  public AbstractInvokeMethodNode(
+      SourceSection sourceSection, ExpressionNode[] argumentNodes, boolean argsRequireInference) {
     super(sourceSection);
     this.argumentNodes = argumentNodes;
+    this.argsRequireInference = argsRequireInference;
   }
 
   @TruffleBoundary
@@ -44,10 +47,13 @@ public abstract class AbstractInvokeMethodNode extends ExpressionNode {
   @ExplodeLoop
   protected Object[] evalArgs(
       VirtualFrame frame, @Nullable Method method, Object owner, @Nullable Object receiver) {
-    // TODO: optimize this away when the call does not contain any implicit new args
-    var methodSlot = getMethodSlot(frame.getFrameDescriptor());
-    var prevMethod = frame.getAuxiliarySlot(methodSlot);
-    frame.setAuxiliarySlot(methodSlot, method);
+    int methodSlot = -1;
+    Object prevMethod = null;
+    if (argsRequireInference) {
+      methodSlot = getMethodSlot(frame.getFrameDescriptor());
+      prevMethod = frame.getAuxiliarySlot(methodSlot);
+      frame.setAuxiliarySlot(methodSlot, method);
+    }
 
     var args = new Object[2 + argumentNodes.length];
     args[0] = receiver;
@@ -58,7 +64,9 @@ public abstract class AbstractInvokeMethodNode extends ExpressionNode {
         args[2 + i] = argumentNodes[i].executeGeneric(frame);
       }
     } finally {
-      frame.setAuxiliarySlot(methodSlot, prevMethod);
+      if (argsRequireInference) {
+        frame.setAuxiliarySlot(methodSlot, prevMethod);
+      }
     }
 
     return args;

pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeQualifiedMethodNode.java

@@ -31,8 +31,9 @@ public abstract sealed class AbstractInvokeQualifiedMethodNode
       Identifier methodName,
       ExpressionNode[] argumentNodes,
       boolean needsConst,
-      ExpressionNode getReceiverNode) {
-    super(sourceSection, methodName, argumentNodes, needsConst);
+      ExpressionNode getReceiverNode,
+      boolean argsRequireInference) {
+    super(sourceSection, methodName, argumentNodes, needsConst, argsRequireInference);
     this.getReceiverNode = getReceiverNode;
   }
 

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeLexicalClassMethodNode.java

@@ -31,8 +31,9 @@ public final class InvokeLexicalClassMethodNode extends AbstractInvokeLexicalMet
       Identifier methodName,
       int levelsUp,
       ExpressionNode[] argumentNodes,
-      boolean needsConst) {
-    super(sourceSection, methodName, levelsUp, argumentNodes, needsConst);
+      boolean needsConst,
+      boolean argsRequireInference) {
+    super(sourceSection, methodName, levelsUp, argumentNodes, needsConst, argsRequireInference);
   }
 
   @Override

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeLexicalObjectMethodNode.java

@@ -30,8 +30,9 @@ public final class InvokeLexicalObjectMethodNode extends AbstractInvokeLexicalMe
       Identifier methodName,
       int levelsUp,
       ExpressionNode[] argumentNodes,
-      boolean needsConst) {
-    super(sourceSection, methodName, levelsUp, argumentNodes, needsConst);
+      boolean needsConst,
+      boolean argsRequireInference) {
+    super(sourceSection, methodName, levelsUp, argumentNodes, needsConst, argsRequireInference);
   }
 
   @Override

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeMethodDirectNode.java

@@ -34,9 +34,9 @@ public final class InvokeMethodDirectNode extends AbstractInvokeMethodNode {
       SourceSection sourceSection,
       ClassMethod method,
       ExpressionNode receiverNode,
-      ExpressionNode[] argumentNodes) {
-
-    super(sourceSection, argumentNodes);
+      ExpressionNode[] argumentNodes,
+      boolean argsRequireInference) {
+    super(sourceSection, argumentNodes, argsRequireInference);
     this.method = method;
     this.owner = method.getOwner();
     this.receiverNode = receiverNode;

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeMethodVirtualNode.java

@@ -52,9 +52,9 @@ public abstract class InvokeMethodVirtualNode extends AbstractInvokeMethodNode {
       Identifier methodName,
       ExpressionNode[] argumentNodes,
       MemberLookupMode lookupMode,
-      boolean needsConst) {
-
-    super(sourceSection, argumentNodes);
+      boolean needsConst,
+      boolean argsRequireInference) {
+    super(sourceSection, argumentNodes, argsRequireInference);
     this.methodName = methodName;
     this.lookupMode = lookupMode;
     this.needsConst = needsConst;
@@ -64,8 +64,9 @@ public abstract class InvokeMethodVirtualNode extends AbstractInvokeMethodNode {
       SourceSection sourceSection,
       Identifier methodName,
       ExpressionNode[] argumentNodes,
-      MemberLookupMode lookupMode) {
-    this(sourceSection, methodName, argumentNodes, lookupMode, false);
+      MemberLookupMode lookupMode,
+      boolean argsRequireInference) {
+    this(sourceSection, methodName, argumentNodes, lookupMode, false, argsRequireInference);
   }
 
   /**
@@ -145,7 +146,14 @@ public abstract class InvokeMethodVirtualNode extends AbstractInvokeMethodNode {
   @Override
   public WrapperNode createWrapper(ProbeNode probe) {
     return new InvokeMethodVirtualNodeWrapper(
-        sourceSection, methodName, argumentNodes, lookupMode, needsConst, this, probe);
+        sourceSection,
+        methodName,
+        argumentNodes,
+        lookupMode,
+        needsConst,
+        argsRequireInference,
+        this,
+        probe);
   }
 
   private void checkConst(ClassMethod method) {

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeQualifiedClassMethodNode.java

@@ -28,8 +28,15 @@ public final class InvokeQualifiedClassMethodNode extends AbstractInvokeQualifie
       Identifier methodName,
       ExpressionNode[] argumentNodes,
       boolean needsConst,
-      ExpressionNode getReceiverNode) {
-    super(sourceSection, methodName, argumentNodes, needsConst, getReceiverNode);
+      ExpressionNode getReceiverNode,
+      boolean argsRequireInference) {
+    super(
+        sourceSection,
+        methodName,
+        argumentNodes,
+        needsConst,
+        getReceiverNode,
+        argsRequireInference);
   }
 
   @Override

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeQualifiedObjectMethodNode.java

@@ -30,8 +30,15 @@ public final class InvokeQualifiedObjectMethodNode extends AbstractInvokeQualifi
       Identifier methodName,
       ExpressionNode[] argumentNodes,
       boolean needsConst,
-      ExpressionNode getReceiverNode) {
-    super(sourceSection, methodName, argumentNodes, needsConst, getReceiverNode);
+      ExpressionNode getReceiverNode,
+      boolean argsRequireInference) {
+    super(
+        sourceSection,
+        methodName,
+        argumentNodes,
+        needsConst,
+        getReceiverNode,
+        argsRequireInference);
   }
 
   @Override

pkl-core/src/main/java/org/pkl/core/ast/expression/member/InvokeSuperMethodNode.java

@@ -35,9 +35,9 @@ public abstract class InvokeSuperMethodNode extends AbstractInvokeMethodNode {
       SourceSection sourceSection,
       Identifier methodName,
       ExpressionNode[] argumentNodes,
-      boolean needsConst) {
-
-    super(sourceSection, argumentNodes);
+      boolean needsConst,
+      boolean argsRequireInference) {
+    super(sourceSection, argumentNodes, argsRequireInference);
     this.needsConst = needsConst;
 
     assert !methodName.isLocalMethod();

pkl-core/src/main/java/org/pkl/core/ast/internal/ToStringNode.java

@@ -88,6 +88,7 @@ public abstract class ToStringNode extends UnaryExpressionNode {
         Identifier.TO_STRING,
         new ExpressionNode[] {},
         MemberLookupMode.EXPLICIT_RECEIVER,
+        false,
         null,
         null);
   }

pkl-core/src/test/files/LanguageSnippetTests/input/methods/inferParameterType.pkl

@@ -24,8 +24,11 @@ local quux = (a: Foo) -> a.x
 
 
 qualified = this.bar(new { x = 1 })
-unqualifiedLexical = bar(new { x = 1 })
-unqualifiedThis = new Qux { x = 7 }.bar(new { x = 1 })
+unqualifiedLexical {
+  call = bar(new { x = 1 })
+}
+unqualifiedThis1 = bar(new { x = 1 })
+unqualifiedThis2 = new Qux { x = 7 }.bar(new { x = 1 })
 `super` = new Foo { x = 2 }.superBar(new Foo { x = 3 })
 nestedMethodCalls = outerMethod(bar(new { x = 8 }), new { x = 9 })
 objectMethod {

pkl-core/src/test/files/LanguageSnippetTests/output/methods/inferParameterType.pcf

@@ -1,6 +1,9 @@
 qualified = 1
-unqualifiedLexical = 1
-unqualifiedThis = 8
+unqualifiedLexical {
+  call = 1
+}
+unqualifiedThis1 = 1
+unqualifiedThis2 = 8
 `super` = 16
 nestedMethodCalls = 26
 objectMethod {