Skip to main content

mz_regexp/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use mz_repr::adt::regex::Regex;
11
12pub fn regexp_split_to_array<'a>(text: &'a str, regexp: &Regex) -> Vec<&'a str> {
13    // Postgres regex split handling differs a bit from spec regex split, so we can't use
14    // regexp.split here. See: https://www.postgresql.org/docs/15/functions-matching.html:
15    // > the regexp split functions ignore zero-length matches that occur at the start or end
16    // > of the string or immediately after a previous match
17
18    let mut finder = regexp.find_iter(text);
19    let mut last = 0;
20    let mut found = Vec::new();
21    loop {
22        match finder.next() {
23            None => {
24                if last <= text.len() {
25                    let s = &text[last..];
26                    found.push(s);
27                }
28                break;
29            }
30            Some(m) => {
31                // Ignore zero length matches at start and end of string.
32                if m.end() > 0 && m.start() < text.len() {
33                    let matched = &text[last..m.start()];
34                    last = m.end();
35                    found.push(matched);
36                }
37            }
38        }
39    }
40    found
41}
42
43/// How many chunks [`regexp_split_to_array`] would return, without building them.
44///
45/// One chunk per kept match plus the trailing remainder, keeping a match exactly when it is not
46/// zero-length at the start or end of the string. This must mirror the split's rule so a caller
47/// sizing an allocation refuses only what the split itself would refuse.
48/// `tests::test_regexp_split_array_count` pins the two together.
49pub fn regexp_split_to_array_count(text: &str, regexp: &Regex) -> usize {
50    1 + regexp
51        .find_iter(text)
52        .filter(|m| m.end() > 0 && m.start() < text.len())
53        .count()
54}
55
56#[cfg(test)]
57mod tests {
58    use mz_repr::adt::regex::Regex;
59
60    use crate::{regexp_split_to_array, regexp_split_to_array_count};
61
62    fn build_regex(needle: &str, flags: &str) -> Result<Regex, anyhow::Error> {
63        let mut case_insensitive = false;
64        // Note: Postgres accepts it when both flags are present, taking the last one. We do the same.
65        for f in flags.chars() {
66            match f {
67                'i' => {
68                    case_insensitive = true;
69                }
70                'c' => {
71                    case_insensitive = false;
72                }
73                _ => anyhow::bail!("unexpected regex flags"),
74            }
75        }
76        Regex::new(needle, case_insensitive).map_err(|e| anyhow::anyhow!("{}", e))
77    }
78
79    // Assert equivalency to postgres and generate TestCases.
80    #[mz_ore::test]
81    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
82    fn test_pg_regexp_split_array() {
83        let Ok(postgres_url) = std::env::var("POSTGRES_URL") else {
84            return;
85        };
86        let mut client = postgres::Client::connect(&postgres_url, postgres::NoTls).unwrap();
87
88        let inputs = vec!["", " ", "  ", "12 34", "12  34", " 12 34 "];
89        let regexps = vec!["", "\\s", "\\s+", "\\s*"];
90        for input in inputs {
91            for re in &regexps {
92                let regex = build_regex(re, "").unwrap();
93                // This test cross-checks against the sync `postgres` crate,
94                // while `mz_postgres_util` wrappers target async tokio-postgres.
95                let pg: Vec<String> = client
96                    .query_one("select regexp_split_to_array($1, $2)", &[&input, re])
97                    .unwrap()
98                    .get(0);
99                let mz = regexp_split_to_array(input, &regex);
100                assert_eq!(pg, mz);
101                // Generate TestCases for static use.
102                println!(
103                    r#"TestCase {{
104                text: "{input}",
105                regexp: "{}",
106                expect: &{pg:?},
107            }},"#,
108                    re.replace('\\', "\\\\"),
109                );
110            }
111        }
112    }
113
114    #[mz_ore::test]
115    #[cfg_attr(miri, ignore)] // too slow
116    fn test_regexp_split_array() {
117        // Expected outputs generated from postgres.
118        struct TestCase {
119            text: &'static str,
120            regexp: &'static str,
121            expect: &'static [&'static str],
122        }
123        let tests = vec![
124            TestCase {
125                text: "",
126                regexp: "",
127                expect: &[""],
128            },
129            TestCase {
130                text: "",
131                regexp: "\\s",
132                expect: &[""],
133            },
134            TestCase {
135                text: "",
136                regexp: "\\s+",
137                expect: &[""],
138            },
139            TestCase {
140                text: "",
141                regexp: "\\s*",
142                expect: &[""],
143            },
144            TestCase {
145                text: " ",
146                regexp: "",
147                expect: &[" "],
148            },
149            TestCase {
150                text: " ",
151                regexp: "\\s",
152                expect: &["", ""],
153            },
154            TestCase {
155                text: " ",
156                regexp: "\\s+",
157                expect: &["", ""],
158            },
159            TestCase {
160                text: " ",
161                regexp: "\\s*",
162                expect: &["", ""],
163            },
164            TestCase {
165                text: "  ",
166                regexp: "",
167                expect: &[" ", " "],
168            },
169            TestCase {
170                text: "  ",
171                regexp: "\\s",
172                expect: &["", "", ""],
173            },
174            TestCase {
175                text: "  ",
176                regexp: "\\s+",
177                expect: &["", ""],
178            },
179            TestCase {
180                text: "  ",
181                regexp: "\\s*",
182                expect: &["", ""],
183            },
184            TestCase {
185                text: "12 34",
186                regexp: "",
187                expect: &["1", "2", " ", "3", "4"],
188            },
189            TestCase {
190                text: "12 34",
191                regexp: "\\s",
192                expect: &["12", "34"],
193            },
194            TestCase {
195                text: "12 34",
196                regexp: "\\s+",
197                expect: &["12", "34"],
198            },
199            TestCase {
200                text: "12 34",
201                regexp: "\\s*",
202                expect: &["1", "2", "3", "4"],
203            },
204            TestCase {
205                text: "12  34",
206                regexp: "",
207                expect: &["1", "2", " ", " ", "3", "4"],
208            },
209            TestCase {
210                text: "12  34",
211                regexp: "\\s",
212                expect: &["12", "", "34"],
213            },
214            TestCase {
215                text: "12  34",
216                regexp: "\\s+",
217                expect: &["12", "34"],
218            },
219            TestCase {
220                text: "12  34",
221                regexp: "\\s*",
222                expect: &["1", "2", "3", "4"],
223            },
224            TestCase {
225                text: " 12 34 ",
226                regexp: "",
227                expect: &[" ", "1", "2", " ", "3", "4", " "],
228            },
229            TestCase {
230                text: " 12 34 ",
231                regexp: "\\s",
232                expect: &["", "12", "34", ""],
233            },
234            TestCase {
235                text: " 12 34 ",
236                regexp: "\\s+",
237                expect: &["", "12", "34", ""],
238            },
239            TestCase {
240                text: " 12 34 ",
241                regexp: "\\s*",
242                expect: &["", "1", "2", "3", "4", ""],
243            },
244        ];
245        for tc in tests {
246            let regex = build_regex(tc.regexp, "").unwrap();
247            let result = regexp_split_to_array(tc.text, &regex);
248            if tc.expect != result {
249                println!(
250                    "input: `{}`, regex: `{}`, got: {:?}, expect: {:?}",
251                    tc.text, tc.regexp, result, tc.expect
252                );
253            }
254        }
255    }
256
257    /// The count must agree with the split for every input, so a caller sizing an allocation from
258    /// it refuses exactly what the split would have built. Covers no match, matches at each
259    /// boundary, and empty-string matches, where the naive "matches + 1" is wrong.
260    #[mz_ore::test]
261    #[cfg_attr(miri, ignore)] // too slow
262    fn test_regexp_split_array_count() {
263        let texts = ["", " ", "  ", "abc", "aaa", "12 34", " 12 34 "];
264        let regexps = ["", "\\s", "\\s+", "\\s*", "a", "a*", "b*", "x"];
265        for text in texts {
266            for re in regexps {
267                let regex = build_regex(re, "").unwrap();
268                assert_eq!(
269                    regexp_split_to_array_count(text, &regex),
270                    regexp_split_to_array(text, &regex).len(),
271                    "text: `{text}`, regexp: `{re}`",
272                );
273            }
274        }
275    }
276}