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    pub(crate) parameters: NodeRef,
209    pub(crate) mappings: NodeRef,
210    pub(crate) conditions: NodeRef,
211    pub(crate) resources: NodeRef,
212    pub(crate) outputs: NodeRef,
213    pub(crate) rules: NodeRef,
214    pub(crate) template_metadata: NodeRef,
215    pub(crate) format_version: Option<String>,
216    pub(crate) description: Option<String>,
217    pub(crate) transforms: Vec<String>,
218    pub(crate) raw_top_level_keys: Vec<String>,
219    pub(crate) diagnostics: Vec<ParseDefect>,
220    pub(crate) globals: NodeRef,
221}
222
223#[derive(Debug)]
224pub struct ParseError {
225    pub message: String,
226    pub line: Option<u32>,
227    pub column: Option<u32>,
228}
229
230impl fmt::Display for ParseError {
231    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
232        match (self.line, self.column) {
233            (Some(l), Some(c)) => write!(f, "{}:{}: {}", l, c, self.message),
234            _ => write!(f, "{}", self.message),
235        }
236    }
237}
238
239impl error::Error for ParseError {}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn arena_alloc_get_roundtrip() {
247        let mut arena = Arena::new();
248        let r = arena.alloc(SpannedNode {
249            node: Node::String("hello".into()),
250            span: SourceSpan { start_line: 1, start_column: 1, end_line: 1, end_column: 5 },
251            path: "test".into(),
252        });
253        assert_eq!(arena.as_str(r), Some("hello"));
254    }
255
256    #[test]
257    fn arena_multiple_nodes() {
258        let mut arena = Arena::new();
259        let refs: Vec<NodeRef> = (0..100)
260            .map(|i| arena.alloc(SpannedNode { node: Node::Int(i), span: UNKNOWN_SPAN, path: format!("{}", i) }))
261            .collect();
262        for (i, r) in refs.iter().enumerate() {
263            match arena.node(*r) {
264                Node::Int(v) => assert_eq!(*v, i as i64),
265                _ => panic!("wrong type"),
266            }
267        }
268    }
269
270    #[test]
271    fn node_accessor_map_get() {
272        let mut arena = Arena::new();
273        let child =
274            arena.alloc(SpannedNode { node: Node::String("value".into()), span: UNKNOWN_SPAN, path: "map/key".into() });
275        let map = arena.alloc(SpannedNode {
276            node: Node::Map(vec![("key".into(), child)]),
277            span: UNKNOWN_SPAN,
278            path: "map".into(),
279        });
280        assert_eq!(arena.map_get(map, "key"), Some(child));
281        assert_eq!(arena.map_get(map, "missing"), None);
282    }
283
284    #[test]
285    fn source_span_unknown_sentinel() {
286        assert_eq!(UNKNOWN_SPAN.start_line, u32::MAX);
287        assert_eq!(UNKNOWN_SPAN.end_column, u32::MAX);
288    }
289
290    #[test]
291    fn arena_out_of_bounds_returns_sentinel() {
292        let arena = Arena::new();
293        let node = arena.get(999);
294        assert!(matches!(node.node, Node::Null));
295        assert_eq!(node.span, UNKNOWN_SPAN);
296        assert!(!arena.is_valid(999));
297        assert_eq!(arena.as_map(999), None, "out-of-bounds as_map should return None");
298        assert_eq!(arena.as_list(999), None, "out-of-bounds as_list should return None");
299    }
300}