cmd/compile/internal/noder: shape-based stenciling for unified IR

This CL switches unified IR to use shape-based stenciling with runtime
dictionaries, like the existing non-unified frontend. Specifically,
when instantiating generic functions and types `X[T]`, we now also
instantiated shaped variants `X[shapify(T)]` that can be shared by
`T`'s with common underlying types.

For example, for generic function `F`, `F[int](args...)` will be
rewritten to `F[go.shape.int](&.dict.F[int], args...)`.

For generic type `T` with method `M` and value `t` of type `T[int]`,
`t.M(args...)` will be rewritten to `T[go.shape.int].M(t,
&.dict.T[int], args...)`.

Two notable distinctions from the non-unified frontend:

1. For simplicity, currently shaping is limited to simply converting
type arguments to their underlying type. Subsequent CLs will implement
more aggressive shaping.

2. For generic types, a single dictionary is generated to be shared by
all methods, rather than separate dictionaries for each method. I
originally went with this design because I have an idea of changing
interface calls to pass the itab pointer via the closure
register (which should have zero overhead), and then the interface
wrappers for generic methods could use the *runtime.itab to find the
runtime dictionary that corresponds to the dynamic type. This would
allow emitting fewer method wrappers.

However, this choice does have the consequence that currently even if
a method is unused and its code is pruned by the linker, it may have
produced runtime dictionary entries that need to be kept alive anyway.

I'm open to changing this to generate per-method dictionaries, though
this would require changing the unified IR export data format; so it
would be best to make this decision before Go 1.20.

The other option is making the linker smarter about pruning unneeded
dictionary entries, like how it already prunes itab entries. For
example, the runtime dictionary for `T[int]` could have a `R_DICTTYPE`
meta-relocation against symbol `.dicttype.T[go.shape.int]` that
declares it's a dictionary associated with that type; and then each
method on `T[go.shape.T]` could have `R_DICTUSE` meta-relocations
against `.dicttype.T[go.shape.T]+offset` indicating which fields
within dictionaries of that type need to be preserved.

Change-Id: I369580b1d93d19640a4b5ecada4f6231adcce3fd
Reviewed-on: https://go-review.googlesource.com/c/go/+/421821
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Keith Randall <khr@golang.org>
Run-TryBot: Matthew Dempsky <mdempsky@google.com>
Reviewed-by: Cuong Manh Le <cuong.manhle.vn@gmail.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

1 files changed, 290 insertions, 49 deletions

diff --git a/src/cmd/compile/internal/noder/writer.go b/src/cmd/compile/internal/noder/writer.go
index 52fa542f6b..a90b2d3bbd 100644
--- a/src/cmd/compile/internal/noder/writer.go
+++ b/src/cmd/compile/internal/noder/writer.go
@@ -172,12 +172,37 @@ type writerDict struct {
 	// derivedIdx maps a Type to its corresponding index within the
 	// derived slice, if present.
 	derivedIdx map[types2.Type]pkgbits.Index
+
+	// These slices correspond to entries in the runtime dictionary.
+	typeParamMethodExprs []writerMethodExprInfo
+	subdicts             []objInfo
+	rtypes               []typeInfo
+	itabs                []itabInfo
+}
+
+type itabInfo struct {
+	typ   typeInfo
+	iface typeInfo
+}
+
+// typeParamIndex returns the index of the given type parameter within
+// the dictionary. This may differ from typ.Index() when there are
+// implicit type parameters due to defined types declared within a
+// generic function or method.
+func (dict *writerDict) typeParamIndex(typ *types2.TypeParam) int {
+	for idx, implicit := range dict.implicits {
+		if implicit.Type().(*types2.TypeParam) == typ {
+			return idx
+		}
+	}
+
+	return len(dict.implicits) + typ.Index()
 }
 
 // A derivedInfo represents a reference to an encoded generic Go type.
 type derivedInfo struct {
 	idx    pkgbits.Index
-	needed bool
+	needed bool // TODO(mdempsky): Remove.
 }
 
 // A typeInfo represents a reference to an encoded Go type.
@@ -234,6 +259,75 @@ func (info objInfo) equals(other objInfo) bool {
 	return true
 }
 
