From 2fdcfef9aa492f4597f6c7ac6eac6c2e2befe9e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Koz=C5=82owski?= <kubukoz@gmail.com>
Date: Thu, 18 Aug 2022 18:08:55 +0200
Subject: [PATCH] Support specifying service at operation name (#82)

---
 .../playground/CommandResultReporter.scala    |  4 +-
 .../scala/playground/CompletionProvider.scala | 17 ++--
 .../scala/playground/DiagnosticProvider.scala |  1 +
 .../playground/DocumentSymbolProvider.scala   |  2 +-
 .../playground/MultiServiceResolver.scala     | 50 +++++++++-
 .../main/scala/playground/QueryCompiler.scala | 32 ++++++-
 core/src/main/scala/playground/run.scala      | 16 ++--
 .../main/scala/playground/smithyql/AST.scala  | 32 ++++---
 .../main/scala/playground/smithyql/DSL.scala  |  2 +-
 .../scala/playground/smithyql/Formatter.scala |  7 ++
 .../scala/playground/smithyql/Parser.scala    | 35 ++++++-
 .../playground/smithyql/RangeIndex.scala      |  4 +-
 .../MultiServiceResolverTests.scala           | 28 +++++-
 .../playground/smithyql/Arbitraries.scala     |  4 +
 .../playground/smithyql/FormattingTests.scala |  9 ++
 .../playground/smithyql/ParserTests.scala     | 30 ++++++
 .../playground/smithyql/PrettyPrint.scala     | 94 -------------------
 .../scala/playground/lsp/LanguageServer.scala | 31 +++---
 .../scala/playground/lsp/converters.scala     | 62 ++++++++----
 19 files changed, 284 insertions(+), 176 deletions(-)
 delete mode 100644 core/src/test/scala/playground/smithyql/PrettyPrint.scala

diff --git a/core/src/main/scala/playground/CommandResultReporter.scala b/core/src/main/scala/playground/CommandResultReporter.scala
index 6551e394..f0488f77 100644
--- a/core/src/main/scala/playground/CommandResultReporter.scala
+++ b/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)
         )
 
diff --git a/core/src/main/scala/playground/CompletionProvider.scala b/core/src/main/scala/playground/CompletionProvider.scala
index 1f36a418..397a745f 100644
--- a/core/src/main/scala/playground/CompletionProvider.scala
+++ b/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
diff --git a/core/src/main/scala/playground/DiagnosticProvider.scala b/core/src/main/scala/playground/DiagnosticProvider.scala
index d932a051..5afbdf84 100644
--- a/core/src/main/scala/playground/DiagnosticProvider.scala
+++ b/core/src/main/scala/playground/DiagnosticProvider.scala
@@ -144,6 +144,7 @@ object DiagnosticProvider {
           range,
           DiagnosticSeverity.Information,
           tags = Set.empty,
+          relatedInfo = Nil,
         )
 
     }
diff --git a/core/src/main/scala/playground/DocumentSymbolProvider.scala b/core/src/main/scala/playground/DocumentSymbolProvider.scala
index 5d3d17b8..7d466de5 100644
--- a/core/src/main/scala/playground/DocumentSymbolProvider.scala
+++ b/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(
diff --git a/core/src/main/scala/playground/MultiServiceResolver.scala b/core/src/main/scala/playground/MultiServiceResolver.scala
index 85c41702..848ba54d 100644
--- a/core/src/main/scala/playground/MultiServiceResolver.scala
+++ b/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
     }
 
 }
diff --git a/core/src/main/scala/playground/QueryCompiler.scala b/core/src/main/scala/playground/QueryCompiler.scala
index 811d5f96..5742f1e4 100644
--- a/core/src/main/scala/playground/QueryCompiler.scala
+++ b/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
diff --git a/core/src/main/scala/playground/run.scala b/core/src/main/scala/playground/run.scala
index 0a9eafa8..b707014a 100644
--- a/core/src/main/scala/playground/run.scala
+++ b/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 =>
diff --git a/core/src/main/scala/playground/smithyql/AST.scala b/core/src/main/scala/playground/smithyql/AST.scala
index c82b2ed4..25c2ff3b 100644
--- a/core/src/main/scala/playground/smithyql/AST.scala
+++ b/core/src/main/scala/playground/smithyql/AST.scala
@@ -12,10 +12,10 @@ 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 {
@@ -23,6 +23,7 @@ object AST {
 }
 
 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
 }
diff --git a/core/src/main/scala/playground/smithyql/DSL.scala b/core/src/main/scala/playground/smithyql/DSL.scala
index 6d2a7de3..a66a91ed 100644
--- a/core/src/main/scala/playground/smithyql/DSL.scala
+++ b/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: _*),
     )
 
