Correctly parse comments if a use clause is missing

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

@@ -2,6 +2,7 @@ package playground.smithyql
 
 import cats.Id
 import cats.implicits._
+import cats.data.NonEmptyList
 
 object DSL {
 
@@ -17,6 +18,16 @@ object DSL {
 
   }
 
+  implicit class QueryOps(val q: Query[Id]) extends AnyVal {
+
+    def useService(path0: String, pathRest: String*)(service: String): Query[Id] = q.copy[Id](
+      useClause = Some(
+        UseClause(QualifiedIdentifier(NonEmptyList(path0, pathRest.toList), service))
+      )
+    )
+
+  }
+
   def struct(
     args: (String, InputNode[Id])*
   ): Struct[Id] = Struct[Id](Struct.Fields.fromSeq[Id](args.map(_.leftMap(Struct.Key(_)))))

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

@@ -62,7 +62,6 @@ object SmithyQLParser {
       .repSep0(whitespace)
       .surroundedBy(whitespace)
 
-    // Should be kept in sync with withComments0
     def withComments[A](
       p: Parser[A]
     ): Parser[T[A]] = ((comments ~ Parser.index).with1 ~ p ~ (Parser.index ~ comments)).map {
@@ -76,20 +75,6 @@ object SmithyQLParser {
         )
     }
 
-    // Duplication of withComments for Parser0
-    def withComments0[A](
-      p: Parser0[A]
-    ): Parser0[T[A]] = ((comments ~ Parser.index) ~ p ~ (Parser.index ~ comments)).map {
-      case (((commentsBefore, indexBefore), v), (indexAfter, commentsAfter)) =>
-        val range = SourceRange(Position(indexBefore), Position(indexAfter))
-        WithSource(
-          commentsLeft = commentsBefore,
-          commentsRight = commentsAfter,
-          range = range,
-          value = v,
-        )
-    }
-
     private[SmithyQLParser] val rawIdentifier =
       (Rfc5234.alpha ~ Parser.charsWhile0(_.isLetterOrDigit))
         .map { case (ch, s) => s.prepended(ch) }
@@ -243,9 +228,22 @@ object SmithyQLParser {
         }
     }
 
-    (tokens.withComments0(useClause.?).map(_.sequence).with1 ~
-      ident.map(_.map(OperationName(_))) ~ struct ~ tokens.comments).map {
-      case (((useClause, opName), input), commentsAfter) =>
+    val useClauseWithSource: Parser0[Option[WithSource[UseClause]]] =
+      (tokens.comments ~
+        Parser.index ~ useClause ~ Parser.index).backtrack.?.map {
+        _.map { case (((commentsBefore, indexBefore), useClause), indexAfter) =>
+          WithSource(
+            commentsBefore,
+            Nil,
+            SourceRange(Position(indexBefore), Position(indexAfter)),
+            useClause,
+          )
+        }
+      }
+
+    (useClauseWithSource.with1 ~
+      ident.map(_.map(OperationName(_))) ~ struct ~ tokens.comments)
+      .map { case (((useClause, opName), input), commentsAfter) =>
         Query(
           useClause,
           opName,
@@ -256,7 +254,7 @@ object SmithyQLParser {
             value = input,
           ),
         )
-    }
+      }
   }
 
 }

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

@@ -234,7 +234,8 @@ object WithSource {
       listed = _.values.allComments(_.flatMap(_.allComments(comments))),
     )
 
