Brummel/RustAst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add AST analysis pass for purity and recursion

Browse Source 
This commit introduces a new AST analysis pass that identifies function
purity and recursion. This information is then used by the optimizer and
specializer to make more informed decisions, particularly regarding
inlining.

The `Analyzer` struct and its associated `Analysis` struct are
responsible for traversing the AST and collecting this data.

Key changes include:
- A new `analyzer` module is added to `ast::compiler`.
- `Analyzer::analyze` performs a two-pass traversal to collect
  global-to-lambda mappings and then analyze purity and recursion.
- The `Optimizer` and `Specializer` are updated to accept and utilize
  the `Analysis` data.
- Recursion checks in `Optimizer` and `Specializer` are replaced with
  checks against the pre-computed `Analysis.is_recursive` set.
- The `Environment` now stores and passes the `Analysis` results to the
  compiler stages.
This commit is contained in:
Michael Schimmel
2026-02-22 15:00:20 +01:00
parent 5a017fb932
commit 8f7947bde1
6 changed files with 267 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/analyzer.rs
+173
View File
@@ -0,0 +1,173 @@
use crate::ast::compiler::bound_nodes::{Address, BoundKind, TypedNode};
use crate::ast::types::{Identity, Purity, StaticType};
use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone, Default)]
pub struct Analysis {
pub purity: HashMap<Identity, Purity>,
pub is_recursive: HashSet<Identity>,
}
pub struct Analyzer<'a> {
global_purity: &'a HashMap<u32, Purity>,
results: Analysis,
/// Stack of currently visiting lambdas to detect direct recursion.
lambda_stack: Vec<Identity>,
/// Map of global index to its Lambda identity if known.
globals_to_lambdas: HashMap<u32, Identity>,
}
impl<'a> Analyzer<'a> {
pub fn analyze(node: &TypedNode, global_purity: &'a HashMap<u32, Purity>) -> Analysis {
let mut analyzer = Self {
global_purity,
results: Analysis::default(),
lambda_stack: Vec::new(),
globals_to_lambdas: HashMap::new(),
};
// First pass: map globals to their lambda identities
analyzer.collect_globals(node);
// Second pass: full analysis
analyzer.visit(node);
analyzer.results
}
fn collect_globals(&mut self, node: &TypedNode) {
match &node.kind {
BoundKind::DefGlobal { global_index, value, .. } => {
if let BoundKind::Lambda { .. } = &value.kind {
self.globals_to_lambdas.insert(*global_index, value.identity.clone());
}
self.collect_globals(value);
}
BoundKind::Block { exprs } => {
for e in exprs { self.collect_globals(e); }
}
_ => {
// Simplified traversal for global collection
node.kind.for_each_child(|child| self.collect_globals(child));
}
}
}
fn visit(&mut self, node: &TypedNode) -> Purity {
let purity = match &node.kind {
BoundKind::Constant(_) | BoundKind::Nop | BoundKind::Parameter { .. } => Purity::Pure,
BoundKind::Get { addr, .. } => match addr {
Address::Global(idx) => self.global_purity.get(idx).cloned().unwrap_or(Purity::Pure),
_ => Purity::Pure, // Locals are considered pure access in this model
},
BoundKind::Set { .. } => Purity::Impure,
BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => {
self.visit(value)
}
BoundKind::If { cond, then_br, else_br } => {
let p_cond = self.visit(cond);
let p_then = self.visit(then_br);
let p_else = else_br.as_ref().map(|e| self.visit(e)).unwrap_or(Purity::Pure);
p_cond.min(p_then).min(p_else)
}
BoundKind::Lambda { body, .. } => {
self.lambda_stack.push(node.identity.clone());
self.visit(body);
self.lambda_stack.pop();
Purity::Pure // Creating a lambda is pure
}
BoundKind::Call { callee, args } => {
let p_callee = self.visit(callee);
let p_args = self.visit(args);
// Detect recursion
if let BoundKind::Get { addr: Address::Global(idx), .. } = &callee.kind
&& let Some(lambda_id) = self.globals_to_lambdas.get(idx)
&& self.lambda_stack.contains(lambda_id)
{
self.results.is_recursive.insert(lambda_id.clone());
// Also mark the call itself if needed
self.results.is_recursive.insert(node.identity.clone());
}
// For purity, we'd need to know the function's purity.
// For now, if it's a call, we conservatively check if it's a known pure global.
let p_func = if let BoundKind::Get { addr: Address::Global(idx), .. } = &callee.kind {
self.global_purity.get(idx).cloned().unwrap_or(Purity::Impure)
} else {
Purity::Impure
};
p_callee.min(p_args).min(p_func)
}
BoundKind::Block { exprs } => {
let mut p = Purity::Pure;
for e in exprs {
p = p.min(self.visit(e));
}
p
}
BoundKind::Tuple { elements } => {
let mut p = Purity::Pure;
for e in elements {
p = p.min(self.visit(e));
}
p
}
BoundKind::Record { fields } => {
let mut p = Purity::Pure;
for (k, v) in fields {
p = p.min(self.visit(k)).min(self.visit(v));
}
p
}
_ => Purity::Impure,
};
self.results.purity.insert(node.identity.clone(), purity);
purity
}
}
/// Extension trait to make traversal easier
trait NodeExt {
fn for_each_child<F: FnMut(&TypedNode)>(&self, f: F);
}
impl NodeExt for BoundKind<StaticType> {
fn for_each_child<F: FnMut(&TypedNode)>(&self, mut f: F) {
match self {
BoundKind::If { cond, then_br, else_br } => {
f(cond); f(then_br); if let Some(e) = else_br { f(e); }
}
BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } | BoundKind::Set { value, .. } => {
f(value);
}
BoundKind::Lambda { params, body, .. } => {
f(params); f(body);
}
BoundKind::Call { callee, args } => {
f(callee); f(args);
}
BoundKind::Block { exprs } => {
for e in exprs { f(e); } // Block
}
BoundKind::Tuple { elements } => {
for e in elements { f(e); } // Tuple
}
BoundKind::Record { fields } => {
for (k, v) in fields { f(k); f(v); }
}
_ => {}
}
}
}
src/ast/compiler/mod.rs
+1
View File
@@ -1,3 +1,4 @@
pub mod analyzer;
pub mod binder;
pub mod bound_nodes;
pub mod dumper;
src/ast/compiler/optimizer.rs
+21 -51
View File
@@ -1,3 +1,4 @@
use crate::ast::compiler::analyzer::Analysis;
use crate::ast::compiler::bound_nodes::{Address, BoundKind, TypedNode};
use crate::ast::nodes::Node;
use crate::ast::types::{Purity, StaticType, Value};
@@ -14,6 +15,7 @@ pub struct Optimizer {
pub globals: Option<Rc<RefCell<Vec<Value>>>>,
pub global_purity: Option<Rc<RefCell<HashMap<u32, Purity>>>>,
pub lambda_registry: Option<Rc<RefCell<HashMap<u32, TypedNode>>>>,
pub analysis: Analysis,
}
/// Global state for tracking the optimization path (recursion, depth).
@@ -64,9 +66,15 @@ impl Optimizer {
globals: None,
global_purity: None,
lambda_registry: None,
analysis: Analysis::default(),
}
}
pub fn with_analysis(mut self, analysis: Analysis) -> Self {
self.analysis = analysis;
self
}
pub fn with_globals(mut self, globals: Rc<RefCell<Vec<Value>>>) -> Self {
self.globals = Some(globals);
self
@@ -100,34 +108,6 @@ impl Optimizer {
current
}
fn is_recursive(
&self,
node: &TypedNode,
global_idx: Option<u32>,
local_slot: Option<u32>,
identity: Option<&crate::ast::types::Identity>,
) -> bool {
let mut info = UsageInfo::default();
self.collect_usage(node, &mut info);
if let Some(idx) = global_idx
&& info.used_globals.contains(&idx)
{
return true;
}
if let Some(slot) = local_slot
&& info.used_locals.contains(&slot)
{
return true;
}
if let Some(id) = identity
&& info.used_identities.contains(id)
{
return true;
}
false
}
fn visit_node(
&self,
node: TypedNode,
@@ -257,8 +237,8 @@ impl Optimizer {
&& positional_count.is_some()
&& path.inlining_depth < 5
{
// STRICT RECURSION CHECK: Don't inline if recursive
if !self.is_recursive(body, None, None, Some(&callee.identity))
// USE STATIC ANALYSIS: Don't inline if the function is known to be recursive
if !self.analysis.is_recursive.contains(&callee.identity)
&& path.enter_lambda(&callee.identity)
{
path.inlining_depth += 1;
@@ -298,13 +278,8 @@ impl Optimizer {
&& upvalues.is_empty()
&& positional_count.is_some()
{
// STRICT RECURSION CHECK: Don't inline if recursive (idx or identity)
if !self.is_recursive(
body,
Some(*idx),
None,
Some(&lambda_node.identity),
) {
// USE STATIC ANALYSIS: Don't inline if the function is known to be recursive
if !self.analysis.is_recursive.contains(&lambda_node.identity) {
let mut inner_sub = SubstitutionMap::new();
path.inlining_stack.insert(*idx);
path.inlining_depth += 1;
@@ -332,13 +307,9 @@ impl Optimizer {
&& (closure.upvalues.is_empty()
|| self.purity_of(&closure.function_node) >= Purity::SideEffectFree)
{
// STRICT RECURSION CHECK: Don't inline if closure body contains its own identity
if !self.is_recursive(
&closure.function_node,
None,
None,
Some(&closure.function_node.identity),
) && path.enter_lambda(&closure.function_node.identity)
// USE STATIC ANALYSIS: Don't inline if the function is known to be recursive
if !self.analysis.is_recursive.contains(&closure.function_node.identity)
&& path.enter_lambda(&closure.function_node.identity)
{
let mut closure_sub = SubstitutionMap::new();
for (i, cell) in closure.upvalues.iter().enumerate() {
@@ -755,7 +726,7 @@ impl Optimizer {
}
}
fn is_inlinable_value(&self, val: &Value, global_idx: Option<u32>) -> bool {
fn is_inlinable_value(&self, val: &Value, _global_idx: Option<u32>) -> bool {
match val {
Value::Int(_)
| Value::Float(_)
@@ -767,12 +738,7 @@ impl Optimizer {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
// A closure is only inlinable if it's not recursive
closure.upvalues.is_empty()
&& !self.is_recursive(
&closure.function_node,
global_idx,
None,
Some(&closure.function_node.identity),
)
&& !self.analysis.is_recursive.contains(&closure.function_node.identity)
} else {
false
}
@@ -782,6 +748,10 @@ impl Optimizer {
}
fn purity_of(&self, node: &TypedNode) -> Purity {
if let Some(p) = self.analysis.purity.get(&node.identity) {
return *p;
}
match &node.kind {
BoundKind::Constant(_) | BoundKind::Parameter { .. } | BoundKind::Nop => Purity::Pure,
// Defining a lambda is pure; only calling it may have side effects.
src/ast/compiler/specializer.rs
+9
View File
@@ -1,3 +1,4 @@
use crate::ast::compiler::analyzer::Analysis;
use crate::ast::compiler::bound_nodes::{Address, BoundKind, BoundNode, TypedNode};
use crate::ast::nodes::Node;
use crate::ast::types::{Signature, StaticType, Value};
@@ -22,6 +23,7 @@ pub type MonoCache = HashMap<MonoCacheKey, (Value, StaticType)>;
pub struct Specializer {
pub cache: Rc<RefCell<MonoCache>>,
pub analysis: Analysis,
registry: Option<Rc<dyn FunctionRegistry>>,
compiler: Option<CompileFunc>,
rtl_lookup: Option<RtlLookupFunc>,
@@ -33,9 +35,11 @@ impl Specializer {
compiler: Option<CompileFunc>,
rtl_lookup: Option<RtlLookupFunc>,
cache: Option<Rc<RefCell<MonoCache>>>,
analysis: Analysis,
) -> Self {
Self {
cache: cache.unwrap_or_else(|| Rc::new(RefCell::new(HashMap::new()))),
analysis,
registry,
compiler,
rtl_lookup,
@@ -245,6 +249,11 @@ impl Specializer {
// 6. Resolve Function Definition
if let Some(func_node) = self.registry.as_ref().and_then(|r| r.resolve(address)) {
// Check for recursion from pre-pass.
if self.analysis.is_recursive.contains(&func_node.identity) {
return (new_callee, new_args, original_ty);
}
// Check constraints (no closures with state)
if let BoundKind::Lambda { upvalues, .. } = &func_node.kind {
if !upvalues.is_empty() {
src/ast/environment.rs
+15 -2
View File
@@ -1,3 +1,4 @@
use crate::ast::compiler::analyzer::{Analysis, Analyzer};
use crate::ast::compiler::binder::Binder;
use crate::ast::compiler::{TypeChecker, TypedNode};
use crate::ast::nodes::{Node, Symbol, UntypedKind};
@@ -29,6 +30,7 @@ pub struct Environment {
pub monomorph_cache: Rc<RefCell<MonoCache>>,
pub debug_mode: bool,
pub optimization: bool,
pub last_analysis: RefCell<Analysis>,
}
struct EnvFunctionRegistry {
@@ -96,6 +98,7 @@ impl Environment {
monomorph_cache: Rc::new(RefCell::new(HashMap::new())),
debug_mode: false,
optimization: true,
last_analysis: RefCell::new(Analysis::default()),
};
env.register_stdlib();
env
@@ -204,6 +207,10 @@ impl Environment {
// 7. Collect Typed Lambdas
LambdaCollector::collect(&typed_ast, &mut self.typed_function_registry.borrow_mut());
// 8. Analyze (Purity, Recursion)
let analysis = Analyzer::analyze(&typed_ast, &self.global_purity.borrow());
*self.last_analysis.borrow_mut() = analysis;
Ok(typed_ast)
}
@@ -216,7 +223,8 @@ impl Environment {
let optimizer = Optimizer::new(self.optimization)
.with_globals(self.global_values.clone())
.with_purity(self.global_purity.clone())
.with_registry(self.typed_function_registry.clone());
.with_registry(self.typed_function_registry.clone())
.with_analysis(self.last_analysis.borrow().clone());
let optimized = optimizer.optimize(specialized);
// 3. TCO (Always performed, converts to ExecNode)
@@ -302,7 +310,9 @@ impl Environment {
let global_types = self.global_types.clone();
let global_purity = self.global_purity.clone();
let optimization = self.optimization;
let analysis = self.last_analysis.borrow().clone();
let compiler_analysis = analysis.clone();
let compiler = Rc::new(
move |func_template: BoundNode,
arg_types: &[StaticType]|
@@ -323,6 +333,7 @@ impl Environment {
None,
Some(sub_rtl_lookup),
Some(mono_cache.clone()),
compiler_analysis.clone(),
);
let specialized_ast = sub_specializer.specialize(retyped_ast);
@@ -330,7 +341,8 @@ impl Environment {
// 3. Optimize (Phase 2: Cracking & Folding)
let optimizer = Optimizer::new(optimization)
.with_globals(global_values.clone())
.with_purity(global_purity.clone());
.with_purity(global_purity.clone())
.with_analysis(compiler_analysis.clone());
let optimized_ast = optimizer.optimize(specialized_ast);
// 4. TCO (converts to ExecNode)
@@ -359,6 +371,7 @@ impl Environment {
Some(compiler),
Some(rtl_lookup),
Some(self.monomorph_cache.clone()),
analysis,
);
specializer.specialize(node)