diff --git a/core/src/main/scala/playground/smithyql/Formatter.scala b/core/src/main/scala/playground/smithyql/Formatter.scala
index c1d7cd7e..a6e654eb 100644
--- a/core/src/main/scala/playground/smithyql/Formatter.scala
+++ b/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)
diff --git a/core/src/main/scala/playground/smithyql/Parser.scala b/core/src/main/scala/playground/smithyql/Parser.scala
index a3e7127e..124ce598 100644
--- a/core/src/main/scala/playground/smithyql/Parser.scala
+++ b/core/src/main/scala/playground/smithyql/Parser.scala
@@ -36,11 +36,19 @@ object SmithyQLParser {
           case e                                  => e.toString
         }
 
+      def prep(s: String): String = s.replace(' ', '·').replace("\n", "\\n\n")
+
       s"$valid${Console.RED}$failed${Console.RESET} - expected ${underlying
           .expected
           .map(showExpectation)
-          .mkString_("/")} at offset ${underlying.failedAtOffset}, got ${Console.YELLOW}\"${failed
-          .take(10)}\"${Console.RESET} instead"
+          .mkString_("/")} after ${Console.BLUE}${prep(
+          text.take(
+            underlying.failedAtOffset
+          )
+        )}${Console.RESET}, got ${Console.YELLOW}\"${prep(
+          failed
+            .take(10)
+        )}\"${Console.RESET} instead"
     }
 
   }
