Merge pull request #2345 from nspcc-dev/compiler-defer-duplication

compiler: do not traverse `defer` function literals twice

pkg/compiler/codegen.go

@@ -803,8 +803,20 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		return nil
 
 	case *ast.FuncLit:
-		l := c.newLabel()
-		c.newLambda(l, n)
+		var found bool
+		var l uint16
+		for _, fs := range c.lambda {
+			if fs.decl.Body == n.Body {
+				found = true
+				l = fs.label
+				break
+			}
+		}
+		if !found {
+			l = c.newLabel()
+			c.newLambda(l, n)
+		}
+
 		buf := make([]byte, 4)
 		binary.LittleEndian.PutUint16(buf, l)
 		emit.Instruction(c.prog.BinWriter, opcode.PUSHA, buf)

pkg/compiler/defer_test.go

@@ -126,6 +126,22 @@ func TestDefer(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, 0, vm.Estack().Len(), "stack contains unexpected items")
 	})
+	t.Run("CodeDuplication", func(t *testing.T) {
+		src := `package main
+		var i int
+		func Main() {
+			defer func() {
+				var j int
+				i += j
+			}()
+			if i == 1 { return }
+			if i == 2 { return }
+			if i == 3 { return }
+			if i == 4 { return }
+			if i == 5 { return }
+		}`
+		checkCallCount(t, src, 0 /* defer body + Main */, 2, -1)
+	})
 }
 
 func TestRecover(t *testing.T) {

pkg/compiler/inline_test.go

@@ -33,7 +33,7 @@ func checkCallCount(t *testing.T, src string, expectedCall, expectedInitSlot, ex
 			actualCall++
 		case opcode.INITSLOT:
 			actualInitSlot++
-			if ctx.IP() == mainStart {
+			if ctx.IP() == mainStart && expectedLocalsMain >= 0 {
 				require.Equal(t, expectedLocalsMain, int(param[0]))
 			}
 		}