Omit unnecessary brackets for single-valued container annotations

Author: JonasKunz

src/main/java/co/elastic/indytransformer/AdviceTransformationPlan.java

@@ -202,19 +202,20 @@ private List<ValueToReturn<?>> transformArgumentAssignments(List<Parameter> ente
             returns.add(new ValueToReturn<>(writtenArg, (method, index) -> {
                 Utils.addImports(method, Advice.AssignReturned.class, Advice.AssignReturned.ToArguments.ToArgument.class);
                 SingleMemberAnnotationExpr toArguments = findOrCreateRepeatableWrapperAnnotation(method, "AssignReturned.ToArguments");
-                ArrayInitializerExpr toArgumentsInitializerList = (ArrayInitializerExpr) (toArguments).getMemberValue();
 
+                AnnotationExpr result;
                 Name annoName = new Name("ToArgument");
                 IntegerLiteralExpr argIndexExpr = new IntegerLiteralExpr(String.valueOf(argIndex));
                 if (index == -1) {
-                    toArgumentsInitializerList.getValues().add(new SingleMemberAnnotationExpr(annoName, argIndexExpr));
+                    result = new SingleMemberAnnotationExpr(annoName, argIndexExpr);
                 } else {
-                    toArgumentsInitializerList.getValues().add(new NormalAnnotationExpr(
+                    result = new NormalAnnotationExpr(
                             annoName, new NodeList<>(
                             new MemberValuePair("value", argIndexExpr),
                             new MemberValuePair("index", new IntegerLiteralExpr(String.valueOf(index)))
-                    )));
+                    ));
                 }
+                addToArrayAnnotationValue(toArguments, result);
             }));
         }
         return returns;
@@ -279,7 +280,6 @@ private List<ValueToReturn<Parameter>> transformFieldAssignments(List<Parameter>
                         Utils.addImports(method, Advice.AssignReturned.class, Advice.AssignReturned.ToFields.ToField.class);
 
                         SingleMemberAnnotationExpr toFields = findOrCreateRepeatableWrapperAnnotation(method, "AssignReturned.ToFields");
-                        ArrayInitializerExpr toFieldsInitializerList = (ArrayInitializerExpr) (toFields).getMemberValue();
 
                         AnnotationExpr result;
                         Name annoName = new Name("ToField");
@@ -292,12 +292,29 @@ private List<ValueToReturn<Parameter>> transformFieldAssignments(List<Parameter>
                                     new MemberValuePair("index", new IntegerLiteralExpr(String.valueOf(index)))
                             ));
                         }
-                        toFieldsInitializerList.getValues().add(result);
+                        addToArrayAnnotationValue(toFields, result);
                     });
                 })
                 .toList();
     }
 
+    private static void addToArrayAnnotationValue(SingleMemberAnnotationExpr containerAnnotation, AnnotationExpr valueToAdd) {
+        Expression originalValue = containerAnnotation.getMemberValue();
+        ArrayInitializerExpr newArrayValue = new ArrayInitializerExpr();
+        if (originalValue instanceof ArrayInitializerExpr originalArray) {
+            newArrayValue = originalArray;
+        } else {
+            newArrayValue.getValues().add(originalValue);
+        }
+        newArrayValue.getValues().add(valueToAdd);
+        if (newArrayValue.getValues().size() == 1) {
+            // No need to wrap in array initializer expression if we have exactly one element
+            containerAnnotation.setMemberValue(newArrayValue.getValues().get(0));
+        } else {
+            containerAnnotation.setMemberValue(newArrayValue);
+        }
+    }
+
 
     private void removeParameterOrVariableIfNeverUsed(MethodDeclaration method, Node paramOrVarDeclaration) {
         AtomicBoolean isUsed = new AtomicBoolean(false);