leo-passes 3.5.0

// 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/>.

use crate::{AleoConstructor, AleoExpr, AleoReg, CompilerState};

use crate::{Bytecode, CompiledPrograms};
use leo_ast::{Function, Location, Program, ProgramId, Variant};
use leo_span::Symbol;

use snarkvm::prelude::Network;

use indexmap::{IndexMap, IndexSet};
use itertools::Itertools;
use std::str::FromStr;

pub struct CodeGeneratingVisitor<'a> {
    pub state: &'a CompilerState,
    /// A counter to track the next available register.
    pub next_register: u64,
    /// Reference to the current function.
    pub current_function: Option<&'a Function>,
    /// Mapping of local variables to registers.
    /// Because these are local, we can identify them using only a `Symbol` (i.e. a path is not necessary here).
    pub variable_mapping: IndexMap<Symbol, AleoExpr>,
    /// Mapping of composites to a Boolean indicate whether the composite is a record or not.
    /// A composite here is identified by its `Location`.
    pub composite_mapping: IndexMap<Location, bool>,
    /// Mapping of global identifiers to their associated names.
    pub global_mapping: IndexMap<Location, String>,
    /// The variant of the function we are currently traversing.
    pub variant: Option<Variant>,
    /// A reference to program. This is needed to look up external programs.
    pub program: &'a Program,
    /// The program ID of the current program.
    pub program_id: Option<ProgramId>,
    /// A reference to the finalize caller.
    /// Because `async transition`s  can only appear in the top level program scope at this stage,
    /// it's safe to keep this a `Symbol` (i.e. a path is not necessary).
    pub finalize_caller: Option<Symbol>,
    /// A counter to track the next available label.
    pub next_label: u64,
    /// The depth of the current conditional block.
    pub conditional_depth: u64,
    /// Internal record input registers of the current function.
    /// This is necessary as if we output them, we need to clone them.
    pub internal_record_inputs: IndexSet<AleoExpr>,
}

/// This function checks whether the constructor is well-formed.
/// If an upgrade configuration is provided, it checks that the constructor matches the configuration.
pub(crate) fn check_snarkvm_constructor<N: Network>(actual: &AleoConstructor) -> snarkvm::prelude::Result<()> {
    use snarkvm::synthesizer::program::Constructor as SVMConstructor;
    // Parse the constructor as a snarkVM constructor.
    SVMConstructor::<N>::from_str(actual.to_string().trim())?;

    Ok(())
}

