1use config::{any_err, get_or_create_sub_module};
2use fancy_regex::{Matches, Regex};
3use mlua::{Lua, UserData, UserDataMethods};
4
5struct RegexWrap(Regex);
6
7impl UserData for RegexWrap {
8 fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
9 methods.add_method("captures", |lua, this, haystack: String| {
10 match this.0.captures(&haystack).map_err(any_err)? {
11 Some(c) => {
12 let result = lua.create_table()?;
13
14 let names = this.0.capture_names();
15 for ((idx, cap), name) in c.iter().enumerate().zip(names) {
16 if let Some(cap) = cap {
17 let s = cap.as_str();
18 result.set(idx, s.to_string())?;
19 if let Some(name) = name {
20 result.set(name, s.to_string())?;
21 }
22 }
23 }
24
25 Ok(Some(result))
26 }
27 None => Ok(None),
28 }
29 });
30
31 methods.add_method("is_match", |_, this, haystack: String| {
32 this.0.is_match(&haystack).map_err(any_err)
33 });
34
35 methods.add_method("find", |_, this, haystack: String| {
36 Ok(this
37 .0
38 .find(&haystack)
39 .map_err(any_err)?
40 .map(|m| m.as_str().to_string()))
41 });
42
43 methods.add_method("find_all", |_, this, haystack: String| {
44 let mut result = vec![];
45
46 for m in this.0.find_iter(&haystack) {
47 let s = m.map_err(any_err)?;
48 result.push(s.as_str().to_string());
49 }
50 Ok(result)
51 });
52
53 methods.add_method("replace", |_, this, (haystack, rep): (String, String)| {
54 Ok(this.0.replace(&haystack, &rep).to_string())
55 });
56
57 methods.add_method(
58 "replace_all",
59 |_, this, (haystack, rep): (String, String)| {
60 Ok(this
61 .0
62 .try_replacen(&haystack, 0, &rep)
63 .map_err(any_err)?
64 .to_string())
65 },
66 );
67
68 methods.add_method(
69 "replacen",
70 |_, this, (haystack, limit, rep): (String, usize, String)| {
71 Ok(this
72 .0
73 .try_replacen(&haystack, limit, &rep)
74 .map_err(any_err)?
75 .to_string())
76 },
77 );
78
79 methods.add_method("split", |_, this, haystack: String| {
80 split_into_vec(&this.0, &haystack).map_err(any_err)
81 });
82 }
83}
84
85pub fn register(lua: &Lua) -> anyhow::Result<()> {
86 let regex_mod = get_or_create_sub_module(lua, "regex")?;
87
88 regex_mod.set(
89 "compile",
90 lua.create_function(move |_, pattern: String| {
91 let re = Regex::new(&pattern).map_err(any_err)?;
92 Ok(RegexWrap(re))
93 })?,
94 )?;
95
96 regex_mod.set(
97 "escape",
98 lua.create_function(move |_, pattern: String| Ok(regex::escape(&pattern)))?,
99 )?;
100
101 Ok(())
102}
103
104struct Split<'r, 'h> {
105 finder: Matches<'r, 'h>,
106 last: usize,
107 haystack: &'h str,
108}
109
110impl<'r, 'h> Split<'r, 'h> {
111 fn split(re: &'r Regex, haystack: &'h str) -> Self {
112 Self {
113 finder: re.find_iter(haystack),
114 last: 0,
115 haystack,
116 }
117 }
118}
119
120impl<'h> Iterator for Split<'_, 'h> {
121 type Item = Result<&'h str, fancy_regex::Error>;
122
123 fn next(&mut self) -> Option<Result<&'h str, fancy_regex::Error>> {
124 match self.finder.next() {
125 None => {
126 let len = self.haystack.len();
127 if self.last > len {
128 None
129 } else {
130 let span = &self.haystack[self.last..len];
131 self.last = len + 1; Some(Ok(span))
133 }
134 }
135 Some(Ok(m)) => {
136 let span = &self.haystack[self.last..m.start()];
137 self.last = m.end();
138 Some(Ok(span))
139 }
140 Some(Err(e)) => Some(Err(e)),
141 }
142 }
143}
144
145#[allow(clippy::result_large_err)]
146fn split_into_vec(re: &Regex, haystack: &str) -> Result<Vec<String>, fancy_regex::Error> {
147 let mut result = vec![];
148 for m in Split::split(re, haystack) {
149 let m = m?;
150 result.push(m.to_string());
151 }
152 Ok(result)
153}
154
155#[cfg(test)]
156mod test {
157 use super::*;
158
159 #[test]
160 fn fancy_split() {
161 let re = Regex::new("[ \t]+").unwrap();
162 let hay = "a b \t c\td e";
163 let fields = split_into_vec(&re, hay).unwrap();
164 assert_eq!(fields, vec!["a", "b", "c", "d", "e"]);
165 }
166}