Handle map keys with idRef

milesziemer · Jan 19, 2024 · b8a7aad · b8a7aad
1 parent e536ea2
commit b8a7aad
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 25 deletions.
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/IdRefShapeReferences.java b/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/IdRefShapeReferences.java
@@ -80,11 +80,8 @@ private static NodeQuery buildNodeQuery(PathFinder.Path path) {
                 break;
             }
             switch (relationship.getRelationshipType()) {
-                case MEMBER_TARGET:
                 case MAP_KEY:
-                case TRAIT:
-                case MIXIN:
-                default:
+                    query.anyMemberName();
                     break;
                 case MAP_VALUE:
                     query.anyMember();
@@ -98,6 +95,9 @@ private static NodeQuery buildNodeQuery(PathFinder.Path path) {
                     MemberShape member = (MemberShape) relationship.getNeighborShape().get();
                     query.member(member.getMemberName());
                     break;
+                default:
+                    // Other relationship types don't produce meaningful edges to search the node.
+                    break;
             }
         }
         return query;

diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/NodeQuery.java b/smithy-model/src/main/java/software/amazon/smithy/model/neighbor/NodeQuery.java
@@ -58,6 +58,16 @@ NodeQuery anyElement() {
         return this;
     }
 
+    NodeQuery anyMemberName() {
+        queries.add(node -> {
+            if (node == null || !node.isObjectNode()) {
+                return Stream.empty();
+            }
+            return node.expectObjectNode().getMembers().keySet().stream();
+        });
+        return this;
+    }
+
     List<Node> execute(Node node) {
         if (queries.isEmpty()) {
             return ListUtils.of();
@@ -72,6 +82,6 @@ List<Node> execute(Node node) {
 
     @FunctionalInterface
     interface Query {
-        Stream<Node> run(Node node);
+        Stream<? extends Node> run(Node node);
     }
 }
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/neighbor/through-idref.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/neighbor/through-idref.smithy
@@ -19,13 +19,14 @@ operation GetFoo {
         eight: Eight
         nine: Nine
         ten: Ten
+        eleven: Eleven
     }
 }
 
 // --
 @trait
 structure withIdRefOnMember {
-    @idRef(failWhenMissing: true, selector: "*")
+    @idRef(failWhenMissing: true)
     ref: String
 }
 
@@ -40,7 +41,7 @@ structure withIdRefOnMemberTarget {
     ref: OnTarget
 }
 
-@idRef(failWhenMissing: true, selector: "*")
+@idRef(failWhenMissing: true)
 string OnTarget
 
 @withIdRefOnMemberTarget(ref: Ref2)
@@ -55,7 +56,7 @@ structure withIdRefOnNestedStructureMember {
 }
 
 structure Nested {
-    @idRef(failWhenMissing: true, selector: "*")
+    @idRef(failWhenMissing: true)
     member: String
 }
 
@@ -75,7 +76,7 @@ list withIdRefOnListMemberTarget {
     member: OnListMemberTarget
 }
 
-@idRef(failWhenMissing: true, selector: "*")
+@idRef(failWhenMissing: true)
 string OnListMemberTarget
 
 @withIdRefOnListMemberTarget([
@@ -87,7 +88,7 @@ structure Ref4 {}
 
 // --
 @trait
-@idRef(failWhenMissing: true, selector: "*")
+@idRef(failWhenMissing: true)
 string withIdRefOnSelf
 
 @withIdRefOnSelf(Ref5)
@@ -126,7 +127,7 @@ structure withIdRefThroughMixin with [ThroughMixin] {}
 
 @mixin
 structure ThroughMixin {
-    @idRef(failWhenMissing: true, selector: "*")
+    @idRef(failWhenMissing: true)
     ref: String
 }
 
@@ -142,7 +143,7 @@ map withIdRefOnMapValue {
     value: OnMap
 }
 
-@idRef(failWhenMissing: true, selector: "*")
+@idRef(failWhenMissing: true)
 string OnMap
 
 @withIdRefOnMapValue({
@@ -169,18 +170,15 @@ structure Ten {}
 structure Ref10 {}
 
 // --
-// NOTE: This actually doesn't work because map key of 'Ref11' it says isn't a valid
-// shape id.
-
-//@trait
-//map withIdRefOnMapKey {
-//    key: OnMap
-//    value: String
-//}
+@trait
+map withIdRefOnMapKey {
+    key: OnMap
+    value: String
+}
 
-//@withIdRefOnMapKey({
-//    Ref11: "foo"
-//})
-//structure Eleven {}
+@withIdRefOnMapKey({
+    "com.foo#Ref11": "foo"
+})
+structure Eleven {}
 
-//structure Ref11 {}
+structure Ref11 {}