Allow creating specific tree types

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

@@ -44,6 +44,27 @@ static TokenTree of(IdlTokenizer tokenizer) {
         return capturingTokenizer.getRoot();
     }
 
+    /**
+     * Create a TokenTree of the given {@link TreeType} from an {@link IdlTokenizer}.
+     *
+     * <p>The following example shows how to create a {@link TokenTree} for a trait.</p>
+     *
+     * <pre>{@code
+     * IdlTokenizer tokenizer = IdlTokenizer.create("@someTrait");
+     * TokenTree traitTree = TokenTree.of(tokenizer, TreeType.TRAIT);
+     * }</pre>
+     *
+     * @param tokenizer Tokenizer to traverse.
+     * @param type Type of tree to create.
+     * @return Returns the created tree.
+     */
+    static TokenTree of(IdlTokenizer tokenizer, TreeType type) {
+        CapturingTokenizer capturingTokenizer = new CapturingTokenizer(tokenizer);
+        type.parse(capturingTokenizer);
+        // The root of the tree is always IDL with children appended, so the first child is the one we want.
+        return capturingTokenizer.getRoot().getChildren().get(0);
+    }
+
     /**
      * Create a leaf tree from a single token.
      *

smithy-syntax/src/test/java/software/amazon/smithy/syntax/TokenTreeTest.java

@@ -33,6 +33,18 @@ public void createsFromTokenizer() {
         assertThat(tree.getChildren().get(2).getType(), equalTo(TreeType.SHAPE_SECTION));
     }
 
+    @Test
+    public void createsFromTokenizerAndType() {
+        IdlTokenizer tokenizer = IdlTokenizer.create("@foo");
+        TokenTree tree = TokenTree.of(tokenizer, TreeType.TRAIT);
+
+        assertThat(tree.getType(), is(TreeType.TRAIT));
+        assertThat(tree.getChildren(), hasSize(2));
+        assertThat(tree.getError(), nullValue());
+        assertThat(tree.getChildren().get(0).getType(), equalTo(TreeType.TOKEN));
+        assertThat(tree.getChildren().get(1).getType(), equalTo(TreeType.SHAPE_ID));
+    }
+
     @Test
     public void createsFromCapturedToken() {
         IdlTokenizer tokenizer = IdlTokenizer.create("foo");

smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java

@@ -1860,9 +1860,7 @@ private static void assertTreeIsInvalid(TokenTree tree) {
     }
 
     private static TokenTree getTree(TreeType type, String forText) {
-        CapturingTokenizer tokenizer = new CapturingTokenizer(IdlTokenizer.create(forText));
-        type.parse(tokenizer);
-        // The root of the tree is always IDL with children appended, so the first child is the one we want.
-        return tokenizer.getRoot().getChildren().get(0);
+        IdlTokenizer tokenizer = IdlTokenizer.create(forText);
+        return TokenTree.of(tokenizer, type);
     }
 }