Skip to main content

ts_typed_ast/
lib.rs

1use std::marker::PhantomData;
2use std::num::NonZeroU16;
3
4use tree_sitter::{Node, TreeCursor};
5
6mod generate;
7pub use generate::generate;
8
9pub trait AstNode<'tree>: Sized {
10    fn can_cast(kind: u16) -> bool;
11
12    fn cast(node: Node<'tree>) -> Option<Self>;
13
14    fn node(&self) -> Node<'tree>;
15
16    fn utf8_text<'a>(&self, source: &'a [u8]) -> Result<&'a str, std::str::Utf8Error> {
17        self.node().utf8_text(source)
18    }
19}
20
21pub struct MissingNodeChildError<'tree> {
22    pub node: tree_sitter::Node<'tree>,
23    pub field_id: u16,
24}
25
26impl<'tree> MissingNodeChildError<'tree> {
27    pub fn new(node: tree_sitter::Node<'tree>, field_id: u16) -> Self {
28        Self { node, field_id }
29    }
30}
31
32impl<'tree> std::fmt::Debug for MissingNodeChildError<'tree> {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        f.debug_struct("MissingNodeChildError")
35            .field("node", &self.node.id() as _)
36            .field("field_id", &self.field_id as _)
37            .finish()
38    }
39}
40
41impl<'tree> std::fmt::Display for MissingNodeChildError<'tree> {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        let field_name = self.node.language().field_name_for_id(self.field_id);
44        write!(f, "missing child node for")?;
45        if let Some(field_name) = field_name {
46            write!(f, " field '{field_name}'")?;
47        } else {
48            write!(f, " unknown field")?;
49        }
50        let position = self.node.start_position();
51        write!(f, " at {}:{}", position.row, position.column)
52    }
53}
54
55impl<'tree> std::error::Error for MissingNodeChildError<'tree> {}
56
57#[doc(hidden)]
58pub enum Children<'tree, T: AstNode<'tree>> {
59    Empty,
60    Walking {
61        cursor: TreeCursor<'tree>,
62        field_id: NonZeroU16,
63        _marker: PhantomData<T>,
64    },
65}
66
67impl<'tree, T: AstNode<'tree>> Children<'tree, T> {
68    #[doc(hidden)]
69    pub fn new(node: tree_sitter::Node<'tree>, field_id: NonZeroU16) -> Self {
70        // TODO probably faster to hardcode the ID in the build phase.
71        let mut cursor = node.walk();
72        cursor.goto_first_child();
73        Self::Walking {
74            cursor,
75            field_id,
76            _marker: Default::default(),
77        }
78    }
79}
80
81impl<'tree, T: AstNode<'tree>> Iterator for Children<'tree, T> {
82    type Item = T;
83
84    fn next(&mut self) -> Option<Self::Item> {
85        match std::mem::replace(self, Children::Empty) {
86            Children::Empty => None,
87            Children::Walking {
88                mut cursor,
89                field_id,
90                _marker,
91            } => loop {
92                let result = if cursor.field_id() == Some(field_id) {
93                    T::cast(cursor.node())
94                } else {
95                    None
96                };
97                if !cursor.goto_next_sibling() {
98                    return result;
99                }
100                if result.is_some() {
101                    _ = std::mem::replace(
102                        self,
103                        Children::Walking {
104                            cursor,
105                            field_id,
106                            _marker,
107                        },
108                    );
109                    return result;
110                }
111            },
112        }
113    }
114}