Codegen direct index access and map member model (#5720)

* Add codegen support for direct index access and map member model

* Add codegen tests

Author: davidh44

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointRulesClientTestSpec.java

@@ -20,6 +20,7 @@
 import com.fasterxml.jackson.core.TreeNode;
 import com.fasterxml.jackson.jr.stree.JrsArray;
 import com.fasterxml.jackson.jr.stree.JrsObject;
+import com.fasterxml.jackson.jr.stree.JrsString;
 import com.fasterxml.jackson.jr.stree.JrsValue;
 import com.squareup.javapoet.AnnotationSpec;
 import com.squareup.javapoet.ClassName;
@@ -68,6 +69,7 @@
 import software.amazon.awssdk.core.rules.testing.util.EmptyPublisher;
 import software.amazon.awssdk.core.sync.RequestBody;
 import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.utils.ImmutableMap;
 import software.amazon.awssdk.utils.Validate;
 
 public class EndpointRulesClientTestSpec implements ClassSpec {
@@ -781,8 +783,26 @@ private CodeBlock createMemberValue(MemberModel memberModel, TreeNode valueNode)
         }
 
         if (memberModel.isMap()) {
-            // Not necessary at the moment
-            throw new RuntimeException("Don't know how to create map member.");
+            Iterator<Map.Entry<String, JrsValue>> fieldsIter = ((JrsObject) valueNode).fields();
+            MemberModel keyModel = memberModel.getMapModel().getKeyModel();
+            MemberModel valueModel = memberModel.getMapModel().getValueModel();
+
+            b.add("$T.of(", ImmutableMap.class);
+            boolean firstKey = true;
+            while (fieldsIter.hasNext()) {
+                Map.Entry<String, JrsValue> entry = fieldsIter.next();
+                JrsString keyNode = new JrsString(entry.getKey());
+
+                if (!firstKey) {
+                    b.add(", ");
+                }
+                firstKey = false;
+                b.add(createMemberValue(keyModel, keyNode));
+                b.add(", ");
+                b.add(createMemberValue(valueModel, entry.getValue()));
+            }
+            b.add(")");
+            return b.build();
         }
 
         return createModelClass(model.getShapes().get(memberModel.getC2jShape()), valueNode);

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

@@ -18,7 +18,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 import software.amazon.awssdk.utils.ToString;
 
 /**
@@ -64,7 +63,6 @@ public RuleExpression visitVariableReferenceExpression(VariableReferenceExpressi
             type = tableBuilder.param(variableName);
         }
         if (type == null) {
-            errors.add(String.format("Undefined variable `%s`", variableName));
             return e;
         }
         return e.toBuilder()
@@ -90,12 +88,6 @@ public RuleExpression visitFunctionCallExpression(FunctionCallExpression e) {
             return expr;
         }
 
-        List<RuleType> args = expr.arguments().stream().map(RuleExpression::type).collect(Collectors.toList());
-        if (!function.matches(args)) {
-            addError("Arguments for function `%s` doesn't match, expected: `%s`, got: `%s`",
-                     name,
-                     function.arguments().values(), args);
-        }
         return expr.toBuilder().type(function.returns()).build();
     }
 
@@ -160,7 +152,13 @@ public RuleExpression visitMemberAccessExpression(MemberAccessExpression e) {
             addError("Cannot find type for expression `%s`", expr.source());
         } else {
             String name = expr.name();
-            RuleType memberType = type.property(name);
+            RuleType memberType;
+            if (expr.directIndex()) {
+                memberType = type;
+            } else {
+                memberType = type.property(name);
+            }
+
             if (memberType == null) {
                 addError("Cannot find a member with name `%s` for type `%s`", name, type);
             } else {

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

@@ -141,7 +141,10 @@ public Void visitVariableReferenceExpression(VariableReferenceExpression e) {
     @Override
     public Void visitMemberAccessExpression(MemberAccessExpression e) {
         e.source().accept(this);
-        builder.add(".$L()", e.name());
+        if (!e.directIndex()) {
+            builder.add(".$L()", e.name());
+        }
+
         return null;
     }
 
@@ -197,25 +200,25 @@ public Void visitLetExpression(LetExpression expr) {
             RuleType type = symbolTable.locals().get(key);
             builder.addStatement("$T $L = null", type.javaType(), key);
         }
-        builder.add("if (");
-        boolean isFirst = true;
+
+        int count = 0;
         for (Map.Entry<String, RuleExpression> kvp : expr.bindings().entrySet()) {
             String k = kvp.getKey();
             RuleExpression v = kvp.getValue();
-            if (!isFirst) {
-                builder.add(" && ");
-            }
+            builder.add("if (");
             builder.add("($L = ", k);
             v.accept(this);
             builder.add(") != null");
-            isFirst = false;
+
+            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.beginControlFlow(")");
-        builder.add("locals = locals.toBuilder()");
-        expr.bindings().forEach((k, v) -> {
-            builder.add(".$1L($1L)", k);
-        });
-        builder.addStatement(".build()");
         return null;
     }
 

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

@@ -316,6 +316,15 @@ private static RuleExpression getAttrExpression(List<TreeNode> argv) {
             tokenizer.expectAtEof("indexed access");
             return indexedAccessBuilder.build();
         }
+        if (tokenizer.isDirectIndexedAccess()) {
+            IndexedAccessExpression.Builder indexedAccessBuilder = IndexedAccessExpression.builder();
+            tokenizer.consumeDirectIndexed(i -> indexedAccessBuilder
+                .source(memberAccessBuilder.directIndex(true).build())
+                .index(i)
+                .build());
+            tokenizer.expectAtEof("indexed access");
+            return indexedAccessBuilder.build();
+        }
         if (tokenizer.isIdentifier()) {
             tokenizer.consumeIdentifier(memberAccessBuilder::name);
             tokenizer.expectAtEof("member access");

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

@@ -25,11 +25,13 @@ public final class MemberAccessExpression implements RuleExpression {
     private final RuleType type;
     private final RuleExpression source;
     private final String name;
+    private final boolean directIndex;
 
     MemberAccessExpression(Builder builder) {
         this.type = builder.type;
         this.source = Validate.paramNotNull(builder.source, "source");
-        this.name = Validate.paramNotNull(builder.name, "name");
+        this.directIndex = builder.directIndex != null ? builder.directIndex : false;
+        this.name = this.directIndex ? builder.name : Validate.paramNotNull(builder.name, "name");
     }
 
     public static Builder builder() {
@@ -68,6 +70,10 @@ public String name() {
         return name;
     }
 
+    public boolean directIndex() {
+        return directIndex;
+    }
+
     public Builder toBuilder() {
         return new Builder(this);
     }
@@ -88,6 +94,9 @@ public boolean equals(Object o) {
 
         MemberAccessExpression that = (MemberAccessExpression) o;
 
+        if (directIndex != that.directIndex) {
+            return false;
+        }
         if (!Objects.equals(type, that.type)) {
             return false;
         }
@@ -102,13 +111,15 @@ public int hashCode() {
         int result = type != null ? type.hashCode() : 0;
         result = 31 * result + source.hashCode();
         result = 31 * result + name.hashCode();
+        result = 31 * result + Boolean.hashCode(directIndex);
         return result;
     }
 
     public static class Builder {
         private RuleType type;
         private RuleExpression source;
         private String name;
+        private Boolean directIndex;
 
         private Builder() {
         }
@@ -117,6 +128,7 @@ private Builder(MemberAccessExpression expr) {
             this.type = expr.type;
             this.source = expr.source;
             this.name = expr.name;
+            this.directIndex = expr.directIndex;
         }
 
         public Builder source(RuleExpression source) {
@@ -134,6 +146,11 @@ public Builder type(RuleType type) {
             return this;
         }
 
+        public Builder directIndex(Boolean directIndex) {
+            this.directIndex = directIndex;
+            return this;
+        }
+
         public MemberAccessExpression build() {
             return new MemberAccessExpression(this);
         }

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

@@ -76,6 +76,7 @@ public RuleExpression visitMemberAccessExpression(MemberAccessExpression e) {
             .builder()
             .name(e.name())
             .type(e.type())
+            .directIndex(e.directIndex())
             .source(e.source().accept(this));
         return builder.build();
     }

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

@@ -211,6 +211,19 @@ public void consumeIndexed(BiConsumer<String, Integer> consumer) {
         index += 4;
     }
 
+    // e.g., [123]
+    public boolean isDirectIndexedAccess() {
+        return matches(TokenKind.OPEN_SQUARE, TokenKind.NUMBER, TokenKind.CLOSE_SQUARE);
+    }
+
+    public void consumeDirectIndexed(Consumer<Integer> consumer) {
+        if (!isDirectIndexedAccess()) {
+            throw new IllegalStateException("not at direct indexed");
+        }
+        consumer.accept(Integer.parseInt(tokens.get(index + 1).value));
+        index += 3;
+    }
+
     // e.g., {url#scheme}
     public boolean isNamedAccess() {
         return matches(TokenKind.OPEN_CURLY, TokenKind.IDENTIFIER, TokenKind.HASH, TokenKind.IDENTIFIER, TokenKind.CLOSE_CURLY);

codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules2/TokenizerTest.java

@@ -36,11 +36,19 @@ public void recognizesIndexedExpression() {
         assertTrue(tokenizer.isIndexedAccess());
         tokenizer.consumeIndexed((name, index) -> {
             assertEquals("foobar", name);
-            assertEquals(index, 123);
+            assertEquals(123, index);
         });
         assertTrue(tokenizer.atEof());
     }
 
+    @Test
+    public void recognizesDirectIndexedExpression() {
+        Tokenizer tokenizer = new Tokenizer("[123]");
+        assertTrue(tokenizer.isDirectIndexedAccess());
+        tokenizer.consumeDirectIndexed(i -> assertEquals(123, i));
+        assertTrue(tokenizer.atEof());
+    }
+
     @Test
     public void recognizesNamedAccessExpression() {
         Tokenizer tokenizer = new Tokenizer("{url#authority}");

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/query/customization-endpoint-know-prop.config

@@ -1,5 +1,12 @@
 {
   "endpointAuthSchemeConfig": {
     "knownEndpointProperties": "software.module.test.KNOWN_PROPS"
-  }
+  },
+    "endpointParameters": {
+      "ArnList": {
+            "required": false,
+            "documentation": "Parameter from the customization config",
+            "type": "StringArray"
+      }
+    }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/query/customization.config

@@ -10,6 +10,11 @@
       "required": false,
       "documentation": "Parameter from the customization config",
       "type": "StringArray"
+    },
+    "ArnList": {
+          "required": false,
+          "documentation": "Parameter from the customization config",
+          "type": "StringArray"
     }
   },
   "customOperationContextParams": [