From 15b732e136daaca86066f228070ded35a8341abb Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Thu, 3 Oct 2024 16:28:19 +0200
Subject: [PATCH] feat(ssh_config): Add rename support; Improvements

---
 .../ssh_config/handlers/completions_tag.go    |  5 +-
 handlers/ssh_config/handlers/rename.go        | 42 ++++++++++++++
 handlers/ssh_config/indexes/indexes.go        |  9 ++-
 .../ssh_config/indexes/indexes_handlers.go    | 39 ++++++++++++-
 handlers/ssh_config/indexes/indexes_test.go   | 24 ++++++--
 .../lsp/text-document-code-action.go          |  1 -
 .../lsp/text-document-definition.go           |  4 +-
 .../lsp/text-document-prepare-rename.go       | 45 ++++++++++++++
 .../ssh_config/lsp/text-document-rename.go    | 58 +++++++++++++++++++
 root-handler/text-document-prepare-rename.go  |  4 +-
 root-handler/text-document-rename.go          |  4 +-
 11 files changed, 218 insertions(+), 17 deletions(-)
 create mode 100644 handlers/ssh_config/handlers/rename.go
 create mode 100644 handlers/ssh_config/lsp/text-document-prepare-rename.go
 create mode 100644 handlers/ssh_config/lsp/text-document-rename.go

