logicaffeine_compile/optimize/
supercompile.rs

1//! Supercompiler — a unified optimization pass that subsumes constant folding,
2//! propagation, dead code elimination, and compile-time function evaluation
3//! in a single framework.
4//!
5//! The supercompiler performs online partial evaluation via:
6//!   1. **Driving** — symbolic execution one step at a time, maintaining an
7//!      abstract store mapping variables to known values
8//!   2. **Folding** — memoizing function evaluation results to avoid
9//!      recomputation and detect infinite recursion
10//!   3. **Generalization** — widening variables modified in loops to ensure
11//!      termination of the analysis
12//!
13//! Scope: pure integer/boolean code only. Collections, IO, and escape blocks
14//! are treated conservatively (passed through unchanged).
15
16use std::collections::HashMap;
17
18use crate::arena::Arena;
19use crate::ast::stmt::{BinaryOpKind, Block, Expr, Literal, Stmt};
20use crate::intern::{Interner, Symbol};
21
22const MAX_INLINE_DEPTH: usize = 64;
23const MAX_INLINE_STEPS: usize = 10_000;
24
25#[derive(Debug, Clone)]
26enum Value {
27    Int(i64),
28    Float(f64),
29    Bool(bool),
30    Text(Symbol),
31    Nothing,
32}
33
34impl PartialEq for Value {
35    fn eq(&self, other: &Self) -> bool {
36        match (self, other) {
37            (Value::Int(a), Value::Int(b)) => a == b,
38            (Value::Float(a), Value::Float(b)) => a.to_bits() == b.to_bits(),
39            (Value::Bool(a), Value::Bool(b)) => a == b,
40            (Value::Text(a), Value::Text(b)) => a == b,
41            (Value::Nothing, Value::Nothing) => true,
42            _ => false,
43        }
44    }
45}
46
47impl Eq for Value {}
48
49impl std::hash::Hash for Value {
50    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
51        std::mem::discriminant(self).hash(state);
52        match self {
53            Value::Int(n) => n.hash(state),
54            Value::Float(f) => f.to_bits().hash(state),
55            Value::Bool(b) => b.hash(state),
56            Value::Text(s) => s.hash(state),
57            Value::Nothing => {}
58        }
59    }
60}
61
62struct FuncDef<'a> {
63    params: Vec<Symbol>,
64    body: Block<'a>,
65}
66
67#[derive(Clone)]
68struct Configuration<'a> {
69    expr: &'a Expr<'a>,
70    store_snapshot: HashMap<Symbol, Value>,
71}
72
73struct History<'a> {
74    entries: Vec<Configuration<'a>>,
75}
76
77impl<'a> History<'a> {
78    fn new() -> Self {
79        Self { entries: Vec::new() }
80    }
81
82    fn push(&mut self, config: Configuration<'a>) {
83        if self.entries.len() >= 16 {
84            self.entries.remove(0);
85        }
86        self.entries.push(config);
87    }
88
89    fn check_embedding(&self, new_expr: &Expr<'a>) -> Option<&Configuration<'a>> {
90        self.entries.iter().find(|c| embeds(c.expr, new_expr))
91    }
92}
93
94struct SuperEnv<'a> {
95    store: HashMap<Symbol, Value>,
96    funcs: HashMap<Symbol, FuncDef<'a>>,
97    memo: HashMap<(Symbol, Vec<Value>), Option<Value>>,
98    steps: usize,
99    history: History<'a>,
100}
101
102impl<'a> SuperEnv<'a> {
103    fn new() -> Self {
104        Self {
105            store: HashMap::new(),
106            funcs: HashMap::new(),
107            memo: HashMap::new(),
108            steps: 0,
109            history: History::new(),
110        }
111    }
112}
113
114fn is_pure_body(stmts: &[Stmt]) -> bool {
115    stmts.iter().all(is_pure_stmt)
116}
117
118fn is_pure_stmt(stmt: &Stmt) -> bool {
119    match stmt {
120        Stmt::Let { value, .. } => is_pure_expr(value),
121        Stmt::Set { value, .. } => is_pure_expr(value),
122        Stmt::Return { value } => value.map_or(true, is_pure_expr),
123        Stmt::If { cond, then_block, else_block } => {
124            is_pure_expr(cond)
125                && is_pure_body(then_block)
126                && else_block.map_or(true, |eb| is_pure_body(eb))
127        }
128        Stmt::While { cond, body, .. } => is_pure_expr(cond) && is_pure_body(body),
129        Stmt::Call { .. } => true,
130        _ => false,
131    }
132}
133
134fn is_pure_expr(expr: &Expr) -> bool {
135    match expr {
136        Expr::Literal(_) | Expr::Identifier(_) | Expr::OptionNone => true,
137        Expr::BinaryOp { left, right, .. } => is_pure_expr(left) && is_pure_expr(right),
138        Expr::Not { operand } => is_pure_expr(operand),
139        Expr::Call { args, .. } => args.iter().all(|a| is_pure_expr(a)),
140        Expr::Length { collection } => is_pure_expr(collection),
141        Expr::FieldAccess { object, .. } => is_pure_expr(object),
142        Expr::Index { collection, index } => is_pure_expr(collection) && is_pure_expr(index),
143        _ => false,
144    }
145}
146
147fn collect_pure_funcs<'a>(stmts: &[Stmt<'a>]) -> HashMap<Symbol, FuncDef<'a>> {
148    let mut funcs = HashMap::new();
149    for stmt in stmts {
150        if let Stmt::FunctionDef { name, params, body, is_native, .. } = stmt {
151            if *is_native {
152                continue;
153            }
154            if is_pure_body(body) {
155                let param_syms: Vec<Symbol> = params.iter().map(|(s, _)| *s).collect();
156                funcs.insert(*name, FuncDef { params: param_syms, body });
157            }
158        }
159    }
160    funcs
161}
162
163// =============================================================================
164// Value-level evaluation (for inline evaluation of pure functions)
165// =============================================================================
166
167fn eval_expr_to_value(
168    expr: &Expr,
169    locals: &HashMap<Symbol, Value>,
170    funcs: &HashMap<Symbol, FuncDef>,
171    memo: &mut HashMap<(Symbol, Vec<Value>), Option<Value>>,
172    steps: &mut usize,
173    depth: usize,
174) -> Option<Value> {
175    if *steps >= MAX_INLINE_STEPS || depth >= MAX_INLINE_DEPTH {
176        return None;
177    }
178    *steps += 1;
179
180    match expr {
181        Expr::Literal(Literal::Number(n)) => Some(Value::Int(*n)),
182        Expr::Literal(Literal::Float(f)) => Some(Value::Float(*f)),
183        Expr::Literal(Literal::Boolean(b)) => Some(Value::Bool(*b)),
184        Expr::Literal(Literal::Text(s)) => Some(Value::Text(*s)),
185        Expr::Literal(Literal::Nothing) => Some(Value::Nothing),
186        Expr::Identifier(sym) => locals.get(sym).cloned(),
187        Expr::BinaryOp { op, left, right } => {
188            let lv = eval_expr_to_value(left, locals, funcs, memo, steps, depth)?;
189            let rv = eval_expr_to_value(right, locals, funcs, memo, steps, depth)?;
190            eval_binop(*op, &lv, &rv)
191        }
192        Expr::Not { operand } => {
193            if let Value::Bool(b) = eval_expr_to_value(operand, locals, funcs, memo, steps, depth)? {
194                Some(Value::Bool(!b))
195            } else {
196                None
197            }
198        }
199        Expr::Call { function, args } => {
200            let func = funcs.get(function)?;
201            if args.len() != func.params.len() {
202                return None;
203            }
204            let mut arg_vals = Vec::with_capacity(args.len());
205            for arg in args {
206                arg_vals.push(eval_expr_to_value(arg, locals, funcs, memo, steps, depth)?);
207            }
208
209            let memo_key = (*function, arg_vals.clone());
210            if let Some(cached) = memo.get(&memo_key) {
211                return cached.clone();
212            }
213
214            // Mark as in-progress to detect infinite recursion
215            memo.insert(memo_key.clone(), None);
216
217            let mut call_locals = HashMap::new();
218            for (param, val) in func.params.iter().zip(arg_vals.iter()) {
219                call_locals.insert(*param, val.clone());
220            }
221
222            let result = eval_block_to_value(func.body, &mut call_locals, funcs, memo, steps, depth + 1);
223            memo.insert(memo_key, result.clone());
224            result
225        }
226        _ => None,
227    }
228}
229
230fn eval_binop(op: BinaryOpKind, lv: &Value, rv: &Value) -> Option<Value> {
231    match (op, lv, rv) {
232        // Constant-fold integer arithmetic ONLY when the result fits i64; on overflow
233        // return None so the now-exact runtime computes the promoted (BigInt) value.
234        // The folder must never bake a WRAPPED constant the runtime would not produce.
235        (BinaryOpKind::Add, Value::Int(a), Value::Int(b)) => a.checked_add(*b).map(Value::Int),
236        (BinaryOpKind::Subtract, Value::Int(a), Value::Int(b)) => a.checked_sub(*b).map(Value::Int),
237        (BinaryOpKind::Multiply, Value::Int(a), Value::Int(b)) => a.checked_mul(*b).map(Value::Int),
238        (BinaryOpKind::Divide, Value::Int(a), Value::Int(b)) if *b != 0 => a.checked_div(*b).map(Value::Int),
239        (BinaryOpKind::Modulo, Value::Int(a), Value::Int(b)) if *b != 0 => a.checked_rem(*b).map(Value::Int),
240        (BinaryOpKind::Shl, Value::Int(a), Value::Int(b)) if *b >= 0 && *b < 64 => {
241            Some(Value::Int(a.wrapping_shl(*b as u32)))
242        }
243        (BinaryOpKind::Shr, Value::Int(a), Value::Int(b)) if *b >= 0 && *b < 64 => {
244            Some(Value::Int(a.wrapping_shr(*b as u32)))
245        }
246        (BinaryOpKind::BitXor, Value::Int(a), Value::Int(b)) => Some(Value::Int(a ^ b)),
247        (BinaryOpKind::And, Value::Int(a), Value::Int(b)) => Some(Value::Int(a & b)),
248        (BinaryOpKind::Or, Value::Int(a), Value::Int(b)) => Some(Value::Int(a | b)),
249        (BinaryOpKind::Eq, Value::Int(a), Value::Int(b)) => Some(Value::Bool(a == b)),
250        (BinaryOpKind::NotEq, Value::Int(a), Value::Int(b)) => Some(Value::Bool(a != b)),
251        (BinaryOpKind::Lt, Value::Int(a), Value::Int(b)) => Some(Value::Bool(a < b)),
252        (BinaryOpKind::Gt, Value::Int(a), Value::Int(b)) => Some(Value::Bool(a > b)),
253        (BinaryOpKind::LtEq, Value::Int(a), Value::Int(b)) => Some(Value::Bool(a <= b)),
254        (BinaryOpKind::GtEq, Value::Int(a), Value::Int(b)) => Some(Value::Bool(a >= b)),
255        (BinaryOpKind::Add, Value::Float(a), Value::Float(b)) => Some(Value::Float(a + b)),
256        (BinaryOpKind::Subtract, Value::Float(a), Value::Float(b)) => Some(Value::Float(a - b)),
257        (BinaryOpKind::Multiply, Value::Float(a), Value::Float(b)) => Some(Value::Float(a * b)),
258        (BinaryOpKind::Divide, Value::Float(a), Value::Float(b)) if *b != 0.0 => Some(Value::Float(a / b)),
259        (BinaryOpKind::And, Value::Bool(a), Value::Bool(b)) => Some(Value::Bool(*a && *b)),
260        (BinaryOpKind::Or, Value::Bool(a), Value::Bool(b)) => Some(Value::Bool(*a || *b)),
261        (BinaryOpKind::Eq, Value::Text(a), Value::Text(b)) => Some(Value::Bool(a == b)),
262        (BinaryOpKind::NotEq, Value::Text(a), Value::Text(b)) => Some(Value::Bool(a != b)),
263        _ => None,
264    }
265}
266
267enum EvalResult {
268    Continue,
269    Return(Value),
270}
271
272fn eval_block_to_value(
273    stmts: &[Stmt],
274    locals: &mut HashMap<Symbol, Value>,
275    funcs: &HashMap<Symbol, FuncDef>,
276    memo: &mut HashMap<(Symbol, Vec<Value>), Option<Value>>,
277    steps: &mut usize,
278    depth: usize,
279) -> Option<Value> {
280    for stmt in stmts {
281        match eval_stmt_to_value(stmt, locals, funcs, memo, steps, depth)? {
282            EvalResult::Continue => {}
283            EvalResult::Return(v) => return Some(v),
284        }
285    }
286    Some(Value::Nothing)
287}
288
289fn eval_stmt_to_value(
290    stmt: &Stmt,
291    locals: &mut HashMap<Symbol, Value>,
292    funcs: &HashMap<Symbol, FuncDef>,
293    memo: &mut HashMap<(Symbol, Vec<Value>), Option<Value>>,
294    steps: &mut usize,
295    depth: usize,
296) -> Option<EvalResult> {
297    if *steps >= MAX_INLINE_STEPS {
298        return None;
299    }
300    *steps += 1;
301
302    match stmt {
303        Stmt::Let { var, value, .. } => {
304            let v = eval_expr_to_value(value, locals, funcs, memo, steps, depth)?;
305            locals.insert(*var, v);
306            Some(EvalResult::Continue)
307        }
308        Stmt::Set { target, value } => {
309            let v = eval_expr_to_value(value, locals, funcs, memo, steps, depth)?;
310            locals.insert(*target, v);
311            Some(EvalResult::Continue)
312        }
313        Stmt::Return { value } => {
314            if let Some(v) = value {
315                let result = eval_expr_to_value(v, locals, funcs, memo, steps, depth)?;
316                Some(EvalResult::Return(result))
317            } else {
318                Some(EvalResult::Return(Value::Nothing))
319            }
320        }
321        Stmt::If { cond, then_block, else_block } => {
322            let cv = eval_expr_to_value(cond, locals, funcs, memo, steps, depth)?;
323            if let Value::Bool(b) = cv {
324                if b {
325                    for s in *then_block {
326                        match eval_stmt_to_value(s, locals, funcs, memo, steps, depth)? {
327                            EvalResult::Continue => {}
328                            EvalResult::Return(v) => return Some(EvalResult::Return(v)),
329                        }
330                    }
331                } else if let Some(eb) = else_block {
332                    for s in *eb {
333                        match eval_stmt_to_value(s, locals, funcs, memo, steps, depth)? {
334                            EvalResult::Continue => {}
335                            EvalResult::Return(v) => return Some(EvalResult::Return(v)),
336                        }
337                    }
338                }
339                Some(EvalResult::Continue)
340            } else {
341                None
342            }
343        }
344        Stmt::While { cond, body, .. } => {
345            loop {
346                let cv = eval_expr_to_value(cond, locals, funcs, memo, steps, depth)?;
347                match cv {
348                    Value::Bool(false) => break,
349                    Value::Bool(true) => {
350                        for s in *body {
351                            match eval_stmt_to_value(s, locals, funcs, memo, steps, depth)? {
352                                EvalResult::Continue => {}
353                                EvalResult::Return(v) => return Some(EvalResult::Return(v)),
354                            }
355                        }
356                    }
357                    _ => return None,
358                }
359            }
360            Some(EvalResult::Continue)
361        }
362        _ => None,
363    }
364}
365
366// =============================================================================
367// AST-level supercompilation (driving + constant propagation + inline)
368// =============================================================================
369
370fn value_to_literal(val: &Value) -> Option<Literal> {
371    match val {
372        Value::Int(n) => Some(Literal::Number(*n)),
373        Value::Float(f) => Some(Literal::Float(*f)),
374        Value::Bool(b) => Some(Literal::Boolean(*b)),
375        Value::Text(s) => Some(Literal::Text(*s)),
376        Value::Nothing => Some(Literal::Nothing),
377    }
378}
379
380fn value_to_expr<'a>(val: &Value, arena: &'a Arena<Expr<'a>>) -> &'a Expr<'a> {
381    match value_to_literal(val) {
382        Some(lit) => arena.alloc(Expr::Literal(lit)),
383        None => unreachable!(),
384    }
385}
386
387fn expr_to_value(expr: &Expr, store: &HashMap<Symbol, Value>) -> Option<Value> {
388    match expr {
389        Expr::Literal(Literal::Number(n)) => Some(Value::Int(*n)),
390        Expr::Literal(Literal::Float(f)) => Some(Value::Float(*f)),
391        Expr::Literal(Literal::Boolean(b)) => Some(Value::Bool(*b)),
392        Expr::Literal(Literal::Text(s)) => Some(Value::Text(*s)),
393        Expr::Literal(Literal::Nothing) => Some(Value::Nothing),
394        Expr::Identifier(sym) => store.get(sym).cloned(),
395        _ => None,
396    }
397}
398
399fn try_eval_expr<'a>(
400    expr: &'a Expr<'a>,
401    env: &mut SuperEnv<'a>,
402    depth: usize,
403) -> Option<Value> {
404    eval_expr_to_value(
405        expr,
406        &env.store,
407        &env.funcs,
408        &mut env.memo,
409        &mut env.steps,
410        depth,
411    )
412}
413
414fn drive_expr<'a>(
415    expr: &'a Expr<'a>,
416    env: &mut SuperEnv<'a>,
417    expr_arena: &'a Arena<Expr<'a>>,
418    depth: usize,
419) -> &'a Expr<'a> {
420    if depth >= MAX_INLINE_DEPTH {
421        return expr;
422    }
423
424    match expr {
425        Expr::Identifier(sym) => {
426            if let Some(val) = env.store.get(sym) {
427                // Only propagate Copy-type values (Int, Float, Bool, Nothing).
428                // Text values are owned (String) in generated Rust — propagating
429                // them changes move semantics and can eliminate use-after-move errors.
430                match val {
431                    Value::Int(_) | Value::Float(_) | Value::Bool(_) | Value::Nothing => {
432                        if let Some(lit) = value_to_literal(val) {
433                            return expr_arena.alloc(Expr::Literal(lit));
434                        }
435                    }
436                    Value::Text(s) => {
437                        return expr_arena.alloc(Expr::Literal(Literal::Text(*s)));
438                    }
439                }
440            }
441            expr
442        }
443        Expr::BinaryOp { op, left, right } => {
444            let dl = drive_expr(left, env, expr_arena, depth);
445            let dr = drive_expr(right, env, expr_arena, depth);
446            // Try value-level evaluation
447            if let (Some(lv), Some(rv)) = (expr_to_value(dl, &env.store), expr_to_value(dr, &env.store)) {
448                if let Some(result) = eval_binop(*op, &lv, &rv) {
449                    return value_to_expr(&result, expr_arena);
450                }
451            }
452            if std::ptr::eq(dl, *left) && std::ptr::eq(dr, *right) {
453                expr
454            } else {
455                expr_arena.alloc(Expr::BinaryOp { op: *op, left: dl, right: dr })
456            }
457        }
458        Expr::Not { operand } => {
459            let d = drive_expr(operand, env, expr_arena, depth);
460            if let Expr::Literal(Literal::Boolean(b)) = d {
461                return expr_arena.alloc(Expr::Literal(Literal::Boolean(!b)));
462            }
463            if std::ptr::eq(d, *operand) { expr } else { expr_arena.alloc(Expr::Not { operand: d }) }
464        }
465        Expr::Call { function, args } => {
466            let driven_args: Vec<&'a Expr<'a>> = args.iter()
467                .map(|a| drive_expr(a, env, expr_arena, depth))
468                .collect();
469            // Don't replace calls in the AST — preserve function call sites.
470            // Call evaluation is done at the store level in drive_stmt for Let/Set.
471            let changed = driven_args.iter().zip(args.iter()).any(|(d, o)| !std::ptr::eq(*d, *o));
472            if changed {
473                expr_arena.alloc(Expr::Call { function: *function, args: driven_args })
474            } else {
475                expr
476            }
477        }
478        Expr::Length { collection } => {
479            let d = drive_expr(collection, env, expr_arena, depth);
480            if std::ptr::eq(d, *collection) { expr } else { expr_arena.alloc(Expr::Length { collection: d }) }
481        }
482        // Index/Slice: preserve original variable names for codegen peephole patterns.
483        // The swap, vec-fill, and index-lowering patterns depend on seeing variable
484        // references (e.g., `j`, `j+1`) rather than propagated constants.
485        // The fold pass handles compile-time index evaluation separately.
486        Expr::Index { .. } | Expr::Slice { .. } => expr,
487        Expr::Contains { collection, value } => {
488            let dc = drive_expr(collection, env, expr_arena, depth);
489            let dv = drive_expr(value, env, expr_arena, depth);
490            if std::ptr::eq(dc, *collection) && std::ptr::eq(dv, *value) {
491                expr
492            } else {
493                expr_arena.alloc(Expr::Contains { collection: dc, value: dv })
494            }
495        }
496        Expr::FieldAccess { object, field } => {
497            let d = drive_expr(object, env, expr_arena, depth);
498            if std::ptr::eq(d, *object) { expr } else { expr_arena.alloc(Expr::FieldAccess { object: d, field: *field }) }
499        }
500        Expr::OptionSome { value } => {
501            let d = drive_expr(value, env, expr_arena, depth);
502            if std::ptr::eq(d, *value) { expr } else { expr_arena.alloc(Expr::OptionSome { value: d }) }
503        }
504        Expr::Copy { expr: inner } => {
505            let d = drive_expr(inner, env, expr_arena, depth);
506            if std::ptr::eq(d, *inner) { expr } else { expr_arena.alloc(Expr::Copy { expr: d }) }
507        }
508        Expr::Give { value } => {
509            let d = drive_expr(value, env, expr_arena, depth);
510            if std::ptr::eq(d, *value) { expr } else { expr_arena.alloc(Expr::Give { value: d }) }
511        }
512        Expr::List(elems) => {
513            let driven: Vec<&'a Expr<'a>> = elems.iter()
514                .map(|e| drive_expr(e, env, expr_arena, depth))
515                .collect();
516            let changed = driven.iter().zip(elems.iter()).any(|(d, o)| !std::ptr::eq(*d, *o));
517            if changed { expr_arena.alloc(Expr::List(driven)) } else { expr }
518        }
519        Expr::Tuple(elems) => {
520            let driven: Vec<&'a Expr<'a>> = elems.iter()
521                .map(|e| drive_expr(e, env, expr_arena, depth))
522                .collect();
523            let changed = driven.iter().zip(elems.iter()).any(|(d, o)| !std::ptr::eq(*d, *o));
524            if changed { expr_arena.alloc(Expr::Tuple(driven)) } else { expr }
525        }
526        // Leaves and complex expressions passed through
527        _ => expr,
528    }
529}
530
531fn drive_stmt<'a>(
532    stmt: Stmt<'a>,
533    env: &mut SuperEnv<'a>,
534    expr_arena: &'a Arena<Expr<'a>>,
535    stmt_arena: &'a Arena<Stmt<'a>>,
536    interner: &mut Interner,
537    depth: usize,
538) -> Option<Stmt<'a>> {
539    match stmt {
540        Stmt::Let { var, ty, value, mutable } => {
541            let driven = drive_expr(value, env, expr_arena, depth);
542            if !mutable {
543                // Try to resolve the driven expression to a value for the store
544                if let Some(val) = expr_to_value(driven, &env.store) {
545                    env.store.insert(var, val);
546                } else {
547                    // Expression isn't a simple literal/identifier — try evaluating
548                    // calls at the value level for store tracking
549                    if let Some(val) = try_eval_expr(driven, env, depth) {
550                        env.store.insert(var, val);
551                    }
552                }
553            }
554            Some(Stmt::Let { var, ty, value: driven, mutable })
555        }
556        Stmt::Set { target, value } => {
557            let driven = drive_expr(value, env, expr_arena, depth);
558            // Update store if we can resolve the new value
559            if let Some(val) = expr_to_value(driven, &env.store) {
560                env.store.insert(target, val);
561            } else {
562                // Value is unknown; remove from store
563                env.store.remove(&target);
564            }
565            Some(Stmt::Set { target, value: driven })
566        }
567        Stmt::If { cond, then_block, else_block } => {
568            let driven_cond = drive_expr(cond, env, expr_arena, depth);
569            // If condition is known, eliminate dead branch
570            if let Expr::Literal(Literal::Boolean(b)) = driven_cond {
571                if *b {
572                    // Take then branch only
573                    let driven_then = drive_block(then_block, env, expr_arena, stmt_arena, interner, depth);
574                    Some(Stmt::If {
575                        cond: driven_cond,
576                        then_block: driven_then,
577                        else_block: None,
578                    })
579                } else {
580                    // Take else branch only
581                    if let Some(eb) = else_block {
582                        let driven_else = drive_block(eb, env, expr_arena, stmt_arena, interner, depth);
583                        Some(Stmt::If {
584                            cond: driven_cond,
585                            then_block: stmt_arena.alloc_slice(vec![]),
586                            else_block: Some(driven_else),
587                        })
588                    } else {
589                        None // Entire If eliminated
590                    }
591                }
592            } else {
593                // Condition unknown: drive both branches with snapshot/restore
594                let snapshot = env.store.clone();
595                let driven_then = drive_block(then_block, env, expr_arena, stmt_arena, interner, depth);
596                let then_store = env.store.clone();
597
598                env.store = snapshot;
599                let driven_else = else_block.map(|eb| {
600                    drive_block(eb, env, expr_arena, stmt_arena, interner, depth)
601                });
602                let else_store = env.store.clone();
603
604                // Join: keep only values that agree in both branches
605                let mut joined = HashMap::new();
606                for (sym, then_val) in &then_store {
607                    if let Some(else_val) = else_store.get(sym) {
608                        if then_val == else_val {
609                            joined.insert(*sym, then_val.clone());
610                        }
611                    }
612                }
613                env.store = joined;
614
615                Some(Stmt::If {
616                    cond: driven_cond,
617                    then_block: driven_then,
618                    else_block: driven_else,
619                })
620            }
621        }
622        Stmt::While { cond, body, decreasing } => {
623            let modified = collect_modified_vars_block(body);
624            let snapshot = env.store.clone();
625
626            // Snapshot configuration before entering loop
627            let pre_config = Configuration {
628                expr: cond,
629                store_snapshot: snapshot.clone(),
630            };
631
632            // Widen modified variables
633            for sym in &modified {
634                env.store.remove(sym);
635            }
636
637            let driven_cond = drive_expr(cond, env, expr_arena, depth);
638            // While-false elimination: if condition is statically false, remove loop
639            if let Expr::Literal(Literal::Boolean(false)) = driven_cond {
640                env.store = snapshot;
641                return None;
642            }
643
644            // Check embedding: if the driven condition embeds in a historical one,
645            // the whistle blows — use MSG to generalize and widen precisely.
646            if let Some(prev) = env.history.check_embedding(driven_cond) {
647                let generalized = msg(prev.expr, driven_cond, expr_arena, interner);
648                // Remove store entries for MSG-introduced variables to ensure termination
649                for i in 0..generalized.num_substitutions {
650                    let name = format!("__msg_{}", i);
651                    let sym = interner.intern(&name);
652                    env.store.remove(&sym);
653                }
654            }
655
656            // Push configuration to history
657            env.history.push(pre_config);
658
659            let driven_body = drive_block(body, env, expr_arena, stmt_arena, interner, depth);
660
661            // After loop, preserve unmodified variables from snapshot
662            let mut post_loop_store = HashMap::new();
663            for (sym, val) in &snapshot {
664                if !modified.contains(sym) {
665                    post_loop_store.insert(*sym, val.clone());
666                }
667            }
668            env.store = post_loop_store;
669
670            Some(Stmt::While { cond: driven_cond, body: driven_body, decreasing })
671        }
672        Stmt::Repeat { pattern, iterable, body } => {
673            let driven_iter = drive_expr(iterable, env, expr_arena, depth);
674            // Loop variable and modified vars are unknown
675            if let crate::ast::stmt::Pattern::Identifier(sym) = &pattern {
676                env.store.remove(sym);
677            }
678            let modified = collect_modified_vars_block(body);
679            for sym in &modified {
680                env.store.remove(sym);
681            }
682            let driven_body = drive_block(body, env, expr_arena, stmt_arena, interner, depth);
683            Some(Stmt::Repeat { pattern, iterable: driven_iter, body: driven_body })
684        }
685        Stmt::Show { object, recipient } => {
686            let driven = drive_expr(object, env, expr_arena, depth);
687            Some(Stmt::Show { object: driven, recipient })
688        }
689        Stmt::Return { value } => {
690            let driven = value.map(|v| drive_expr(v, env, expr_arena, depth));
691            Some(Stmt::Return { value: driven })
692        }
693        Stmt::Call { function, args } => {
694            let driven_args: Vec<&'a Expr<'a>> = args.into_iter()
695                .map(|a| drive_expr(a, env, expr_arena, depth))
696                .collect();
697            Some(Stmt::Call { function, args: driven_args })
698        }
699        Stmt::Push { value, collection } => {
700            let driven_val = drive_expr(value, env, expr_arena, depth);
701            // Push invalidates known value of collection
702            if let Expr::Identifier(sym) = collection {
703                env.store.remove(sym);
704            }
705            Some(Stmt::Push { value: driven_val, collection })
706        }
707        Stmt::SetIndex { collection, index, value } => {
708            // Don't drive index — codegen peephole patterns depend on original variable names
709            if let Expr::Identifier(sym) = collection {
710                env.store.remove(sym);
711            }
712            Some(Stmt::SetIndex { collection, index, value })
713        }
714        Stmt::SetField { object, field, value } => {
715            let dv = drive_expr(value, env, expr_arena, depth);
716            if let Expr::Identifier(sym) = object {
717                env.store.remove(sym);
718            }
719            Some(Stmt::SetField { object, field, value: dv })
720        }
721        Stmt::RuntimeAssert { condition, hard } => {
722            let driven = drive_expr(condition, env, expr_arena, depth);
723            Some(Stmt::RuntimeAssert { condition: driven , hard })
724        }
725        Stmt::FunctionDef { name, params, generics, body, return_type, is_native, native_path, is_exported, export_target, opt_flags } => {
726            // Drive inside function body with fresh store
727            let snapshot = env.store.clone();
728            env.store.clear();
729            let driven_body = drive_block(body, env, expr_arena, stmt_arena, interner, depth);
730            env.store = snapshot;
731            Some(Stmt::FunctionDef {
732                name, params, generics,
733                body: driven_body,
734                return_type, is_native, native_path, is_exported, export_target, opt_flags,
735            })
736        }
737        // Zone blocks: pass through unchanged. Propagating inside zones can
738        // eliminate zone escape patterns that the Rust compiler needs to detect.
739        Stmt::Zone { .. } => Some(stmt),
740        // Pass through everything else unchanged
741        other => Some(other),
742    }
743}
744
745fn drive_block<'a>(
746    block: &'a [Stmt<'a>],
747    env: &mut SuperEnv<'a>,
748    expr_arena: &'a Arena<Expr<'a>>,
749    stmt_arena: &'a Arena<Stmt<'a>>,
750    interner: &mut Interner,
751    depth: usize,
752) -> &'a [Stmt<'a>] {
753    let mut result = Vec::with_capacity(block.len());
754    for stmt in block.iter().cloned() {
755        if let Some(driven) = drive_stmt(stmt, env, expr_arena, stmt_arena, interner, depth) {
756            result.push(driven);
757        }
758    }
759    stmt_arena.alloc_slice(result)
760}
761
762fn collect_modified_vars_block(stmts: &[Stmt]) -> Vec<Symbol> {
763    let mut modified = Vec::new();
764    for stmt in stmts {
765        collect_modified_vars_stmt(stmt, &mut modified);
766    }
767    modified
768}
769
770fn collect_modified_vars_stmt(stmt: &Stmt, out: &mut Vec<Symbol>) {
771    match stmt {
772        Stmt::Set { target, .. } => out.push(*target),
773        Stmt::Let { var, mutable: true, .. } => out.push(*var),
774        Stmt::If { then_block, else_block, .. } => {
775            for s in *then_block {
776                collect_modified_vars_stmt(s, out);
777            }
778            if let Some(eb) = else_block {
779                for s in *eb {
780                    collect_modified_vars_stmt(s, out);
781                }
782            }
783        }
784        Stmt::While { body, .. } => {
785            for s in *body {
786                collect_modified_vars_stmt(s, out);
787            }
788        }
789        Stmt::Repeat { body, .. } => {
790            for s in *body {
791                collect_modified_vars_stmt(s, out);
792            }
793        }
794        Stmt::Inspect { arms, .. } => {
795            for arm in arms {
796                for s in arm.body {
797                    collect_modified_vars_stmt(s, out);
798                }
799            }
800        }
801        Stmt::Zone { body, .. } => {
802            for s in *body {
803                collect_modified_vars_stmt(s, out);
804            }
805        }
806        _ => {}
807    }
808}
809
810// =============================================================================
811// Public API
812// =============================================================================
813
814pub fn supercompile_stmts<'a>(
815    stmts: Vec<Stmt<'a>>,
816    expr_arena: &'a Arena<Expr<'a>>,
817    stmt_arena: &'a Arena<Stmt<'a>>,
818    interner: &mut Interner,
819) -> Vec<Stmt<'a>> {
820    let funcs = collect_pure_funcs(&stmts);
821    let mut env = SuperEnv::new();
822    env.funcs = funcs;
823
824    let mut result = Vec::with_capacity(stmts.len());
825    for stmt in stmts {
826        if let Some(driven) = drive_stmt(stmt, &mut env, expr_arena, stmt_arena, interner, 0) {
827            result.push(driven);
828        }
829    }
830    result
831}
832
833/// Homeomorphic embedding check: `e1 ◁ e2`.
834///
835/// Returns true if `e1` is homeomorphically embedded in `e2`, meaning `e1` can
836/// be obtained from `e2` by erasing some constructors (diving) or both share
837/// the same constructor with all children embedded (coupling).
838pub fn embeds<'a>(e1: &Expr<'a>, e2: &Expr<'a>) -> bool {
839    match (e1, e2) {
840        // Coupling: same constructor, all children embed
841        (Expr::BinaryOp { op: op1, left: l1, right: r1 },
842         Expr::BinaryOp { op: op2, left: l2, right: r2 }) if op1 == op2 =>
843            embeds(l1, l2) && embeds(r1, r2),
844        (Expr::Call { function: f1, args: a1 },
845         Expr::Call { function: f2, args: a2 }) if f1 == f2 && a1.len() == a2.len() =>
846            a1.iter().zip(a2.iter()).all(|(x, y)| embeds(x, y)),
847        (Expr::Not { operand: e1_inner }, Expr::Not { operand: e2_inner }) =>
848            embeds(e1_inner, e2_inner),
849        // Base: literals and identifiers embed in themselves
850        (Expr::Literal(l1), Expr::Literal(l2)) => l1 == l2,
851        (Expr::Identifier(s1), Expr::Identifier(s2)) => s1 == s2,
852        // Diving: e1 embeds in a subterm of e2
853        (_, Expr::BinaryOp { left, right, .. }) =>
854            embeds(e1, left) || embeds(e1, right),
855        (_, Expr::Call { args, .. }) =>
856            args.iter().any(|a| embeds(e1, a)),
857        (_, Expr::Not { operand }) =>
858            embeds(e1, operand),
859        _ => false,
860    }
861}
862
863/// Result of Most Specific Generalization (MSG).
864pub struct MsgResult<'a> {
865    /// The generalized expression with fresh variables replacing differing parts.
866    pub expr: &'a Expr<'a>,
867    /// Number of fresh variables (substitutions) introduced.
868    pub num_substitutions: usize,
869}
870
871/// Compute the Most Specific Generalization (MSG) of two expressions.
872///
873/// Given `e1` and `e2`, finds the most specific expression `g` such that
874/// both `e1` and `e2` are instances of `g`. Differing subexpressions are
875/// replaced with fresh variables named `__msg_0`, `__msg_1`, etc.
876pub fn msg<'a>(
877    e1: &'a Expr<'a>,
878    e2: &'a Expr<'a>,
879    arena: &'a Arena<Expr<'a>>,
880    interner: &mut Interner,
881) -> MsgResult<'a> {
882    let mut counter = 0;
883    let expr = msg_inner(e1, e2, arena, interner, &mut counter);
884    MsgResult { expr, num_substitutions: counter }
885}
886
887fn msg_inner<'a>(
888    e1: &'a Expr<'a>,
889    e2: &'a Expr<'a>,
890    arena: &'a Arena<Expr<'a>>,
891    interner: &mut Interner,
892    counter: &mut usize,
893) -> &'a Expr<'a> {
894    match (e1, e2) {
895        // Same literal or same identifier → preserve
896        (Expr::Literal(l1), Expr::Literal(l2)) if l1 == l2 => e1,
897        (Expr::Identifier(s1), Expr::Identifier(s2)) if s1 == s2 => e1,
898        // Same BinaryOp constructor → recurse into children
899        (Expr::BinaryOp { op: op1, left: l1, right: r1 },
900         Expr::BinaryOp { op: op2, left: l2, right: r2 }) if op1 == op2 => {
901            let left = msg_inner(l1, l2, arena, interner, counter);
902            let right = msg_inner(r1, r2, arena, interner, counter);
903            arena.alloc(Expr::BinaryOp { op: *op1, left, right })
904        }
905        // Same Call constructor → recurse into args
906        (Expr::Call { function: f1, args: a1 },
907         Expr::Call { function: f2, args: a2 }) if f1 == f2 && a1.len() == a2.len() => {
908            let args: Vec<&'a Expr<'a>> = a1.iter().zip(a2.iter())
909                .map(|(x, y)| msg_inner(x, y, arena, interner, counter))
910                .collect();
911            arena.alloc(Expr::Call { function: *f1, args })
912        }
913        // Same Not constructor → recurse
914        (Expr::Not { operand: e1_inner }, Expr::Not { operand: e2_inner }) => {
915            let inner = msg_inner(e1_inner, e2_inner, arena, interner, counter);
916            arena.alloc(Expr::Not { operand: inner })
917        }
918        // Differing → introduce fresh variable
919        _ => {
920            let name = format!("__msg_{}", counter);
921            *counter += 1;
922            let sym = interner.intern(&name);
923            arena.alloc(Expr::Identifier(sym))
924        }
925    }
926}
logicaffeine_compile/optimize/supercompile.rs

logicaffeine_compile/optimize/
supercompile.rs
