1use crate::{
4 ir::{
5 attribute::DenseI64ArrayAttribute, operation::OperationBuilder, Identifier, Location,
6 Operation, Region, Type, Value,
7 },
8 Context,
9};
10
11pub fn condition<'c>(
13 condition: Value<'c, '_>,
14 values: &[Value<'c, '_>],
15 location: Location<'c>,
16) -> Operation<'c> {
17 OperationBuilder::new("scf.condition", location)
18 .add_operands(&[condition])
19 .add_operands(values)
20 .build()
21 .expect("valid operation")
22}
23
24pub fn execute_region<'c>(
26 result_types: &[Type<'c>],
27 region: Region<'c>,
28 location: Location<'c>,
29) -> Operation<'c> {
30 OperationBuilder::new("scf.execute_region", location)
31 .add_results(result_types)
32 .add_regions([region])
33 .build()
34 .expect("valid operation")
35}
36
37pub fn r#for<'c>(
39 start: Value<'c, '_>,
40 end: Value<'c, '_>,
41 step: Value<'c, '_>,
42 region: Region<'c>,
43 location: Location<'c>,
44) -> Operation<'c> {
45 OperationBuilder::new("scf.for", location)
46 .add_operands(&[start, end, step])
47 .add_regions([region])
48 .build()
49 .expect("valid operation")
50}
51
52pub fn r#if<'c>(
54 condition: Value<'c, '_>,
55 result_types: &[Type<'c>],
56 then_region: Region<'c>,
57 else_region: Region<'c>,
58 location: Location<'c>,
59) -> Operation<'c> {
60 OperationBuilder::new("scf.if", location)
61 .add_operands(&[condition])
62 .add_results(result_types)
63 .add_regions([then_region, else_region])
64 .build()
65 .expect("valid operation")
66}
67
68pub fn index_switch<'c>(
70 context: &'c Context,
71 condition: Value<'c, '_>,
72 result_types: &[Type<'c>],
73 cases: DenseI64ArrayAttribute<'c>,
74 regions: Vec<Region<'c>>,
75 location: Location<'c>,
76) -> Operation<'c> {
77 OperationBuilder::new("scf.index_switch", location)
78 .add_operands(&[condition])
79 .add_results(result_types)
80 .add_attributes(&[(Identifier::new(context, "cases"), cases.into())])
81 .add_regions_vec(regions)
82 .build()
83 .expect("valid operation")
84}
85
86pub fn r#while<'c>(
88 initial_values: &[Value<'c, '_>],
89 result_types: &[Type<'c>],
90 before_region: Region<'c>,
91 after_region: Region<'c>,
92 location: Location<'c>,
93) -> Operation<'c> {
94 OperationBuilder::new("scf.while", location)
95 .add_operands(initial_values)
96 .add_results(result_types)
97 .add_regions([before_region, after_region])
98 .build()
99 .expect("valid operation")
100}
101
102pub fn r#yield<'c>(values: &[Value<'c, '_>], location: Location<'c>) -> Operation<'c> {
104 OperationBuilder::new("scf.yield", location)
105 .add_operands(values)
106 .build()
107 .expect("valid operation")
108}
109
110#[cfg(test)]
111mod tests {
112 use super::*;
113 use crate::{
114 dialect::{arith, func},
115 ir::{
116 attribute::{FloatAttribute, IntegerAttribute, StringAttribute, TypeAttribute},
117 r#type::{FunctionType, IntegerType, Type},
118 Attribute, Block, Module,
119 },
120 test::load_all_dialects,
121 Context,
122 };
123
124 #[test]
125 fn compile_execute_region() {
126 let context = Context::new();
127 load_all_dialects(&context);
128
129 let location = Location::unknown(&context);
130 let module = Module::new(location);
131 let index_type = Type::index(&context);
132
133 module.body().append_operation(func::func(
134 &context,
135 StringAttribute::new(&context, "foo"),
136 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
137 {
138 let block = Block::new(&[]);
139
140 block.append_operation(execute_region(
141 &[index_type],
142 {
143 let block = Block::new(&[]);
144
145 let value = block.append_operation(arith::constant(
146 &context,
147 IntegerAttribute::new(index_type, 0).into(),
148 location,
149 ));
150
151 block.append_operation(r#yield(
152 &[value.result(0).unwrap().into()],
153 location,
154 ));
155
156 let region = Region::new();
157 region.append_block(block);
158 region
159 },
160 location,
161 ));
162
163 block.append_operation(func::r#return(&[], location));
164
165 let region = Region::new();
166 region.append_block(block);
167 region
168 },
169 &[],
170 location,
171 ));
172
173 assert!(module.as_operation().verify());
174 insta::assert_snapshot!(module.as_operation());
175 }
176
177 #[test]
178 fn compile_for() {
179 let context = Context::new();
180 load_all_dialects(&context);
181
182 let location = Location::unknown(&context);
183 let module = Module::new(location);
184
185 module.body().append_operation(func::func(
186 &context,
187 StringAttribute::new(&context, "foo"),
188 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
189 {
190 let block = Block::new(&[]);
191
192 let start = block.append_operation(arith::constant(
193 &context,
194 Attribute::parse(&context, "0 : index").unwrap(),
195 location,
196 ));
197
198 let end = block.append_operation(arith::constant(
199 &context,
200 Attribute::parse(&context, "8 : index").unwrap(),
201 location,
202 ));
203
204 let step = block.append_operation(arith::constant(
205 &context,
206 Attribute::parse(&context, "1 : index").unwrap(),
207 location,
208 ));
209
210 block.append_operation(r#for(
211 start.result(0).unwrap().into(),
212 end.result(0).unwrap().into(),
213 step.result(0).unwrap().into(),
214 {
215 let block = Block::new(&[(Type::index(&context), location)]);
216 block.append_operation(r#yield(&[], location));
217
218 let region = Region::new();
219 region.append_block(block);
220 region
221 },
222 location,
223 ));
224
225 block.append_operation(func::r#return(&[], location));
226
227 let region = Region::new();
228 region.append_block(block);
229 region
230 },
231 &[],
232 location,
233 ));
234
235 assert!(module.as_operation().verify());
236 insta::assert_snapshot!(module.as_operation());
237 }
238
239 mod r#if {
240 use super::*;
241
242 #[test]
243 fn compile() {
244 let context = Context::new();
245 load_all_dialects(&context);
246
247 let location = Location::unknown(&context);
248 let module = Module::new(location);
249 let index_type = Type::index(&context);
250
251 module.body().append_operation(func::func(
252 &context,
253 StringAttribute::new(&context, "foo"),
254 TypeAttribute::new(FunctionType::new(&context, &[], &[index_type]).into()),
255 {
256 let block = Block::new(&[]);
257
258 let condition = block.append_operation(arith::constant(
259 &context,
260 IntegerAttribute::new(IntegerType::new(&context, 1).into(), 0).into(),
261 location,
262 ));
263
264 let result = block.append_operation(r#if(
265 condition.result(0).unwrap().into(),
266 &[index_type],
267 {
268 let block = Block::new(&[]);
269
270 let result = block.append_operation(arith::constant(
271 &context,
272 IntegerAttribute::new(index_type, 42).into(),
273 location,
274 ));
275
276 block.append_operation(r#yield(
277 &[result.result(0).unwrap().into()],
278 location,
279 ));
280
281 let region = Region::new();
282 region.append_block(block);
283 region
284 },
285 {
286 let block = Block::new(&[]);
287
288 let result = block.append_operation(arith::constant(
289 &context,
290 IntegerAttribute::new(index_type, 13).into(),
291 location,
292 ));
293
294 block.append_operation(r#yield(
295 &[result.result(0).unwrap().into()],
296 location,
297 ));
298
299 let region = Region::new();
300 region.append_block(block);
301 region
302 },
303 location,
304 ));
305
306 block.append_operation(func::r#return(
307 &[result.result(0).unwrap().into()],
308 location,
309 ));
310
311 let region = Region::new();
312 region.append_block(block);
313 region
314 },
315 &[],
316 location,
317 ));
318
319 assert!(module.as_operation().verify());
320 insta::assert_snapshot!(module.as_operation());
321 }
322
323 #[test]
324 fn compile_one_sided() {
325 let context = Context::new();
326 load_all_dialects(&context);
327
328 let location = Location::unknown(&context);
329 let module = Module::new(location);
330
331 module.body().append_operation(func::func(
332 &context,
333 StringAttribute::new(&context, "foo"),
334 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
335 {
336 let block = Block::new(&[]);
337
338 let condition = block.append_operation(arith::constant(
339 &context,
340 IntegerAttribute::new(IntegerType::new(&context, 1).into(), 0).into(),
341 location,
342 ));
343
344 block.append_operation(r#if(
345 condition.result(0).unwrap().into(),
346 &[],
347 {
348 let block = Block::new(&[]);
349
350 block.append_operation(r#yield(&[], location));
351
352 let region = Region::new();
353 region.append_block(block);
354 region
355 },
356 Region::new(),
357 location,
358 ));
359
360 block.append_operation(func::r#return(&[], location));
361
362 let region = Region::new();
363 region.append_block(block);
364 region
365 },
366 &[],
367 location,
368 ));
369
370 assert!(module.as_operation().verify());
371 insta::assert_snapshot!(module.as_operation());
372 }
373 }
374
375 #[test]
376 fn compile_index_switch() {
377 let context = Context::new();
378 load_all_dialects(&context);
379
380 let location = Location::unknown(&context);
381 let module = Module::new(location);
382
383 module.body().append_operation(func::func(
384 &context,
385 StringAttribute::new(&context, "foo"),
386 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
387 {
388 let block = Block::new(&[]);
389
390 let condition = block.append_operation(arith::constant(
391 &context,
392 IntegerAttribute::new(Type::index(&context), 0).into(),
393 location,
394 ));
395
396 block.append_operation(index_switch(
397 &context,
398 condition.result(0).unwrap().into(),
399 &[],
400 DenseI64ArrayAttribute::new(&context, &[0, 1]),
401 vec![
402 {
403 let block = Block::new(&[]);
404
405 block.append_operation(r#yield(&[], location));
406
407 let region = Region::new();
408 region.append_block(block);
409 region
410 },
411 {
412 let block = Block::new(&[]);
413
414 block.append_operation(r#yield(&[], location));
415
416 let region = Region::new();
417 region.append_block(block);
418 region
419 },
420 {
421 let block = Block::new(&[]);
422
423 block.append_operation(r#yield(&[], location));
424
425 let region = Region::new();
426 region.append_block(block);
427 region
428 },
429 ],
430 location,
431 ));
432
433 block.append_operation(func::r#return(&[], location));
434
435 let region = Region::new();
436 region.append_block(block);
437 region
438 },
439 &[],
440 location,
441 ));
442
443 assert!(module.as_operation().verify());
444 insta::assert_snapshot!(module.as_operation());
445 }
446
447 mod r#while {
448 use super::*;
449
450 #[test]
451 fn compile() {
452 let context = Context::new();
453 load_all_dialects(&context);
454
455 let location = Location::unknown(&context);
456 let module = Module::new(location);
457 let index_type = Type::index(&context);
458
459 module.body().append_operation(func::func(
460 &context,
461 StringAttribute::new(&context, "foo"),
462 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
463 {
464 let block = Block::new(&[]);
465
466 let initial = block.append_operation(arith::constant(
467 &context,
468 IntegerAttribute::new(index_type, 0).into(),
469 location,
470 ));
471
472 block.append_operation(r#while(
473 &[initial.result(0).unwrap().into()],
474 &[index_type],
475 {
476 let block = Block::new(&[(index_type, location)]);
477
478 let condition = block.append_operation(arith::constant(
479 &context,
480 IntegerAttribute::new(IntegerType::new(&context, 1).into(), 0)
481 .into(),
482 location,
483 ));
484
485 let result = block.append_operation(arith::constant(
486 &context,
487 IntegerAttribute::new(Type::index(&context), 42).into(),
488 location,
489 ));
490
491 block.append_operation(super::condition(
492 condition.result(0).unwrap().into(),
493 &[result.result(0).unwrap().into()],
494 location,
495 ));
496
497 let region = Region::new();
498 region.append_block(block);
499 region
500 },
501 {
502 let block = Block::new(&[(index_type, location)]);
503
504 let result = block.append_operation(arith::constant(
505 &context,
506 IntegerAttribute::new(index_type, 42).into(),
507 location,
508 ));
509
510 block.append_operation(r#yield(
511 &[result.result(0).unwrap().into()],
512 location,
513 ));
514
515 let region = Region::new();
516 region.append_block(block);
517 region
518 },
519 location,
520 ));
521
522 block.append_operation(func::r#return(&[], location));
523
524 let region = Region::new();
525 region.append_block(block);
526 region
527 },
528 &[],
529 location,
530 ));
531
532 assert!(module.as_operation().verify());
533 insta::assert_snapshot!(module.as_operation());
534 }
535
536 #[test]
537 fn compile_with_different_argument_and_result_types() {
538 let context = Context::new();
539 load_all_dialects(&context);
540
541 let location = Location::unknown(&context);
542 let module = Module::new(location);
543 let index_type = Type::index(&context);
544 let float_type = Type::float64(&context);
545
546 module.body().append_operation(func::func(
547 &context,
548 StringAttribute::new(&context, "foo"),
549 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
550 {
551 let block = Block::new(&[]);
552
553 let initial = block.append_operation(arith::constant(
554 &context,
555 IntegerAttribute::new(index_type, 0).into(),
556 location,
557 ));
558
559 block.append_operation(r#while(
560 &[initial.result(0).unwrap().into()],
561 &[float_type],
562 {
563 let block = Block::new(&[(index_type, location)]);
564
565 let condition = block.append_operation(arith::constant(
566 &context,
567 IntegerAttribute::new(IntegerType::new(&context, 1).into(), 0)
568 .into(),
569 location,
570 ));
571
572 let result = block.append_operation(arith::constant(
573 &context,
574 FloatAttribute::new(&context, float_type, 42.0).into(),
575 location,
576 ));
577
578 block.append_operation(super::condition(
579 condition.result(0).unwrap().into(),
580 &[result.result(0).unwrap().into()],
581 location,
582 ));
583
584 let region = Region::new();
585 region.append_block(block);
586 region
587 },
588 {
589 let block = Block::new(&[(float_type, location)]);
590
591 let result = block.append_operation(arith::constant(
592 &context,
593 IntegerAttribute::new(Type::index(&context), 42).into(),
594 location,
595 ));
596
597 block.append_operation(r#yield(
598 &[result.result(0).unwrap().into()],
599 location,
600 ));
601
602 let region = Region::new();
603 region.append_block(block);
604 region
605 },
606 location,
607 ));
608
609 block.append_operation(func::r#return(&[], location));
610
611 let region = Region::new();
612 region.append_block(block);
613 region
614 },
615 &[],
616 location,
617 ));
618
619 assert!(module.as_operation().verify());
620 insta::assert_snapshot!(module.as_operation());
621 }
622
623 #[test]
624 fn compile_with_multiple_arguments_and_results() {
625 let context = Context::new();
626 load_all_dialects(&context);
627
628 let location = Location::unknown(&context);
629 let module = Module::new(location);
630 let index_type = Type::index(&context);
631
632 module.body().append_operation(func::func(
633 &context,
634 StringAttribute::new(&context, "foo"),
635 TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
636 {
637 let block = Block::new(&[]);
638
639 let initial = block.append_operation(arith::constant(
640 &context,
641 IntegerAttribute::new(index_type, 0).into(),
642 location,
643 ));
644
645 block.append_operation(r#while(
646 &[
647 initial.result(0).unwrap().into(),
648 initial.result(0).unwrap().into(),
649 ],
650 &[index_type, index_type],
651 {
652 let block =
653 Block::new(&[(index_type, location), (index_type, location)]);
654
655 let condition = block.append_operation(arith::constant(
656 &context,
657 IntegerAttribute::new(IntegerType::new(&context, 1).into(), 0)
658 .into(),
659 location,
660 ));
661
662 let result = block.append_operation(arith::constant(
663 &context,
664 IntegerAttribute::new(Type::index(&context), 42).into(),
665 location,
666 ));
667
668 block.append_operation(super::condition(
669 condition.result(0).unwrap().into(),
670 &[
671 result.result(0).unwrap().into(),
672 result.result(0).unwrap().into(),
673 ],
674 location,
675 ));
676
677 let region = Region::new();
678 region.append_block(block);
679 region
680 },
681 {
682 let block =
683 Block::new(&[(index_type, location), (index_type, location)]);
684
685 let result = block.append_operation(arith::constant(
686 &context,
687 IntegerAttribute::new(index_type, 42).into(),
688 location,
689 ));
690
691 block.append_operation(r#yield(
692 &[
693 result.result(0).unwrap().into(),
694 result.result(0).unwrap().into(),
695 ],
696 location,
697 ));
698
699 let region = Region::new();
700 region.append_block(block);
701 region
702 },
703 location,
704 ));
705
706 block.append_operation(func::r#return(&[], location));
707
708 let region = Region::new();
709 region.append_block(block);
710 region
711 },
712 &[],
713 location,
714 ));
715
716 assert!(module.as_operation().verify());
717 insta::assert_snapshot!(module.as_operation());
718 }
719 }
720}