Implement RangeIndex

kubukoz · Nov 10, 2024 · 75e5b2a · 75e5b2a
1 parent 12101f6
commit 75e5b2a
Show file tree

Hide file tree

Showing 5 changed files with 148 additions and 53 deletions.
diff --git a/modules/ast/src/test/scala/playground/Assertions.scala b/modules/ast/src/test/scala/playground/Assertions.scala
@@ -23,7 +23,7 @@ object Assertions extends Expectations.Helpers {
         val stringWithResets = d.show()(conf).linesWithSeparators.map(Console.RESET + _).mkString
 
         failure(
-          s"Diff failed:\n${Console.RESET}(${conf.left("expected")}, ${conf.right("actual")})\n\n" + stringWithResets
+          s"Diff failed:\n${Console.RESET}(${conf.left("actual")}, ${conf.right("expected")})\n\n" + stringWithResets
         )
     }
 

diff --git a/modules/core/src/main/scala/playground/smithyql/RangeIndex.scala b/modules/core/src/main/scala/playground/smithyql/RangeIndex.scala
@@ -1,73 +1,151 @@
 package playground.smithyql
 
+import cats.kernel.Monoid
 import cats.syntax.all.*
+import org.polyvariant.treesitter4s.Node
+import util.chaining.*
 
 trait RangeIndex {
 
   def findAtPosition(
     pos: Position
-  ): NodeContext
+  ): Option[NodeContext]
 
 }
 
 object RangeIndex {
 
-  // def build(parsed: playground.generated.nodes.SourceFile) = ???
+  given Monoid[RangeIndex] = Monoid.instance(
+    _ => None,
+    (a, b) => pos => a.findAtPosition(pos).orElse(b.findAtPosition(pos)),
+  )
 
-  def build(
-    sf: SourceFile[WithSource]
-  ): RangeIndex =
-    new RangeIndex {
-
-      private val allRanges: List[ContextRange] = {
-        val path = NodeContext.EmptyPath
-
-        val preludeRanges: List[ContextRange] = sf
-          .prelude
-          .useClauses
-          .toNel
-          .foldMap { useClauses =>
-            val newBase = path.inPrelude
-
-            ContextRange(useClauses.map(_.range).reduceLeft(_.fakeUnion(_)), newBase) ::
-              sf.prelude
-                .useClauses
-                .mapWithIndex {
-                  (
-                    uc,
-                    i,
-                  ) =>
-                    findInUseClause(uc, newBase.inUseClause(i))
-                }
-                .combineAll
+  def build(parsed: playground.generated.nodes.SourceFile): RangeIndex = fromRanges {
+
+    extension (node: Node) {
+      def range: SourceRange = SourceRange(Position(node.startByte), Position(node.endByte))
+    }
+
+    val root = NodeContext.EmptyPath
+
+    val preludeRanges = parsed
+      .prelude
+      .toList
+      .flatMap { prelude =>
+        val newBase = root.inPrelude
+        ContextRange(prelude.range, newBase) ::
+          prelude.use_clause.zipWithIndex.map { (useClause, i) =>
+            ContextRange(useClause.range, newBase.inUseClause(i))
           }
+      }
 
-        val queryRanges = sf.queries(WithSource.unwrap).zipWithIndex.flatMap { case (rq, index) =>
-          findInQuery(rq.query, path.inQuery(index))
-        }
+    def inputNodeRanges(node: playground.generated.nodes.InputNode, base: NodeContext)
+      : List[ContextRange] =
+      node match {
+        case playground.generated.nodes.String_(string) =>
+          ContextRange(string.range.shrink1, base.inQuotes) :: Nil
+
+        case playground.generated.nodes.List_(list) =>
+          ContextRange(list.range.shrink1, base.inCollectionEntry(None)) ::
+            list.list_fields.zipWithIndex.flatMap { (inputNode, i) =>
+              ContextRange(inputNode.range, base.inCollectionEntry(Some(i))) ::
+                inputNodeRanges(inputNode, base.inCollectionEntry(Some(i)))
+            }
+
+        case playground.generated.nodes.Struct(struct) =>
+          ContextRange(struct.range, base) ::
+            ContextRange(struct.range.shrink1, base.inStructBody) ::
+            struct.bindings.toList.flatMap { binding =>
+              (binding.key, binding.value).tupled.toList.flatMap { (key, value) =>
+                ContextRange(
+                  value.range,
+                  base
+                    .inStructBody
+                    .inStructValue(key.source),
+                ) :: inputNodeRanges(value, base.inStructBody.inStructValue(key.source))
+              }
+            }
+
+        case _ => Nil
+      }
 
-        preludeRanges ++ queryRanges
+    val queryRanges = parsed.statements.zipWithIndex.flatMap { (stat, statementIndex) =>
+      stat.run_query.toList.flatMap { runQuery =>
+        ContextRange(runQuery.range, root.inQuery(statementIndex)) :: runQuery
+          .operation_name
+          .toList
+          .flatMap { operationName =>
+            ContextRange(operationName.range, root.inQuery(statementIndex).inOperationName) :: Nil
+          } ++
+          runQuery.input.toList.flatMap { input =>
+            inputNodeRanges(
+              playground.generated.nodes.InputNode(input).toOption.get /* todo */,
+              root.inQuery(statementIndex).inOperationInput,
+            )
+
+          }
       }
 
-      // Console
-      //   .err
-      //   .println(
-      //     s"""Found ${allRanges.size} ranges for query ${q.operationName.value.text}:
-      //        |${allRanges
-      //         .map(_.render)
-      //         .mkString("\n")}""".stripMargin
-      //   )
-
-      def findAtPosition(
-        pos: Position
-      ): NodeContext = allRanges
+    }
+
+    preludeRanges ++ queryRanges
+  }
+
+  def build(
+    sf: SourceFile[WithSource]
+  ): RangeIndex = fromRanges {
+    val path = NodeContext.EmptyPath
+
+    val preludeRanges: List[ContextRange] = sf
+      .prelude
+      .useClauses
+      .toNel
+      .foldMap { useClauses =>
+        val newBase = path.inPrelude
+
+        ContextRange(useClauses.map(_.range).reduceLeft(_.fakeUnion(_)), newBase) ::
+          sf.prelude
+            .useClauses
+            .mapWithIndex {
+              (
+                uc,
+                i,
+              ) =>
+                findInUseClause(uc, newBase.inUseClause(i))
+            }
+            .combineAll
+      }
+
+    val queryRanges = sf.queries(WithSource.unwrap).zipWithIndex.flatMap { case (rq, index) =>
+      findInQuery(rq.query, path.inQuery(index))
+    }
+
+    preludeRanges ++ queryRanges
+
+    // Console
+    //   .err
+    //   .println(
+    //     s"""Found ${allRanges.size} ranges for query ${q.operationName.value.text}:
+    //        |${allRanges
+    //         .map(_.render)
+    //         .mkString("\n")}""".stripMargin
+    //   )
+  }
+
+  def fromRanges(allRanges: List[ContextRange]): RangeIndex =
+    pos =>
+      allRanges
         .filter(_.range.contains(pos))
+        .tap { ranges =>
+          println()
+          println("=======")
+          println(s"all ranges: ${allRanges.map(_.render).mkString(", ")}")
+          println(s"ranges for position ${pos.index}: ${ranges.map(_.render).mkString(", ")}")
+          println("=======")
+          println()
+        }
         .maxByOption(_.ctx.length)
         .map(_.ctx)
-        // By default, we're on root level
-        .getOrElse(NodeContext.EmptyPath)
-
-    }
 
   private def findInQuery(
     q: WithSource[Query[WithSource]],

diff --git a/modules/core/src/test/scala/playground/smithyql/AtPositionTests.scala b/modules/core/src/test/scala/playground/smithyql/AtPositionTests.scala
@@ -1,5 +1,7 @@
 package playground.smithyql
 
+import cats.syntax.all.*
+import org.polyvariant.treesitter4s.TreeSitterAPI
 import playground.Assertions.*
 import playground.Diffs.given
 import playground.smithyql.parser.SourceParser
@@ -24,15 +26,22 @@ object AtPositionTests extends FunSuite {
     text: String
   ): NodeContext = {
     val (extracted, position) = extractCursor(text)
+    val parsedTs = playground
+      .generated
+      .nodes
+      .SourceFile
+      .unsafeApply(TreeSitterAPI.make("smithyql").parse(extracted).rootNode.get)
+
     val parsed =
       SourceParser[SourceFile]
         .parse(extracted)
         .toTry
         .get
 
     RangeIndex
-      .build(parsed)
+      .build(parsedTs)
       .findAtPosition(position)
+      .getOrElse(NodeContext.EmptyPath)
   }
 
   // tests for before/after/between queries

diff --git a/modules/language-support/src/main/scala/playground/language/CompletionProvider.scala b/modules/language-support/src/main/scala/playground/language/CompletionProvider.scala
@@ -3,6 +3,7 @@ package playground.language
 import cats.Id
 import cats.kernel.Order.catsKernelOrderingForOrder
 import cats.syntax.all.*
+import org.polyvariant.treesitter4s.TreeSitterAPI
 import playground.MultiServiceResolver
 import playground.ServiceIndex
 import playground.smithyql.NodeContext
@@ -169,16 +170,22 @@ object CompletionProvider {
     (
       doc,
       pos,
-    ) =>
+    ) => {
+      val parsedTs = playground
+        .generated
+        .nodes
+        .SourceFile
+        .unsafeApply(TreeSitterAPI.make("smithyql").parse(doc).rootNode.get)
+
       SourceParser[SourceFile].parse(doc) match {
         case Left(_) =>
           // we can try to deal with this later
           Nil
 
         case Right(sf) =>
-          val matchingNode = RangeIndex
-            .build(sf)
+          val matchingNode = RangeIndex.build(parsedTs)
             .findAtPosition(pos)
+            .getOrElse(NodeContext.EmptyPath)
 
           // System.err.println("matchingNode: " + matchingNode.render)
 
@@ -208,6 +215,7 @@ object CompletionProvider {
           }
 
       }
+    }
   }
 
 }
diff --git a/...yql/parser/v3/TreeSitterParserTests.scala → ...yql/parser/v3/TreeSitterParserTests.scala b/...yql/parser/v3/TreeSitterParserTests.scala → ...yql/parser/v3/TreeSitterParserTests.scala