internal/analysis: don't rewrite loops whose names the new code would capture

cmd/loopvec-toolexec/compile.go

@@ -321,6 +321,7 @@ func check(fset *token.FileSet, files []*ast.File, c compileArgs, cfg *importcfg
 		Defs:       make(map[*ast.Ident]types.Object),
 		Uses:       make(map[*ast.Ident]types.Object),
 		Selections: make(map[*ast.SelectorExpr]*types.Selection),
+		Scopes:     make(map[ast.Node]*types.Scope),
 	}
 	var errs []error
 	goarch := os.Getenv("GOARCH")

internal/analysis/analyzer.go

@@ -42,9 +42,10 @@ func Analyze(file *ast.File, info *types.Info, opts Options) []*loopir.Plan {
 			continue
 		}
 		results := loopir.NamedResults(fn, info)
+		reserved := loopir.UsesReservedNames(fn)
 		ast.Inspect(fn.Body, func(n ast.Node) bool {
 			if stmt, ok := n.(ast.Stmt); ok {
-				if plan, _ := analyzeLoop(stmt, info, results, opts.Plan()); plan != nil {
+				if plan, _ := analyzeLoop(stmt, info, info.Scopes[file], results, reserved, opts.Plan()); plan != nil {
 					plans = append(plans, plan)
 				}
 			}
@@ -55,8 +56,10 @@ func Analyze(file *ast.File, info *types.Info, opts Options) []*loopir.Plan {
 }
 
 // analyzeLoop lowers stmt, which must be a for or range statement, and plans
-// it. It returns the reason when the loop is not vectorizable.
-func analyzeLoop(stmt ast.Stmt, info *types.Info, results map[types.Object]bool, plan loopir.Options) (*loopir.Plan, loopir.Reason) {
+// it. It returns the reason when the loop is not vectorizable. reserved is set
+// when the enclosing function uses a name that begins with an underscore (see
+// loopir.UsesReservedNames); file is the scope of the file stmt is in.
+func analyzeLoop(stmt ast.Stmt, info *types.Info, file *types.Scope, results map[types.Object]bool, reserved bool, plan loopir.Options) (*loopir.Plan, loopir.Reason) {
 	switch stmt.(type) {
 	case *ast.RangeStmt, *ast.ForStmt:
 	default:
@@ -66,5 +69,21 @@ func analyzeLoop(stmt ast.Stmt, info *types.Info, results map[types.Object]bool,
 	if reason != "" {
 		return nil, reason
 	}
-	return loopir.NewPlan(loop, plan)
+	p, reason := loopir.NewPlan(loop, plan)
+	if p == nil {
+		return nil, reason
+	}
+	// The names are checked last, so that a loop that cannot be vectorized
+	// anyway keeps the reason that matters.
+	if reserved {
+		return nil, loopir.ReasonReservedName
+	}
+	var scope *types.Scope
+	if file != nil {
+		scope = file.Innermost(stmt.Pos())
+	}
+	if loopir.ShadowsEmittedName(scope, stmt.Pos()) || loopir.HeaderShadowsName(stmt) {
+		return nil, loopir.ReasonShadowedName
+	}
+	return p, reason
 }

internal/analysis/explain.go

@@ -32,6 +32,7 @@ func AnalyzeExplain(fset *token.FileSet, file *ast.File, info *types.Info, opts
 		}
 		isMethod := fn.Recv != nil
 		results := loopir.NamedResults(fn, info)
+		reserved := loopir.UsesReservedNames(fn)
 		ast.Inspect(fn.Body, func(n ast.Node) bool {
 			var body *ast.BlockStmt
 			switch stmt := n.(type) {
@@ -54,7 +55,7 @@ func AnalyzeExplain(fset *token.FileSet, file *ast.File, info *types.Info, opts
 				return true
 			}
 
-			plan, reason := analyzeLoop(n.(ast.Stmt), info, results, opts.Plan())
+			plan, reason := analyzeLoop(n.(ast.Stmt), info, info.Scopes[file], results, reserved, opts.Plan())
 			c.Vectorized = plan != nil
 			c.Reason = reason
 			out = append(out, c)

internal/loopir/lower.go

@@ -7,6 +7,7 @@ import (
 	"go/token"
 	"go/types"
 	"math"
+	"slices"
 	"strconv"
 	"strings"
 )
@@ -286,6 +287,12 @@ func (l *lowerer) bound(e ast.Expr) (bound, bool) {
 	}
 	b := bound{text: types.ExprString(e)}
 	b.value, b.isConst = l.constant(e)
+	if b.isConst {
+		// A named constant or a constant expression (s0, 1 - 1) is spelled as its
+		// value, so the checks that read a constant start or limit see it and the
+		// emitted code never indexes at a negative constant they would have caught.
+		b.text = strconv.FormatInt(b.value, 10)
+	}
 	if call, ok := ast.Unparen(e).(*ast.CallExpr); ok {
 		b.slice, _ = l.lenOf(call)
 	}
@@ -370,6 +377,121 @@ func isIntType(t types.Type) bool {
 	return ok && (basic.Kind() == types.Int || basic.Kind() == types.UntypedInt)
 }
 
+// UsesReservedNames reports whether fn mentions an identifier that begins with
+// an underscore, other than the blank identifier: every name loopvec adds does,
+// so code it inserts into fn could capture one.
+func UsesReservedNames(fn *ast.FuncDecl) bool {
+	found := false
+	ast.Inspect(fn, func(n ast.Node) bool {
+		if id, ok := n.(*ast.Ident); ok && len(id.Name) > 1 && id.Name[0] == '_' {
+			found = true
+		}
+		return !found
+	})
+	return found
+}
+
+// emittedNames are the predeclared and package names the rewritten code uses:
+// len, min, max and copy, uint64 for shift counts, the element types of
+// reductions (byte and rune among them, as go/types spells those), and the
+// simd and unsafe packages.
+var emittedNames = []string{
+	"len", "min", "max", "copy", "uint64",
+	"int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "float32", "float64",
+	"byte", "rune",
+	"simd", "unsafe",
+}
+
+// HeaderShadowsName reports whether a variable stmt, a for or range statement,
+// declares in its header could capture a name of the rewritten code. The
+// scalar loop that finishes a reduction declares the same variables as stmt and
+// evaluates the bound inside their scope, while stmt evaluates its range
+// operand, or the start it counts down from, outside: an index named len would
+// break len(a), and for n := range n or for n := n - 1; n >= 0; n-- would
+// compare the index with itself. So stmt must not declare one of emittedNames,
+// nor a name used by its range operand or its initial values.
+func HeaderShadowsName(stmt ast.Stmt) bool {
+	var vars, outside []ast.Expr
+	switch s := stmt.(type) {
+	case *ast.RangeStmt:
+		if s.Tok == token.DEFINE {
+			vars = []ast.Expr{s.Key, s.Value}
+		}
+		outside = []ast.Expr{s.X}
+	case *ast.ForStmt:
+		if init, ok := s.Init.(*ast.AssignStmt); ok && init.Tok == token.DEFINE {
+			vars, outside = init.Lhs, init.Rhs
+		}
+	}
+	for _, v := range vars {
+		id, ok := v.(*ast.Ident)
+		if !ok || id.Name == "_" {
+			continue
+		}
+		if slices.Contains(emittedNames, id.Name) {
+			return true
+		}
+		for _, e := range outside {
+			if mentions(e, id.Name) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// mentions reports whether e uses an identifier named name.
+func mentions(e ast.Expr, name string) bool {
+	found := false
+	ast.Inspect(e, func(n ast.Node) bool {
+		if id, ok := n.(*ast.Ident); ok && id.Name == name {
+			found = true
+		}
+		return !found
+	})
+	return found
+}
+
+// helperNames are the predeclared names that _loopvecOverlap, which loopvec
+// declares at package level, uses besides emittedNames.
+var helperNames = []string{"any", "bool", "false", "int", "uintptr"}
+
+// ShadowsEmittedName reports whether, at pos in scope, a name the rewritten code
+// uses means something else than the code needs: a predeclared name declared
+// again in an enclosing scope, the package or an import, or simd or unsafe naming
+// something other than those packages. scope is the innermost scope containing
+// pos, or nil when the type information has no scopes, which counts as shadowed.
+// The names only the overlap helper uses matter at package level, where it is
+// declared.
+func ShadowsEmittedName(scope *types.Scope, pos token.Pos) bool {
+	if scope == nil {
+		return true
+	}
+	pkg := scope
+	for pkg.Parent() != nil && pkg.Parent() != types.Universe {
+		pkg = pkg.Parent()
+	}
+	for _, name := range helperNames {
+		if pkg.Lookup(name) != nil {
+			return true
+		}
+	}
+	for _, name := range emittedNames {
+		_, obj := scope.LookupParent(name, pos)
+		switch name {
+		case "simd", "unsafe":
+			if pkg, ok := obj.(*types.PkgName); obj != nil && (!ok || pkg.Imported().Path() != name) {
+				return true
+			}
+		default:
+			if obj != types.Universe.Lookup(name) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // NamedResults returns the named result variables of fn and of every function
 // literal inside it.
 func NamedResults(fn *ast.FuncDecl, info *types.Info) map[types.Object]bool {
@@ -532,6 +654,25 @@ func (l *lowerer) body(node ast.Stmt, block *ast.BlockStmt, trip Trip, step int6
 		}
 	}
 
+	// A constant offset that puts an index below zero makes the scalar loop panic
+	// when it runs: on its first iteration when the start is a constant and the
+	// first index is negative, or on every iteration when the limit is a constant
+	// and even the last index is. It is left alone, since the vector loop would
+	// index or slice at a negative constant, which does not compile.
+	start, errStart := strconv.ParseInt(trip.Start, 10, 64)
+	limit, errLimit := strconv.ParseInt(trip.Limit, 10, 64)
+	negative := func(offset string) bool {
+		off, err := strconv.ParseInt(offset, 10, 64)
+		return err == nil && (errStart == nil && start+off < 0 || errLimit == nil && limit+off <= 0)
+	}
+	for _, stmt := range body {
+		for _, leaf := range Leaves(stmtValue(stmt)) {
+			if load, ok := leaf.(*Load); ok && negative(load.Off) {
+				return nil, ReasonNegativeIndex
+			}
+		}
+	}
+
 	loop := &Loop{
 		Node: node,
 		Ind:  Induction{Var: l.iv, Step: step, Trip: trip},

internal/loopir/reason.go

@@ -42,6 +42,12 @@ const (
 	// iteration writes (a[i] = a[i-1] + x): iteration i reads what iteration i-1
 	// wrote, so the iterations cannot run as a vector.
 	ReasonCarriedDependence Reason = "an iteration reads an element an earlier iteration wrote"
+	// ReasonNegativeIndex covers a slice read at a constant offset that puts an
+	// index below zero: the first index, with a constant start (a[i-1] in a loop
+	// from 0), or every index, with a constant limit (a[i-20] in a loop that stops
+	// before 16). The scalar loop panics when it runs, and the vector loop would
+	// index or slice at a negative constant, which does not compile.
+	ReasonNegativeIndex Reason = "a constant offset puts an index below zero"
 	// ReasonUnsupportedIf covers an if that is not a select: a condition that is
 	// not a comparison of loop values (joined by && and ||), an init statement, or
 	// a branch that is not one assignment to an element of the same slice.
@@ -96,6 +102,18 @@ const (
 	// propagate the NaN and order the zeros. arm64 matches Go, so the result
 	// would depend on the machine.
 	ReasonFloatMinMax Reason = "float min and max differ from Go's on some architectures (NaN and signed zero)"
+	// ReasonReservedName covers a loop in a function that uses an identifier
+	// beginning with an underscore, other than _ itself. Every name loopvec adds
+	// begins with one (_i, _count, _w1, _v1, _vcFloat32s0, _sx_f, ...), and code
+	// it adds before or inside the loop could shadow the function's own and change
+	// what it refers to.
+	ReasonReservedName Reason = "the function uses a name that begins with an underscore, as loopvec's own names do"
+	// ReasonShadowedName covers a loop where a predeclared name the rewritten
+	// code uses (len, min, max, copy, uint64, an element type) is declared again
+	// in a scope around it, its package or an import, or where simd or unsafe
+	// names something other than those packages: the code would call or convert
+	// with the wrong thing.
+	ReasonShadowedName Reason = "a name the rewritten code relies on (len, min, simd, ...) means something else here"
 	// ReasonMethodSkipped covers a loop inside a method, skipped unless
 	// -methods is passed (see analysis.Analyze's doc comment).
 	ReasonMethodSkipped Reason = "loops inside methods are skipped without -methods"
@@ -113,7 +131,7 @@ const (
 // and "lower" for everything about the loop's own shape.
 func (r Reason) Stage() string {
 	switch r {
-	case ReasonMethodSkipped:
+	case ReasonMethodSkipped, ReasonReservedName, ReasonShadowedName:
 		return "analysis"
 	case ReasonUnsupportedOp, ReasonFloatMinMax, ReasonFloatReduction, ReasonMixedTypes, ReasonTooManySlices, ReasonShortLoop:
 		return "plan"

internal/rewrite/rewriter.go

@@ -74,6 +74,10 @@ func File(fset *token.FileSet, file *ast.File, info *types.Info, src []byte, opt
 	}
 	needsOverlapHelper := false
 	needsSimd := false
+	// rewritten holds the functions given a rewritten loop, by their index
+	// among the file's function declarations. Rewriting changes only loop
+	// bodies and appends the overlap helper, so the indexes stay valid.
+	rewritten := map[int]bool{}
 
 	for _, plan := range plans {
 		stmtNode := plan.Loop.Node
@@ -106,6 +110,11 @@ func File(fset *token.FileSet, file *ast.File, info *types.Info, src []byte, opt
 			end:   stmtNode.End(),
 			text:  repl,
 		})
+		for i, fd := range funcDecls(file) {
+			if fd.Pos() <= stmtNode.Pos() && stmtNode.End() <= fd.End() {
+				rewritten[i] = true
+			}
+		}
 		if pre != "" {
 			extraStmts = append(extraStmts, struct {
 				pos  token.Pos
@@ -179,7 +188,7 @@ func File(fset *token.FileSet, file *ast.File, info *types.Info, src []byte, opt
 	// Rename blank identifier parameters to avoid a gotip GOEXPERIMENT=simd
 	// compiler bug where _ params in simd-using functions produce
 	// "cannot use _ as value or type".
-	formatted = fixBlankParams(fset, formatted)
+	formatted = fixBlankParams(fset, formatted, rewritten)
 
 	return Result{Src: formatted, Rewrites: len(replacements), DefinesHelper: definesHelper}, nil
 }
@@ -307,23 +316,22 @@ func ensureOverlapHelper(src []byte) []byte {
 // fixBlankParams renames blank identifier (_) function parameters in src to
 // _p0, _p1, … to work around a gotip GOEXPERIMENT=simd compiler bug:
 // blank params in functions that contain simd code cause the SIMD lowering
-// pass to emit "cannot use _ as value or type". The rename is safe because
-// the parameters are never referenced by name in the body.
-func fixBlankParams(fset *token.FileSet, src []byte) []byte {
+// pass to emit "cannot use _ as value or type". Only the functions loopvec
+// rewrote, given by their funcDecls index in rewritten, are renamed: they
+// mention no name that begins with an underscore, so the new names cannot
+// capture anything, while in another function a _p0 parameter could hide a
+// package-level _p0.
+func fixBlankParams(fset *token.FileSet, src []byte, rewritten map[int]bool) []byte {
 	file, err := parser.ParseFile(fset, "", src, parser.ParseComments)
 	if err != nil {
 		return src
 	}
 
 	counter := 0
 	changed := false
-	ast.Inspect(file, func(n ast.Node) bool {
-		fd, ok := n.(*ast.FuncDecl)
-		if !ok {
-			return true
-		}
-		if fd.Type.Params == nil {
-			return true
+	for i, fd := range funcDecls(file) {
+		if fd.Type.Params == nil || !rewritten[i] {
+			continue
 		}
 		for _, field := range fd.Type.Params.List {
 			for _, name := range field.Names {
@@ -334,8 +342,7 @@ func fixBlankParams(fset *token.FileSet, src []byte) []byte {
 				}
 			}
 		}
-		return true
-	})
+	}
 
 	if !changed {
 		return src
@@ -348,6 +355,17 @@ func fixBlankParams(fset *token.FileSet, src []byte) []byte {
 	return buf.Bytes()
 }
 
+// funcDecls returns the function declarations of file in source order.
+func funcDecls(file *ast.File) []*ast.FuncDecl {
+	var fds []*ast.FuncDecl
+	for _, decl := range file.Decls {
+		if fd, ok := decl.(*ast.FuncDecl); ok {
+			fds = append(fds, fd)
+		}
+	}
+	return fds
+}
+
 // ensureImport adds an import of path to the source if not already present.
 func ensureImport(src []byte, path string) ([]byte, error) {
 	fset := token.NewFileSet()