Fix formatting comments in operation errors list (#2283)

Fixes: #2261

For operation errors like:
```
errors // foo
: // bar
[
    // baz
    MyError
]
```
you'd expect formatting of:
```
// foo
// bar
errors: [
    // baz
    MyError
]
```
Prior to this commit, we were handling the cases of `foo` and `bar`,
but inadvertently doing the same thing for `baz`. This is because
we were looking for comments to pull out to above `errors` in the
direct children of OPERATION_ERRORS, which include all WS within the
`[]`.

To make it work as expected, we need to only pull out comments in the
first two positions (`foo`, `bar`), and leave the rest for BracketFormatter
to format. BracketFormatter was expecting to operate on all the children
of a given TreeCursor, so I modified it to act on an arbitrary stream of
cursors, and added a way to get all remaining siblings after a cursor.

smithy-syntax/src/main/java/software/amazon/smithy/syntax/BracketFormatter.java

@@ -38,7 +38,37 @@ static Function<TreeCursor, Stream<Doc>> extractor(
             Function<TreeCursor, Doc> visitor,
             Function<TreeCursor, Stream<TreeCursor>> mapper
     ) {
-        return new Extractor(visitor, mapper);
+        return extractor(visitor, mapper, TreeCursor::children);
+    }
+
+    static Function<TreeCursor, Stream<Doc>> extractor(
+            Function<TreeCursor, Doc> visitor,
+            Function<TreeCursor, Stream<TreeCursor>> mapper,
+            Function<TreeCursor, Stream<TreeCursor>> childrenSupplier
+    ) {
+        return (cursor) -> childrenSupplier.apply(cursor)
+                .flatMap(c -> {
+                    TreeType type = c.getTree().getType();
+                    return type == TreeType.WS || type == TreeType.COMMENT ? Stream.of(c) : mapper.apply(c);
+                })
+                .flatMap(c -> {
+                    // If the child extracts WS, then filter it down to just comments.
+                    return c.getTree().getType() == TreeType.WS
+                            ? c.getChildrenByType(TreeType.COMMENT).stream()
+                            : Stream.of(c);
+                })
+                .map(visitor)
+                .filter(doc -> doc != Doc.empty());
+    }
+
+    static Function<TreeCursor, Stream<TreeCursor>> byTypeMapper(TreeType childType) {
+        return child -> child.getTree().getType() == childType
+                ? Stream.of(child)
+                : Stream.empty();
+    }
+
+    static Function<TreeCursor, Stream<TreeCursor>> siblingChildrenSupplier() {
+        return cursor -> cursor.remainingSiblings().stream();
     }
 
     // Brackets children of childType between open and closed brackets. If the children can fit together
@@ -47,9 +77,7 @@ static Function<TreeCursor, Stream<Doc>> extractByType(
             TreeType childType,
             Function<TreeCursor, Doc> visitor
     ) {
-        return extractor(visitor, child -> child.getTree().getType() == childType
-                                           ? Stream.of(child)
-                                           : Stream.empty());
+        return extractor(visitor, byTypeMapper(childType));
     }
 
     BracketFormatter open(Doc open) {
@@ -138,36 +166,4 @@ private Doc getBracketLineBreakDoc() {
         }
         return lineDoc;
     }
-
-    private static final class Extractor implements Function<TreeCursor, Stream<Doc>> {
-        private final Function<TreeCursor, Stream<TreeCursor>> mapper;
-        private final Function<TreeCursor, Doc> visitor;
-
-        private Extractor(
-                Function<TreeCursor, Doc> visitor,
-                Function<TreeCursor, Stream<TreeCursor>> mapper
-        ) {
-            this.visitor = visitor;
-            this.mapper = mapper;
-        }
-
-        @Override
-        public Stream<Doc> apply(TreeCursor cursor) {
-            SmithyBuilder.requiredState("childExtractor", mapper);
-            SmithyBuilder.requiredState("visitor", visitor);
-            return cursor.children()
-                    .flatMap(c -> {
-                        TreeType type = c.getTree().getType();
-                        return type == TreeType.WS || type == TreeType.COMMENT ? Stream.of(c) : mapper.apply(c);
-                    })
-                    .flatMap(c -> {
-                        // If the child extracts WS, then filter it down to just comments.
-                        return c.getTree().getType() == TreeType.WS
-                               ? c.getChildrenByType(TreeType.COMMENT).stream()
-                               : Stream.of(c);
-                    })
-                    .map(visitor)
-                    .filter(doc -> doc != Doc.empty());
-        }
-    }
 }

