From 67d9d94fb3151ceee82f681c3d25baa3028e9df7 Mon Sep 17 00:00:00 2001
From: yuyi <yuyi98@163.com>
Date: Thu, 25 Jun 2020 02:41:26 +0800
Subject: [PATCH] map: fix m[a]++ etc

---
 vlib/builtin/map.v      | 19 +++++++++++++++++++
 vlib/builtin/map_test.v | 33 ++++++++++++++++++++++++++++++---
 vlib/v/gen/cgen.v       | 27 ++++++++++++++++++++-------
 3 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/vlib/builtin/map.v b/vlib/builtin/map.v
index 9eee563bc4..166dd1c1d3 100644
--- a/vlib/builtin/map.v
+++ b/vlib/builtin/map.v
@@ -359,6 +359,25 @@ fn (m map) get3(key string, zero voidptr) voidptr {
 	return zero
 }
 
+fn (mut m map) get2(key string, zero voidptr) voidptr {
+	for {
+		mut index,mut meta := m.key_to_index(key)
+		for {
+			if meta == m.metas[index] {
+				kv_index := m.metas[index + 1]
+				if fast_string_eq(key, m.key_values.keys[kv_index]) {
+					return voidptr(m.key_values.values + kv_index * u32(m.value_bytes))
+				}
+			}
+			index += 2
+			meta += probe_inc
+			if meta > m.metas[index] { break }
+		}
+		// set zero if not found
+		m.set(key, zero)
+	}
+}
+
 fn (m map) exists(key string) bool {
 	mut index,mut meta := m.key_to_index(key)
 	for {
diff --git a/vlib/builtin/map_test.v b/vlib/builtin/map_test.v
index b4eca25481..11bf73d4a2 100644
--- a/vlib/builtin/map_test.v
+++ b/vlib/builtin/map_test.v
@@ -239,13 +239,40 @@ fn test_map_assign() {
 		's': u16(5)
 		't': 3
 	}
-	d := Mstruct1 {
+	_ := Mstruct1 {
 		{ 'p': 12 }
 	}
-	e := Mstruct2 {
+	_ := Mstruct2 {
 		{ 'q': 1.7 }
 	}
-	f := Mstruct3 {
+	_ := Mstruct3 {
 		{ 'r': u16(6), 's': 5 }
 	}
 }
+
+fn test_postfix_op_directly() {
+	mut a := map[string]int
+	a['aaa']++
+	assert a['aaa'] == 1
+	a['aaa']++
+	assert a['aaa'] == 2
+	a['bbb']--
+	assert a['bbb'] == -1
+	a['bbb']--
+	assert a['bbb'] == -2
+}
+
+fn test_map_push_directly() {
+	mut a := map[string][]string
+	a['aaa'] << ['a', 'b', 'c']
+	assert a['aaa'].len == 3
+	assert a['aaa'] == ['a', 'b', 'c']
+}
+
+fn test_assign_directly() {
+	mut a := map[string]int
+	a['aaa'] += 4
+	assert a['aaa'] == 4
+	a['aaa'] -= 2
+	assert a['aaa'] == 2
+}
diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v
index 290f130d1d..cafdf7fee8 100644
--- a/vlib/v/gen/cgen.v
+++ b/vlib/v/gen/cgen.v
@@ -60,6 +60,8 @@ mut:
 	is_sql               bool // Inside `sql db{}` statement, generating sql instead of C (e.g. `and` instead of `&&` etc)
 	optionals            []string // to avoid duplicates TODO perf, use map
 	inside_ternary       int // ?: comma separated statements on a single line
+	inside_map_postfix   bool // inside map++/-- postfix expr
+	inside_map_infix     bool // inside map<</+=/-= infix expr
 	ternary_names        map[string]string
 	ternary_level_names  map[string][]string
 	stmt_path_pos        []int
@@ -1485,7 +1487,13 @@ fn (mut g Gen) expr(node ast.Expr) {
 			g.index_expr(node)
 		}
 		ast.InfixExpr {
-			g.infix_expr(node)
+			if node.op in [.left_shift, .plus_assign, .minus_assign] {
+				g.inside_map_infix = true
+				g.infix_expr(node)
+				g.inside_map_infix = false
+			} else {
+				g.infix_expr(node)
+			}
 		}
 		ast.IntegerLiteral {
 			if node.val.starts_with('0o') {
@@ -1530,7 +1538,9 @@ fn (mut g Gen) expr(node ast.Expr) {
 			g.write(')')
 		}
 		ast.PostfixExpr {
+			g.inside_map_postfix = true
 			g.expr(node.expr)
+			g.inside_map_postfix = false
 			g.write(node.op.str())
 		}
 		ast.PrefixExpr {
@@ -2232,14 +2242,17 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) {
 				g.write(', ')
 				g.expr(node.index)
 				g.write(', &($elem_type_str[]) { ')
-			} else {
-				/*
+			} else if g.inside_map_postfix || g.inside_map_infix {
+				zero := g.type_default(info.value_type)
 				g.write('(*($elem_type_str*)map_get2(')
+				if !left_is_ptr {
+					g.write('&')
+				}
 				g.expr(node.left)
-					g.write(', ')
-					g.expr(node.index)
-					g.write('))')
-				*/
+				g.write(', ')
+				g.expr(node.index)
+				g.write(', &($elem_type_str[]){ $zero }))\n')
+			} else {
 				zero := g.type_default(info.value_type)
 				g.write('(*($elem_type_str*)map_get3(')
 				g.expr(node.left)