Repository: goplus/gop
Pull request #589 (2020-09-11 01:15:31)

Map two value assign
#495


Patch:
diff --git a/ast/ast_ql.go b/ast/ast_ql.go
index 7c515552..33eb035b 100644
--- a/ast/ast_ql.go
+++ b/ast/ast_ql.go
@@ -426,3 +426,10 @@ func (p *ForPhraseStmt) End() token.Pos {
 func (*ForPhraseStmt) stmtNode() {}
 
 // -----------------------------------------------------------------------------
+
+// A TwoValueIndexExpr node represents a two-value assignment expression (v, ok := m["key"])
+type TwoValueIndexExpr struct {
+	*IndexExpr
+}
+
+// -----------------------------------------------------------------------------
diff --git a/cl/expr.go b/cl/expr.go
index 54395f88..f4202416 100644
--- a/cl/expr.go
+++ b/cl/expr.go
@@ -72,7 +72,9 @@ func compileExpr(ctx *blockCtx, expr ast.Expr) func() {
 	case *ast.ErrWrapExpr:
 		return compileErrWrapExpr(ctx, v)
 	case *ast.IndexExpr:
-		return compileIndexExpr(ctx, v)
+		return compileIndexExpr(ctx, v, false)
+	case *ast.TwoValueIndexExpr:
+		return compileIndexExpr(ctx, v.IndexExpr, true)
 	case *ast.SliceExpr:
 		return compileSliceExpr(ctx, v)
 	case *ast.CompositeLit:
@@ -881,7 +883,7 @@ func compileIndexExprLHS(ctx *blockCtx, v *ast.IndexExpr, mode compileMode) {
 			logIllTypeMapIndexPanic(ctx, v, t, typIdx)
 		}
 		if ctx.indirect {
-			ctx.out.MapIndex().AddrOp(kindOf(typElem), exec.OpAssign)
+			ctx.out.MapIndex(false).AddrOp(kindOf(typElem), exec.OpAssign)
 		} else {
 			ctx.out.SetMapIndex()
 		}
@@ -946,7 +948,7 @@ func compileIdx(ctx *blockCtx, v ast.Expr, nlast int, kind reflect.Kind) int {
 	return -1
 }
 
-func compileIndexExpr(ctx *blockCtx, v *ast.IndexExpr) func() { // x[i]
+func compileIndexExpr(ctx *blockCtx, v *ast.IndexExpr, twoValue bool) func() { // x[i]
 	var kind reflect.Kind
 	var typElem reflect.Type
 	exprX := compileExpr(ctx, v.X)
@@ -999,7 +1001,10 @@ func compileIndexExpr(ctx *blockCtx, v *ast.IndexExpr) func() { // x[i]
 			if t := i.(iValue).Type(); t != typIdx {
 				logIllTypeMapIndexPanic(ctx, v, t, typIdx)
 			}
-			ctx.out.MapIndex()
+			if twoValue {
+				ctx.infer.Push(&goValue{t: exec.TyBool})
+			}
+			ctx.out.MapIndex(twoValue)
 		default:
 			log.Panicln("compileIndexExpr: unknown -", typ)
 		}
diff --git a/cl/stmt.go b/cl/stmt.go
index 4cfd81f1..c75bc038 100644
--- a/cl/stmt.go
+++ b/cl/stmt.go
@@ -593,7 +593,11 @@ func compileAssignStmt(ctx *blockCtx, expr *ast.AssignStmt) {
 		log.Panicln("compileAssignStmt internal error: infer stack is not empty.")
 	}
 	if len(expr.Rhs) == 1 {
-		compileExpr(ctx, expr.Rhs[0])()
+		rhsExpr := expr.Rhs[0]
+		if ie, ok := rhsExpr.(*ast.IndexExpr); ok && len(expr.Lhs) == 2 {
+			rhsExpr = &ast.TwoValueIndexExpr{IndexExpr: ie}
+		}
+		compileExpr(ctx, rhsExpr)()
 		v := ctx.infer.Get(-1).(iValue)
 		n := v.NumValues()
 		if n != 1 {
diff --git a/exec.spec/interface.go b/exec.spec/interface.go
index 373e03c8..c2a4bed8 100644
--- a/exec.spec/interface.go
+++ b/exec.spec/interface.go
@@ -336,7 +336,7 @@ type Builder interface {
 	Append(typ reflect.Type, arity int) Builder
 
 	// MapIndex instr
-	MapIndex() Builder
+	MapIndex(twoValue bool) Builder
 
 	// SetMapIndex instr
 	SetMapIndex() Builder
diff --git a/exec/bytecode/goinstr_lstcompr.go b/exec/bytecode/goinstr_lstcompr.go
index 76cf53be..dda65fe1 100644
--- a/exec/bytecode/goinstr_lstcompr.go
+++ b/exec/bytecode/goinstr_lstcompr.go
@@ -278,16 +278,20 @@ func execMapIndex(i Instr, p *Context) {
 	n := len(p.data)
 	key := reflect.ValueOf(p.data[n-1])
 	v := reflect.ValueOf(p.data[n-2])
-	switch i & bitsOperand {
+	switch op := i & bitsOperand; op {
 	case 1: // value mapData $key $setMapIndex
 		v.SetMapIndex(key, reflect.ValueOf(p.data[n-3]))
 		p.PopN(3)
 	default: // mapData $key $mapIndex
 		value := v.MapIndex(key)
-		if !value.IsValid() {
+		valid := value.IsValid()
+		if !valid {
 			value = reflect.Zero(v.Type().Elem())
 		}
 		p.Ret(2, value.Interface())
+		if op == 2 {
+			p.Push(valid)
+		}
 	}
 }
 
@@ -488,8 +492,12 @@ func (p *Builder) Make(typ reflect.Type, arity int) *Builder {
 }
 
 // MapIndex instr
-func (p *Builder) MapIndex() *Builder {
-	p.code.data = append(p.code.data, opMapIndex<<bitsOpShift)
+func (p *Builder) MapIndex(twoValue bool) *Builder {
+	op := Instr(0)
+	if twoValue {
+		op = 2
+	}
+	p.code.data = append(p.code.data, (opMapIndex<<bitsOpShift)|op)
 	return p
 }
 
diff --git a/exec/bytecode/interface.go b/exec/bytecode/interface.go
index 4f27c73a..88abc589 100644
--- a/exec/bytecode/interface.go
+++ b/exec/bytecode/interface.go
@@ -379,8 +379,8 @@ func (p *iBuilder) Make(typ reflect.Type, arity int) exec.Builder {
 }
 
 // MapIndex instr
-func (p *iBuilder) MapIndex() exec.Builder {
-	((*Builder)(p)).MapIndex()
+func (p *iBuilder) MapIndex(twoValue bool) exec.Builder {
+	((*Builder)(p)).MapIndex(twoValue)
 	return p
 }
 
diff --git a/exec/golang/interface.go b/exec/golang/interface.go
index 02e0d5fc..55357244 100644
--- a/exec/golang/interface.go
+++ b/exec/golang/interface.go
@@ -395,7 +395,7 @@ func (p *iBuilder) Make(typ reflect.Type, arity int) exec.Builder {
 }
 
 // MapIndex instr
-func (p *iBuilder) MapIndex() exec.Builder {
+func (p *iBuilder) MapIndex(twoValue bool) exec.Builder {
 	((*Builder)(p)).MapIndex()
 	return p
 }


Test patch:
diff --git a/cl/expr_test.go b/cl/expr_test.go
index 673a65fe..91c9b38d 100644
--- a/cl/expr_test.go
+++ b/cl/expr_test.go
@@ -1230,7 +1230,6 @@ func TestMethodCases(t *testing.T) {
 
 // -----------------------------------------------------------------------------
 
-// -----------------------------------------------------------------------------
 var testStarExprClauses = map[string]testData{
 	"star expr": {`
 				func A(a *int, c *struct {
@@ -1354,3 +1353,14 @@ func testScripts(t *testing.T, testName string, scripts map[string]testData) {
 }
 
 // -----------------------------------------------------------------------------
+
+func TestTwoValueExpr(t *testing.T) {
+	clause := `m:={2:3,1:2}
+			if v,ok:=m[m[1]];ok{
+				println(1,v,ok)
+			}
+			if v,ok:=m[m[3]];!ok{
+				println(3,v,ok)
+			}`
+	cltest.Expect(t, clause, "1 3 true\n3 0 false\n")
+}
diff --git a/exec/bytecode/goinstr_lstcompr_test.go b/exec/bytecode/goinstr_lstcompr_test.go
index c46f451f..3d492e1f 100644
--- a/exec/bytecode/goinstr_lstcompr_test.go
+++ b/exec/bytecode/goinstr_lstcompr_test.go
@@ -88,18 +88,21 @@ func TestMapIndex(t *testing.T) {
 		StoreVar(m).
 		LoadVar(m).
 		Push("go+").
-		MapIndex().
+		MapIndex(false).
 		LoadVar(m).
 		Push("xsw").
-		MapIndex().
+		MapIndex(true).
 		Resolve()
 
 	ctx := NewContext(code)
 	ctx.Exec(0, code.Len())
-	if v := ctx.Get(-1); v != 1.0 {
+	if v := ctx.Get(-1); v != true {
+		t.Fatal("{`Hello`: 3.2, `xsw`: 1.0}[`xsw`] != 1.0,true, ret =", v)
+	}
+	if v := ctx.Get(-2); v != 1.0 {
 		t.Fatal("{`Hello`: 3.2, `xsw`: 1.0}[`xsw`] != 1.0, ret =", v)
 	}
-	if v := ctx.Get(-2); v != 0.0 {
+	if v := ctx.Get(-3); v != 0.0 {
 		t.Fatal("{`Hello`: 3.2, `xsw`: 1.0}[`go+`] != 1.0, ret =", v)
 	}
 }
diff --git a/exec/bytecode/gopkg_test.go b/exec/bytecode/gopkg_test.go
index d4fbfc79..e8d2a185 100644
--- a/exec/bytecode/gopkg_test.go
+++ b/exec/bytecode/gopkg_test.go
@@ -354,13 +354,13 @@ func TestMapField(t *testing.T) {
 	code := b.
 		LoadGoVar(x).
 		Push(0).
-		MapIndex().
+		MapIndex(false).
 		LoadGoVar(x).
 		Push(5).
 		SetMapIndex(). // pkg.M[5] = pkg.M[0]
 		LoadGoVar(x).
 		Push(5).
-		MapIndex().
+		MapIndex(false).
 		LoadField(typ, []int{1}). // pkg.M[5]
 		Resolve()
 	ctx := NewContext(code)
@@ -396,7 +396,7 @@ func TestMapBadStoreField(t *testing.T) {
 		Push(-1).
 		LoadGoVar(x).
 		Push(0).
-		MapIndex().
+		MapIndex(false).
 		StoreField(typ, []int{1}). // pkg.M[0].Y = -1 , cannot assign
 		Resolve()
 
diff --git a/exec/golang/testdata/05.maplit/maplit.gop b/exec/golang/testdata/05.maplit/maplit.gop
index a72c6c3c..6b31e9a8 100644
--- a/exec/golang/testdata/05.maplit/maplit.gop
+++ b/exec/golang/testdata/05.maplit/maplit.gop
@@ -19,3 +19,12 @@ println({1: 1.4, 3: "Go+"})             // map[int]interface{}
 delete({1: 1.4, 3: "Go+"}, 1)           // delete map[int]interface{}
 
 println("empty map:", {}) // map[string]interface{}
+
+m := {2: 3, 1: 2}
+v, ok := m[m[1]]
+if ok {
+	println(1, v, ok)
+}
+if v, ok = m[m[3]]; !ok {
+	println(1, v, ok)
+}