Skip to main content

linkrs_query/executor/expression/functions/builtin/
regex.rs

1//! Implementation of regular expression functions
2
3use crate::executor::expression::{ExpressionError, ExpressionErrorType};
4use linkrs_core::value::NullType;
5use linkrs_core::Value;
6
7/// Enumeration of regular expression functions
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum RegexFunction {
10    RegexMatch,
11    RegexReplace,
12    RegexFind,
13    RegexpFullMatch,
14    RegexpExtract,
15    RegexpExtractAll,
16    RegexpSplitToArray,
17}
18
19impl RegexFunction {
20    pub fn name(&self) -> &str {
21        match self {
22            RegexFunction::RegexMatch => "regex_match",
23            RegexFunction::RegexReplace => "regex_replace",
24            RegexFunction::RegexFind => "regex_find",
25            RegexFunction::RegexpFullMatch => "regexp_full_match",
26            RegexFunction::RegexpExtract => "regexp_extract",
27            RegexFunction::RegexpExtractAll => "regexp_extract_all",
28            RegexFunction::RegexpSplitToArray => "regexp_split_to_array",
29        }
30    }
31
32    pub fn arity(&self) -> usize {
33        match self {
34            RegexFunction::RegexMatch => 2,
35            RegexFunction::RegexReplace => 3,
36            RegexFunction::RegexFind => 2,
37            RegexFunction::RegexpFullMatch => 2,
38            RegexFunction::RegexpExtract => 2,
39            RegexFunction::RegexpExtractAll => 2,
40            RegexFunction::RegexpSplitToArray => 2,
41        }
42    }
43
44    pub fn is_variadic(&self) -> bool {
45        false
46    }
47
48    pub fn description(&self) -> &str {
49        match self {
50            RegexFunction::RegexMatch => "regular expression matching (math.)",
51            RegexFunction::RegexReplace => "regular expression substitution",
52            RegexFunction::RegexFind => "regular expression lookup (computing)",
53            RegexFunction::RegexpFullMatch => "Check if string fully matches regex pattern",
54            RegexFunction::RegexpExtract => "Extract first match of regex pattern",
55            RegexFunction::RegexpExtractAll => "Extract all matches of regex pattern",
56            RegexFunction::RegexpSplitToArray => "Split string by regex pattern into array",
57        }
58    }
59
60    pub fn execute(&self, args: &[Value]) -> Result<Value, ExpressionError> {
61        if !self.is_variadic() && args.len() != self.arity() {
62            return Err(ExpressionError::invalid_arity(
63                self.name(),
64                self.arity(),
65                args.len(),
66            ));
67        }
68        match self {
69            RegexFunction::RegexMatch => match (&args[0], &args[1]) {
70                (Value::String(s), Value::String(pattern)) => {
71                    let regex = regex::Regex::new(pattern).map_err(|_| {
72                        ExpressionError::new(
73                            ExpressionErrorType::InvalidOperation,
74                            format!("Invalid regular expression: {}", pattern),
75                        )
76                    })?;
77                    Ok(Value::Bool(regex.is_match(s)))
78                }
79                (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
80                _ => Err(ExpressionError::type_error(
81                    "The regex_match function requires the string type",
82                )),
83            },
84            RegexFunction::RegexReplace => match (&args[0], &args[1], &args[2]) {
85                (Value::String(s), Value::String(pattern), Value::String(replacement)) => {
86                    let regex = regex::Regex::new(pattern).map_err(|_| {
87                        ExpressionError::new(
88                            ExpressionErrorType::InvalidOperation,
89                            format!("Invalid regular expression: {}", pattern),
90                        )
91                    })?;
92                    Ok(Value::string(regex.replace_all(s, replacement.as_str())))
93                }
94                (Value::Null(_), _, _) | (_, Value::Null(_), _) | (_, _, Value::Null(_)) => {
95                    Ok(Value::Null(NullType::Null))
96                }
97                _ => Err(ExpressionError::type_error(
98                    "The regex_replace function requires the string type",
99                )),
100            },
101            RegexFunction::RegexFind => match (&args[0], &args[1]) {
102                (Value::String(s), Value::String(pattern)) => {
103                    let regex = regex::Regex::new(pattern).map_err(|_| {
104                        ExpressionError::new(
105                            ExpressionErrorType::InvalidOperation,
106                            format!("Invalid regular expression: {}", pattern),
107                        )
108                    })?;
109                    if let Some(matched) = regex.find(s) {
110                        Ok(Value::string(matched.as_str()))
111                    } else {
112                        Ok(Value::Null(NullType::Null))
113                    }
114                }
115                (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
116                _ => Err(ExpressionError::type_error(
117                    "The regex_find function requires the string type",
118                )),
119            },
120            RegexFunction::RegexpFullMatch => match (&args[0], &args[1]) {
121                (Value::String(s), Value::String(pattern)) => {
122                    let regex = regex::Regex::new(pattern).map_err(|_| {
123                        ExpressionError::new(
124                            ExpressionErrorType::InvalidOperation,
125                            format!("Invalid regular expression: {}", pattern),
126                        )
127                    })?;
128                    Ok(Value::Bool(
129                        regex.is_match(s) && regex.find(s).is_some_and(|m| m.as_str() == s),
130                    ))
131                }
132                (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
133                _ => Err(ExpressionError::type_error(
134                    "regexp_full_match requires string arguments",
135                )),
136            },
137            RegexFunction::RegexpExtract => match (&args[0], &args[1]) {
138                (Value::String(s), Value::String(pattern)) => {
139                    let regex = regex::Regex::new(pattern).map_err(|_| {
140                        ExpressionError::new(
141                            ExpressionErrorType::InvalidOperation,
142                            format!("Invalid regular expression: {}", pattern),
143                        )
144                    })?;
145                    if let Some(matched) = regex.find(s) {
146                        Ok(Value::string(matched.as_str()))
147                    } else {
148                        Ok(Value::Null(NullType::Null))
149                    }
150                }
151                (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
152                _ => Err(ExpressionError::type_error(
153                    "regexp_extract requires string arguments",
154                )),
155            },
156            RegexFunction::RegexpExtractAll => match (&args[0], &args[1]) {
157                (Value::String(s), Value::String(pattern)) => {
158                    let regex = regex::Regex::new(pattern).map_err(|_| {
159                        ExpressionError::new(
160                            ExpressionErrorType::InvalidOperation,
161                            format!("Invalid regular expression: {}", pattern),
162                        )
163                    })?;
164                    let matches: Vec<Value> = regex
165                        .find_iter(s)
166                        .map(|m| Value::string(m.as_str()))
167                        .collect();
168                    Ok(Value::list(linkrs_core::value::list::List {
169                        values: matches,
170                    }))
171                }
172                (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
173                _ => Err(ExpressionError::type_error(
174                    "regexp_extract_all requires string arguments",
175                )),
176            },
177            RegexFunction::RegexpSplitToArray => match (&args[0], &args[1]) {
178                (Value::String(s), Value::String(pattern)) => {
179                    let regex = regex::Regex::new(pattern).map_err(|_| {
180                        ExpressionError::new(
181                            ExpressionErrorType::InvalidOperation,
182                            format!("Invalid regular expression: {}", pattern),
183                        )
184                    })?;
185                    let parts: Vec<Value> = regex.split(s).map(Value::string).collect();
186                    Ok(Value::list(linkrs_core::value::list::List {
187                        values: parts,
188                    }))
189                }
190                (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
191                _ => Err(ExpressionError::type_error(
192                    "regexp_split_to_array requires string arguments",
193                )),
194            },
195        }
196    }
197
198    /// Execute a function (with caching)
199    ///
200    /// The caching function has been removed; the `execute` method is called directly.
201    pub fn execute_with_cache(
202        &self,
203        args: &[Value],
204        _cache: &mut (),
205    ) -> Result<Value, ExpressionError> {
206        self.execute(args)
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn test_regexp_full_match() {
216        let func = RegexFunction::RegexpFullMatch;
217        let result = func
218            .execute(&[Value::string("hello123"), Value::string(r"^[a-z]+\d+$")])
219            .expect("Execution should succeed");
220        assert_eq!(result, Value::Bool(true));
221    }
222
223    #[test]
224    fn test_regexp_full_match_no_match() {
225        let func = RegexFunction::RegexpFullMatch;
226        let result = func
227            .execute(&[Value::string("hello world"), Value::string(r"^\d+$")])
228            .expect("Execution should succeed");
229        assert_eq!(result, Value::Bool(false));
230    }
231
232    #[test]
233    fn test_regexp_extract() {
234        let func = RegexFunction::RegexpExtract;
235        let result = func
236            .execute(&[Value::string("abc123def"), Value::string(r"\d+")])
237            .expect("Execution should succeed");
238        assert_eq!(result, Value::string("123"));
239    }
240
241    #[test]
242    fn test_regexp_extract_no_match() {
243        let func = RegexFunction::RegexpExtract;
244        let result = func
245            .execute(&[Value::string("abc"), Value::string(r"\d+")])
246            .expect("Execution should succeed");
247        assert_eq!(result, Value::Null(NullType::Null));
248    }
249
250    #[test]
251    fn test_regexp_extract_all() {
252        let func = RegexFunction::RegexpExtractAll;
253        let result = func
254            .execute(&[Value::string("abc123def456"), Value::string(r"\d+")])
255            .expect("Execution should succeed");
256        assert_eq!(
257            result,
258            Value::list(linkrs_core::value::list::List {
259                values: vec![Value::string("123"), Value::string("456")]
260            })
261        );
262    }
263
264    #[test]
265    fn test_regexp_split_to_array() {
266        let func = RegexFunction::RegexpSplitToArray;
267        let result = func
268            .execute(&[Value::string("abc123def456"), Value::string(r"\d+")])
269            .expect("Execution should succeed");
270        assert_eq!(
271            result,
272            Value::list(linkrs_core::value::list::List {
273                values: vec![
274                    Value::string("abc"),
275                    Value::string("def"),
276                    Value::string("")
277                ]
278            })
279        );
280    }
281}