seqc/

capture_analysis.rs

//! Capture Analysis for Closures
//!
//! This module handles the analysis of closure captures - determining which values
//! from the creation site need to be captured in a closure's environment.
//!
//! The key insight is that closures bridge two stack effects:
//! - **Body effect**: what the quotation body actually needs to execute
//! - **Call effect**: what the call site will provide when the closure is invoked
//!
//! The difference between these determines what must be captured at creation time.
//!
//! ## Example
//!
//! ```text
//! : add-to ( Int -- [Int -- Int] )
//!   [ add ] ;
//! ```
//!
//! Here:
//! - Body needs: `(Int Int -- Int)` (add requires two integers)
//! - Call provides: `(Int -- Int)` (caller provides one integer)
//! - Captures: `[Int]` (one integer captured from creation site)

use crate::types::{Effect, StackType, Type};

/// Calculate capture types for a closure
///
/// Given:
/// - `body_effect`: what the quotation body needs (e.g., `Int Int -- Int`)
/// - `call_effect`: what the call site will provide (e.g., `Int -- Int`)
///
/// Returns:
/// - `captures`: types to capture from creation stack (e.g., `[Int]`)
///
/// # Capture Ordering
///
/// Captures are returned bottom-to-top (deepest value first), matching the
/// runtime's env layout and the order the closure body pushes them:
///
/// ```text
/// Stack at creation: ( ...rest bottom top )
/// push_closure pops top-down, then reverses:
///   env[0] = bottom (caller's deepest capture)
///   env[N-1] = top   (caller's shallowest capture)
/// Closure function pushes env[0], env[1], ..., env[N-1] in order, so
/// the body stack looks like ( ...rest bottom top ) — identical to the
/// caller's visual order.
/// ```
///
/// Single captures are unaffected (a one-element vector reversed is itself).
///
/// # Errors
///
/// Returns an error if the call site provides more values than the body needs.
pub fn calculate_captures(body_effect: &Effect, call_effect: &Effect) -> Result<Vec<Type>, String> {
    // Extract concrete types from stack types (bottom to top)
    let body_inputs = extract_concrete_types(&body_effect.inputs);
    let call_inputs = extract_concrete_types(&call_effect.inputs);

    // Validate: call site shouldn't provide MORE than body needs
    if call_inputs.len() > body_inputs.len() {
        return Err(format!(
            "Closure signature error: call site provides {} values but body only needs {}",
            call_inputs.len(),
            body_inputs.len()
        ));
    }

    // Calculate how many to capture (from bottom of stack)
    let capture_count = body_inputs.len() - call_inputs.len();

    // Verify the topmost body inputs (the non-captured ones) align with
    // what the call site provides. If they don't match, the body is
    // incompatible with the combinator regardless of captures.
    let body_provided = &body_inputs[capture_count..];
    for (i, (body_type, call_type)) in body_provided.iter().zip(call_inputs.iter()).enumerate() {
        if body_type != call_type {
            // Type variables (like Acc, T from row polymorphism) won't match
            // concrete types here — that's expected, because the body's types
            // are inferred from a seeded row-variable stack. Skip the check
            // for type variables; they'll be verified by downstream unification.
            let is_var = matches!(body_type, Type::Var(_)) || matches!(call_type, Type::Var(_));
            if !is_var {
                return Err(format!(
                    "Closure capture error: body input at position {} (from top) is {}, \
                     but combinator provides {}. The non-captured inputs must match.",
                    i, body_type, call_type
                ));
            }
        }
    }

    // Captures are the first N types (bottom of stack)
    // Example: body needs [Int, String] (bottom to top), call provides [String]
    // Captures: [Int] (the bottom type)
    Ok(body_inputs[0..capture_count].to_vec())
}

