5 files changed, 55 insertions, 17 deletions
diff --git a/tools/go_marshal/analysis/analysis_unsafe.go b/tools/go_marshal/analysis/analysis_unsafe.go
index cd55cf5cb..7a3d6bbea 100644
--- a/tools/go_marshal/analysis/analysis_unsafe.go
+++ b/tools/go_marshal/analysis/analysis_unsafe.go
@@ -81,7 +81,7 @@ func RandomizeValue(x interface{}) {
 // This is used for zeroing padding fields after calling RandomizeValue.
 func reflectZeroPaddingFields(r reflect.Type, data []byte, zero bool) {
 	if zero {
-		for i, _ := range data {
+		for i := range data {
 			data[i] = 0
 		}
 	}
diff --git a/tools/go_marshal/gomarshal/generator.go b/tools/go_marshal/gomarshal/generator.go
index 4a53d25be..fa642c88a 100644
--- a/tools/go_marshal/gomarshal/generator.go
+++ b/tools/go_marshal/gomarshal/generator.go
@@ -148,7 +148,7 @@ func (g *Generator) writeTypeChecks(ms map[string]struct{}) error {
 	}
 
 	msl := make([]string, 0, len(ms))
-	for m, _ := range ms {
+	for m := range ms {
 		msl = append(msl, m)
 	}
 	sort.Strings(msl)
@@ -213,10 +213,11 @@ type sliceAPI struct {
 type marshallableType struct {
 	spec  *ast.TypeSpec
 	slice *sliceAPI
+	recv  string
 }
 
-func newMarshallableType(fset *token.FileSet, tagLine *ast.Comment, spec *ast.TypeSpec) marshallableType {
-	mt := marshallableType{
+func newMarshallableType(fset *token.FileSet, tagLine *ast.Comment, spec *ast.TypeSpec) *marshallableType {
+	mt := &marshallableType{
 		spec:  spec,
 		slice: nil,
 	}
@@ -261,12 +262,31 @@ func newMarshallableType(fset *token.FileSet, tagLine *ast.Comment, spec *ast.Ty
 
 // collectMarshallableTypes walks the parsed AST and collects a list of type
 // declarations for which we need to generate the Marshallable interface.
-func (g *Generator) collectMarshallableTypes(a *ast.File, f *token.FileSet) []marshallableType {
-	var types []marshallableType
+func (g *Generator) collectMarshallableTypes(a *ast.File, f *token.FileSet) map[*ast.TypeSpec]*marshallableType {
+	recv := make(map[string]string) // Type name to recevier name.
+	types := make(map[*ast.TypeSpec]*marshallableType)
 	for _, decl := range a.Decls {
 		gdecl, ok := decl.(*ast.GenDecl)
 		// Type declaration?
 		if !ok || gdecl.Tok != token.TYPE {
+			// Is this a function declaration? We remember receiver names.
+			d, ok := decl.(*ast.FuncDecl)
+			if ok && d.Recv != nil && len(d.Recv.List) == 1 {
+				// Accept concrete methods & pointer methods.
+				ident, ok := d.Recv.List[0].Type.(*ast.Ident)
+				if !ok {
+					var st *ast.StarExpr
+					st, ok = d.Recv.List[0].Type.(*ast.StarExpr)
+					if ok {
+						ident, ok = st.X.(*ast.Ident)
+					}
+				}
+				// The receiver name may be not present.
+				if ok && len(d.Recv.List[0].Names) == 1 {
+					// Recover the type receiver name in this case.
+					recv[ident.Name] = d.Recv.List[0].Names[0].Name
+				}
+			}
 			debugfAt(f.Position(decl.Pos()), "Skipping declaration since it's not a type declaration.\n")
 			continue
 		}
@@ -305,10 +325,20 @@ func (g *Generator) collectMarshallableTypes(a *ast.File, f *token.FileSet) []ma
 				// don't support it.
 				abortAt(f.Position(t.Pos()), fmt.Sprintf("Marshalling codegen was requested on type '%s', but go-marshal doesn't support this kind of declaration.\n", t.Name))
 			}
-			types = append(types, newMarshallableType(f, tagLine, t))
-
+			types[t] = newMarshallableType(f, tagLine, t)
 		}
 	}
+	// Update the types with the last seen receiver. As long as the
+	// receiver name is consistent for the type, then we will generate
+	// code that is still consistent with itself.
+	for t, mt := range types {
+		r, ok := recv[t.Name.Name]
+		if !ok {
+			mt.recv = receiverName(t) // Default.
+			continue
+		}
+		mt.recv = r // Last seen.
+	}
 	return types
 }
 
@@ -345,8 +375,8 @@ func (g *Generator) collectImports(a *ast.File, f *token.FileSet) map[string]imp
 
 }
 
-func (g *Generator) generateOne(t marshallableType, fset *token.FileSet) *interfaceGenerator {
-	i := newInterfaceGenerator(t.spec, fset)
+func (g *Generator) generateOne(t *marshallableType, fset *token.FileSet) *interfaceGenerator {
+	i := newInterfaceGenerator(t.spec, t.recv, fset)
 	switch ty := t.spec.Type.(type) {
 	case *ast.StructType:
 		i.validateStruct(t.spec, ty)
@@ -376,8 +406,8 @@ func (g *Generator) generateOne(t marshallableType, fset *token.FileSet) *interf
 
 // generateOneTestSuite generates a test suite for the automatically generated
 // implementations type t.
-func (g *Generator) generateOneTestSuite(t marshallableType) *testGenerator {
-	i := newTestGenerator(t.spec)
+func (g *Generator) generateOneTestSuite(t *marshallableType) *testGenerator {
+	i := newTestGenerator(t.spec, t.recv)
 	i.emitTests(t.slice)
 	return i
 }
@@ -417,7 +447,15 @@ func (g *Generator) Run() error {
 	for i, a := range asts {
 		// Collect type declarations marked for code generation and generate
 		// Marshallable interfaces.
+		var sortedTypes []*marshallableType
 		for _, t := range g.collectMarshallableTypes(a, fsets[i]) {
+			sortedTypes = append(sortedTypes, t)
+		}
+		sort.Slice(sortedTypes, func(x, y int) bool {
+			// Sort by type name, which should be unique within a package.
+			return sortedTypes[x].spec.Name.String() < sortedTypes[y].spec.Name.String()
+		})
+		for _, t := range sortedTypes {
 			impl := g.generateOne(t, fsets[i])
 			// Collect Marshallable types referenced by the generated code.
 			for ref := range impl.ms {
diff --git a/tools/go_marshal/gomarshal/generator_interfaces.go b/tools/go_marshal/gomarshal/generator_interfaces.go
index 36447b86b..65f5ea34d 100644
--- a/tools/go_marshal/gomarshal/generator_interfaces.go
+++ b/tools/go_marshal/gomarshal/generator_interfaces.go
@@ -54,10 +54,10 @@ func (g *interfaceGenerator) typeName() string {
 }
 
 // newinterfaceGenerator creates a new interface generator.
-func newInterfaceGenerator(t *ast.TypeSpec, fset *token.FileSet) *interfaceGenerator {
+func newInterfaceGenerator(t *ast.TypeSpec, r string, fset *token.FileSet) *interfaceGenerator {
 	g := &interfaceGenerator{
 		t:  t,
-		r:  receiverName(t),
+		r:  r,
 		f:  fset,
 		is: make(map[string]struct{}),
 		ms: make(map[string]struct{}),
diff --git a/tools/go_marshal/gomarshal/generator_interfaces_struct.go b/tools/go_marshal/gomarshal/generator_interfaces_struct.go
index fe76d3785..5f6306b8f 100644
--- a/tools/go_marshal/gomarshal/generator_interfaces_struct.go
+++ b/tools/go_marshal/gomarshal/generator_interfaces_struct.go
@@ -38,7 +38,7 @@ func (g *interfaceGenerator) areFieldsPackedExpression() (string, bool) {
 	}
 
 	cs := make([]string, 0, len(g.as))
-	for accessor, _ := range g.as {
+	for accessor := range g.as {
 		cs = append(cs, fmt.Sprintf("%s.Packed()", accessor))
 	}
 	// Sort expressions for determinstic build outputs.
diff --git a/tools/go_marshal/gomarshal/generator_tests.go b/tools/go_marshal/gomarshal/generator_tests.go
index 631295373..6cf00843f 100644
--- a/tools/go_marshal/gomarshal/generator_tests.go
+++ b/tools/go_marshal/gomarshal/generator_tests.go
@@ -53,10 +53,10 @@ type testGenerator struct {
 	decl *importStmt
 }
 
-func newTestGenerator(t *ast.TypeSpec) *testGenerator {
+func newTestGenerator(t *ast.TypeSpec, r string) *testGenerator {
 	g := &testGenerator{
 		t:       t,
-		r:       receiverName(t),
+		r:       r,
 		imports: newImportTable(),
 	}