Forgot to add captures.rs

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<Identity, Vec<Identity>>) -> BoundNode {
+        Self::transform(node, capture_map)
+    }
+
+    fn transform(mut node: BoundNode, capture_map: &HashMap<Identity, Vec<Identity>>) -> 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
+    }
+}