oxilean_codegen/opt_specialize/
functions.rs

1//! Auto-generated module
2//!
3//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
4
5use crate::lcnf::*;
6use std::collections::{HashMap, HashSet};
7
8use super::types::{
9    NumericSpecializer, SizeBudget, SpecAnalysisCache, SpecCallSite, SpecClosureArg, SpecConstArg,
10    SpecConstantFoldingHelper, SpecDepGraph, SpecDominatorTree, SpecExtCache, SpecExtConstFolder,
11    SpecExtDepGraph, SpecExtDomTree, SpecExtLiveness, SpecExtPassConfig, SpecExtPassPhase,
12    SpecExtPassRegistry, SpecExtPassStats, SpecExtWorklist, SpecLivenessInfo, SpecPassConfig,
13    SpecPassPhase, SpecPassRegistry, SpecPassStats, SpecTypeArg, SpecWorklist, SpecializationCache,
14    SpecializationConfig, SpecializationKey, SpecializationPass, SpecializationStats,
15};
16
17/// Generate a short suffix for a type
18pub(super) fn type_suffix(ty: &LcnfType) -> String {
19    match ty {
20        LcnfType::Nat => "nat".to_string(),
21        LcnfType::Int => "int".to_string(),
22        LcnfType::Object => "obj".to_string(),
23        LcnfType::Unit => "unit".to_string(),
24        LcnfType::Erased => "e".to_string(),
25        LcnfType::LcnfString => "str".to_string(),
26        LcnfType::Var(name) => name.clone(),
27        LcnfType::Ctor(name, _) => name.clone(),
28        LcnfType::Fun(_, _) => "fn".to_string(),
29        LcnfType::Irrelevant => "irr".to_string(),
30    }
31}
32/// Analyze call sites in a function body to find specialization opportunities
33pub(super) fn find_specialization_sites(
34    expr: &LcnfExpr,
35    known_constants: &HashMap<LcnfVarId, LcnfLit>,
36    known_functions: &HashMap<LcnfVarId, String>,
37    decl_names: &HashSet<String>,
38) -> Vec<SpecCallSite> {
39    let mut sites = Vec::new();
40    let mut call_idx = 0;
41    find_spec_sites_inner(
42        expr,
43        known_constants,
44        known_functions,
45        decl_names,
46        &mut sites,
47        &mut call_idx,
48    );
49    sites
50}
51pub(super) fn find_spec_sites_inner(
52    expr: &LcnfExpr,
53    known_constants: &HashMap<LcnfVarId, LcnfLit>,
54    known_functions: &HashMap<LcnfVarId, String>,
55    decl_names: &HashSet<String>,
56    sites: &mut Vec<SpecCallSite>,
57    call_idx: &mut usize,
58) {
59    match expr {
60        LcnfExpr::Let {
61            id, value, body, ..
62        } => {
63            let mut extended_consts = known_constants.clone();
64            if let LcnfLetValue::Lit(lit) = value {
65                extended_consts.insert(*id, lit.clone());
66            }
67            let mut extended_fns = known_functions.clone();
68            if let LcnfLetValue::FVar(fvar) = value {
69                if let Some(fname) = known_functions.get(fvar) {
70                    extended_fns.insert(*id, fname.clone());
71                }
72            }
73            if let LcnfLetValue::App(func, args) = value {
74                let callee_name = match func {
75                    LcnfArg::Var(v) => known_functions.get(v).cloned(),
76                    _ => None,
77                };
78                if let Some(ref callee) = callee_name {
79                    if decl_names.contains(callee.as_str()) {
80                        let const_args: Vec<SpecConstArg> = args
81                            .iter()
82                            .map(|arg| match arg {
83                                LcnfArg::Lit(LcnfLit::Nat(n)) => SpecConstArg::Nat(*n),
84                                LcnfArg::Lit(LcnfLit::Str(s)) => SpecConstArg::Str(s.clone()),
85                                LcnfArg::Var(v) => {
86                                    if let Some(lit) = extended_consts.get(v) {
87                                        match lit {
88                                            LcnfLit::Nat(n) => SpecConstArg::Nat(*n),
89                                            LcnfLit::Int(_) => SpecConstArg::Unknown,
90                                            LcnfLit::Str(s) => SpecConstArg::Str(s.clone()),
91                                        }
92                                    } else {
93                                        SpecConstArg::Unknown
94                                    }
95                                }
96                                _ => SpecConstArg::Unknown,
97                            })
98                            .collect();
99                        let closure_args: Vec<SpecClosureArg> = args
100                            .iter()
101                            .enumerate()
102                            .map(|(i, arg)| {
103                                let known_fn = match arg {
104                                    LcnfArg::Var(v) => extended_fns.get(v).cloned(),
105                                    _ => None,
106                                };
107                                SpecClosureArg {
108                                    known_fn,
109                                    param_idx: i,
110                                }
111                            })
112                            .collect();
113                        let callee_var = match func {
114                            LcnfArg::Var(v) => Some(*v),
115                            _ => None,
116                        };
117                        sites.push(SpecCallSite {
118                            callee: callee.clone(),
119                            call_idx: *call_idx,
120                            type_args: vec![],
121                            const_args,
122                            closure_args,
123                            callee_var,
124                        });
125                        *call_idx += 1;
126                    }
127                }
128            }
129            find_spec_sites_inner(
130                body,
131                &extended_consts,
132                &extended_fns,
133                decl_names,
134                sites,
135                call_idx,
136            );
137        }
138        LcnfExpr::Case { alts, default, .. } => {
139            for alt in alts {
140                find_spec_sites_inner(
141                    &alt.body,
142                    known_constants,
143                    known_functions,
144                    decl_names,
145                    sites,
146                    call_idx,
147                );
148            }
149            if let Some(def) = default {
150                find_spec_sites_inner(
151                    def,
152                    known_constants,
153                    known_functions,
154                    decl_names,
155                    sites,
156                    call_idx,
157                );
158            }
159        }
160        LcnfExpr::TailCall(func, args) => {
161            let callee_name = match func {
162                LcnfArg::Var(v) => known_functions.get(v).cloned(),
163                _ => None,
164            };
165            if let Some(callee) = callee_name {
166                if decl_names.contains(callee.as_str()) {
167                    let const_args: Vec<SpecConstArg> = args
168                        .iter()
169                        .map(|arg| match arg {
170                            LcnfArg::Lit(LcnfLit::Nat(n)) => SpecConstArg::Nat(*n),
171                            LcnfArg::Lit(LcnfLit::Str(s)) => SpecConstArg::Str(s.clone()),
172                            LcnfArg::Var(v) => {
173                                if let Some(lit) = known_constants.get(v) {
174                                    match lit {
175                                        LcnfLit::Nat(n) => SpecConstArg::Nat(*n),
176                                        LcnfLit::Int(_) => SpecConstArg::Unknown,
177                                        LcnfLit::Str(s) => SpecConstArg::Str(s.clone()),
178                                    }
179                                } else {
180                                    SpecConstArg::Unknown
181                                }
182                            }
183                            _ => SpecConstArg::Unknown,
184                        })
185                        .collect();
186                    let closure_args: Vec<SpecClosureArg> = args
187                        .iter()
188                        .enumerate()
189                        .map(|(i, arg)| {
190                            let known_fn = match arg {
191                                LcnfArg::Var(v) => known_functions.get(v).cloned(),
192                                _ => None,
193                            };
194                            SpecClosureArg {
195                                known_fn,
196                                param_idx: i,
197                            }
198                        })
199                        .collect();
200                    let callee_var = match func {
201                        LcnfArg::Var(v) => Some(*v),
202                        _ => None,
203                    };
204                    sites.push(SpecCallSite {
205                        callee,
206                        call_idx: *call_idx,
207                        type_args: vec![],
208                        const_args,
209                        closure_args,
210                        callee_var,
211                    });
212                    *call_idx += 1;
213                }
214            }
215        }
216        LcnfExpr::Return(_) | LcnfExpr::Unreachable => {}
217    }
218}
219/// Count the number of LCNF instructions in an expression
220pub(super) fn count_instructions(expr: &LcnfExpr) -> usize {
221    match expr {
222        LcnfExpr::Let { body, .. } => 1 + count_instructions(body),
223        LcnfExpr::Case { alts, default, .. } => {
224            let alts_size: usize = alts.iter().map(|a| count_instructions(&a.body)).sum();
225            let def_size = default.as_ref().map(|d| count_instructions(d)).unwrap_or(0);
226            1 + alts_size + def_size
227        }
228        LcnfExpr::Return(_) | LcnfExpr::TailCall(_, _) | LcnfExpr::Unreachable => 1,
229    }
230}
231/// Analyze whether a function parameter is always called with the same closure
232pub(super) fn analyze_closure_uniformity(
233    decl: &LcnfFunDecl,
234    param_idx: usize,
235    sites: &[SpecCallSite],
236) -> Option<String> {
237    let mut known_fn: Option<String> = None;
238    for site in sites {
239        if site.callee != decl.name {
240            continue;
241        }
242        if param_idx >= site.closure_args.len() {
243            return None;
244        }
245        match &site.closure_args[param_idx].known_fn {
246            Some(fn_name) => {
247                if let Some(ref existing) = known_fn {
248                    if existing != fn_name {
249                        return None;
250                    }
251                } else {
252                    known_fn = Some(fn_name.clone());
253                }
254            }
255            None => return None,
256        }
257    }
258    known_fn
259}
260/// Check whether a parameter is used as a function (called) in the body
261pub(super) fn is_called_as_function(expr: &LcnfExpr, param_id: LcnfVarId) -> bool {
262    match expr {
263        LcnfExpr::Let { value, body, .. } => {
264            let called_here = matches!(
265                value, LcnfLetValue::App(LcnfArg::Var(v), _) if * v == param_id
266            );
267            called_here || is_called_as_function(body, param_id)
268        }
269        LcnfExpr::Case { alts, default, .. } => {
270            alts.iter()
271                .any(|a| is_called_as_function(&a.body, param_id))
272                || default
273                    .as_ref()
274                    .is_some_and(|d| is_called_as_function(d, param_id))
275        }
276        LcnfExpr::TailCall(LcnfArg::Var(v), _) => *v == param_id,
277        _ => false,
278    }
279}
280/// Main entry point: specialize functions in a module
281pub fn specialize_module(module: &mut LcnfModule, config: &SpecializationConfig) {
282    let mut pass = SpecializationPass::new(config.clone());
283    pass.run(module);
284}
285/// Specialize a single function for numeric operations
286pub fn specialize_numeric(decl: &LcnfFunDecl) -> Option<LcnfFunDecl> {
287    let specializer = NumericSpecializer::new();
288    if !specializer.is_numeric_op(&decl.name) {
289        return None;
290    }
291    let mut spec = decl.clone();
292    spec.name = format!("{}_u64", decl.name);
293    for param in &mut spec.params {
294        param.ty = specializer.specialize_nat_to_u64(&param.ty);
295    }
296    spec.ret_type = specializer.specialize_nat_to_u64(&spec.ret_type);
297    Some(spec)
298}
299/// Check if a function is worth specializing based on heuristics
300pub fn is_worth_specializing(decl: &LcnfFunDecl, config: &SpecializationConfig) -> bool {
301    let size = count_instructions(&decl.body);
302    if size > config.size_threshold {
303        return false;
304    }
305    let has_poly = decl
306        .params
307        .iter()
308        .any(|p| matches!(p.ty, LcnfType::Var(_) | LcnfType::Object | LcnfType::Erased));
309    let has_fn_param = decl
310        .params
311        .iter()
312        .any(|p| matches!(p.ty, LcnfType::Fun(_, _)));
313    has_poly || (has_fn_param && config.specialize_closures)
314}
315#[cfg(test)]
316mod tests {
317    use super::*;
318    pub(super) fn make_var(n: u64) -> LcnfVarId {
319        LcnfVarId(n)
320    }
321    pub(super) fn make_param(n: u64, name: &str, ty: LcnfType) -> LcnfParam {
322        LcnfParam {
323            id: LcnfVarId(n),
324            name: name.to_string(),
325            ty,
326            erased: false,
327            borrowed: false,
328        }
329    }
330    pub(super) fn make_simple_let(id: u64, value: LcnfLetValue, body: LcnfExpr) -> LcnfExpr {
331        LcnfExpr::Let {
332            id: LcnfVarId(id),
333            name: format!("x{}", id),
334            ty: LcnfType::Nat,
335            value,
336            body: Box::new(body),
337        }
338    }
339    pub(super) fn make_decl(name: &str, params: Vec<LcnfParam>, body: LcnfExpr) -> LcnfFunDecl {
340        LcnfFunDecl {
341            name: name.to_string(),
342            original_name: None,
343            params,
344            ret_type: LcnfType::Nat,
345            body,
346            is_recursive: false,
347            is_lifted: false,
348            inline_cost: 1,
349        }
350    }
351    #[test]
352    pub(super) fn test_config_default() {
353        let config = SpecializationConfig::default();
354        assert_eq!(config.max_specializations, 8);
355        assert!(config.specialize_closures);
356        assert!(config.specialize_numerics);
357        assert_eq!(config.size_threshold, 200);
358    }
359    #[test]
360    pub(super) fn test_spec_key_trivial() {
361        let key = SpecializationKey {
362            original: "foo".to_string(),
363            type_args: vec![SpecTypeArg::Poly],
364            const_args: vec![SpecConstArg::Unknown],
365            closure_args: vec![SpecClosureArg {
366                known_fn: None,
367                param_idx: 0,
368            }],
369        };
370        assert!(key.is_trivial());
371    }
372    #[test]
373    pub(super) fn test_spec_key_non_trivial_type() {
374        let key = SpecializationKey {
375            original: "foo".to_string(),
376            type_args: vec![SpecTypeArg::Concrete(LcnfType::Nat)],
377            const_args: vec![],
378            closure_args: vec![],
379        };
380        assert!(!key.is_trivial());
381    }
382    #[test]
383    pub(super) fn test_spec_key_non_trivial_const() {
384        let key = SpecializationKey {
385            original: "foo".to_string(),
386            type_args: vec![],
387            const_args: vec![SpecConstArg::Nat(42)],
388            closure_args: vec![],
389        };
390        assert!(!key.is_trivial());
391    }
392    #[test]
393    pub(super) fn test_spec_key_mangled_name() {
394        let key = SpecializationKey {
395            original: "List.map".to_string(),
396            type_args: vec![SpecTypeArg::Concrete(LcnfType::Nat)],
397            const_args: vec![SpecConstArg::Unknown],
398            closure_args: vec![],
399        };
400        let name = key.mangled_name();
401        assert!(name.starts_with("List.map"));
402        assert!(name.contains("_T0_nat"));
403    }
404    #[test]
405    pub(super) fn test_spec_key_mangled_name_with_const() {
406        let key = SpecializationKey {
407            original: "repeat".to_string(),
408            type_args: vec![],
409            const_args: vec![SpecConstArg::Nat(3)],
410            closure_args: vec![],
411        };
412        let name = key.mangled_name();
413        assert!(name.contains("_C0_N3"));
414    }
415    #[test]
416    pub(super) fn test_spec_key_mangled_name_with_closure() {
417        let key = SpecializationKey {
418            original: "List.map".to_string(),
419            type_args: vec![],
420            const_args: vec![],
421            closure_args: vec![SpecClosureArg {
422                known_fn: Some("double".to_string()),
423                param_idx: 0,
424            }],
425        };
426        let name = key.mangled_name();
427        assert!(name.contains("_Fdouble"));
428    }
429    #[test]
430    pub(super) fn test_type_suffix() {
431        assert_eq!(type_suffix(&LcnfType::Nat), "nat");
432        assert_eq!(type_suffix(&LcnfType::Object), "obj");
433        assert_eq!(type_suffix(&LcnfType::Unit), "unit");
434        assert_eq!(type_suffix(&LcnfType::LcnfString), "str");
435    }
436    #[test]
437    pub(super) fn test_cache_operations() {
438        let mut cache = SpecializationCache::new();
439        let key = SpecializationKey {
440            original: "foo".to_string(),
441            type_args: vec![SpecTypeArg::Concrete(LcnfType::Nat)],
442            const_args: vec![],
443            closure_args: vec![],
444        };
445        assert!(cache.lookup(&key).is_none());
446        cache.insert(key.clone(), "foo_nat".to_string(), 10);
447        assert_eq!(cache.lookup(&key), Some("foo_nat"));
448        assert_eq!(cache.specialization_count("foo"), 1);
449        assert_eq!(cache.total_growth, 10);
450    }
451    #[test]
452    pub(super) fn test_size_budget() {
453        let mut budget = SizeBudget::new(100, 2.0);
454        assert!(budget.can_afford(50));
455        assert!(budget.can_afford(100));
456        assert!(!budget.can_afford(101));
457        budget.spend(50);
458        assert!(budget.can_afford(50));
459        assert!(!budget.can_afford(51));
460        assert_eq!(budget.remaining(), 50);
461    }
462    #[test]
463    pub(super) fn test_numeric_specializer() {
464        let specializer = NumericSpecializer::new();
465        assert!(specializer.is_numeric_op("Nat.add"));
466        assert!(specializer.is_numeric_op("Nat.mul"));
467        assert!(!specializer.is_numeric_op("List.map"));
468    }
469    #[test]
470    pub(super) fn test_numeric_type_specialization() {
471        let specializer = NumericSpecializer::new();
472        let ty = LcnfType::Fun(vec![LcnfType::Nat, LcnfType::Nat], Box::new(LcnfType::Nat));
473        let spec = specializer.specialize_nat_to_u64(&ty);
474        assert_eq!(
475            spec,
476            LcnfType::Fun(vec![LcnfType::Nat, LcnfType::Nat], Box::new(LcnfType::Nat))
477        );
478    }
479    #[test]
480    pub(super) fn test_specialize_numeric() {
481        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
482        let decl = make_decl(
483            "Nat.add",
484            vec![
485                make_param(0, "a", LcnfType::Nat),
486                make_param(1, "b", LcnfType::Nat),
487            ],
488            body,
489        );
490        let result = specialize_numeric(&decl);
491        assert!(result.is_some());
492        let spec = result.expect("spec should be Some/Ok");
493        assert_eq!(spec.name, "Nat.add_u64");
494    }
495    #[test]
496    pub(super) fn test_specialize_numeric_non_numeric() {
497        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
498        let decl = make_decl("List.map", vec![make_param(0, "f", LcnfType::Object)], body);
499        let result = specialize_numeric(&decl);
500        assert!(result.is_none());
501    }
502    #[test]
503    pub(super) fn test_is_worth_specializing_polymorphic() {
504        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
505        let decl = make_decl(
506            "id",
507            vec![make_param(0, "x", LcnfType::Var("a".to_string()))],
508            body,
509        );
510        let config = SpecializationConfig::default();
511        assert!(is_worth_specializing(&decl, &config));
512    }
513    #[test]
514    pub(super) fn test_is_worth_specializing_concrete() {
515        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
516        let decl = make_decl(
517            "add",
518            vec![
519                make_param(0, "a", LcnfType::Nat),
520                make_param(1, "b", LcnfType::Nat),
521            ],
522            body,
523        );
524        let config = SpecializationConfig::default();
525        assert!(!is_worth_specializing(&decl, &config));
526    }
527    #[test]
528    pub(super) fn test_is_worth_specializing_higher_order() {
529        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
530        let fn_ty = LcnfType::Fun(vec![LcnfType::Nat], Box::new(LcnfType::Nat));
531        let decl = make_decl("apply", vec![make_param(0, "f", fn_ty)], body);
532        let config = SpecializationConfig::default();
533        assert!(is_worth_specializing(&decl, &config));
534    }
535    #[test]
536    pub(super) fn test_is_called_as_function() {
537        let body = make_simple_let(
538            5,
539            LcnfLetValue::App(
540                LcnfArg::Var(make_var(0)),
541                vec![LcnfArg::Lit(LcnfLit::Nat(1))],
542            ),
543            LcnfExpr::Return(LcnfArg::Var(make_var(5))),
544        );
545        assert!(is_called_as_function(&body, make_var(0)));
546        assert!(!is_called_as_function(&body, make_var(1)));
547    }
548    #[test]
549    pub(super) fn test_count_instructions() {
550        let expr = make_simple_let(
551            1,
552            LcnfLetValue::Lit(LcnfLit::Nat(42)),
553            make_simple_let(
554                2,
555                LcnfLetValue::Lit(LcnfLit::Nat(10)),
556                LcnfExpr::Return(LcnfArg::Var(make_var(2))),
557            ),
558        );
559        assert_eq!(count_instructions(&expr), 3);
560    }
561    #[test]
562    pub(super) fn test_substitute_constant() {
563        let mut expr = make_simple_let(
564            1,
565            LcnfLetValue::FVar(make_var(0)),
566            LcnfExpr::Return(LcnfArg::Var(make_var(1))),
567        );
568        let pass = SpecializationPass::new(SpecializationConfig::default());
569        pass.substitute_constant(&mut expr, make_var(0), &LcnfLit::Nat(42));
570        if let LcnfExpr::Let { value, .. } = &expr {
571            assert_eq!(*value, LcnfLetValue::Lit(LcnfLit::Nat(42)));
572        } else {
573            panic!("Expected Let");
574        }
575    }
576    #[test]
577    pub(super) fn test_substitute_constant_in_return() {
578        let mut expr = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
579        let pass = SpecializationPass::new(SpecializationConfig::default());
580        pass.substitute_constant(&mut expr, make_var(0), &LcnfLit::Nat(99));
581        assert_eq!(expr, LcnfExpr::Return(LcnfArg::Lit(LcnfLit::Nat(99))));
582    }
583    #[test]
584    pub(super) fn test_specialize_module_empty() {
585        let mut module = LcnfModule::default();
586        let config = SpecializationConfig::default();
587        specialize_module(&mut module, &config);
588        assert!(module.fun_decls.is_empty());
589    }
590    #[test]
591    pub(super) fn test_specialize_module_simple() {
592        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
593        let decl = make_decl(
594            "id",
595            vec![make_param(0, "x", LcnfType::Var("a".to_string()))],
596            body,
597        );
598        let mut module = LcnfModule {
599            fun_decls: vec![decl],
600            extern_decls: vec![],
601            name: "test".to_string(),
602            metadata: LcnfModuleMetadata::default(),
603        };
604        let config = SpecializationConfig::default();
605        specialize_module(&mut module, &config);
606        assert!(!module.fun_decls.is_empty());
607    }
608    #[test]
609    pub(super) fn test_closure_uniformity_analysis() {
610        let body = LcnfExpr::Return(LcnfArg::Lit(LcnfLit::Nat(0)));
611        let decl = make_decl("apply", vec![make_param(0, "f", LcnfType::Object)], body);
612        let sites = vec![
613            SpecCallSite {
614                callee: "apply".to_string(),
615                call_idx: 0,
616                type_args: vec![],
617                const_args: vec![],
618                closure_args: vec![SpecClosureArg {
619                    known_fn: Some("double".to_string()),
620                    param_idx: 0,
621                }],
622                callee_var: None,
623            },
624            SpecCallSite {
625                callee: "apply".to_string(),
626                call_idx: 1,
627                type_args: vec![],
628                const_args: vec![],
629                closure_args: vec![SpecClosureArg {
630                    known_fn: Some("double".to_string()),
631                    param_idx: 0,
632                }],
633                callee_var: None,
634            },
635        ];
636        let result = analyze_closure_uniformity(&decl, 0, &sites);
637        assert_eq!(result, Some("double".to_string()));
638    }
639    #[test]
640    pub(super) fn test_closure_uniformity_non_uniform() {
641        let body = LcnfExpr::Return(LcnfArg::Lit(LcnfLit::Nat(0)));
642        let decl = make_decl("apply", vec![make_param(0, "f", LcnfType::Object)], body);
643        let sites = vec![
644            SpecCallSite {
645                callee: "apply".to_string(),
646                call_idx: 0,
647                type_args: vec![],
648                const_args: vec![],
649                closure_args: vec![SpecClosureArg {
650                    known_fn: Some("double".to_string()),
651                    param_idx: 0,
652                }],
653                callee_var: None,
654            },
655            SpecCallSite {
656                callee: "apply".to_string(),
657                call_idx: 1,
658                type_args: vec![],
659                const_args: vec![],
660                closure_args: vec![SpecClosureArg {
661                    known_fn: Some("triple".to_string()),
662                    param_idx: 0,
663                }],
664                callee_var: None,
665            },
666        ];
667        let result = analyze_closure_uniformity(&decl, 0, &sites);
668        assert!(result.is_none());
669    }
670    #[test]
671    pub(super) fn test_find_specialization_sites() {
672        let body = make_simple_let(
673            1,
674            LcnfLetValue::App(
675                LcnfArg::Var(make_var(10)),
676                vec![LcnfArg::Lit(LcnfLit::Nat(42))],
677            ),
678            LcnfExpr::Return(LcnfArg::Var(make_var(1))),
679        );
680        let known_consts: HashMap<LcnfVarId, LcnfLit> = HashMap::new();
681        let mut known_fns: HashMap<LcnfVarId, String> = HashMap::new();
682        known_fns.insert(make_var(10), "target_fn".to_string());
683        let mut decl_names = HashSet::new();
684        decl_names.insert("target_fn".to_string());
685        let sites = find_specialization_sites(&body, &known_consts, &known_fns, &decl_names);
686        assert_eq!(sites.len(), 1);
687        assert_eq!(sites[0].callee, "target_fn");
688        assert!(matches!(sites[0].const_args[0], SpecConstArg::Nat(42)));
689    }
690    #[test]
691    pub(super) fn test_create_specialization() {
692        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
693        let decl = make_decl(
694            "my_fn",
695            vec![
696                make_param(0, "x", LcnfType::Nat),
697                make_param(1, "y", LcnfType::Nat),
698            ],
699            body,
700        );
701        let key = SpecializationKey {
702            original: "my_fn".to_string(),
703            type_args: vec![],
704            const_args: vec![SpecConstArg::Nat(10), SpecConstArg::Unknown],
705            closure_args: vec![],
706        };
707        let mut pass = SpecializationPass::new(SpecializationConfig::default());
708        let result = pass.create_specialization(&decl, &key);
709        assert!(result.is_some());
710        let spec = result.expect("spec should be Some/Ok");
711        assert!(spec.decl.name.contains("my_fn"));
712        assert!(spec.decl.name.contains("_C0_N10"));
713        assert_eq!(spec.decl.params.len(), 1);
714    }
715    #[test]
716    pub(super) fn test_stats_default() {
717        let stats = SpecializationStats::default();
718        assert_eq!(stats.type_specializations, 0);
719        assert_eq!(stats.const_specializations, 0);
720        assert_eq!(stats.closure_specializations, 0);
721    }
722    #[test]
723    pub(super) fn test_pass_fresh_id() {
724        let mut pass = SpecializationPass::new(SpecializationConfig::default());
725        let id1 = pass.fresh_id();
726        let id2 = pass.fresh_id();
727        assert_ne!(id1, id2);
728    }
729    #[test]
730    pub(super) fn test_substitute_in_tailcall() {
731        let mut expr = LcnfExpr::TailCall(
732            LcnfArg::Var(make_var(10)),
733            vec![LcnfArg::Var(make_var(0)), LcnfArg::Var(make_var(1))],
734        );
735        let pass = SpecializationPass::new(SpecializationConfig::default());
736        pass.substitute_constant(&mut expr, make_var(0), &LcnfLit::Nat(7));
737        if let LcnfExpr::TailCall(_, args) = &expr {
738            assert_eq!(args[0], LcnfArg::Lit(LcnfLit::Nat(7)));
739            assert_eq!(args[1], LcnfArg::Var(make_var(1)));
740        } else {
741            panic!("Expected TailCall");
742        }
743    }
744    #[test]
745    pub(super) fn test_is_called_in_case() {
746        let body = LcnfExpr::Case {
747            scrutinee: make_var(1),
748            scrutinee_ty: LcnfType::Nat,
749            alts: vec![LcnfAlt {
750                ctor_name: "True".to_string(),
751                ctor_tag: 0,
752                params: vec![],
753                body: make_simple_let(
754                    5,
755                    LcnfLetValue::App(
756                        LcnfArg::Var(make_var(0)),
757                        vec![LcnfArg::Lit(LcnfLit::Nat(1))],
758                    ),
759                    LcnfExpr::Return(LcnfArg::Var(make_var(5))),
760                ),
761            }],
762            default: None,
763        };
764        assert!(is_called_as_function(&body, make_var(0)));
765        assert!(!is_called_as_function(&body, make_var(2)));
766    }
767    #[test]
768    pub(super) fn test_tailcall_specialization_site() {
769        let expr = LcnfExpr::TailCall(
770            LcnfArg::Var(make_var(10)),
771            vec![LcnfArg::Lit(LcnfLit::Nat(5))],
772        );
773        let mut known_fns: HashMap<LcnfVarId, String> = HashMap::new();
774        known_fns.insert(make_var(10), "recurse".to_string());
775        let mut decl_names = HashSet::new();
776        decl_names.insert("recurse".to_string());
777        let known_consts: HashMap<LcnfVarId, LcnfLit> = HashMap::new();
778        let sites = find_specialization_sites(&expr, &known_consts, &known_fns, &decl_names);
779        assert_eq!(sites.len(), 1);
780        assert_eq!(sites[0].callee, "recurse");
781    }
782    #[test]
783    pub(super) fn test_recursive_specialization_disabled() {
784        let body = LcnfExpr::Return(LcnfArg::Var(make_var(0)));
785        let mut decl = make_decl("rec_fn", vec![make_param(0, "n", LcnfType::Nat)], body);
786        decl.is_recursive = true;
787        let key = SpecializationKey {
788            original: "rec_fn".to_string(),
789            type_args: vec![],
790            const_args: vec![SpecConstArg::Nat(5)],
791            closure_args: vec![],
792        };
793        let mut pass = SpecializationPass::new(SpecializationConfig {
794            allow_recursive: false,
795            ..SpecializationConfig::default()
796        });
797        let result = pass.create_specialization(&decl, &key);
798        assert!(result.is_none());
799    }
800}
801#[cfg(test)]
802mod Spec_infra_tests {
803    use super::*;
804    #[test]
805    pub(super) fn test_pass_config() {
806        let config = SpecPassConfig::new("test_pass", SpecPassPhase::Transformation);
807        assert!(config.enabled);
808        assert!(config.phase.is_modifying());
809        assert_eq!(config.phase.name(), "transformation");
810    }
811    #[test]
812    pub(super) fn test_pass_stats() {
813        let mut stats = SpecPassStats::new();
814        stats.record_run(10, 100, 3);
815        stats.record_run(20, 200, 5);
816        assert_eq!(stats.total_runs, 2);
817        assert!((stats.average_changes_per_run() - 15.0).abs() < 0.01);
818        assert!((stats.success_rate() - 1.0).abs() < 0.01);
819        let s = stats.format_summary();
820        assert!(s.contains("Runs: 2/2"));
821    }
822    #[test]
823    pub(super) fn test_pass_registry() {
824        let mut reg = SpecPassRegistry::new();
825        reg.register(SpecPassConfig::new("pass_a", SpecPassPhase::Analysis));
826        reg.register(SpecPassConfig::new("pass_b", SpecPassPhase::Transformation).disabled());
827        assert_eq!(reg.total_passes(), 2);
828        assert_eq!(reg.enabled_count(), 1);
829        reg.update_stats("pass_a", 5, 50, 2);
830        let stats = reg.get_stats("pass_a").expect("stats should exist");
831        assert_eq!(stats.total_changes, 5);
832    }
833    #[test]
834    pub(super) fn test_analysis_cache() {
835        let mut cache = SpecAnalysisCache::new(10);
836        cache.insert("key1".to_string(), vec![1, 2, 3]);
837        assert!(cache.get("key1").is_some());
838        assert!(cache.get("key2").is_none());
839        assert!((cache.hit_rate() - 0.5).abs() < 0.01);
840        cache.invalidate("key1");
841        assert!(!cache.entries["key1"].valid);
842        assert_eq!(cache.size(), 1);
843    }
844    #[test]
845    pub(super) fn test_worklist() {
846        let mut wl = SpecWorklist::new();
847        assert!(wl.push(1));
848        assert!(wl.push(2));
849        assert!(!wl.push(1));
850        assert_eq!(wl.len(), 2);
851        assert_eq!(wl.pop(), Some(1));
852        assert!(!wl.contains(1));
853        assert!(wl.contains(2));
854    }
855    #[test]
856    pub(super) fn test_dominator_tree() {
857        let mut dt = SpecDominatorTree::new(5);
858        dt.set_idom(1, 0);
859        dt.set_idom(2, 0);
860        dt.set_idom(3, 1);
861        assert!(dt.dominates(0, 3));
862        assert!(dt.dominates(1, 3));
863        assert!(!dt.dominates(2, 3));
864        assert!(dt.dominates(3, 3));
865    }
866    #[test]
867    pub(super) fn test_liveness() {
868        let mut liveness = SpecLivenessInfo::new(3);
869        liveness.add_def(0, 1);
870        liveness.add_use(1, 1);
871        assert!(liveness.defs[0].contains(&1));
872        assert!(liveness.uses[1].contains(&1));
873    }
874    #[test]
875    pub(super) fn test_constant_folding() {
876        assert_eq!(SpecConstantFoldingHelper::fold_add_i64(3, 4), Some(7));
877        assert_eq!(SpecConstantFoldingHelper::fold_div_i64(10, 0), None);
878        assert_eq!(SpecConstantFoldingHelper::fold_div_i64(10, 2), Some(5));
879        assert_eq!(
880            SpecConstantFoldingHelper::fold_bitand_i64(0b1100, 0b1010),
881            0b1000
882        );
883        assert_eq!(SpecConstantFoldingHelper::fold_bitnot_i64(0), -1);
884    }
885    #[test]
886    pub(super) fn test_dep_graph() {
887        let mut g = SpecDepGraph::new();
888        g.add_dep(1, 2);
889        g.add_dep(2, 3);
890        g.add_dep(1, 3);
891        assert_eq!(g.dependencies_of(2), vec![1]);
892        let topo = g.topological_sort();
893        assert_eq!(topo.len(), 3);
894        assert!(!g.has_cycle());
895        let pos: std::collections::HashMap<u32, usize> =
896            topo.iter().enumerate().map(|(i, &n)| (n, i)).collect();
897        assert!(pos[&1] < pos[&2]);
898        assert!(pos[&1] < pos[&3]);
899        assert!(pos[&2] < pos[&3]);
900    }
901}
902#[cfg(test)]
903mod specext_pass_tests {
904    use super::*;
905    #[test]
906    pub(super) fn test_specext_phase_order() {
907        assert_eq!(SpecExtPassPhase::Early.order(), 0);
908        assert_eq!(SpecExtPassPhase::Middle.order(), 1);
909        assert_eq!(SpecExtPassPhase::Late.order(), 2);
910        assert_eq!(SpecExtPassPhase::Finalize.order(), 3);
911        assert!(SpecExtPassPhase::Early.is_early());
912        assert!(!SpecExtPassPhase::Early.is_late());
913    }
914    #[test]
915    pub(super) fn test_specext_config_builder() {
916        let c = SpecExtPassConfig::new("p")
917            .with_phase(SpecExtPassPhase::Late)
918            .with_max_iter(50)
919            .with_debug(1);
920        assert_eq!(c.name, "p");
921        assert_eq!(c.max_iterations, 50);
922        assert!(c.is_debug_enabled());
923        assert!(c.enabled);
924        let c2 = c.disabled();
925        assert!(!c2.enabled);
926    }
927    #[test]
928    pub(super) fn test_specext_stats() {
929        let mut s = SpecExtPassStats::new();
930        s.visit();
931        s.visit();
932        s.modify();
933        s.iterate();
934        assert_eq!(s.nodes_visited, 2);
935        assert_eq!(s.nodes_modified, 1);
936        assert!(s.changed);
937        assert_eq!(s.iterations, 1);
938        let e = s.efficiency();
939        assert!((e - 0.5).abs() < 1e-9);
940    }
941    #[test]
942    pub(super) fn test_specext_registry() {
943        let mut r = SpecExtPassRegistry::new();
944        r.register(SpecExtPassConfig::new("a").with_phase(SpecExtPassPhase::Early));
945        r.register(SpecExtPassConfig::new("b").disabled());
946        assert_eq!(r.len(), 2);
947        assert_eq!(r.enabled_passes().len(), 1);
948        assert_eq!(r.passes_in_phase(&SpecExtPassPhase::Early).len(), 1);
949    }
950    #[test]
951    pub(super) fn test_specext_cache() {
952        let mut c = SpecExtCache::new(4);
953        assert!(c.get(99).is_none());
954        c.put(99, vec![1, 2, 3]);
955        let v = c.get(99).expect("v should be present in map");
956        assert_eq!(v, &[1u8, 2, 3]);
957        assert!(c.hit_rate() > 0.0);
958        assert_eq!(c.live_count(), 1);
959    }
960    #[test]
961    pub(super) fn test_specext_worklist() {
962        let mut w = SpecExtWorklist::new(10);
963        w.push(5);
964        w.push(3);
965        w.push(5);
966        assert_eq!(w.len(), 2);
967        assert!(w.contains(5));
968        let first = w.pop().expect("first should be available to pop");
969        assert!(!w.contains(first));
970    }
971    #[test]
972    pub(super) fn test_specext_dom_tree() {
973        let mut dt = SpecExtDomTree::new(5);
974        dt.set_idom(1, 0);
975        dt.set_idom(2, 0);
976        dt.set_idom(3, 1);
977        dt.set_idom(4, 1);
978        assert!(dt.dominates(0, 3));
979        assert!(dt.dominates(1, 4));
980        assert!(!dt.dominates(2, 3));
981        assert_eq!(dt.depth_of(3), 2);
982    }
983    #[test]
984    pub(super) fn test_specext_liveness() {
985        let mut lv = SpecExtLiveness::new(3);
986        lv.add_def(0, 1);
987        lv.add_use(1, 1);
988        assert!(lv.var_is_def_in_block(0, 1));
989        assert!(lv.var_is_used_in_block(1, 1));
990        assert!(!lv.var_is_def_in_block(1, 1));
991    }
992    #[test]
993    pub(super) fn test_specext_const_folder() {
994        let mut cf = SpecExtConstFolder::new();
995        assert_eq!(cf.add_i64(3, 4), Some(7));
996        assert_eq!(cf.div_i64(10, 0), None);
997        assert_eq!(cf.mul_i64(6, 7), Some(42));
998        assert_eq!(cf.and_i64(0b1100, 0b1010), 0b1000);
999        assert_eq!(cf.fold_count(), 3);
1000        assert_eq!(cf.failure_count(), 1);
1001    }
1002    #[test]
1003    pub(super) fn test_specext_dep_graph() {
1004        let mut g = SpecExtDepGraph::new(4);
1005        g.add_edge(0, 1);
1006        g.add_edge(1, 2);
1007        g.add_edge(2, 3);
1008        assert!(!g.has_cycle());
1009        assert_eq!(g.topo_sort(), Some(vec![0, 1, 2, 3]));
1010        assert_eq!(g.reachable(0).len(), 4);
1011        let sccs = g.scc();
1012        assert_eq!(sccs.len(), 4);
1013    }
1014}
oxilean_codegen/opt_specialize/functions.rs

oxilean_codegen/opt_specialize/
functions.rs
