More powerful select

kubukoz · Nov 11, 2024 · 08fbbd6 · 08fbbd6
1 parent 5458c9b
commit 08fbbd6
Show file tree

Hide file tree

Showing 21 changed files with 91 additions and 42 deletions.
diff --git a/modules/parser-gen/src/main/scala/playground/parsergen/ParserGen.scala b/modules/parser-gen/src/main/scala/playground/parsergen/ParserGen.scala
@@ -99,6 +99,9 @@ private def renderUnion(u: Type.Union): String = {
         |
         |  final case class Selector(path: List[$name]) extends Selection[$name] {
         |${selectorMethods.indentTrim(4)}
+        |
+        |    type Self = Selector
+        |    protected val remake = Selector.apply
         |  }
         |}
         |""".stripMargin
@@ -249,6 +252,9 @@ private def renderProduct(p: Type.Product): String = {
         |
         |  final case class Selector(path: List[$name]) extends Selection[$name] {
         |${selectorMethods.indentTrim(4)}
+        |
+        |    type Self = Selector
+        |    protected val remake = Selector.apply
         |  }
         |}
         |""".stripMargin

diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Binding.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Binding.scala
@@ -35,5 +35,8 @@ object Binding {
   final case class Selector(path: List[Binding]) extends Selection[Binding] {
     def key: Identifier.Selector = Identifier.Selector(path.flatMap(_.key))
     def value: InputNode.Selector = InputNode.Selector(path.flatMap(_.value))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Boolean_.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Boolean_.scala
@@ -20,5 +20,8 @@ object Boolean_ {
 
   final case class Selector(path: List[Boolean_]) extends Selection[Boolean_] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Comment.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Comment.scala
@@ -20,5 +20,8 @@ object Comment {
 
   final case class Selector(path: List[Comment]) extends Selection[Comment] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Identifier.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Identifier.scala
@@ -20,5 +20,8 @@ object Identifier {
 
   final case class Selector(path: List[Identifier]) extends Selection[Identifier] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/InputNode.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/InputNode.scala
@@ -40,5 +40,8 @@ object InputNode {
     def number : Number.Selector = Number.Selector(path.flatMap(_.asNumber))
     def string : String_.Selector = String_.Selector(path.flatMap(_.asString))
     def struct : Struct.Selector = Struct.Selector(path.flatMap(_.asStruct))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/List_.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/List_.scala
@@ -30,5 +30,8 @@ object List_ {
 
   final case class Selector(path: List[List_]) extends Selection[List_] {
     def list_fields: InputNode.Selector = InputNode.Selector(path.flatMap(_.list_fields))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Null_.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Null_.scala
@@ -20,5 +20,8 @@ object Null_ {
 
   final case class Selector(path: List[Null_]) extends Selection[Null_] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Number.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Number.scala
@@ -20,5 +20,8 @@ object Number {
 
   final case class Selector(path: List[Number]) extends Selection[Number] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/OperationName.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/OperationName.scala
@@ -35,5 +35,8 @@ object OperationName {
   final case class Selector(path: List[OperationName]) extends Selection[OperationName] {
     def identifier: QualifiedIdentifier.Selector = QualifiedIdentifier.Selector(path.flatMap(_.identifier))
     def name: Identifier.Selector = Identifier.Selector(path.flatMap(_.name))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Prelude.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Prelude.scala
@@ -32,5 +32,8 @@ object Prelude {
 
   final case class Selector(path: List[Prelude]) extends Selection[Prelude] {
     def use_clause: UseClause.Selector = UseClause.Selector(path.flatMap(_.use_clause))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/QualifiedIdentifier.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/QualifiedIdentifier.scala
@@ -35,5 +35,8 @@ object QualifiedIdentifier {
   final case class Selector(path: List[QualifiedIdentifier]) extends Selection[QualifiedIdentifier] {
     def namespace: Identifier.Selector = Identifier.Selector(path.flatMap(_.namespace))
     def selection: Identifier.Selector = Identifier.Selector(path.flatMap(_.selection))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/RunQuery.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/RunQuery.scala
@@ -35,5 +35,8 @@ object RunQuery {
   final case class Selector(path: List[RunQuery]) extends Selection[RunQuery] {
     def input: Struct.Selector = Struct.Selector(path.flatMap(_.input))
     def operation_name: OperationName.Selector = OperationName.Selector(path.flatMap(_.operation_name))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/SourceFile.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/SourceFile.scala
@@ -35,5 +35,8 @@ object SourceFile {
   final case class Selector(path: List[SourceFile]) extends Selection[SourceFile] {
     def prelude: Prelude.Selector = Prelude.Selector(path.flatMap(_.prelude))
     def statements: TopLevelStatement.Selector = TopLevelStatement.Selector(path.flatMap(_.statements))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/String_.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/String_.scala
@@ -20,5 +20,8 @@ object String_ {
 
   final case class Selector(path: List[String_]) extends Selection[String_] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Struct.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Struct.scala
@@ -30,5 +30,8 @@ object Struct {
 
   final case class Selector(path: List[Struct]) extends Selection[Struct] {
     def bindings: Binding.Selector = Binding.Selector(path.flatMap(_.bindings))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/TopLevelStatement.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/TopLevelStatement.scala
@@ -32,5 +32,8 @@ object TopLevelStatement {
 
   final case class Selector(path: List[TopLevelStatement]) extends Selection[TopLevelStatement] {
     def run_query: RunQuery.Selector = RunQuery.Selector(path.flatMap(_.run_query))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/UseClause.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/UseClause.scala
@@ -35,5 +35,8 @@ object UseClause {
   final case class Selector(path: List[UseClause]) extends Selection[UseClause] {
     def identifier: QualifiedIdentifier.Selector = QualifiedIdentifier.Selector(path.flatMap(_.identifier))
     def whitespace: Whitespace.Selector = Whitespace.Selector(path.flatMap(_.whitespace))
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/generated/nodes/Whitespace.scala b/modules/treesitter/src/main/scala/playground/generated/nodes/Whitespace.scala
@@ -20,5 +20,8 @@ object Whitespace {
 
   final case class Selector(path: List[Whitespace]) extends Selection[Whitespace] {
 
+
+    type Self = Selector
+    protected val remake = Selector.apply
   }
 }
diff --git a/modules/treesitter/src/main/scala/playground/treesitter4s/std/Selection.scala b/modules/treesitter/src/main/scala/playground/treesitter4s/std/Selection.scala
@@ -1,5 +1,10 @@
 package playground.treesitter4s.std
 
 trait Selection[A] {
+  type Self <: Selection[A]
   def path: List[A]
+  protected def remake: List[A] => Self
+
+  def transform(f: List[A] => List[A]): Self = remake(f(path))
+  def find(f: A => Boolean): Self = transform(_.find(f).toList)
 }
diff --git a/modules/treesitter/src/test/scala/playground/smithyql/parser/v3/TreeSitterParserTests.scala b/modules/treesitter/src/test/scala/playground/smithyql/parser/v3/TreeSitterParserTests.scala
@@ -50,29 +50,20 @@ object TreeSitterParserTests extends FunSuite {
     val in = parse("""use service foo.bar.baz.bax#Baz
                      |GetBaz { a = { x = 42 } }""".stripMargin)
 
-    // this ain't pretty huh
-    // watch out for the upcoming lookup DSL
     val valueOfX =
-      in.statements
-        .head
-        .run_query
-        .get
-        .input
-        .get
-        .bindings
-        .find(_.key.get.source == "a")
-        .get
-        .value
-        .get
-        .asStruct
-        .get
-        .bindings
-        .find(_.key.get.source == "x")
-        .get
-        .value
-        .get
-        .asNumber
-        .get
+      in.select(
+        _.statements
+          .run_query
+          .input
+          .bindings
+          .find(_.key.get.source == "a")
+          .value
+          .struct
+          .bindings
+          .find(_.key.get.source == "x")
+          .value
+          .number
+      ).head
         .source
         .toInt
 
@@ -86,26 +77,19 @@ object TreeSitterParserTests extends FunSuite {
     // this ain't pretty huh
     // watch out for the upcoming lookup DSL
     val valueOfX =
-      in.statements
-        .head
-        .run_query
-        .get
-        .input
-        .get
-        .bindings
-        .find(_.key.get.source == "a")
-        .get
-        .value
-        .get
-        .asStruct
-        .get
-        .bindings
-        .find(_.key.get.source == "x")
-        .get
-        .value
-        .get
-        .asNumber
-        .get
+      in.select(
+        _.statements
+          .run_query
+          .input
+          .bindings
+          .find(_.key.get.source == "a")
+          .value
+          .struct
+          .bindings
+          .find(_.key.get.source == "x")
+          .value
+          .number
+      ).head
         .source
         .toInt