diff --git a/handlers/ssh_config/handlers/completions_tag.go b/handlers/ssh_config/handlers/completions_tag.go
index 724dd05..61e922a 100644
--- a/handlers/ssh_config/handlers/completions_tag.go
+++ b/handlers/ssh_config/handlers/completions_tag.go
@@ -5,6 +5,7 @@ import (
 	"config-lsp/common/formatting"
 	sshconfig "config-lsp/handlers/ssh_config"
 	"config-lsp/handlers/ssh_config/ast"
+	"config-lsp/handlers/ssh_config/indexes"
 	"config-lsp/utils"
 	"fmt"
 
@@ -18,9 +19,9 @@ func getTagCompletions(
 ) ([]protocol.CompletionItem, error) {
 	return utils.MapMapToSlice(
 		d.Indexes.Tags,
-		func(name string, block *ast.SSHMatchBlock) protocol.CompletionItem {
+		func(name string, info indexes.SSHIndexTagInfo) protocol.CompletionItem {
 			kind := protocol.CompletionItemKindModule
-			text := renderMatchBlock(block)
+			text := renderMatchBlock(info.Block)
 			return protocol.CompletionItem{
 				Label: name,
 				Kind:  &kind,
diff --git a/handlers/ssh_config/handlers/rename.go b/handlers/ssh_config/handlers/rename.go
new file mode 100644
index 0000000..6813e66
--- /dev/null
+++ b/handlers/ssh_config/handlers/rename.go
@@ -0,0 +1,42 @@
+package handlers
+
+import (
+	sshconfig "config-lsp/handlers/ssh_config"
+	"errors"
+
+	protocol "github.com/tliron/glsp/protocol_3_16"
+)
+
+func RenameTag(
+	params *protocol.RenameParams,
+	d *sshconfig.SSHDocument,
+	oldName string,
+	newName string,
+) (*protocol.WorkspaceEdit, error) {
+	changes := make([]protocol.TextEdit, 0)
+
+	// tag rename
+	info, found := d.Indexes.Tags[oldName]
+	if !found {
+		return nil, errors.New("Tag could not be found")
+	}
+
+	changes = append(changes, protocol.TextEdit{
+		Range:   info.EntryValue.ToLSPRange(),
+		NewText: newName,
+	})
+
+	// Rename all occurrences
+	for _, option := range d.Indexes.TagImports[oldName] {
+		changes = append(changes, protocol.TextEdit{
+			Range:   option.OptionValue.ToLSPRange(),
+			NewText: newName,
+		})
+	}
+
+	return &protocol.WorkspaceEdit{
+		Changes: map[protocol.DocumentUri][]protocol.TextEdit{
+			params.TextDocument.URI: changes,
+		},
+	}, nil
+}
diff --git a/handlers/ssh_config/indexes/indexes.go b/handlers/ssh_config/indexes/indexes.go
index e75d0e2..9d0bf78 100644
--- a/handlers/ssh_config/indexes/indexes.go
+++ b/handlers/ssh_config/indexes/indexes.go
@@ -4,6 +4,7 @@ import (
 	"config-lsp/common"
 	"config-lsp/handlers/ssh_config/ast"
 	"config-lsp/handlers/ssh_config/fields"
+	matchparser "config-lsp/handlers/ssh_config/match-parser"
 )
 
 type ValidPath string
@@ -35,6 +36,11 @@ type SSHIndexIgnoredUnknowns struct {
 	IgnoredOptions map[fields.NormalizedOptionName]SSHIndexIgnoredUnknownInfo
 }
 
+type SSHIndexTagInfo struct {
+	EntryValue *matchparser.MatchValue
+	Block      *ast.SSHMatchBlock
+}
+
 type SSHIndexes struct {
 	AllOptionsPerName map[fields.NormalizedOptionName](map[ast.SSHBlock]([]*ast.SSHOption))
 
@@ -51,5 +57,6 @@ type SSHIndexes struct {
 	// This is a map of <line> to <option>
 	UnknownOptions map[uint32]ast.AllOptionInfo
 
-	Tags map[string]*ast.SSHMatchBlock
+	Tags       map[string]SSHIndexTagInfo
+	TagImports map[string](map[ast.SSHBlock]*ast.SSHOption)
 }
diff --git a/handlers/ssh_config/indexes/indexes_handlers.go b/handlers/ssh_config/indexes/indexes_handlers.go
index 59d4050..0e287a5 100644
--- a/handlers/ssh_config/indexes/indexes_handlers.go
+++ b/handlers/ssh_config/indexes/indexes_handlers.go
@@ -18,12 +18,14 @@ func NewSSHIndexes() *SSHIndexes {
 		Includes:          make([]*SSHIndexIncludeLine, 0),
 		IgnoredOptions:    make(map[ast.SSHBlock]SSHIndexIgnoredUnknowns),
 		UnknownOptions:    make(map[uint32]ast.AllOptionInfo),
-		Tags:              make(map[string]*ast.SSHMatchBlock),
+		Tags:              make(map[string]SSHIndexTagInfo),
+		TagImports:        make(map[string](map[ast.SSHBlock]*ast.SSHOption)),
 	}
 }
 
 var includeOption = fields.CreateNormalizedName("Include")
 var matchOption = fields.CreateNormalizedName("Match")
+var tagOption = fields.CreateNormalizedName("Tag")
 
 func CreateIndexes(config ast.SSHConfig) (*SSHIndexes, []common.LSPError) {
 	errs := make([]common.LSPError, 0)
@@ -106,14 +108,45 @@ func CreateIndexes(config ast.SSHConfig) (*SSHIndexes, []common.LSPError) {
 					// Tag already exists
 					errs = append(errs, common.LSPError{
 						Range: entry.LocationRange,
-						Err:   fmt.Errorf("Tag %s has already been defined on line %d", name, existingBlock.Start.Line+1),
+						Err:   fmt.Errorf("Tag %s has already been defined on line %d", name, existingBlock.Block.Start.Line+1),
 					})
 
 					continue
 				}
 
 				// Add tag
-				indexes.Tags[name] = block
+				indexes.Tags[name] = SSHIndexTagInfo{
+					EntryValue: value,
+					Block:      block,
+				}
+			}
+		}
+	}
+
+	// Add Tag imports
+	for block, options := range indexes.AllOptionsPerName[tagOption] {
+		for _, option := range options {
+			if option.OptionValue == nil || option.OptionValue.Value.Value == "" {
+				continue
+			}
+
+			tagName := option.OptionValue.Value.Value
+
+			if tagImport, found := indexes.TagImports[tagName]; found {
+				if definedOption, found := tagImport[block]; found {
+					errs = append(errs, common.LSPError{
+						Range: option.OptionValue.LocationRange,
+						Err:   fmt.Errorf("Tag %s has already been imported on line %d", tagName, definedOption.Start.Line+1),
+					})
+
+					continue
+				} else {
+					tagImport[block] = option
+				}
+			} else {
+				indexes.TagImports[tagName] = map[ast.SSHBlock]*ast.SSHOption{
+					block: option,
+				}
 			}
 		}
 	}
diff --git a/handlers/ssh_config/indexes/indexes_test.go b/handlers/ssh_config/indexes/indexes_test.go
index 3c81f18..b391c5d 100644
--- a/handlers/ssh_config/indexes/indexes_test.go
+++ b/handlers/ssh_config/indexes/indexes_test.go
@@ -3,7 +3,6 @@ package indexes
 import (
 	"config-lsp/handlers/ssh_config/ast"
 	"config-lsp/utils"
-	"fmt"
 	"testing"
 )
 
@@ -128,6 +127,7 @@ Match tagged good_ip
 
 Match tagged myuser
 	User root
+	Tag good_ip
 `)
 
 	config := ast.NewSSHConfig()
@@ -150,14 +150,26 @@ Match tagged myuser
 
 	rawFirstMatch, _ := config.Options.Get(uint32(0))
 	firstMatch := rawFirstMatch.(*ast.SSHMatchBlock)
-	println(fmt.Sprintf("%v", indexes.Tags["good_ip"]))
-	if !(indexes.Tags["good_ip"].Start.Line == firstMatch.Start.Line) {
+	if !(indexes.Tags["good_ip"].Block.Start.Line == firstMatch.Start.Line) {
 		t.Errorf("Expected first tag to be 'good_ip', but got %v", indexes.Tags)
 	}
 
-	rawSecondMatch, _ := config.Options.Get(uint32(3))
-	secondMatch := rawSecondMatch.(*ast.SSHMatchBlock)
-	if !(indexes.Tags["myuser"].Start.Line == secondMatch.Start.Line) {
+	_, secondBlock := config.FindOption(uint32(3))
+	secondMatch := secondBlock.(*ast.SSHMatchBlock)
+	if !(indexes.Tags["myuser"].Block.Start.Line == secondMatch.Start.Line) {
 		t.Errorf("Expected second tag to be 'myuser', but got %v", indexes.Tags)
 	}
+
+	if !(len(indexes.TagImports) == 1) {
+		t.Errorf("Expected 1 tag import, but got %v", indexes.TagImports)
+	}
+
+	if !(len(indexes.TagImports["good_ip"]) == 1) {
+		t.Errorf("Expected 1 tag import for 'good_ip', but got %v", indexes.TagImports["good_ip"])
+	}
+
+	tagOption, _ := config.FindOption(uint32(5))
+	if !(indexes.TagImports["good_ip"][secondBlock].Start.Line == tagOption.Start.Line) {
+		t.Errorf("Expected first tag import to be 'good_ip', but got %v", indexes.TagImports)
+	}
 }
diff --git a/handlers/ssh_config/lsp/text-document-code-action.go b/handlers/ssh_config/lsp/text-document-code-action.go
index 1d1a0e9..9631b1a 100644
--- a/handlers/ssh_config/lsp/text-document-code-action.go
+++ b/handlers/ssh_config/lsp/text-document-code-action.go
@@ -9,7 +9,6 @@ import (
 )
 
 func TextDocumentCodeAction(context *glsp.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
-	println("TextDocumentCodeAction sshconfig")
 	d := sshconfig.DocumentParserMap[params.TextDocument.URI]
 	actions := handlers.FetchCodeActions(d, params)
 
diff --git a/handlers/ssh_config/lsp/text-document-definition.go b/handlers/ssh_config/lsp/text-document-definition.go
index a2e4332..a0d3288 100644
--- a/handlers/ssh_config/lsp/text-document-definition.go
+++ b/handlers/ssh_config/lsp/text-document-definition.go
@@ -18,11 +18,11 @@ func TextDocumentDefinition(context *glsp.Context, params *protocol.DefinitionPa
 	option, _ := d.Config.FindOption(line)
 
 	if option != nil && option.Key.Key == tagOption && option.OptionValue != nil {
-		if block, found := d.Indexes.Tags[option.OptionValue.Value.Value]; found {
+		if info, found := d.Indexes.Tags[option.OptionValue.Value.Value]; found {
 			return []protocol.Location{
 				{
 					URI:   params.TextDocument.URI,
-					Range: block.ToLSPRange(),
+					Range: info.Block.ToLSPRange(),
 				},
 			}, nil
 		}
diff --git a/handlers/ssh_config/lsp/text-document-prepare-rename.go b/handlers/ssh_config/lsp/text-document-prepare-rename.go
new file mode 100644
index 0000000..e800b19
--- /dev/null
+++ b/handlers/ssh_config/lsp/text-document-prepare-rename.go
@@ -0,0 +1,45 @@
+package lsp
+
+import (
+	"config-lsp/common"
+	sshconfig "config-lsp/handlers/ssh_config"
+	"config-lsp/handlers/ssh_config/ast"
+
+	"github.com/tliron/glsp"
+	protocol "github.com/tliron/glsp/protocol_3_16"
+)
+
+func TextDocumentPrepareRename(context *glsp.Context, params *protocol.PrepareRenameParams) (any, error) {
+	d := sshconfig.DocumentParserMap[params.TextDocument.URI]
+	line := params.Position.Line
+
+	option, block := d.Config.FindOption(line)
+
+	if option == nil || option.Key == nil {
+		// Empty line
+		return nil, nil
+	}
+
+	if option.Key.Key == tagOption && option.OptionValue != nil {
+		return option.OptionValue.ToLSPRange(), nil
+	}
+
+	if option.Key.Key == matchOption {
+		matchBlock := block.(*ast.SSHMatchBlock)
+		entry := matchBlock.MatchValue.GetEntryAtPosition(common.LSPCharacterAsIndexPosition(params.Position.Character))
+
+		if entry == nil {
+			return nil, nil
+		}
+
+		value := entry.GetValueAtPosition(common.LSPCharacterAsIndexPosition(params.Position.Character))
+
+		if value == nil {
+			return nil, nil
+		}
+
+		return value.ToLSPRange(), nil
+	}
+
+	return nil, nil
+}
diff --git a/handlers/ssh_config/lsp/text-document-rename.go b/handlers/ssh_config/lsp/text-document-rename.go
new file mode 100644
index 0000000..5e6dd9f
--- /dev/null
+++ b/handlers/ssh_config/lsp/text-document-rename.go
@@ -0,0 +1,58 @@
+package lsp
+
+import (
+	"config-lsp/common"
+	"config-lsp/common/formatting"
+	sshconfig "config-lsp/handlers/ssh_config"
+	"config-lsp/handlers/ssh_config/ast"
+	"config-lsp/handlers/ssh_config/fields"
+	"config-lsp/handlers/ssh_config/handlers"
+
+	"github.com/tliron/glsp"
+	protocol "github.com/tliron/glsp/protocol_3_16"
+)
+
+var nameTemplate = formatting.FormatTemplate("/!'%s/!'")
+var matchOption = fields.CreateNormalizedName("Match")
+
+func TextDocumentRename(context *glsp.Context, params *protocol.RenameParams) (*protocol.WorkspaceEdit, error) {
+	d := sshconfig.DocumentParserMap[params.TextDocument.URI]
+	index := common.LSPCharacterAsIndexPosition(params.Position.Character)
+	line := params.Position.Line
+
+	option, block := d.Config.FindOption(line)
+
+	if option != nil && option.OptionValue != nil && option.OptionValue.Value.Value != "" {
+		newName := nameTemplate.Format(formatting.DefaultFormattingOptions, params.NewName)
+
+		if option.Key.Key == tagOption {
+			return handlers.RenameTag(
+				params,
+				d,
+				option.OptionValue.Value.Value,
+				newName,
+			)
+		}
+
+		if option.Key.Key == matchOption {
+			matchBlock := block.(*ast.SSHMatchBlock)
+
+			entry := matchBlock.MatchValue.GetEntryAtPosition(index)
+
+			if entry != nil {
+				value := entry.GetValueAtPosition(index)
+
+				if value != nil {
+					return handlers.RenameTag(
+						params,
+						d,
+						value.Value.Value,
+						newName,
+					)
+				}
+			}
+		}
+	}
+
+	return nil, nil
+}
diff --git a/root-handler/text-document-prepare-rename.go b/root-handler/text-document-prepare-rename.go
index 710fc9b..19b6fd4 100644
--- a/root-handler/text-document-prepare-rename.go
+++ b/root-handler/text-document-prepare-rename.go
@@ -2,6 +2,8 @@ package roothandler
 
 import (
 	aliases "config-lsp/handlers/aliases/lsp"
+	sshconfig "config-lsp/handlers/ssh_config/lsp"
+
 	"github.com/tliron/glsp"
 
 	protocol "github.com/tliron/glsp/protocol_3_16"
@@ -26,7 +28,7 @@ func TextDocumentPrepareRename(context *glsp.Context, params *protocol.PrepareRe
 	case LanguageSSHDConfig:
 		return nil, nil
 	case LanguageSSHConfig:
-		return nil, nil
+		return sshconfig.TextDocumentPrepareRename(context, params)
 	case LanguageFstab:
 		return nil, nil
 	case LanguageWireguard:
diff --git a/root-handler/text-document-rename.go b/root-handler/text-document-rename.go
index 9bdfafd..084d4f8 100644
--- a/root-handler/text-document-rename.go
+++ b/root-handler/text-document-rename.go
@@ -2,6 +2,8 @@ package roothandler
 
 import (
 	aliases "config-lsp/handlers/aliases/lsp"
+	sshconfig "config-lsp/handlers/ssh_config/lsp"
+
 	"github.com/tliron/glsp"
 	protocol "github.com/tliron/glsp/protocol_3_16"
 )
@@ -25,7 +27,7 @@ func TextDocumentRename(context *glsp.Context, params *protocol.RenameParams) (*
 	case LanguageSSHDConfig:
 		return nil, nil
 	case LanguageSSHConfig:
-		return nil, nil
+		return sshconfig.TextDocumentRename(context, params)
 	case LanguageFstab:
 		return nil, nil
 	case LanguageWireguard: