From 4b797ee95acb44d0f562e66a5096adb5c24e1089 Mon Sep 17 00:00:00 2001
From: Miles Ziemer <mziemer@amazon.com>
Date: Tue, 15 Aug 2023 11:43:30 -0400
Subject: [PATCH] Allow creating specific tree types

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.
---
 .../amazon/smithy/syntax/TokenTree.java       | 21 +++++++++++++++++++
 .../amazon/smithy/syntax/TokenTreeTest.java   | 12 +++++++++++
 .../amazon/smithy/syntax/TreeTypeTest.java    |  6 ++----
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java
index 2967bde648a..6eb84d00e1a 100644
--- a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/TokenTree.java
+++ b/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.
      *
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 569886180f5..d16ebed91dc 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
@@ -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");
diff --git a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java b/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java
index def758ddff9..3f5f87640fa 100644
--- a/smithy-syntax/src/test/java/software/amazon/smithy/syntax/TreeTypeTest.java
+++ b/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);
     }
 }