llmcc_core/
block.rs

1use strum_macros::{Display, EnumIter, EnumString, FromRepr};
2
3use crate::context::CompileUnit;
4use crate::declare_arena;
5use crate::ir::HirNode;
6
7declare_arena!([
8    blk_root: BlockRoot<'tcx>,
9    blk_func: BlockFunc<'tcx>,
10    blk_class: BlockClass<'tcx>,
11    blk_impl: BlockImpl<'tcx>,
12    blk_stmt: BlockStmt<'tcx>,
13    blk_call: BlockCall<'tcx>,
14    blk_enum: BlockEnum<'tcx>,
15    blk_field: BlockField<'tcx>,
16    blk_const: BlockConst<'tcx>,
17]);
18
19#[derive(
20    Debug, Clone, Copy, PartialEq, Eq, Hash, EnumIter, EnumString, FromRepr, Display, Default,
21)]
22#[strum(serialize_all = "snake_case")]
23pub enum BlockKind {
24    #[default]
25    Undefined,
26    Root,
27    Func,
28    Stmt,
29    Call,
30    Class,
31    Enum,
32    Const,
33    Impl,
34    Field,
35    Scope,
36}
37
38#[derive(Debug, Clone)]
39pub enum BasicBlock<'blk> {
40    Undefined,
41    Root(&'blk BlockRoot<'blk>),
42    Func(&'blk BlockFunc<'blk>),
43    Stmt(&'blk BlockStmt<'blk>),
44    Call(&'blk BlockCall<'blk>),
45    Enum(&'blk BlockEnum<'blk>),
46    Class(&'blk BlockClass<'blk>),
47    Impl(&'blk BlockImpl<'blk>),
48    Const(&'blk BlockConst<'blk>),
49    Field(&'blk BlockField<'blk>),
50    Block,
51}
52
53impl<'blk> BasicBlock<'blk> {
54    pub fn format_block(&self, _unit: CompileUnit<'blk>) -> String {
55        let block_id = self.block_id();
56        let kind = self.kind();
57        let name = self
58            .base()
59            .and_then(|base| base.opt_get_name())
60            .unwrap_or("");
61
62        // Include file_name for Root blocks
63        if let BasicBlock::Root(root) = self {
64            if let Some(file_name) = &root.file_name {
65                return format!("{}:{} {} ({})", kind, block_id, name, file_name);
66            }
67        }
68
69        format!("{}:{} {}", kind, block_id, name)
70    }
71
72    /// Get the base block information regardless of variant
73    pub fn base(&self) -> Option<&BlockBase<'blk>> {
74        match self {
75            BasicBlock::Undefined | BasicBlock::Block => None,
76            BasicBlock::Root(block) => Some(&block.base),
77            BasicBlock::Func(block) => Some(&block.base),
78            BasicBlock::Class(block) => Some(&block.base),
79            BasicBlock::Impl(block) => Some(&block.base),
80            BasicBlock::Stmt(block) => Some(&block.base),
81            BasicBlock::Call(block) => Some(&block.base),
82            BasicBlock::Enum(block) => Some(&block.base),
83            BasicBlock::Const(block) => Some(&block.base),
84            BasicBlock::Field(block) => Some(&block.base),
85        }
86    }
87
88    /// Get the block ID
89    pub fn block_id(&self) -> BlockId {
90        self.base().unwrap().id
91    }
92
93    /// Get the block kind
94    pub fn kind(&self) -> BlockKind {
95        self.base().map(|base| base.kind).unwrap_or_default()
96    }
97
98    /// Get the HIR node
99    pub fn node(&self) -> &HirNode<'blk> {
100        self.base().map(|base| &base.node).unwrap()
101    }
102
103    pub fn opt_node(&self) -> Option<&HirNode<'blk>> {
104        self.base().map(|base| &base.node)
105    }
106
107    /// Get the children block IDs
108    pub fn children(&self) -> &[BlockId] {
109        self.base()
110            .map(|base| base.children.as_slice())
111            .unwrap_or(&[])
112    }
113
114    pub fn child_count(&self) -> usize {
115        self.children().len()
116    }
117
118    /// Check if this is a specific kind of block
119    pub fn is_kind(&self, kind: BlockKind) -> bool {
120        self.kind() == kind
121    }
122}
123
124#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, Default)]
125pub struct BlockId(pub u32);
126
127impl std::fmt::Display for BlockId {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        write!(f, "{}", self.0)
130    }
131}
132
133impl BlockId {
134    pub fn new(id: u32) -> Self {
135        Self(id)
136    }
137
138    pub fn as_u32(self) -> u32 {
139        self.0
140    }
141
142    pub const ROOT_PARENT: BlockId = BlockId(u32::MAX);
143
144    pub fn is_root_parent(self) -> bool {
145        self.0 == u32::MAX
146    }
147}
148
149#[derive(
150    Debug, Clone, Copy, PartialEq, Eq, Hash, EnumIter, EnumString, FromRepr, Display, Default,
151)]
152#[strum(serialize_all = "snake_case")]
153pub enum BlockRelation {
154    #[default]
155    Unknown,
156    DependedBy,
157    DependsOn,
158}
159
160#[derive(Debug, Clone)]
161pub struct BlockBase<'blk> {
162    pub id: BlockId,
163    pub node: HirNode<'blk>,
164    pub kind: BlockKind,
165    pub parent: Option<BlockId>,
166    pub children: Vec<BlockId>,
167}
168
169impl<'blk> BlockBase<'blk> {
170    pub fn new(
171        id: BlockId,
172        node: HirNode<'blk>,
173        kind: BlockKind,
174        parent: Option<BlockId>,
175        children: Vec<BlockId>,
176    ) -> Self {
177        Self {
178            id,
179            node,
180            kind,
181            parent,
182            children,
183        }
184    }
185
186    pub fn opt_get_name(&self) -> Option<&str> {
187        self.node
188            .as_scope()
189            .and_then(|scope| scope.ident.as_ref())
190            .map(|ident| ident.name.as_str())
191    }
192
193    pub fn add_child(&mut self, child_id: BlockId) {
194        if !self.children.contains(&child_id) {
195            self.children.push(child_id);
196        }
197    }
198
199    pub fn remove_child(&mut self, child_id: BlockId) {
200        self.children.retain(|&id| id != child_id);
201    }
202}
203
204#[derive(Debug, Clone)]
205pub struct BlockRoot<'blk> {
206    pub base: BlockBase<'blk>,
207    pub file_name: Option<String>,
208}
209
210impl<'blk> BlockRoot<'blk> {
211    pub fn new(base: BlockBase<'blk>, file_name: Option<String>) -> Self {
212        Self { base, file_name }
213    }
214
215    pub fn from_hir(
216        id: BlockId,
217        node: HirNode<'blk>,
218        parent: Option<BlockId>,
219        children: Vec<BlockId>,
220        file_name: Option<String>,
221    ) -> Self {
222        let base = BlockBase::new(id, node, BlockKind::Root, parent, children);
223        Self::new(base, file_name)
224    }
225}
226
227#[derive(Debug, Clone)]
228pub struct BlockFunc<'blk> {
229    pub base: BlockBase<'blk>,
230    pub name: String,
231    pub parameters: Option<BlockId>,
232    pub returns: Option<BlockId>,
233    pub stmts: Option<Vec<BlockId>>,
234}
235
236impl<'blk> BlockFunc<'blk> {
237    pub fn new(base: BlockBase<'blk>, name: String) -> Self {
238        Self {
239            base,
240            name,
241            parameters: None,
242            returns: None,
243            stmts: None,
244        }
245    }
246
247    pub fn from_hir(
248        id: BlockId,
249        node: HirNode<'blk>,
250        parent: Option<BlockId>,
251        children: Vec<BlockId>,
252    ) -> Self {
253        let base = BlockBase::new(id, node, BlockKind::Func, parent, children);
254        let name = base.opt_get_name().unwrap_or("").to_string();
255        Self::new(base, name)
256    }
257}
258
259#[derive(Debug, Clone)]
260pub struct BlockStmt<'blk> {
261    pub base: BlockBase<'blk>,
262}
263
264impl<'blk> BlockStmt<'blk> {
265    pub fn new(base: BlockBase<'blk>) -> Self {
266        Self { base }
267    }
268
269    pub fn from_hir(
270        id: BlockId,
271        node: HirNode<'blk>,
272        parent: Option<BlockId>,
273        children: Vec<BlockId>,
274    ) -> Self {
275        let base = BlockBase::new(id, node, BlockKind::Stmt, parent, children);
276        Self::new(base)
277    }
278}
279
280#[derive(Debug, Clone)]
281pub struct BlockCall<'blk> {
282    pub base: BlockBase<'blk>,
283}
284
285impl<'blk> BlockCall<'blk> {
286    pub fn new(base: BlockBase<'blk>) -> Self {
287        Self { base }
288    }
289
290    pub fn from_hir(
291        id: BlockId,
292        node: HirNode<'blk>,
293        parent: Option<BlockId>,
294        children: Vec<BlockId>,
295    ) -> Self {
296        let base = BlockBase::new(id, node, BlockKind::Call, parent, children);
297        Self::new(base)
298    }
299}
300
301#[derive(Debug, Clone)]
302pub struct BlockClass<'blk> {
303    pub base: BlockBase<'blk>,
304    pub name: String,
305}
306
307impl<'blk> BlockClass<'blk> {
308    pub fn new(base: BlockBase<'blk>, name: String) -> Self {
309        Self { base, name }
310    }
311
312    pub fn from_hir(
313        id: BlockId,
314        node: HirNode<'blk>,
315        parent: Option<BlockId>,
316        children: Vec<BlockId>,
317    ) -> Self {
318        let base = BlockBase::new(id, node, BlockKind::Class, parent, children);
319        let name = base.opt_get_name().unwrap_or("").to_string();
320        Self::new(base, name)
321    }
322}
323
324#[derive(Debug, Clone)]
325pub struct BlockImpl<'blk> {
326    pub base: BlockBase<'blk>,
327    pub name: String,
328}
329
330impl<'blk> BlockImpl<'blk> {
331    pub fn new(base: BlockBase<'blk>, name: String) -> Self {
332        Self { base, name }
333    }
334
335    pub fn from_hir(
336        id: BlockId,
337        node: HirNode<'blk>,
338        parent: Option<BlockId>,
339        children: Vec<BlockId>,
340    ) -> Self {
341        let base = BlockBase::new(id, node, BlockKind::Impl, parent, children);
342        let name = base.opt_get_name().unwrap_or("").to_string();
343        Self::new(base, name)
344    }
345}
346
347#[derive(Debug, Clone)]
348pub struct BlockEnum<'blk> {
349    pub base: BlockBase<'blk>,
350    pub name: String,
351    pub fields: Vec<BlockId>,
352}
353
354impl<'blk> BlockEnum<'blk> {
355    pub fn new(base: BlockBase<'blk>, name: String) -> Self {
356        Self {
357            base,
358            name,
359            fields: Vec::new(),
360        }
361    }
362
363    pub fn from_hir(
364        id: BlockId,
365        node: HirNode<'blk>,
366        parent: Option<BlockId>,
367        children: Vec<BlockId>,
368    ) -> Self {
369        let base = BlockBase::new(id, node, BlockKind::Enum, parent, children);
370        let name = base.opt_get_name().unwrap_or("").to_string();
371        Self::new(base, name)
372    }
373
374    pub fn add_field(&mut self, field_id: BlockId) {
375        self.fields.push(field_id);
376    }
377}
378
379#[derive(Debug, Clone)]
380pub struct BlockConst<'blk> {
381    pub base: BlockBase<'blk>,
382    pub name: String,
383}
384
385impl<'blk> BlockConst<'blk> {
386    pub fn new(base: BlockBase<'blk>, name: String) -> Self {
387        Self { base, name }
388    }
389
390    pub fn from_hir(
391        id: BlockId,
392        node: HirNode<'blk>,
393        parent: Option<BlockId>,
394        children: Vec<BlockId>,
395    ) -> Self {
396        let base = BlockBase::new(id, node, BlockKind::Const, parent, children);
397        let name = base.opt_get_name().unwrap_or("").to_string();
398        Self::new(base, name)
399    }
400}
401
402#[derive(Debug, Clone)]
403pub struct BlockField<'blk> {
404    pub base: BlockBase<'blk>,
405    pub name: String,
406}
407
408impl<'blk> BlockField<'blk> {
409    pub fn new(base: BlockBase<'blk>, name: String) -> Self {
410        Self { base, name }
411    }
412
413    pub fn from_hir(
414        id: BlockId,
415        node: HirNode<'blk>,
416        parent: Option<BlockId>,
417        children: Vec<BlockId>,
418    ) -> Self {
419        let base = BlockBase::new(id, node, BlockKind::Field, parent, children);
420        let name = base.opt_get_name().unwrap_or("").to_string();
421        Self::new(base, name)
422    }
423}