dissolve-python 0.3.0

// Copyright (C) 2024 Jelmer Vernooij <jelmer@samba.org>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Optimized visitor implementation that reduces memory allocations and clones

use crate::core::ReplaceInfo;
// Removed unused imports
use rustpython_ast as ast;
use std::borrow::Cow;
use std::collections::HashMap;

/// Optimized parameter mapper that reduces string clones
pub struct OptimizedParameterMapper<'a> {
    /// String cache for parameter names (these are often repeated)
    param_cache: &'a mut HashMap<String, String>,
}

impl<'a> OptimizedParameterMapper<'a> {
    pub fn new(param_cache: &'a mut HashMap<String, String>) -> Self {
        Self { param_cache }
    }

    /// Map function call arguments to parameters with minimal cloning
    pub fn map_call_arguments(
        &mut self,
        args: &[ast::Expr],
        keywords: &[ast::Keyword],
        replace_info: &ReplaceInfo,
        func: &ast::Expr,
        expr_to_string: impl Fn(&ast::Expr) -> String,
    ) -> HashMap<String, String> {
        let mut arg_map = HashMap::new();

        self.map_positional_args(&mut arg_map, args, replace_info, func, &expr_to_string);
        self.map_keyword_args(&mut arg_map, keywords, &expr_to_string);

        arg_map
    }

    /// Map positional arguments with optimized string handling
    fn map_positional_args(
        &mut self,
        arg_map: &mut HashMap<String, String>,
        args: &[ast::Expr],
        replace_info: &ReplaceInfo,
        func: &ast::Expr,
        expr_to_string: &impl Fn(&ast::Expr) -> String,
    ) {
        let is_method_call = matches!(func, ast::Expr::Attribute(_));

        if is_method_call {
            self.map_method_call_args(arg_map, args, replace_info, func, expr_to_string);
        } else {
            self.map_function_call_args(arg_map, args, replace_info, expr_to_string);
        }
    }

    /// Map arguments for method calls (handle 'self' parameter)
    fn map_method_call_args(
        &mut self,
        arg_map: &mut HashMap<String, String>,
        args: &[ast::Expr],
        replace_info: &ReplaceInfo,
        func: &ast::Expr,
        expr_to_string: &impl Fn(&ast::Expr) -> String,
    ) {
        if let ast::Expr::Attribute(attr_expr) = func {
            let self_str = expr_to_string(&attr_expr.value);

            // Use cached string for 'self' if possible
            let self_key = self.get_cached_string("self");
            arg_map.insert(self_key, self_str);

            // Map remaining arguments (skip 'self' parameter)
            let method_params: Vec<_> = replace_info
                .parameters
                .iter()
                .filter(|p| p.name != "self")
                .collect();

            for (i, arg) in args.iter().enumerate() {
                if let Some(param) = method_params.get(i) {
                    let arg_str = expr_to_string(arg);
                    let param_key = self.get_cached_string(&param.name);
                    arg_map.insert(param_key, arg_str);
                }
            }
        }
    }

    /// Map arguments for regular function calls
    fn map_function_call_args(
        &mut self,
        arg_map: &mut HashMap<String, String>,
        args: &[ast::Expr],
        replace_info: &ReplaceInfo,
        expr_to_string: &impl Fn(&ast::Expr) -> String,
    ) {
        for (i, arg) in args.iter().enumerate() {
            if let Some(param) = replace_info.parameters.get(i) {
                let arg_str = expr_to_string(arg);
                let param_key = self.get_cached_string(&param.name);
                arg_map.insert(param_key, arg_str);
            }
        }
    }

    /// Map keyword arguments
    fn map_keyword_args(
        &mut self,
        arg_map: &mut HashMap<String, String>,
        keywords: &[ast::Keyword],
        expr_to_string: &impl Fn(&ast::Expr) -> String,
    ) {
        for keyword in keywords {
            if let Some(arg_name) = &keyword.arg {
                let arg_str = expr_to_string(&keyword.value);
                let arg_key = self.get_cached_string(arg_name);
                arg_map.insert(arg_key, arg_str);
            }
        }
    }

    /// Get a cached string to avoid repeated allocations
    fn get_cached_string(&mut self, s: &str) -> String {
        if let Some(cached) = self.param_cache.get(s) {
            cached.clone()
        } else {
            let owned = s.to_string();
            self.param_cache.insert(owned.clone(), owned.clone());
            owned
        }
    }
}

