lsp: Implement selection ranges and linked editing range (#1722)

Selection ranges provide smart selection of code that can expand and shrink based
on knowledge of the AST rather than simple text properties like hyphens or whitespace.
I didn't think this would be particularly important, but having it used it for some time
while working on this feature, this is definitely a feature that I now wouldn't want to
be without. While not likely to be used much outside of rules, ranges also expand in both
package declarations and imports.

Linked editing range OTOH was mostly implemented to try out the feature, and at this point
has not proven to be particularly useful, and the feature is therefore kept behind a new
`REGAL_EXPERIMENTAL` flag for now. More info about that in the docs included in this PR.

Also:
- Remove the "log only once" logic in the development bundle loader, as it's no longer flooding the logs
- Use shared schema for all `TextDocumentPosition` params handlers
- Document all environment variables available for development and debugging purposes

Signed-off-by: Anders Eknert <anders@eknert.com>

Author: anderseknert

.vscode/settings.json

@@ -12,5 +12,10 @@
   "go.buildFlags": [
     "-tags",
     "e2e"
-  ]
+  ],
+  "[rego]": {
+    // This is better handled by Regal's selection ranges implementation
+    "editor.smartSelect.selectLeadingAndTrailingWhitespace": false,
+    "editor.smartSelect.selectSubwords": false
+  }
 }

bundle/bundle.go

