From beb8f1c8c8d6b696cfd137f429a6c74cc95b8b1b Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Fri, 16 Oct 2020 16:45:12 -0700
Subject: [PATCH] More cleanup

---
 rust/compiler-ext/src/lib.rs | 22 ++++++++--------------
 rust/tvm/src/lib.rs          |  7 ++++---
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/rust/compiler-ext/src/lib.rs b/rust/compiler-ext/src/lib.rs
index 346f40fa2ce0..5f83f7b3b8ef 100644
--- a/rust/compiler-ext/src/lib.rs
+++ b/rust/compiler-ext/src/lib.rs
@@ -18,25 +18,19 @@
  */
 
 use env_logger;
-use tvm;
-use tvm::runtime::function::register_override;
+use tvm::export;
 
-fn test_fn() -> Result<(), tvm::Error> {
-    println!("Hello Greg from Rust!");
-    Ok(())
+fn diagnostics() -> Result<(), tvm::Error> {
+    tvm::ir::diagnostics::codespan::init()
 }
 
-fn test_fn2(message: tvm::runtime::string::String) -> Result<(), tvm::Error> {
-    println!("The message: {}", message);
-    Ok(())
-}
-
-tvm::export!(test_fn, test_fn2);
+export!(diagnostics);
 
 #[no_mangle]
-fn compiler_ext_initialize() -> i32 {
+extern fn compiler_ext_initialize() -> i32 {
     let _ = env_logger::try_init();
-    tvm_export("rust_ext").expect("failed to initialize Rust compiler_ext");
-    log::debug!("done!");
+    tvm_export("rust_ext")
+        .expect("failed to initialize the Rust compiler extensions.");
+    log::debug!("Loaded the Rust compiler extension.");
     return 0;
 }
diff --git a/rust/tvm/src/lib.rs b/rust/tvm/src/lib.rs
index d193f09803ae..ec80ece1e37a 100644
--- a/rust/tvm/src/lib.rs
+++ b/rust/tvm/src/lib.rs
@@ -50,10 +50,11 @@ pub use runtime::version;
 
 #[macro_export]
 macro_rules! export {
-    ($($fn_names:expr),*) => {
+    ($($fn_name:expr),*) => {
         pub fn tvm_export(ns: &str) -> Result<(), tvm::Error> {
             $(
-                register_override($fn_name, concat!($ns, stringfy!($fn_name)), true)?;
+                let name = String::from(ns) + ::std::stringify!($fn_name);
+                tvm::runtime::function::register_override($fn_name, name, true)?;
             )*
             Ok(())
         }
@@ -65,7 +66,7 @@ macro_rules! export_mod {
     ($ns:expr, $($mod_name:expr),*) => {
         pub fn tvm_mod_export() -> Result<(), tvm::Error> {
             $(
-                $mod_names::tvm_export($ns)?;
+                $mod_name::tvm_export($ns)?;
             )*
             Ok(())
         }