Support specifying service at operation name (#82)

core/src/main/scala/playground/CommandResultReporter.scala

@@ -62,13 +62,13 @@ object CommandResultReporter {
           requestCounter.updateAndGet(_ + 1).flatTap { requestId =>
             Feedback[F]
               .logOutput(
-                s"// Calling ${parsed.operationName.text} ($requestId)"
+                s"// Calling ${parsed.operationName.operationName.text} ($requestId)"
               )
           }
 
       def onQuerySuccess(parsed: Query[Id], requestId: RequestId, out: InputNode[Id]): F[Unit] =
         Feedback[F].logOutput(
-          s"// Succeeded ${parsed.operationName.text} ($requestId), response:\n"
+          s"// Succeeded ${parsed.operationName.operationName.text} ($requestId), response:\n"
             + writeOutput(out)
         )
 

core/src/main/scala/playground/CompletionProvider.scala

@@ -46,13 +46,14 @@ object CompletionProvider {
         .foldMap(e => Map(OperationName[Id](e.name) -> NonEmptyList.one(serviceId)))
     }
 
-    val completeOperationName = servicesById
+    val completeOperationName
+      : Map[QualifiedIdentifier, List[QualifiedIdentifier] => List[CompletionItem]] = servicesById
       .map { case (serviceId, service) =>
-        serviceId -> { (useClauseIdent: Option[QualifiedIdentifier]) =>
+        serviceId -> { (presentServiceIdentifiers: List[QualifiedIdentifier]) =>
           val needsUseClause =
             MultiServiceResolver
               .resolveService(
-                useClauseIdent,
+                presentServiceIdentifiers,
                 servicesById,
               )
               .isLeft
@@ -76,7 +77,7 @@ object CompletionProvider {
         }
       }
 
-    val completeAnyOperationName = completeOperationName.toList.map(_._2).flatSequence.apply(None)
+    val completeAnyOperationName = completeOperationName.toList.map(_._2).flatSequence.apply(Nil)
 
     val completionsByEndpoint
       : Map[QualifiedIdentifier, Map[OperationName[Id], CompletionResolver[Any]]] = servicesById
@@ -97,7 +98,7 @@ object CompletionProvider {
       serviceId match {
         case Some(serviceId) =>
           completeOperationName(serviceId)(
-            q.useClause.value.map(_.identifier.value)
+            q.mapK(WithSource.unwrap).collectServiceIdentifiers
           )
         case _ => completeAnyOperationName
       }
@@ -120,7 +121,7 @@ object CompletionProvider {
           val serviceIdOpt =
             MultiServiceResolver
               .resolveService(
-                q.useClause.value.map(_.identifier.value),
+                q.mapK(WithSource.unwrap).collectServiceIdentifiers,
                 serviceIdsById,
               )
               .toOption
@@ -141,7 +142,9 @@ object CompletionProvider {
               case NodeContext.PathEntry.AtOperationInput ^^: ctx =>
                 serviceIdOpt match {
                   case Some(serviceId) =>
-                    completionsByEndpoint(serviceId)(q.operationName.value.mapK(WithSource.unwrap))
+                    completionsByEndpoint(serviceId)(
+                      q.operationName.value.operationName.value.mapK(WithSource.unwrap)
+                    )
                       .getCompletions(ctx)
 
                   case None => Nil

core/src/main/scala/playground/DiagnosticProvider.scala

@@ -144,6 +144,7 @@ object DiagnosticProvider {
           range,
           DiagnosticSeverity.Information,
           tags = Set.empty,
+          relatedInfo = Nil,
         )
 
     }

core/src/main/scala/playground/DocumentSymbolProvider.scala

@@ -16,7 +16,7 @@ object DocumentSymbolProvider {
       case Left(_) => Nil
       case Right(q) =>
         findInUseClause(q.useClause) ++
-          findInOperation(q.operationName, q.input)
+          findInOperation(q.operationName.value.operationName, q.input)
     }
 
   private def findInUseClause(

core/src/main/scala/playground/MultiServiceResolver.scala

@@ -9,10 +9,23 @@ import playground.smithyql.WithSource
 object MultiServiceResolver {
 
   def resolveService[A](
-    useClauseIdentifier: Option[QualifiedIdentifier],
+    identifiers: List[QualifiedIdentifier],
     services: Map[QualifiedIdentifier, A],
   ): Either[ResolutionFailure, A] =
-    useClauseIdentifier match {
+    identifiers match {
+      case Nil         => resolveFromOne(None, services)
+      case body :: Nil => resolveFromOne(Some(body) /* once told me */, services)
+      case more =>
+        ResolutionFailure
+          .ConflictingServiceReference(more)
+          .asLeft
+    }
+
+  private def resolveFromOne[A](
+    ident: Option[QualifiedIdentifier],
+    services: Map[QualifiedIdentifier, A],
+  ): Either[ResolutionFailure, A] =
+    ident match {
       case None if services.sizeIs == 1 => services.head._2.asRight
       case None => ResolutionFailure.AmbiguousService(services.keySet.toList).asLeft
 
@@ -32,16 +45,45 @@ object ResolutionFailure {
   final case class AmbiguousService(knownServices: List[QualifiedIdentifier])
     extends ResolutionFailure
 
+  final case class ConflictingServiceReference(references: List[QualifiedIdentifier])
+    extends ResolutionFailure
+
   final case class UnknownService(
     unknownId: QualifiedIdentifier,
     knownServices: List[QualifiedIdentifier],
   ) extends ResolutionFailure
 
+  def toCompilationError(rf: ResolutionFailure, q: Query[WithSource]): CompilationError = {
+    val err = CompilationErrorDetails.fromResolutionFailure(rf)
+
+    CompilationError
+      .error(
+        err,
+        defaultRange(q),
+      )
+      .copy(relatedInfo =
+        q.operationName
+          .value
+          .identifier
+          .map { qsr =>
+            DiagnosticRelatedInformation(
+              RelativeLocation(
+                DocumentReference.SameFile,
+                qsr.range,
+              ),
+              err,
+            )
+          }
+          .toList
+      )
+  }
+
   // Returns the preferred range for diagnostics about resolution failure
-  def diagnosticRange(q: Query[WithSource]): SourceRange =
+  private def defaultRange(q: Query[WithSource]): SourceRange =
     q.useClause.value match {
-      case None         => q.operationName.range
+      case None         => q.operationName.value.operationName.range
       case Some(clause) => clause.identifier.range
+      // todo: involve the optional range in q.operationName's service reference
     }
 
 }

core/src/main/scala/playground/QueryCompiler.scala

@@ -131,6 +131,7 @@ final case class CompilationError(
   range: SourceRange,
   severity: DiagnosticSeverity,
   tags: Set[DiagnosticTag],
+  relatedInfo: List[DiagnosticRelatedInformation],
 ) {
   def deprecated: CompilationError = copy(tags = tags + DiagnosticTag.Deprecated)
 
@@ -167,18 +168,36 @@ object CompilationError {
     range = range,
     severity = severity,
     tags = Set.empty,
+    relatedInfo = Nil,
   )
 
 }
 
+final case class DiagnosticRelatedInformation(
+  location: RelativeLocation,
+  message: CompilationErrorDetails,
+)
+
+final case class RelativeLocation(document: DocumentReference, range: SourceRange)
+  extends Product
+  with Serializable
+
+sealed trait DocumentReference extends Product with Serializable
+
+object DocumentReference {
+  case object SameFile extends DocumentReference
+}
+
 sealed trait CompilationErrorDetails extends Product with Serializable {
 
   def render: String =
     this match {
-      case Message(text)        => text
-      case DeprecatedItem(info) => "Deprecated" + CompletionItem.deprecationString(info)
-      case InvalidUUID          => "Invalid UUID"
-      case InvalidBlob          => "Invalid blob, expected base64-encoded string"
+      case Message(text)                  => text
+      case DeprecatedItem(info)           => "Deprecated" + CompletionItem.deprecationString(info)
+      case InvalidUUID                    => "Invalid UUID"
+      case InvalidBlob                    => "Invalid blob, expected base64-encoded string"
+      case ConflictingServiceReference(_) => "Conflicting service references"
+
       case NumberOutOfRange(value, expectedType) => s"Number out of range for $expectedType: $value"
       case EnumFallback(enumName) =>
         s"""Matching enums by value is deprecated and may be removed in the future. Use $enumName instead.""".stripMargin
@@ -239,6 +258,8 @@ object CompilationErrorDetails {
       CompilationErrorDetails.AmbiguousService(knownServices)
     case ResolutionFailure.UnknownService(unknownId, knownServices) =>
       CompilationErrorDetails.UnknownService(unknownId, knownServices)
+    case ResolutionFailure.ConflictingServiceReference(refs) =>
+      CompilationErrorDetails.ConflictingServiceReference(refs)
 
   }
 
@@ -247,6 +268,9 @@ object CompilationErrorDetails {
   final case class UnknownService(id: QualifiedIdentifier, knownServices: List[QualifiedIdentifier])
     extends CompilationErrorDetails
 
+  final case class ConflictingServiceReference(refs: List[QualifiedIdentifier])
+    extends CompilationErrorDetails
+
   final case class AmbiguousService(
     known: List[QualifiedIdentifier]
   ) extends CompilationErrorDetails

core/src/main/scala/playground/run.scala

@@ -146,7 +146,7 @@ private class ServiceCompiler[Alg[_[_, _, _, _, _]], Op[_, _, _, _, _]](
   private def operationNotFound(q: Query[WithSource]): CompilationError = CompilationError.error(
     CompilationErrorDetails
       .OperationNotFound(
-        q.operationName.value.mapK(WithSource.unwrap),
+        q.operationName.value.operationName.value.mapK(WithSource.unwrap),
         endpoints.keys.map(OperationName[Id](_)).toList,
       ),
     q.operationName.range,
@@ -196,7 +196,7 @@ private class ServiceCompiler[Alg[_[_, _, _, _, _]], Op[_, _, _, _, _]](
   def compile(q: Query[WithSource]): Ior[Throwable, CompiledInput] = {
     val compiled =
       endpoints
-        .get(q.operationName.value.text)
+        .get(q.operationName.value.operationName.value.text)
         .toRightIor(NonEmptyList.one(operationNotFound(q)))
         .flatTap { case (e, _) => deprecatedOperationCheck(q, e) }
         .flatMap(_._2.apply(q.input)) <& deprecationWarnings(q)
@@ -213,13 +213,13 @@ private class MultiServiceCompiler[Alg[_[_, _, _, _, _]], Op[_, _, _, _, _]](
   private def getService(
     q: Query[WithSource]
   ): Either[Throwable, Compiler[Ior[Throwable, *]]] = MultiServiceResolver
-    .resolveService(q.useClause.value.map(_.identifier.value), services)
+    .resolveService(
+      q.mapK(WithSource.unwrap).collectServiceIdentifiers,
+      services,
+    )
     .leftMap { rf =>
       CompilationFailed.one(
-        CompilationError.error(
-          CompilationErrorDetails.fromResolutionFailure(rf),
-          ResolutionFailure.diagnosticRange(q),
-        )
+        ResolutionFailure.toCompilationError(rf, q)
       )
     }
 
@@ -323,7 +323,7 @@ object Runner {
     new Resolver[F] {
       def get(q: Query[WithSource]): IorNel[Issue, Runner[F]] = MultiServiceResolver
         .resolveService(
-          q.useClause.value.map(_.identifier.value),
+          q.mapK(WithSource.unwrap).collectServiceIdentifiers,
           runners,
         )
         .leftMap(rf =>

core/src/main/scala/playground/smithyql/AST.scala

@@ -12,17 +12,18 @@ import cats.Id
 import playground.ServiceNameExtractor
 import smithy4s.Service
 import cats.kernel.Order
+import cats.Apply
 
 sealed trait AST[F[_]] extends Product with Serializable {
   def mapK[G[_]: Functor](fk: F ~> G): AST[G]
-  def kind: NodeKind
 }
 
 object AST {
   implicit val astIdEq: Eq[AST[Id]] = Eq.fromUniversalEquals
 }
 
 sealed trait InputNode[F[_]] extends AST[F] {
+  def kind: NodeKind
 
   def fold[A](
     struct: Struct[F] => A,
@@ -46,9 +47,6 @@ sealed trait InputNode[F[_]] extends AST[F] {
 
 final case class OperationName[F[_]](text: String) extends AST[F] {
   def mapK[G[_]: Functor](fk: F ~> G): OperationName[G] = copy()
-
-  def kind: NodeKind = NodeKind.OperationName
-
 }
 
 final case class QualifiedIdentifier(segments: NonEmptyList[String], selection: String) {
@@ -85,23 +83,38 @@ object QualifiedIdentifier {
 
 final case class UseClause[F[_]](identifier: F[QualifiedIdentifier]) extends AST[F] {
   def mapK[G[_]: Functor](fk: F ~> G): UseClause[G] = UseClause(fk(identifier))
-  def kind: NodeKind = NodeKind.UseClause
+}
+
+final case class QueryOperationName[F[_]](
+  identifier: Option[F[QualifiedIdentifier]],
+  operationName: F[OperationName[F]],
+) extends AST[F] {
+
+  def mapK[G[_]: Functor](fk: F ~> G): QueryOperationName[G] = QueryOperationName(
+    identifier.map(fk(_)),
+    fk(operationName).map(_.mapK(fk)),
+  )
+
 }
 
 final case class Query[F[_]](
   useClause: F[Option[UseClause[F]]],
-  operationName: F[OperationName[F]],
+  operationName: F[QueryOperationName[F]],
   input: F[Struct[F]],
 ) extends AST[F] {
 
-  def kind: NodeKind = NodeKind.Query
-
   def mapK[G[_]: Functor](fk: F ~> G): Query[G] = Query(
     fk(useClause).map(_.map(_.mapK(fk))),
     fk(operationName).map(_.mapK(fk)),
     fk(input).map(_.mapK(fk)),
   )
 
+  def collectServiceIdentifiers(implicit F: Apply[F]): F[List[F[QualifiedIdentifier]]] =
+    (
+      useClause.map(_.map(_.identifier)),
+      operationName.map(_.identifier),
+    ).mapN(_.toList ++ _.toList)
+
 }
 
 final case class Struct[F[_]](
@@ -197,9 +210,6 @@ object NodeKind {
   case object IntLiteral extends NodeKind
   case object NullLiteral extends NodeKind
   case object StringLiteral extends NodeKind
-  case object Query extends NodeKind
   case object Listed extends NodeKind
   case object Bool extends NodeKind
-  case object UseClause extends NodeKind
-  case object OperationName extends NodeKind
 }

core/src/main/scala/playground/smithyql/DSL.scala

@@ -12,7 +12,7 @@ object DSL {
       args: (String, InputNode[Id])*
     ): Query[Id] = Query[Id](
       useClause = None,
-      operationName = OperationName(s),
+      operationName = QueryOperationName[Id](None, OperationName(s)),
       input = struct(args: _*),
     )
 

core/src/main/scala/playground/smithyql/Formatter.scala

@@ -12,6 +12,13 @@ object Formatter {
 
   def writeAst(ast: AST[WithSource]): Doc =
     ast match {
+      case qo: QueryOperationName[WithSource] =>
+        // Comments inside this whole node are not allowed, so we ignore them
+        qo.identifier
+          .map(_.value)
+          .fold(Doc.empty)(renderIdent(_) + Doc.char('.')) +
+          writeAst(qo.operationName.value)
+
       case o: OperationName[WithSource] => renderOperationName(o)
       case q: Query[WithSource]         => writeQuery(q)
       case u: UseClause[WithSource]     => renderUseClause(u)