-    q.operationName.allComments(_ => Nil) ++
+    q.useClause.foldMap(u => u.commentsLeft ++ u.commentsRight) ++
+      q.operationName.allComments(_ => Nil) ++
       q.input
         .allComments(
           _.fold(

core/src/test/scala/playground/smithyql/CommentParsingTests.scala

@@ -19,12 +19,12 @@ object CommentParsingTests extends SimpleIOSuite with Checkers {
     else
       CheckConfig.default
 
-  /*
   pureTest("Comments from entire query are retained while parsing") {
     assert.eql(
       SmithyQLParser.parseFull(Examples.fullOfComments).map(WithSource.allQueryComments),
       Right(
         List(
+          Comment(" before use clause"),
           Comment(" before op"),
           Comment(" after op"),
           Comment("before key"),
@@ -52,6 +52,7 @@ object CommentParsingTests extends SimpleIOSuite with Checkers {
       result.map(WithSource.allQueryComments),
       Right(
         List(
+          Comment(" before use clause"),
           Comment(" before op"),
           Comment(" after op"),
           Comment(" before key"),
@@ -68,7 +69,6 @@ object CommentParsingTests extends SimpleIOSuite with Checkers {
       ),
     )
   }
-   */
 
   implicit val showQuery: Show[Query[WithSource]] = Show.fromToString
   implicit val showStruct: Show[Struct[WithSource]] = Show.fromToString

core/src/test/scala/playground/smithyql/Examples.scala

@@ -2,6 +2,8 @@ package playground.smithyql
 
 object Examples {
   val fullOfComments = """
+// before use clause
+use service some.api#Service
 // before op
 op
 // after op

core/src/test/scala/playground/smithyql/FormattingTests.scala

@@ -8,7 +8,15 @@ import DSL._
 
 object FormattingTests extends SimpleIOSuite with Checkers {
 
-  def formattingTest(label: String)(q: Query[WithSource])(expected: String) =
+  def formattingTest(
+    label: TestName
+  )(
+    q: => Query[WithSource]
+  )(
+    expected: String
+  )(
+    implicit loc: SourceLocation
+  ) =
     pureTest(label) {
       val result = playground.smithyql.Formatter.format(q, 80)
       assert.eql(result, expected)
@@ -76,4 +84,28 @@ object FormattingTests extends SimpleIOSuite with Checkers {
       |}
       |""".stripMargin)
 
+  formattingTest("use service clause with lots of comments") {
+    parse("""//before clause
+    use service com.example#Service
+
+    // after clause
+    hello { }""")
+  }("""// before clause
+      |use service com.example#Service
+      |
+      |// after clause
+      |hello {
+      |
+      |}
+      |""".stripMargin)
+
+  formattingTest("no service clause with comment on the call") {
+    parse("""//before call
+    hello { }""")
+  }("""// before call
+      |hello {
+      |
+      |}
+      |""".stripMargin)
+
 }

core/src/test/scala/playground/smithyql/ParserTests.scala

@@ -3,7 +3,6 @@ package playground.smithyql
 import cats.Id
 import playground.smithyql.Query
 import weaver._
-import cats.data.NonEmptyList
 
 object ParserTests extends FunSuite {
 
@@ -31,27 +30,15 @@ object ParserTests extends FunSuite {
   parsingTest("simple call, sparse with underscore", " hello_world { } ")("hello_world".call())
 
   parsingTest("use service", "use service com.example#Demo hello {}")(
-    Query[Id](
-      Some(UseClause(QualifiedIdentifier(NonEmptyList.of("com", "example"), "Demo"))),
-      OperationName("hello"),
-      Struct[Id](Struct.Fields(Nil)),
-    )
+    "hello".call().useService("com", "example")("Demo")
   )
 
   parsingTest("use service with numbers", "use service com1.example2#Demo3 hello {}")(
-    Query[Id](
-      Some(UseClause(QualifiedIdentifier(NonEmptyList.of("com1", "example2"), "Demo3"))),
-      OperationName("hello"),
-      Struct[Id](Struct.Fields(Nil)),
-    )
+    "hello".call().useService("com1", "example2")("Demo3")
   )
 
   parsingTest("use service with underscore", "use service com.aws#Kinesis_2022 hello {}")(
-    Query[Id](
-      Some(UseClause(QualifiedIdentifier(NonEmptyList.of("com", "aws"), "Kinesis_2022"))),
-      OperationName("hello"),
-      Struct[Id](Struct.Fields(Nil)),
-    )
+    "hello".call().useService("com", "aws")("Kinesis_2022")
   )
 
   val simpleResult = "hello".call("world" -> "bar")
@@ -285,9 +272,11 @@ hero = { //foo
     "Comments literally everywhere",
     Examples.fullOfComments,
   )(
-    "op".call(
-      "firstKey" -> "firstValue",
-      "secondKey" -> "secondValue",
-    )
+    "op"
+      .call(
+        "firstKey" -> "firstValue",
+        "secondKey" -> "secondValue",
+      )
+      .useService("some", "api")("Service")
   )
 }