Add regex module

Author: khvzak

Cargo.toml

@@ -17,10 +17,13 @@ send = ["mlua/send"]
 vendored = ["mlua/vendored"]
 
 json = ["mlua/serde", "dep:ouroboros", "dep:serde", "dep:serde_json"]
+regex = ["dep:regex", "dep:ouroboros", "dep:quick_cache"]
 
 [dependencies]
 mlua = { version = "0.11" }
 ouroboros = { version = "0.18", optional = true }
 serde = { version = "1.0", optional = true }
 serde_json = { version = "1.0", optional = true }
 owo-colors = "4"
+regex = { version = "1.0", optional = true }
+quick_cache = { version = "0.6", optional = true }

src/lib.rs

@@ -15,3 +15,5 @@ pub mod testing;
 
 #[cfg(feature = "json")]
 pub mod json;
+#[cfg(feature = "regex")]
+pub mod regex;

src/regex.rs

@@ -0,0 +1,210 @@
+use std::ops::Deref;
+use std::result::Result as StdResult;
+use std::sync::LazyLock;
+
+use mlua::{Lua, MetaMethod, Result, String as LuaString, Table, UserData, UserDataMethods, Value, Variadic};
+use ouroboros::self_referencing;
+use quick_cache::sync::Cache;
+
+// A reasonable cache size for regexes. This can be adjusted as needed.
+const REGEX_CACHE_SIZE: usize = 256;
+
+#[derive(Clone, Debug)]
+pub struct Regex(regex::bytes::Regex);
+
+impl Deref for Regex {
+    type Target = regex::bytes::Regex;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+// Global cache for regexes shared across all Lua states.
+static CACHE: LazyLock<Cache<String, Regex>> = LazyLock::new(|| Cache::new(REGEX_CACHE_SIZE));
+
+impl Regex {
+    /// Creates a new cached regex or retrieves it from the cache if it already exists.
+    pub fn new(_: &Lua, re: &str) -> StdResult<Self, regex::Error> {
+        if let Some(re) = CACHE.get(re) {
+            return Ok(re);
+        }
+        let regex = Self(regex::bytes::Regex::new(&re)?);
+        CACHE.insert(re.to_string(), regex.clone());
+        Ok(regex)
+    }
+}
+
+impl UserData for Regex {
+    fn register(registry: &mut mlua::UserDataRegistry<Self>) {
+        registry.add_method("is_match", |_, this, text: LuaString| {
+            Ok(this.0.is_match(&text.as_bytes()))
+        });
+
+        registry.add_method("match", |lua, this, text: LuaString| {
+            let text = (*text.as_bytes()).into();
+            let caps = Captures::try_new(text, |text| this.0.captures(text).ok_or(()));
+            match caps {
+                Ok(caps) => Ok(Value::UserData(lua.create_userdata(caps)?)),
+                Err(_) => Ok(Value::Nil),
+            }
+        });
+
+        // Returns low level information about raw offsets of each submatch.
+        registry.add_method("captures_read", |lua, this, text: LuaString| {
+            let mut locs = this.capture_locations();
+            match this.captures_read(&mut locs, &text.as_bytes()) {
+                Some(_) => Ok(Value::UserData(lua.create_userdata(CaptureLocations(locs))?)),
+                None => Ok(Value::Nil),
+            }
+        });
+
+        registry.add_method("split", |lua, this, text: LuaString| {
+            lua.create_sequence_from(this.split(&text.as_bytes()).map(LuaString::wrap))
+        });
+
+        registry.add_method("splitn", |lua, this, (text, limit): (LuaString, usize)| {
+            lua.create_sequence_from(this.splitn(&text.as_bytes(), limit).map(LuaString::wrap))
+        });
+
+        registry.add_method("replace", |lua, this, (text, rep): (LuaString, LuaString)| {
+            lua.create_string(this.replace(&text.as_bytes(), &*rep.as_bytes()))
+        });
+    }
+}
+
+#[self_referencing]
+struct Captures {
+    text: Box<[u8]>,
+
+    #[borrows(text)]
+    #[covariant]
+    caps: regex::bytes::Captures<'this>,
+}
+
+impl UserData for Captures {
+    fn register(registry: &mut mlua::UserDataRegistry<Self>) {
+        registry.add_meta_method(MetaMethod::Index, |lua, this, key: Value| match key {
+            Value::String(s) => {
+                let name = s.to_string_lossy();
+                this.borrow_caps()
+                    .name(&name)
+                    .map(|v| lua.create_string(v.as_bytes()))
+                    .transpose()
+            }
+            Value::Integer(i) => this
+                .borrow_caps()
+                .get(i as usize)
+                .map(|v| lua.create_string(v.as_bytes()))
+                .transpose(),
+            _ => Ok(None),
+        })
+    }
+}
+
+struct CaptureLocations(regex::bytes::CaptureLocations);
+
+impl UserData for CaptureLocations {
+    fn register(registry: &mut mlua::UserDataRegistry<Self>) {
+        // Returns the total number of capture groups.
+        registry.add_method("len", |_, this, ()| Ok(this.0.len()));
+
+        // Returns the start and end positions of the Nth capture group.
+        registry.add_method("get", |_, this, i: usize| match this.0.get(i) {
+            // We add 1 to the start position because Lua is 1-indexed.
+            // End position is non-inclusive, so we don't need to add 1.
+            Some((start, end)) => Ok(Variadic::from_iter([start + 1, end])),
+            None => Ok(Variadic::new()),
+        });
+    }
+}
+
+struct RegexSet(regex::bytes::RegexSet);
+
+impl Deref for RegexSet {
+    type Target = regex::bytes::RegexSet;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl UserData for RegexSet {
+    fn register(registry: &mut mlua::UserDataRegistry<Self>) {
+        registry.add_function("new", |_, patterns: Vec<String>| {
+            let set = lua_try!(regex::bytes::RegexSet::new(patterns).map(RegexSet));
+            Ok(Ok(set))
+        });
+
+        registry.add_method("is_match", |_, this, text: LuaString| {
+            Ok(this.is_match(&text.as_bytes()))
+        });
+
+        registry.add_method("len", |_, this, ()| Ok(this.len()));
+
+        registry.add_method("matches", |_, this, text: LuaString| {
+            Ok(this
+                .matches(&text.as_bytes())
+                .iter()
+                .map(|i| i + 1)
+                .collect::<Vec<_>>())
+        });
+    }
+}
+
+/// Compiles a regular expression.
+///
+/// Once compiled, it can be used repeatedly to search, split or replace substrings in a text.
+fn regex_new(lua: &Lua, re: LuaString) -> Result<StdResult<Regex, String>> {
+    let re = re.to_str()?;
+    Ok(Ok(lua_try!(Regex::new(lua, &re))))
+}
+
+/// Escapes a string so that it can be used as a literal in a regular expression.
+fn regex_escape(_: &Lua, text: LuaString) -> Result<String> {
+    Ok(regex::escape(&text.to_str()?))
+}
+
+/// Returns true if there is a match for the regex anywhere in the given text.
+fn regex_is_match(lua: &Lua, (re, text): (LuaString, LuaString)) -> Result<StdResult<bool, String>> {
+    let re = re.to_str()?;
+    let re = lua_try!(Regex::new(lua, &re));
+    Ok(Ok(re.is_match(&text.as_bytes())))
+}
+
+/// Returns all matches of the regex in the given text or nil if there is no match.
+fn regex_match(lua: &Lua, (re, text): (LuaString, LuaString)) -> Result<StdResult<Value, String>> {
+    let re = re.to_str()?;
+    let re = lua_try!(Regex::new(lua, &re));
+    match re.captures(&text.as_bytes()) {
+        Some(caps) => {
+            let mut it = caps.iter().map(|om| om.map(|m| LuaString::wrap(m.as_bytes())));
+            let first = it.next().unwrap();
+            let table = lua.create_sequence_from(it)?;
+            table.raw_set(0, first)?;
+            Ok(Ok(Value::Table(table)))
+        }
+        None => Ok(Ok(Value::Nil)),
+    }
+}
+
+/// A loader for the `regex` module.
+fn loader(lua: &Lua) -> Result<Table> {
+    let t = lua.create_table()?;
+    t.set("new", lua.create_function(regex_new)?)?;
+    t.set("escape", lua.create_function(regex_escape)?)?;
+    t.set("is_match", lua.create_function(regex_is_match)?)?;
+    t.set("match", lua.create_function(regex_match)?)?;
+    t.set("RegexSet", lua.create_proxy::<RegexSet>()?)?;
+    Ok(t)
+}
+
+/// Registers the `regex` module in the given Lua state.
+pub fn register(lua: &Lua, name: Option<&str>) -> Result<Table> {
+    let name = name.unwrap_or("@regex");
+    let value = loader(lua)?;
+    lua.register_module(name, &value)?;
+    Ok(value)
+}

tests/lib.rs

@@ -13,6 +13,8 @@ fn run_file(modname: &str) -> Result<()> {
 
     #[cfg(feature = "json")]
     mlua_stdlib::json::register(&lua, None)?;
+    #[cfg(feature = "regex")]
+    mlua_stdlib::regex::register(&lua, None)?;
 
     // Add `testing` global variable (an instance of the testing framework)
     let testing = testing.call_function::<Table>("new", modname)?;
@@ -50,3 +52,5 @@ include_tests! {
 
 #[cfg(feature = "json")]
 include_tests!(json);
+#[cfg(feature = "regex")]
+include_tests!(regex);

tests/lua/regex_tests.lua

@@ -0,0 +1,87 @@
+local regex = require("@regex")
+
+-- Test basic regex functionality
+testing:test("regex_basic", function(t)
+    local re = regex.new(".*(?P<gr1>abc)")
+
+    t.assert(re:is_match("123abc321"), "is_match() should have matches")
+    t.assert(not re:is_match("123"), "is_match() should not have matches")
+
+    local matches = re:match("123abc321")
+    t.assert_eq(matches[0], "123abc", "zero capture group should match the whole text")
+    t.assert_eq(matches[1], "abc", "first capture group should match `abc`")
+    t.assert_eq(matches["gr1"], "abc", "named capture group should match `abc`")
+    t.assert_eq(matches[true], nil, "bad key should have no match")
+
+    -- Test split
+    local re_split = regex.new("[,.]")
+    local vec = re_split:split("abc.qwe,rty.asd")
+    t.assert_eq(#vec, 4, "vec len should be 4")
+    t.assert(
+        vec[1] == "abc" and vec[2] == "qwe" and vec[3] == "rty" and vec[4] == "asd",
+        "vec must be 'abc', 'qwe', 'rty', 'asd'"
+    )
+
+    vec = re_split:splitn("abc,bcd,cde", 2)
+    t.assert_eq(#vec, 2, "vec len should be 2")
+    t.assert(vec[1] == "abc" and vec[2] == "bcd,cde", "vec must be 'abc', 'bcd,cde'")
+
+    -- Test invalid regex
+    local re_invalid, err = regex.new("(")
+    t.assert_eq(re_invalid, nil, "re is not nil")
+    t.assert(string.find(err, "regex parse error"), "err must contain 'regex parse error'")
+
+    -- Test replace
+    local re_replace = regex.new("(?P<last>[^,\\s]+),\\s+(?P<first>\\S+)")
+    local str = re_replace:replace("Smith, John", "$first $last")
+    t.assert_eq(str, "John Smith", "str must be 'John Smith'")
+end)
+
+-- Test regex shortcuts (escape, is_match, match functions)
+testing:test("regex_shortcuts", function(t)
+    -- Test escape
+    t.assert_eq(regex.escape("a*b"), "a\\*b", "escaped regex must be 'a\\*b'")
+
+    -- Test "is_match"
+    t.assert(regex.is_match("\\b\\w{13}\\b", "I categorically deny having ..."), "is_match should have matches")
+    t.assert(not regex.is_match("abc", "bca"), "is_match should not have matches")
+    local is_match, err = regex.is_match("(", "")
+    t.assert(is_match == nil and string.find(err, "regex parse error") ~= nil, "is_match should return error")
+
+    -- Test "match"
+    local matches = regex.match("^(\\d{4})-(\\d{2})-(\\d{2})$", "2014-05-01")
+    t.assert_eq(matches[0], "2014-05-01", "zero capture group should match the whole text")
+    t.assert_eq(matches[1], "2014", "first capture group should match year")
+    t.assert_eq(matches[2], "05", "second capture group should match month")
+    t.assert_eq(matches[3], "01", "third capture group should match day")
+    matches, err = regex.match("(", "")
+    t.assert(matches == nil and string.find(err, "regex parse error") ~= nil, "match should return error")
+end)
+
+-- Test RegexSet functionality
+testing:test("regex_set", function(t)
+    local set = regex.RegexSet.new({ "\\w+", "\\d+", "\\pL+", "foo", "bar", "barfoo", "foobar" })
+    t.assert_eq(set:len(), 7, "len should be 7")
+    t.assert(set:is_match("foobar"), "is_match should have matches")
+    t.assert_eq(table.concat(set:matches("foobar"), ","), "1,3,4,5,7", "matches should return 1,3,4,5,7")
+end)
+
+-- Test capture locations
+testing:test("capture_locations", function(t)
+    local re = regex.new("\\d+(abc)\\d+")
+
+    local str = "123abc321"
+    local locs = re:captures_read(str)
+    t.assert(locs, "locs is nil")
+    t.assert_eq(locs:len(), 2, "locs len is not 2")
+    local i, j = locs:get(0)
+    t.assert(i == 1 and j == 9, "locs:get(0) is not 1, 9")
+    i, j = locs:get(1)
+    t.assert(i == 4 and j == 6, "locs:get(1) is not 4, 6")
+    t.assert_eq(str:sub(i, j), "abc", "str:sub(i, j) is not 'abc'")
+    t.assert_eq(locs:get(2), nil, "locs:get(2) is nil")
+
+    -- Test no match
+    locs = re:captures_read("123")
+    t.assert_eq(locs, nil, "locs is not nil")
+end)