Repository: smithy-lang/smithy
Pull request #2084 (2023-12-23 16:35:47)

Fix tree node start and end locations
Some tree types are guaranteed to have certain children, even if those children are empty. For example, SHAPE_STATEMENT always has a TRAIT_STATEMENT as its first child. Since trees with children have their start/end line and column locations computed from the respective locations of their first and last child, trees with 0 children ("empty" trees) will have a location of (0,0)-(0,0). If an empty child is the first or last child of its parent, the 0 location will bubble up to that parent.

This commit changes the way start/end line and column locations are computed for trees with children, such that empty children are ignored. This required a change to use statement sorting, which relied on the 0 location bubbling up to the parent to indicate to the formatter whether the use statement has a line comment.

Additional tests were added for the formatter and tree location computation. TokenTreeLeaf toString was also updated to include tree location, useful for debugging purposes.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.



Patch:
diff --git a/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/validators/AwsSpecialCaseEndpointValidator.java b/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/validators/AwsSpecialCaseEndpointValidator.java
index a5591c87a..5294b04f6 100644
--- a/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/validators/AwsSpecialCaseEndpointValidator.java
+++ b/smithy-aws-endpoints/src/main/java/software/amazon/smithy/rulesengine/aws/validators/AwsSpecialCaseEndpointValidator.java
@@ -115,7 +115,7 @@ public final class AwsSpecialCaseEndpointValidator extends AbstractValidator {
                     "InvalidEndpointPatternScheme"));
         }
 
