cmd/compile: load results into registers on open defer return path

When a function panics then recovers, it needs to return to the
caller with named results having the correct values. For
in-register results, we need to load them into registers at the
defer return path.

For non-open-coded defers, we already generate correct code, as
the defer return path is part of the SSA CFG and contains the
instructions that are the same as an ordinary return statement,
including putting the results to the right places.

For open-coded defers, we have a special code generation that
emits a disconnected block that currently contains only the
deferreturn call and a RET instruction. It leaves the result
registers unset. This CL adds instructions that load the result
registers on that path.

Updates #40724.

Change-Id: I1f60514da644fd5fb4b4871a1153c62f42927282
Reviewed-on: https://go-review.googlesource.com/c/go/+/307231
Trust: Cherry Zhang <cherryyz@google.com>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Austin Clements <austin@google.com>

src/cmd/compile/internal/amd64/galign.go

@@ -23,4 +23,5 @@ func Init(arch *ssagen.ArchInfo) {
 	arch.SSAMarkMoves = ssaMarkMoves
 	arch.SSAGenValue = ssaGenValue
 	arch.SSAGenBlock = ssaGenBlock
+	arch.LoadRegResults = loadRegResults
 }

src/cmd/compile/internal/amd64/ssa.go

@@ -1348,3 +1348,19 @@ func ssaGenBlock(s *ssagen.State, b, next *ssa.Block) {
 		b.Fatalf("branch not implemented: %s", b.LongString())
 	}
 }
+
+func loadRegResults(s *ssagen.State, f *ssa.Func) {
+	for _, o := range f.OwnAux.ABIInfo().OutParams() {
+		n := o.Name.(*ir.Name)
+		rts, offs := o.RegisterTypesAndOffsets()
+		for i := range o.Registers {
+			p := s.Prog(loadByType(rts[i]))
+			p.From.Type = obj.TYPE_MEM
+			p.From.Name = obj.NAME_AUTO
+			p.From.Sym = n.Linksym()
+			p.From.Offset = n.FrameOffset() + offs[i]
+			p.To.Type = obj.TYPE_REG
+			p.To.Reg = ssa.ObjRegForAbiReg(o.Registers[i], f.Config)
+		}
+	}
+}

src/cmd/compile/internal/ssa/op.go

