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_span::Symbol;
23
24use crate::{Module, ProgramId, Stub};
25use indexmap::IndexMap;
26use serde::Serialize;
27use std::fmt;
28/// Stores the Leo program abstract syntax tree.
29#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
30pub struct Program {
31    /// A map from module paths to module definitions.
32    #[serde(with = "module_map")]
33    pub modules: IndexMap<Vec<Symbol>, Module>,
34    /// A map from import names (including the `.aleo` if present) to import definitions.
35    pub imports: IndexMap<Symbol, ProgramId>,
36    /// A map from program stub names to program stub scopes.
37    pub stubs: IndexMap<Symbol, Stub>,
38    /// A map from program names to program scopes.
39    pub program_scopes: IndexMap<Symbol, ProgramScope>,
40}
41
42impl fmt::Display for Program {
43    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
44        for (_, stub) in self.stubs.iter() {
45            writeln!(f, "{stub}")?;
46        }
47        for (_, module) in self.modules.iter() {
48            writeln!(f, "{module}")?;
49        }
50        for (_, import_id) in self.imports.iter() {
51            writeln!(f, "import {import_id};")?;
52        }
53        for (_, program_scope) in self.program_scopes.iter() {
54            writeln!(f, "{program_scope}")?;
55        }
56        Ok(())
57    }
58}
59
60impl Default for Program {
61    /// Constructs an empty program node.
62    fn default() -> Self {
63        Self {
64            modules: IndexMap::new(),
65            imports: IndexMap::new(),
66            stubs: IndexMap::new(),
67            program_scopes: IndexMap::new(),
68        }
69    }
70}
71
72/// Helper function to recursively filter keys from AST JSON.
73pub fn remove_key_from_json(value: serde_json::Value, key: &str) -> serde_json::Value {
74    match value {
75        serde_json::Value::Object(map) => serde_json::Value::Object(
76            map.into_iter().filter(|(k, _)| k != key).map(|(k, v)| (k, remove_key_from_json(v, key))).collect(),
77        ),
78        serde_json::Value::Array(values) => {
79            serde_json::Value::Array(values.into_iter().map(|v| remove_key_from_json(v, key)).collect())
80        }
81        _ => value,
82    }
83}
84
85/// Helper function to normalize AST JSON into a form compatible with TGC.
86///
87/// This function traverses the original JSON value and produces a new one under the following rules:
88///
89/// 1. Remove empty object mappings from JSON arrays.
90/// 2. If a JSON array contains exactly two elements and one is an empty object
91///    mapping while the other is not, lift the non-empty element.
92pub fn normalize_json_value(value: serde_json::Value) -> serde_json::Value {
93    match value {
94        serde_json::Value::Array(vec) => {
95            let orig_length = vec.len();
96
97            let mut new_vec: Vec<serde_json::Value> = vec
98                .into_iter()
99                .filter(|v| !matches!(v, serde_json::Value::Object(map) if map.is_empty()))
100                .map(normalize_json_value)
101                .collect();
102
103            if orig_length == 2 && new_vec.len() == 1 {
104                new_vec.pop().unwrap()
105            } else {
106                serde_json::Value::Array(new_vec)
107            }
108        }
109        serde_json::Value::Object(map) => {
110            serde_json::Value::Object(map.into_iter().map(|(k, v)| (k, normalize_json_value(v))).collect())
111        }
112        _ => value,
113    }
114}
115
116/// Serde helpers for `IndexMap<Vec<Symbol>, V>` maps keyed by module paths.
117///
118/// JSON object keys must be strings, so the `Vec<Symbol>` path is joined into a single
119/// `::`-separated string when serializing.
120pub(crate) mod module_map {
121    use leo_span::{Symbol, with_session_globals};
122
123    use indexmap::IndexMap;
124    use serde::{Serialize, Serializer};
125
126    pub fn serialize<S, V>(map: &IndexMap<Vec<Symbol>, V>, serializer: S) -> Result<S::Ok, S::Error>
127    where
128        S: Serializer,
129        V: Serialize,
130    {
131        let joined: IndexMap<String, &V> = with_session_globals(|globals| {
132            map.iter()
133                .map(|(path, value)| {
134                    let key = path.iter().map(|sym| sym.as_str(globals, str::to_owned)).collect::<Vec<_>>().join("::");
135                    (key, value)
136                })
137                .collect()
138        });
139        joined.serialize(serializer)
140    }
141
142    #[cfg(test)]
143    mod tests {
144        use super::*;
145        use leo_span::create_session_if_not_set_then;
146
147        #[derive(Debug, PartialEq, Eq, Serialize)]
148        struct Wrapper(#[serde(with = "super")] IndexMap<Vec<Symbol>, u32>);
149
150        #[test]
151        fn serializes_single_and_multi_segment_keys() {
152            create_session_if_not_set_then(|_| {
153                let mut map = IndexMap::new();
154                map.insert(vec![Symbol::intern("utils")], 1);
155                map.insert(vec![Symbol::intern("utils"), Symbol::intern("math")], 2);
156                let wrapper = Wrapper(map);
157
158                let json = serde_json::to_value(&wrapper).unwrap();
159                assert_eq!(json, serde_json::json!({ "utils": 1, "utils::math": 2 }));
160            });
161        }
162    }
163}