From 8261fc2a8d42dfdc97df2ec26cad5dc1ead0293e Mon Sep 17 00:00:00 2001
From: visualfc <visualfc@gmail.com>
Date: Fri, 23 Feb 2024 12:09:46 +0800
Subject: [PATCH] cl: types record ast.ForPhrase

---
 cl/compile.go            |  3 +++
 cl/expr.go               |  3 +++
 x/typesutil/gopinfo.go   |  3 +++
 x/typesutil/info_test.go | 20 ++++++++++++++++++++
 4 files changed, 29 insertions(+)

diff --git a/cl/compile.go b/cl/compile.go
index c22a7735f..f342c276e 100644
--- a/cl/compile.go
+++ b/cl/compile.go
@@ -158,6 +158,9 @@ type Recorder interface {
 	//     *ast.ForStmt
 	//     *ast.RangeStmt
 	//     *ast.ForPhraseStmt
+	//     *ast.ForPhrase
+	//     *ast.LambdaExpr
+	//     *ast.LambdaExpr2
 	//
 	Scope(ast.Node, *types.Scope)
 }
diff --git a/cl/expr.go b/cl/expr.go
index 3d66cd8ef..f88da3f28 100644
--- a/cl/expr.go
+++ b/cl/expr.go
@@ -1263,6 +1263,9 @@ func compileComprehensionExpr(ctx *blockCtx, v *ast.ComprehensionExpr, twoValue
 		compileExpr(ctx, forStmt.X)
 		cb.RangeAssignThen(forStmt.TokPos)
 		defNames(ctx, defineNames, cb.Scope())
+		if rec := ctx.recorder(); rec != nil {
+			rec.Scope(forStmt, cb.Scope())
+		}
 		if forStmt.Cond != nil {
 			cb.If()
 			if forStmt.Init != nil {
diff --git a/x/typesutil/gopinfo.go b/x/typesutil/gopinfo.go
index c3b2e8961..c5c04f139 100644
--- a/x/typesutil/gopinfo.go
+++ b/x/typesutil/gopinfo.go
@@ -122,6 +122,9 @@ type Info struct {
 	//     *ast.ForStmt
 	//     *ast.RangeStmt
 	//     *ast.ForPhraseStmt
+	//     *ast.ForPhrase
+	//     *ast.LambdaExpr
+	//     *ast.LambdaExpr2
 	//
 	Scopes map[ast.Node]*types.Scope
 
diff --git a/x/typesutil/info_test.go b/x/typesutil/info_test.go
index 87e083ba1..8fd27d792 100644
--- a/x/typesutil/info_test.go
+++ b/x/typesutil/info_test.go
@@ -1758,6 +1758,24 @@ func TestScopesInfo(t *testing.T) {
 		{`package p25; func test(fn func(int)int){};func _(){test( x => { y := x*x; return y } ) }`, []string{
 			"file:test", "func:fn", "func:", "lambda:x y",
 		}},
+		{`package p26; func _(){ b := {for x <- ["1", "3", "5", "7", "11"], x == "5"}; _ = b }`, []string{
+			"file:", "func:b", "for phrase:x",
+		}},
+		{`package p27; func _(){ b, ok := {i for i, x <- ["1", "3", "5", "7", "11"], x == "5"}; _ = b; _ = ok }`, []string{
+			"file:", "func:b ok", "for phrase:i x",
+		}},
+		{`package p28; func _(){ a := [x*x for x <- [1, 3.4, 5] if x > 2 ]; _ = a }`, []string{
+			"file:", "func:a", "for phrase:x",
+		}},
+		{`package p29; func _(){ arr := [1, 2, 3, 4.1, 5, 6];x := [[a, b] for a <- arr, a < b for b <- arr, b > 2]; _ = x }`, []string{
+			"file:", "func:arr x", "for phrase:b", "for phrase:a",
+		}},
+		{`package p30; func _(){ y := {x: i for i, x <- ["1", "3", "5", "7", "11"]}; _ = y }`, []string{
+			"file:", "func:y", "for phrase:i x",
+		}},
+		{`package p31; func _(){ z := {v: k for k, v <- {"Hello": 1, "Hi": 3, "xsw": 5, "Go+": 7}, v > 3}; _ = z }`, []string{
+			"file:", "func:z", "for phrase:k v",
+		}},
 	}
 
 	for _, test := range tests {
@@ -1802,6 +1820,8 @@ func TestScopesInfo(t *testing.T) {
 				kind = "lambda"
 			case *ast.LambdaExpr2:
 				kind = "lambda"
+			case *ast.ForPhrase:
+				kind = "for phrase"
 			default:
 				kind = fmt.Sprintf("<unknown node kind> %T", node)
 			}