From e2627f1c6c17f7a9c2dc986d57f5836880d0dcef Mon Sep 17 00:00:00 2001
From: Jaakko Paju <jaakko.paju@wolt.com>
Date: Sun, 4 Feb 2024 03:09:42 +0200
Subject: [PATCH] Improve tree-sitter queries for Scala (#9475)

- Simplify function highlighting
- Highlight extension methods
- Textobject query (mia/maa) for class/trait constructor parameters/arguments
- Textobject query (mif/maf) for Scala 3 braceless lambdas
---
 runtime/queries/scala/highlights.scm  | 27 ++++++---------------------
 runtime/queries/scala/textobjects.scm | 10 ++++++++--
 2 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/runtime/queries/scala/highlights.scm b/runtime/queries/scala/highlights.scm
index 4f90bfda66b2..e21a3909d1cf 100644
--- a/runtime/queries/scala/highlights.scm
+++ b/runtime/queries/scala/highlights.scm
@@ -53,28 +53,13 @@
 (var_declaration
   name: (identifier) @variable)
 
-; method definition
+; function definitions/declarations
 
-(class_definition
-  body: (template_body
-    [
-      (function_definition
-        name: (identifier) @function.method)
-      (function_declaration
-        name: (identifier) @function.method)
-    ]))
-(trait_definition
-  body: (template_body
-    [
-      (function_definition
-        name: (identifier) @function.method)
-      (function_declaration
-        name: (identifier) @function.method)
-    ]))
-(object_definition
-  body: (template_body
-    (function_definition
-      name: (identifier) @function.method)))
+(function_declaration
+    name: (identifier) @function.method)
+
+(function_definition
+      name: (identifier) @function.method)
 
 ; imports/exports
 
diff --git a/runtime/queries/scala/textobjects.scm b/runtime/queries/scala/textobjects.scm
index 6e551c417bf7..21286b3ef4f2 100644
--- a/runtime/queries/scala/textobjects.scm
+++ b/runtime/queries/scala/textobjects.scm
@@ -1,12 +1,15 @@
 ; Function queries
 
 (function_definition
-  body: (_) @function.inside) @function.around
+  body: (_) @function.inside) @function.around ; Does not include end marker
 
-; Does not match block lambdas or Scala 3 braceless lambdas
 (lambda_expression
   (_) @function.inside) @function.around
 
+; Scala 3 braceless lambda
+(colon_argument
+  (_) @function.inside) @function.around
+
 
 ; Class queries
 
@@ -32,6 +35,9 @@
 (parameters
   ((_) @parameter.inside . ","? @parameter.around) @parameter.around)
 
+(class_parameters
+  ((_) @parameter.inside . ","? @parameter.around) @parameter.around)
+
 (parameter_types
   ((_) @parameter.inside . ","? @parameter.around) @parameter.around)