dissolve_python/
optimized_visitor.rs

1// Copyright (C) 2024 Jelmer Vernooij <jelmer@samba.org>
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//    http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Optimized visitor implementation that reduces memory allocations and clones
16
17use crate::core::ReplaceInfo;
18// Removed unused imports
19use rustpython_ast as ast;
20use std::borrow::Cow;
21use std::collections::HashMap;
22
23/// Optimized parameter mapper that reduces string clones
24pub struct OptimizedParameterMapper<'a> {
25    /// String cache for parameter names (these are often repeated)
26    param_cache: &'a mut HashMap<String, String>,
27}
28
29impl<'a> OptimizedParameterMapper<'a> {
30    pub fn new(param_cache: &'a mut HashMap<String, String>) -> Self {
31        Self { param_cache }
32    }
33
34    /// Map function call arguments to parameters with minimal cloning
35    pub fn map_call_arguments(
36        &mut self,
37        args: &[ast::Expr],
38        keywords: &[ast::Keyword],
39        replace_info: &ReplaceInfo,
40        func: &ast::Expr,
41        expr_to_string: impl Fn(&ast::Expr) -> String,
42    ) -> HashMap<String, String> {
43        let mut arg_map = HashMap::new();
44
45        self.map_positional_args(&mut arg_map, args, replace_info, func, &expr_to_string);
46        self.map_keyword_args(&mut arg_map, keywords, &expr_to_string);
47
48        arg_map
49    }
50
51    /// Map positional arguments with optimized string handling
52    fn map_positional_args(
53        &mut self,
54        arg_map: &mut HashMap<String, String>,
55        args: &[ast::Expr],
56        replace_info: &ReplaceInfo,
57        func: &ast::Expr,
58        expr_to_string: &impl Fn(&ast::Expr) -> String,
59    ) {
60        let is_method_call = matches!(func, ast::Expr::Attribute(_));
61
62        if is_method_call {
63            self.map_method_call_args(arg_map, args, replace_info, func, expr_to_string);
64        } else {
65            self.map_function_call_args(arg_map, args, replace_info, expr_to_string);
66        }
67    }
68
69    /// Map arguments for method calls (handle 'self' parameter)
70    fn map_method_call_args(
71        &mut self,
72        arg_map: &mut HashMap<String, String>,
73        args: &[ast::Expr],
74        replace_info: &ReplaceInfo,
75        func: &ast::Expr,
76        expr_to_string: &impl Fn(&ast::Expr) -> String,
77    ) {
78        if let ast::Expr::Attribute(attr_expr) = func {
79            let self_str = expr_to_string(&attr_expr.value);
80
81            // Use cached string for 'self' if possible
82            let self_key = self.get_cached_string("self");
83            arg_map.insert(self_key, self_str);
84
85            // Map remaining arguments (skip 'self' parameter)
86            let method_params: Vec<_> = replace_info
87                .parameters
88                .iter()
89                .filter(|p| p.name != "self")
90                .collect();
91
92            for (i, arg) in args.iter().enumerate() {
93                if let Some(param) = method_params.get(i) {
94                    let arg_str = expr_to_string(arg);
95                    let param_key = self.get_cached_string(&param.name);
96                    arg_map.insert(param_key, arg_str);
97                }
98            }
99        }
100    }
101
102    /// Map arguments for regular function calls
103    fn map_function_call_args(
104        &mut self,
105        arg_map: &mut HashMap<String, String>,
106        args: &[ast::Expr],
107        replace_info: &ReplaceInfo,
108        expr_to_string: &impl Fn(&ast::Expr) -> String,
109    ) {
110        for (i, arg) in args.iter().enumerate() {
111            if let Some(param) = replace_info.parameters.get(i) {
112                let arg_str = expr_to_string(arg);
113                let param_key = self.get_cached_string(&param.name);
114                arg_map.insert(param_key, arg_str);
115            }
116        }
117    }
118
119    /// Map keyword arguments
120    fn map_keyword_args(
121        &mut self,
122        arg_map: &mut HashMap<String, String>,
123        keywords: &[ast::Keyword],
124        expr_to_string: &impl Fn(&ast::Expr) -> String,
125    ) {
126        for keyword in keywords {
127            if let Some(arg_name) = &keyword.arg {
128                let arg_str = expr_to_string(&keyword.value);
129                let arg_key = self.get_cached_string(arg_name);
130                arg_map.insert(arg_key, arg_str);
131            }
132        }
133    }
134
135    /// Get a cached string to avoid repeated allocations
136    fn get_cached_string(&mut self, s: &str) -> String {
137        if let Some(cached) = self.param_cache.get(s) {
138            cached.clone()
139        } else {
140            let owned = s.to_string();
141            self.param_cache.insert(owned.clone(), owned.clone());
142            owned
143        }
144    }
145}
146
147/// Optimized string replacement that uses Cow<str> to minimize allocations
148pub fn optimize_parameter_substitution<'a>(
149    template: &'a str,
150    arg_map: &HashMap<String, String>,
151) -> Cow<'a, str> {
152    // First pass: check if any substitutions are needed
153    let needs_substitution = arg_map
154        .keys()
155        .any(|param_name| template.contains(&format!("{{{}}}", param_name)));
156
157    if !needs_substitution {
158        return Cow::Borrowed(template);
159    }
160
161    // Perform substitutions on owned string
162    let mut result = template.to_string();
163    for (param_name, arg_value) in arg_map {
164        let placeholder = format!("{{{}}}", param_name);
165        result = result.replace(&placeholder, arg_value);
166    }
167
168    Cow::Owned(result)
169}
170
171/// Optimized module name extraction that avoids clones
172pub fn extract_source_module(old_name: &str) -> Option<&str> {
173    old_name
174        .find('.')
175        .map(|first_dot_pos| &old_name[..first_dot_pos])
176}
177
178/// Optimized module prefix stripping
179pub fn strip_module_prefix<'a>(
180    replacement: &'a str,
181    source_module: Option<&str>,
182    current_module: &str,
183) -> Cow<'a, str> {
184    match source_module {
185        Some(src_mod) if src_mod == current_module => {
186            let prefix = format!("{}.", src_mod);
187            if replacement.contains(&prefix) {
188                Cow::Owned(replacement.replace(&prefix, ""))
189            } else {
190                Cow::Borrowed(replacement)
191            }
192        }
193        _ => Cow::Borrowed(replacement),
194    }
195}
196
197/// Optimized replacement building that minimizes string allocations
198pub struct OptimizedReplacementBuilder<'a> {
199    template: &'a str,
200    substitutions: HashMap<&'a str, String>,
201}
202
203impl<'a> OptimizedReplacementBuilder<'a> {
204    pub fn new(template: &'a str) -> Self {
205        Self {
206            template,
207            substitutions: HashMap::new(),
208        }
209    }
210
211    /// Add a substitution using borrowed key to avoid allocation
212    pub fn add_substitution(&mut self, key: &'a str, value: String) {
213        self.substitutions.insert(key, value);
214    }
215
216    /// Build the final replacement string with minimal allocations
217    pub fn build(self) -> String {
218        let mut result = self.template.to_string();
219
220        for (key, value) in self.substitutions {
221            let placeholder = format!("{{{}}}", key);
222            result = result.replace(&placeholder, &value);
223        }
224
225        result
226    }
227
228    /// Build with cow optimization - returns borrowed if no substitutions needed
229    pub fn build_cow(self) -> Cow<'a, str> {
230        if self.substitutions.is_empty() {
231            return Cow::Borrowed(self.template);
232        }
233
234        // Check if any placeholders exist
235        let has_placeholders = self
236            .substitutions
237            .keys()
238            .any(|key| self.template.contains(&format!("{{{}}}", key)));
239
240        if !has_placeholders {
241            return Cow::Borrowed(self.template);
242        }
243
244        Cow::Owned(self.build())
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use crate::core::ParameterInfo;
252
253    #[test]
254    fn test_optimized_parameter_substitution() {
255        let mut arg_map = HashMap::new();
256        arg_map.insert("name".to_string(), "John".to_string());
257        arg_map.insert("age".to_string(), "30".to_string());
258
259        // Template with substitutions
260        let result1 = optimize_parameter_substitution("Hello {name}, age {age}", &arg_map);
261        assert!(matches!(result1, Cow::Owned(_)));
262        assert_eq!(result1, "Hello John, age 30");
263
264        // Template without substitutions
265        let result2 = optimize_parameter_substitution("Hello World", &arg_map);
266        assert!(matches!(result2, Cow::Borrowed(_)));
267        assert_eq!(result2, "Hello World");
268    }
269
270    #[test]
271    fn test_extract_source_module() {
272        assert_eq!(extract_source_module("module.Class.method"), Some("module"));
273        assert_eq!(extract_source_module("simple_func"), None);
274        assert_eq!(extract_source_module("a.b.c.d"), Some("a"));
275    }
276
277    #[test]
278    fn test_strip_module_prefix() {
279        let result1 = strip_module_prefix("module.function()", Some("module"), "module");
280        assert_eq!(result1, "function()");
281
282        let result2 =
283            strip_module_prefix("other_module.function()", Some("module"), "current_module");
284        assert_eq!(result2, "other_module.function()");
285    }
286
287    #[test]
288    fn test_optimized_replacement_builder() {
289        let mut builder = OptimizedReplacementBuilder::new("Hello {name}, age {age}");
290        builder.add_substitution("name", "John".to_string());
291        builder.add_substitution("age", "30".to_string());
292
293        let result = builder.build_cow();
294        assert_eq!(result, "Hello John, age 30");
295    }
296
297    #[test]
298    fn test_replacement_builder_no_placeholders() {
299        let builder = OptimizedReplacementBuilder::new("Hello World");
300        let result = builder.build_cow();
301
302        // Should return borrowed since no placeholders
303        assert!(matches!(result, Cow::Borrowed(_)));
304        assert_eq!(result, "Hello World");
305    }
306}
dissolve_python/optimized_visitor.rs

dissolve_python/
optimized_visitor.rs
