add missing codegen rules

stepancheg · Apr 15, 2019 · 1a0c2f4 · 1a0c2f4
1 parent dde2c0b
commit 1a0c2f4
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 15 deletions.
diff --git a/protobuf-codegen/src/enums.rs b/protobuf-codegen/src/enums.rs
@@ -58,16 +58,21 @@ impl<'a> EnumGen<'a> {
         customize: &Customize,
         _root_scope: &RootScope,
     ) -> EnumGen<'a> {
-        EnumGen {
-            enum_with_scope,
-            type_name: enum_with_scope.rust_name().to_path(),
-            lite_runtime: enum_with_scope
+
+        let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
+            enum_with_scope
                 .get_scope()
                 .get_file_descriptor()
                 .options
                 .get_message()
                 .get_optimize_for()
-                == file_options::OptimizeMode::LITE_RUNTIME,
+                == file_options::OptimizeMode::LITE_RUNTIME
+        });
+
+        EnumGen {
+            enum_with_scope,
+            type_name: enum_with_scope.rust_name().to_path(),
+            lite_runtime,
             customize: customize.clone(),
         }
     }

diff --git a/protobuf-codegen/src/lib.rs b/protobuf-codegen/src/lib.rs
@@ -146,6 +146,13 @@ fn gen_file(
     let scope = FileScope {
         file_descriptor: file,
     }.to_scope();
+    let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
+        file
+            .options
+            .get_message()
+            .get_optimize_for()
+            == file_options::OptimizeMode::LITE_RUNTIME
+    });
 
     let mut v = Vec::new();
 
@@ -168,8 +175,7 @@ fn gen_file(
 
         write_extensions(file, &root_scope, &mut w);
 
-        let optimize_mode = file.options.get_message().get_optimize_for();
-        if optimize_mode != file_options::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
@@ -49,19 +49,20 @@ 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()
+                .options
+                .get_message()
+                .get_optimize_for()
+                == file_options::OptimizeMode::LITE_RUNTIME
+        });
         MessageGen {
             message: message,
             root_scope: root_scope,
             type_name: message.rust_name().to_path(),
             fields: fields,
-            lite_runtime: customize.lite_runtime.unwrap_or_else(|| {
-                message
-                    .get_file_descriptor()
-                    .options
-                    .get_message()
-                    .get_optimize_for()
-                    == FileOptions_OptimizeMode::LITE_RUNTIME
-            }),
+            lite_runtime,
             customize,
         }
     }