/// Extract concrete types from a stack type (bottom to top order)
///
/// This function traverses a `StackType` and returns a vector of concrete types
/// in bottom-to-top order (deepest stack element first).
///
/// # Example
///
/// ```text
/// Input: Cons { rest: Cons { rest: Empty, top: Int }, top: String }
/// Output: [Int, String]  (bottom to top)
/// ```
///
/// # Row Variables
///
/// Row variables (like `..a`) are skipped - this function only extracts
/// concrete types. This is appropriate for capture analysis where we need
/// to know the actual types being captured.
///
/// # Performance
///
/// Uses recursion to build the vector in the correct order without needing
/// to clone the entire stack structure or reverse the result.
pub(crate) fn extract_concrete_types(stack: &StackType) -> Vec<Type> {
    // Use recursion to build the vector in bottom-to-top order
    fn collect(stack: &StackType, result: &mut Vec<Type>) {
        match stack {
            StackType::Cons { rest, top } => {
                // First recurse to collect types below, then add this type
                collect(rest, result);
                result.push(top.clone());
            }
            StackType::Empty | StackType::RowVar(_) => {
                // Base case: nothing more to collect
            }
        }
    }

    let mut types = Vec::new();
    collect(stack, &mut types);
    types
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{Effect, StackType, Type};

    fn make_stack(types: &[Type]) -> StackType {
        let mut stack = StackType::Empty;
        for t in types {
            stack = StackType::Cons {
                rest: Box::new(stack),
                top: t.clone(),
            };
        }
        stack
    }

    fn make_effect(inputs: &[Type], outputs: &[Type]) -> Effect {
        Effect {
            inputs: make_stack(inputs),
            outputs: make_stack(outputs),
            effects: Vec::new(),
        }
    }

    #[test]
    fn test_extract_empty_stack() {
        let types = extract_concrete_types(&StackType::Empty);
        assert!(types.is_empty());
    }

    #[test]
    fn test_extract_single_type() {
        let stack = make_stack(&[Type::Int]);
        let types = extract_concrete_types(&stack);
        assert_eq!(types, vec![Type::Int]);
    }

    #[test]
    fn test_extract_multiple_types() {
        let stack = make_stack(&[Type::Int, Type::String, Type::Bool]);
        let types = extract_concrete_types(&stack);
        assert_eq!(types, vec![Type::Int, Type::String, Type::Bool]);
    }

    #[test]
    fn test_calculate_no_captures() {
        // Body needs (Int -- Int), call provides (Int -- Int)
        let body = make_effect(&[Type::Int], &[Type::Int]);
        let call = make_effect(&[Type::Int], &[Type::Int]);

        let captures = calculate_captures(&body, &call).unwrap();
        assert!(captures.is_empty());
    }

    #[test]
    fn test_calculate_one_capture() {
        // Body needs (Int Int -- Int), call provides (Int -- Int)
        // Should capture one Int
        let body = make_effect(&[Type::Int, Type::Int], &[Type::Int]);
        let call = make_effect(&[Type::Int], &[Type::Int]);

        let captures = calculate_captures(&body, &call).unwrap();
        assert_eq!(captures, vec![Type::Int]);
    }

    #[test]
    fn test_calculate_multiple_captures() {
        // Body needs (Int String Bool -- Bool), call provides (Bool -- Bool)
        // Should capture [Int, String] (bottom to top)
        let body = make_effect(&[Type::Int, Type::String, Type::Bool], &[Type::Bool]);
        let call = make_effect(&[Type::Bool], &[Type::Bool]);

        let captures = calculate_captures(&body, &call).unwrap();
        assert_eq!(captures, vec![Type::Int, Type::String]);
    }

    #[test]
    fn test_calculate_all_captured() {
        // Body needs (Int String -- Int), call provides ( -- Int)
        // Should capture [Int, String]
        let body = make_effect(&[Type::Int, Type::String], &[Type::Int]);
        let call = make_effect(&[], &[Type::Int]);

        let captures = calculate_captures(&body, &call).unwrap();
        assert_eq!(captures, vec![Type::Int, Type::String]);
    }

    #[test]
    fn test_calculate_error_too_many_call_inputs() {
        // Body needs (Int -- Int), call provides (Int Int -- Int)
        // Error: call provides more than body needs
        let body = make_effect(&[Type::Int], &[Type::Int]);
        let call = make_effect(&[Type::Int, Type::Int], &[Type::Int]);

        let result = calculate_captures(&body, &call);
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .contains("provides 2 values but body only needs 1")
        );
    }
}