Allow creating specific tree types (#1925)

Adds method to TokenTree that allows creating a tree with a specific
TreeType. Previously, there was no way to do this since the parse
methods on TreeType are package private. This addition supports
use cases where new trees need to be added to a root tree.

A test was added for this method, and other tests were updated to
use it.

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);
     }
 }