Skip to content
4 min read

PR: Preserve Comments in mapsloop

Fixing a Go analyzer that dropped comments while refactoring map loops, and the AST position detail that made the difference.

  • golang
  • tooling
  • open-source

PR: golang/tools#660255 fixes Issue: golang/go#72958

Background

I discovered this issue while browsing through the Go issue tracker.

Problem Analysis

The issue reported that the mapsloop analyzer was discarding comments when refactoring map operations. The root cause was straightforward: the refactoring process wasn’t collecting and preserving existing comments.

Solution Approach

The fix involved collecting comments using the allComments function and inserting them before the refactored code. However, this required some code reorganization first.

Code Restructuring

The allComments function was originally defined inside the minmax analyzer. To reuse it in mapsloop, I needed to extract it to a shared location. Alan Donovan reviewed the code and recommended extracting rather than duplicating the function—a cleaner approach that follows the DRY principle.

Edge Case Discovery

My initial implementation collected all comments and placed them at the start node position regardless of whether mrhs was nil. This caused incorrect behavior when mrhs != nil.

The issue: When mrhs != nil (i.e., when there’s a preceding make assignment), the start node points to the wrong location:

m := make(map[string]string)
 ^  // start node position - too late!

Placing comments here would put them in the middle of the statement rather than before it.

The fix: Use curPrev.Node() instead, which points to the node before the range statement, placing comments before the entire refactored block.

Final Implementation

Here’s the complete diff:

--- a/gopls/internal/analysis/modernize/maps.go
+++ b/gopls/internal/analysis/modernize/maps.go
@@ -156,16 +156,35 @@ func mapsloop(pass *analysis.Pass) {
                        start, end token.Pos
                )
                if mrhs != nil {
-                       // Replace RHS of preceding m=... assignment (and loop) with expression.
-                       start, end = mrhs.Pos(), rng.End()
-                       newText = fmt.Appendf(nil, "%s%s(%s)",
+                       // Replace assignment and loop with expression.
+                       //
+                       //   m = make(...)
+                       //   for k, v := range x { /* comments */ m[k] = v }
+                       //
+                       //   ->
+                       //
+                       //   /* comments */
+                       //   m = maps.Copy(x)
+                       curPrev, _ := curRange.PrevSibling()
+                       start, end = curPrev.Node().Pos(), rng.End()
+                       newText = fmt.Appendf(nil, "%s%s = %s%s(%s)",
+                               allComments(file, start, end),
+                               analysisinternal.Format(pass.Fset, m),
                                prefix,
                                funcName,
                                analysisinternal.Format(pass.Fset, x))
                } else {
                        // Replace loop with call statement.
+                       //
+                       //   for k, v := range x { /* comments */ m[k] = v }
+                       //
+                       //   ->
+                       //
+                       //   /* comments */
+                       //   maps.Copy(m, x)
                        start, end = rng.Pos(), rng.End()
-                       newText = fmt.Appendf(nil, "%s%s(%s, %s)",
+                       newText = fmt.Appendf(nil, "%s%s%s(%s, %s)",
+                               allComments(file, start, end),
                                prefix,
                                funcName,
                                analysisinternal.Format(pass.Fset, m),

The diff shows two key improvements:

  1. Comment preservation: Uses allComments(file, start, end) to collect comments from the original code
  2. Correct positioning: When mrhs != nil, uses curPrev.Node().Pos() as the start position to ensure comments are placed before the entire statement block
  3. Enhanced documentation: Adds inline comments showing the before/after transformation, making the code more maintainable

Key Takeaways

  1. Don’t Repeat Yourself (DRY): When you find yourself about to copy a function, extract it to a shared location instead. A small refactor up front keeps the codebase maintainable and avoids divergent copies.

  2. Test edge cases thoroughly: My initial implementation worked for the simple case (mrhs == nil) but failed when mrhs != nil. Always exercise every code path before calling it done.

  3. AST positioning matters: When working with Go’s AST, a single wrong node position can place your changes in entirely the wrong location. The difference between rng.Pos() and curPrev.Node().Pos() is subtle but critical.

  4. Document transformations: The inline before/after examples in the code comments help future maintainers understand what the refactoring does at a glance.

← Back