koto_parser/
ast.rs

1use koto_lexer::Span;
2use std::{fmt, num::TryFromIntError};
3
4use crate::{ConstantPool, Node, error::*};
5
6/// The index type used by nodes in the [Ast]
7#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
8pub struct AstIndex(u32);
9
10impl From<AstIndex> for u32 {
11    fn from(value: AstIndex) -> Self {
12        value.0
13    }
14}
15
16impl From<AstIndex> for usize {
17    fn from(value: AstIndex) -> Self {
18        value.0 as usize
19    }
20}
21
22impl From<u32> for AstIndex {
23    fn from(value: u32) -> Self {
24        Self(value)
25    }
26}
27
28impl From<&u32> for AstIndex {
29    fn from(value: &u32) -> Self {
30        Self(*value)
31    }
32}
33
34impl TryFrom<usize> for AstIndex {
35    type Error = TryFromIntError;
36
37    fn try_from(value: usize) -> std::result::Result<Self, Self::Error> {
38        Ok(Self(u32::try_from(value)?))
39    }
40}
41
42impl fmt::Display for AstIndex {
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        write!(f, "{}", self.0)
45    }
46}
47
48/// A [Node] in the [Ast], along with its corresponding [Span]
49#[derive(Clone, Debug, Default)]
50pub struct AstNode {
51    /// The node itself
52    pub node: Node,
53    /// The index of the node's corresponding [Span]
54    ///
55    /// The span is stored in the [Ast], and retrieved via [Ast::span].
56    pub span: AstIndex,
57}
58
59/// A Koto program represented as an Abstract Syntax Tree
60///
61/// This is produced by the parser, and consumed by the compiler.
62#[derive(Debug, Default, Clone)]
63pub struct Ast {
64    nodes: Vec<AstNode>,
65    spans: Vec<Span>,
66    constants: ConstantPool,
67}
68
69impl Ast {
70    /// Initializes an Ast with the given capacity
71    pub fn with_capacity(capacity: usize) -> Self {
72        Self {
73            nodes: Vec::with_capacity(capacity),
74            spans: Vec::with_capacity(capacity),
75            constants: ConstantPool::default(),
76        }
77    }
78
79    /// Pushes a node and corresponding span onto the tree
80    pub fn push(&mut self, node: Node, span: Span) -> Result<AstIndex> {
81        // We could potentially achieve some compression by
82        // using a set for the spans, for now a Vec will do.
83        self.spans.push(span);
84        let span_index = AstIndex::try_from(self.spans.len() - 1)
85            .map_err(|_| Error::new(InternalError::AstCapacityOverflow.into(), span))?;
86
87        self.nodes.push(AstNode {
88            node,
89            span: span_index,
90        });
91        AstIndex::try_from(self.nodes.len() - 1)
92            .map_err(|_| Error::new(InternalError::AstCapacityOverflow.into(), span))
93    }
94
95    /// Returns a node for a given node index
96    pub fn node(&self, index: AstIndex) -> &AstNode {
97        &self.nodes[usize::from(index)]
98    }
99
100    /// Returns a span for a given span index
101    pub fn span(&self, index: AstIndex) -> &Span {
102        &self.spans[usize::from(index)]
103    }
104
105    /// Returns the constant pool referred to by the AST
106    pub fn constants(&self) -> &ConstantPool {
107        &self.constants
108    }
109
110    /// Moves the constants out of the AST
111    ///
112    /// This is used when building a `Chunk` after compilation.
113    /// The constants get transferred to the `Chunk` once the AST has been converted into bytecode.
114    pub fn consume_constants(self) -> ConstantPool {
115        self.constants
116    }
117
118    pub(crate) fn set_constants(&mut self, constants: ConstantPool) {
119        self.constants = constants
120    }
121
122    /// Returns the root node in the tree
123    pub fn entry_point(&self) -> Option<AstIndex> {
124        if self.nodes.is_empty() {
125            None
126        } else {
127            AstIndex::try_from(self.nodes.len() - 1).ok()
128        }
129    }
130
131    /// Used in testing to validate the tree's contents
132    pub fn nodes(&self) -> &[AstNode] {
133        &self.nodes
134    }
135}
koto_parser/ast.rs

koto_parser/
ast.rs
