Fix macro error handling (#811)

This is a follow-up for #807.

Author: raviqqe

macro/build.rs

@@ -12,6 +12,9 @@ fn main() -> Result<(), Box<dyn Error>> {
         llvm_config("--includedir", &version_variable)?
     );
 
+    // For testing
+    println!("cargo:rustc-env=ENV_VAR_FOR_MLIR_DIALECT_DIR=foo/bar/baz");
+
     Ok(())
 }
 

macro/src/dialect.rs

@@ -31,32 +31,14 @@ pub fn generate_dialect(input: DialectInput) -> Result<TokenStream, Error> {
 
     parser = parser.add_include_directory(LLVM_INCLUDE_DIRECTORY);
 
-    let get_path = |path: &str| {
-        let path = if matches!(
-            Path::new(path).components().next(),
-            Some(Component::CurDir | Component::ParentDir)
-        ) {
-            path.into()
-        } else {
-            Path::new(LLVM_INCLUDE_DIRECTORY).join(path)
-        };
-        path.display().to_string()
-    };
-
     for path in input.directories() {
-        let path = get_path(path);
-        parser = parser.add_include_directory(&path);
+        parser = parser.add_include_directory(&resolve_include_directory(path));
     }
 
     for (env_var, span) in input.directory_env_vars() {
-        let path = match env::var(env_var) {
-            Ok(path) => path,
-            Err(err) => {
-                return Err(syn::Error::new(span.clone(), err.to_string()).into())
-            }
-        };
-        let path = get_path(&path);
-        parser = parser.add_include_directory(&path);
+        parser = parser.add_include_directory(&resolve_include_directory(
+            &env::var(env_var).map_err(|error| syn::Error::new(*span, error.to_string()))?,
+        ));
     }
 
     if input.files().count() > 0 {
@@ -112,6 +94,19 @@ fn generate_dialect_module(
     })
 }
 
+fn resolve_include_directory(path: &str) -> String {
+    if matches!(
+        Path::new(path).components().next(),
+        Some(Component::CurDir | Component::ParentDir)
+    ) {
+        path.into()
+    } else {
+        Path::new(LLVM_INCLUDE_DIRECTORY).join(path)
+    }
+    .display()
+    .to_string()
+}
+
 fn create_syn_error(error: impl Display) -> syn::Error {
     syn::Error::new(Span::call_site(), format!("{error}"))
 }

macro/src/dialect/error.rs

@@ -1,14 +1,14 @@
 mod ods;
 
 pub use self::ods::OdsError;
+use proc_macro::TokenStream;
+use quote::quote;
 use std::{
     error,
     fmt::{self, Display, Formatter},
     io,
     string::FromUtf8Error,
 };
-use quote::quote;
-use proc_macro::TokenStream;
 use tblgen::{
     SourceInfo,
     error::{SourceError, TableGenError},

macro/src/dialect/input.rs

@@ -1,8 +1,8 @@
 mod input_field;
 
 use self::input_field::InputField;
-use std::ops::Deref;
 use proc_macro2::Span;
+use std::ops::Deref;
 use syn::{Token, parse::Parse, punctuated::Punctuated};
 
 pub struct DialectInput {
@@ -26,9 +26,9 @@ impl DialectInput {
     }
 
     pub fn directory_env_vars(&self) -> impl Iterator<Item = (&str, &Span)> {
-        self.directory_env_vars.iter().map(|(value, span)|{
-            (value.deref(), span)
-        })
+        self.directory_env_vars
+            .iter()
+            .map(|(value, span)| (value.deref(), span))
     }
 }
 
@@ -49,7 +49,10 @@ impl Parse for DialectInput {
                     directories = field.into_iter().map(|literal| literal.value()).collect()
                 }
                 InputField::DirectoryEnvVars(field) => {
-                    directory_env_vars = field.into_iter().map(|literal| (literal.value(), literal.span())).collect()
+                    directory_env_vars = field
+                        .into_iter()
+                        .map(|literal| (literal.value(), literal.span()))
+                        .collect()
                 }
             }
         }

macro/src/lib.rs

@@ -28,9 +28,7 @@ use syn::parse_macro_input;
 #[proc_macro]
 pub fn dialect(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DialectInput);
-    dialect::generate_dialect(input).unwrap_or_else(|error| {
-        error.to_compile_error().into()
-    })
+    dialect::generate_dialect(input).unwrap_or_else(|error| error.to_compile_error())
 }
 
 #[proc_macro]