Fix forward reference resolution of service shapes

smithy-model/src/main/java/software/amazon/smithy/model/loader/IdlModelParser.java

@@ -22,11 +22,10 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import java.util.StringJoiner;
+import java.util.function.Consumer;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.ArrayNode;
 import software.amazon.smithy.model.node.Node;
@@ -546,17 +545,9 @@ private void parseOperationStatement(ShapeId id, SourceLocation location) {
         ObjectNode node = IdlNodeParser.parseObjectNode(this);
         LoaderUtils.checkForAdditionalProperties(node, id, OPERATION_PROPERTY_NAMES, modelFile.events());
         modelFile.onShape(builder);
-        node.getStringMember("input").ifPresent(input -> {
-            modelFile.addForwardReference(input.getValue(), builder::input);
-        });
-        node.getStringMember("output").ifPresent(output -> {
-            modelFile.addForwardReference(output.getValue(), builder::output);
-        });
-        node.getArrayMember("errors").ifPresent(errors -> {
-            for (StringNode value : errors.getElementsAs(StringNode.class)) {
-                modelFile.addForwardReference(value.getValue(), builder::addError);
-            }
-        });
+        optionalId(node, "input", builder::input);
+        optionalId(node, "output", builder::output);
+        optionalIdList(node, "errors", builder::addError);
     }
 
     private void parseServiceStatement(ShapeId id, SourceLocation location) {
@@ -566,21 +557,23 @@ private void parseServiceStatement(ShapeId id, SourceLocation location) {
         LoaderUtils.checkForAdditionalProperties(shapeNode, id, SERVICE_PROPERTY_NAMES, modelFile.events());
         builder.version(shapeNode.expectStringMember(VERSION_KEY).getValue());
         modelFile.onShape(builder);
-        optionalIdList(shapeNode, id.getNamespace(), OPERATIONS_KEY).forEach(builder::addOperation);
-        optionalIdList(shapeNode, id.getNamespace(), RESOURCES_KEY).forEach(builder::addResource);
+        optionalIdList(shapeNode, OPERATIONS_KEY, builder::addOperation);
+        optionalIdList(shapeNode, RESOURCES_KEY, builder::addResource);
     }
 
-    private static Optional<ShapeId> optionalId(ObjectNode node, String namespace, String name) {
-        return node.getStringMember(name).map(stringNode -> stringNode.expectShapeId(namespace));
+    private void optionalId(ObjectNode node, String name, Consumer<ShapeId> consumer) {
+        if (node.getMember(name).isPresent()) {
+            modelFile.addForwardReference(node.expectStringMember(name).getValue(), consumer);
+        }
     }
 
-    private static List<ShapeId> optionalIdList(ObjectNode node, String namespace, String name) {
-        return node.getArrayMember(name)
-                .map(array -> array.getElements().stream()
-                        .map(Node::expectStringNode)
-                        .map(s -> s.expectShapeId(namespace))
-                        .collect(Collectors.toList()))
-                .orElseGet(Collections::emptyList);
+    private void optionalIdList(ObjectNode node, String name, Consumer<ShapeId> consumer) {
+        if (node.getMember(name).isPresent()) {
+            ArrayNode value = node.expectArrayMember(name);
+            for (StringNode element : value.getElementsAs(StringNode.class)) {
+                modelFile.addForwardReference(element.getValue(), consumer);
+            }
+        }
     }
 
     private void parseResourceStatement(ShapeId id, SourceLocation location) {
@@ -590,16 +583,15 @@ private void parseResourceStatement(ShapeId id, SourceLocation location) {
         ObjectNode shapeNode = IdlNodeParser.parseObjectNode(this);
 
         LoaderUtils.checkForAdditionalProperties(shapeNode, id, RESOURCE_PROPERTY_NAMES, modelFile.events());
-        optionalId(shapeNode, id.getNamespace(), PUT_KEY).ifPresent(builder::put);
-        optionalId(shapeNode, id.getNamespace(), CREATE_KEY).ifPresent(builder::create);
-        optionalId(shapeNode, id.getNamespace(), READ_KEY).ifPresent(builder::read);
-        optionalId(shapeNode, id.getNamespace(), UPDATE_KEY).ifPresent(builder::update);
-        optionalId(shapeNode, id.getNamespace(), DELETE_KEY).ifPresent(builder::delete);
-        optionalId(shapeNode, id.getNamespace(), LIST_KEY).ifPresent(builder::list);
-        optionalIdList(shapeNode, id.getNamespace(), OPERATIONS_KEY).forEach(builder::addOperation);
-        optionalIdList(shapeNode, id.getNamespace(), RESOURCES_KEY).forEach(builder::addResource);
-        optionalIdList(shapeNode, id.getNamespace(), COLLECTION_OPERATIONS_KEY)
-                .forEach(builder::addCollectionOperation);
+        optionalId(shapeNode, PUT_KEY, builder::put);
+        optionalId(shapeNode, CREATE_KEY, builder::create);
+        optionalId(shapeNode, READ_KEY, builder::read);
+        optionalId(shapeNode, UPDATE_KEY, builder::update);
+        optionalId(shapeNode, DELETE_KEY, builder::delete);
+        optionalId(shapeNode, LIST_KEY, builder::list);
+        optionalIdList(shapeNode, OPERATIONS_KEY, builder::addOperation);
+        optionalIdList(shapeNode, RESOURCES_KEY, builder::addResource);
+        optionalIdList(shapeNode, COLLECTION_OPERATIONS_KEY, builder::addCollectionOperation);
 
         // Load identifiers and resolve forward references.
         shapeNode.getObjectMember(IDENTIFIERS_KEY).ifPresent(ids -> {

smithy-model/src/test/java/software/amazon/smithy/model/loader/IdlModelLoaderTest.java

@@ -22,12 +22,14 @@
 import static org.hamcrest.Matchers.not;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ResourceShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
@@ -37,6 +39,7 @@
 import software.amazon.smithy.model.validation.ValidatedResult;
 import software.amazon.smithy.model.validation.ValidatedResultException;
 import software.amazon.smithy.model.validation.ValidationEvent;
+import software.amazon.smithy.utils.ListUtils;
 
 public class IdlModelLoaderTest {
     @Test
@@ -154,4 +157,29 @@ public void warnsWhenInvalidSyntacticShapeIdIsFound() {
                            .count(),
                    equalTo(1L));
     }
+
+    @Test
+    public void properlyLoadsOperationsWithUseStatements() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("use-operations/service.smithy"))
+                .addImport(getClass().getResource("use-operations/nested.smithy"))
+                .addImport(getClass().getResource("use-operations/other.smithy"))
+                .assemble()
+                .unwrap();
+
+        // Spot check for a specific "use" shape.
+        assertThat(model.expectShape(ShapeId.from("smithy.example#Local"), OperationShape.class).getInput(),
+                   equalTo(Optional.of(ShapeId.from("smithy.example.nested#A"))));
+
+        assertThat(model.expectShape(ShapeId.from("smithy.example#Local"), OperationShape.class).getErrors(),
+                   equalTo(ListUtils.of(ShapeId.from("smithy.example.nested#C"))));
+
+        Map<String, ShapeId> identifiers = model.expectShape(
+                ShapeId.from("smithy.example.nested#Resource"),
+                ResourceShape.class
+        ).getIdentifiers();
+
+        assertThat(identifiers.get("s"), equalTo(ShapeId.from("smithy.api#String")));
+        assertThat(identifiers.get("x"), equalTo(ShapeId.from("smithy.example.other#X")));
+    }
 }

smithy-model/src/test/resources/software/amazon/smithy/model/loader/use-operations/nested.smithy

@@ -0,0 +1,28 @@
+namespace smithy.example.nested
+
+use smithy.example.other#Hello2
+use smithy.example.other#X
+use smithy.example.other#GetHello2
+
+operation Hello {}
+
+resource Resource {
+    // A "use" identifier.
+    identifiers: {
+        x: X,
+        s: String,
+    },
+
+    // A "use" lifecycle operation.
+    read: GetHello2,
+
+    // A "use" instance operation.
+    operations: [Hello2]
+}
+
+structure A {}
+
+structure B {}
+
+@error("client")
+structure C {}

smithy-model/src/test/resources/software/amazon/smithy/model/loader/use-operations/other.smithy

@@ -0,0 +1,20 @@
+namespace smithy.example.other
+
+operation Hello2 {
+    input: Hello2Input,
+}
+
+structure Hello2Input {
+    @required
+    x: X,
+
+    @required
+    s: String, // this should resolve to smithy.api#String
+}
+
+string X
+
+@readonly
+operation GetHello2 {
+    input: Hello2Input,
+}

smithy-model/src/test/resources/software/amazon/smithy/model/loader/use-operations/service.smithy

@@ -0,0 +1,28 @@
+namespace smithy.example
+
+use smithy.example.nested#Hello
+use smithy.example.nested#A
+use smithy.example.nested#B
+use smithy.example.nested#C
+use smithy.example.nested#Resource
+
+service Foo {
+    version: "2020-06-11",
+
+    // A "use" resource.
+    resources: [Resource],
+
+    // "use" instance operations.
+    operations: [Hello, Local]
+}
+
+operation Local {
+    // "use" input
+    input: A,
+
+    // "use" output
+    output: B,
+
+    // "use" errors
+    errors: [C]
+}