Skip to main content

template_model/
ir.rs

1use crate::consts::*;
2pub(crate) use crate::defect::ParseDefect;
3pub(crate) use crate::span::{SourceSpan, UNKNOWN_SPAN};
4use std::collections::HashMap;
5use std::error;
6use std::fmt;
7
8pub type NodeRef = u32;
9
10pub const NULL_REF: NodeRef = u32::MAX;
11
12static SENTINEL_NODE: SpannedNode = SpannedNode { node: Node::Null, span: UNKNOWN_SPAN, path: String::new() };
13
14#[derive(Debug)]
15pub struct Arena {
16    nodes: Vec<SpannedNode>,
17}
18
19impl Arena {
20    pub fn new() -> Self {
21        Self { nodes: Vec::new() }
22    }
23
24    pub fn alloc(&mut self, node: SpannedNode) -> NodeRef {
25        let idx = self.nodes.len() as u32;
26        self.nodes.push(node);
27        idx
28    }
29
30    pub fn get(&self, r: NodeRef) -> &SpannedNode {
31        match self.nodes.get(r as usize) {
32            Some(node) => node,
33            None => {
34                log::warn!("NodeRef {} out of bounds (arena size {}), returning sentinel", r, self.nodes.len());
35                &SENTINEL_NODE
36            }
37        }
38    }
39
40    pub fn is_valid(&self, r: NodeRef) -> bool {
41        (r as usize) < self.nodes.len()
42    }
43
44    pub fn node(&self, r: NodeRef) -> &Node {
45        &self.get(r).node
46    }
47
48    pub fn span(&self, r: NodeRef) -> SourceSpan {
49        self.get(r).span
50    }
51
52    pub fn as_str(&self, r: NodeRef) -> Option<&str> {
53        match &self.get(r).node {
54            Node::String(s) => Some(s),
55            _ => None,
56        }
57    }
58
59    pub fn as_map(&self, r: NodeRef) -> Option<&[(String, NodeRef)]> {
60        if r == NULL_REF || (r as usize) >= self.nodes.len() {
61            return None;
62        }
63        match &self.nodes[r as usize].node {
64            Node::Map(entries) => Some(entries),
65            _ => None,
66        }
67    }
68
69    pub fn map_get(&self, r: NodeRef, key: &str) -> Option<NodeRef> {
70        self.as_map(r)?.iter().find(|(k, _)| k == key).map(|(_, v)| *v)
71    }
72
73    pub fn as_list(&self, r: NodeRef) -> Option<&[NodeRef]> {
74        if r == NULL_REF || (r as usize) >= self.nodes.len() {
75            return None;
76        }
77        match &self.nodes[r as usize].node {
78            Node::List(items) => Some(items),
79            _ => None,
80        }
81    }
82
83    pub fn len(&self) -> usize {
84        self.nodes.len()
85    }
86
87    pub fn map_remove(&mut self, r: NodeRef, key: &str) {
88        if let Node::Map(entries) = &mut self.nodes[r as usize].node {
89            entries.retain(|(k, _)| k != key);
90        }
91    }
92
93    pub fn map_insert(&mut self, r: NodeRef, key: String, val: NodeRef) {
94        if let Node::Map(entries) = &mut self.nodes[r as usize].node {
95            entries.push((key, val));
96        }
97    }
98
99    /// Overwrites the node at `r` with `node`. Used by the language-extensions
100    /// transform to install a rewritten section over the ref the model holds,
101    /// without threading a new ref back to every caller.
102    pub fn set(&mut self, r: NodeRef, node: SpannedNode) {
103        if (r as usize) < self.nodes.len() {
104            self.nodes[r as usize] = node;
105        }
106    }
107}
108
109impl Default for Arena {
110    fn default() -> Self {
111        Self::new()
112    }
113}
114
115#[derive(Debug, Clone)]
116pub struct SpannedNode {
117    pub node: Node,
118    pub span: SourceSpan,
119    pub path: String,
120}
121
122#[derive(Debug, Clone, PartialEq)]
123pub enum Node {
124    Null,
125    Bool(bool),
126    Int(i64),
127    Float(f64),
128    String(String),
129    List(Vec<NodeRef>),
130    Map(Vec<(String, NodeRef)>),
131    Intrinsic(IntrinsicFn),
132}
133
134#[derive(Debug, Clone, PartialEq)]
135pub enum IntrinsicFn {
136    Ref(String),
137    GetAtt(String, String),
138    Sub(String, Option<Vec<(String, NodeRef)>>),
139    Join(NodeRef, NodeRef),
140    Select(NodeRef, NodeRef),
141    If(String, NodeRef, NodeRef),
142    IfExpr(NodeRef, NodeRef, NodeRef),
143    FindInMap(NodeRef, NodeRef, NodeRef, Option<NodeRef>),
144    Split(NodeRef, NodeRef),
145    Base64(NodeRef),
146    Cidr(NodeRef, NodeRef, NodeRef),
147    GetAZs(NodeRef),
148    ImportValue(NodeRef),
149    GetStackOutput(Vec<(String, NodeRef)>),
150    Transform(String, Vec<(String, NodeRef)>),
151    And(Vec<NodeRef>),
152    Or(Vec<NodeRef>),
153    Not(NodeRef),
154    Equals(NodeRef, NodeRef),
155    ToJsonString(NodeRef),
156    Length(NodeRef),
157    ForEach(String, String, NodeRef, NodeRef),
158    ValueOf(String, String),
159    ValueOfAll(String, String),
160    RefAll(String),
161    Contains(NodeRef, NodeRef),
162    EachMemberEquals(NodeRef, NodeRef),
163    EachMemberIn(NodeRef, NodeRef),
164}
165
166pub type GlobalIndex = HashMap<String, NodeRef>;
167
168/// Returns the CloudFormation function name as it appears in templates (e.g. `"Fn::GetAtt"`).
169pub fn cfn_function_name(intrinsic: &IntrinsicFn) -> &'static str {
170    match intrinsic {
171        IntrinsicFn::Ref(name) if name.starts_with(CONDITION_REF_PREFIX) => FN_CONDITION,
172        IntrinsicFn::Ref(_) => FN_REF,
173        IntrinsicFn::GetAtt(_, _) => FN_GET_ATT,
174        IntrinsicFn::Sub(_, _) => FN_SUB,
175        IntrinsicFn::Join(_, _) => FN_JOIN,
176        IntrinsicFn::Select(_, _) => FN_SELECT,
177        IntrinsicFn::If(_, _, _) | IntrinsicFn::IfExpr(_, _, _) => FN_IF,
178        IntrinsicFn::FindInMap(_, _, _, _) => FN_FIND_IN_MAP,
179        IntrinsicFn::Split(_, _) => FN_SPLIT,
180        IntrinsicFn::Base64(_) => FN_BASE64,
181        IntrinsicFn::Cidr(_, _, _) => FN_CIDR,
182        IntrinsicFn::GetAZs(_) => FN_GET_AZS,
183        IntrinsicFn::ImportValue(_) => FN_IMPORT_VALUE,
184        IntrinsicFn::GetStackOutput(_) => FN_GET_STACK_OUTPUT,
185        IntrinsicFn::Transform(_, _) => FN_TRANSFORM,
186        IntrinsicFn::And(_) => FN_AND,
187        IntrinsicFn::Or(_) => FN_OR,
188        IntrinsicFn::Not(_) => FN_NOT,
189        IntrinsicFn::Equals(_, _) => FN_EQUALS,
190        IntrinsicFn::ToJsonString(_) => FN_TO_JSON_STRING,
191        IntrinsicFn::Length(_) => FN_LENGTH,
192        IntrinsicFn::ForEach(_, _, _, _) => FN_FOR_EACH,
193        IntrinsicFn::ValueOf(_, _) => FN_VALUE_OF,
194        IntrinsicFn::ValueOfAll(_, _) => FN_VALUE_OF_ALL,
195        IntrinsicFn::RefAll(_) => FN_REF_ALL,
196        IntrinsicFn::Contains(_, _) => FN_CONTAINS,
197        IntrinsicFn::EachMemberEquals(_, _) => FN_EACH_MEMBER_EQUALS,
198        IntrinsicFn::EachMemberIn(_, _) => FN_EACH_MEMBER_IN,
199    }
200}
201pub type SourceSpanIndex = HashMap<String, SourceSpan>;
202
203#[derive(Debug)]
204pub struct TemplateIR {
205    pub(crate) arena: Arena,
206    pub(crate) global_index: GlobalIndex,
207    pub(crate) span_index: SourceSpanIndex,
208    /// The template's top-level map, from which every section below hangs.
209    pub(crate) root: NodeRef,
210    pub(crate) parameters: NodeRef,
211    pub(crate) mappings: NodeRef,
212    pub(crate) conditions: NodeRef,
213    pub(crate) resources: NodeRef,
214    pub(crate) outputs: NodeRef,
215    pub(crate) rules: NodeRef,
216    pub(crate) template_metadata: NodeRef,
217    pub(crate) format_version: Option<String>,
218    pub(crate) description: Option<String>,
219    pub(crate) transforms: Vec<String>,
220    pub(crate) raw_top_level_keys: Vec<String>,
221    pub(crate) diagnostics: Vec<ParseDefect>,
222    pub(crate) globals: NodeRef,
223}
224
225#[derive(Debug)]
226pub struct ParseError {
227    pub message: String,
228    pub line: Option<u32>,
229    pub column: Option<u32>,
230}
231
232impl fmt::Display for ParseError {
233    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
234        match (self.line, self.column) {
235            (Some(l), Some(c)) => write!(f, "{}:{}: {}", l, c, self.message),
236            _ => write!(f, "{}", self.message),
237        }
238    }
239}
240
241impl error::Error for ParseError {}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246
247    #[test]
248    fn arena_alloc_get_roundtrip() {
249        let mut arena = Arena::new();
250        let r = arena.alloc(SpannedNode {
251            node: Node::String("hello".into()),
252            span: SourceSpan { start_line: 1, start_column: 1, end_line: 1, end_column: 5 },
253            path: "test".into(),
254        });
255        assert_eq!(arena.as_str(r), Some("hello"));
256    }
257
258    #[test]
259    fn arena_multiple_nodes() {
260        let mut arena = Arena::new();
261        let refs: Vec<NodeRef> = (0..100)
262            .map(|i| arena.alloc(SpannedNode { node: Node::Int(i), span: UNKNOWN_SPAN, path: format!("{}", i) }))
263            .collect();
264        for (i, r) in refs.iter().enumerate() {
265            match arena.node(*r) {
266                Node::Int(v) => assert_eq!(*v, i as i64),
267                _ => panic!("wrong type"),
268            }
269        }
270    }
271
272    #[test]
273    fn node_accessor_map_get() {
274        let mut arena = Arena::new();
275        let child =
276            arena.alloc(SpannedNode { node: Node::String("value".into()), span: UNKNOWN_SPAN, path: "map/key".into() });
277        let map = arena.alloc(SpannedNode {
278            node: Node::Map(vec![("key".into(), child)]),
279            span: UNKNOWN_SPAN,
280            path: "map".into(),
281        });
282        assert_eq!(arena.map_get(map, "key"), Some(child));
283        assert_eq!(arena.map_get(map, "missing"), None);
284    }
285
286    #[test]
287    fn source_span_unknown_sentinel() {
288        assert_eq!(UNKNOWN_SPAN.start_line, u32::MAX);
289        assert_eq!(UNKNOWN_SPAN.end_column, u32::MAX);
290    }
291
292    #[test]
293    fn arena_out_of_bounds_returns_sentinel() {
294        let arena = Arena::new();
295        let node = arena.get(999);
296        assert!(matches!(node.node, Node::Null));
297        assert_eq!(node.span, UNKNOWN_SPAN);
298        assert!(!arena.is_valid(999));
299        assert_eq!(arena.as_map(999), None, "out-of-bounds as_map should return None");
300        assert_eq!(arena.as_list(999), None, "out-of-bounds as_list should return None");
301    }
302}