fix: handle assign statements on map elements

A new token MapAssign is added to implement assignements involving
a map, slice or array, an index key expression and a value expression.

9 files changed, 116 insertions, 79 deletions

diff --git a/comp/compiler.go b/comp/compiler.go
index 47c4a88..8a136e1 100644
--- a/comp/compiler.go
+++ b/comp/compiler.go
@@ -130,8 +130,13 @@ func (c *Compiler) Generate(tokens parser.Tokens) (err error) {
 		case lang.Index:
 			showStack()
 			pop()
-			push(&symbol.Symbol{Type: pop().Type.Elem()})
-			emit(t, vm.Index)
+			s := pop()
+			if s.Type.Rtype.Kind() == reflect.Map {
+				emit(t, vm.MapIndex)
+			} else {
+				emit(t, vm.Index)
+			}
+			push(&symbol.Symbol{Type: s.Type.Elem()})
 
 		case lang.Greater:
 			push(&symbol.Symbol{Type: booleanOpType(pop(), pop())})
@@ -232,6 +237,18 @@ func (c *Compiler) Generate(tokens parser.Tokens) (err error) {
 			}
 			emit(t, vm.Vassign)
 
+		case lang.MapAssign:
+			s := stack[len(stack)-3]
+			switch s.Type.Rtype.Kind() {
+			case reflect.Array, reflect.Slice:
+				emit(t, vm.IndexSet)
+			case reflect.Map:
+				emit(t, vm.MapSet)
+			default:
+				return errorf("not a map or array: %s", s.Name)
+			}
+			stack = stack[:len(stack)-3]
+
 		case lang.Equal:
 			push(&symbol.Symbol{Type: booleanOpType(pop(), pop())})
 			emit(t, vm.Equal)
diff --git a/interp/interpreter_test.go b/interp/interpreter_test.go
index 6c92102..0be2409 100644
--- a/interp/interpreter_test.go
+++ b/interp/interpreter_test.go
@@ -218,8 +218,13 @@ func TestStruct(t *testing.T) {
 }
 
 func TestMap(t *testing.T) {
+	src0 := `type M map[string]bool;`
 	run(t, []etest{
-		{src: `type M map[string]bool; var m M; m`, res: `map[]`}, // #00
+		{src: src0 + `var m M; m`, res: `map[]`},                                     // #00
+		{src: `m := map[string]bool{"foo": true}; m["foo"]`, res: `true`},            // #01
+		{src: src0 + `m := M{"xx": true}; m`, res: `map[xx:true]`},                   // #02
+		{src: src0 + `var m = M{"xx": true}; m`, res: `map[xx:true]`},                // #03
+		{src: src0 + `var m = M{"xx": true}; m["xx"] = false`, res: `map[xx:false]`}, // #04
 	})
 }
 
diff --git a/lang/token.go b/lang/token.go
index 4f5da35..46a5bb4 100644
--- a/lang/token.go
+++ b/lang/token.go
@@ -59,6 +59,7 @@ const (
 	AndNotAssign // &^=
 	Inc          // ++
 	Dec          // --
+	MapAssign
 
 	// Unary operations.
 	Plus     // unary +
diff --git a/lang/token_string.go b/lang/token_string.go
index 44215b3..7f06092 100644
--- a/lang/token_string.go
+++ b/lang/token_string.go
@@ -51,63 +51,64 @@ func _() {
 	_ = x[AndNotAssign-40]
 	_ = x[Inc-41]
 	_ = x[Dec-42]
-	_ = x[Plus-43]
-	_ = x[Minus-44]
-	_ = x[Addr-45]
-	_ = x[Deref-46]
-	_ = x[BitComp-47]
-	_ = x[Arrow-48]
-	_ = x[Ellipsis-49]
-	_ = x[Not-50]
-	_ = x[Tilde-51]
-	_ = x[Comma-52]
-	_ = x[Semicolon-53]
-	_ = x[Colon-54]
-	_ = x[ParenBlock-55]
-	_ = x[BracketBlock-56]
-	_ = x[BraceBlock-57]
-	_ = x[Break-58]
-	_ = x[Case-59]
-	_ = x[Chan-60]
-	_ = x[Const-61]
-	_ = x[Continue-62]
-	_ = x[Default-63]
-	_ = x[Defer-64]
-	_ = x[Else-65]
-	_ = x[Fallthrough-66]
-	_ = x[For-67]
-	_ = x[Func-68]
-	_ = x[Go-69]
-	_ = x[Goto-70]
-	_ = x[If-71]
-	_ = x[Import-72]
-	_ = x[Interface-73]
-	_ = x[Map-74]
-	_ = x[Package-75]
-	_ = x[Range-76]
-	_ = x[Return-77]
-	_ = x[Select-78]
-	_ = x[Struct-79]
-	_ = x[Switch-80]
-	_ = x[Type-81]
-	_ = x[Var-82]
-	_ = x[Call-83]
-	_ = x[CallX-84]
-	_ = x[Composite-85]
-	_ = x[EqualSet-86]
-	_ = x[Grow-87]
-	_ = x[Index-88]
-	_ = x[JumpFalse-89]
-	_ = x[JumpSetFalse-90]
-	_ = x[JumpSetTrue-91]
-	_ = x[Label-92]
-	_ = x[New-93]
-	_ = x[MaxTok-94]
+	_ = x[MapAssign-43]
+	_ = x[Plus-44]
+	_ = x[Minus-45]
+	_ = x[Addr-46]
+	_ = x[Deref-47]
+	_ = x[BitComp-48]
+	_ = x[Arrow-49]
+	_ = x[Ellipsis-50]
+	_ = x[Not-51]
+	_ = x[Tilde-52]
+	_ = x[Comma-53]
+	_ = x[Semicolon-54]
+	_ = x[Colon-55]
+	_ = x[ParenBlock-56]
+	_ = x[BracketBlock-57]
+	_ = x[BraceBlock-58]
+	_ = x[Break-59]
+	_ = x[Case-60]
+	_ = x[Chan-61]
+	_ = x[Const-62]
+	_ = x[Continue-63]
+	_ = x[Default-64]
+	_ = x[Defer-65]
+	_ = x[Else-66]
+	_ = x[Fallthrough-67]
+	_ = x[For-68]
+	_ = x[Func-69]
+	_ = x[Go-70]
+	_ = x[Goto-71]
+	_ = x[If-72]
+	_ = x[Import-73]
+	_ = x[Interface-74]
+	_ = x[Map-75]
+	_ = x[Package-76]
+	_ = x[Range-77]
+	_ = x[Return-78]
+	_ = x[Select-79]
+	_ = x[Struct-80]
+	_ = x[Switch-81]
+	_ = x[Type-82]
+	_ = x[Var-83]
+	_ = x[Call-84]
+	_ = x[CallX-85]
+	_ = x[Composite-86]
+	_ = x[EqualSet-87]
+	_ = x[Grow-88]
+	_ = x[Index-89]
+	_ = x[JumpFalse-90]
+	_ = x[JumpSetFalse-91]
+	_ = x[JumpSetTrue-92]
+	_ = x[Label-93]
+	_ = x[New-94]
+	_ = x[MaxTok-95]
 }
 
-const _Token_name = "IllegalCommentIdentCharFloatImagIntStringAddSubMulQuoRemAndOrXorShlShrAndNotPeriodEqualGreaterGreaterEqualLandLessLessEqualLorNotEqualDefineAssignAddAssignSubAssignMulAssignQuoAssignRemAssignAndAssignOrAssignXorAssignShlAssignShrAssignAndNotAssignIncDecPlusMinusAddrDerefBitCompArrowEllipsisNotTildeCommaSemicolonColonParenBlockBracketBlockBraceBlockBreakCaseChanConstContinueDefaultDeferElseFallthroughForFuncGoGotoIfImportInterfaceMapPackageRangeReturnSelectStructSwitchTypeVarCallCallXCompositeEqualSetGrowIndexJumpFalseJumpSetFalseJumpSetTrueLabelNewMaxTok"
+const _Token_name = "IllegalCommentIdentCharFloatImagIntStringAddSubMulQuoRemAndOrXorShlShrAndNotPeriodEqualGreaterGreaterEqualLandLessLessEqualLorNotEqualDefineAssignAddAssignSubAssignMulAssignQuoAssignRemAssignAndAssignOrAssignXorAssignShlAssignShrAssignAndNotAssignIncDecMapAssignPlusMinusAddrDerefBitCompArrowEllipsisNotTildeCommaSemicolonColonParenBlockBracketBlockBraceBlockBreakCaseChanConstContinueDefaultDeferElseFallthroughForFuncGoGotoIfImportInterfaceMapPackageRangeReturnSelectStructSwitchTypeVarCallCallXCompositeEqualSetGrowIndexJumpFalseJumpSetFalseJumpSetTrueLabelNewMaxTok"
 
-var _Token_index = [...]uint16{0, 7, 14, 19, 23, 28, 32, 35, 41, 44, 47, 50, 53, 56, 59, 61, 64, 67, 70, 76, 82, 87, 94, 106, 110, 114, 123, 126, 134, 140, 146, 155, 164, 173, 182, 191, 200, 208, 217, 226, 235, 247, 250, 253, 257, 262, 266, 271, 278, 283, 291, 294, 299, 304, 313, 318, 328, 340, 350, 355, 359, 363, 368, 376, 383, 388, 392, 403, 406, 410, 412, 416, 418, 424, 433, 436, 443, 448, 454, 460, 466, 472, 476, 479, 483, 488, 497, 505, 509, 514, 523, 535, 546, 551, 554, 560}
+var _Token_index = [...]uint16{0, 7, 14, 19, 23, 28, 32, 35, 41, 44, 47, 50, 53, 56, 59, 61, 64, 67, 70, 76, 82, 87, 94, 106, 110, 114, 123, 126, 134, 140, 146, 155, 164, 173, 182, 191, 200, 208, 217, 226, 235, 247, 250, 253, 262, 266, 271, 275, 280, 287, 292, 300, 303, 308, 313, 322, 327, 337, 349, 359, 364, 368, 372, 377, 385, 392, 397, 401, 412, 415, 419, 421, 425, 427, 433, 442, 445, 452, 457, 463, 469, 475, 481, 485, 488, 492, 497, 506, 514, 518, 523, 532, 544, 555, 560, 563, 569}
 
 func (i Token) String() string {
 	idx := int(i) - 0
diff --git a/parser/README.md b/parser/README.md
index 9f40e77..ccf580d 100644
--- a/parser/README.md
+++ b/parser/README.md
@@ -59,6 +59,7 @@ Go language support:
 - [x] iota expression
 - [ ] defer statement
 - [ ] recover statement
+- [ ] range clause
 - [ ] go statement
 - [x] if statement (including else and else if)
 - [x] for statement
diff --git a/parser/expr.go b/parser/expr.go
index 0a87999..35e54ff 100644
--- a/parser/expr.go
+++ b/parser/expr.go
@@ -145,7 +145,14 @@ func (p *Parser) parseExpr(in Tokens, typeStr string) (out Tokens, err error) {
 				return out, err
 			}
 			out = append(out, toks...)
-			ops = append(ops, scanner.Token{Tok: lang.Index, Pos: t.Pos})
+			if i < len(in)-2 && in[i+1].Tok == lang.Assign {
+				// A brace block followed by assign implies a MapAssing token,
+				// as assignement to a map element cannot be implemented through a normal Assign.
+				ops = append(ops, scanner.Token{Tok: lang.MapAssign, Pos: t.Pos})
+				i++
+			} else {
+				ops = append(ops, scanner.Token{Tok: lang.Index, Pos: t.Pos})
+			}
 
 		case lang.Struct:
 			typ, err := p.parseTypeExpr(in[i : i+2])
diff --git a/parser/parse.go b/parser/parse.go
index c049b87..32ee0e3 100644
--- a/parser/parse.go
+++ b/parser/parse.go
@@ -59,7 +59,7 @@ func (p *Parser) Scan(s string, endSemi bool) (Tokens, error) {
 	return p.Scanner.Scan(s, endSemi)
 }
 
-// Parse performs syntax analysis on s and return Tokens or an error.
+// Parse performs syntax analysis on s and returns Tokens or an error.
 func (p *Parser) Parse(src string) (out Tokens, err error) {
 	in, err := p.Scan(src, true)
 	if err != nil {
diff --git a/vm/op_string.go b/vm/op_string.go
index e5a00b6..e838628 100644
--- a/vm/op_string.go
+++ b/vm/op_string.go
@@ -31,29 +31,30 @@ func _() {
 	_ = x[Grow-20]
 	_ = x[Index-21]
 	_ = x[IndexSet-22]
-	_ = x[MapSet-23]
-	_ = x[Jump-24]
-	_ = x[JumpTrue-25]
-	_ = x[JumpFalse-26]
-	_ = x[JumpSetTrue-27]
-	_ = x[JumpSetFalse-28]
-	_ = x[Lower-29]
-	_ = x[Loweri-30]
-	_ = x[Mul-31]
-	_ = x[New-32]
-	_ = x[Negate-33]
-	_ = x[Not-34]
-	_ = x[Pop-35]
-	_ = x[Push-36]
-	_ = x[Return-37]
-	_ = x[Sub-38]
-	_ = x[Subi-39]
-	_ = x[Swap-40]
+	_ = x[MapIndex-23]
+	_ = x[MapSet-24]
+	_ = x[Jump-25]
+	_ = x[JumpTrue-26]
+	_ = x[JumpFalse-27]
+	_ = x[JumpSetTrue-28]
+	_ = x[JumpSetFalse-29]
+	_ = x[Lower-30]
+	_ = x[Loweri-31]
+	_ = x[Mul-32]
+	_ = x[New-33]
+	_ = x[Negate-34]
+	_ = x[Not-35]
+	_ = x[Pop-36]
+	_ = x[Push-37]
+	_ = x[Return-38]
+	_ = x[Sub-39]
+	_ = x[Subi-40]
+	_ = x[Swap-41]
 }
 
-const _Op_name = "NopAddAddrAssignFassignVassignCallCalliCallXDerefDupFdupFnewEqualEqualSetExitFieldFieldSetFieldFsetGreaterGrowIndexIndexSetMapSetJumpJumpTrueJumpFalseJumpSetTrueJumpSetFalseLowerLoweriMulNewNegateNotPopPushReturnSubSubiSwap"
+const _Op_name = "NopAddAddrAssignFassignVassignCallCalliCallXDerefDupFdupFnewEqualEqualSetExitFieldFieldSetFieldFsetGreaterGrowIndexIndexSetMapIndexMapSetJumpJumpTrueJumpFalseJumpSetTrueJumpSetFalseLowerLoweriMulNewNegateNotPopPushReturnSubSubiSwap"
 
-var _Op_index = [...]uint8{0, 3, 6, 10, 16, 23, 30, 34, 39, 44, 49, 52, 56, 60, 65, 73, 77, 82, 90, 99, 106, 110, 115, 123, 129, 133, 141, 150, 161, 173, 178, 184, 187, 190, 196, 199, 202, 206, 212, 215, 219, 223}
+var _Op_index = [...]uint8{0, 3, 6, 10, 16, 23, 30, 34, 39, 44, 49, 52, 56, 60, 65, 73, 77, 82, 90, 99, 106, 110, 115, 123, 131, 137, 141, 149, 158, 169, 181, 186, 192, 195, 198, 204, 207, 210, 214, 220, 223, 227, 231}
 
 func (i Op) String() string {
 	idx := int(i) - 0
diff --git a/vm/vm.go b/vm/vm.go
index 709b459..2baf384 100644
--- a/vm/vm.go
+++ b/vm/vm.go
@@ -42,6 +42,7 @@ const (
 	Grow                   // -- ; sp += $1
 	Index                  // a i -- a[i] ;
 	IndexSet               // a i v -- a; a[i] = v
+	MapIndex               // a i -- a[i]
 	MapSet                 // a i v -- a; a[i] = v
 	Jump                   // -- ; ip += $1
 	JumpTrue               // cond -- ; if cond { ip += $1 }
@@ -267,6 +268,9 @@ func (m *Machine) Run() (err error) {
 		case IndexSet:
 			mem[sp-3].Value.Index(int(mem[sp-2].Int())).Set(mem[sp-1].Value)
 			mem = mem[:sp-2]
+		case MapIndex:
+			mem[sp-2].Value = mem[sp-2].MapIndex(mem[sp-1].Value)
+			mem = mem[:sp-1]
 		case MapSet:
 			mem[sp-3].SetMapIndex(mem[sp-2].Value, mem[sp-1].Value)
 			mem = mem[:sp-2]