Update the generator to the google-internal version.

- a few changes to the plugin interface
- handle situation where an imported package is called "proto"
- use the file base name for the package name if a package is not specified

R=rsc
CC=dsymonds1
http://codereview.appspot.com/1641042

compiler/generator/generator.go

@@ -33,10 +33,6 @@
 	The code generator for the plugin for the Google protocol buffer compiler.
 	It generates Go code from the protocol buffer description files read by the
 	main routine.
-
-	Not supported yet:
-		services
-		options
 */
 package generator
 
@@ -57,12 +53,16 @@ import (
 // such as to produce RPC stubs.
 type Plugin interface {
 	// Name identifies the plugin.
-	Name()	string
-	// Generate produces the code generated by the plugin for this file, except for the imports,
-	// by calling the generator's methods P, In, and Out.
-	Generate(g *Generator, file *FileDescriptor)
+	Name() string
+	// Init is called once after data structures are built but before
+	// code generation begins.
+	Init(g *Generator)
+	// Generate produces the code generated by the plugin for this file,
+	// except for the imports, by calling the generator's methods P, In, and Out.
+	Generate(file *FileDescriptor)
 	// GenerateImports produces the import declarations for this file.
-	GenerateImports(g *Generator, file *FileDescriptor)
+	// It is called after Generate.
+	GenerateImports(file *FileDescriptor)
 }
 
 var plugins []Plugin
@@ -72,13 +72,12 @@ var plugins []Plugin
 func RegisterPlugin(p Plugin) {
 	n := len(plugins)
 	if cap(plugins) == n {
-		nplugins := make([]Plugin, n, n+10)	// very unlikely to need more than this
+		nplugins := make([]Plugin, n, n+10) // very unlikely to need more than this
 		copy(nplugins, plugins)
 		plugins = nplugins
 	}
-	plugins = plugins[0:n+1]
+	plugins = plugins[0 : n+1]
 	plugins[n] = p
-	log.Stderr("installed plugin:", p.Name())
 }
 
 // Each type we import as a protocol buffer (other than FileDescriptorProto) needs
@@ -180,7 +179,7 @@ func (e *EnumDescriptor) integerValueAsString(name string) string {
 type ExtensionDescriptor struct {
 	common
 	*descriptor.FieldDescriptorProto
-	parent   *Descriptor // The containing message, if any.
+	parent *Descriptor // The containing message, if any.
 }
 
 // TypeName returns the elements of the dotted type name.
@@ -204,31 +203,24 @@ func (e *ExtensionDescriptor) TypeName() (s []string) {
 // Those slices are constructed by WrapTypes.
 type FileDescriptor struct {
 	*descriptor.FileDescriptorProto
-	desc []*Descriptor           // All the messages defined in this file.
-	enum []*EnumDescriptor       // All the enums defined in this file.
-	ext  []*ExtensionDescriptor  // All the top-level extensions defined in this file.
+	desc []*Descriptor          // All the messages defined in this file.
+	enum []*EnumDescriptor      // All the enums defined in this file.
+	ext  []*ExtensionDescriptor // All the top-level extensions defined in this file.
 }
 
 // PackageName is the package name we'll use in the generated code to refer to this file.
 func (d *FileDescriptor) PackageName() string { return uniquePackageOf(d.FileDescriptorProto) }
 
 // The package named defined in the input for this file, possibly dotted.
+// If the file does not define a package, use the base of the file name.
 func (d *FileDescriptor) originalPackageName() string {
-	return proto.GetString(d.Package)
-}
-
-// Whether the proto library needs importing.
-// This will be true if there are any enums, extensions, or messages with extension ranges.
-func (d *FileDescriptor) needProtoImport() bool {
-	if len(d.enum) > 0 || len(d.ext) > 0 {
-		return true
+	// Does the file have a package clause?
+	pkg := proto.GetString(d.Package)
+	if pkg != "" {
+		return pkg
 	}
-	for _, desc := range d.desc {
-		if len(desc.ext) > 0 || len(desc.ExtensionRange) > 0 {
-			return true
-		}
-	}
-	return false
+	// Use the file base name.
+	return BaseName(proto.GetString(d.Name))
 }
 
 // Object is an interface abstracting the abilities shared by enums and messages.
@@ -256,6 +248,12 @@ type Generator struct {
 	Request  *plugin.CodeGeneratorRequest  // The input.
 	Response *plugin.CodeGeneratorResponse // The output.
 
+	Param        map[string]string // Command-line parameters.
+	ImportPrefix string            // String to prefix to imported package file names.
+	ImportMap    map[string]string // Mapping from import name to generated name
+
+	ProtoPkg string // The name under which we import the library's package proto.
+
 	packageName      string            // What we're calling ourselves.
 	allFiles         []*FileDescriptor // All files in the tree
 	genFiles         []*FileDescriptor // Those files we will generate output for.
@@ -290,6 +288,29 @@ func (g *Generator) Fail(msgs ...string) {
 	os.Exit(1)
 }
 
+// CommandLineParameters breaks the comma-separated list of key=value pairs
+// in the parameter (a member of the request protobuf) into a key/value map.
+// It then sets file name mappings defined by those entries.
+func (g *Generator) CommandLineParameters(parameter string) {
+	g.Param = make(map[string]string)
+	for _, p := range strings.Split(parameter, ",", 0) {
+		if i := strings.Index(p, "="); i < 0 {
+			g.Param[p] = ""
+		} else {
+			g.Param[p[0:i]] = p[i+1:]
+		}
+	}
+
+	g.ImportMap = make(map[string]string)
+	for k, v := range g.Param {
+		if k == "import_prefix" {
+			g.ImportPrefix = v
+		} else if len(k) > 0 && k[0] == 'M' {
+			g.ImportMap[k[1:]] = v
+		}
+	}
+}
+
 // DefaultPackageName returns the package name printed for the object.
 // If its file is in a different package, it returns the package name we're using for this file, plus ".".
 // Otherwise it returns the empty string.
@@ -303,17 +324,40 @@ func (g *Generator) DefaultPackageName(obj Object) string {
 
 // For each input file, the unique package name to use, underscored.
 var uniquePackageName = make(map[*descriptor.FileDescriptorProto]string)
+// Package names already registered.  Key is the name from the .proto file;
+// value is the name that appears in the generated code.
+var pkgNamesInUse = make(map[string]bool)
+
+// Create and remember a guaranteed unique package name for this file descriptor.
+// Pkg is the candidate name.  If f is nil, it's a builtin package like "proto" and
+// has no file descriptor.
+func RegisterUniquePackageName(pkg string, f *FileDescriptor) string {
+	for pkgNamesInUse[pkg] {
+		// It's a duplicate; must rename.
+		pkg += "X"
+	}
+	// Install it.
+	pkgNamesInUse[pkg] = true
+	pkg = strings.Map(DotToUnderscore, pkg)
+	if f != nil {
+		uniquePackageName[f.FileDescriptorProto] = pkg
+	}
+	return pkg
+}
 
-// SetPackageNames Sets the package name for this run.
+// SetPackageNames sets the package name for this run.
 // The package name must agree across all files being generated.
 // It also defines unique package names for all imported files.
 func (g *Generator) SetPackageNames() {
-	inUse := make(map[string]bool)
-	pkg := proto.GetString(g.genFiles[0].Package)
-	g.packageName = strings.Map(DotToUnderscore, pkg)
-	inUse[pkg] = true
+	// Register the name for this package.  It will be the first name
+	// registered so is guaranteed to be unmodified.
+	pkg := g.genFiles[0].originalPackageName()
+	g.packageName = RegisterUniquePackageName(pkg, g.genFiles[0])
+	// Register the proto package name.  It might collide with the
+	// name of a package we import.
+	g.ProtoPkg = RegisterUniquePackageName("proto", nil)
 	for _, f := range g.genFiles {
-		thisPkg := proto.GetString(f.Package)
+		thisPkg := f.originalPackageName()
 		if thisPkg != pkg {
 			g.Fail("inconsistent package names:", thisPkg, pkg)
 		}
@@ -327,20 +371,7 @@ AllFiles:
 				continue AllFiles
 			}
 		}
-		truePkg := proto.GetString(f.Package)
-		pkg := truePkg
-		for {
-			_, present := inUse[pkg]
-			if present {
-				// It's a duplicate; must rename.
-				pkg += "X"
-				continue
-			}
-			break
-		}
-		// Install it.
-		inUse[pkg] = true
-		uniquePackageName[f.FileDescriptorProto] = strings.Map(DotToUnderscore, pkg)
+		RegisterUniquePackageName(f.originalPackageName(), f)
 	}
 }
 
@@ -350,10 +381,6 @@ AllFiles:
 func (g *Generator) WrapTypes() {
 	g.allFiles = make([]*FileDescriptor, len(g.Request.ProtoFile))
 	for i, f := range g.Request.ProtoFile {
-		pkg := proto.GetString(f.Package)
-		if pkg == "" {
-			g.Fail(proto.GetString(f.Name), "is missing a package declaration")
-		}
 		// We must wrap the descriptors before we wrap the enums
 		descs := wrapDescriptors(f)
 		g.buildNestedDescriptors(descs)
@@ -483,13 +510,19 @@ func wrapExtensions(file *descriptor.FileDescriptorProto) []*ExtensionDescriptor
 	return sl
 }
 
-// BuildTypeNameMap builds the map from fully qualified type names to objects. 
+// BuildTypeNameMap builds the map from fully qualified type names to objects.
 // The key names for the map come from the input data, which puts a period at the beginning.
 // It should be called after SetPackageNames and before GenerateAllFiles.
 func (g *Generator) BuildTypeNameMap() {
 	g.typeNameToObject = make(map[string]Object)
 	for _, f := range g.allFiles {
-		dottedPkg := "." + f.originalPackageName() + "."
+		// The names in this loop are defined by the proto world, not us, so the
+		// package name may be empty.  If so, the dotted package name of X will
+		// be ".X"; otherwise it will be ".pkg.X".
+		dottedPkg := "." + proto.GetString(f.Package)
+		if dottedPkg != "." {
+			dottedPkg += "."
+		}
 		for _, enum := range f.enum {
 			name := dottedPkg + dottedSlice(enum.TypeName())
 			g.typeNameToObject[name] = enum
@@ -521,8 +554,16 @@ func (g *Generator) P(str ...interface{}) {
 			g.WriteString(s)
 		case *string:
 			g.WriteString(*s)
+		case bool:
+			g.WriteString(fmt.Sprintf("%t", s))
+		case *bool:
+			g.WriteString(fmt.Sprintf("%t", *s))
 		case *int32:
 			g.WriteString(fmt.Sprintf("%d", *s))
+		case float64:
+			g.WriteString(fmt.Sprintf("%g", s))
+		case *float64:
+			g.WriteString(fmt.Sprintf("%g", *s))
 		default:
 			g.Fail(fmt.Sprintf("unknown type in printer: %T", v))
 		}
@@ -542,10 +583,14 @@ func (g *Generator) Out() {
 
 // GenerateAllFiles generates the output for all the files we're outputting.
 func (g *Generator) GenerateAllFiles() {
+	// Initialize the plugins
+	for _, p := range plugins {
+		p.Init(g)
+	}
+	// Generate the output.
 	for i, file := range g.genFiles {
 		g.Reset()
 		g.generate(file)
-		g.runPlugins(file)
 		g.Response.File[i] = plugin.NewCodeGeneratorResponse_File()
 		g.Response.File[i].Name = proto.String(goFileName(*file.Name))
 		g.Response.File[i].Content = proto.String(g.String())
@@ -555,7 +600,7 @@ func (g *Generator) GenerateAllFiles() {
 // Run all the plugins associated with the file.
 func (g *Generator) runPlugins(file *FileDescriptor) {
 	for _, p := range plugins {
-		p.Generate(g, file)
+		p.Generate(file)
 	}
 }
 
@@ -588,6 +633,9 @@ func (g *Generator) generate(file *FileDescriptor) {
 	}
 	g.generateInitFunction()
 
+	// Run the plugins before the imports so we know which imports are necessary.
+	g.runPlugins(file)
+
 	// Generate header and imports last, though they appear first in the output.
 	rem := g.Buffer
 	g.Buffer = new(bytes.Buffer)
@@ -607,16 +655,25 @@ func (g *Generator) generateHeader() {
 
 // Generate the header, including package definition and imports
 func (g *Generator) generateImports() {
-	if g.file.needProtoImport() {
-		g.P(`import "goprotobuf.googlecode.com/hg/proto"`)
-	}
+	// We almost always need a proto import.  Rather than computing when we
+	// do, which is tricky when there's a plugin, just import it and
+	// reference it later.
+	g.P("import " + g.ProtoPkg + " " + Quote(g.ImportPrefix+"net/proto2/go/proto"))
 	for _, s := range g.file.Dependency {
 		// Need to find the descriptor for this file
 		for _, fd := range g.allFiles {
+			// Do not import our own package.
+			if fd.PackageName() == g.packageName {
+				continue
+			}
 			if proto.GetString(fd.Name) == s {
 				filename := goFileName(s)
+				if substitution, ok := g.ImportMap[s]; ok {
+					filename = substitution
+				}
+				filename = g.ImportPrefix + filename
 				if strings.HasSuffix(filename, ".go") {
-					filename = filename[0:len(filename)-3]
+					filename = filename[0 : len(filename)-3]
 				}
 				if _, ok := g.usedPackages[fd.PackageName()]; ok {
 					g.P("import ", fd.PackageName(), " ", Quote(filename))
@@ -630,9 +687,12 @@ func (g *Generator) generateImports() {
 	g.P()
 	// TODO: may need to worry about uniqueness across plugins
 	for _, p := range plugins {
-		p.GenerateImports(g, g.file)
+		p.GenerateImports(g.file)
 		g.P()
 	}
+	g.P("// Reference proto import to suppress error if it's not otherwise used.")
+	g.P("var _ = ", g.ProtoPkg, ".GetString")
+	g.P()
 }
 
 // Generate the enum definitions for this EnumDescriptor.
@@ -652,7 +712,7 @@ func (g *Generator) generateEnum(enum *EnumDescriptor) {
 	g.P(")")
 	g.P("var ", ccTypeName, "_name = map[int32] string {")
 	g.In()
-	generated := make(map[int32] bool)	// avoid duplicate values
+	generated := make(map[int32]bool) // avoid duplicate values
 	for _, e := range enum.Value {
 		duplicate := ""
 		if _, present := generated[*e.Number]; present {
@@ -871,15 +931,15 @@ func (g *Generator) generateMessage(message *Descriptor) {
 	// Extension support methods
 	if len(message.ExtensionRange) > 0 {
 		g.P()
-		g.P("var extRange_", ccTypeName, " = []proto.ExtensionRange{")
+		g.P("var extRange_", ccTypeName, " = []", g.ProtoPkg, ".ExtensionRange{")
 		g.In()
 		for _, r := range message.ExtensionRange {
-			end := fmt.Sprint(*r.End - 1)  // make range inclusive on both ends
-			g.P("proto.ExtensionRange{", r.Start, ", ", end, "},")
+			end := fmt.Sprint(*r.End - 1) // make range inclusive on both ends
+			g.P(g.ProtoPkg+".ExtensionRange{", r.Start, ", ", end, "},")
 		}
 		g.Out()
 		g.P("}")
-		g.P("func (*", ccTypeName, ") ExtensionRangeArray() []proto.ExtensionRange {")
+		g.P("func (*", ccTypeName, ") ExtensionRangeArray() []", g.ProtoPkg, ".ExtensionRange {")
 		g.In()
 		g.P("return extRange_", ccTypeName)
 		g.Out()
@@ -949,7 +1009,7 @@ func (g *Generator) generateExtension(ext *ExtensionDescriptor) {
 	tag := g.goTag(field, wireType)
 	g.RecordTypeUse(*ext.Extendee)
 
-	g.P("var ", ccTypeName, " = &proto.ExtensionDesc{")
+	g.P("var ", ccTypeName, " = &", g.ProtoPkg, ".ExtensionDesc{")
 	g.In()
 	g.P("ExtendedType: (", extendedType, ")(nil),")
 	g.P("ExtensionType: (", fieldType, ")(nil),")
@@ -977,7 +1037,7 @@ func (g *Generator) generateEnumRegistration(enum *EnumDescriptor) {
 	typeName := enum.TypeName()
 	// The full type name, CamelCased.
 	ccTypeName := CamelCaseSlice(typeName)
-	g.P("proto.RegisterEnum(", Quote(pkg+ccTypeName), ", ", ccTypeName+"_name, ", ccTypeName+"_value)")
+	g.P(g.ProtoPkg+".RegisterEnum(", Quote(pkg+ccTypeName), ", ", ccTypeName+"_name, ", ccTypeName+"_value)")
 }
 
 // And now lots of helper functions.
@@ -1042,10 +1102,26 @@ func isRepeated(field *descriptor.FieldDescriptorProto) bool {
 }
 
 // DotToUnderscore is the mapping function used to generate Go names from package names,
-// which can be dotted in the input .proto file.
+// which can be dotted in the input .proto file.  It maps dots to underscores.
+// Because we also get here from package names generated from file names, it also maps
+// minus signs to underscores.
 func DotToUnderscore(rune int) int {
-	if rune == '.' {
+	switch rune {
+	case '.', '-':
 		return '_'
 	}
 	return rune
 }
+
+// BaseName returns the last path element of the name, with the last dotted suffix removed.
+func BaseName(name string) string {
+	// First, find the last element
+	if i := strings.LastIndex(name, "/"); i >= 0 {
+		name = name[i+1:]
+	}
+	// Now drop the suffix
+	if i := strings.LastIndex(name, "."); i >= 0 {
+		name = name[0:i]
+	}
+	return name
+}

compiler/main.go

@@ -66,6 +66,8 @@ func main() {
 		g.Fail("no files to generate")
 	}
 
+	g.CommandLineParameters(proto.GetString(g.Request.Parameter))
+
 	// Create a wrapped version of the Descriptors and EnumDescriptors that
 	// point to the file that defines them.
 	g.WrapTypes()