Skip to main content

leo_ast/program/
mod.rs

1// Copyright (C) 2019-2026 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17//! A Leo program consists of import statements and program scopes.
18
19mod program_scope;
20pub use program_scope::*;
21
22use leo_errors::Result;
23use leo_span::Symbol;
24
25use crate::{Module, ProgramId, Stub};
26use indexmap::IndexMap;
27use serde::{Deserialize, Serialize};
28use std::fmt;
29/// Stores the Leo program abstract syntax tree.
30#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
31pub struct Program {
32    /// A map from module paths to module definitions.
33    #[serde(with = "module_map")]
34    pub modules: IndexMap<Vec<Symbol>, Module>,
35    /// A map from import names (including the `.aleo` if present) to import definitions.
36    pub imports: IndexMap<Symbol, ProgramId>,
37    /// A map from program stub names to program stub scopes.
38    pub stubs: IndexMap<Symbol, Stub>,
39    /// A map from program names to program scopes.
40    pub program_scopes: IndexMap<Symbol, ProgramScope>,
41}
42
43impl fmt::Display for Program {
44    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
45        for (_, stub) in self.stubs.iter() {
46            writeln!(f, "{stub}")?;
47        }
48        for (_, module) in self.modules.iter() {
49            writeln!(f, "{module}")?;
50        }
51        for (_, import_id) in self.imports.iter() {
52            writeln!(f, "import {import_id};")?;
53        }
54        for (_, program_scope) in self.program_scopes.iter() {
55            writeln!(f, "{program_scope}")?;
56        }
57        Ok(())
58    }
59}
60
61impl Default for Program {
62    /// Constructs an empty program node.
63    fn default() -> Self {
64        Self {
65            modules: IndexMap::new(),
66            imports: IndexMap::new(),
67            stubs: IndexMap::new(),
68            program_scopes: IndexMap::new(),
69        }
70    }
71}
72
73impl Program {
74    /// Serializes the ast into a JSON string.
75    pub fn to_json_string(&self) -> Result<String> {
76        Ok(serde_json::to_string_pretty(&self).map_err(|e| crate::errors::failed_to_convert_ast_to_json_string(&e))?)
77    }
78
79    // Converts the ast into a JSON value.
80    // Note that there is no corresponding `from_json_value` function
81    // since we modify JSON values leaving them unable to be converted
82    // back into Programs.
83    pub fn to_json_value(&self) -> Result<serde_json::Value> {
84        Ok(serde_json::to_value(self).map_err(|e| crate::errors::failed_to_convert_ast_to_json_value(&e))?)
85    }
86
87    /// Serializes the ast into a JSON file.
88    pub fn to_json_file(&self, path: std::path::PathBuf, file_name: &str) -> Result<()> {
89        write_ast_json(self, path, file_name)
90    }
91
92    /// Serializes the ast into a JSON value and removes keys from object mappings before writing to a file.
93    pub fn to_json_file_without_keys(
94        &self,
95        path: std::path::PathBuf,
96        file_name: &str,
97        excluded_keys: &[&str],
98    ) -> Result<()> {
99        write_ast_json_filtered(self, path, file_name, excluded_keys)
100    }
101
102    /// Deserializes the JSON string into a ast.
103    pub fn from_json_string(json: &str) -> Result<Self> {
104        let ast: Program =
105            serde_json::from_str(json).map_err(|e| crate::errors::failed_to_read_json_string_to_ast(&e))?;
106        Ok(ast)
107    }
108
109    /// Deserializes the JSON string into a ast from a file.
110    pub fn from_json_file(path: std::path::PathBuf) -> Result<Self> {
111        let data = std::fs::read_to_string(&path).map_err(|e| crate::errors::failed_to_read_json_file(&path, &e))?;
112        Self::from_json_string(&data)
113    }
114}
115
116/// Helper function to recursively filter keys from AST JSON.
117pub fn remove_key_from_json(value: serde_json::Value, key: &str) -> serde_json::Value {
118    match value {
119        serde_json::Value::Object(map) => serde_json::Value::Object(
120            map.into_iter().filter(|(k, _)| k != key).map(|(k, v)| (k, remove_key_from_json(v, key))).collect(),
121        ),
122        serde_json::Value::Array(values) => {
123            serde_json::Value::Array(values.into_iter().map(|v| remove_key_from_json(v, key)).collect())
124        }
125        _ => value,
126    }
127}
128
129/// Helper function to normalize AST JSON into a form compatible with TGC.
130///
131/// This function traverses the original JSON value and produces a new one under the following rules:
132///
133/// 1. Remove empty object mappings from JSON arrays.
134/// 2. If a JSON array contains exactly two elements and one is an empty object
135///    mapping while the other is not, lift the non-empty element.
136pub fn normalize_json_value(value: serde_json::Value) -> serde_json::Value {
137    match value {
138        serde_json::Value::Array(vec) => {
139            let orig_length = vec.len();
140
141            let mut new_vec: Vec<serde_json::Value> = vec
142                .into_iter()
143                .filter(|v| !matches!(v, serde_json::Value::Object(map) if map.is_empty()))
144                .map(normalize_json_value)
145                .collect();
146
147            if orig_length == 2 && new_vec.len() == 1 {
148                new_vec.pop().unwrap()
149            } else {
150                serde_json::Value::Array(new_vec)
151            }
152        }
153        serde_json::Value::Object(map) => {
154            serde_json::Value::Object(map.into_iter().map(|(k, v)| (k, normalize_json_value(v))).collect())
155        }
156        _ => value,
157    }
158}
159
160/// Serializes an AST node (`Program` or `Library`) into a pretty JSON file, spans included.
161pub fn write_ast_json<T: Serialize>(value: &T, mut path: std::path::PathBuf, file_name: &str) -> Result<()> {
162    path.push(file_name);
163    let file = std::fs::File::create(&path).map_err(|e| crate::errors::failed_to_create_ast_json_file(&path, &e))?;
164    let writer = std::io::BufWriter::new(file);
165    Ok(serde_json::to_writer_pretty(writer, value)
166        .map_err(|e| crate::errors::failed_to_write_ast_to_json_file(&path, &e))?)
167}
168
169/// Serializes an AST node to a JSON file, stripping `excluded_keys` and normalizing first.
170pub fn write_ast_json_filtered<T: Serialize>(
171    value: &T,
172    mut path: std::path::PathBuf,
173    file_name: &str,
174    excluded_keys: &[&str],
175) -> Result<()> {
176    path.push(file_name);
177    let file = std::fs::File::create(&path).map_err(|e| crate::errors::failed_to_create_ast_json_file(&path, &e))?;
178    let writer = std::io::BufWriter::new(file);
179
180    let mut value = serde_json::to_value(value).map_err(|e| crate::errors::failed_to_convert_ast_to_json_value(&e))?;
181    for key in excluded_keys {
182        value = remove_key_from_json(value, key);
183    }
184    value = normalize_json_value(value);
185
186    Ok(serde_json::to_writer_pretty(writer, &value)
187        .map_err(|e| crate::errors::failed_to_write_ast_to_json_file(&path, &e))?)
188}
189
190/// Serde helpers for `IndexMap<Vec<Symbol>, V>` maps keyed by module paths.
191///
192/// JSON object keys must be strings, so the `Vec<Symbol>` path is joined into a single
193/// `::`-separated string when serializing and split back when deserializing.
194pub(crate) mod module_map {
195    use leo_span::{Symbol, with_session_globals};
196
197    use indexmap::IndexMap;
198    use serde::{Deserialize, Deserializer, Serialize, Serializer};
199
200    pub fn serialize<S, V>(map: &IndexMap<Vec<Symbol>, V>, serializer: S) -> Result<S::Ok, S::Error>
201    where
202        S: Serializer,
203        V: Serialize,
204    {
205        let joined: IndexMap<String, &V> = with_session_globals(|globals| {
206            map.iter()
207                .map(|(path, value)| {
208                    let key = path.iter().map(|sym| sym.as_str(globals, str::to_owned)).collect::<Vec<_>>().join("::");
209                    (key, value)
210                })
211                .collect()
212        });
213        joined.serialize(serializer)
214    }
215
216    pub fn deserialize<'de, D, V>(deserializer: D) -> Result<IndexMap<Vec<Symbol>, V>, D::Error>
217    where
218        D: Deserializer<'de>,
219        V: Deserialize<'de>,
220    {
221        Ok(IndexMap::<String, V>::deserialize(deserializer)?
222            .into_iter()
223            .map(|(path, value)| (path.split("::").map(Symbol::intern).collect(), value))
224            .collect())
225    }
226
227    #[cfg(test)]
228    mod tests {
229        use super::*;
230        use leo_span::create_session_if_not_set_then;
231
232        #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
233        struct Wrapper(#[serde(with = "super")] IndexMap<Vec<Symbol>, u32>);
234
235        #[test]
236        fn round_trips_single_and_multi_segment_keys() {
237            create_session_if_not_set_then(|_| {
238                let mut map = IndexMap::new();
239                map.insert(vec![Symbol::intern("utils")], 1);
240                map.insert(vec![Symbol::intern("utils"), Symbol::intern("math")], 2);
241                let wrapper = Wrapper(map);
242
243                let json = serde_json::to_value(&wrapper).unwrap();
244                assert_eq!(json, serde_json::json!({ "utils": 1, "utils::math": 2 }));
245
246                let restored: Wrapper = serde_json::from_value(json).unwrap();
247                assert_eq!(restored, wrapper);
248            });
249        }
250    }
251}