diff --git a/src/ast/compiler/captures.rs b/src/ast/compiler/captures.rs new file mode 100644 index 0000000..5c53a9c --- /dev/null +++ b/src/ast/compiler/captures.rs @@ -0,0 +1,133 @@ +use crate::ast::compiler::bound_nodes::{BoundKind, BoundNode}; +use crate::ast::types::Identity; +use std::collections::HashMap; + +pub struct CapturePass; + +impl CapturePass { + pub fn apply(node: BoundNode, capture_map: &HashMap>) -> BoundNode { + Self::transform(node, capture_map) + } + + fn transform(mut node: BoundNode, capture_map: &HashMap>) -> BoundNode { + match node.kind { + BoundKind::Define { + name, + addr, + kind, + value, + .. + } => { + let captured_by = capture_map.get(&node.identity).cloned().unwrap_or_default(); + node.kind = BoundKind::Define { + name, + addr, + kind, + value: Box::new(Self::transform(*value, capture_map)), + captured_by, + }; + } + + BoundKind::If { + cond, + then_br, + else_br, + } => { + node.kind = BoundKind::If { + cond: Box::new(Self::transform(*cond, capture_map)), + then_br: Box::new(Self::transform(*then_br, capture_map)), + else_br: else_br.map(|e| Box::new(Self::transform(*e, capture_map))), + }; + } + + BoundKind::Set { addr, value } => { + node.kind = BoundKind::Set { + addr, + value: Box::new(Self::transform(*value, capture_map)), + }; + } + + BoundKind::Destructure { pattern, value } => { + node.kind = BoundKind::Destructure { + pattern: Box::new(Self::transform(*pattern, capture_map)), + value: Box::new(Self::transform(*value, capture_map)), + }; + } + + BoundKind::Lambda { + params, + upvalues, + body, + positional_count, + } => { + node.kind = BoundKind::Lambda { + params: std::rc::Rc::new(Self::transform(params.as_ref().clone(), capture_map)), + upvalues, + body: std::rc::Rc::new(Self::transform(body.as_ref().clone(), capture_map)), + positional_count, + }; + } + + BoundKind::Call { callee, args } => { + node.kind = BoundKind::Call { + callee: Box::new(Self::transform(*callee, capture_map)), + args: Box::new(Self::transform(*args, capture_map)), + }; + } + + BoundKind::Again { args } => { + node.kind = BoundKind::Again { + args: Box::new(Self::transform(*args, capture_map)), + }; + } + + BoundKind::Block { exprs } => { + node.kind = BoundKind::Block { + exprs: exprs + .into_iter() + .map(|e| Self::transform(e, capture_map)) + .collect(), + }; + } + + BoundKind::Tuple { elements } => { + node.kind = BoundKind::Tuple { + elements: elements + .into_iter() + .map(|e| Self::transform(e, capture_map)) + .collect(), + }; + } + + BoundKind::Record { fields } => { + node.kind = BoundKind::Record { + fields: fields + .into_iter() + .map(|(k, v)| { + ( + Self::transform(k, capture_map), + Self::transform(v, capture_map), + ) + }) + .collect(), + }; + } + + BoundKind::Expansion { + original_call, + bound_expanded, + } => { + node.kind = BoundKind::Expansion { + original_call, + bound_expanded: Box::new(Self::transform(*bound_expanded, capture_map)), + }; + } + + BoundKind::Nop + | BoundKind::Constant(_) + | BoundKind::Get { .. } + | BoundKind::Extension(_) => {} + } + node + } +}