add missing codegen rules

stepancheg · Apr 23, 2019 · de13d72 · de13d72
1 parent acdc94a
commit de13d72
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/protobuf-codegen/src/enums.rs b/protobuf-codegen/src/enums.rs
@@ -69,15 +69,19 @@ impl<'a> EnumGen<'a> {
                 false,
                 root_scope)
         };
-        EnumGen {
-            enum_with_scope,
-            type_name: rust_name,
-            lite_runtime: enum_with_scope
+        let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
+            enum_with_scope
                 .get_scope()
                 .get_file_descriptor()
                 .get_options()
-                .get_optimize_for() ==
-                FileOptions_OptimizeMode::LITE_RUNTIME,
+                .get_optimize_for()
+                == FileOptions_OptimizeMode::LITE_RUNTIME
+        });
+
+        EnumGen {
+            enum_with_scope,
+            type_name: rust_name,
+            lite_runtime: lite_runtime,
             customize: customize.clone(),
         }
     }

diff --git a/protobuf-codegen/src/lib.rs b/protobuf-codegen/src/lib.rs
@@ -114,6 +114,12 @@ fn gen_file(
     customize.update_with(&customize_from_rustproto_for_file(file.get_options()));
 
     let scope = FileScope { file_descriptor: file }.to_scope();
+    let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
+        file
+            .get_options()
+            .get_optimize_for()
+            == FileOptions_OptimizeMode::LITE_RUNTIME
+    });
 
     let mut v = Vec::new();
 
@@ -142,7 +148,7 @@ fn gen_file(
 
         write_extensions(file, &root_scope, &mut w);
 
-        if file.get_options().get_optimize_for() != FileOptions_OptimizeMode::LITE_RUNTIME {
+        if !lite_runtime {
             w.write_line("");
             write_file_descriptor_data(file, &mut w);
         }

diff --git a/protobuf-codegen/src/message.rs b/protobuf-codegen/src/message.rs
@@ -37,6 +37,13 @@ impl<'a> MessageGen<'a> {
             .into_iter()
             .map(|field| FieldGen::parse(field, root_scope, &customize))
             .collect();
+        let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
+            message
+                .get_file_descriptor()
+                .get_options()
+                .get_optimize_for()
+                == FileOptions_OptimizeMode::LITE_RUNTIME
+        });
         MessageGen {
             message: message,
             root_scope: root_scope,