Fix custom options behaviour

6 files changed, 161 insertions, 38 deletions

diff --git a/src/ProtoGen/EnumGenerator.cs b/src/ProtoGen/EnumGenerator.cs
index aaed53c5..65eeb7d3 100644
--- a/src/ProtoGen/EnumGenerator.cs
+++ b/src/ProtoGen/EnumGenerator.cs
@@ -5,6 +5,7 @@ namespace Google.ProtocolBuffers.ProtoGen {
     internal EnumGenerator(EnumDescriptor descriptor) : base(descriptor) {
     }
 
+    // TODO(jonskeet): Write out enum descriptors? Can be retrieved from file...
     public void Generate(TextGenerator writer) {
       writer.WriteLine("{0} enum {1} {{", ClassAccessLevel, Descriptor.Name);
       writer.Indent();
diff --git a/src/ProtoGen/ExtensionGenerator.cs b/src/ProtoGen/ExtensionGenerator.cs
index c3412fe9..6b44ed73 100644
--- a/src/ProtoGen/ExtensionGenerator.cs
+++ b/src/ProtoGen/ExtensionGenerator.cs
@@ -1,17 +1,18 @@
-using System;
-using System.Collections.Generic;
-using System.Text;
 using Google.ProtocolBuffers.Descriptors;
 
 namespace Google.ProtocolBuffers.ProtoGen {
   internal class ExtensionGenerator : SourceGeneratorBase<FieldDescriptor>, ISourceGenerator {
-    internal ExtensionGenerator(FieldDescriptor descriptor) : base(descriptor) {
-    }
 
-    public void Generate(TextGenerator writer) {
-      string name = Descriptor.CSharpOptions.PropertyName;
+    private readonly string scope;
+    private readonly string type;
+    private readonly string name;
 
-      string type;
+    internal ExtensionGenerator(FieldDescriptor descriptor) : base(descriptor) {
+      if (Descriptor.ExtensionScope != null) {
+        scope = GetClassName(Descriptor.ExtensionScope);
+      } else {
+        scope = DescriptorUtil.GetFullUmbrellaClassName(Descriptor.File);
+      }
       switch (Descriptor.MappedType) {
         case MappedType.Message:
           type = GetClassName(Descriptor.MessageType);
@@ -23,16 +24,28 @@ namespace Google.ProtocolBuffers.ProtoGen {
           type = DescriptorUtil.GetMappedTypeName(Descriptor.MappedType);
           break;
       }
+      name = Descriptor.CSharpOptions.PropertyName;
+    }
 
+    public void Generate(TextGenerator writer) {
       writer.WriteLine ("public const int {0} = {1};", GetFieldConstantName(Descriptor), Descriptor.FieldNumber);
       if (Descriptor.IsRepeated) {
-        writer.WriteLine("{0} static readonly", ClassAccessLevel);
-        writer.WriteLine("    pb::GeneratedExtensionBase<scg::IList<{0}>> {1} =", type, name);
-        writer.WriteLine("    pb::GeneratedRepeatExtension<{0}>.CreateInstance(Descriptor.Extensions[{1}]);", type, Descriptor.Index);
+        writer.WriteLine("{0} static pb::GeneratedExtensionBase<scg::IList<{1}>> {2};", ClassAccessLevel, type, name);
       } else {
-        writer.WriteLine("{0} static readonly pb::GeneratedExtensionBase<{1}> {2} =", ClassAccessLevel, type, name);
-        writer.WriteLine("    pb::GeneratedSingleExtension<{0}>.CreateInstance(Descriptor.Extensions[{1}]);", type, Descriptor.Index);
+        writer.WriteLine("{0} static pb::GeneratedExtensionBase<{1}> {2};", ClassAccessLevel, type, name);
       }
     }
+
+    internal void GenerateStaticVariableInitializers(TextGenerator writer) {
+      if (Descriptor.IsRepeated) {
+        writer.WriteLine("{0}.{1} = pb::GeneratedRepeatExtension<{2}>.CreateInstance({0}.Descriptor.Extensions[{3}]);", scope, name, type, Descriptor.Index);
+      } else {
+        writer.WriteLine("{0}.{1} = pb::GeneratedSingleExtension<{2}>.CreateInstance({0}.Descriptor.Extensions[{3}]);", scope, name, type, Descriptor.Index);
+      }
+    }
+
+    internal void GenerateExtensionRegistrationCode(TextGenerator writer) {
+      writer.WriteLine("registry.Add({0}.{1});", scope, name);
+    }
   }
 }
diff --git a/src/ProtoGen/MessageGenerator.cs b/src/ProtoGen/MessageGenerator.cs
index 4b983003..53859231 100644
--- a/src/ProtoGen/MessageGenerator.cs
+++ b/src/ProtoGen/MessageGenerator.cs
@@ -39,25 +39,42 @@ namespace Google.ProtocolBuffers.ProtoGen {
 
       // The descriptor for this type.
       string access = Descriptor.File.CSharpOptions.NestClasses ? "private" : "internal";
-      writer.WriteLine("{0} static readonly pbd::MessageDescriptor internal__{1}__Descriptor", access, identifier);
+      writer.WriteLine("{0} static pbd::MessageDescriptor internal__{1}__Descriptor;", access, identifier);
+      writer.WriteLine("{0} static pb::FieldAccess.FieldAccessorTable<{1}, {1}.Builder> internal__{2}__FieldAccessorTable;",
+          access, FullClassName, identifier);
+
+      // Generate static members for all nested types.
+      foreach (MessageDescriptor nestedMessage in Descriptor.NestedTypes) {
+        new MessageGenerator(nestedMessage).GenerateStaticVariables(writer);
+      }
+    }
+
+    internal void GenerateStaticVariableInitializers(TextGenerator writer) {
+      string identifier = GetUniqueFileScopeIdentifier(Descriptor);
+
+      writer.Write("internal__{0}__Descriptor = ", identifier);
       if (Descriptor.ContainingType == null) {
-        writer.WriteLine("    = Descriptor.MessageTypes[{0}];", Descriptor.Index);
+        writer.WriteLine("Descriptor.MessageTypes[{0}];", Descriptor.Index);
       } else {
-        writer.WriteLine("    = internal__{0}__Descriptor.NestedTypes[{1}];", GetUniqueFileScopeIdentifier(Descriptor.ContainingType), Descriptor.Index);
+        writer.WriteLine("internal__{0}__Descriptor.NestedTypes[{1}];", GetUniqueFileScopeIdentifier(Descriptor.ContainingType), Descriptor.Index);
       }
-      writer.WriteLine("{0} static pb::FieldAccess.FieldAccessorTable<{1}, {1}.Builder> internal__{2}__FieldAccessorTable",
-          access, FullClassName, identifier);
-      writer.WriteLine("    = new pb::FieldAccess.FieldAccessorTable<{0}, {0}.Builder>(internal__{1}__Descriptor,",
-          FullClassName, identifier);
+
+      writer.WriteLine("internal__{0}__FieldAccessorTable = ", identifier);
+      writer.WriteLine("    new pb::FieldAccess.FieldAccessorTable<{1}, {1}.Builder>(internal__{0}__Descriptor,",
+          identifier, FullClassName);
       writer.Print("        new string[] { ");
       foreach (FieldDescriptor field in Descriptor.Fields) {
         writer.Write("\"{0}\", ", field.CSharpOptions.PropertyName);
       }
       writer.WriteLine("});");
 
-      // Generate static members for all nested types.
+      // Generate static member initializers for all nested types.
       foreach (MessageDescriptor nestedMessage in Descriptor.NestedTypes) {
-        new MessageGenerator(nestedMessage).GenerateStaticVariables(writer);
+        new MessageGenerator(nestedMessage).GenerateStaticVariableInitializers(writer);
+      }
+
+      foreach (FieldDescriptor extension in Descriptor.Extensions) {
+        new ExtensionGenerator(extension).GenerateStaticVariableInitializers(writer);
       }
     }
 
@@ -119,6 +136,17 @@ namespace Google.ProtocolBuffers.ProtoGen {
 
       GenerateParseFromMethods(writer);
       GenerateBuilder(writer);
+
+      // Force the static initialization code for the file to run, since it may
+      // initialize static variables declared in this class.
+      writer.WriteLine("static {0}() {{", ClassName);
+      // Note that the variable is needed just so we can access the property
+      writer.WriteLine("  pbd::FileDescriptor descriptor = {0}.Descriptor;", DescriptorUtil.GetFullUmbrellaClassName(Descriptor));
+      writer.WriteLine("}");
+
+      writer.Outdent();
+      writer.WriteLine("}");
+      writer.WriteLine();
     }
 
     private void GenerateMessageSerializationMethods(TextGenerator writer) {
@@ -297,9 +325,6 @@ namespace Google.ProtocolBuffers.ProtoGen {
       }
       writer.Outdent();
       writer.WriteLine("}");
-      writer.Outdent();
-      writer.WriteLine("}");
-      writer.WriteLine();
     }
 
     private void GenerateCommonBuilderMethods(TextGenerator writer) {
@@ -472,5 +497,14 @@ namespace Google.ProtocolBuffers.ProtoGen {
       writer.WriteLine("}");
       writer.WriteLine();
     }
+
+    internal void GenerateExtensionRegistrationCode(TextGenerator writer) {
+      foreach (FieldDescriptor extension in Descriptor.Extensions) {
+        new ExtensionGenerator(extension).GenerateExtensionRegistrationCode(writer);
+      }
+      foreach (MessageDescriptor nestedMessage in Descriptor.NestedTypes) {
+        new MessageGenerator(nestedMessage).GenerateExtensionRegistrationCode(writer);
+      }
+    }
   }
 }
diff --git a/src/ProtoGen/SourceFileGenerator.cs b/src/ProtoGen/SourceFileGenerator.cs
index 61321753..b5bd5efb 100644
--- a/src/ProtoGen/SourceFileGenerator.cs
+++ b/src/ProtoGen/SourceFileGenerator.cs
@@ -1,8 +1,4 @@
-using System;
-using System.Collections.Generic;
 using System.IO;
-using System.Text;
-using Google.ProtocolBuffers.Descriptors;
 
 namespace Google.ProtocolBuffers.ProtoGen {
   /// <summary>
diff --git a/src/ProtoGen/SourceGeneratorBase.cs b/src/ProtoGen/SourceGeneratorBase.cs
index dea9def2..1604f11c 100644
--- a/src/ProtoGen/SourceGeneratorBase.cs
+++ b/src/ProtoGen/SourceGeneratorBase.cs
@@ -31,7 +31,7 @@ namespace Google.ProtocolBuffers.ProtoGen {
     }
 
     internal static string GetFieldConstantName(FieldDescriptor field) {
-      return NameHelpers.UnderscoresToPascalCase(GetFieldName(field)) + "FieldNumber";
+      return field.CSharpOptions.PropertyName + "FieldNumber";
     }
 
     private static string ToCSharpName(string name, FileDescriptor file) {
diff --git a/src/ProtoGen/UmbrellaClassGenerator.cs b/src/ProtoGen/UmbrellaClassGenerator.cs
index 938507b5..732b6f8d 100644
--- a/src/ProtoGen/UmbrellaClassGenerator.cs
+++ b/src/ProtoGen/UmbrellaClassGenerator.cs
@@ -1,4 +1,6 @@
 using System;
+using System.Collections;
+using System.Collections.Generic;
 using Google.ProtocolBuffers.DescriptorProtos;
 using Google.ProtocolBuffers.Descriptors;
 
@@ -13,19 +15,49 @@ namespace Google.ProtocolBuffers.ProtoGen {
       : base(descriptor) {
     }
 
+    // Recursively searches the given message to see if it contains any extensions.
+    private static bool UsesExtensions(IMessage message) {
+      // We conservatively assume that unknown fields are extensions.
+      if (message.UnknownFields.FieldDictionary.Count > 0) {
+        return true;
+      }
+
+      foreach (KeyValuePair<FieldDescriptor, object> keyValue in message.AllFields) {
+        FieldDescriptor field = keyValue.Key;
+        if (field.IsExtension) {
+          return true;
+        }
+        if (field.MappedType == MappedType.Message) {
+          if (field.IsRepeated) {
+            foreach (IMessage subMessage in (IEnumerable)keyValue.Value) {
+              if (UsesExtensions(subMessage)) {
+                return true;
+              }
+            }
+          } else {
+            if (UsesExtensions((IMessage)keyValue.Value)) {
+              return true;
+            }
+          }
+        }
+      }
+      return false;
+    }
+
     public string UmbrellaClassName {
       get { throw new NotImplementedException(); }
     }
 
     public void Generate(TextGenerator writer) {
       WriteIntroduction(writer);
-      WriteDescriptor(writer);
+      WriteExtensionRegistration(writer);
       WriteChildren(writer, "Extensions", Descriptor.Extensions);
       writer.WriteLine("#region Static variables");
       foreach (MessageDescriptor message in Descriptor.MessageTypes) {
         new MessageGenerator(message).GenerateStaticVariables(writer);
       }
       writer.WriteLine("#endregion");
+      WriteDescriptor(writer);
       // The class declaration either gets closed before or after the children are written.
       if (!Descriptor.CSharpOptions.NestClasses) {
         writer.Outdent();
@@ -60,14 +92,32 @@ namespace Google.ProtocolBuffers.ProtoGen {
       writer.Indent();
     }
 
+    private void WriteExtensionRegistration(TextGenerator writer) {
+      writer.WriteLine("#region Extension registration");
+      writer.WriteLine("public static void RegisterAllExtensions(pb::ExtensionRegistry registry) {");
+      writer.Indent();
+      foreach (FieldDescriptor extension in Descriptor.Extensions) {
+        new ExtensionGenerator(extension).GenerateExtensionRegistrationCode(writer);
+      }
+      foreach (MessageDescriptor message in Descriptor.MessageTypes) {
+        new MessageGenerator(message).GenerateExtensionRegistrationCode(writer);
+      }
+      writer.Outdent();
+      writer.WriteLine("}");
+      writer.WriteLine("#endregion");
+    }
+
     private void WriteDescriptor(TextGenerator writer) {
       writer.WriteLine("#region Descriptor");
 
       writer.WriteLine("public static pbd::FileDescriptor Descriptor {");
       writer.WriteLine("  get { return descriptor; }");
       writer.WriteLine("}");
-      writer.WriteLine("private static readonly pbd::FileDescriptor descriptor = pbd::FileDescriptor.InternalBuildGeneratedFileFrom(");
-      writer.WriteLine("    global::System.Convert.FromBase64String(");
+      writer.WriteLine("private static pbd::FileDescriptor descriptor;");
+      writer.WriteLine();
+      writer.WriteLine("static {0}() {{", Descriptor.CSharpOptions.UmbrellaClassname);
+      writer.Indent();
+      writer.WriteLine("byte[] descriptorData = global::System.Convert.FromBase64String(");
       writer.Indent();
       writer.Indent();
 
@@ -79,15 +129,44 @@ namespace Google.ProtocolBuffers.ProtoGen {
         writer.WriteLine("\"{0}\" + ", base64.Substring(0, 60));
         base64 = base64.Substring(60);
       }
-      writer.WriteLine("\"{0}\"),", base64);
+      writer.WriteLine("\"{0}\");", base64);
+      writer.Outdent();
+      writer.Outdent();
+      writer.WriteLine("pbd::FileDescriptor.InternalDescriptorAssigner assigner = delegate(pbd::FileDescriptor root) {");
+      writer.Indent();
+      writer.WriteLine("descriptor = root;");
+      foreach (MessageDescriptor message in Descriptor.MessageTypes) {
+        new MessageGenerator(message).GenerateStaticVariableInitializers(writer);
+      }
+      foreach (FieldDescriptor extension in Descriptor.Extensions) {
+        new ExtensionGenerator(extension).GenerateStaticVariableInitializers(writer);
+      }
 
-      writer.WriteLine("new pbd::FileDescriptor[] {");
-      foreach (FileDescriptor dependency in Descriptor.Dependencies) {
-        writer.WriteLine("  {0}.Descriptor, ", DescriptorUtil.GetFullUmbrellaClassName(dependency));
+      if (UsesExtensions(Descriptor.Proto)) {
+        // Must construct an ExtensionRegistry containing all possible extensions
+        // and return it.
+        writer.WriteLine("pb::ExtensionRegistry registry = pb::ExtensionRegistry.CreateInstance();");
+        writer.WriteLine("RegisterAllExtensions(registry);");
+        foreach (FileDescriptor dependency in Descriptor.Dependencies) {
+          writer.WriteLine("{0}.RegisterAllExtensions(registry);", DescriptorUtil.GetFullUmbrellaClassName(dependency));
+        }
+        writer.WriteLine("return registry;");
+      } else {
+        writer.WriteLine("return null;");
       }
-      writer.WriteLine("});");
       writer.Outdent();
+      writer.WriteLine("};");
+
+      // -----------------------------------------------------------------
+      // Invoke internalBuildGeneratedFileFrom() to build the file.
+      writer.WriteLine("pbd::FileDescriptor.InternalBuildGeneratedFileFrom(descriptorData,");
+      writer.WriteLine("    new pbd::FileDescriptor[] {");
+      foreach (FileDescriptor dependency in Descriptor.Dependencies) {
+        writer.WriteLine("    {0}.Descriptor, ", DescriptorUtil.GetFullUmbrellaClassName(dependency));
+      }
+      writer.WriteLine("    }, assigner);");
       writer.Outdent();
+      writer.WriteLine("}");
       writer.WriteLine("#endregion");
       writer.WriteLine();
     }