perf: use maps.Copy, avoid iterators+collect with csync.Map.Copy

Carlos Alexandro Becker created 2 months ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/agent/tools/mcp/init.go |  3 +--
internal/app/lsp_events.go       |  3 +--
internal/csync/maps.go           | 12 ++++++++----
internal/csync/versionedmap.go   |  5 +++++
internal/lsp/client.go           |  2 +-
5 files changed, 16 insertions(+), 9 deletions(-)

Detailed changes

internal/agent/tools/mcp/init.go

@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
-	"maps"
 	"net/http"
 	"os"
 	"os/exec"
@@ -98,7 +97,7 @@ func SubscribeEvents(ctx context.Context) <-chan pubsub.Event[Event] {
 
 // GetStates returns the current state of all MCP clients
 func GetStates() map[string]ClientInfo {
-	return maps.Collect(states.Seq2())
+	return states.Copy()
 }
 
 // GetState returns the state of a specific MCP client

internal/app/lsp_events.go

@@ -2,7 +2,6 @@ package app
 
 import (
 	"context"
-	"maps"
 	"time"
 
 	"github.com/charmbracelet/crush/internal/csync"
@@ -49,7 +48,7 @@ func SubscribeLSPEvents(ctx context.Context) <-chan pubsub.Event[LSPEvent] {
 
 // GetLSPStates returns the current state of all LSP clients
 func GetLSPStates() map[string]LSPClientInfo {
-	return maps.Collect(lspStates.Seq2())
+	return lspStates.Copy()
 }
 
 // GetLSPState returns the state of a specific LSP client

internal/csync/maps.go

@@ -96,12 +96,16 @@ func (m *Map[K, V]) Take(key K) (V, bool) {
 	return v, ok
 }
 
+// Copy returns a copy of the inner map.
+func (m *Map[K, V]) Copy() map[K]V {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return maps.Clone(m.inner)
+}
+
 // Seq2 returns an iter.Seq2 that yields key-value pairs from the map.
 func (m *Map[K, V]) Seq2() iter.Seq2[K, V] {
-	dst := make(map[K]V)
-	m.mu.RLock()
-	maps.Copy(dst, m.inner)
-	m.mu.RUnlock()
+	dst := m.Copy()
 	return func(yield func(K, V) bool) {
 		for k, v := range dst {
 			if !yield(k, v) {

internal/csync/versionedmap.go

@@ -40,6 +40,11 @@ func (m *VersionedMap[K, V]) Seq2() iter.Seq2[K, V] {
 	return m.m.Seq2()
 }
 
+// Copy returns a copy of the inner map.
+func (m *VersionedMap[K, V]) Copy() map[K]V {
+	return m.m.Copy()
+}
+
 // Len returns the number of items in the map.
 func (m *VersionedMap[K, V]) Len() int {
 	return m.m.Len()

internal/lsp/client.go

@@ -350,7 +350,7 @@ func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnos
 
 // GetDiagnostics returns all diagnostics for all files.
 func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic {
-	return maps.Collect(c.diagnostics.Seq2())
+	return c.diagnostics.Copy()
 }
 
 // OpenFileOnDemand opens a file only if it's not already open.