melior/dialect/
func.rs

1//! `func` dialect.
2
3use crate::{
4    ir::{
5        attribute::{FlatSymbolRefAttribute, StringAttribute, TypeAttribute},
6        operation::OperationBuilder,
7        r#type::FunctionType,
8        Attribute, Identifier, Location, Operation, Region, Type, Value,
9    },
10    Context,
11};
12
13/// Create a `func.call` operation.
14pub fn call<'c>(
15    context: &'c Context,
16    function: FlatSymbolRefAttribute<'c>,
17    arguments: &[Value<'c, '_>],
18    result_types: &[Type<'c>],
19    location: Location<'c>,
20) -> Operation<'c> {
21    OperationBuilder::new("func.call", location)
22        .add_attributes(&[(Identifier::new(context, "callee"), function.into())])
23        .add_operands(arguments)
24        .add_results(result_types)
25        .build()
26        .expect("valid operation")
27}
28
29/// Create a `func.call_indirect` operation.
30pub fn call_indirect<'c>(
31    function: Value<'c, '_>,
32    arguments: &[Value<'c, '_>],
33    result_types: &[Type<'c>],
34    location: Location<'c>,
35) -> Operation<'c> {
36    OperationBuilder::new("func.call_indirect", location)
37        .add_operands(&[function])
38        .add_operands(arguments)
39        .add_results(result_types)
40        .build()
41        .expect("valid operation")
42}
43
44/// Create a `func.constant` operation.
45pub fn constant<'c>(
46    context: &'c Context,
47    function: FlatSymbolRefAttribute<'c>,
48    r#type: FunctionType<'c>,
49    location: Location<'c>,
50) -> Operation<'c> {
51    OperationBuilder::new("func.constant", location)
52        .add_attributes(&[(Identifier::new(context, "value"), function.into())])
53        .add_results(&[r#type.into()])
54        .build()
55        .expect("valid operation")
56}
57
58/// Create a `func.func` operation.
59pub fn func<'c>(
60    context: &'c Context,
61    name: StringAttribute<'c>,
62    r#type: TypeAttribute<'c>,
63    region: Region<'c>,
64    attributes: &[(Identifier<'c>, Attribute<'c>)],
65    location: Location<'c>,
66) -> Operation<'c> {
67    OperationBuilder::new("func.func", location)
68        .add_attributes(&[
69            (Identifier::new(context, "sym_name"), name.into()),
70            (Identifier::new(context, "function_type"), r#type.into()),
71        ])
72        .add_attributes(attributes)
73        .add_regions([region])
74        .build()
75        .expect("valid operation")
76}
77
78/// Create a `func.return` operation.
79pub fn r#return<'c>(operands: &[Value<'c, '_>], location: Location<'c>) -> Operation<'c> {
80    OperationBuilder::new("func.return", location)
81        .add_operands(operands)
82        .build()
83        .expect("valid operation")
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::{
90        ir::{Block, Module, Type},
91        test::create_test_context,
92    };
93
94    #[test]
95    fn compile_call() {
96        let context = create_test_context();
97
98        let location = Location::unknown(&context);
99        let module = Module::new(location);
100        let index_type = Type::index(&context);
101        let function_type = FunctionType::new(&context, &[index_type], &[index_type]);
102
103        let function = func(
104            &context,
105            StringAttribute::new(&context, "foo"),
106            TypeAttribute::new(function_type.into()),
107            {
108                let block = Block::new(&[(index_type, location)]);
109
110                let value = block
111                    .append_operation(call(
112                        &context,
113                        FlatSymbolRefAttribute::new(&context, "foo"),
114                        &[block.argument(0).unwrap().into()],
115                        &[index_type],
116                        location,
117                    ))
118                    .result(0)
119                    .unwrap()
120                    .into();
121                block.append_operation(r#return(&[value], location));
122
123                let region = Region::new();
124                region.append_block(block);
125                region
126            },
127            &[],
128            Location::unknown(&context),
129        );
130
131        module.body().append_operation(function);
132
133        assert!(module.as_operation().verify());
134        insta::assert_snapshot!(module.as_operation());
135    }
136
137    #[test]
138    fn compile_call_indirect() {
139        let context = create_test_context();
140
141        let location = Location::unknown(&context);
142        let module = Module::new(location);
143        let index_type = Type::index(&context);
144        let function_type = FunctionType::new(&context, &[index_type], &[index_type]);
145
146        let function = func(
147            &context,
148            StringAttribute::new(&context, "foo"),
149            TypeAttribute::new(function_type.into()),
150            {
151                let block = Block::new(&[(index_type, location)]);
152
153                let function = block.append_operation(constant(
154                    &context,
155                    FlatSymbolRefAttribute::new(&context, "foo"),
156                    function_type,
157                    location,
158                ));
159                let value = block
160                    .append_operation(call_indirect(
161                        function.result(0).unwrap().into(),
162                        &[block.argument(0).unwrap().into()],
163                        &[index_type],
164                        location,
165                    ))
166                    .result(0)
167                    .unwrap()
168                    .into();
169                block.append_operation(r#return(&[value], location));
170
171                let region = Region::new();
172                region.append_block(block);
173                region
174            },
175            &[],
176            Location::unknown(&context),
177        );
178
179        module.body().append_operation(function);
180
181        assert!(module.as_operation().verify());
182        insta::assert_snapshot!(module.as_operation());
183    }
184
185    #[test]
186    fn compile_function() {
187        let context = create_test_context();
188
189        let location = Location::unknown(&context);
190        let module = Module::new(location);
191
192        let integer_type = Type::index(&context);
193
194        let function = {
195            let block = Block::new(&[(integer_type, location)]);
196
197            block.append_operation(r#return(&[block.argument(0).unwrap().into()], location));
198
199            let region = Region::new();
200            region.append_block(block);
201
202            func(
203                &context,
204                StringAttribute::new(&context, "foo"),
205                TypeAttribute::new(
206                    FunctionType::new(&context, &[integer_type], &[integer_type]).into(),
207                ),
208                region,
209                &[],
210                Location::unknown(&context),
211            )
212        };
213
214        module.body().append_operation(function);
215
216        assert!(module.as_operation().verify());
217        insta::assert_snapshot!(module.as_operation());
218    }
219
220    #[test]
221    fn compile_external_function() {
222        let context = create_test_context();
223
224        let location = Location::unknown(&context);
225        let module = Module::new(location);
226
227        let integer_type = Type::index(&context);
228
229        let function = func(
230            &context,
231            StringAttribute::new(&context, "foo"),
232            TypeAttribute::new(
233                FunctionType::new(&context, &[integer_type], &[integer_type]).into(),
234            ),
235            Region::new(),
236            &[(
237                Identifier::new(&context, "sym_visibility"),
238                StringAttribute::new(&context, "private").into(),
239            )],
240            Location::unknown(&context),
241        );
242
243        module.body().append_operation(function);
244
245        assert!(module.as_operation().verify());
246        insta::assert_snapshot!(module.as_operation());
247    }
248}
melior/dialect/func.rs

melior/dialect/
func.rs