@@ -153,6 +153,9 @@ func (a *AuxCall) Reg(i *regInfo, c *Config) *regInfo {
 func (a *AuxCall) ABI() *abi.ABIConfig {
 	return a.abiInfo.Config()
 }
+func (a *AuxCall) ABIInfo() *abi.ABIParamResultInfo {
+	return a.abiInfo
+}
 func (a *AuxCall) ResultReg(c *Config) *regInfo {
 	if a.abiInfo.OutRegistersUsed() == 0 {
 		return a.reg
@@ -171,6 +174,8 @@ func (a *AuxCall) ResultReg(c *Config) *regInfo {
 	return a.reg
 }
 
+// For ABI register index r, returns the (dense) register number used in
+// SSA backend.
 func archRegForAbiReg(r abi.RegIndex, c *Config) uint8 {
 	var m int8
 	if int(r) < len(c.intParamRegs) {
@@ -181,6 +186,13 @@ func archRegForAbiReg(r abi.RegIndex, c *Config) uint8 {
 	return uint8(m)
 }
 
+// For ABI register index r, returns the register number used in the obj
+// package (assembler).
+func ObjRegForAbiReg(r abi.RegIndex, c *Config) int16 {
+	m := archRegForAbiReg(r, c)
+	return c.registers[m].objNum
+}
+
 // ArgWidth returns the amount of stack needed for all the inputs
 // and outputs of a function or method, including ABI-defined parameter
 // slots and ABI-defined spill slots for register-resident parameters.

src/cmd/compile/internal/ssagen/arch.go

@@ -39,4 +39,9 @@ type ArchInfo struct {
 	// SSAGenBlock emits end-of-block Progs. SSAGenValue should be called
 	// for all values in the block before SSAGenBlock.
 	SSAGenBlock func(s *State, b, next *ssa.Block)
+
+	// LoadRegResults emits instructions that loads register-assigned results
+	// into registers. They are already in memory (PPARAMOUT nodes).
+	// Used in open-coded defer return path.
+	LoadRegResults func(s *State, f *ssa.Func)
 }

src/cmd/compile/internal/ssagen/ssa.go

@@ -475,8 +475,7 @@ func buildssa(fn *ir.Func, worker int) *ssa.Func {
 	case s.hasOpenDefers && (base.Ctxt.Flag_shared || base.Ctxt.Flag_dynlink) && base.Ctxt.Arch.Name == "386":
 		// Don't support open-coded defers for 386 ONLY when using shared
 		// libraries, because there is extra code (added by rewriteToUseGot())
-		// preceding the deferreturn/ret code that is generated by gencallret()
-		// that we don't track correctly.
+		// preceding the deferreturn/ret code that we don't track correctly.
 		s.hasOpenDefers = false
 	}
 	if s.hasOpenDefers && len(s.curfn.Exit) > 0 {
@@ -6409,16 +6408,6 @@ func (s *state) addNamedValue(n *ir.Name, v *ssa.Value) {
 	s.f.NamedValues[loc] = append(values, v)
 }
 
-// Generate a disconnected call to a runtime routine and a return.
-func gencallret(pp *objw.Progs, sym *obj.LSym) *obj.Prog {
-	p := pp.Prog(obj.ACALL)
-	p.To.Type = obj.TYPE_MEM
-	p.To.Name = obj.NAME_EXTERN
-	p.To.Sym = sym
-	p = pp.Prog(obj.ARET)
-	return p
-}
-
 // Branch is an unresolved branch.
 type Branch struct {
 	P *obj.Prog  // branch instruction
@@ -6707,7 +6696,20 @@ func genssa(f *ssa.Func, pp *objw.Progs) {
 		// deferreturn and a return. This will be used to during panic
 		// recovery to unwind the stack and return back to the runtime.
 		s.pp.NextLive = s.livenessMap.DeferReturn
-		gencallret(pp, ir.Syms.Deferreturn)
+		p := pp.Prog(obj.ACALL)
+		p.To.Type = obj.TYPE_MEM
+		p.To.Name = obj.NAME_EXTERN
+		p.To.Sym = ir.Syms.Deferreturn
+
+		// Load results into registers. So when a deferred function
+		// recovers a panic, it will return to caller with right results.
+		// The results are already in memory, because they are not SSA'd
+		// when the function has defers (see canSSAName).
+		if f.OwnAux.ABIInfo().OutRegistersUsed() != 0 {
+			Arch.LoadRegResults(&s, f)
+		}
+
+		pp.Prog(obj.ARET)
 	}
 
 	if inlMarks != nil {

test/abi/defer_recover_results.go

@@ -0,0 +1,36 @@
+// run
+
+// Copyright 2021 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Test that when a function recovers from a panic, it
+// returns the correct results to the caller (in particular,
+// setting the result registers correctly).
+
+package main
+
+type S struct {
+	x uint8
+	y uint16
+	z uint32
+	w float64
+}
+
+var a0, b0, c0, d0 = 10, "hello", S{1, 2, 3, 4}, [2]int{111, 222}
+
+//go:noinline
+//go:registerparams
+func F() (a int, b string, _ int, c S, d [2]int) {
+	a, b, c, d = a0, b0, c0, d0
+	defer func() { recover() }()
+	panic("XXX")
+	return
+}
+
+func main() {
+	a1, b1, zero, c1, d1 := F()
+	if a1 != a0 || b1 != b0 || c1 != c0 || d1 != d0 || zero != 0 { // unnamed result gets zero value
+		panic("FAIL")
+	}
+}