Repository: smithy-lang/smithy (java)

Issue:
Referencing members in traits fails parsing by Smithy assembler
I am attempting to create a new trait which uses the `idRef` trait. The new trait lets members in structs be linked to members of other structs.  When I reference a struct member from the trait using the `@traitName(StructName$memberName)` syntax I get the following parsing error when loading the model using the Smithy assembler: `Expected RPAREN(')') but found DOLLAR('$')`.  Is this the expected behavior for referencing members within traits, or should it be possible?

This is a minimal example which demonstrates the behavior:

```smithy
$version: "2.0"

namespace example

@trait
@idRef(failWhenMissing: true, selector: ":test(structure) > member")
string link

structure StructA {
  @required
  id: String
}

structure StructB {
  @required
  @link(StructA$id) // invalid line
  structAId: String
}
```

This model appears to be valid when parsed by both the Smithy LSP and Intellij Smithy Plugin syntax highlighting.

I also noticed that I can replace `@link(StructA$id)` with a fully qualified string like `@link("example#StructA$Id")` and it will parse and assemble correctly.

Pull request description:
Fix trait parse error for shape IDs
closes #2021

*Issue #, if available:*

*Description of changes:*


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-model/src/main/java/software/amazon/smithy/model/loader/IdlNodeParser.java b/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlNodeParser.java
index 2aee9178a..fa35f6556 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlNodeParser.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlNodeParser.java
@@ -74,7 +74,7 @@ final class IdlNodeParser {
                 return result;
             case IDENTIFIER:
                 String shapeId = loader.internString(IdlShapeIdParser.expectAndSkipShapeId(tokenizer));
-                return parseIdentifier(loader, shapeId, location);
+                return createIdentifier(loader, shapeId, location);
             case NUMBER:
                 Number number = tokenizer.getCurrentTokenNumberValue();
                 tokenizer.next();
@@ -95,10 +95,10 @@ final class IdlNodeParser {
      * @param location   Source location to assign to the identifier.
      * @return Returns the parsed identifier.
      */
-    static Node parseIdentifier(IdlModelLoader loader, String identifier, SourceLocation location) {
+    static Node createIdentifier(IdlModelLoader loader, String identifier, SourceLocation location) {
         Keyword keyword = Keyword.from(identifier);
         return keyword == null
-               ? parseSyntacticShapeId(loader, identifier, location)
+               ? createSyntacticShapeId(loader, identifier, location)
                : keyword.createNode(location);
     }
 
@@ -138,7 +138,7 @@ final class IdlNodeParser {
         }
     }
 
-    private static Node parseSyntacticShapeId(
+    private static Node createSyntacticShapeId(
             IdlModelLoader loader,
             String identifier,
             SourceLocation location
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlTraitParser.java b/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlTraitParser.java
index 0d8006e2b..82b52de57 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlTraitParser.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlTraitParser.java
@@ -193,19 +193,30 @@ final class IdlTraitParser {
                 }
             case IDENTIFIER:
             default:
-                String identifier = loader.internString(tokenizer.getCurrentTokenLexeme());
-                tokenizer.next();
+                // Handle: `foo`, `foo$bar`, `foo.bar#baz`, `foo.bar#baz$bam`, `foo: bam`
+                String identifier = loader.internString(IdlShapeIdParser.expectAndSkipShapeId(tokenizer));
                 tokenizer.skipWsAndDocs();
-                if (tokenizer.getCurrentToken() == IdlToken.COLON) {
+                if (tokenizer.getCurrentToken() == IdlToken.RPAREN || isItDefinitelyShapeId(identifier)) {
+                    return IdlNodeParser.createIdentifier(loader, identifier, location);
+                } else {
+                    tokenizer.expect(IdlToken.COLON);
                     tokenizer.next();
                     tokenizer.skipWsAndDocs();
                     return parseStructuredTrait(loader, new StringNode(identifier, location));
-                } else {
-                    return IdlNodeParser.parseIdentifier(loader, identifier, location);
                 }
         }
     }
 
+    private static boolean isItDefinitelyShapeId(String identifier) {
+        for (int i = 0; i < identifier.length(); i++) {
+            char c = identifier.charAt(i);
+            if (c == '.' || c == '$' || c == '#') {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private static ObjectNode parseStructuredTrait(IdlModelLoader loader, StringNode firstKey) {
         IdlInternalTokenizer tokenizer = loader.getTokenizer();
         loader.increaseNestingLevel();


Test patch:
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/annotation-unclosed-object1.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/annotation-unclosed-object1.smithy
index 87564f759..74952506c 100644
--- a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/annotation-unclosed-object1.smithy
+++ b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/annotation-unclosed-object1.smithy
@@ -1,4 +1,4 @@
-// Syntax error at line 4, column 8: Expected RPAREN(')') but found IDENTIFIER('MyString') | Model
+// Syntax error at line 4, column 8: Expected COLON(':') but found IDENTIFIER('MyString') | Model
 namespace com.foo
 @foo(
 string MyString
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/invalid-trait-shape-id-key.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/invalid-trait-shape-id-key.smithy
index 6d3ac8643..83e4c9e55 100644
--- a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/invalid-trait-shape-id-key.smithy
+++ b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/invalid/traits/invalid-trait-shape-id-key.smithy
@@ -1,4 +1,4 @@
-// Syntax error at line 4, column 30: Expected RPAREN(')') but found DOT('.') | Model
+// Syntax error at line 4, column 42: Expected RPAREN(')') but found COLON(':') | Model
 namespace smithy.example
 
 @externalDocumentation(smithy.example#foo: "bar")
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/valid/traits/reference-shape-from-trait.json b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/valid/traits/reference-shape-from-trait.json
new file mode 100644
index 000000000..e49630c64
--- /dev/null
+++ b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/valid/traits/reference-shape-from-trait.json
@@ -0,0 +1,53 @@
+{
+    "smithy": "2.0",
+    "shapes": {
+        "smithy.example#StructA": {
+            "type": "structure",
+            "members": {
+                "id": {
+                    "target": "smithy.api#String"
+                }
+            }
+        },
+        "smithy.example#StructB": {
+            "type": "structure",
+            "members": {
+                "structAId": {
+                    "target": "smithy.api#String",
+                    "traits": {
+                        "smithy.example#link": "smithy.example#StructA$id"
+                    }
+                }
+            }
+        },
+        "smithy.example#StructC": {
+            "type": "structure",
+            "members": {
+                "structAId": {
+                    "target": "smithy.api#String",
+                    "traits": {
+                        "smithy.example#link": "smithy.example#StructA$id"
+                    }
+                }
+            }
+        },
+        "smithy.example#StructD": {
+            "type": "structure",
+            "members": {
+                "structAId": {
+                    "target": "smithy.api#String",
+                    "traits": {
+                        "smithy.example#link": "smithy.example#StructA"
+                    }
+                }
+            }
+        },
+        "smithy.example#link": {
+            "type": "string",
+            "traits": {
+                "smithy.api#idRef": {},
+                "smithy.api#trait": {}
+            }
+        }
+    }
+}
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/loader/valid/traits/reference-shape-from-trait.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/valid/traits/reference-shape-from-trait.smithy
new file mode 100644
index 000000000..a23e03dcf
--- /dev/null
+++ b/smithy-model/src/test/resources/software/amazon/smithy/model/loader/valid/traits/reference-shape-from-trait.smithy
@@ -0,0 +1,26 @@
+$version: "2.0"
+
+namespace smithy.example
+
+@trait
+@idRef
+string link
+
+structure StructA {
+    id: String
+}
+
+structure StructB {
+    @link(StructA$id)
+    structAId: String
+}
+
+structure StructC {
+    @link(smithy.example#StructA$id)
+    structAId: String
+}
+
+structure StructD {
+    @link(smithy.example#StructA)
+    structAId: String
+}