rewriteresults.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package rewriteresults rewrites local variables returned directly by a
     6  // function to use the corresponding result parameter's storage.
     7  package rewriteresults
     8  
     9  import (
    10  	"cmd/compile/internal/base"
    11  	"cmd/compile/internal/ir"
    12  	"cmd/compile/internal/types"
    13  	"fmt"
    14  	"os"
    15  )
    16  
    17  // Funcs applies the rewriteresults pass to fns.
    18  func Funcs(fns []*ir.Func) {
    19  	if base.Flag.N != 0 || base.Debug.RewriteResults == 0 {
    20  		return
    21  	}
    22  
    23  	for _, fn := range fns {
    24  		rewrite(fn)
    25  	}
    26  }
    27  
    28  func rewrite(fn *ir.Func) {
    29  	if fn == nil || len(fn.Body) == 0 {
    30  		return
    31  	}
    32  
    33  	var returns []*ir.ReturnStmt
    34  	hasDefer := false
    35  	ir.VisitList(fn.Body, func(n ir.Node) {
    36  		switch n := n.(type) {
    37  		case *ir.ReturnStmt:
    38  			returns = append(returns, n)
    39  		case *ir.GoDeferStmt:
    40  			if n.Op() == ir.ODEFER {
    41  				hasDefer = true
    42  			}
    43  		}
    44  	})
    45  	if hasDefer || len(returns) == 0 {
    46  		return
    47  	}
    48  
    49  	results := fn.Type().Results()
    50  	for _, ret := range returns {
    51  		if len(ret.Results) == 0 || len(ret.Results) != len(results) {
    52  			return
    53  		}
    54  	}
    55  
    56  	// candidates maps each local variable to the result slot whose storage
    57  	// it can use.
    58  	candidates := make(map[*ir.Name]*ir.Name)
    59  	conflicts := make(map[*ir.Name]bool)
    60  	for i, result := range results {
    61  		// If the result is already named, source already has access to
    62  		// its storage; leave those functions alone for now.
    63  		if !isAnonymousResult(result) {
    64  			continue
    65  		}
    66  		out := result.Nname.(*ir.Name)
    67  
    68  		var local *ir.Name
    69  		for _, ret := range returns {
    70  			n, ok := ret.Results[i].(*ir.Name)
    71  			if !ok {
    72  				continue
    73  			}
    74  			if !isCandidateLocal(n, result) {
    75  				continue
    76  			}
    77  			if local == nil {
    78  				local = n
    79  			} else if local != n {
    80  				local = nil
    81  				break
    82  			}
    83  		}
    84  		if local == nil {
    85  			continue
    86  		}
    87  		if prev, ok := candidates[local]; ok && prev != out {
    88  			conflicts[local] = true
    89  			continue
    90  		}
    91  		candidates[local] = out
    92  	}
    93  	if len(candidates) == 0 {
    94  		return
    95  	}
    96  	for local := range conflicts {
    97  		delete(candidates, local)
    98  	}
    99  	if len(candidates) == 0 {
   100  		return
   101  	}
   102  
   103  	captured := make(map[*ir.Name]bool)
   104  	ir.VisitList(fn.Body, func(n ir.Node) {
   105  		if n, ok := n.(*ir.ClosureExpr); ok {
   106  			for _, cv := range n.Func.ClosureVars {
   107  				captured[cv.Canonical()] = true
   108  			}
   109  		}
   110  	})
   111  	for local := range candidates {
   112  		if captured[local] {
   113  			delete(candidates, local)
   114  		}
   115  	}
   116  	if len(candidates) == 0 {
   117  		return
   118  	}
   119  	for local, out := range candidates {
   120  		if local.Addrtaken() {
   121  			out.SetAddrtaken(true)
   122  		}
   123  		out.SetUsed(true)
   124  		out.SetEsc(local.Esc())
   125  	}
   126  
   127  	if base.Debug.RewriteResults > 1 {
   128  		for local, out := range candidates {
   129  			fmt.Fprintf(os.Stderr, "rewriteresults: %v: %v => %v\n", ir.FuncName(fn), local, out)
   130  		}
   131  	}
   132  
   133  	var edit func(ir.Node) ir.Node
   134  	edit = func(n ir.Node) ir.Node {
   135  		switch n := n.(type) {
   136  		case nil:
   137  			return nil
   138  		case *ir.Name:
   139  			if out, ok := candidates[n]; ok {
   140  				return out
   141  			}
   142  			return n
   143  		}
   144  
   145  		ir.EditChildren(n, edit)
   146  		return n
   147  	}
   148  
   149  	for i, n := range fn.Body {
   150  		fn.Body[i] = edit(n)
   151  	}
   152  }
   153  
   154  func isCandidateLocal(n *ir.Name, result *types.Field) bool {
   155  	return n.Class == ir.PAUTO &&
   156  		!n.AutoTemp() &&
   157  		n.Esc() != ir.EscHeap &&
   158  		isBareDecl(n) &&
   159  		types.Identical(n.Type(), result.Type)
   160  }
   161  
   162  func isBareDecl(n *ir.Name) bool {
   163  	// Locals declared without an explicit initializer have no defining
   164  	// assignment; their nil-RHS assignment only supplies the zero value.
   165  	return n.Defn == nil
   166  }
   167  
   168  func isAnonymousResult(result *types.Field) bool {
   169  	return result.Sym == nil || result.Sym.IsBlank()
   170  }
   171
View as plain text