smithy-syntax/src/main/java/software/amazon/smithy/syntax/FormatVisitor.java

@@ -263,15 +263,29 @@ Doc visit(TreeCursor cursor) {
             }
 
             case OPERATION_ERRORS: {
-                return skippedComments(cursor, false)
-                        .append(Doc.text("errors: "))
+                // Pull out any comments that come after 'errors' but before the opening '[' so they
+                // can be placed before 'errors: ['
+                Doc comments = Doc.empty();
+                TreeCursor child = cursor.getFirstChild(); // 'errors'
+                if (child.getNextSibling().getTree().getType() == TreeType.WS) {
+                    comments = comments.append(skippedComments(child.getNextSibling(), false));
+                    child = child.getNextSibling(); // skip ws
+                }
+                child = child.getNextSibling(); // ':'
+                if (child.getNextSibling().getTree().getType() == TreeType.WS) {
+                    comments = comments.append(skippedComments(child.getNextSibling(), false));
+                    child = child.getNextSibling();
+                }
+                return comments.append(Doc.text("errors: ")
                         .append(new BracketFormatter()
                                         .open(Formatter.LBRACKET)
                                         .close(Formatter.RBRACKET)
-                                        .extractChildren(cursor, BracketFormatter.extractByType(TreeType.SHAPE_ID,
-                                                                                                this::visit))
+                                        .extractChildren(child, BracketFormatter.extractor(
+                                                this::visit,
+                                                BracketFormatter.byTypeMapper(TreeType.SHAPE_ID),
+                                                BracketFormatter.siblingChildrenSupplier()))
                                         .forceLineBreaks() // always put each error on separate lines.
-                                        .write());
+                                        .write()));
             }
 
             case MIXINS: {

smithy-syntax/src/main/java/software/amazon/smithy/syntax/TreeCursor.java

@@ -318,6 +318,19 @@ public TreeCursor findAt(int line, int column) {
         }
     }
 
+    /**
+     * @return All siblings after this cursor, in order.
+     */
+    public List<TreeCursor> remainingSiblings() {
+        List<TreeCursor> remaining = new ArrayList<>();
+        TreeCursor nextSibling = getNextSibling();
+        while (nextSibling != null) {
+            remaining.add(nextSibling);
+            nextSibling = nextSibling.getNextSibling();
+        }
+        return remaining;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {

smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/operation-errors-comments.formatted.smithy

@@ -0,0 +1,27 @@
+$version: "2.0"
+
+namespace smithy.example
+
+operation Foo {
+    errors: [
+        // a
+        E1
+        // c
+        E2
+        // d
+    ]
+}
+
+operation Bar {
+    // a
+    // b
+    errors: [
+        // c
+    ]
+}
+
+@error("client")
+structure E1 {}
+
+@error("client")
+structure E2 {}

smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/operation-errors-comments.smithy

@@ -0,0 +1,27 @@
+$version: "2.0"
+
+namespace smithy.example
+
+operation Foo {
+    errors: [
+        // a
+        E1
+        // c
+        E2
+        // d
+    ]
+}
+
+operation Bar {
+    errors // a
+    : // b
+    [
+        // c
+    ]
+}
+
+@error("client")
+structure E1 {}
+
+@error("client")
+structure E2 {}