leo-passes 4.4.1

// Copyright (C) 2019-2026 Provable Inc.
// This file is part of the Leo library.

// The Leo library is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// The Leo library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.

//! Transforms a captured `Block` into a standalone async `Function` plus a
//! corresponding call expression.
//!
//! This pass analyzes symbol accesses inside the block, determines which
//! variables must become parameters, and synthesizes the necessary `Input`s
//! and call-site arguments. Tuple and tuple-field accesses are normalized so
//! that each accessed element becomes a unique parameter, with full-tuple
//! reconstruction when needed.
//!
//! The original block is then reconstructed with all symbol references
//! replaced by these synthesized parameters. The result is a function
//! that encapsulates the block's logic and a call expression that invokes it.
//!
//! # Example
//! ```leo
//! // Original block
//! let a: i32 = 1;
//! let b: i32 = 2;
//! let c: (i32, i32) = (3, 4);
//! {
//!     let x = a + b;
//!     let y = c.0 + c.1;
//!     ..
//! }
//!
//! // Rewritten as a function + call expression (assume `variant` is `AsyncFunction` here)
//! async function generated_async(a: i32, b: i32, "c.0": i32, "c.1": i32) {
//!     let x = a + b;
//!     let y = "c.0" + "c.1";
//!     ..
//! }
//!
//! // Call
//! generated_async(a, b, c.0, c.1);
//! ```

use crate::{CompilerState, Replacer, SymbolAccessCollector};

use leo_ast::{
    AstReconstructor,
    AstVisitor,
    Block,
    CallExpression,
    Expression,
    Function,
    Identifier,
    Input,
    Location,
    Node,
    Path,
    TupleAccess,
    TupleExpression,
    TupleType,
    TypeKind,
    TypeNode,
    Variant,
};
use leo_span::{Span, Symbol};

use indexmap::IndexMap;

pub struct BlockToFunctionRewriter<'a> {
    state: &'a mut CompilerState,
    current_unit: Symbol,
}

impl<'a> BlockToFunctionRewriter<'a> {
    pub fn new(state: &'a mut CompilerState, current_unit: Symbol) -> Self {
        Self { state, current_unit }
    }
}

