From 2de799737bcd9aedd55ccb15bc97df188460e287 Mon Sep 17 00:00:00 2001
From: Yota Toyama <raviqqe@gmail.com>
Date: Wed, 18 Oct 2023 16:23:46 +1100
Subject: [PATCH] Fix

---
 melior/src/diagnostic.rs       |  4 +++-
 melior/src/execution_engine.rs | 16 +++++++--------
 melior/src/ir/module.rs        | 16 +++++++++------
 melior/src/pass/manager.rs     | 37 ++++++++++++++++------------------
 4 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/melior/src/diagnostic.rs b/melior/src/diagnostic.rs
index 7ea1905167..4978de291a 100644
--- a/melior/src/diagnostic.rs
+++ b/melior/src/diagnostic.rs
@@ -79,10 +79,12 @@ impl<'a> Display for Diagnostic<'a> {
 #[cfg(test)]
 mod tests {
     use crate::{ir::Module, Context};
+    use std::ffi::CString;
 
     #[test]
     fn handle_diagnostic() {
         let mut message = None;
+        let source = CString::new("foo").unwrap();
         let context = Context::new();
 
         context.attach_diagnostic_handler(|diagnostic| {
@@ -90,7 +92,7 @@ mod tests {
             true
         });
 
-        Module::parse(&context, "foo");
+        Module::parse(&context, &source);
 
         assert_eq!(
             message.unwrap(),
diff --git a/melior/src/execution_engine.rs b/melior/src/execution_engine.rs
index 367ba8a240..f1c44db0c0 100644
--- a/melior/src/execution_engine.rs
+++ b/melior/src/execution_engine.rs
@@ -98,13 +98,11 @@ impl Drop for ExecutionEngine {
 mod tests {
     use super::*;
     use crate::{pass, test::create_test_context};
+    use std::ffi::CString;
 
     #[test]
     fn invoke_packed() {
-        let context = create_test_context();
-
-        let mut module = Module::parse(
-            &context,
+        let source = CString::new(
             r#"
             module {
                 func.func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
@@ -115,6 +113,9 @@ mod tests {
             "#,
         )
         .unwrap();
+        let context = create_test_context();
+
+        let mut module = Module::parse(&context, &source).unwrap();
 
         let pass_manager = pass::PassManager::new(&context);
         pass_manager.add_pass(pass::conversion::create_func_to_llvm());
@@ -150,10 +151,7 @@ mod tests {
 
     #[test]
     fn dump_to_object_file() {
-        let context = create_test_context();
-
-        let mut module = Module::parse(
-            &context,
+        let source = CString::new(
             r#"
             module {
                 func.func @add(%arg0 : i32) -> i32 {
@@ -164,7 +162,9 @@ mod tests {
             "#,
         )
         .unwrap();
+        let context = create_test_context();
 
+        let mut module = Module::parse(&context, &source).unwrap();
         let pass_manager = pass::PassManager::new(&context);
         pass_manager.add_pass(pass::conversion::create_func_to_llvm());
 
diff --git a/melior/src/ir/module.rs b/melior/src/ir/module.rs
index 0bf8d654c6..644b5e1a4e 100644
--- a/melior/src/ir/module.rs
+++ b/melior/src/ir/module.rs
@@ -7,7 +7,7 @@ use mlir_sys::{
     mlirModuleCreateEmpty, mlirModuleCreateParse, mlirModuleDestroy, mlirModuleFromOperation,
     mlirModuleGetBody, mlirModuleGetContext, mlirModuleGetOperation, MlirModule,
 };
-use std::{ffi::CString, marker::PhantomData};
+use std::{ffi::CStr, marker::PhantomData};
 
 /// A module.
 #[derive(Debug)]
@@ -23,9 +23,8 @@ impl<'c> Module<'c> {
     }
 
     /// Parses a module.
-    pub fn parse(context: &Context, source: &str) -> Option<Self> {
-        let source = CString::new(source).unwrap();
-        let source = StringRef::from_c_str(&source);
+    pub fn parse(context: &Context, source: &'c CStr) -> Option<Self> {
+        let source = StringRef::from_c_str(source);
 
         unsafe { Self::from_option_raw(mlirModuleCreateParse(context.to_raw(), source.to_raw())) }
     }
@@ -94,6 +93,7 @@ mod tests {
         ir::{operation::OperationBuilder, Block, Region},
         test::create_test_context,
     };
+    use std::ffi::CString;
 
     #[test]
     fn new() {
@@ -107,12 +107,16 @@ mod tests {
 
     #[test]
     fn parse() {
-        assert!(Module::parse(&Context::new(), "module{}").is_some());
+        let source = CString::new("module{}").unwrap();
+
+        assert!(Module::parse(&Context::new(), &source).is_some());
     }
 
     #[test]
     fn parse_none() {
-        assert!(Module::parse(&Context::new(), "module{").is_none());
+        let source = CString::new("module{").unwrap();
+
+        assert!(Module::parse(&Context::new(), &source).is_none());
     }
 
     #[test]
diff --git a/melior/src/pass/manager.rs b/melior/src/pass/manager.rs
index f2bceb921e..1930acc058 100644
--- a/melior/src/pass/manager.rs
+++ b/melior/src/pass/manager.rs
@@ -87,6 +87,7 @@ mod tests {
     };
     use indoc::indoc;
     use pretty_assertions::assert_eq;
+    use std::ffi::CString;
 
     #[test]
     fn new() {
@@ -130,20 +131,18 @@ mod tests {
 
     #[test]
     fn run_on_function() {
+        let source = CString::new(indoc!(
+            "
+            func.func @foo(%arg0 : i32) -> i32 {
+                %res = arith.addi %arg0, %arg0 : i32
+                return %res : i32
+            }
+            "
+        ))
+        .unwrap();
         let context = create_test_context();
 
-        let mut module = Module::parse(
-            &context,
-            indoc!(
-                "
-                func.func @foo(%arg0 : i32) -> i32 {
-                    %res = arith.addi %arg0, %arg0 : i32
-                    return %res : i32
-                }
-                "
-            ),
-        )
-        .unwrap();
+        let mut module = Module::parse(&context, &source).unwrap();
 
         let manager = PassManager::new(&context);
         manager.add_pass(pass::transform::create_print_op_stats());
@@ -153,12 +152,8 @@ mod tests {
 
     #[test]
     fn run_on_function_in_nested_module() {
-        let context = create_test_context();
-
-        let mut module = Module::parse(
-            &context,
-            indoc!(
-                "
+        let source = CString::new(indoc!(
+            "
                 func.func @foo(%arg0 : i32) -> i32 {
                     %res = arith.addi %arg0, %arg0 : i32
                     return %res : i32
@@ -171,9 +166,11 @@ mod tests {
                     }
                 }
                 "
-            ),
-        )
+        ))
         .unwrap();
+        let context = create_test_context();
+
+        let mut module = Module::parse(&context, &source).unwrap();
 
         let manager = PassManager::new(&context);
         manager