Support multi-symbol replace_load correctly (#1366)

The logic for removing old loading symbols replaced by `replace_load` was flawed. It would leave behind old (redundant) loads, iff the `replace_load` was to replace two symbols from the same location. This change simplifies how loads are replaced by simply keeping track of non-replaced loads rather than trying to remove from the list with index manipulation.

I've added a few more test cases. Most of these new tests will fail on the old code.

Author: sbarfurth

.bazelversion

@@ -0,0 +1 @@
+8.2.1

edit/edit.go

@@ -1192,17 +1192,16 @@ func ReplaceLoad(stmts []build.Expr, location string, from, to []string) []build
 			continue
 		}
 
+		var loadTo, loadFrom []*build.Ident
 		for i, to := range load.To {
-			if toSymbols[to.Name] {
-				if i < len(load.From)-1 {
-					load.From = append(load.From[:i], load.From[i+1:]...)
-					load.To = append(load.To[:i], load.To[i+1:]...)
-				} else {
-					load.From = load.From[:i]
-					load.To = load.To[:i]
-				}
+			// Only add the load to the statement if it will NOT be replaced by a new load.
+			if !toSymbols[to.Name] {
+				loadTo = append(loadTo, load.To[i])
+				loadFrom = append(loadFrom, load.From[i])
 			}
 		}
+		load.To = loadTo
+		load.From = loadFrom
 
 		if len(load.To) > 0 {
 			all = append(all, load)

edit/edit_test.go

@@ -162,42 +162,105 @@ load("other loc", "symbol")`,
 }
 
 func TestReplaceLoad(t *testing.T) {
-	tests := []struct{ input, expected string }{
+	tests := []struct {
+		name     string
+		input    string
+		location string
+		from     []string
+		to       []string
+		expected string
+	}{
 		{
-			``,
-			`load("new_location", "symbol")`,
+			name:     "add_symbol",
+			input:    ``,
+			location: "new_location",
+			from:     []string{"symbol"},
+			to:       []string{"symbol"},
+			expected: `load("new_location", "symbol")`,
 		},
 		{
-			`load("location", "symbol")`,
-			`load("new_location", "symbol")`,
+			name:     "replace_location",
+			input:    `load("location", "symbol")`,
+			location: "new_location",
+			from:     []string{"symbol"},
+			to:       []string{"symbol"},
+			expected: `load("new_location", "symbol")`,
 		},
 		{
-			`load("location", "other", "symbol")`,
-			`load("location", "other")
+			name:     "replace_location_one_of_multiple",
+			input:    `load("location", "other", "symbol")`,
+			location: "new_location",
+			from:     []string{"symbol"},
+			to:       []string{"symbol"},
+			expected: `load("location", "other")
 load("new_location", "symbol")`,
 		},
 		{
-			`load("location", symbol = "other")`,
-			`load("new_location", "symbol")`,
+			name:     "replace_location_alias",
+			input:    `load("location", symbol = "other")`,
+			location: "new_location",
+			from:     []string{"symbol"},
+			to:       []string{"symbol"},
+			expected: `load("new_location", "symbol")`,
 		},
 		{
-			`load("other loc", "symbol")
+			name: "collapse_duplicate_symbol",
+			input: `load("other loc", "symbol")
 load("location", "symbol")`,
-			`load("new_location", "symbol")`,
+			location: "new_location",
+			from:     []string{"symbol"},
+			to:       []string{"symbol"},
+			expected: `load("new_location", "symbol")`,
+		},
+		{
+			name:     "replace_multiple_same_location",
+			input:    `load("location", "symbol_a", "symbol_b", "symbol_c")`,
+			location: "new_location",
+			from:     []string{"symbol_a", "symbol_b", "symbol_c"},
+			to:       []string{"symbol_a", "symbol_b", "symbol_c"},
+			expected: `load("new_location", "symbol_a", "symbol_b", "symbol_c")`,
+		},
+		{
+			name:     "replace_multiple_same_location_out_of_order",
+			input:    `load("location", "symbol_a", "symbol_b", "symbol_c")`,
+			location: "new_location",
+			from:     []string{"symbol_c", "symbol_a", "symbol_b"},
+			to:       []string{"symbol_c", "symbol_a", "symbol_b"},
+			expected: `load("new_location", "symbol_a", "symbol_b", "symbol_c")`,
+		},
+		{
+			name:     "replace_multiple_same_location_partial",
+			input:    `load("location", "symbol_a", "symbol_b", "symbol_c")`,
+			location: "new_location",
+			from:     []string{"symbol_a", "symbol_b"},
+			to:       []string{"symbol_a", "symbol_b"},
+			expected: `load("location", "symbol_c")
+load("new_location", "symbol_a", "symbol_b")`,
+		},
+		{
+			name:     "replace_multiple_same_location_partial_out_of_order",
+			input:    `load("location", "symbol_a", "symbol_b", "symbol_c")`,
+			location: "new_location",
+			from:     []string{"symbol_b", "symbol_a"},
+			to:       []string{"symbol_b", "symbol_a"},
+			expected: `load("location", "symbol_c")
+load("new_location", "symbol_a", "symbol_b")`,
 		},
 	}
 
 	for _, tst := range tests {
-		bld, err := build.Parse("BUILD", []byte(tst.input))
-		if err != nil {
-			t.Error(err)
-			continue
-		}
-		bld.Stmt = ReplaceLoad(bld.Stmt, "new_location", []string{"symbol"}, []string{"symbol"})
-		got := strings.TrimSpace(string(build.Format(bld)))
-		if got != tst.expected {
-			t.Errorf("maybeReplaceLoad(%s): got %s, expected %s", tst.input, got, tst.expected)
-		}
+		t.Run(tst.name, func(t *testing.T) {
+			bld, err := build.Parse("BUILD", []byte(tst.input))
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			bld.Stmt = ReplaceLoad(bld.Stmt, tst.location, tst.from, tst.to)
+			got := strings.TrimSpace(string(build.Format(bld)))
+			if got != tst.expected {
+				t.Errorf("ReplaceLoad(%s): got %s, expected %s", tst.input, got, tst.expected)
+			}
+		})
 	}
 }