impl CodeGeneratingVisitor<'_> {
    pub(crate) fn visit_package(&mut self) -> CompiledPrograms {
        let import_bytecodes = self
            .state
            .ast
            .as_repr()
            .stubs
            .values()
            .filter_map(|stub| {
                match stub {
                    leo_ast::Stub::FromLeo { program, .. } => {
                        let program_name = program
                            .program_scopes
                            .first()
                            .expect("programs must have a single program scope at this time.")
                            .0;

                        // Get transitive imports for this program
                        let transitive_imports = self
                            .state
                            .symbol_table
                            .get_transitive_imports(program_name)
                            .into_iter()
                            .map(|sym| sym.to_string())
                            .collect::<Vec<_>>();

                        // Generate this stub’s Aleo program text
                        let mut bytecode = self.visit_program(program);
                        bytecode.imports.extend(transitive_imports);

                        Some(Bytecode { program_name: program_name.to_string(), bytecode: bytecode.to_string() })
                    }
                    leo_ast::Stub::FromAleo { program, .. } => {
                        // Nothing to return here because the bytecode is already computed.
                        // However, we still need to collect global items into the appropriate mappings.
                        for (name, _) in &program.mappings {
                            self.global_mapping
                                .insert(Location::new(program.stub_id.name.name, vec![*name]), name.to_string());
                        }
                        for (name, leo_ast::Composite { is_record, .. }) in &program.composites {
                            self.composite_mapping
                                .insert(Location::new(program.stub_id.name.name, vec![*name]), *is_record);
                        }
                        None
                    }
                }
            })
            .collect();

        let primary_bytecode = self.visit_program(self.state.ast.as_repr()).to_string();

        CompiledPrograms { primary_bytecode, import_bytecodes }
    }

    pub(crate) fn next_register(&mut self) -> AleoReg {
        self.next_register += 1;
        AleoReg::R(self.next_register - 1)
    }

    /// Converts a path into a legal Aleo identifier, if possible.
    ///
    /// # Behavior
    /// - If the path is a single valid Leo identifier (`[a-zA-Z][a-zA-Z0-9_]*`), it's returned as-is.
    /// - If the last segment matches `Name::[args]` (e.g. `Vec3::[3, 4]`), it's converted to a legal identifier using hashing.
    /// - If the path has multiple segments, and all segments are valid identifiers except the last one (which may be `Name::[args]`),
    ///   it's hashed using the last segment as base.
    /// - Returns `None` if:
    ///   - The path is empty
    ///   - Any segment other than the last is not a valid identifier
    ///   - The last segment is invalid and not legalizable
    ///
    /// # Parameters
    /// - `path`: A slice of `Symbol`s representing a path to an item.
    ///
    /// # Returns
    /// - `Some(String)`: A valid Leo identifier.
    /// - `None`: If the path is invalid or cannot be legalized.
    pub(crate) fn legalize_path(path: &[Symbol]) -> Option<String> {
        /// Checks if a string is a legal Leo identifier: [a-zA-Z][a-zA-Z0-9_]*
        fn is_legal_identifier(s: &str) -> bool {
            let mut chars = s.chars();
            matches!(chars.next(), Some(c) if c.is_ascii_alphabetic())
                && chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
                && s.len() <= 31
        }

        /// Generates a hashed Leo identifier from the full path, using the given base segment.
        fn generate_hashed_name(path: &[Symbol], base: &str) -> String {
            use base62::encode;
            use sha2::{Digest, Sha256};
            use std::fmt::Write;

            let full_path = path.iter().format("::").to_string();

            let mut hasher = Sha256::new();
            hasher.update(full_path.as_bytes());
            let hash = hasher.finalize();

            let hash_number = u64::from_be_bytes(hash[..8].try_into().unwrap());
            let hash_base62 = encode(hash_number);

            let fixed_suffix_len = 2 + hash_base62.len(); // "__" + hash
            let max_ident_len = 31 - fixed_suffix_len;

            let mut result = String::new();
            write!(&mut result, "{base:.max_ident_len$}__{hash_base62}").unwrap();
            result
        }

        let last = path.last()?.to_string();

        // Validate all segments except the last
        if path.len() > 1 && !path[..path.len() - 1].iter().all(|sym| is_legal_identifier(&sym.to_string())) {
            return None;
        }

        // === Case 1: Single, legal identifier ===
        if path.len() == 1 && is_legal_identifier(&last) {
            return Some(last);
        }

        // === Case 2: Storage-generated identifiers ===
        //
        // These occur after lowering storage variables into mappings.
        // They can take the forms:
        //   - `some_var__` (for singleton storage)
        //   - `some_var__len__` (for vector length mapping)
        //
        // Both can exceed 31 chars once the "__" or "__len__" suffix is added,
        // so we truncate the base portion to preserve the suffix and make space
        // for the hash. The truncation lengths below were chosen so that:
        // ```
        // total_length = prefix_len + suffix_len + "__" + hash_len <= 31
        // ```
        // where hash_len = 11.
        if let Some(prefix) = last.strip_suffix("__len__") {
            // "some_very_long_storage_variable__len__"
            //  - Keep at most 13 chars from the prefix
            //  - Produces something like: "some_very_lon__len__CpUbpLTf1Ow"
            let truncated_prefix = &prefix[..13.min(prefix.len())];
            return Some(generate_hashed_name(path, &(truncated_prefix.to_owned() + "__len")));
        }

        if let Some(prefix) = last.strip_suffix("__") {
            // "some_very_long_storage_variable__"
            //  - Keep at most 18 chars from the prefix
            //  - Produces something like: "some_very_long_sto__Hn1pThQeV3"
            let truncated_prefix = &prefix[..18.min(prefix.len())];
            return Some(generate_hashed_name(path, &(truncated_prefix.to_owned() + "__")));
        }

        // === Case 3: Matches special form like `path::to::Name::[3, 4]` ===
        let re = regex::Regex::new(r#"^([a-zA-Z_][\w]*)(?:::\[.*?\])?$"#).unwrap();

        if let Some(captures) = re.captures(&last) {
            let ident = captures.get(1)?.as_str();

            // The produced name here will be of the form: `<last>__AYMqiUeJeQN`.
            return Some(generate_hashed_name(path, ident));
        }

        // === Case 4: Matches special form like `path::to::Name?` (last always ends with `?`) ===
        if last.ends_with("?\"") {
            // Because the last segment of `path` always ends with `?` in case 3, we can guarantee
            // that there will be no conflicts with case 2 (which doesn't allow `?` anywhere).
            //
            // The produced name here will be of the form: `Optional__JZCpIGdQvEZ`.
            // The suffix after the `__` cannot conflict with the suffix in case 2 because of the `?`
            return Some(generate_hashed_name(path, "Optional"));
        }

        // Last segment is neither legal nor matches special pattern
        None
    }
}