-        if (!isValidURL(endpoint)) {
+        if (!isValidUrl(endpoint)) {
             events.add(error(
                     serviceShape, location,
                     String.format("Endpoint `%s` should be a valid URL.",
@@ -126,16 +126,14 @@ public final class AwsSpecialCaseEndpointValidator extends AbstractValidator {
         return events;
     }
 
-    private boolean isValidURL(String endpointPattern) {
+    private boolean isValidUrl(String endpointPattern) {
         String url = endpointPattern
                 .replace("{", "")
                 .replace("}", "");
         try {
             new URL(url).toURI();
             return true;
-        } catch (MalformedURLException e) {
-            return false;
-        } catch (URISyntaxException e) {
+        } catch (MalformedURLException | URISyntaxException e) {
             return false;
         }
     }
diff --git a/smithy-diff/src/main/java/software/amazon/smithy/diff/evaluators/RemovedShape.java b/smithy-diff/src/main/java/software/amazon/smithy/diff/evaluators/RemovedShape.java
index c0abb8a9d..8b4913433 100644
--- a/smithy-diff/src/main/java/software/amazon/smithy/diff/evaluators/RemovedShape.java
+++ b/smithy-diff/src/main/java/software/amazon/smithy/diff/evaluators/RemovedShape.java
@@ -19,11 +19,15 @@ import java.util.List;
 import java.util.stream.Collectors;
 import software.amazon.smithy.diff.Differences;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeType;
+import software.amazon.smithy.model.traits.EnumTrait;
 import software.amazon.smithy.model.traits.PrivateTrait;
+import software.amazon.smithy.model.validation.Severity;
 import software.amazon.smithy.model.validation.ValidationEvent;
 
 /**
- * Creates an ERROR event when a non-private shape is removed.
+ * Creates an ERROR event when a non-private non-scalar shape is removed.
+ * Creates a WARNING event when a non-private scalar shape is removed.
  */
 public final class RemovedShape extends AbstractDiffEvaluator {
     @Override
@@ -31,7 +35,15 @@ public final class RemovedShape extends AbstractDiffEvaluator {
         return differences.removedShapes()
                 .filter(shape -> !shape.hasTrait(PrivateTrait.class))
                 .filter(shape -> !isMemberOfRemovedShape(shape, differences))
-                .map(shape -> error(shape, String.format("Removed %s `%s`", shape.getType(), shape.getId())))
+                .map(shape -> isInconsequentialType(shape)
+                        ? ValidationEvent.builder()
+                            .severity(Severity.WARNING)
+                            .message(String.format("Removed %s `%s`", shape.getType(), shape.getId()))
+                            .shapeId(shape.getId())
+                            .id(getEventId() + ".ScalarShape")
+                            .sourceLocation(shape.getSourceLocation())
+                            .build()
+                        : error(shape, String.format("Removed %s `%s`", shape.getType(), shape.getId())))
                 .collect(Collectors.toList());
     }
 
@@ -40,4 +52,20 @@ public final class RemovedShape extends AbstractDiffEvaluator {
                 .filter(member -> !differences.getNewModel().getShapeIds().contains(member.getContainer()))
                 .isPresent();
     }
+
+    private boolean isInconsequentialType(Shape shape) {
+        ShapeType shapeType = shape.getType();
+        return shapeType == ShapeType.BIG_DECIMAL
+                || shapeType == ShapeType.BIG_INTEGER
+                || shapeType == ShapeType.BLOB
+                || shapeType == ShapeType.BOOLEAN
+                || shapeType == ShapeType.BYTE
+                || shapeType == ShapeType.DOUBLE
+                || shapeType == ShapeType.FLOAT
+                || shapeType == ShapeType.SHORT
+                || shapeType == ShapeType.TIMESTAMP
+                || shapeType == ShapeType.LONG
+                || ((shapeType == ShapeType.STRING) && !shape.hasTrait(EnumTrait.class))
+                || (shapeType == ShapeType.INTEGER);
+    }
 }
diff --git a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/SortUseStatements.java b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/SortUseStatements.java
index b9e29bc06..ac7ea465c 100644
--- a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/SortUseStatements.java
+++ b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/SortUseStatements.java
@@ -172,22 +172,16 @@ final class SortUseStatements implements Function<TokenTree, TokenTree> {
 
             if (lineComment != null) {
                 br.appendChild(sp);
-                // Modify the position of the comment so that the formatter knows it's a trailing line comment.
-                CapturedToken updatedLineComment = lineComment.getTree()
-                        .tokens()
-                        .findFirst()
-                        .get()
-                        .toBuilder()
-                        // Set the line and column to zero so the formatter knows it's on the same line as the
-                        // statement, making it an end-of-line comment.
-                        .startLine(0)
-                        .endLine(0)
-                        .build();
-                TokenTree commentTree = TokenTree.of(TreeType.COMMENT);
-                commentTree.appendChild(TokenTree.of(updatedLineComment));
-                br.appendChild(commentTree);
+                br.appendChild(lineComment.getTree());
             } else {
-                br.appendChild(TokenTree.of(CapturedToken.builder().token(IdlToken.NEWLINE).lexeme("\n").build()));
+                br.appendChild(TokenTree.of(CapturedToken.builder()
+                        .token(IdlToken.NEWLINE)
+                        .lexeme("\n")
+                        // Set the start line to be the start line of the use statement
+                        // so the formatter  knows it's a line comment.
+                        .startLine(id.getStartLine())
+                        .endLine(id.getStartLine() + 1)
+                        .build()));
             }
 
             TokenTree ws = TokenTree.of(TreeType.WS);
diff --git a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeLeaf.java b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeLeaf.java
index 562e05f94..68a66c0e4 100644
--- a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeLeaf.java
+++ b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeLeaf.java
@@ -79,7 +79,9 @@ final class TokenTreeLeaf implements TokenTree {
         if (token.getErrorMessage() != null) {
             return token.getIdlToken() + "(" + token.getErrorMessage() + ')';
         } else {
-            return token.getIdlToken().getDebug(token.getLexeme());
+            return token.getIdlToken().getDebug(token.getLexeme())
+                    + " (" + getStartLine() + ", " + getStartColumn() + ")"
+                    + " - (" + getEndLine() + ", " + getEndColumn() + ")";
         }
     }
 
diff --git a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeNode.java b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeNode.java
index 3b1b2466c..ba13b52f1 100644
--- a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeNode.java
+++ b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTreeNode.java
@@ -109,22 +109,50 @@ class TokenTreeNode implements TokenTree {
 
     @Override
     public final int getStartLine() {
-        return getChildren().isEmpty() ? 0 : getChildren().get(0).getStartLine();
+        // Start line of 0 indicates an empty child, so ignore it.
+        for (TokenTree child : getChildren()) {
+            int startLine = child.getStartLine();
+            if (startLine > 0) {
+                return startLine;
+            }
+        }
+        return 0;
     }
 
     @Override
     public final int getStartColumn() {
-        return getChildren().isEmpty() ? 0 : getChildren().get(0).getStartColumn();
+        // Start column of 0 indicates an empty child, so ignore it.
+        for (TokenTree child : getChildren()) {
+            int startColumn = child.getStartColumn();
+            if (startColumn > 0) {
+                return startColumn;
+            }
+        }
+        return 0;
     }
 
     @Override
     public final int getEndLine() {
-        return children.isEmpty() ? getStartLine() : children.get(children.size() - 1).getEndLine();
+        // End line of 0 indicates an empty child, so ignore it.
+        for (int i = children.size() - 1; i >= 0; i--) {
+            int endLine = children.get(i).getEndLine();
+            if (endLine > 0) {
+                return endLine;
+            }
+        }
+        return 0;
     }
 
     @Override
     public final int getEndColumn() {
-        return children.isEmpty() ? getStartColumn() : children.get(children.size() - 1).getEndColumn();
+        // End column of 0 indicates an empty child, so ignore it.
+        for (int i = children.size() - 1; i >= 0; i--) {
+            int endColumn = children.get(i).getEndColumn();
+            if (endColumn > 0) {
+                return endColumn;
+            }
+        }
+        return 0;
     }
 
     @Override


Test patch:
diff --git a/smithy-aws-protocol-tests/model/awsJson1_1/main.smithy b/smithy-aws-protocol-tests/model/awsJson1_1/main.smithy
index 19e660cc3..380b2ab8a 100644
--- a/smithy-aws-protocol-tests/model/awsJson1_1/main.smithy
+++ b/smithy-aws-protocol-tests/model/awsJson1_1/main.smithy
@@ -14,7 +14,7 @@ use smithy.test#httpResponseTests
     cloudFormationName: "JsonProtocol",
     cloudTrailEventSource: "jsonprotocol.amazonaws.com",
 )
-@sigv4(name: "foo")
+@sigv4(name: "jsonprotocol")
 @awsJson1_1
 @title("Sample Json 1.1 Protocol Service")
 service JsonProtocol {
diff --git a/smithy-diff/src/test/java/software/amazon/smithy/diff/ModelDiffTest.java b/smithy-diff/src/test/java/software/amazon/smithy/diff/ModelDiffTest.java
index cbc22a416..307dacbe9 100644
--- a/smithy-diff/src/test/java/software/amazon/smithy/diff/ModelDiffTest.java
+++ b/smithy-diff/src/test/java/software/amazon/smithy/diff/ModelDiffTest.java
@@ -11,6 +11,7 @@ import org.junit.jupiter.api.Test;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.StringShape;
+import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.traits.SensitiveTrait;
 import software.amazon.smithy.model.validation.Severity;
 import software.amazon.smithy.model.validation.ValidatedResult;
@@ -77,7 +78,7 @@ public class ModelDiffTest {
     @Test
     public void findsBreakingChanges() {
         Model oldModel = Model.builder()
-                .addShape(StringShape.builder().id("smithy.example#Str").build())
+                .addShape(StructureShape.builder().id("smithy.example#Str").build())
                 .build();
         Model newModel = Model.builder().build();
         ModelDiff.Result result = ModelDiff.builder().oldModel(oldModel).newModel(newModel).compare();
diff --git a/smithy-diff/src/test/java/software/amazon/smithy/diff/evaluators/RemovedShapeTest.java b/smithy-diff/src/test/java/software/amazon/smithy/diff/evaluators/RemovedShapeTest.java
index 7cae1ff63..04b187e6c 100644
--- a/smithy-diff/src/test/java/software/amazon/smithy/diff/evaluators/RemovedShapeTest.java
+++ b/smithy-diff/src/test/java/software/amazon/smithy/diff/evaluators/RemovedShapeTest.java
@@ -23,10 +23,26 @@ import java.util.List;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.diff.ModelDiff;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.BigDecimalShape;
+import software.amazon.smithy.model.shapes.BigIntegerShape;
+import software.amazon.smithy.model.shapes.BlobShape;
+import software.amazon.smithy.model.shapes.BooleanShape;
+import software.amazon.smithy.model.shapes.ByteShape;
+import software.amazon.smithy.model.shapes.DoubleShape;
+import software.amazon.smithy.model.shapes.EnumShape;
+import software.amazon.smithy.model.shapes.FloatShape;
+import software.amazon.smithy.model.shapes.IntEnumShape;
+import software.amazon.smithy.model.shapes.IntegerShape;
 import software.amazon.smithy.model.shapes.ListShape;
+import software.amazon.smithy.model.shapes.LongShape;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShortShape;
 import software.amazon.smithy.model.shapes.StringShape;
+import software.amazon.smithy.model.shapes.StructureShape;
+import software.amazon.smithy.model.shapes.TimestampShape;
+import software.amazon.smithy.model.traits.EnumDefinition;
+import software.amazon.smithy.model.traits.EnumTrait;
 import software.amazon.smithy.model.traits.PrivateTrait;
 import software.amazon.smithy.model.validation.Severity;
 import software.amazon.smithy.model.validation.ValidationEvent;
@@ -34,8 +50,8 @@ import software.amazon.smithy.model.validation.ValidationEvent;
 public class RemovedShapeTest {
     @Test
     public void detectsShapeRemoval() {
-        Shape shapeA1 = StringShape.builder().id("foo.baz#Baz").build();
-        Shape shapeB1 = StringShape.builder().id("foo.baz#Bam").build();
+        Shape shapeA1 = StructureShape.builder().id("foo.baz#Baz").build();
+        Shape shapeB1 = StructureShape.builder().id("foo.baz#Bam").build();
         Model modelA = Model.assembler().addShapes(shapeA1, shapeB1).assemble().unwrap();
         Model modelB = Model.assembler().addShapes(shapeB1).assemble().unwrap();
         List<ValidationEvent> events = ModelDiff.compare(modelA, modelB);
@@ -45,6 +61,61 @@ public class RemovedShapeTest {
         assertThat(TestHelper.findEvents(events, Severity.ERROR).size(), equalTo(1));
     }
 
+    @Test
+    public void emitsWarningsForScalarShapes() {
+        Shape[] scalarShapes = new Shape[] {
+                IntegerShape.builder().id("foo.baz#BazOne").build(),
+                BigDecimalShape.builder().id("foo.baz#BazTwo").build(),
+                BigIntegerShape.builder().id("foo.baz#BazThree").build(),
+                BlobShape.builder().id("foo.baz#BazFour").build(),
+                BooleanShape.builder().id("foo.baz#BazFive").build(),
+                ByteShape.builder().id("foo.baz#BazSix").build(),
+                DoubleShape.builder().id("foo.baz#BazSeven").build(),
+                FloatShape.builder().id("foo.baz#BazEight").build(),
+                ShortShape.builder().id("foo.baz#BazNine").build(),
+                TimestampShape.builder().id("foo.baz#BazTen").build(),
+                LongShape.builder().id("foo.baz#BazEleven").build(),
+                StringShape.builder().id("foo.baz#BazTwelve").build()
+        };
+        Model modelA = Model.assembler().addShapes(scalarShapes).assemble().unwrap();
+        Model modelB = Model.assembler().assemble().unwrap();
+        List<ValidationEvent> events = ModelDiff.compare(modelA, modelB);
+
+        assertThat(TestHelper.findEvents(events, "RemovedShape.ScalarShape").size(), equalTo(12));
+        assertThat("Scalar removals should be WARNING severity",
+                events.stream().allMatch(event -> Severity.WARNING.equals(event.getSeverity())));
+    }
+
+    @Test
+    public void emitsErrorForEnumString() {
+        Shape shapeA1 = StringShape.builder()
+                .id("foo.baz#Baz")
+                .addTrait(EnumTrait.builder().addEnum(EnumDefinition.builder().value("val").build()).build())
+                .build();
+        Model modelA = Model.assembler().addShapes(shapeA1).assemble().unwrap();
+        Model modelB = Model.assembler().addShapes().assemble().unwrap();
+        List<ValidationEvent> events = ModelDiff.compare(modelA, modelB);
+
+        assertThat(TestHelper.findEvents(events, "RemovedShape").size(), equalTo(1));
+        assertThat(TestHelper.findEvents(events, shapeA1.getId()).size(), equalTo(1));
+        assertThat(TestHelper.findEvents(events, Severity.ERROR).size(), equalTo(1));
+    }
+
+    @Test
+    public void emitsErrorForIntEnum() {
+        Shape shapeA1 = IntEnumShape.builder()
+                .id("foo.baz#Baz")
+                .addMember("FOO", 1)
+                .build();
+        Model modelA = Model.assembler().addShapes(shapeA1).assemble().unwrap();
+        Model modelB = Model.assembler().addShapes().assemble().unwrap();
+        List<ValidationEvent> events = ModelDiff.compare(modelA, modelB);
+
+        assertThat(TestHelper.findEvents(events, "RemovedShape").size(), equalTo(1));
+        assertThat(TestHelper.findEvents(events, shapeA1.getId()).size(), equalTo(1));
+        assertThat(TestHelper.findEvents(events, Severity.ERROR).size(), equalTo(1));
+    }
+
     @Test
     public void doesNotEmitForPrivateShapes() {
         Shape shape = StringShape.builder().id("foo.baz#Baz").addTrait(new PrivateTrait()).build();
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/validation/testrunner/SmithyTestSuite.java b/smithy-model/src/main/java/software/amazon/smithy/model/validation/testrunner/SmithyTestSuite.java
index 544cc5409..943c4d4e1 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/validation/testrunner/SmithyTestSuite.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/validation/testrunner/SmithyTestSuite.java
@@ -66,13 +66,13 @@ public final class SmithyTestSuite {
      * of the given {@code contextClass}, and that model discovery should be
      * used using the given {@code contextClass}.
      *
-     * <p>Each returns {@code Object[]} contains the filename of the test as
+     * <p>Each returned {@code Object[]} contains the filename of the test as
      * the first argument, followed by a {@code Callable<SmithyTestCase.Result>}
      * as the second argument. All a parameterized test needs to do is call
      * {@code call} on the provided {@code Callable} to execute the test and
      * fail if the test case is invalid.
      *
-     * <p>For example, the following can used as a unit test:
+     * <p>For example, the following can be used as a unit test:
      *
      * <pre>{@code
      * import java.util.concurrent.Callable;
@@ -107,7 +107,49 @@ public final class SmithyTestSuite {
                 .parameterizedTestSource();
     }
 
-    private Stream<Object[]> parameterizedTestSource() {
+    /**
+     * Factory method used to create a JUnit 5 {@code ParameterizedTest}
+     * {@code MethodSource}.
+     *
+     * <p>Test cases need to be added to the test suite before calling this,
+     * for example by using {@link #addTestCasesFromDirectory(Path)}.
+     *
+     * <p>Each returned {@code Object[]} contains the filename of the test as
+     * the first argument, followed by a {@code Callable<SmithyTestCase.Result>}
+     * as the second argument. All a parameterized test needs to do is call
+     * {@code call} on the provided {@code Callable} to execute the test and
+     * fail if the test case is invalid.
+     *
+     * <p>For example, the following can be used as a unit test:
+     *
+     * <pre>{@code
+     * import java.util.concurrent.Callable;
+     * import java.util.stream.Stream;
+     * import org.junit.jupiter.params.ParameterizedTest;
+     * import org.junit.jupiter.params.provider.MethodSource;
+     * import software.amazon.smithy.model.validation.testrunner.SmithyTestCase;
+     * import software.amazon.smithy.model.validation.testrunner.SmithyTestSuite;
+     *
+     * public class TestRunnerTest {
+     *     \@ParameterizedTest(name = "\{0\}")
+     *     \@MethodSource("source")
+     *     public void testRunner(String filename, Callable&lt;SmithyTestCase.Result&gt; callable) throws Exception {
+     *         callable.call();
+     *     }
+     *
+     *     public static Stream&lt;?&gt; source() {
+     *         ModelAssembler assembler = Model.assembler(TestRunnerTest.class.getClassLoader());
+     *         return SmithyTestSuite.runner()
+     *                 .setModelAssemblerFactory(assembler::copy)
+     *                 .addTestCasesFromUrl(TestRunnerTest.class.getResource("errorfiles"))
+     *                 .parameterizedTestSource();
+     *     }
+     * }
+     * }</pre>
+     *
+     * @return Returns the Stream that should be used as a JUnit 5 {@code MethodSource} return value.
+     */
+    public Stream<Object[]> parameterizedTestSource() {
         return cases.stream().map(testCase -> {
             Callable<SmithyTestCase.Result> callable = createTestCaseCallable(testCase);
             Callable<SmithyTestCase.Result> wrappedCallable = () -> callable.call().unwrap();
diff --git a/smithy-openapi/src/test/java/software/amazon/smithy/openapi/fromsmithy/protocols/AwsRestJson1ProtocolTest.java b/smithy-openapi/src/test/java/software/amazon/smithy/openapi/fromsmithy/protocols/AwsRestJson1ProtocolTest.java
index 4b588f984..35209a359 100644
--- a/smithy-openapi/src/test/java/software/amazon/smithy/openapi/fromsmithy/protocols/AwsRestJson1ProtocolTest.java
+++ b/smithy-openapi/src/test/java/software/amazon/smithy/openapi/fromsmithy/protocols/AwsRestJson1ProtocolTest.java
@@ -215,6 +215,7 @@ public class AwsRestJson1ProtocolTest {
 
     @ParameterizedTest
     @MethodSource("protocolHeaderCases")
+    @SuppressWarnings("unchecked")
     public void assertProtocolHeaders(
             String operationId,
             Set<String> expectedRequestHeaders,
diff --git a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java
index d16ebed91..e8832fda9 100644
--- a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java
+++ b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java
@@ -85,4 +85,54 @@ public class TokenTreeTest {
         assertThat(tree.replaceChild(child1, child2), is(false));
         assertThat(tree.zipper().getFirstChild(TreeType.COMMA).getTree(), is(child2));
     }
+
+    @Test
+    public void emptyTreeLocationsDontBubbleUpToParent() {
+        // If a tree's first/last child is empty, meaning it has a location of
+        // (0, 0) - (0, 0), that should not be considered the starting/ending
+        // location of the tree.
+        IdlTokenizer shapeStatementTokenizer = IdlTokenizer.create("structure Foo {}");
+        TokenTree shapeStatement = TokenTree.of(shapeStatementTokenizer, TreeType.SHAPE_STATEMENT);
+        assertThat(shapeStatement.getStartLine(), equalTo(1));
+        assertThat(shapeStatement.getStartColumn(), equalTo(1));
+        assertThat(shapeStatement.getEndLine(), equalTo(1));
+        assertThat(shapeStatement.getEndColumn(), equalTo(17));
+
+        // There are no traits, so this will be an empty tree with location
+        // (0, 0) - (0, 0).
+        TokenTree traitStatements = shapeStatement.getChildren().get(0);
+        assertThat(traitStatements.getStartLine(), equalTo(0));
+        assertThat(traitStatements.getStartColumn(), equalTo(0));
+        assertThat(traitStatements.getEndLine(), equalTo(0));
+        assertThat(traitStatements.getEndColumn(), equalTo(0));
+
+        // The SHAPE_STATEMENT location should come from the SHAPE child.
+        TokenTree shape = shapeStatement.getChildren().get(1);
+        assertThat(shape.getStartLine(), equalTo(shapeStatement.getStartLine()));
+        assertThat(shape.getStartColumn(), equalTo(shapeStatement.getStartColumn()));
+        assertThat(shape.getEndLine(), equalTo(shapeStatement.getEndLine()));
+        assertThat(shape.getEndColumn(), equalTo(shapeStatement.getEndColumn()));
+
+        IdlTokenizer shapeSectionTokenizer = IdlTokenizer.create("namespace com.foo\n");
+        TokenTree shapeSection = TokenTree.of(shapeSectionTokenizer, TreeType.SHAPE_SECTION);
+        assertThat(shapeSection.getStartLine(), equalTo(1));
+        assertThat(shapeSection.getStartColumn(), equalTo(1));
+        assertThat(shapeSection.getEndLine(), equalTo(2));
+        assertThat(shapeSection.getEndColumn(), equalTo(1));
+
+        // There are shape statements, so this will be an empty tree with location
+        // (0, 0) - (0, 0).
+        TokenTree shapeStatements = shapeSection.getChildren().get(2);
+        assertThat(shapeStatements.getStartLine(), equalTo(0));
+        assertThat(shapeStatements.getStartColumn(), equalTo(0));
+        assertThat(shapeStatements.getEndLine(), equalTo(0));
+        assertThat(shapeStatements.getEndColumn(), equalTo(0));
+
+        // The SHAPE_SECTION location should come from the NAMESPACE_STATEMENT child.
+        TokenTree namespaceStatement = shapeSection.getChildren().get(0);
+        assertThat(namespaceStatement.getStartLine(), equalTo(shapeSection.getStartLine()));
+        assertThat(namespaceStatement.getStartColumn(), equalTo(shapeSection.getStartColumn()));
+        assertThat(namespaceStatement.getEndLine(), equalTo(shapeSection.getEndLine()));
+        assertThat(namespaceStatement.getEndColumn(), equalTo(shapeSection.getEndColumn()));
+    }
 }
diff --git a/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/br-comments-with-multiple-use.formatted.smithy b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/br-comments-with-multiple-use.formatted.smithy
new file mode 100644
index 000000000..e9c3d2ee2
--- /dev/null
+++ b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/br-comments-with-multiple-use.formatted.smithy
@@ -0,0 +1,12 @@
+$version: "2.0"
+
+namespace smithy.example
+
+use smithy.api#length // def
+use smithy.api#sensitive // abc
+
+// Comment
+/// A
+@sensitive
+@length(min: 1)
+string A
diff --git a/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/br-comments-with-multiple-use.smithy b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/br-comments-with-multiple-use.smithy
new file mode 100644
index 000000000..1f3f37a88
--- /dev/null
+++ b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/br-comments-with-multiple-use.smithy
@@ -0,0 +1,12 @@
+$version: "2.0"
+
+namespace smithy.example
+
+use smithy.api#sensitive // abc
+use smithy.api#length // def
+
+// Comment
+/// A
+@sensitive
+@length(min: 1)
+string A