Merge pull request #2348 from nspcc-dev/compiler-defer-conditional

compiler: properly process `defer` in conditional statements

pkg/compiler/codegen.go

@@ -1024,10 +1024,14 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		binary.LittleEndian.PutUint16(param[0:], catch)
 		binary.LittleEndian.PutUint16(param[4:], finally)
 		emit.Instruction(c.prog.BinWriter, opcode.TRYL, param)
+		index := c.scope.newLocal(fmt.Sprintf("defer@%d", n.Call.Pos()))
+		emit.Opcodes(c.prog.BinWriter, opcode.PUSH1)
+		c.emitStoreByIndex(varLocal, index)
 		c.scope.deferStack = append(c.scope.deferStack, deferInfo{
 			catchLabel:   catch,
 			finallyLabel: finally,
 			expr:         n.Call,
+			localIndex:   index,
 		})
 		return nil
 
@@ -1330,13 +1334,21 @@ func (c *codegen) isCallExprSyscall(e ast.Expr) bool {
 // 2. `recover` can or can not handle a possible exception.
 // Thus we use the following approach:
 // 1. Throwed exception is saved in a static field X, static fields Y and is set to true.
-// 2. CATCH and FINALLY blocks are the same, and both contain the same CALLs.
-// 3. In CATCH block we set Y to true and emit default return values if it is the last defer.
-// 4. Execute FINALLY block only if Y is false.
+// 2. For each defer local there is a dedicated local variable which is set to 1 if `defer` statement
+//    is encountered during an actual execution.
+// 3. CATCH and FINALLY blocks are the same, and both contain the same CALLs.
+// 4. Right before the CATCH block check a variable from (2). If it is null, jump to the end of CATCH+FINALLY block.
+// 5. In CATCH block we set Y to true and emit default return values if it is the last defer.
+// 6. Execute FINALLY block only if Y is false.
 func (c *codegen) processDefers() {
 	for i := len(c.scope.deferStack) - 1; i >= 0; i-- {
 		stmt := c.scope.deferStack[i]
 		after := c.newLabel()
+
+		c.emitLoadByIndex(varLocal, c.scope.deferStack[i].localIndex)
+		emit.Opcodes(c.prog.BinWriter, opcode.ISNULL)
+		emit.Jmp(c.prog.BinWriter, opcode.JMPIFL, after)
+
 		emit.Jmp(c.prog.BinWriter, opcode.ENDTRYL, after)
 
 		c.setLabel(stmt.catchLabel)

pkg/compiler/defer_test.go

@@ -2,8 +2,10 @@ package compiler_test
 
 import (
 	"math/big"
+	"strings"
 	"testing"
 
+	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 	"github.com/stretchr/testify/require"
 )
 
@@ -144,6 +146,92 @@ func TestDefer(t *testing.T) {
 	})
 }
 
+func TestConditionalDefer(t *testing.T) {
+	type testCase struct {
+		a      []bool
+		result int64
+	}
+
+	t.Run("no panic", func(t *testing.T) {
+		src := `package foo
+		var i int
+		func Main(a []bool) int { return f(a[0], a[1], a[2]) + i }
+		func g() { i += 10 }
+		func f(a bool, b bool, c bool) int {
+			if a { defer func() { i += 1 }() }
+			if b { defer g() }
+			if c { defer func() { i += 100 }() }
+			return 0
+		}`
+		testCases := []testCase{
+			{[]bool{false, false, false}, 0},
+			{[]bool{false, false, true}, 100},
+			{[]bool{false, true, false}, 10},
+			{[]bool{false, true, true}, 110},
+			{[]bool{true, false, false}, 1},
+			{[]bool{true, false, true}, 101},
+			{[]bool{true, true, false}, 11},
+			{[]bool{true, true, true}, 111},
+		}
+		for _, tc := range testCases {
+			args := []stackitem.Item{stackitem.Make(tc.a[0]), stackitem.Make(tc.a[1]), stackitem.Make(tc.a[2])}
+			evalWithArgs(t, src, nil, args, big.NewInt(tc.result))
+		}
+	})
+	t.Run("panic between ifs", func(t *testing.T) {
+		src := `package foo
+		var i int
+		func Main(a []bool) int { if a[1] { defer func() { recover() }() }; return f(a[0], a[1]) + i }
+		func f(a, b bool) int {
+			if a { defer func() { i += 1; recover() }() }
+			panic("totally expected")
+			if b { defer func() { i += 100; recover() }() }
+			return 0
+		}`
+
+		args := []stackitem.Item{stackitem.Make(false), stackitem.Make(false)}
+		v := vmAndCompile(t, src)
+		v.Estack().PushVal(args)
+		err := v.Run()
+		require.Error(t, err)
+		require.True(t, strings.Contains(err.Error(), "totally expected"))
+
+		testCases := []testCase{
+			{[]bool{false, true}, 0},
+			{[]bool{true, false}, 1},
+			{[]bool{true, true}, 1},
+		}
+		for _, tc := range testCases {
+			args := []stackitem.Item{stackitem.Make(tc.a[0]), stackitem.Make(tc.a[1])}
+			evalWithArgs(t, src, nil, args, big.NewInt(tc.result))
+		}
+	})
+	t.Run("panic in conditional", func(t *testing.T) {
+		src := `package foo
+		var i int
+		func Main(a []bool) int { if a[1] { defer func() { recover() }() }; return f(a[0], a[1]) + i }
+		func f(a, b bool) int {
+			if a {
+				defer func() { i += 1; recover() }()
+				panic("somewhat expected")
+			}
+			if b { defer func() { i += 100; recover() }() }
+			return 0
+		}`
+
+		testCases := []testCase{
+			{[]bool{false, false}, 0},
+			{[]bool{false, true}, 100},
+			{[]bool{true, false}, 1},
+			{[]bool{true, true}, 1},
+		}
+		for _, tc := range testCases {
+			args := []stackitem.Item{stackitem.Make(tc.a[0]), stackitem.Make(tc.a[1])}
+			evalWithArgs(t, src, nil, args, big.NewInt(tc.result))
+		}
+	})
+}
+
 func TestRecover(t *testing.T) {
 	t.Run("Panic", func(t *testing.T) {
 		src := `package foo

pkg/compiler/func_scope.go

@@ -53,6 +53,7 @@ type deferInfo struct {
 	catchLabel   uint16
 	finallyLabel uint16
 	expr         *ast.CallExpr
+	localIndex   int
 }
 
 const (

pkg/vm/vm.go

@@ -1594,6 +1594,9 @@ func (v *VM) call(ctx *Context, offset int) {
 	newCtx.retCount = -1
 	newCtx.local = nil
 	newCtx.arguments = nil
+	// If memory for `elems` is reused, we can end up
+	// with incorrect exception context state in the caller.
+	newCtx.tryStack.elems = nil
 	initStack(&newCtx.tryStack, "exception", nil)
 	newCtx.NEF = ctx.NEF
 	v.istack.PushItem(newCtx)

pkg/vm/vm_test.go

@@ -1027,6 +1027,22 @@ func TestTRY(t *testing.T) {
 		// add 5 to the exception, mul to the result of inner finally (2)
 		getTRYTestFunc(47, inner, append(add5, byte(opcode.MUL)), add9)(t)
 	})
+	t.Run("nested, in throw and catch in call", func(t *testing.T) {
+		catchP := []byte{byte(opcode.PUSH10), byte(opcode.ADD)}
+		inner := getTRYProgram(throw, catchP, []byte{byte(opcode.PUSH2)})
+		outer := getTRYProgram([]byte{byte(opcode.CALL), 0}, []byte{byte(opcode.PUSH3)}, []byte{byte(opcode.PUSH4)})
+		outer = append(outer, byte(opcode.RET))
+		outer[4] = byte(len(outer) - 3) // CALL argument at 3 (TRY) + 1 (CALL opcode)
+		outer = append(outer, inner...)
+		outer = append(outer, byte(opcode.RET))
+
+		v := load(outer)
+		runVM(t, v)
+		require.Equal(t, 3, v.Estack().Len())
+		require.Equal(t, big.NewInt(4), v.Estack().Pop().Value())  // outer FINALLY
+		require.Equal(t, big.NewInt(2), v.Estack().Pop().Value())  // inner FINALLY
+		require.Equal(t, big.NewInt(23), v.Estack().Pop().Value()) // inner THROW + CATCH
+	})
 }
 
 func TestMEMCPY(t *testing.T) {