@@ -134,8 +142,8 @@ object SmithyQLParser {
     // doesn't accept comments
     val qualifiedIdent: Parser[QualifiedIdentifier] =
       (
-        rawIdent.repSep(tokens.dot.surroundedBy(tokens.whitespace)),
-        tokens.hash *> rawIdent,
+        rawIdent.repSep(tokens.dot.surroundedBy(tokens.whitespace).backtrack),
+        tokens.hash.surroundedBy(tokens.whitespace) *> rawIdent,
       ).mapN(QualifiedIdentifier.apply)
 
     val useClause: Parser[UseClause[T]] = {
@@ -255,8 +263,25 @@ object SmithyQLParser {
           )
       }
 
+    val queryOperationName: Parser[T[QueryOperationName[WithSource]]] = {
+
+      val serviceRef =
+        tokens
+          .withRange(
+            qualifiedIdent.backtrack <* tokens.whitespace
+          ) <* tokens.dot
+
+      val operationName = tokens.withRange(rawIdent).map(_.map(OperationName[WithSource](_)))
+
+      tokens.withComments {
+        ((serviceRef <* tokens.whitespace).?.with1 ~ operationName).map {
+          QueryOperationName.apply[WithSource].tupled
+        }
+      }
+    }
+
     (useClauseWithSource.with1 ~
-      ident.map(_.map(OperationName[WithSource](_))) ~ struct ~ tokens.comments)
+      queryOperationName ~ struct ~ tokens.comments)
       .map { case (((useClause, opName), input), commentsAfter) =>
         Query(
           useClause,
diff --git a/core/src/main/scala/playground/smithyql/RangeIndex.scala b/core/src/main/scala/playground/smithyql/RangeIndex.scala
index 20990fa3..7c949051 100644
--- a/core/src/main/scala/playground/smithyql/RangeIndex.scala
+++ b/core/src/main/scala/playground/smithyql/RangeIndex.scala
@@ -42,10 +42,10 @@ object RangeIndex {
       .toList
 
   private def findInOperationName(
-    operationName: WithSource[OperationName[WithSource]]
+    operationName: WithSource[QueryOperationName[WithSource]]
   ): List[ContextRange] =
     ContextRange(
-      operationName.range,
+      operationName.value.operationName.range,
       NodeContext.Root.inOperationName,
     ) :: Nil
 
diff --git a/core/src/test/scala/playground/MultiServiceResolverTests.scala b/core/src/test/scala/playground/MultiServiceResolverTests.scala
index 087b36d3..d6ad2200 100644
--- a/core/src/test/scala/playground/MultiServiceResolverTests.scala
+++ b/core/src/test/scala/playground/MultiServiceResolverTests.scala
@@ -14,7 +14,7 @@ object MultiServiceResolverTests extends SimpleIOSuite with Checkers {
         name: String,
       ) =>
         val result = MultiServiceResolver.resolveService(
-          None,
+          Nil,
           Map(
             ident -> name
           ),
@@ -32,7 +32,7 @@ object MultiServiceResolverTests extends SimpleIOSuite with Checkers {
         otherServices: Map[QualifiedIdentifier, String],
       ) =>
         val result = MultiServiceResolver.resolveService(
-          Some(useClauseIdent),
+          List(useClauseIdent),
           otherServices ++ Map(
             useClauseIdent -> name
           ),
@@ -42,6 +42,26 @@ object MultiServiceResolverTests extends SimpleIOSuite with Checkers {
     }
   }
 
+  test("resolveService with any amount of services, a use clause and a service ref") {
+    forall {
+      (
+        useClauseIdent: QualifiedIdentifier,
+        serviceRef: QualifiedIdentifier,
+        services: Map[QualifiedIdentifier, String],
+      ) =>
+        val result = MultiServiceResolver.resolveService(
+          List(useClauseIdent, serviceRef),
+          services,
+        )
+
+        assert(
+          result == Left(
+            ResolutionFailure.ConflictingServiceReference(List(useClauseIdent, serviceRef))
+          )
+        )
+    }
+  }
+
   test("resolveService with any amount of services and a mismatching clause") {
     forall {
       (
@@ -51,7 +71,7 @@ object MultiServiceResolverTests extends SimpleIOSuite with Checkers {
         val ident = useClauseIdent
 
         val result = MultiServiceResolver.resolveService(
-          Some(useClauseIdent),
+          List(useClauseIdent),
           services - useClauseIdent,
         )
 
@@ -75,7 +95,7 @@ object MultiServiceResolverTests extends SimpleIOSuite with Checkers {
         val allServices = services + extraService1 + extraService2
 
         val result = MultiServiceResolver.resolveService(
-          useClauseIdentifier = None,
+          Nil,
           services = allServices,
         )
 
diff --git a/core/src/test/scala/playground/smithyql/Arbitraries.scala b/core/src/test/scala/playground/smithyql/Arbitraries.scala
index d42ad20e..f9ecad95 100644
--- a/core/src/test/scala/playground/smithyql/Arbitraries.scala
+++ b/core/src/test/scala/playground/smithyql/Arbitraries.scala
@@ -106,6 +106,10 @@ object Arbitraries {
 
   implicit val arbStruct: Arbitrary[Struct[WithSource]] = Arbitrary(genStruct(genInputNode(1)))
 
+  implicit val arbQueryOperationName: Arbitrary[QueryOperationName[WithSource]] = Arbitrary(
+    Gen.resultOf(QueryOperationName.apply[WithSource])
+  )
+
   implicit val arbQuery: Arbitrary[Query[WithSource]] = Arbitrary {
     Gen.resultOf(Query.apply[WithSource])
   }
diff --git a/core/src/test/scala/playground/smithyql/FormattingTests.scala b/core/src/test/scala/playground/smithyql/FormattingTests.scala
index e60ab388..e898462a 100644
--- a/core/src/test/scala/playground/smithyql/FormattingTests.scala
+++ b/core/src/test/scala/playground/smithyql/FormattingTests.scala
@@ -106,4 +106,13 @@ object FormattingTests extends SimpleIOSuite with Checkers {
       |
       |}
       |""".stripMargin)
+
+  formattingTest("no service clause with comment on the call and explicit service ref") {
+    parse("""//before call
+    a.b#C.hello { }""")
+  }("""// before call
+      |a.b#C.hello {
+      |
+      |}
+      |""".stripMargin)
 }
diff --git a/core/src/test/scala/playground/smithyql/ParserTests.scala b/core/src/test/scala/playground/smithyql/ParserTests.scala
index cd638765..500f8a62 100644
--- a/core/src/test/scala/playground/smithyql/ParserTests.scala
+++ b/core/src/test/scala/playground/smithyql/ParserTests.scala
@@ -3,6 +3,7 @@ package playground.smithyql
 import cats.Id
 import playground.smithyql.Query
 import weaver._
+import cats.implicits._
 
 object ParserTests extends FunSuite {
 
@@ -47,10 +48,39 @@ object ParserTests extends FunSuite {
     "hello".call().useService("com1", "example2")("Demo3")
   )
 
+  parsingTest("use service with whitespace", "use service com1 . example2 # Demo3 hello {}")(
+    "hello".call().useService("com1", "example2")("Demo3")
+  )
+
   parsingTest("use service with underscore", "use service com.aws#Kinesis_2022 hello {}")(
     "hello".call().useService("com", "aws")("Kinesis_2022")
   )
 
+  parsingTest("per-operation service reference", "com.aws#Kinesis.hello {}")(
+    Query[Id](
+      useClause = None,
+      operationName = QueryOperationName[Id](
+        QualifiedIdentifier.of("com", "aws", "Kinesis").some,
+        OperationName("hello"),
+      ),
+      struct(),
+    )
+  )
+
+  parsingTest(
+    "per-operation service reference, with whitespace",
+    "com . aws # Kinesis . hello {}",
+  )(
+    Query[Id](
+      useClause = None,
+      operationName = QueryOperationName[Id](
+        QualifiedIdentifier.of("com", "aws", "Kinesis").some,
+        OperationName("hello"),
+      ),
+      struct(),
+    )
+  )
+
   val simpleResult = "hello".call("world" -> "bar")
 
   parsingTest("one parameter single-line", """hello { world = "bar" }""")(
diff --git a/core/src/test/scala/playground/smithyql/PrettyPrint.scala b/core/src/test/scala/playground/smithyql/PrettyPrint.scala
deleted file mode 100644
index 955806d8..00000000
--- a/core/src/test/scala/playground/smithyql/PrettyPrint.scala
+++ /dev/null
@@ -1,94 +0,0 @@
-package playground.smithyql
-
-import cats.Id
-import scala.annotation.nowarn
-
-// Some pretty-printing utils for diagnostics
-object PrettyPrint {
-
-  def printList[A](
-    l: List[A]
-  )(
-    ppContent: A => String
-  ): String = l.map(ppContent).mkString("List(", ", ", ")")
-
-  def escapeString(
-    s: String
-  ) = Formatter.renderStringLiteral(s).replace("\\", "\\\\").replace("\"", "\\\"")
-
-//
-  def prettyPrintWithComments[A](withSource: WithSource[A])(ppContent: A => String): String =
-    s"""WithSource(
-      commentsLeft = ${printList(withSource.commentsLeft)(c =>
-        s"Comment(${escapeString(c.text)})"
-      )},
-      commentsRight = ${printList(withSource.commentsRight)(c =>
-        s"Comment(${escapeString(c.text)})"
-      )},
-      value = ${ppContent(withSource.value)},
-    )"""
-
-//
-  def prettyPrint(q: Query[WithSource]): String = {
-    @nowarn() // idc
-    def prettyPrintNode(node: InputNode[WithSource]): String =
-      node match {
-        case StringLiteral(ss) => s"StringLiteral(${escapeString(ss)})"
-        case BooleanLiteral(b) => s"BooleanLiteral(${b.toString})"
-        case IntLiteral(ii)    => s"IntLiteral(${ii.toString})"
-        case s @ Struct(_)     => prettyPrintStruct(s)
-        case Listed(values)    => ??? // todo
-      }
-
-    def prettyPrintStruct(s: Struct[WithSource]): String =
-      s"""Struct[WithSource](
-        ${prettyPrintWithComments(s.fields)(
-          _.value
-            .map { case (k, v) =>
-              s"${prettyPrintWithComments(k)(kk => s"Key(${escapeString(kk.text)})")} -> ${prettyPrintWithComments(v)(prettyPrintNode)},\n"
-            }
-            .mkString("Map(", "\n", ")")
-        )}
-      )"""
-
-    s"""Query[WithSource](operationName = ${prettyPrintWithComments(q.operationName)(n =>
-        s"OperationName(${escapeString(n.text)})"
-      )}, input = ${prettyPrintWithComments(q.input)(prettyPrintStruct)})"""
-  }
-
-  case class Structure(keys: Map[String, Structure]) {
-
-    def render(depth: Int): String = {
-      val indent = "  " * depth
-      keys
-        .map { case (k, v) => s"$indent$k:\n${v.render(depth + 1)}" }
-        .mkString("\n")
-    }
-
-  }
-
-  def empty = Structure(Map.empty)
-
-  def just(k: String) = Structure(Map(k -> empty))
-
-  // I literally don't care
-  @nowarn()
-  def toStructure: InputNode[Id] => Structure =
-    _.fold(
-      struct =
-        fields =>
-          Structure(
-            fields
-              .fields
-              .value
-              .map { case (k, v) => k.text -> toStructure(v) }
-              .toMap
-          ),
-      string = s => Structure(Map("string" -> just(s.value))),
-      int = i => Structure(Map("int" -> just(i.value.toString))),
-      listed = list => ???, /* todo */
-      bool = b => Structure(Map("bool" -> just(b.value.toString))),
-      nul = _ => Structure(Map("null" -> empty)),
-    )
-
-}
diff --git a/lsp/src/main/scala/playground/lsp/LanguageServer.scala b/lsp/src/main/scala/playground/lsp/LanguageServer.scala
index e63a1159..012dfeb8 100644
--- a/lsp/src/main/scala/playground/lsp/LanguageServer.scala
+++ b/lsp/src/main/scala/playground/lsp/LanguageServer.scala
@@ -190,22 +190,25 @@ object LanguageServer {
 
       def diagnostic(
         params: DocumentDiagnosticParams
-      ): F[DocumentDiagnosticReport] = TextDocumentManager[F]
-        .get(params.getTextDocument().getUri())
-        .map { documentText =>
-          val diags = diagnosticProvider.getDiagnostics(
-            params.getTextDocument().getUri(),
-            documentText,
-          )
+      ): F[DocumentDiagnosticReport] = {
+        val documentUri = params.getTextDocument().getUri()
+        TextDocumentManager[F]
+          .get(documentUri)
+          .map { documentText =>
+            val diags = diagnosticProvider.getDiagnostics(
+              params.getTextDocument().getUri(),
+              documentText,
+            )
 
-          new DocumentDiagnosticReport(
-            new RelatedFullDocumentDiagnosticReport(
-              diags
-                .map(converters.toLSP.diagnostic(documentText, _))
-                .asJava
+            new DocumentDiagnosticReport(
+              new RelatedFullDocumentDiagnosticReport(
+                diags
+                  .map(converters.toLSP.diagnostic(documentText, documentUri, _))
+                  .asJava
+              )
             )
-          )
-        }
+          }
+      }
 
       def codeLens(
         params: CodeLensParams
diff --git a/lsp/src/main/scala/playground/lsp/converters.scala b/lsp/src/main/scala/playground/lsp/converters.scala
index 618edd82..bd6679fd 100644
--- a/lsp/src/main/scala/playground/lsp/converters.scala
+++ b/lsp/src/main/scala/playground/lsp/converters.scala
@@ -21,6 +21,8 @@ import playground.smithyql.TextEdit
 
 import scala.jdk.CollectionConverters._
 import scala.util.chaining._
+import playground.DocumentReference.SameFile
+import playground.RelativeLocation
 
 object converters {
 
@@ -109,27 +111,49 @@ object converters {
         })
     }
 
-    def diagnostic(doc: String, diag: CompilationError): lsp4j.Diagnostic = new lsp4j.Diagnostic()
-      .tap(_.setRange(toLSP.range(doc, diag.range)))
-      .tap(_.setMessage(diag.err.render))
-      .tap(_.setSeverity(diag.severity match {
-        case DiagnosticSeverity.Error       => lsp4j.DiagnosticSeverity.Error
-        case DiagnosticSeverity.Information => lsp4j.DiagnosticSeverity.Information
-        case DiagnosticSeverity.Warning     => lsp4j.DiagnosticSeverity.Warning
-      }))
-      .tap(
-        _.setTags(
-          diag
-            .tags
-            .map { tag =>
-              tag match {
-                case DiagnosticTag.Deprecated => lsp4j.DiagnosticTag.Deprecated
-                case DiagnosticTag.Unused     => lsp4j.DiagnosticTag.Unnecessary
+    def diagnostic(doc: String, documentUri: String, diag: CompilationError): lsp4j.Diagnostic =
+      new lsp4j.Diagnostic()
+        .tap(_.setRange(toLSP.range(doc, diag.range)))
+        .tap(_.setMessage(diag.err.render))
+        .tap(_.setSeverity(diag.severity match {
+          case DiagnosticSeverity.Error       => lsp4j.DiagnosticSeverity.Error
+          case DiagnosticSeverity.Information => lsp4j.DiagnosticSeverity.Information
+          case DiagnosticSeverity.Warning     => lsp4j.DiagnosticSeverity.Warning
+        }))
+        .tap(
+          _.setRelatedInformation(
+            diag
+              .relatedInfo
+              .map { info =>
+                new lsp4j.DiagnosticRelatedInformation(
+                  location(doc, documentUri, info.location),
+                  info.message.render,
+                )
               }
-            }
-            .toList
-            .asJava
+              .asJava
+          )
         )
+        .tap(
+          _.setTags(
+            diag
+              .tags
+              .map { tag =>
+                tag match {
+                  case DiagnosticTag.Deprecated => lsp4j.DiagnosticTag.Deprecated
+                  case DiagnosticTag.Unused     => lsp4j.DiagnosticTag.Unnecessary
+                }
+              }
+              .toList
+              .asJava
+          )
+        )
+
+    private def location(doc: String, documentUri: String, loc: RelativeLocation): lsp4j.Location =
+      new lsp4j.Location(
+        loc.document match {
+          case SameFile => documentUri
+        },
+        range(doc, loc.range),
       )
 
     def codeLens(documentText: String, lens: CodeLens): lsp4j.CodeLens =