hir/
node_map.rs

1use std::collections::HashMap;
2
3use span::Span;
4
5use crate::item::Item;
6use crate::Param;
7use crate::{Expression, HirId, Module, PathDef, Statement, Type, item::Field};
8
9#[derive(Clone, Copy, Debug)]
10pub enum HirNodeKind<'hir> {
11    Expr(&'hir Expression<'hir>),
12    Item(&'hir Item<'hir>),
13    Param(&'hir Param<'hir>),
14    PathDef(&'hir PathDef),
15    Stmt(&'hir Statement<'hir>),
16    Field(&'hir Field<'hir>),
17    Ty(&'hir Type<'hir>),
18    Module(&'hir Module<'hir>),
19}
20
21impl<'hir> HirNodeKind<'hir> {
22    /// Expects this node to be a [Item] variant.
23    ///
24    /// # Panics
25    /// - If the node is NOT a [Item] variant
26    ///
27    /// [Item]: HirNodeKind::Item
28    pub fn expect_item(self) -> &'hir Item<'hir> {
29        self.as_item().unwrap_or_else(|| {
30            unreachable!("Expected item")
31        })
32    }
33
34    pub fn as_item(self) -> Option<&'hir Item<'hir>> {
35        match self {
36            Self::Item(item) => Some(item),
37            _ => None
38        }
39    }
40
41    pub fn get_span(&self) -> Option<Span> {
42        Some(match self {
43            HirNodeKind::Expr(expression) => expression.span,
44            HirNodeKind::Item(item) => item.span,
45            HirNodeKind::Param(p) => p.span,
46            HirNodeKind::Stmt(statement) => statement.span,
47            HirNodeKind::Field(field) => field.span,
48            HirNodeKind::PathDef(_) |
49            HirNodeKind::Ty(_) => return None,
50            HirNodeKind::Module(module) => module.span,
51        })
52    }
53
54    pub fn get_name(&self) -> &'static str {
55        match self {
56            HirNodeKind::Expr(_) => "Expr",
57            HirNodeKind::Item(_) => "Item",
58            HirNodeKind::Param(_) => "Param",
59            HirNodeKind::PathDef(_) => "PathDef",
60            HirNodeKind::Stmt(_) => "Stmt",
61            HirNodeKind::Field(_) => "Field",
62            HirNodeKind::Ty(_) => "Ty",
63            HirNodeKind::Module(_) => "Module",
64        }
65    }
66}
67
68impl<'hir> From<&'hir Expression<'hir>> for HirNodeKind<'hir> {
69    fn from(value: &'hir Expression<'hir>) -> Self { HirNodeKind::Expr(value) }
70}
71
72#[derive(Default)]
73pub struct NodeMap<'hir> {
74    map: HashMap<HirId, HirNodeKind<'hir>>,
75    id_count: usize,
76}
77
78impl<'hir> NodeMap<'hir> {
79    pub fn define(&mut self, id: HirId, val: impl Into<HirNodeKind<'hir>>) {
80        self.map.insert(id, val.into());
81    }
82
83    pub fn get_by_id(&self, id: &HirId) -> Option<HirNodeKind<'hir>> { self.map.get(id).copied() }
84
85    pub fn get_next_id(&mut self) -> HirId {
86        self.id_count += 1;
87        HirId(self.id_count)
88    }
89}