1 files changed, 9 insertions, 5 deletions
diff --git a/src/google/protobuf/compiler/java/java_generator_factory.cc b/src/google/protobuf/compiler/java/java_generator_factory.cc
index 92ef851b..3218b410 100644
--- a/src/google/protobuf/compiler/java/java_generator_factory.cc
+++ b/src/google/protobuf/compiler/java/java_generator_factory.cc
@@ -35,6 +35,7 @@
 #include <google/protobuf/compiler/java/java_context.h>
 #include <google/protobuf/compiler/java/java_enum_field.h>
 #include <google/protobuf/compiler/java/java_extension.h>
+#include <google/protobuf/compiler/java/java_extension_lite.h>
 #include <google/protobuf/compiler/java/java_field.h>
 #include <google/protobuf/compiler/java/java_helpers.h>
 #include <google/protobuf/compiler/java/java_message.h>
@@ -58,17 +59,20 @@ ImmutableGeneratorFactory::~ImmutableGeneratorFactory() {}
 
 MessageGenerator* ImmutableGeneratorFactory::NewMessageGenerator(
     const Descriptor* descriptor) const {
-  if (descriptor->file()->options().optimize_for() ==
-      FileOptions::LITE_RUNTIME) {
-    return new ImmutableMessageLiteGenerator(descriptor, context_);
-  } else {
+  if (HasDescriptorMethods(descriptor, context_->EnforceLite())) {
     return new ImmutableMessageGenerator(descriptor, context_);
+  } else {
+    return new ImmutableMessageLiteGenerator(descriptor, context_);
   }
 }
 
 ExtensionGenerator* ImmutableGeneratorFactory::NewExtensionGenerator(
     const FieldDescriptor* descriptor) const {
-  return new ImmutableExtensionGenerator(descriptor, context_);
+  if (HasDescriptorMethods(descriptor->file(), context_->EnforceLite())) {
+    return new ImmutableExtensionGenerator(descriptor, context_);
+  } else {
+    return new ImmutableExtensionLiteGenerator(descriptor, context_);
+  }
 }
 
 ServiceGenerator* ImmutableGeneratorFactory::NewServiceGenerator(