Skip to main content

melior/ir/
module.rs

1use super::{operation::OperationRefMut, BlockRef, Location, Operation, OperationRef};
2use crate::{
3    context::{Context, ContextRef},
4    string_ref::StringRef,
5};
6use mlir_sys::{
7    mlirModuleCreateEmpty, mlirModuleCreateParse, mlirModuleDestroy, mlirModuleFromOperation,
8    mlirModuleGetBody, mlirModuleGetContext, mlirModuleGetOperation, MlirModule,
9};
10use std::{ffi::CString, marker::PhantomData};
11
12/// A module.
13#[derive(Debug)]
14pub struct Module<'c> {
15    raw: MlirModule,
16    _context: PhantomData<&'c Context>,
17}
18
19impl<'c> Module<'c> {
20    /// Creates a module.
21    pub fn new(location: Location) -> Self {
22        unsafe { Self::from_raw(mlirModuleCreateEmpty(location.to_raw())) }
23    }
24
25    /// Parses a module.
26    pub fn parse(context: &Context, source: &str) -> Option<Self> {
27        // TODO Use a string not null-terminated.
28        // Somehow, we still need a null-terminated string for a source.
29        let source = CString::new(source).unwrap();
30        let source = StringRef::from_c_str(&source);
31
32        unsafe { Self::from_option_raw(mlirModuleCreateParse(context.to_raw(), source.to_raw())) }
33    }
34
35    /// Converts a module into an operation.
36    pub fn as_operation(&self) -> OperationRef<'c, '_> {
37        unsafe { OperationRef::from_raw(mlirModuleGetOperation(self.raw)) }
38    }
39
40    /// Converts a module into a mutable operation.
41    pub fn as_operation_mut(&mut self) -> OperationRefMut<'c, '_> {
42        unsafe { OperationRefMut::from_raw(mlirModuleGetOperation(self.raw)) }
43    }
44
45    /// Returns a context.
46    pub fn context(&self) -> ContextRef<'c> {
47        unsafe { ContextRef::from_raw(mlirModuleGetContext(self.raw)) }
48    }
49
50    /// Returns a block of a module body.
51    pub fn body(&self) -> BlockRef<'c, '_> {
52        unsafe { BlockRef::from_raw(mlirModuleGetBody(self.raw)) }
53    }
54
55    /// Converts an operation into a module.
56    pub fn from_operation(operation: Operation) -> Option<Self> {
57        unsafe { Self::from_option_raw(mlirModuleFromOperation(operation.into_raw())) }
58    }
59
60    /// Creates a module from a raw object.
61    ///
62    /// # Safety
63    ///
64    /// A raw object must be valid.
65    pub unsafe fn from_raw(raw: MlirModule) -> Self {
66        Self {
67            raw,
68            _context: Default::default(),
69        }
70    }
71
72    /// Creates an optional module from a raw object.
73    ///
74    /// # Safety
75    ///
76    /// A raw object must be valid.
77    pub unsafe fn from_option_raw(raw: MlirModule) -> Option<Self> {
78        if raw.ptr.is_null() {
79            None
80        } else {
81            Some(Self::from_raw(raw))
82        }
83    }
84
85    /// Converts a module into a raw object.
86    pub const fn to_raw(&self) -> MlirModule {
87        self.raw
88    }
89}
90
91impl Drop for Module<'_> {
92    fn drop(&mut self) {
93        unsafe { mlirModuleDestroy(self.raw) };
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use crate::{
101        ir::{attribute::StringAttribute, operation::OperationBuilder, Block, Region},
102        test::create_test_context,
103    };
104
105    #[test]
106    fn new() {
107        Module::new(Location::new(&Context::new(), "foo", 42, 42));
108    }
109
110    #[test]
111    fn context() {
112        Module::new(Location::new(&Context::new(), "foo", 42, 42)).context();
113    }
114
115    #[test]
116    fn parse() {
117        assert!(Module::parse(&Context::new(), "module{}").is_some());
118    }
119
120    #[test]
121    fn parse_none() {
122        assert!(Module::parse(&Context::new(), "module{").is_none());
123    }
124
125    #[test]
126    fn from_operation() {
127        let context = create_test_context();
128
129        let region = Region::new();
130        region.append_block(Block::new(&[]));
131
132        let module = Module::from_operation(
133            OperationBuilder::new("builtin.module", Location::unknown(&context))
134                .add_regions([region])
135                .build()
136                .unwrap(),
137        )
138        .unwrap();
139
140        assert!(module.as_operation().verify());
141        assert_eq!(module.as_operation().to_string(), "module {\n}\n")
142    }
143
144    #[test]
145    fn from_operation_fail() {
146        let context = create_test_context();
147
148        assert!(Module::from_operation(
149            OperationBuilder::new("func.func", Location::unknown(&context),)
150                .build()
151                .unwrap()
152        )
153        .is_none());
154    }
155
156    #[test]
157    fn set_attribute() {
158        let context = create_test_context();
159
160        let mut module = Module::new(Location::unknown(&context));
161
162        module
163            .as_operation_mut()
164            .set_attribute("sym_name", StringAttribute::new(&context, "foo").into());
165
166        assert!(module.as_operation().verify());
167    }
168}