impl BlockToFunctionRewriter<'_> {
    pub fn rewrite_block(
        &mut self,
        input: &Block,
        function_name: Symbol,
        function_variant: Variant,
    ) -> (Function, Expression) {
        // Collect all symbol accesses in the block.
        let mut access_collector = SymbolAccessCollector::new(self.state);
        access_collector.visit_block(input);

        // Stores mapping from accessed symbol (and optional index) to the expression used in replacement.
        let mut replacements: IndexMap<(Symbol, Option<usize>), Expression> = IndexMap::new();

        // Helper to create a fresh `Identifier`.
        let make_identifier = |slf: &mut Self, symbol: Symbol| Identifier {
            name: symbol,
            span: Span::default(),
            id: slf.state.node_builder.next_id(),
        };

        // Generates a set of `Input`s and corresponding call-site `Expression`s for a given symbol access.
        //
        // This function handles both:
        // - Direct variable accesses (e.g., `foo`)
        // - Tuple element accesses (e.g., `foo.0`)
        //
        // For tuple accesses:
        // - If a single element (e.g. `foo.0`) is accessed, it generates a synthetic input like `"foo.0"`.
        // - If the whole tuple (e.g. `foo`) is accessed, it ensures all elements are covered by:
        //     - Reusing existing inputs from `replacements` if already generated via prior field access.
        //     - Creating new inputs and arguments for any missing elements.
        // - The entire tuple is reconstructed in `replacements` using the individual elements as a `TupleExpression`.
        //
        // This function also ensures deduplication by consulting the `replacements` map:
        // - If a given `(symbol, index)` has already been processed, no duplicate input or argument is generated.
        // - This prevents repeated parameters for accesses like both `foo` and `foo.0`.
        //
        // # Parameters
        // - `symbol`: The symbol being accessed.
        // - `var_type`: The type of the symbol (may be a tuple or base type).
        // - `index_opt`: `Some(index)` for a tuple field (e.g., `.0`), or `None` for full-variable access.
        //
        // # Returns
        // A `Vec<(Input, Expression)>`, where:
        // - `Input` is a parameter for the generated function.
        // - `Expression` is the call-site argument expression used to invoke that parameter.
        let mut make_inputs_and_arguments = |slf: &mut Self,
                                             symbol: Symbol,
                                             var_type: &TypeKind,
                                             index_opt: Option<usize>|
         -> Vec<(Input, Expression)> {
            if replacements.contains_key(&(symbol, index_opt)) {
                return vec![]; // No new input needed; argument already exists
            }

            match index_opt {
                Some(index) => {
                    let TypeKind::Tuple(TupleType { elements }) = var_type else {
                        // The type checker has already emitted an error for this invalid access;
                        // return no inputs so compilation can continue to report all diagnostics.
                        return vec![];
                    };

                    // The type checker has already emitted an error for this out-of-bounds access;
                    // return no inputs so compilation can continue to report all diagnostics.
                    if index >= elements.len() {
                        return vec![];
                    }

                    let synthetic_name = format!("\"{symbol}.{index}\"");
                    let synthetic_symbol = Symbol::intern(&synthetic_name);
                    let identifier = make_identifier(slf, synthetic_symbol);

                    let input = Input {
                        identifier,
                        mode: leo_ast::Mode::None,
                        type_: TypeNode::new(&slf.state.types, elements[index].clone(), Span::default()),
                        span: Span::default(),
                        id: slf.state.node_builder.next_id(),
                    };

                    replacements.insert((symbol, Some(index)), Path::from(identifier).to_local().into());

                    vec![(
                        input,
                        TupleAccess {
                            tuple: Path::from(make_identifier(slf, symbol)).to_local().into(),
                            index: index.into(),
                            span: Span::default(),
                            id: slf.state.node_builder.next_id(),
                        }
                        .into(),
                    )]
                }

                None => match var_type {
                    TypeKind::Tuple(TupleType { elements }) => {
                        let mut inputs_and_arguments = Vec::with_capacity(elements.len());
                        let mut tuple_elements = Vec::with_capacity(elements.len());

                        for (i, element_type) in elements.iter().enumerate() {
                            let key = (symbol, Some(i));

                            // Skip if this field is already handled
                            if let Some(existing_expr) = replacements.get(&key) {
                                tuple_elements.push(existing_expr.clone());
                                continue;
                            }

                            // Otherwise, synthesize identifier and input
                            let synthetic_name = format!("\"{symbol}.{i}\"");
                            let synthetic_symbol = Symbol::intern(&synthetic_name);
                            let identifier = make_identifier(slf, synthetic_symbol);

                            let input = Input {
                                identifier,
                                mode: leo_ast::Mode::None,
                                type_: TypeNode::new(&slf.state.types, element_type.clone(), Span::default()),
                                span: Span::default(),
                                id: slf.state.node_builder.next_id(),
                            };

                            let expr: Expression = Path::from(identifier).to_local().into();

                            replacements.insert(key, expr.clone());
                            tuple_elements.push(expr.clone());
                            inputs_and_arguments.push((
                                input,
                                TupleAccess {
                                    tuple: Path::from(make_identifier(slf, symbol)).to_local().into(),
                                    index: i.into(),
                                    span: Span::default(),
                                    id: slf.state.node_builder.next_id(),
                                }
                                .into(),
                            ));
                        }

                        // Now insert the full tuple (even if all fields were already there).
                        replacements.insert(
                            (symbol, None),
                            Expression::Tuple(TupleExpression {
                                elements: tuple_elements,
                                span: Span::default(),
                                id: slf.state.node_builder.next_id(),
                            }),
                        );

                        inputs_and_arguments
                    }

                    _ => {
                        let identifier = make_identifier(slf, symbol);
                        let input = Input {
                            identifier,
                            mode: leo_ast::Mode::None,
                            type_: TypeNode::new(&slf.state.types, var_type.clone(), Span::default()),
                            span: Span::default(),
                            id: slf.state.node_builder.next_id(),
                        };

                        replacements.insert((symbol, None), Path::from(identifier).to_local().into());

                        let argument = Path::from(make_identifier(slf, symbol)).to_local().into();
                        vec![(input, argument)]
                    }
                },
            }
        };

        // Resolve symbol accesses into inputs and call arguments.
        let (inputs, arguments): (Vec<_>, Vec<_>) = access_collector
            .symbol_accesses
            .iter()
            .filter_map(|(path, index)| {
                // Skip globals and variables that are local to this block or to one of its children.

                // Skip globals.
                if path.is_global() {
                    return None;
                }

                // Skip variables that are local to this block or to one of its children.
                let local_var_name = path.expect_local_symbol(); // Not global, so must be local 
                if self.state.symbol_table.is_local_to_or_in_child_scope(input.id(), local_var_name) {
                    return None;
                }

                // All other variables become parameters to the function being built.
                let var = self.state.symbol_table.lookup_local(local_var_name)?;
                // The variable's type is `None` only when an earlier type-checking diagnostic
                // already prevented it from being typed (e.g. a tuple-pattern arity mismatch
                // like `let (a, b, c): (u32, u32) = ...` leaves the surplus identifier
                // un-typed). Skip such variables so the rewriter does not panic before the
                // user-facing diagnostic surfaces.
                let ty = var.type_.clone()?;
                Some(make_inputs_and_arguments(self, local_var_name, &ty, *index))
            })
            .flatten()
            .unzip();

        // Replacement logic used to patch the block.
        let replace_expr = |expr: &Expression| -> Expression {
            match expr {
                Expression::Path(path) => {
                    replacements.get(&(path.identifier().name, None)).cloned().unwrap_or_else(|| expr.clone())
                }

                Expression::TupleAccess(ta) => {
                    if let Expression::Path(path) = &ta.tuple {
                        replacements
                            .get(&(path.identifier().name, Some(ta.index.value())))
                            .cloned()
                            .unwrap_or_else(|| expr.clone())
                    } else {
                        expr.clone()
                    }
                }

                _ => expr.clone(),
            }
        };

        // Reconstruct the block with replaced references.
        let mut replacer = Replacer::new(replace_expr, true /* refresh IDs */, self.state);
        let new_block = replacer.reconstruct_block(input.clone()).0;

        // Define the new function.
        let function = Function {
            is_exported: None,
            annotations: vec![],
            variant: function_variant,
            identifier: make_identifier(self, function_name),
            const_parameters: vec![],
            input: inputs,
            output: vec![],              // No returns supported yet.
            output_type: TypeKind::Unit, // No returns supported yet.
            block: new_block,
            span: input.span,
            id: self.state.node_builder.next_id(),
        };

        // Create the call expression to invoke the function.
        let call_to_function = CallExpression {
            function: Path::from(make_identifier(self, function_name))
                .to_global(Location::new(self.current_unit, vec![function_name])),
            const_arguments: vec![],
            arguments,
            span: input.span,
            id: self.state.node_builder.next_id(),
        };

        (function, call_to_function.into())
    }
}