Fix type names (#110)

raviqqe · May 8, 2023 · fcb3ed7 · fcb3ed7
1 parent 64630aa
commit fcb3ed7
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/macro/src/type_check_functions.rs b/macro/src/type_check_functions.rs
@@ -4,15 +4,19 @@ use proc_macro2::Ident;
 use quote::quote;
 use std::error::Error;
 
+const FLOAT_8E5M2_PATTERN: &str = "8_e_5_m_2";
+
 pub fn generate(identifiers: &[Ident]) -> Result<TokenStream, Box<dyn Error>> {
     let mut stream = TokenStream::new();
 
     for identifier in identifiers {
-        let name = identifier
-            .to_string()
-            .strip_prefix("mlirTypeIsA")
-            .unwrap()
-            .to_case(Case::Snake);
+        let name = map_type_name(
+            &identifier
+                .to_string()
+                .strip_prefix("mlirTypeIsA")
+                .unwrap()
+                .to_case(Case::Snake),
+        );
 
         let function_name = Ident::new(&format!("is_{}", &name), identifier.span());
         let document = format!(" Returns `true` if a type is `{}`.", name);
@@ -27,3 +31,10 @@ pub fn generate(identifiers: &[Ident]) -> Result<TokenStream, Box<dyn Error>> {
 
     Ok(stream)
 }
+
+fn map_type_name(name: &str) -> String {
+    match name {
+        "bf_16" | "f_16" | "f_32" | "f_64" => name.replace('_', ""),
+        name => name.replace(FLOAT_8E5M2_PATTERN, &FLOAT_8E5M2_PATTERN.replace('_', "")),
+    }
+}