melior/ir/block/
block_like.rs

1use super::{BlockArgument, BlockRef, TypeLike};
2use crate::{
3    Error,
4    ir::{Location, Operation, OperationRef, RegionRef, Type, Value, operation::OperationRefMut},
5};
6use core::fmt::Display;
7use mlir_sys::{
8    MlirBlock, mlirBlockAddArgument, mlirBlockAppendOwnedOperation, mlirBlockEraseArgument,
9    mlirBlockGetArgument, mlirBlockGetFirstOperation, mlirBlockGetNextInRegion,
10    mlirBlockGetNumArguments, mlirBlockGetNumPredecessors, mlirBlockGetNumSuccessors,
11    mlirBlockGetParentOperation, mlirBlockGetParentRegion, mlirBlockGetPredecessor,
12    mlirBlockGetSuccessor, mlirBlockGetTerminator, mlirBlockInsertArgument,
13    mlirBlockInsertOwnedOperation, mlirBlockInsertOwnedOperationAfter,
14    mlirBlockInsertOwnedOperationBefore,
15};
16
17/// A trait for block-like types.
18// TODO Split this trait into `BlockLike` and `BlockLikeMut`.
19pub trait BlockLike<'c, 'a>: Display + 'a {
20    /// Converts a block into a raw object.
21    fn to_raw(&self) -> MlirBlock;
22
23    /// Returns an argument at a position.
24    fn argument(&self, index: usize) -> Result<BlockArgument<'c, 'a>, Error> {
25        unsafe {
26            if index < self.argument_count() {
27                Ok(BlockArgument::from_raw(mlirBlockGetArgument(
28                    self.to_raw(),
29                    index as isize,
30                )))
31            } else {
32                Err(Error::PositionOutOfBounds {
33                    name: "block argument",
34                    value: self.to_string(),
35                    index,
36                })
37            }
38        }
39    }
40
41    /// Returns a number of arguments.
42    fn argument_count(&self) -> usize {
43        unsafe { mlirBlockGetNumArguments(self.to_raw()) as usize }
44    }
45
46    /// Returns a reference to the first operation.
47    fn first_operation(&self) -> Option<OperationRef<'c, 'a>> {
48        unsafe { OperationRef::from_option_raw(mlirBlockGetFirstOperation(self.to_raw())) }
49    }
50
51    /// Returns a mutable reference to the first operation.
52    fn first_operation_mut(&self) -> Option<OperationRefMut<'c, 'a>> {
53        unsafe { OperationRefMut::from_option_raw(mlirBlockGetFirstOperation(self.to_raw())) }
54    }
55
56    /// Returns a reference to a terminator operation.
57    fn terminator(&self) -> Option<OperationRef<'c, 'a>> {
58        unsafe { OperationRef::from_option_raw(mlirBlockGetTerminator(self.to_raw())) }
59    }
60
61    /// Returns a mutable reference to a terminator operation.
62    fn terminator_mut(&self) -> Option<OperationRefMut<'c, 'a>> {
63        unsafe { OperationRefMut::from_option_raw(mlirBlockGetTerminator(self.to_raw())) }
64    }
65
66    /// Returns a parent region.
67    // TODO Store lifetime of regions in blocks, or create another type like
68    // `InsertedBlockRef`?
69    fn parent_region(&self) -> Option<RegionRef<'c, 'a>> {
70        unsafe { RegionRef::from_option_raw(mlirBlockGetParentRegion(self.to_raw())) }
71    }
72
73    /// Returns a parent operation.
74    fn parent_operation(&self) -> Option<OperationRef<'c, 'a>> {
75        unsafe { OperationRef::from_option_raw(mlirBlockGetParentOperation(self.to_raw())) }
76    }
77
78    /// Adds an argument.
79    fn add_argument(&self, r#type: Type<'c>, location: Location<'c>) -> Value<'c, 'a> {
80        unsafe {
81            Value::from_raw(mlirBlockAddArgument(
82                self.to_raw(),
83                r#type.to_raw(),
84                location.to_raw(),
85            ))
86        }
87    }
88
89    /// Appends an operation.
90    fn append_operation(&self, operation: Operation<'c>) -> OperationRef<'c, 'a> {
91        unsafe {
92            let operation = operation.into_raw();
93
94            mlirBlockAppendOwnedOperation(self.to_raw(), operation);
95
96            OperationRef::from_raw(operation)
97        }
98    }
99
100    /// Inserts an operation.
101    // TODO How can we make those update functions take `&mut self`?
102    // TODO Use cells?
103    fn insert_operation(&self, position: usize, operation: Operation<'c>) -> OperationRef<'c, 'a> {
104        unsafe {
105            let operation = operation.into_raw();
106
107            mlirBlockInsertOwnedOperation(self.to_raw(), position as isize, operation);
108
109            OperationRef::from_raw(operation)
110        }
111    }
112
113    /// Inserts an operation after another.
114    fn insert_operation_after(
115        &self,
116        one: OperationRef<'c, 'a>,
117        other: Operation<'c>,
118    ) -> OperationRef<'c, 'a> {
119        unsafe {
120            let other = other.into_raw();
121
122            mlirBlockInsertOwnedOperationAfter(self.to_raw(), one.to_raw(), other);
123
124            OperationRef::from_raw(other)
125        }
126    }
127
128    /// Inserts an operation before another.
129    fn insert_operation_before(
130        &self,
131        one: OperationRef<'c, 'a>,
132        other: Operation<'c>,
133    ) -> OperationRef<'c, 'a> {
134        unsafe {
135            let other = other.into_raw();
136
137            mlirBlockInsertOwnedOperationBefore(self.to_raw(), one.to_raw(), other);
138
139            OperationRef::from_raw(other)
140        }
141    }
142
143    /// Inserts an argument at the given position.
144    ///
145    /// Existing arguments at and after `index` are shifted. Any previously
146    /// obtained `BlockArgument` handles for those positions become stale.
147    fn insert_argument(
148        &self,
149        index: usize,
150        r#type: Type<'c>,
151        location: Location<'c>,
152    ) -> Value<'c, 'a> {
153        unsafe {
154            Value::from_raw(mlirBlockInsertArgument(
155                self.to_raw(),
156                index as isize,
157                r#type.to_raw(),
158                location.to_raw(),
159            ))
160        }
161    }
162
163    /// Erases the argument at the given position.
164    ///
165    /// # Safety
166    ///
167    /// This invalidates any `BlockArgument` or `Value` handles that were
168    /// obtained for the erased argument. Using such handles after this call
169    /// is undefined behavior.
170    unsafe fn erase_argument(&self, index: usize) {
171        unsafe { mlirBlockEraseArgument(self.to_raw(), index as u32) }
172    }
173
174    /// Returns the number of successors.
175    fn successor_count(&self) -> usize {
176        unsafe { mlirBlockGetNumSuccessors(self.to_raw()) as usize }
177    }
178
179    /// Returns a successor block at a position.
180    fn successor(&self, index: usize) -> Result<BlockRef<'c, 'a>, Error> {
181        if index < self.successor_count() {
182            Ok(unsafe { BlockRef::from_raw(mlirBlockGetSuccessor(self.to_raw(), index as isize)) })
183        } else {
184            Err(Error::PositionOutOfBounds {
185                name: "block successor",
186                value: self.to_string(),
187                index,
188            })
189        }
190    }
191
192    /// Returns the number of predecessors.
193    fn predecessor_count(&self) -> usize {
194        unsafe { mlirBlockGetNumPredecessors(self.to_raw()) as usize }
195    }
196
197    /// Returns a predecessor block at a position.
198    fn predecessor(&self, index: usize) -> Result<BlockRef<'c, 'a>, Error> {
199        if index < self.predecessor_count() {
200            Ok(unsafe {
201                BlockRef::from_raw(mlirBlockGetPredecessor(self.to_raw(), index as isize))
202            })
203        } else {
204            Err(Error::PositionOutOfBounds {
205                name: "block predecessor",
206                value: self.to_string(),
207                index,
208            })
209        }
210    }
211
212    /// Returns a next block in a region.
213    fn next_in_region(&self) -> Option<BlockRef<'c, 'a>> {
214        unsafe { BlockRef::from_option_raw(mlirBlockGetNextInRegion(self.to_raw())) }
215    }
216}