cmd/govim: extract text property removal to own func

cmd/govim/highlight.go

@@ -62,19 +62,10 @@ func (v *vimstate) redefineHighlights(diags []types.Diagnostic, force bool) erro
 		return nil
 	}
 
+	v.removeTextProps(types.DiagnosticTextPropID)
+
 	v.BatchStart()
 	defer v.BatchCancelIfNotEnded()
-	for bufnr, buf := range v.buffers {
-		if !buf.Loaded {
-			continue // vim removes properties when a buffer is unloaded
-		}
-		v.BatchChannelCall("prop_remove", struct {
-			ID    int `json:"id"`
-			BufNr int `json:"bufnr"`
-			All   int `json:"all"`
-		}{0, bufnr, 1})
-	}
-
 	for _, d := range diags {
 		// Do not add textprops to unknown buffers
 		if d.Buf < 0 {
@@ -106,3 +97,26 @@ func (v *vimstate) redefineHighlights(diags []types.Diagnostic, force bool) erro
 	v.BatchEnd()
 	return nil
 }
+
+func (v *vimstate) removeTextProps(id types.TextPropID) {
+	var didStart bool
+	if didStart = v.BatchStartIfNeeded(); didStart {
+		defer v.BatchCancelIfNotEnded()
+	}
+
+	for bufnr, buf := range v.buffers {
+		if !buf.Loaded {
+			continue // vim removes properties when a buffer is unloaded
+		}
+		v.BatchChannelCall("prop_remove", struct {
+			ID    int `json:"id"`
+			BufNr int `json:"bufnr"`
+			All   int `json:"all"`
+		}{int(id), bufnr, 1})
+	}
+
+	if didStart {
+		// prop_remove returns number of removed properties per call
+		v.BatchEnd()
+	}
+}

cmd/govim/internal/types/types.go

@@ -290,3 +290,10 @@ var SeverityHoverHighlight = map[Severity]config.Highlight{
 	SeverityInfo: config.HighlightHoverInfo,
 	SeverityHint: config.HighlightHoverHint,
 }
+
+// TextPropID is the govim internal mapping of ID used when adding/removing text properties
+type TextPropID int
+
+const (
+	DiagnosticTextPropID = 0
+)

cmd/govim/vimstate.go

@@ -113,18 +113,7 @@ func (v *vimstate) setConfig(args ...json.RawMessage) (interface{}, error) {
 	if !vimconfig.EqualBool(v.config.HighlightDiagnostics, preConfig.HighlightDiagnostics) {
 		if v.config.HighlightDiagnostics == nil || !*v.config.HighlightDiagnostics {
 			// HighlightDiagnostics is now not on - remove existing text properties
-			v.BatchStart()
-			for bufnr, buf := range v.buffers {
-				if !buf.Loaded {
-					continue // vim removes properties when a buffer is unloaded
-				}
-				v.BatchChannelCall("prop_remove", struct {
-					ID    int `json:"id"`
-					BufNr int `json:"bufnr"`
-					All   int `json:"all"`
-				}{0, bufnr, 1})
-			}
-			v.BatchEnd()
+			v.removeTextProps(types.DiagnosticTextPropID)
 		} else {
 			if err := v.redefineHighlights(v.diagnostics(), true); err != nil {
 				return nil, fmt.Errorf("failed to update diagnostic highlights: %v", err)
@@ -203,6 +192,14 @@ func (v *vimstate) BatchStart() {
 	v.currBatch = &batch{}
 }
 
+func (v *vimstate) BatchStartIfNeeded() bool {
+	if v.currBatch != nil {
+		return false
+	}
+	v.currBatch = &batch{}
+	return true
+}
+
 type batchResult func() json.RawMessage
 
 type AssertExpr string