Repository: goplus/gop
Pull request #944 (2021-11-26 02:49:28)

cl: support lambda assign


Patch:
diff --git a/cl/compile.go b/cl/compile.go
index d89cdefc..65071302 100644
--- a/cl/compile.go
+++ b/cl/compile.go
@@ -899,7 +899,17 @@ func loadVars(ctx *blockCtx, v *ast.ValueSpec, global bool) {
 			compileExpr(ctx, v.Values[0], true)
 		} else {
 			for _, val := range v.Values {
-				compileExpr(ctx, val)
+				switch e := val.(type) {
+				case *ast.LambdaExpr, *ast.LambdaExpr2:
+					if len(v.Values) == 1 {
+						sig := checkLambdaFuncType(ctx, e, typ, clLambaAssign, v.Names[0])
+						compileLambda(ctx, e, sig)
+					} else {
+						panic(ctx.newCodeErrorf(e.Pos(), "lambda unsupport multiple assignment"))
+					}
+				default:
+					compileExpr(ctx, val)
+				}
 			}
 		}
 		cb.EndInit(nv)
diff --git a/cl/expr.go b/cl/expr.go
index 3ae52203..67551bce 100644
--- a/cl/expr.go
+++ b/cl/expr.go
@@ -434,11 +434,11 @@ func compileCallExpr(ctx *blockCtx, v *ast.CallExpr, flags int) {
 		switch expr := arg.(type) {
 		case *ast.LambdaExpr:
 			fn.initWith(fnt, i, len(expr.Lhs))
-			sig := checkLambdaArgumentType(ctx, expr, v.Fun, fn.arg(i, true), len(expr.Rhs))
+			sig := checkLambdaFuncType(ctx, expr, fn.arg(i, true), clLambaArgument, v.Fun)
 			compileLambdaExpr(ctx, expr, sig)
 		case *ast.LambdaExpr2:
 			fn.initWith(fnt, i, len(expr.Lhs))
-			sig := checkLambdaArgumentType(ctx, expr, v.Fun, fn.arg(i, true), -1)
+			sig := checkLambdaFuncType(ctx, expr, fn.arg(i, true), clLambaArgument, v.Fun)
 			compileLambdaExpr2(ctx, expr, sig)
 		case *ast.CompositeLit:
 			fn.initWith(fnt, i, -1)
@@ -450,41 +450,42 @@ func compileCallExpr(ctx *blockCtx, v *ast.CallExpr, flags int) {
 	ctx.cb.CallWith(len(v.Args), ellipsis, v)
 }
 
-// check lambda func argument type, check results len if chkResults >= 0.
-func checkLambdaArgumentType(ctx *blockCtx, lambda ast.Expr, fun ast.Expr, argTyp types.Type, chkResults int) *types.Signature {
-	typ := argTyp
+type clLambaFlag string
+
+const (
+	clLambaAssign   clLambaFlag = "assignment"
+	clLambaField    clLambaFlag = "field value"
+	clLambaArgument clLambaFlag = "argument"
+)
+
+// check lambda func type
+func checkLambdaFuncType(ctx *blockCtx, lambda ast.Expr, ftyp types.Type, flag clLambaFlag, toNode ast.Node) *types.Signature {
+	typ := ftyp
 retry:
 	switch t := typ.(type) {
 	case *types.Signature:
-		if chkResults < 0 || chkResults == t.Results().Len() {
-			return t
+		if l, ok := lambda.(*ast.LambdaExpr); ok {
+			if len(l.Rhs) != t.Results().Len() {
+				break
+			}
 		}
+		return t
 	case *types.Named:
 		typ = t.Underlying()
 		goto retry
 	}
-	pos := ctx.Position(lambda.Pos())
-	src, _ := ctx.LoadExpr(fun)
-	err := newCodeErrorf(&pos, "cannot use lambda literal as type %v in argument to %v", argTyp, src)
+	src, _ := ctx.LoadExpr(toNode)
+	err := ctx.newCodeErrorf(lambda.Pos(), "cannot use lambda literal as type %v in %v to %v", ftyp, flag, src)
 	panic(err)
 }
 
-// check lambda field type, check results len if chkResults >= 0.
-func checkLambdaFieldType(ctx *blockCtx, lambda ast.Expr, fieldType types.Type, chkResults int) *types.Signature {
-	typ := fieldType
-retry:
-	switch t := typ.(type) {
-	case *types.Signature:
-		if chkResults < 0 || chkResults == t.Results().Len() {
-			return t
-		}
-	case *types.Named:
-		typ = t.Underlying()
-		goto retry
+func compileLambda(ctx *blockCtx, lambda ast.Expr, sig *types.Signature) {
+	switch expr := lambda.(type) {
+	case *ast.LambdaExpr:
+		compileLambdaExpr(ctx, expr, sig)
+	case *ast.LambdaExpr2:
+		compileLambdaExpr2(ctx, expr, sig)
 	}
-	pos := ctx.Position(lambda.Pos())
-	err := newCodeErrorf(&pos, "cannot use lambda literal as type %v in field value", fieldType)
-	panic(err)
 }
 
 func compileLambdaParams(ctx *blockCtx, pos token.Pos, lhs []*ast.Ident, in *types.Tuple) []*types.Var {
@@ -613,12 +614,9 @@ func compileStructLitInKeyVal(ctx *blockCtx, elts []ast.Expr, t *types.Struct, t
 			log.Panicln("TODO: struct member not found -", name)
 		}
 		switch expr := kv.Value.(type) {
-		case *ast.LambdaExpr:
-			sig := checkLambdaFieldType(ctx, expr, t.Field(idx).Type(), len(expr.Rhs))
-			compileLambdaExpr(ctx, expr, sig)
-		case *ast.LambdaExpr2:
-			sig := checkLambdaFieldType(ctx, expr, t.Field(idx).Type(), -1)
-			compileLambdaExpr2(ctx, expr, sig)
+		case *ast.LambdaExpr, *ast.LambdaExpr2:
+			sig := checkLambdaFuncType(ctx, expr, t.Field(idx).Type(), clLambaField, kv.Key)
+			compileLambda(ctx, expr, sig)
 		default:
 			compileExpr(ctx, kv.Value)
 		}
diff --git a/cl/stmt.go b/cl/stmt.go
index d79f5c7c..845f4498 100644
--- a/cl/stmt.go
+++ b/cl/stmt.go
@@ -212,7 +212,18 @@ func compileAssignStmt(ctx *blockCtx, expr *ast.AssignStmt) {
 		compileExprLHS(ctx, lhs)
 	}
 	for _, rhs := range expr.Rhs {
-		compileExpr(ctx, rhs, twoValue)
+		switch e := rhs.(type) {
+		case *ast.LambdaExpr, *ast.LambdaExpr2:
+			if len(expr.Lhs) == 1 && len(expr.Rhs) == 1 {
+				typ := ctx.cb.Get(-1).Type.(interface{ Elem() types.Type }).Elem()
+				sig := checkLambdaFuncType(ctx, e, typ, clLambaAssign, expr.Lhs[0])
+				compileLambda(ctx, e, sig)
+			} else {
+				panic(ctx.newCodeErrorf(e.Pos(), "lambda unsupport multiple assignment"))
+			}
+		default:
+			compileExpr(ctx, rhs, twoValue)
+		}
 	}
 	if tok == token.ASSIGN {
 		ctx.cb.AssignWith(len(expr.Lhs), len(expr.Rhs), expr)


Test patch:
diff --git a/cl/compile_test.go b/cl/compile_test.go
index 78413739..299a24c4 100644
--- a/cl/compile_test.go
+++ b/cl/compile_test.go
@@ -2792,26 +2792,41 @@ func main() {
 	gopClTest(t, `
 type Fn func() (int, error)
 func Do(fn Fn) {
-	v, err := fn()
-	println(v, err)
 }
 
 Do => (100, nil)
 `, `package main
 
-import fmt "fmt"
-
 type Fn func() (int, error)
 
 func Do(fn Fn) {
-	v, err := fn()
-	fmt.Println(v, err)
 }
 func main() {
 	Do(func() (int, error) {
 		return 100, nil
 	})
 }
+`)
+	gopClTest(t, `
+var fn func(int) (int,error) = x => (x*x, nil)
+`, `package main
+
+var fn func(int) (int, error) = func(x int) (int, error) {
+	return x * x, nil
+}
+`)
+	gopClTest(t, `
+var fn func(int) (int,error)
+fn = x => (x*x, nil)
+`, `package main
+
+var fn func(int) (int, error)
+
+func main() {
+	fn = func(x int) (int, error) {
+		return x * x, nil
+	}
+}
 `)
 }
 
@@ -2838,8 +2853,7 @@ func main() {
 `)
 	gopClTest(t, `
 func Do(fn func() (int, error)) {
-	v, err := fn()
-	println(v, err)
+	// ...
 }
 
 Do => {
@@ -2847,11 +2861,7 @@ Do => {
 }
 `, `package main
 
-import fmt "fmt"
-
 func Do(fn func() (int, error)) {
-	v, err := fn()
-	fmt.Println(v, err)
 }
 func main() {
 	Do(func() (int, error) {
@@ -2904,8 +2914,6 @@ func main() {
 	gopClTest(t, `
 type Fn func() (int, error)
 func Do(fn Fn) {
-	v, err := fn()
-	println(v, err)
 }
 
 Do => {
@@ -2913,19 +2921,40 @@ Do => {
 }
 `, `package main
 
-import fmt "fmt"
-
 type Fn func() (int, error)
 
 func Do(fn Fn) {
-	v, err := fn()
-	fmt.Println(v, err)
 }
 func main() {
 	Do(func() (int, error) {
 		return 100, nil
 	})
 }
+`)
+	gopClTest(t, `
+var fn func(int) (int,error) = x => {
+	return x * x, nil
+}
+`, `package main
+
+var fn func(int) (int, error) = func(x int) (int, error) {
+	return x * x, nil
+}
+`)
+	gopClTest(t, `
+var fn func(int) (int,error)
+fn = x => {
+	return x * x, nil
+}
+`, `package main
+
+var fn func(int) (int, error)
+
+func main() {
+	fn = func(x int) (int, error) {
+		return x * x, nil
+	}
+}
 `)
 }
 
diff --git a/cl/error_msg_test.go b/cl/error_msg_test.go
index 6f408652..86138ac3 100644
--- a/cl/error_msg_test.go
+++ b/cl/error_msg_test.go
@@ -72,7 +72,7 @@ func main() {
 	foo((x, y, z) => {})
 }
 `)
-	codeErrorTest(t, "./bar.gop:6:8: cannot use lambda literal as type int in field value", `
+	codeErrorTest(t, "./bar.gop:6:8: cannot use lambda literal as type int in field value to Plot", `
 type Foo struct {
 	Plot int
 }
@@ -80,7 +80,7 @@ foo := &Foo{
 	Plot: x => (x * 2, x * x),
 }
 `)
-	codeErrorTest(t, "./bar.gop:6:8: cannot use lambda literal as type int in field value", `
+	codeErrorTest(t, "./bar.gop:6:8: cannot use lambda literal as type int in field value to Plot", `
 type Foo struct {
 	Plot int
 }
@@ -103,13 +103,32 @@ func foo(func()) {
 foo => (100)
 `)
 	codeErrorTest(t,
-		"./bar.gop:6:8: cannot use lambda literal as type func() int in field value", `
+		"./bar.gop:6:8: cannot use lambda literal as type func() int in field value to Plot", `
 type Foo struct {
 	Plot func() int
 }
 foo := &Foo{
 	Plot: x => (x * 2, x * x),
 }
+`)
+
+	codeErrorTest(t,
+		"./bar.gop:2:18: cannot use lambda literal as type func() in assignment to foo", `
+var foo func() = => (100)
+`)
+	codeErrorTest(t,
+		"./bar.gop:3:7: cannot use lambda literal as type func() in assignment to foo", `
+var foo func()
+foo = => (100)
+`)
+	codeErrorTest(t,
+		"./bar.gop:2:29: lambda unsupport multiple assignment", `
+var foo, foo1 func() = nil, => {}
+`)
+	codeErrorTest(t,
+		"./bar.gop:3:15: lambda unsupport multiple assignment", `
+var foo func()
+_, foo = nil, => {}
 `)
 }