/// Optimized string replacement that uses Cow<str> to minimize allocations
pub fn optimize_parameter_substitution<'a>(
    template: &'a str,
    arg_map: &HashMap<String, String>,
) -> Cow<'a, str> {
    // First pass: check if any substitutions are needed
    let needs_substitution = arg_map
        .keys()
        .any(|param_name| template.contains(&format!("{{{}}}", param_name)));

    if !needs_substitution {
        return Cow::Borrowed(template);
    }

    // Perform substitutions on owned string
    let mut result = template.to_string();
    for (param_name, arg_value) in arg_map {
        let placeholder = format!("{{{}}}", param_name);
        result = result.replace(&placeholder, arg_value);
    }

    Cow::Owned(result)
}

/// Optimized module name extraction that avoids clones
pub fn extract_source_module(old_name: &str) -> Option<&str> {
    old_name
        .find('.')
        .map(|first_dot_pos| &old_name[..first_dot_pos])
}

/// Optimized module prefix stripping
pub fn strip_module_prefix<'a>(
    replacement: &'a str,
    source_module: Option<&str>,
    current_module: &str,
) -> Cow<'a, str> {
    match source_module {
        Some(src_mod) if src_mod == current_module => {
            let prefix = format!("{}.", src_mod);
            if replacement.contains(&prefix) {
                Cow::Owned(replacement.replace(&prefix, ""))
            } else {
                Cow::Borrowed(replacement)
            }
        }
        _ => Cow::Borrowed(replacement),
    }
}

/// Optimized replacement building that minimizes string allocations
pub struct OptimizedReplacementBuilder<'a> {
    template: &'a str,
    substitutions: HashMap<&'a str, String>,
}

impl<'a> OptimizedReplacementBuilder<'a> {
    pub fn new(template: &'a str) -> Self {
        Self {
            template,
            substitutions: HashMap::new(),
        }
    }

    /// Add a substitution using borrowed key to avoid allocation
    pub fn add_substitution(&mut self, key: &'a str, value: String) {
        self.substitutions.insert(key, value);
    }

    /// Build the final replacement string with minimal allocations
    pub fn build(self) -> String {
        let mut result = self.template.to_string();

        for (key, value) in self.substitutions {
            let placeholder = format!("{{{}}}", key);
            result = result.replace(&placeholder, &value);
        }

        result
    }

    /// Build with cow optimization - returns borrowed if no substitutions needed
    pub fn build_cow(self) -> Cow<'a, str> {
        if self.substitutions.is_empty() {
            return Cow::Borrowed(self.template);
        }

        // Check if any placeholders exist
        let has_placeholders = self
            .substitutions
            .keys()
            .any(|key| self.template.contains(&format!("{{{}}}", key)));

        if !has_placeholders {
            return Cow::Borrowed(self.template);
        }

        Cow::Owned(self.build())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::ParameterInfo;

    #[test]
    fn test_optimized_parameter_substitution() {
        let mut arg_map = HashMap::new();
        arg_map.insert("name".to_string(), "John".to_string());
        arg_map.insert("age".to_string(), "30".to_string());

        // Template with substitutions
        let result1 = optimize_parameter_substitution("Hello {name}, age {age}", &arg_map);
        assert!(matches!(result1, Cow::Owned(_)));
        assert_eq!(result1, "Hello John, age 30");

        // Template without substitutions
        let result2 = optimize_parameter_substitution("Hello World", &arg_map);
        assert!(matches!(result2, Cow::Borrowed(_)));
        assert_eq!(result2, "Hello World");
    }

    #[test]
    fn test_extract_source_module() {
        assert_eq!(extract_source_module("module.Class.method"), Some("module"));
        assert_eq!(extract_source_module("simple_func"), None);
        assert_eq!(extract_source_module("a.b.c.d"), Some("a"));
    }

    #[test]
    fn test_strip_module_prefix() {
        let result1 = strip_module_prefix("module.function()", Some("module"), "module");
        assert_eq!(result1, "function()");

        let result2 =
            strip_module_prefix("other_module.function()", Some("module"), "current_module");
        assert_eq!(result2, "other_module.function()");
    }

    #[test]
    fn test_optimized_replacement_builder() {
        let mut builder = OptimizedReplacementBuilder::new("Hello {name}, age {age}");
        builder.add_substitution("name", "John".to_string());
        builder.add_substitution("age", "30".to_string());

        let result = builder.build_cow();
        assert_eq!(result, "Hello John, age 30");
    }

    #[test]
    fn test_replacement_builder_no_placeholders() {
        let builder = OptimizedReplacementBuilder::new("Hello World");
        let result = builder.build_cow();

        // Should return borrowed since no placeholders
        assert!(matches!(result, Cow::Borrowed(_)));
        assert_eq!(result, "Hello World");
    }
}