From 8fa00f826c02c6de0341899f192845837b734d10 Mon Sep 17 00:00:00 2001
From: Marc Vertes <mvertes@free.fr>
Date: Mon, 18 Jan 2021 19:04:05 +0100
Subject: [PATCH] interp: fix map assignment from arithmetic operations

The logic to trigger assigment optimizations has been refactored for
clarity, and to exclude assignments to map entries.

Fixes #981.
---
 _test/issue-981.go | 15 +++++++++++++++
 interp/cfg.go      | 41 ++++++++++++++++++++++-------------------
 2 files changed, 37 insertions(+), 19 deletions(-)
 create mode 100644 _test/issue-981.go

diff --git a/_test/issue-981.go b/_test/issue-981.go
new file mode 100644
index 000000000..d72664304
--- /dev/null
+++ b/_test/issue-981.go
@@ -0,0 +1,15 @@
+package main
+
+import "fmt"
+
+func main() {
+	dp := make(map[int]int)
+	dp[0] = 1
+	for i := 1; i < 10; i++ {
+		dp[i] = dp[i-1] + dp[i-2]
+	}
+	fmt.Printf("%v\n", dp)
+}
+
+// Output:
+// map[0:1 1:1 2:2 3:3 4:5 5:8 6:13 7:21 8:34 9:55]
diff --git a/interp/cfg.go b/interp/cfg.go
index 050432cf8..9f82f22e1 100644
--- a/interp/cfg.go
+++ b/interp/cfg.go
@@ -555,10 +555,21 @@ func (interp *Interpreter) cfg(root *node, importPath string) ([]*node, error) {
 				n.findex = dest.findex
 				n.level = dest.level
 
-				// Propagate type
-				// TODO: Check that existing destination type matches source type
+				// Propagate type.
+				// TODO: Check that existing destination type matches source type.
+
+				// In the following, we attempt to optimize by skipping the assign
+				// operation and setting the source location directly to the destination
+				// location in the frame.
+				//
 				switch {
-				case n.action == aAssign && isCall(src) && dest.typ.cat != interfaceT && !isMapEntry(dest) && !isRecursiveField(dest):
+				case n.action != aAssign:
+					// Do not optimize assign combined with another operator.
+				case isMapEntry(dest):
+					// Setting a map entry needs an additional step, do not optimize.
+					// As we only write, skip the default useless getIndexMap dest action.
+					dest.gen = nop
+				case isCall(src) && dest.typ.cat != interfaceT && !isRecursiveField(dest):
 					// Call action may perform the assignment directly.
 					n.gen = nop
 					src.level = level
@@ -566,32 +577,25 @@ func (interp *Interpreter) cfg(root *node, importPath string) ([]*node, error) {
 					if src.typ.untyped && !dest.typ.untyped {
 						src.typ = dest.typ
 					}
-				case n.action == aAssign && src.action == aRecv:
+				case src.action == aRecv:
 					// Assign by reading from a receiving channel.
 					n.gen = nop
-					src.findex = dest.findex // Set recv address to LHS
+					src.findex = dest.findex // Set recv address to LHS.
 					dest.typ = src.typ
-				case n.action == aAssign && src.action == aCompositeLit && !isMapEntry(dest):
+				case src.action == aCompositeLit:
 					if dest.typ.cat == valueT && dest.typ.rtype.Kind() == reflect.Interface {
-						// Skip optimisation for assigned binary interface or map entry
-						// which require and additional operation to set the value
+						// Skip optimisation for assigned interface.
 						break
 					}
 					if dest.action == aGetIndex {
-						// optimization does not work when assigning to a struct field. Maybe we're not
-						// setting the right frame index or something, and we would end up not writing at
-						// the right place. So disabling it for now.
+						// Optimization does not work when assigning to a struct field.
 						break
 					}
-					// Skip the assign operation entirely, the source frame index is set
-					// to destination index, avoiding extra memory alloc and duplication.
 					n.gen = nop
 					src.findex = dest.findex
 					src.level = level
-				case n.action == aAssign && len(n.child) < 4 && !src.rval.IsValid() && isArithmeticAction(src):
+				case len(n.child) < 4 && !src.rval.IsValid() && isArithmeticAction(src):
 					// Optimize single assignments from some arithmetic operations.
-					// Skip the assign operation entirely, the source frame index is set
-					// to destination index, avoiding extra memory alloc and duplication.
 					src.typ = dest.typ
 					src.findex = dest.findex
 					src.level = level
@@ -600,15 +604,14 @@ func (interp *Interpreter) cfg(root *node, importPath string) ([]*node, error) {
 					// Assign to nil.
 					src.rval = reflect.New(dest.typ.TypeOf()).Elem()
 				}
+
 				n.typ = dest.typ
 				if sym != nil {
 					sym.typ = n.typ
 					sym.recv = src.recv
 				}
 				n.level = level
-				if isMapEntry(dest) {
-					dest.gen = nop // skip getIndexMap
-				}
+
 				if n.anc.kind == constDecl {
 					n.gen = nop
 					n.findex = notInFrame