Skip to main content

dissolve_python/
optimizations.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//! Optimizations and refactoring helpers to reduce allocations
16
17use std::borrow::Cow;
18use std::rc::Rc;
19use std::sync::Arc;
20
21/// String storage that can be either owned or borrowed
22/// This reduces clones when strings don't need to be modified
23pub type CowStr<'a> = Cow<'a, str>;
24
25/// Shared string that avoids cloning for read-only access
26pub type SharedString = Arc<str>;
27
28/// Convert common string operations to use Cow to avoid clones
29pub trait StringOptimizations {
30    /// Get a borrowed or owned string depending on whether modification is needed
31    fn as_cow(&self) -> CowStr;
32
33    /// Share a string across multiple owners without cloning
34    fn to_shared(&self) -> SharedString;
35}
36
37impl StringOptimizations for String {
38    fn as_cow(&self) -> CowStr {
39        Cow::Borrowed(self.as_str())
40    }
41
42    fn to_shared(&self) -> SharedString {
43        Arc::from(self.as_str())
44    }
45}
46
47impl StringOptimizations for &str {
48    fn as_cow(&self) -> CowStr {
49        Cow::Borrowed(self)
50    }
51
52    fn to_shared(&self) -> SharedString {
53        Arc::from(*self)
54    }
55}
56
57/// Cache for commonly used strings to avoid repeated allocations
58pub struct StringCache {
59    cache: std::collections::HashMap<String, Rc<str>>,
60}
61
62impl StringCache {
63    pub fn new() -> Self {
64        Self {
65            cache: std::collections::HashMap::new(),
66        }
67    }
68
69    /// Get or insert a string in the cache
70    pub fn get_or_insert(&mut self, s: &str) -> Rc<str> {
71        if let Some(cached) = self.cache.get(s) {
72            Rc::clone(cached)
73        } else {
74            let rc = Rc::from(s);
75            self.cache.insert(s.to_string(), Rc::clone(&rc));
76            rc
77        }
78    }
79
80    /// Get a string from the cache if it exists
81    pub fn get(&self, s: &str) -> Option<Rc<str>> {
82        self.cache.get(s).map(Rc::clone)
83    }
84}
85
86impl Default for StringCache {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92/// Optimized parameter mapping that reduces clones
93pub struct OptimizedParamMap<'a> {
94    data: std::collections::HashMap<&'a str, &'a str>,
95    owned: std::collections::HashMap<String, String>,
96}
97
98impl<'a> OptimizedParamMap<'a> {
99    pub fn new() -> Self {
100        Self {
101            data: std::collections::HashMap::new(),
102            owned: std::collections::HashMap::new(),
103        }
104    }
105
106    /// Insert a borrowed key-value pair
107    pub fn insert_borrowed(&mut self, key: &'a str, value: &'a str) {
108        self.data.insert(key, value);
109    }
110
111    /// Insert an owned key-value pair (only when necessary)
112    pub fn insert_owned(&mut self, key: String, value: String) {
113        self.owned.insert(key, value);
114    }
115
116    /// Get a value by key (checks both borrowed and owned maps)
117    pub fn get(&self, key: &str) -> Option<&str> {
118        self.data
119            .get(key)
120            .map(|&s| s)
121            .or_else(|| self.owned.get(key).map(String::as_str))
122    }
123
124    /// Check if a key exists
125    pub fn contains_key(&self, key: &str) -> bool {
126        self.data.contains_key(key) || self.owned.contains_key(key)
127    }
128
129    /// Iterate over all key-value pairs
130    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> + '_ {
131        self.data
132            .iter()
133            .map(|(&k, &v)| (k, v))
134            .chain(self.owned.iter().map(|(k, v)| (k.as_str(), v.as_str())))
135    }
136}
137
138impl<'a> Default for OptimizedParamMap<'a> {
139    fn default() -> Self {
140        Self::new()
141    }
142}
143
144/// Helper to avoid cloning in common patterns
145pub mod clone_reduction {
146    use super::*;
147
148    /// Process a string without cloning unless modification is needed
149    pub fn process_string<'a>(s: &'a str, needs_modification: bool) -> Cow<'a, str> {
150        if needs_modification {
151            // Only clone when we actually need to modify
152            let owned = s.to_string();
153            // Apply modifications here
154            Cow::Owned(owned)
155        } else {
156            Cow::Borrowed(s)
157        }
158    }
159
160    /// Share a collection element without cloning
161    pub fn share_element<T: Clone>(collection: &[T], index: usize) -> Option<Rc<T>> {
162        collection.get(index).map(|elem| Rc::new(elem.clone()))
163    }
164
165    /// Use string interning for frequently used strings
166    pub fn intern_string(s: &str, cache: &mut StringCache) -> Rc<str> {
167        cache.get_or_insert(s)
168    }
169}
170
171/// Optimized replacement for HashMap<String, String> that reduces allocations
172pub struct OptimizedStringMap {
173    // Use Arc<str> for keys and values to enable cheap cloning
174    data: std::collections::HashMap<Arc<str>, Arc<str>>,
175}
176
177impl OptimizedStringMap {
178    pub fn new() -> Self {
179        Self {
180            data: std::collections::HashMap::new(),
181        }
182    }
183
184    pub fn insert(&mut self, key: impl Into<Arc<str>>, value: impl Into<Arc<str>>) {
185        self.data.insert(key.into(), value.into());
186    }
187
188    pub fn get(&self, key: &str) -> Option<&str> {
189        self.data.get(key).map(|arc| arc.as_ref())
190    }
191
192    pub fn contains_key(&self, key: &str) -> bool {
193        self.data.contains_key(key)
194    }
195
196    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> + '_ {
197        self.data.iter().map(|(k, v)| (k.as_ref(), v.as_ref()))
198    }
199
200    /// Clone is cheap because Arc<str> is reference counted
201    pub fn cheap_clone(&self) -> Self {
202        Self {
203            data: self.data.clone(),
204        }
205    }
206}
207
208impl Default for OptimizedStringMap {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn test_string_cache() {
220        let mut cache = StringCache::new();
221
222        let s1 = cache.get_or_insert("hello");
223        let s2 = cache.get_or_insert("hello");
224
225        // Should return the same Rc
226        assert!(Rc::ptr_eq(&s1, &s2));
227    }
228
229    #[test]
230    fn test_optimized_param_map() {
231        let mut map = OptimizedParamMap::new();
232
233        map.insert_borrowed("key1", "value1");
234        map.insert_owned("key2".to_string(), "value2".to_string());
235
236        assert_eq!(map.get("key1"), Some("value1"));
237        assert_eq!(map.get("key2"), Some("value2"));
238        assert!(map.contains_key("key1"));
239        assert!(map.contains_key("key2"));
240    }
241
242    #[test]
243    fn test_cow_usage() {
244        let original = "hello";
245
246        // No modification needed - no allocation
247        let borrowed = clone_reduction::process_string(original, false);
248        assert!(matches!(borrowed, Cow::Borrowed(_)));
249
250        // Modification needed - allocation happens
251        let owned = clone_reduction::process_string(original, true);
252        assert!(matches!(owned, Cow::Owned(_)));
253    }
254
255    #[test]
256    fn test_optimized_string_map() {
257        let mut map = OptimizedStringMap::new();
258
259        map.insert("key1", "value1");
260        map.insert("key2", "value2");
261
262        assert_eq!(map.get("key1"), Some("value1"));
263        assert_eq!(map.get("key2"), Some("value2"));
264
265        // Cloning is cheap with Arc
266        let cloned = map.cheap_clone();
267        assert_eq!(cloned.get("key1"), Some("value1"));
268    }
269}