leo_passes/storage_lowering/
program.rs

1// Copyright (C) 2019-2025 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
17use super::StorageLoweringVisitor;
18use leo_ast::{
19    AstReconstructor,
20    Identifier,
21    IntegerType,
22    Mapping,
23    ProgramReconstructor,
24    ProgramScope,
25    Statement,
26    StorageVariable,
27    Type,
28    VectorType,
29};
30use leo_span::Symbol;
31
32impl ProgramReconstructor for StorageLoweringVisitor<'_> {
33    fn reconstruct_program_scope(&mut self, input: ProgramScope) -> ProgramScope {
34        self.program = input.program_id.name.name;
35
36        let storage_variables = input
37            .storage_variables
38            .into_iter()
39            .map(|(id, storage_variable)| (id, self.reconstruct_storage_variable(storage_variable)))
40            .collect::<Vec<_>>();
41
42        // After reconstructing the storage variables, `self.new_mappings` might be non-empty.
43        // Reconstruct old mappings and then append the new mappings to the final list.
44        let mut mappings =
45            input.mappings.into_iter().map(|(id, mapping)| (id, self.reconstruct_mapping(mapping))).collect::<Vec<_>>();
46        mappings.extend(self.new_mappings.clone().into_iter().collect::<Vec<_>>());
47
48        ProgramScope {
49            program_id: input.program_id,
50            consts: input
51                .consts
52                .into_iter()
53                .map(|(i, c)| match self.reconstruct_const(c) {
54                    (Statement::Const(declaration), _) => (i, declaration),
55                    _ => panic!("`reconstruct_const` can only return `Statement::Const`"),
56                })
57                .collect(),
58            structs: input.structs.into_iter().map(|(i, c)| (i, self.reconstruct_struct(c))).collect(),
59            mappings,
60            storage_variables,
61            functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
62            constructor: input.constructor.map(|c| self.reconstruct_constructor(c)),
63            span: input.span,
64        }
65    }
66
67    fn reconstruct_storage_variable(&mut self, input: StorageVariable) -> StorageVariable {
68        // For each storage variable, create one or more backing mappings that represent
69        // how the value is stored in lower-level form.
70        //
71        // Example:
72        //   storage x: field;
73        // becomes:
74        //   mapping x__: bool => field;
75        //
76        // Example (vector):
77        //   storage vec: [u32];
78        // becomes:
79        //   mapping vec__: u32 => u32;       // element mapping
80        //   mapping vec__len__: bool => u32; // length tracking
81
82        let id = || self.state.node_builder.next_id();
83        let name = input.identifier.name.to_string();
84
85        // Common base mapping name: e.g. "x__" or "vec__"
86        let mapping_name = Symbol::intern(&(name.clone() + "__"));
87
88        match &input.type_ {
89            Type::Vector(VectorType { element_type }) => {
90                // === Vector storage ===
91                //
92                // Create two mappings:
93                //   1. `<name>__`:  index → element
94                //   2. `<name>__len__`:  bool → length
95
96                // Mapping for the vector’s contents
97                self.new_mappings.insert(mapping_name, Mapping {
98                    identifier: Identifier::new(mapping_name, id()),
99                    key_type: Type::Integer(IntegerType::U32),
100                    value_type: *element_type.clone(),
101                    span: input.span,
102                    id: id(),
103                });
104
105                // Mapping for the vector’s length
106                let len_name = Symbol::intern(&(name + "__len__"));
107                self.new_mappings.insert(len_name, Mapping {
108                    identifier: Identifier::new(len_name, id()),
109                    key_type: Type::Boolean,
110                    value_type: Type::Integer(IntegerType::U32),
111                    span: input.span,
112                    id: id(),
113                });
114            }
115
116            _ => {
117                // === Singleton storage ===
118                //
119                // Every non-vector storage variable is represented as a single mapping:
120                //   `<name>__`: bool → <type>
121                //
122                // The `bool` key acts as a presence indicator (typically `false`).
123
124                self.new_mappings.insert(mapping_name, Mapping {
125                    identifier: Identifier::new(mapping_name, id()),
126                    key_type: Type::Boolean,
127                    value_type: input.type_.clone(),
128                    span: input.span,
129                    id: id(),
130                });
131            }
132        }
133
134        // Return the original (unmodified) storage variable node.
135        input
136    }
137}