Use generic struct field not instantiated one in Uses

We do not extract instantiated named types, and instead use the generic type. But fields of the underlying struct of an instantiated named types are obtained from the Uses map. We solve this keeping track of which objects should be overridden by which other objects.

Author: owen-mc

extractor/extractor.go

@@ -1531,6 +1531,7 @@ func extractType(tw *trap.Writer, tp types.Type) trap.Label {
 dbscheme.TypeNameTable.Emit(tw, lbl, origintp.Obj().Name())
 underlying := origintp.Underlying()
 extractUnderlyingType(tw, lbl, underlying)
+trackInstantiatedStructFields(tw, tp, origintp)
 
 entitylbl, exists := tw.Labeler.LookupObjectID(origintp.Obj(), lbl)
 if entitylbl == trap.InvalidLabel {
@@ -1902,6 +1903,9 @@ func getObjectBeingUsed(tw *trap.Writer, ident *ast.Ident) types.Object {
 if obj == nil {
 return nil
 }
+if override, ok := tw.ObjectsOverride[obj]; ok {
+return override
+}
 if funcObj, ok := obj.(*types.Func); ok {
 sig := funcObj.Type().(*types.Signature)
 if recv := sig.Recv(); recv != nil {
@@ -1948,3 +1952,33 @@ func tryGetGenericType(tp types.Type) (*types.Named, bool) {
 }
 return nil, false
 }
+
+// trackInstantiatedStructFields tries to give the fields of an instantiated
+// struct type underlying `tp` the same labels as the corresponding fields of
+// the generic struct type. This is so that when we come across the
+// instantiated field in `tw.Package.TypesInfo.Uses` we will get the label for
+// the generic field instead.
+func trackInstantiatedStructFields(tw *trap.Writer, tp, origintp *types.Named) {
+if tp == origintp {
+return
+}
+
+if instantiatedStruct, ok := tp.Underlying().(*types.Struct); ok {
+genericStruct, ok2 := origintp.Underlying().(*types.Struct)
+if !ok2 {
+log.Fatalf(
+"Error: underlying type of instantiated type is a struct but underlying type of generic type is %s",
+origintp.Underlying())
+}
+
+if instantiatedStruct.NumFields() != genericStruct.NumFields() {
+log.Fatalf(
+"Error: instantiated struct %s has different number of fields than the generic version %s (%d != %d)",
+instantiatedStruct, genericStruct, instantiatedStruct.NumFields(), genericStruct.NumFields())
+}
+
+for i := 0; i < instantiatedStruct.NumFields(); i++ {
+tw.ObjectsOverride[instantiatedStruct.Field(i)] = genericStruct.Field(i)
+}
+}
+}

extractor/trap/trapwriter.go

@@ -27,6 +27,7 @@ type Writer struct {
 Package *packages.Package
 TypesOverride map[ast.Expr]types.Type
 TypeParamParent map[*types.TypeParam]Label
+ObjectsOverride map[types.Object]types.Object
 }
 
 func FileFor(path string) (string, error) {
@@ -66,6 +67,7 @@ func NewWriter(path string, pkg *packages.Package) (*Writer, error) {
 pkg,
 make(map[ast.Expr]types.Type),
 make(map[*types.TypeParam]Label),
+make(map[types.Object]types.Object),
 }
 tw.Labeler = newLabeler(tw)
 return tw, nil