From ed43b75631b69aed6c42b0dfbe1a747b4066a87a Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 2 Feb 2022 16:51:10 +0300
Subject: [PATCH] compiler: do not traverse `defer` function literals twice

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/compiler/codegen.go     | 16 ++++++++++++++--
 pkg/compiler/defer_test.go  | 16 ++++++++++++++++
 pkg/compiler/inline_test.go |  2 +-
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index c2a513d91..9f5a09e99 100644
--- a/pkg/compiler/codegen.go
+++ b/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)
diff --git a/pkg/compiler/defer_test.go b/pkg/compiler/defer_test.go
index 37e43c6d8..17275021b 100644
--- a/pkg/compiler/defer_test.go
+++ b/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) {
diff --git a/pkg/compiler/inline_test.go b/pkg/compiler/inline_test.go
index 2b67d46c4..3bbf38415 100644
--- a/pkg/compiler/inline_test.go
+++ b/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]))
 			}
 		}