polars_utils/
regex_cache.rs

1use std::cell::RefCell;
2use std::sync::Arc;
3
4use regex::bytes::Regex as BytesRegex;
5use regex::{Regex, RegexBuilder};
6use regex_syntax::ast::{self, Ast, RepetitionKind, RepetitionRange};
7use regex_syntax::hir::translate::Translator;
8
9use crate::cache::LruCache;
10
11fn get_size_limit() -> Option<usize> {
12    Some(
13        std::env::var("POLARS_REGEX_SIZE_LIMIT")
14            .ok()
15            .filter(|l| !l.is_empty())?
16            .parse()
17            .expect("invalid POLARS_REGEX_SIZE_LIMIT"),
18    )
19}
20
21// Regex compilation is really heavy, and the resulting regexes can be large as
22// well, so we should have a good caching scheme.
23//
24// TODO: add larger global cache which has time-based flush.
25
26/// A cache for compiled regular expressions.
27pub struct RegexCache {
28    cache: LruCache<String, Regex>,
29    is_match_cache: LruCache<String, Regex>,
30    bytes_cache: LruCache<usize, (Arc<BytesRegex>, BytesRegex)>,
31    size_limit: Option<usize>,
32}
33
34impl RegexCache {
35    fn new() -> Self {
36        Self {
37            cache: LruCache::with_capacity(32),
38            is_match_cache: LruCache::with_capacity(32),
39            bytes_cache: LruCache::with_capacity(32),
40            size_limit: get_size_limit(),
41        }
42    }
43
44    pub fn compile(&mut self, re: &str) -> Result<&Regex, regex::Error> {
45        compile_cached(&mut self.cache, &mut self.size_limit, re)
46    }
47
48    /// Like [`Self::compile`], for `is_match` only, see [`is_match_pattern`].
49    pub fn compile_is_match(&mut self, re: &str) -> Result<&Regex, regex::Error> {
50        let Self {
51            cache,
52            is_match_cache,
53            size_limit,
54            ..
55        } = self;
56        let r = is_match_cache.try_get_or_insert_with(re, |re| match is_match_pattern(re) {
57            Some(trimmed) => build_regex(size_limit, &trimmed),
58            None => compile_cached(cache, size_limit, re).cloned(),
59        });
60        Ok(&*r?)
61    }
62
63    /// Borrows a cached clone of the supplied regex, keyed by its Arc pointer.
64    pub fn get_or_insert_bytes(&mut self, re: &Arc<BytesRegex>) -> &BytesRegex {
65        // Retain the original Arc so its address cannot be reused while cached.
66        &self
67            .bytes_cache
68            .get_or_insert_with(&(Arc::as_ptr(re) as usize), |_| {
69                (re.clone(), re.as_ref().clone())
70            })
71            .1
72    }
73}
74
75fn compile_cached<'a>(
76    cache: &'a mut LruCache<String, Regex>,
77    size_limit: &mut Option<usize>,
78    re: &str,
79) -> Result<&'a Regex, regex::Error> {
80    let r = cache.try_get_or_insert_with(re, |re| build_regex(size_limit, re));
81    Ok(&*r?)
82}
83
84fn build_regex(size_limit: &mut Option<usize>, re: &str) -> Result<Regex, regex::Error> {
85    build_within_size_limit(size_limit, |limit| {
86        let mut builder = RegexBuilder::new(re);
87        if let Some(bytes) = limit {
88            builder.size_limit(bytes);
89        }
90        builder.build()
91    })
92}
93
94/// `pat` without the repetitions at either end that can match nothing (`.*foo.*` -> `foo`): same
95/// `is_match`, but since regex 1.13.1 a leading `.*` disables its reverse suffix/inner search.
96pub fn is_match_pattern(pat: &str) -> Option<String> {
97    let ast = ast::parse::Parser::new().parse(pat).ok()?;
98    let Ast::Concat(concat) = &ast else {
99        return None;
100    };
101    let items = &concat.asts;
102    let is_flags = |a: &Ast| matches!(a, Ast::Flags(_));
103    let can_match_empty = |a: &Ast| match a {
104        Ast::Repetition(rep) => match rep.op.kind {
105            RepetitionKind::ZeroOrOne | RepetitionKind::ZeroOrMore => true,
106            RepetitionKind::OneOrMore => false,
107            RepetitionKind::Range(
108                RepetitionRange::Exactly(min)
109                | RepetitionRange::AtLeast(min)
110                | RepetitionRange::Bounded(min, _),
111            ) => min == 0,
112        },
113        _ => false,
114    };
115    let kept = |a: &Ast| !is_flags(a) && !can_match_empty(a);
116    let first = items.iter().position(kept)?;
117    let last = items.iter().rposition(kept)?;
118    if items[..first]
119        .iter()
120        .chain(&items[last + 1..])
121        .all(is_flags)
122    {
123        return None;
124    }
125    // Only rewrite patterns that compile, so errors stay the same.
126    Translator::new().translate(pat, &ast).ok()?;
127    // Slice `pat` rather than print a `Hir`, which doesn't round-trip nested repetitions.
128    let span = |a: &Ast| a.span().start.offset..a.span().end.offset;
129    let mut out = String::new();
130    // Flags apply to what follows them.
131    for a in items[..first].iter().filter(|a| is_flags(a)) {
132        out.push_str(&pat[span(a)]);
133    }
134    out.push_str(&pat[span(&items[first]).start..span(&items[last]).end]);
135    Some(out)
136}
137
138// We do this little loop to only check POLARS_REGEX_SIZE_LIMIT when a regex
139// fails to compile due to the size limit.
140fn build_within_size_limit<R>(
141    size_limit: &mut Option<usize>,
142    build: impl Fn(Option<usize>) -> Result<R, regex::Error>,
143) -> Result<R, regex::Error> {
144    loop {
145        match build(*size_limit) {
146            err @ Err(regex::Error::CompiledTooBig(_)) => {
147                let new_size_limit = get_size_limit();
148                if new_size_limit != *size_limit {
149                    *size_limit = new_size_limit;
150                    continue; // Try to compile again.
151                }
152                break err;
153            },
154            r => break r,
155        }
156    }
157}
158
159thread_local! {
160    static LOCAL_REGEX_CACHE: RefCell<RegexCache> = RefCell::new(RegexCache::new());
161}
162
163pub fn compile_regex(re: &str) -> Result<Regex, regex::Error> {
164    LOCAL_REGEX_CACHE.with_borrow_mut(|cache| cache.compile(re).cloned())
165}
166
167/// Like [`compile_regex`], for `is_match` only, see [`is_match_pattern`].
168pub fn compile_is_match_regex(re: &str) -> Result<Regex, regex::Error> {
169    LOCAL_REGEX_CACHE.with_borrow_mut(|cache| cache.compile_is_match(re).cloned())
170}
171
172pub fn with_regex_cache<R, F: FnOnce(&mut RegexCache) -> R>(f: F) -> R {
173    LOCAL_REGEX_CACHE.with_borrow_mut(f)
174}
175
176#[macro_export]
177macro_rules! cached_regex {
178    () => {};
179
180    ($vis:vis static $name:ident = $regex:expr; $($rest:tt)*) => {
181        #[allow(clippy::disallowed_methods)]
182        $vis static $name: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| regex::Regex::new($regex).unwrap());
183        $crate::regex_cache::cached_regex!($($rest)*);
184    };
185}
186pub use cached_regex;
187
188#[cfg(test)]
189mod tests {
190    use regex::bytes::RegexBuilder as BytesRegexBuilder;
191
192    use super::*;
193
194    #[test]
195    fn bytes_cache_preserves_builder_options() {
196        let mut cache = RegexCache::new();
197        let sensitive = Arc::new(BytesRegexBuilder::new("abc").build().unwrap());
198        let insensitive = Arc::new(
199            BytesRegexBuilder::new("abc")
200                .case_insensitive(true)
201                .build()
202                .unwrap(),
203        );
204
205        for _ in 0..2 {
206            assert!(!cache.get_or_insert_bytes(&sensitive).is_match(b"ABC"));
207            assert!(cache.get_or_insert_bytes(&insensitive).is_match(b"ABC"));
208        }
209    }
210
211    #[test]
212    fn bytes_cache_reuses_cloned_arc() {
213        let mut cache = RegexCache::new();
214        let re = Arc::new(BytesRegexBuilder::new("abc").build().unwrap());
215        let cloned = re.clone();
216        let cached = std::ptr::from_ref(cache.get_or_insert_bytes(&re));
217
218        assert_eq!(
219            cached,
220            std::ptr::from_ref(cache.get_or_insert_bytes(&cloned))
221        );
222        assert_ne!(cached, Arc::as_ptr(&re));
223    }
224
225    #[test]
226    fn bytes_cache_retains_arc_until_eviction() {
227        let mut cache = RegexCache::new();
228        let regex = BytesRegexBuilder::new("abc").build().unwrap();
229        let re = Arc::new(regex.clone());
230        let weak = Arc::downgrade(&re);
231        cache.get_or_insert_bytes(&re);
232        drop(re);
233        assert!(weak.upgrade().is_some());
234
235        for _ in 0..32 {
236            let re = Arc::new(regex.clone());
237            cache.get_or_insert_bytes(&re);
238        }
239        assert!(weak.upgrade().is_none());
240    }
241}
polars_utils/regex_cache.rs

polars_utils/
regex_cache.rs
