From f23cc779705f5e90e326bc0e25240588bcb5ecce Mon Sep 17 00:00:00 2001
From: Andrew J Westlake <kelixes@gmail.com>
Date: Mon, 25 Sep 2017 20:56:24 -0500
Subject: [PATCH] corrected find_enum and find_message for keyword identifiers

---
 codegen/src/keywords.rs | 16 ++++++++++++++++
 codegen/src/types.rs    | 20 ++++++++++----------
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/codegen/src/keywords.rs b/codegen/src/keywords.rs
index b18bc223..3ad4ef9b 100644
--- a/codegen/src/keywords.rs
+++ b/codegen/src/keywords.rs
@@ -89,3 +89,19 @@ pub fn sanitize_keyword(ident: &mut String) {
         }).collect::<Vec<_>>().join(".");
     }
 }
+
+// get the correct identifier (sanitize if needed, otherwise do nothing)
+// (disclaimer) I'm new to rust, so this can probably be better
+pub fn get_ident(ident: &String) -> String {
+    if !ident.contains('.') {
+        if RUST_KEYWORDS.contains(&&ident[..]) {
+            ident.clone() + &"_pb".to_string()
+        }
+        else {
+            ident.clone()
+        }
+    }
+    else {
+        ident.split('.').map(|s| get_ident(&s.to_string())).collect::<Vec<_>>().join(".")
+    }
+}
diff --git a/codegen/src/types.rs b/codegen/src/types.rs
index 200c1a9f..fa83750d 100644
--- a/codegen/src/types.rs
+++ b/codegen/src/types.rs
@@ -4,7 +4,7 @@ use std::fs::File;
 
 use errors::{Result, Error, ErrorKind};
 use parser::file_descriptor;
-use keywords::sanitize_keyword;
+use keywords::{get_ident, sanitize_keyword};
 
 fn sizeof_varint(v: u32) -> usize {
     match v {
@@ -161,11 +161,11 @@ impl FieldType {
             FieldType::Message(ref m) => {
                 let mut found = match m.rfind('.') {
                     Some(p) => {
-                        let package = &m[..p];
-                        let name = &m[(p + 1)..];
-                        msgs.iter().find(|m| m.package == package && m.name == name)
+                        let package = get_ident(&String::from(&m[..p]));
+                        let name = get_ident(&String::from(&m[(p + 1)..]));
+                        msgs.iter().find(|m| get_ident(&m.package) == package && get_ident(&m.name) == name)
                     },
-                    None => msgs.iter().find(|m2| m2.name == &m[..]),
+                    None => msgs.iter().find(|m2| get_ident(&m2.name) == get_ident(&m)),
                 };
 
                 if found.is_none() {
@@ -188,11 +188,11 @@ impl FieldType {
             FieldType::Enum(ref m) => {
                 let mut found = match m.rfind('.') {
                     Some(p) => {
-                        let package = &m[..p];
-                        let name = &m[(p + 1)..];
-                        enums.iter().find(|m| m.package == package && m.name == name)
+                        let package = get_ident(&String::from(&m[..p]));
+                        let name = get_ident(&String::from(&m[(p + 1)..]));
+                        enums.iter().find(|m| get_ident(&m.package) == package && get_ident(&m.name) == name)
                     },
-                    None => enums.iter().find(|m2| m2.name == &m[..]),
+                    None => enums.iter().find(|m2| get_ident(&m2.name) == get_ident(&m)),
                 };
 
                 if found.is_none() {
@@ -1433,7 +1433,7 @@ fn update_mod_file(path: &Path) -> Result<()> {
     let mut file = path.to_path_buf();
     use std::fs::OpenOptions;
     use std::io::prelude::*;
-    
+
     let name = file.file_stem().unwrap().to_string_lossy().to_string();
     file.pop();
     file.push("mod.rs");