Pass locals as arguments in endpoint rules (#6131)

* Pass locals as arguments in endpoint rules

* Add change log entry

* Ensure we always return from a tree method

Author: sugmanue

.changes/next-release/feature-AWSSDKforJavav2-3289c08.json

@@ -0,0 +1,6 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "contributor": "", 
+    "type": "feature", 
+    "description": "Improve the endpoint rules performance by directly passing the needed params instead of using a POJO to keep track of them."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules2/CodeGeneratorVisitor.java

@@ -21,6 +21,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import software.amazon.awssdk.awscore.endpoints.AwsEndpointAttribute;
 import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4AuthScheme;
 import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4aAuthScheme;
@@ -32,14 +33,17 @@ public class CodeGeneratorVisitor extends WalkRuleExpressionVisitor {
     private final RuleRuntimeTypeMirror typeMirror;
     private final SymbolTable symbolTable;
     private final Map<String, KeyTypePair> knownEndpointAttributes;
+    private final Map<String, ComputeScopeTree.Scope> ruleIdToScope;
 
     public CodeGeneratorVisitor(RuleRuntimeTypeMirror typeMirror,
                                 SymbolTable symbolTable,
                                 Map<String, KeyTypePair> knownEndpointAttributes,
+                                Map<String, ComputeScopeTree.Scope> ruleIdToScope,
                                 CodeBlock.Builder builder) {
         this.builder = builder;
         this.symbolTable = symbolTable;
         this.knownEndpointAttributes = knownEndpointAttributes;
+        this.ruleIdToScope = ruleIdToScope;
         this.typeMirror = typeMirror;
     }
 
@@ -196,28 +200,14 @@ public Void visitRuleSetExpression(RuleSetExpression e) {
 
     @Override
     public Void visitLetExpression(LetExpression expr) {
-        for (String key : expr.bindings().keySet()) {
-            RuleType type = symbolTable.locals().get(key);
-            builder.addStatement("$T $L = null", type.javaType(), key);
-        }
-
-        int count = 0;
         for (Map.Entry<String, RuleExpression> kvp : expr.bindings().entrySet()) {
             String k = kvp.getKey();
             RuleExpression v = kvp.getValue();
-            builder.add("if (");
-            builder.add("($L = ", k);
+            RuleType type = symbolTable.locals().get(k);
+            builder.add("$T $L = ", type.javaType(), k);
             v.accept(this);
-            builder.add(") != null");
-
-            builder.beginControlFlow(")");
-            builder.addStatement("locals = locals.toBuilder().$1L($1L).build()", k);
-
-            if (++count < expr.bindings().size()) {
-                builder.nextControlFlow("else");
-                builder.addStatement("return RuleResult.carryOn()");
-                builder.endControlFlow();
-            }
+            builder.addStatement("");
+            builder.beginControlFlow("if ($L != null)", k);
         }
         return null;
     }
@@ -235,40 +225,101 @@ private void conditionsPreamble(RuleSetExpression expr) {
     }
 
     private void conditionsEpilogue(RuleSetExpression expr) {
-        int blocksToClose = expr.conditions().size();
-        for (int idx = 0; idx < blocksToClose; ++idx) {
-            builder.endControlFlow();
+        for (RuleExpression condition : expr.conditions()) {
+            if (condition.kind() == RuleExpression.RuleExpressionKind.LET) {
+                LetExpression let = (LetExpression) condition;
+                for (int x = 0; x < let.bindings().size(); x++) {
+                    builder.endControlFlow();
+                }
+            } else {
+                builder.endControlFlow();
+            }
         }
-        if (!expr.conditions().isEmpty()) {
+        if (needsReturn(expr)) {
             builder.addStatement("return $T.carryOn()", typeMirror.rulesResult().type());
         }
     }
 
+    private boolean needsReturn(RuleSetExpression expr) {
+        // If the expression can be inlined, then it doesn't live in
+        // its own method, no return at the end required
+        if (canBeInlined(expr)) {
+            return false;
+        }
+        // If the expression has conditions all be be wrapped in
+        // if-blocks, thus at the end of the method we need to return
+        // carryOn()
+        if (!expr.conditions().isEmpty()) {
+            return true;
+        }
+        // If the expression doesn't have any conditions, and doesn't
+        // have any children then we need to return carryOn(). This
+        // case SHOULD NOT happen but we assume below that there are
+        // children, thus adding the test here.
+        if (expr.children().isEmpty()) {
+            return true;
+        }
+        // We have children, check the last one.
+        int size = expr.children().size();
+        RuleSetExpression child = expr.children().get(size - 1);
+        // If a tree then we don't need a return.
+        if (child.isTree()) {
+            return false;
+        }
+        // The child is not a tree, so it was inlined. Check if it
+        // does have any conditions, if it so, its body will be inside
+        // a block already so we need to return after it.
+        return !child.conditions().isEmpty();
+    }
+
     private void codegenTreeBody(RuleSetExpression expr) {
         List<RuleSetExpression> children = expr.children();
         int size = children.size();
+        boolean isFirst = true;
         for (int idx = 0; idx < size; ++idx) {
             RuleSetExpression child = children.get(idx);
+            if (canBeInlined(child)) {
+                child.accept(this);
+                continue;
+            }
             boolean isLast = idx == size - 1;
             if (isLast) {
-                builder.addStatement("return $L(params, locals)",
-                                     child.ruleId());
+                builder.addStatement("return $L($L)",
+                                     child.ruleId(),
+                                     callParams(child.ruleId()));
                 continue;
             }
-            boolean isFirst = idx == 0;
+
             if (isFirst) {
-                builder.addStatement("$T result = $L(params, locals)",
+                isFirst = false;
+                builder.addStatement("$T result = $L($L)",
                                      typeMirror.rulesResult().type(),
-                                     child.ruleId());
+                                     child.ruleId(),
+                                     callParams(child.ruleId()));
             } else {
-                builder.addStatement("result = $L(params, locals)",
-                                     child.ruleId());
+                builder.addStatement("result = $L($L)",
+                                     child.ruleId(),
+                                     callParams(child.ruleId()));
             }
             builder.beginControlFlow("if (result.isResolved())")
                    .addStatement("return result")
                    .endControlFlow();
         }
+    }
 
+    private boolean canBeInlined(RuleSetExpression child) {
+        return !child.isTree();
+    }
+
+    private String callParams(String ruleId) {
+        ComputeScopeTree.Scope scope = ruleIdToScope.get(ruleId);
+        String args = scope.usesLocals().stream()
+                           .filter(a -> !scope.defines().contains(a))
+                           .collect(Collectors.joining(", "));
+        if (args.isEmpty()) {
+            return "params";
+        }
+        return "params, " + args;
     }
 
     @Override
@@ -381,7 +432,6 @@ private void addAttributeBlock(String k, RuleExpression v) {
         builder.add(")");
     }
 
-
     public CodeBlock.Builder builder() {
         return builder;
     }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules2/CodegenExpressionBuidler.java

@@ -21,10 +21,16 @@
 public final class CodegenExpressionBuidler {
     private final RuleSetExpression root;
     private final SymbolTable symbolTable;
+    private final Map<String, ComputeScopeTree.Scope> scopesByName;
 
-    public CodegenExpressionBuidler(RuleSetExpression root, SymbolTable symbolTable) {
+    public CodegenExpressionBuidler(
+        RuleSetExpression root,
+        SymbolTable symbolTable,
+        Map<String, ComputeScopeTree.Scope> scopesByName
+    ) {
         this.root = root;
         this.symbolTable = symbolTable;
+        this.scopesByName = scopesByName;
     }
 
     public static CodegenExpressionBuidler from(RuleSetExpression root, RuleRuntimeTypeMirror typeMirror, SymbolTable table) {
@@ -36,10 +42,17 @@ public static CodegenExpressionBuidler from(RuleSetExpression root, RuleRuntimeT
         }
         table = assignTypesVisitor.symbolTable();
         root = assignIdentifier(root);
-        PrepareForCodegenVisitor prepareForCodegenVisitor = new PrepareForCodegenVisitor(table);
-        root = (RuleSetExpression) root.accept(prepareForCodegenVisitor);
-        table = prepareForCodegenVisitor.symbolTable();
-        return new CodegenExpressionBuidler(root, table);
+
+        RenameForCodegenVisitor renameForCodegenVisitor = new RenameForCodegenVisitor(table);
+        root = (RuleSetExpression) root.accept(renameForCodegenVisitor);
+        table = renameForCodegenVisitor.symbolTable();
+
+        ComputeScopeTree computeScopeTree = new ComputeScopeTree(table);
+        root.accept(computeScopeTree);
+
+        PrepareForCodegenVisitor prepareForCodegenVisitor = new PrepareForCodegenVisitor();
+        RuleSetExpression newRoot = (RuleSetExpression) root.accept(prepareForCodegenVisitor);
+        return new CodegenExpressionBuidler(newRoot, table, computeScopeTree.scopesByName());
     }
 
     private static RuleSetExpression assignIdentifier(RuleSetExpression root) {
@@ -51,27 +64,15 @@ public RuleSetExpression root() {
         return root;
     }
 
-    public boolean isParam(String name) {
-        return symbolTable.isParam(name);
-    }
-
-    public boolean isLocal(String name) {
-        return symbolTable.isLocal(name);
-    }
-
     public String regionParamName() {
         return symbolTable.regionParamName();
     }
 
-    public Map<String, RuleType> locals() {
-        return symbolTable.locals();
-    }
-
-    public Map<String, RuleType> params() {
-        return symbolTable.params();
-    }
-
     public SymbolTable symbolTable() {
         return symbolTable;
     }
+
+    public Map<String, ComputeScopeTree.Scope> scopesByName() {
+        return scopesByName;
+    }
 }