1414// You should have received a copy of the GNU General Public License
1515// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
1616
17+ use super :: SsaConstPropagationMode ;
18+
1719use crate :: CompilerState ;
1820
19- use leo_ast:: { Expression , FromStrRadix , Literal , LiteralVariant , Location , Node , NodeID , const_eval:: Value } ;
21+ use leo_ast:: {
22+ Composite ,
23+ CompositeExpression ,
24+ Expression ,
25+ FromStrRadix ,
26+ Literal ,
27+ LiteralVariant ,
28+ Location ,
29+ Node ,
30+ NodeID ,
31+ Type ,
32+ const_eval:: Value ,
33+ } ;
2034use leo_errors:: Formatted ;
2135use leo_span:: { Span , Symbol } ;
2236
@@ -25,16 +39,23 @@ use indexmap::IndexMap;
2539/// Visitor that propagates constant values through the program.
2640pub struct SsaConstPropagationVisitor < ' a > {
2741 pub state : & ' a mut CompilerState ,
42+ /// Which subset of SSA const-prop rewrites this invocation is allowed to perform.
43+ pub mode : SsaConstPropagationMode ,
2844 /// Current program being analyzed.
2945pub program : Symbol ,
3046 /// Maps variable names to their constant values.
3147/// Only variables assigned constant values are tracked here.
3248pub constants : IndexMap < Symbol , Value > ,
33- /// Maps variable names bound to a composite RHS — where every field
34- /// initializer is an atom — to the atom each field was initialized with.
49+ /// Maps variable names bound to a composite RHS - where every field
50+ /// initializer is an atom - to the atom each field was initialized with.
3551/// Used to forward `x.field` to the stored atom without rematerializing
36- /// the enclosing struct — effectively scalarizing short-lived aggregates.
52+ /// the enclosing struct - effectively scalarizing short-lived aggregates.
3753pub atom_fielded_composites : IndexMap < Symbol , IndexMap < Symbol , Expression > > ,
54+ /// Maps local variables that are simple aliases of atom-fielded composites.
55+ pub aliases : IndexMap < Symbol , Symbol > ,
56+ /// Struct definitions visible in the current program scope, including
57+ /// compiler-generated Optional wrappers added after the symbol table.
58+ pub composites : IndexMap < Location , Composite > ,
3859 /// Have we actually modified the program at all?
3960pub changed : bool ,
4061}
@@ -46,6 +67,25 @@ pub fn is_atom(expr: &Expression) -> bool {
4667 matches ! ( expr, Expression :: Path ( _) | Expression :: Literal ( _) )
4768}
4869
70+ pub ( super ) fn optional_composite_atoms ( composite : & CompositeExpression ) -> Option < IndexMap < Symbol , Expression > > {
71+ if composite. base . is_some ( ) || composite. members . len ( ) != 2 {
72+ return None ;
73+ }
74+
75+ let mut fields = IndexMap :: with_capacity ( composite. members . len ( ) ) ;
76+ for member in & composite. members {
77+ let expr = member. expression . as_ref ( ) ?;
78+ if !is_atom ( expr) {
79+ return None ;
80+ }
81+ fields. insert ( member. identifier . name , expr. clone ( ) ) ;
82+ }
83+
84+ fields. get ( & Symbol :: intern ( "is_some" ) ) ?;
85+ fields. get ( & Symbol :: intern ( "val" ) ) ?;
86+ Some ( fields)
87+ }
88+
4989/// Parse a numeric literal string, handling underscores and radix prefixes (0x, 0o, 0b).
5090fn parse_literal_value ( s : & str ) -> Option < i128 > {
5191 let clean = s. replace ( '_' , "" ) ;
@@ -77,20 +117,75 @@ pub fn is_one_literal(lit: &Literal) -> bool {
77117 }
78118}
79119
80- /// Check if two expressions refer to the same SSA variable.
81- /// In SSA form, each variable name is unique, so name equality implies value equality.
82- pub fn same_ssa_atom ( a : & Expression , b : & Expression ) -> bool {
83- match ( a, b) {
84- ( Expression :: Path ( pa) , Expression :: Path ( pb) ) => {
85- let sa = pa. try_local_symbol ( ) ;
86- let sb = pb. try_local_symbol ( ) ;
87- sa == sb && sa. is_some ( )
120+ impl SsaConstPropagationVisitor < ' _ > {
121+ pub ( super ) fn runs_full_const_prop ( & self ) -> bool {
122+ self . mode == SsaConstPropagationMode :: Full
123+ }
124+
125+ pub ( super ) fn forwards_composite_members ( & self ) -> bool {
126+ self . runs_full_const_prop ( )
127+ }
128+
129+ pub ( super ) fn tracks_optional_unwraps ( & self ) -> bool {
130+ self . mode == SsaConstPropagationMode :: LateOptionalUnwrap
131+ }
132+
133+ /// Clear analysis state that is only valid within one SSA function body.
134+ pub ( super ) fn clear_tracked_values ( & mut self ) {
135+ self . constants . clear ( ) ;
136+ self . atom_fielded_composites . clear ( ) ;
137+ self . aliases . clear ( ) ;
138+ }
139+
140+ pub ( super ) fn is_optional_wrapper_type ( & self , ty : & Type ) -> bool {
141+ let Type :: Composite ( composite_ty) = ty else {
142+ return false ;
143+ } ;
144+ let location = composite_ty. path . expect_global_location ( ) ;
145+ let Some ( composite) =
146+ self . state . symbol_table . lookup_struct ( self . program , location) . or_else ( || self . composites . get ( location) )
147+ else {
148+ return false ;
149+ } ;
150+ let [ is_some, val] = composite. members . as_slice ( ) else {
151+ return false ;
152+ } ;
153+ is_some. identifier . name == Symbol :: intern ( "is_some" )
154+ && matches ! ( is_some. type_, Type :: Boolean )
155+ && val. identifier . name == Symbol :: intern ( "val" )
156+ }
157+
158+ pub ( super ) fn is_optional_composite_expression ( & self , composite : & CompositeExpression ) -> bool {
159+ self . state . type_table . get ( & composite. id ( ) ) . is_some_and ( |ty| self . is_optional_wrapper_type ( & ty) )
160+ }
161+
162+ pub ( super ) fn resolve_composite_alias ( & self , mut name : Symbol ) -> Symbol {
163+ for _ in 0 ..self . aliases . len ( ) {
164+ let Some ( alias) = self . aliases . get ( & name) . copied ( ) else {
165+ break ;
166+ } ;
167+ if alias == name {
168+ break ;
169+ }
170+ name = alias;
88171 }
89- _ => false ,
172+ name
173+ }
174+
175+ pub ( super ) fn atom_for_composite_member (
176+ & self ,
177+ expression : & Expression ,
178+ field : Symbol ,
179+ ) -> Option < ( Symbol , Expression ) > {
180+ let Expression :: Path ( path) = expression else {
181+ return None ;
182+ } ;
183+ let original_name = path. try_local_symbol ( ) ?;
184+ let name = self . resolve_composite_alias ( original_name) ;
185+ let atom = self . atom_fielded_composites . get ( & name) ?. get ( & field) ?. clone ( ) ;
186+ Some ( ( name, atom) )
90187 }
91- }
92188
93- impl SsaConstPropagationVisitor < ' _ > {
94189 /// Emit a `StaticAnalyzerError`.
95190pub fn emit_err ( & self , err : Formatted ) {
96191 self . state . handler . emit_err ( err) ;
0 commit comments