+type writerMethodExprInfo struct {
+	typeParamIdx int
+	methodInfo   selectorInfo
+}
+
+// typeParamMethodExprIdx returns the index where the given encoded
+// method expression function pointer appears within this dictionary's
+// type parameters method expressions section, adding it if necessary.
+func (dict *writerDict) typeParamMethodExprIdx(typeParamIdx int, methodInfo selectorInfo) int {
+	newInfo := writerMethodExprInfo{typeParamIdx, methodInfo}
+
+	for idx, oldInfo := range dict.typeParamMethodExprs {
+		if oldInfo == newInfo {
+			return idx
+		}
+	}
+
+	idx := len(dict.typeParamMethodExprs)
+	dict.typeParamMethodExprs = append(dict.typeParamMethodExprs, newInfo)
+	return idx
+}
+
+// subdictIdx returns the index where the given encoded object's
+// runtime dictionary appears within this dictionary's subdictionary
+// section, adding it if necessary.
+func (dict *writerDict) subdictIdx(newInfo objInfo) int {
+	for idx, oldInfo := range dict.subdicts {
+		if oldInfo.equals(newInfo) {
+			return idx
+		}
+	}
+
+	idx := len(dict.subdicts)
+	dict.subdicts = append(dict.subdicts, newInfo)
+	return idx
+}
+
+// rtypeIdx returns the index where the given encoded type's
+// *runtime._type value appears within this dictionary's rtypes
+// section, adding it if necessary.
+func (dict *writerDict) rtypeIdx(newInfo typeInfo) int {
+	for idx, oldInfo := range dict.rtypes {
+		if oldInfo == newInfo {
+			return idx
+		}
+	}
+
+	idx := len(dict.rtypes)
+	dict.rtypes = append(dict.rtypes, newInfo)
+	return idx
+}
+
+// itabIdx returns the index where the given encoded type pair's
+// *runtime.itab value appears within this dictionary's itabs section,
+// adding it if necessary.
+func (dict *writerDict) itabIdx(typInfo, ifaceInfo typeInfo) int {
+	newInfo := itabInfo{typInfo, ifaceInfo}
+
+	for idx, oldInfo := range dict.itabs {
+		if oldInfo == newInfo {
+			return idx
+		}
+	}
+
+	idx := len(dict.itabs)
+	dict.itabs = append(dict.itabs, newInfo)
+	return idx
+}
+
 func (pw *pkgWriter) newWriter(k pkgbits.RelocKind, marker pkgbits.SyncMarker) *writer {
 	return &writer{
 		Encoder: pw.NewEncoder(k, marker),
@@ -412,19 +506,9 @@ func (pw *pkgWriter) typIdx(typ types2.Type, dict *writerDict) typeInfo {
 		w.obj(obj, targs)
 
 	case *types2.TypeParam:
-		index := func() int {
-			for idx, name := range w.dict.implicits {
-				if name.Type().(*types2.TypeParam) == typ {
-					return idx
-				}
-			}
-
-			return len(w.dict.implicits) + typ.Index()
-		}()
-
 		w.derived = true
 		w.Code(pkgbits.TypeTypeParam)
-		w.Len(index)
+		w.Len(w.dict.typeParamIndex(typ))
 
 	case *types2.Array:
 		w.Code(pkgbits.TypeArray)
@@ -757,6 +841,33 @@ func (w *writer) objDict(obj types2.Object, dict *writerDict) {
 		w.Bool(typ.needed)
 	}
 
+	// Write runtime dictionary information.
+	//
+	// N.B., the go/types importer reads up to the section, but doesn't
+	// read any further, so it's safe to change. (See TODO above.)
+
+	w.Len(len(dict.typeParamMethodExprs))
+	for _, info := range dict.typeParamMethodExprs {
+		w.Len(info.typeParamIdx)
+		w.selectorInfo(info.methodInfo)
+	}
+
+	w.Len(len(dict.subdicts))
+	for _, info := range dict.subdicts {
+		w.objInfo(info)
+	}
+
+	w.Len(len(dict.rtypes))
+	for _, info := range dict.rtypes {
+		w.typInfo(info)
+	}
+
+	w.Len(len(dict.itabs))
+	for _, info := range dict.itabs {
+		w.typInfo(info.typ)
+		w.typInfo(info.iface)
+	}
+
 	assert(len(dict.derived) == nderived)
 }
 
@@ -960,17 +1071,21 @@ func (w *writer) funcargs(sig *types2.Signature) {
 
 func (w *writer) funcarg(param *types2.Var, result bool) {
 	if param.Name() != "" || result {
-		w.addLocal(param)
+		w.addLocal(param, true)
 	}
 }
 
 // addLocal records the declaration of a new local variable.
-func (w *writer) addLocal(obj *types2.Var) {
-	w.Sync(pkgbits.SyncAddLocal)
+func (w *writer) addLocal(obj *types2.Var, isParam bool) {
 	idx := len(w.localsIdx)
-	if w.p.SyncMarkers() {
-		w.Int(idx)
+
+	if !isParam {
+		w.Sync(pkgbits.SyncAddLocal)
+		if w.p.SyncMarkers() {
+			w.Int(idx)
+		}
 	}
+
 	if w.localsIdx == nil {
 		w.localsIdx = make(map[*types2.Var]int)
 	}
@@ -1161,7 +1276,7 @@ func (w *writer) assign(expr syntax.Expr) {
 
 			// TODO(mdempsky): Minimize locals index size by deferring
 			// this until the variables actually come into scope.
-			w.addLocal(obj)
+			w.addLocal(obj, false)
 			return
 		}
 	}
@@ -1424,7 +1539,7 @@ func (w *writer) switchStmt(stmt *syntax.SwitchStmt) {
 
 			obj := obj.(*types2.Var)
 			w.typ(obj.Type())
-			w.addLocal(obj)
+			w.addLocal(obj, false)
 		}
 
 		w.stmts(clause.Body)
@@ -1494,7 +1609,8 @@ func (w *writer) expr(expr syntax.Expr) {
 			obj := obj.(*types2.Func)
 
 			w.Code(exprFuncInst)
-			w.obj(obj, targs)
+			w.pos(expr)
+			w.funcInst(obj, targs)
 			return
 		}
 
@@ -1540,9 +1656,9 @@ func (w *writer) expr(expr syntax.Expr) {
 
 		case types2.MethodVal:
 			w.Code(exprMethodVal)
-			w.recvExpr(expr, sel)
+			typ := w.recvExpr(expr, sel)
 			w.pos(expr)
-			w.selector(sel.Obj())
+			w.methodExpr(expr, typ, sel)
 
 		case types2.MethodExpr:
 			w.Code(exprMethodExpr)
@@ -1551,9 +1667,27 @@ func (w *writer) expr(expr syntax.Expr) {
 			assert(ok)
 			assert(tv.IsType())
 
-			w.typ(tv.Type)
+			index := sel.Index()
+			implicits := index[:len(index)-1]
+
+			typ := tv.Type
+			w.typ(typ)
+
+			w.Len(len(implicits))
+			for _, ix := range implicits {
+				w.Len(ix)
+				typ = deref2(typ).Underlying().(*types2.Struct).Field(ix).Type()
+			}
+
+			recv := sel.Obj().(*types2.Func).Type().(*types2.Signature).Recv().Type()
+			if w.Bool(isPtrTo(typ, recv)) { // need deref
+				typ = recv
+			} else if w.Bool(isPtrTo(recv, typ)) { // need addr
+				typ = recv
+			}
+
 			w.pos(expr)
-			w.selector(sel.Obj())
+			w.methodExpr(expr, typ, sel)
 		}
 
 	case *syntax.IndexExpr:
@@ -1695,18 +1829,28 @@ func (w *writer) expr(expr syntax.Expr) {
 		}
 
 		writeFunExpr := func() {
-			if selector, ok := unparen(expr.Fun).(*syntax.SelectorExpr); ok {
+			fun := unparen(expr.Fun)
+
+			if selector, ok := fun.(*syntax.SelectorExpr); ok {
 				if sel, ok := w.p.info.Selections[selector]; ok && sel.Kind() == types2.MethodVal {
-					w.recvExpr(selector, sel)
 					w.Bool(true) // method call
-					w.pos(selector)
-					w.selector(sel.Obj())
+					typ := w.recvExpr(selector, sel)
+					w.methodExpr(selector, typ, sel)
 					return
 				}
 			}
 
-			w.expr(expr.Fun)
 			w.Bool(false) // not a method call (i.e., normal function call)
+
+			if obj, inst := lookupObj(w.p.info, fun); w.Bool(obj != nil && inst.TypeArgs.Len() != 0) {
+				obj := obj.(*types2.Func)
+
+				w.pos(fun)
+				w.funcInst(obj, inst.TypeArgs)
+				return
+			}
+
+			w.expr(fun)
 		}
 
 		sigType := types2.CoreType(tv.Type).(*types2.Signature)
@@ -1742,7 +1886,8 @@ func (w *writer) optExpr(expr syntax.Expr) {
 }
 
 // recvExpr writes out expr.X, but handles any implicit addressing,
-// dereferencing, and field selections.
+// dereferencing, and field selections appropriate for the method
+// selection.
 func (w *writer) recvExpr(expr *syntax.SelectorExpr, sel *types2.Selection) types2.Type {
 	index := sel.Index()
 	implicits := index[:len(index)-1]
@@ -1758,13 +1903,6 @@ func (w *writer) recvExpr(expr *syntax.SelectorExpr, sel *types2.Selection) type
 		w.Len(ix)
 	}
 
-	isPtrTo := func(from, to types2.Type) bool {
-		if from, ok := from.(*types2.Pointer); ok {
-			return types2.Identical(from.Elem(), to)
-		}
-		return false
-	}
-
 	recv := sel.Obj().(*types2.Func).Type().(*types2.Signature).Recv().Type()
 	if w.Bool(isPtrTo(typ, recv)) { // needs deref
 		typ = recv
@@ -1775,6 +1913,84 @@ func (w *writer) recvExpr(expr *syntax.SelectorExpr, sel *types2.Selection) type
 	return typ
 }
 
+// funcInst writes a reference to an instantiated function.
+func (w *writer) funcInst(obj *types2.Func, targs *types2.TypeList) {
+	info := w.p.objInstIdx(obj, targs, w.dict)
+
+	// Type arguments list contains derived types; we can emit a static
+	// call to the shaped function, but need to dynamically compute the
+	// runtime dictionary pointer.
+	if w.Bool(info.anyDerived()) {
+		w.Len(w.dict.subdictIdx(info))
+		return
+	}
+
+	// Type arguments list is statically known; we can emit a static
+	// call with a statically reference to the respective runtime
+	// dictionary.
+	w.objInfo(info)
+}
+
+// methodExpr writes out a reference to the method selected by
+// expr. sel should be the corresponding types2.Selection, and recv
+// the type produced after any implicit addressing, dereferencing, and
+// field selection. (Note: recv might differ from sel.Obj()'s receiver
+// parameter in the case of interface types, and is needed for
+// handling type parameter methods.)
+func (w *writer) methodExpr(expr *syntax.SelectorExpr, recv types2.Type, sel *types2.Selection) {
+	fun := sel.Obj().(*types2.Func)
+	sig := fun.Type().(*types2.Signature)
+
+	w.typ(recv)
+	w.signature(sig)
+	w.pos(expr)
+	w.selector(fun)
+
+	// Method on a type parameter. These require an indirect call
+	// through the current function's runtime dictionary.
+	if typeParam, ok := recv.(*types2.TypeParam); w.Bool(ok) {
+		typeParamIdx := w.dict.typeParamIndex(typeParam)
+		methodInfo := w.p.selectorIdx(fun)
+
+		w.Len(w.dict.typeParamMethodExprIdx(typeParamIdx, methodInfo))
+		return
+	}
+
+	if isInterface(recv) != isInterface(sig.Recv().Type()) {
+		w.p.fatalf(expr, "isInterface inconsistency: %v and %v", recv, sig.Recv().Type())
+	}
+
+	if !isInterface(recv) {
+		if named, ok := deref2(recv).(*types2.Named); ok {
+			obj, targs := splitNamed(named)
+			info := w.p.objInstIdx(obj, targs, w.dict)
+
+			// Method on a derived receiver type. These can be handled by a
+			// static call to the shaped method, but require dynamically
+			// looking up the appropriate dictionary argument in the current
+			// function's runtime dictionary.
+			if w.p.hasImplicitTypeParams(obj) || info.anyDerived() {
+				w.Bool(true) // dynamic subdictionary
+				w.Len(w.dict.subdictIdx(info))
+				return
+			}
+
+			// Method on a fully known receiver type. These can be handled
+			// by a static call to the shaped method, and with a static
+			// reference to the receiver type's dictionary.
+			if targs.Len() != 0 {
+				w.Bool(false) // no dynamic subdictionary
+				w.Bool(true)  // static dictionary
+				w.objInfo(info)
+				return
+			}
+		}
+	}
+
+	w.Bool(false) // no dynamic subdictionary
+	w.Bool(false) // no static dictionary
+}
+
 // multiExpr writes a sequence of expressions, where the i'th value is
 // implicitly converted to dstType(i). It also handles when exprs is a
 // single, multi-valued expression (e.g., the multi-valued argument in
@@ -1930,18 +2146,34 @@ func (w *writer) exprs(exprs []syntax.Expr) {
 // rtype writes information so that the reader can construct an
 // expression of type *runtime._type representing typ.
 func (w *writer) rtype(typ types2.Type) {
+	typ = types2.Default(typ)
+
 	w.Sync(pkgbits.SyncRType)
-	w.typNeeded(typ)
-}
 
-// typNeeded writes a reference to typ, and records that its
-// *runtime._type is needed.
-func (w *writer) typNeeded(typ types2.Type) {
 	info := w.p.typIdx(typ, w.dict)
-	w.typInfo(info)
+	if w.Bool(info.derived) {
+		w.Len(w.dict.rtypeIdx(info))
+	} else {
+		w.typInfo(info)
+	}
+}
+
+func isUntyped(typ types2.Type) bool {
+	basic, ok := typ.(*types2.Basic)
+	return ok && basic.Info()&types2.IsUntyped != 0
+}
+
+func (w *writer) itab(typ, iface types2.Type) {
+	typ = types2.Default(typ)
+	iface = types2.Default(iface)
 
-	if info.derived {
-		w.dict.derived[info.idx].needed = true
+	typInfo := w.p.typIdx(typ, w.dict)
+	ifaceInfo := w.p.typIdx(iface, w.dict)
+	if w.Bool(typInfo.derived || ifaceInfo.derived) {
+		w.Len(w.dict.itabIdx(typInfo, ifaceInfo))
+	} else {
+		w.typInfo(typInfo)
+		w.typInfo(ifaceInfo)
 	}
 }
 
@@ -1949,8 +2181,7 @@ func (w *writer) typNeeded(typ types2.Type) {
 // expressions for converting from src to dst.
 func (w *writer) convRTTI(src, dst types2.Type) {
 	w.Sync(pkgbits.SyncConvRTTI)
-	w.typNeeded(src)
-	w.typNeeded(dst)
+	w.itab(src, dst)
 }
 
 func (w *writer) exprType(iface types2.Type, typ syntax.Expr) {
@@ -1963,9 +2194,13 @@ func (w *writer) exprType(iface types2.Type, typ syntax.Expr) {
 	w.Sync(pkgbits.SyncExprType)
 	w.pos(typ)
 
-	w.typNeeded(tv.Type)
 	if w.Bool(iface != nil && !iface.Underlying().(*types2.Interface).Empty()) {
-		w.typ(iface)
+		w.itab(tv.Type, iface)
+	} else {
+		w.rtype(tv.Type)
+
+		info := w.p.typIdx(tv.Type, w.dict)
+		w.Bool(info.derived)
 	}
 }
 
@@ -2435,3 +2670,9 @@ func asPragmaFlag(p syntax.Pragma) ir.PragmaFlag {
 	}
 	return p.(*pragmas).Flag
 }
+
+// isPtrTo reports whether from is the type *to.
+func isPtrTo(from, to types2.Type) bool {
+	ptr, ok := from.(*types2.Pointer)
+	return ok && types2.Identical(ptr.Elem(), to)
+}