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#[derive(Debug)]
14pub struct Module<'c> {
15 raw: MlirModule,
16 _context: PhantomData<&'c Context>,
17}
18
19impl<'c> Module<'c> {
20 pub fn new(location: Location) -> Self {
22 unsafe { Self::from_raw(mlirModuleCreateEmpty(location.to_raw())) }
23 }
24
25 pub fn parse(context: &Context, source: &str) -> Option<Self> {
27 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 pub fn as_operation(&self) -> OperationRef<'c, '_> {
37 unsafe { OperationRef::from_raw(mlirModuleGetOperation(self.raw)) }
38 }
39
40 pub fn as_operation_mut(&mut self) -> OperationRefMut<'c, '_> {
42 unsafe { OperationRefMut::from_raw(mlirModuleGetOperation(self.raw)) }
43 }
44
45 pub fn context(&self) -> ContextRef<'c> {
47 unsafe { ContextRef::from_raw(mlirModuleGetContext(self.raw)) }
48 }
49
50 pub fn body(&self) -> BlockRef<'c, '_> {
52 unsafe { BlockRef::from_raw(mlirModuleGetBody(self.raw)) }
53 }
54
55 pub fn from_operation(operation: Operation) -> Option<Self> {
57 unsafe { Self::from_option_raw(mlirModuleFromOperation(operation.into_raw())) }
58 }
59
60 pub unsafe fn from_raw(raw: MlirModule) -> Self {
66 Self {
67 raw,
68 _context: Default::default(),
69 }
70 }
71
72 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 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}