@@ -24,10 +24,7 @@ var (
 	//go:embed *
 	regalBundle embed.FS
 
-	lastErrMsg     = atomic.Pointer[string]{}
-	successLogOnce = sync.OnceFunc(func() {
-		fmt.Fprintln(os.Stderr, "Successfully loaded development bundle")
-	})
+	lastErrMsg = atomic.Pointer[string]{}
 )
 
 type devMode struct {
@@ -92,8 +89,6 @@ func (dm *devMode) Reload() {
 	for _, c := range dm.subscribers {
 		c <- struct{}{}
 	}
-
-	successLogOnce()
 }
 
 func (dm *devMode) Subscribe(c chan struct{}) {

bundle/regal/lsp/documenthighlight/documenthighlight.rego

@@ -7,11 +7,11 @@
 #   - https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_documentHighlight
 # schemas:
 #   - input:        schema.regal.lsp.common
-#   - input.params: schema.regal.lsp.documenthighlight
+#   - input.params: schema.regal.lsp.textdocumentposition
 package regal.lsp.documenthighlight
 
 import data.regal.ast
-import data.regal.lsp.completion.location
+import data.regal.lsp.util.find
 import data.regal.lsp.util.location as uloc
 import data.regal.util
 
@@ -22,18 +22,18 @@ result["response"] := items
 # METADATA
 # description: Highlights a function args in position
 items contains item if {
-	[arg, _] := _arg_at_position
+	[arg, _] := find.arg_at_position
 
 	item := {
-		"range": uloc.to_range(util.to_location_object(arg.location)),
+		"range": uloc.parse_range(arg.location),
 		"kind": 2, # Write
 	}
 }
 
 # METADATA
 # description: Highlights function arg references in function body when clicked
 items contains item if {
-	[arg, i] := _arg_at_position
+	[arg, i] := find.arg_at_position
 
 	some expr in ast.found.expressions[sprintf("%d", [i])]
 
@@ -43,23 +43,23 @@ items contains item if {
 	value.value == arg.value
 
 	item := {
-		"range": uloc.to_range(util.to_location_object(value.location)),
+		"range": uloc.parse_range(value.location),
 		"kind": 3, # Read
 	}
 }
 
 # METADATA
 # description: Highlights function arg references in head value body when clicked
 items contains item if {
-	[arg, i] := _arg_at_position
+	[arg, i] := find.arg_at_position
 
 	walk(data.workspace.parsed[input.params.textDocument.uri].rules[i].head.value, [_, value])
 
 	value.type == "var"
 	value.value == arg.value
 
 	item := {
-		"range": uloc.to_range(util.to_location_object(value.location)),
+		"range": uloc.parse_range(value.location),
 		"kind": 3, # Read
 	}
 }
@@ -117,23 +117,6 @@ items contains item if {
 	}
 }
 
-_arg_at_position := [arg, i] if {
-	text := input.regal.file.lines[input.params.position.line]
-	word := location.word_at(text, input.params.position.character)
-
-	some i, rule in data.workspace.parsed[input.params.textDocument.uri].rules
-	some arg in rule.head.args
-
-	arg.type == "var"
-	arg.value == word.text
-
-	loc := util.to_location_object(arg.location)
-
-	input.params.position.line + 1 == loc.row
-	input.params.position.character >= loc.col - 1
-	input.params.position.character <= (loc.col - 1) + count(arg.value)
-}
-
 _find_annotation(module, row) := annotation if {
 	util.to_location_row(module.package.annotations[0].location) == row
 

bundle/regal/lsp/linkededitingrange/linkededitingrange.rego

@@ -0,0 +1,64 @@
+# METADATA
+# description: |
+#   Linked editing ranges allow having a local rename *on type* reflect in multiple places. This arguably a potentialy
+#   confusing feature, and seems to be off by default in at least VS Code, so unless we find strong use-cases for it, we
+#   should not invest much effort into polishing this. Currently we only provide experimental support of linked editing
+#   ranges for:
+#     - function arguments and references to them in the function head and body
+# related_resources:
+#   - https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_linkedEditingRange
+# schemas:
+#   - input:        schema.regal.lsp.common
+#   - input.params: schema.regal.lsp.textdocumentposition
+package regal.lsp.linkededitingrange
+
+import data.regal.ast
+import data.regal.util
+
+import data.regal.lsp.util.find
+import data.regal.lsp.util.location as uloc
+
+# METADATA
+# entrypoint: true
+default result.response.ranges := set()
+
+# This is currently an experimental feature kept behind a flag to not accidedntally result in a poor
+# user experience for anyone who perhaps mistakenly enabled linked editing in their editor. Handler, and
+# the code kept here, to allow us to quickly test out the feature for different use-cases later.
+result.response.ranges := ranges if util.parse_bool(opa.runtime().env.REGAL_EXPERIMENTAL)
+
+# METADATA
+# description: Link a function args in position
+ranges contains range if {
+	[arg, _] := find.arg_at_position
+
+	range := uloc.parse_range(arg.location)
+}
+
+# METADATA
+# description: Link function arg references in function body to arg
+ranges contains range if {
+	[arg, i] := find.arg_at_position
+
+	some expr in ast.found.expressions[sprintf("%d", [i])]
+
+	walk(expr, [_, value])
+
+	value.type == "var"
+	value.value == arg.value
+
+	range := uloc.parse_range(value.location)
+}
+
+# METADATA
+# description: Link function arg references in head value to arg
+ranges contains range if {
+	[arg, i] := find.arg_at_position
+
+	walk(data.workspace.parsed[input.params.textDocument.uri].rules[i].head.value, [_, value])
+
+	value.type == "var"
+	value.value == arg.value
+
+	range := uloc.parse_range(value.location)
+}

bundle/regal/lsp/linkededitingrange/linkededitingrange_test.rego

@@ -0,0 +1,83 @@
+package regal.lsp.linkededitingrange_test
+
+import data.regal.lsp.linkededitingrange
+
+test_linked_editing_range_function_arg if {
+	file_content := `package p
+
+foo(bar, baz) := baz if {
+    bar == baz
+}`
+
+	# querying 'ranges' directly to test without the experimental flag logic
+	ranges := linkededitingrange.ranges with input as text_document_position(2, 6)
+		with input.regal.file.lines as split(file_content, "\n")
+		with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", file_content)
+
+	expected_ranges := {
+		# function arg 'bar' in function head
+		{
+			"start": {"line": 2, "character": 4},
+			"end": {"line": 2, "character": 7},
+		},
+		# function arg 'bar' reference in function body
+		{
+			"start": {"line": 3, "character": 4},
+			"end": {"line": 3, "character": 7},
+		},
+	}
+
+	ranges == expected_ranges
+}
+
+test_linked_editing_range_disabled_without_flag if {
+	file_content := `package p
+
+foo(bar, baz) := baz if {
+    bar == baz
+}`
+
+	# querying 'ranges' directly to test without the experimental flag logic
+	ranges := linkededitingrange.result.response.ranges with input as text_document_position(2, 6)
+		with input.regal.file.lines as split(file_content, "\n")
+		with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", file_content)
+
+	ranges == set()
+}
+
+test_linked_editing_range_enabled_with_flag_set_to_true if {
+	file_content := `package p
+
+foo(bar, baz) := baz if {
+    bar == baz
+}`
+
+	# querying 'ranges' directly to test without the experimental flag logic
+	ranges := linkededitingrange.result.response.ranges with input as text_document_position(2, 6)
+		with input.regal.file.lines as split(file_content, "\n")
+		with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", file_content)
+		with opa.runtime as {"env": {"REGAL_EXPERIMENTAL": "true"}}
+
+	expected_ranges := {
+		# function arg 'bar' in function head
+		{
+			"start": {"line": 2, "character": 4},
+			"end": {"line": 2, "character": 7},
+		},
+		# function arg 'bar' reference in function body
+		{
+			"start": {"line": 3, "character": 4},
+			"end": {"line": 3, "character": 7},
+		},
+	}
+
+	ranges == expected_ranges
+}
+
+text_document_position(line, character) := {"params": {
+	"textDocument": {"uri": "file://p.rego"},
+	"position": {
+		"line": line,
+		"character": character,
+	},
+}}