Extending jmespath expressions

tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/transform/CppEndpointsJmesPathVisitor.java

@@ -1,6 +1,12 @@
 package com.amazonaws.util.awsclientgenerator.transform;
 
-import software.amazon.smithy.utils.Pair;
+import com.amazonaws.util.awsclientgenerator.domainmodels.codegeneration.Shape;
+import com.amazonaws.util.awsclientgenerator.domainmodels.codegeneration.ShapeMember;
+import com.amazonaws.util.awsclientgenerator.domainmodels.codegeneration.cpp.CppViewHelper;
+import java.text.MessageFormat;
+import java.util.Optional;
+import software.amazon.smithy.build.SmithyBuildException;
+import software.amazon.smithy.jmespath.ExpressionVisitor;
 import software.amazon.smithy.jmespath.ast.AndExpression;
 import software.amazon.smithy.jmespath.ast.ComparatorExpression;
 import software.amazon.smithy.jmespath.ast.CurrentExpression;
@@ -16,18 +22,10 @@
 import software.amazon.smithy.jmespath.ast.NotExpression;
 import software.amazon.smithy.jmespath.ast.ObjectProjectionExpression;
 import software.amazon.smithy.jmespath.ast.OrExpression;
-import software.amazon.smithy.build.SmithyBuildException;
 import software.amazon.smithy.jmespath.ast.ProjectionExpression;
 import software.amazon.smithy.jmespath.ast.SliceExpression;
 import software.amazon.smithy.jmespath.ast.Subexpression;
-
-import software.amazon.smithy.jmespath.ExpressionVisitor;
-import java.text.MessageFormat;
-import com.amazonaws.util.awsclientgenerator.domainmodels.codegeneration.Shape;
-import com.amazonaws.util.awsclientgenerator.domainmodels.codegeneration.ShapeMember;
-import com.amazonaws.util.awsclientgenerator.domainmodels.codegeneration.cpp.CppViewHelper;
-
-import java.util.*;
+import software.amazon.smithy.utils.Pair;
 
 class CppEndpointsJmesPathVisitor implements ExpressionVisitor<Pair<String, Shape>> {
     final OperationContextCppCodeGenerator context;
@@ -45,20 +43,19 @@ public Pair<String, Shape> visitObjectProjection(ObjectProjectionExpression expr
 
     @Override
     public Pair<String, Shape> visitProjection(ProjectionExpression expression) {
-        return Optional.of(
-            expression.getLeft().accept(this)).filter( elem -> elem.right.isList()).map(
-                elem -> {
-                    String varName = elem.left + "Elem";
-                    context.rangeBasedForLoop(varName);
-                    context.openVariableScope(varName);
-                    Pair<String, Shape> right = expression.getRight().accept(
-                            new CppEndpointsJmesPathVisitor(context, elem.right.getListMember().getShape()));
-                    context.closeVariableScope();
-                    return Pair.of(
-                        elem.left,
-                        right.right);
-                }
-            ).orElseThrow(() -> new SmithyBuildException("Unsupported JMESPath expression"));
+      return Optional.of(expression.getLeft().accept(this))
+          .filter(elem -> elem.right.isList())
+          .map(elem -> {
+            String varName = elem.left + "Elem";
+            context.rangeBasedForLoop(varName);
+            context.openVariableScope(varName);
+            Pair<String, Shape> right =
+                expression.getRight().accept(new CppEndpointsJmesPathVisitor(
+                    context, elem.right.getListMember().getShape()));
+            context.closeVariableScope();
+            return Pair.of(elem.left, right.right);
+          })
+          .orElse(Pair.of("", this.input));
     }
 
     @Override
@@ -100,10 +97,10 @@ public Pair<String, Shape> visitField(FieldExpression expression) {
         // at the start of each code block
         String varName = expression.getName() + "Elem";
         // if a new scope started, declare variable accessed
-        if (context.isStartOfNewScope() ||
-                context.getVarName().isEmpty()) {
-            context.addVariableInScope(varName);
+        if (context.isStartOfNewScope() || context.getVarName().isEmpty()) {
+          context.addVariableInScope(varName);
         }
+
         // chain accessors
         context.getCppCode()
                 .append(MessageFormat.format(".Get{0}()", CppViewHelper.convertToUpperCamel(expression.getName())));
@@ -112,6 +109,10 @@ public Pair<String, Shape> visitField(FieldExpression expression) {
             // if leaf element, push to result
             if (member.getShape().isString()) {
                 context.addInScopeVariableToResult(Optional.empty());
+
+                // this is where we can refer to the original scope variable
+                // again
+                context.cleanupVariablesCurrentScope();
             }
         }
         return Pair.of(
@@ -146,7 +147,13 @@ public Pair<String, Shape> visitCurrentNode(CurrentExpression expression) {
 
     @Override
     public Pair<String, Shape> visitFlatten(FlattenExpression expression) {
-        throw new SmithyBuildException("Unsupported JMESPath expression");
+      if (expression.getExpression() instanceof ProjectionExpression) {
+        return visitProjection(
+            (ProjectionExpression)expression.getExpression());
+      }
+      // If it's not a ProjectionExpression, proceed with handling
+      // FlattenExpression
+      return expression.accept(this);
     }
 
     @Override
@@ -161,9 +168,9 @@ public Pair<String, Shape> visitLiteral(LiteralExpression expression) {
 
     @Override
     public Pair<String, Shape> visitMultiSelectList(MultiSelectListExpression expression) {
-        throw new SmithyBuildException("Unsupported JMESPath expression");
+      expression.getExpressions().forEach(expr -> { expr.accept(this); });
+      return Pair.of("", this.input);
     }
-
     @Override
     public Pair<String, Shape> visitMultiSelectHash(MultiSelectHashExpression expression) {
         throw new SmithyBuildException("Unsupported JMESPath expression");

tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/transform/OperationContextCppCodeGenerator.java

@@ -1,8 +1,8 @@
 package com.amazonaws.util.awsclientgenerator.transform;
-import lombok.Data;
 import java.text.MessageFormat;
-import java.util.Stack;
 import java.util.Optional;
+import java.util.Stack;
+import lombok.Data;
 import software.amazon.smithy.utils.Pair;
 
 @Data
@@ -41,7 +41,13 @@ public void addVariableInScope(String var)
         }
         this.varName.push(Pair.of(var, false));
     }
-
+
+    public void cleanupVariablesCurrentScope() {
+      while (!this.getVarName().isEmpty() && !this.getVarName().peek().right) {
+        this.getVarName().pop();
+      }
+    }
+
     public void rangeBasedForLoop(String varName)
     {
         this.getCppCode().append(MessageFormat.format("{2}for (auto& {0} : {1})\n", varName, this.getVarName().peek().left, this.getIndentationPrefix()) );
@@ -79,10 +85,7 @@ public void closeVariableScope()
     }
 
     public boolean isStartOfNewScope(){
-        return (
-                this.getCppCode().length() > 2 &&
-                this.getCppCode().substring(this.getCppCode().length() - 2).equals("{\n")
-            );
+      return (!this.getVarName().isEmpty()) && this.getVarName().peek().right;
     }
 
     public void addInScopeVariableToResult(Optional<String> accessorSuffix)