cmd/compile/internal/types2: record types for union subexpressions

This is a port of CL 371757 from go/types to types2, with
minor adjustments for different error handling and AST.

It also names the added API test cases more consistently.
The same renaming was applied to the respective go/types
file.

Updates #50093

Change-Id: Iaa132106a197a207f831525432e62e9d452b17c9
Reviewed-on: https://go-review.googlesource.com/c/go/+/372475
Trust: Robert Griesemer <gri@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>

3 files changed, 53 insertions, 23 deletions

diff --git a/src/cmd/compile/internal/types2/api.go b/src/cmd/compile/internal/types2/api.go
index ed5bced643..4ea3989c39 100644
--- a/src/cmd/compile/internal/types2/api.go
+++ b/src/cmd/compile/internal/types2/api.go
@@ -202,12 +202,6 @@ type Info struct {
 	// identifier z in a variable declaration 'var z int' is found
 	// only in the Defs map, and identifiers denoting packages in
 	// qualified identifiers are collected in the Uses map.
-	//
-	// For binary expressions representing unions in constraint
-	// position or type elements in interfaces, a union type is
-	// recorded for the top-level expression only. For instance,
-	// given the constraint a|b|c, the union type for (a|b)|c
-	// is recorded, but not the union type for a|b.
 	Types map[syntax.Expr]TypeAndValue
 
 	// Instances maps identifiers denoting parameterized types or functions to
diff --git a/src/cmd/compile/internal/types2/api_test.go b/src/cmd/compile/internal/types2/api_test.go
index fc8b5cd4ee..dee7ffbaf7 100644
--- a/src/cmd/compile/internal/types2/api_test.go
+++ b/src/cmd/compile/internal/types2/api_test.go
@@ -348,10 +348,25 @@ func TestTypesInfo(t *testing.T) {
 		{`package u1a; func _[_ interface{~int}]() {}`, `~int`, `~int`},
 		{`package u2a; func _[_ interface{int|string}]() {}`, `int | string`, `int|string`},
 		{`package u3a; func _[_ interface{int|string|~bool}]() {}`, `int | string | ~bool`, `int|string|~bool`},
+		{`package u3a; func _[_ interface{int|string|~bool}]() {}`, `int | string`, `int|string`},
+		{`package u3a; func _[_ interface{int|string|~bool}]() {}`, `~bool`, `~bool`},
+		{`package u3a; func _[_ interface{int|string|~float64|~bool}]() {}`, `int | string | ~float64`, `int|string|~float64`},
+
 		{`package u0b; func _[_ int]() {}`, `int`, `int`},
 		{`package u1b; func _[_ ~int]() {}`, `~int`, `~int`},
 		{`package u2b; func _[_ int|string]() {}`, `int | string`, `int|string`},
 		{`package u3b; func _[_ int|string|~bool]() {}`, `int | string | ~bool`, `int|string|~bool`},
+		{`package u3b; func _[_ int|string|~bool]() {}`, `int | string`, `int|string`},
+		{`package u3b; func _[_ int|string|~bool]() {}`, `~bool`, `~bool`},
+		{`package u3b; func _[_ int|string|~float64|~bool]() {}`, `int | string | ~float64`, `int|string|~float64`},
+
+		{`package u0c; type _ interface{int}`, `int`, `int`},
+		{`package u1c; type _ interface{~int}`, `~int`, `~int`},
+		{`package u2c; type _ interface{int|string}`, `int | string`, `int|string`},
+		{`package u3c; type _ interface{int|string|~bool}`, `int | string | ~bool`, `int|string|~bool`},
+		{`package u3c; type _ interface{int|string|~bool}`, `int | string`, `int|string`},
+		{`package u3c; type _ interface{int|string|~bool}`, `~bool`, `~bool`},
+		{`package u3c; type _ interface{int|string|~float64|~bool}`, `int | string | ~float64`, `int|string|~float64`},
 	}
 
 	for _, test := range tests {
diff --git a/src/cmd/compile/internal/types2/union.go b/src/cmd/compile/internal/types2/union.go
index 98dd6cedc7..6f66260af4 100644
--- a/src/cmd/compile/internal/types2/union.go
+++ b/src/cmd/compile/internal/types2/union.go
@@ -48,23 +48,37 @@ const maxTermCount = 100
 // parseUnion parses uexpr as a union of expressions.
 // The result is a Union type, or Typ[Invalid] for some errors.
 func parseUnion(check *Checker, uexpr syntax.Expr) Type {
-	tlist := flattenUnion(nil, uexpr)
+	blist, tlist := flattenUnion(nil, uexpr)
+	assert(len(blist) == len(tlist)-1)
 
 	var terms []*Term
-	for _, x := range tlist {
-		tilde, typ := parseTilde(check, x)
-		if len(tlist) == 1 && !tilde {
+
+	var u Type
+	for i, x := range tlist {
+		term := parseTilde(check, x)
+		if len(tlist) == 1 && !term.tilde {
 			// Single type. Ok to return early because all relevant
 			// checks have been performed in parseTilde (no need to
 			// run through term validity check below).
-			return typ // typ already recorded through check.typ in parseTilde
+			return term.typ // typ already recorded through check.typ in parseTilde
 		}
 		if len(terms) >= maxTermCount {
-			check.errorf(x, "cannot handle more than %d union terms (implementation limitation)", maxTermCount)
-			check.recordTypeAndValue(uexpr, typexpr, Typ[Invalid], nil)
-			return Typ[Invalid]
+			if u != Typ[Invalid] {
+				check.errorf(x, "cannot handle more than %d union terms (implementation limitation)", maxTermCount)
+				u = Typ[Invalid]
+			}
+		} else {
+			terms = append(terms, term)
+			u = &Union{terms}
+		}
+
+		if i > 0 {
+			check.recordTypeAndValue(blist[i-1], typexpr, u, nil)
 		}
-		terms = append(terms, NewTerm(tilde, typ))
+	}
+
+	if u == Typ[Invalid] {
+		return u
 	}
 
 	// Check validity of terms.
@@ -106,17 +120,17 @@ func parseUnion(check *Checker, uexpr syntax.Expr) Type {
 		}
 	})
 
-	u := &Union{terms}
-	check.recordTypeAndValue(uexpr, typexpr, u, nil)
 	return u
 }
 
-func parseTilde(check *Checker, x syntax.Expr) (tilde bool, typ Type) {
+func parseTilde(check *Checker, tx syntax.Expr) *Term {
+	x := tx
+	var tilde bool
 	if op, _ := x.(*syntax.Operation); op != nil && op.Op == syntax.Tilde {
 		x = op.X
 		tilde = true
 	}
-	typ = check.typ(x)
+	typ := check.typ(x)
 	// Embedding stand-alone type parameters is not permitted (issue #47127).
 	// We don't need this restriction anymore if we make the underlying type of a type
 	// parameter its constraint interface: if we embed a lone type parameter, we will
@@ -126,7 +140,11 @@ func parseTilde(check *Checker, x syntax.Expr) (tilde bool, typ Type) {
 		check.error(x, "cannot embed a type parameter")
 		typ = Typ[Invalid]
 	}
-	return
+	term := NewTerm(tilde, typ)
+	if tilde {
+		check.recordTypeAndValue(tx, typexpr, &Union{[]*Term{term}}, nil)
+	}
+	return term
 }
 
 // overlappingTerm reports the index of the term x in terms which is
@@ -147,10 +165,13 @@ func overlappingTerm(terms []*Term, y *Term) int {
 	return -1
 }
 
-func flattenUnion(list []syntax.Expr, x syntax.Expr) []syntax.Expr {
+// flattenUnion walks a union type expression of the form A | B | C | ...,
+// extracting both the binary exprs (blist) and leaf types (tlist).
+func flattenUnion(list []syntax.Expr, x syntax.Expr) (blist, tlist []syntax.Expr) {
 	if o, _ := x.(*syntax.Operation); o != nil && o.Op == syntax.Or {
-		list = flattenUnion(list, o.X)
+		blist, tlist = flattenUnion(list, o.X)
+		blist = append(blist, o)
 		x = o.Y
 	}
-	return append(list, x)
+	return blist, append(tlist, x)
 }