diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go index 9f5a09e99e..4e7ac96aa8 100644 --- a/pkg/compiler/codegen.go +++ b/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) diff --git a/pkg/compiler/defer_test.go b/pkg/compiler/defer_test.go index 17275021bd..734a8c1b98 100644 --- a/pkg/compiler/defer_test.go +++ b/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 diff --git a/pkg/compiler/func_scope.go b/pkg/compiler/func_scope.go index 2874a6e432..e9a5a7ef69 100644 --- a/pkg/compiler/func_scope.go +++ b/pkg/compiler/func_scope.go @@ -53,6 +53,7 @@ type deferInfo struct { catchLabel uint16 finallyLabel uint16 expr *ast.CallExpr + localIndex int } const (