use aho_corasick::AhoCorasick;
use regex::Regex;
use roto::{Context, Optional, Runtime, Val, roto_method, roto_static_method};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use crate::means_of_production::{Error, MutableStringList};
#[derive(Debug, Clone, Context, Default)]
pub struct IocaineContext {
pub iocaine_patterns: PatternFinderMap,
pub iocaine_regexes: RegexFinderMap,
}
pub type PatternFinder = Arc<AhoCorasick>;
pub type RegexFinder = Arc<Regex>;
pub type PatternFinderMap = Arc<RwLock<HashMap<Arc<str>, PatternFinder>>>;
pub type RegexFinderMap = Arc<RwLock<HashMap<Arc<str>, RegexFinder>>>;
fn register_pattern_finder(runtime: &mut Runtime) -> Result<(), Error> {
runtime.register_clone_type_with_name::<PatternFinderMap>(
"PatternFinderMap",
"A map of pattern finders",
)?;
runtime.register_clone_type_with_name::<Optional<PatternFinder>>(
"PatternFinder",
"A pattern finder",
)?;
#[roto_static_method(runtime, Optional<PatternFinder>)]
fn new(patterns: Val<MutableStringList>) -> Val<Optional<PatternFinder>> {
AhoCorasick::builder()
.ascii_case_insensitive(true)
.build((*patterns).0.borrow().iter())
.map_or_else(
|e| {
tracing::error!("Unable to create AhoCorasick: {e}");
Optional::None
},
|ac| Optional::Some(ac.into()),
)
.into()
}
#[roto_method(runtime, PatternFinderMap)]
fn insert(
finders: Val<PatternFinderMap>,
key: Arc<str>,
patterns: Val<Optional<PatternFinder>>,
) {
if let Optional::Some(p) = patterns.0 {
let _ = finders
.write()
.map(|mut f| f.insert(key, p))
.inspect_err(|e| {
tracing::error!("Unable to lock PatternFinderMap for writing: {e}");
});
}
}
#[roto_method(runtime, PatternFinderMap)]
fn insert_patterns(
finders: Val<PatternFinderMap>,
key: Arc<str>,
patterns: Val<MutableStringList>,
) {
let ac = AhoCorasick::builder()
.ascii_case_insensitive(true)
.build((*patterns).0.borrow().iter());
let ac = match ac {
Ok(v) => v,
Err(e) => {
tracing::error!("Unable to create AhoCorasick: {e}");
return;
}
};
let _ = finders
.write()
.map(|mut f| f.insert(key, ac.into()))
.inspect_err(|e| tracing::error!("Unable to lock PatternFinder for writing: {e}"));
}
#[roto_method(runtime, PatternFinderMap)]
fn get(patterns: Val<PatternFinderMap>, key: Arc<str>) -> Val<Optional<PatternFinder>> {
(*patterns)
.read()
.map_or_else(
|e| {
tracing::error!("Unable to lock PatternFinderMap for reading: {e}");
Optional::None
},
|p| p.get(&key).cloned().into(),
)
.into()
}
#[roto_method(runtime, Optional<PatternFinder>)]
fn is_match(pf: Val<Optional<PatternFinder>>, s: Arc<str>) -> bool {
if let Optional::Some(pattern) = pf.0 {
pattern.is_match(&s.to_string())
} else {
false
}
}
Ok(())
}
fn register_regex_finder(runtime: &mut Runtime) -> Result<(), Error> {
runtime.register_clone_type_with_name::<RegexFinderMap>(
"RegexFinderMap",
"A map of compiled regexes",
)?;
runtime
.register_clone_type_with_name::<Optional<RegexFinder>>("RegexFinder", "A regex finder")?;
#[roto_static_method(runtime, Optional<RegexFinder>)]
fn new(s: Arc<str>) -> Val<Optional<RegexFinder>> {
Regex::new(&s)
.map_or_else(
|e| {
tracing::error!({ regex = s.to_string() }, "Unable to create Regex: {e}");
Optional::None
},
|p| Optional::Some(p.into()),
)
.into()
}
#[roto_method(runtime, RegexFinderMap)]
fn insert(regexes: Val<RegexFinderMap>, key: Arc<str>, rf: Val<Optional<RegexFinder>>) {
if let Optional::Some(re) = rf.0 {
let _ = regexes
.write()
.map(|mut f| f.insert(key, re))
.inspect_err(|e| tracing::error!("Unable to lock RegexFinderMap for writing: {e}"));
}
}
#[roto_method(runtime, RegexFinderMap)]
fn insert_regex(regexes: Val<RegexFinderMap>, key: Arc<str>, pattern: Arc<str>) {
if pattern.is_empty() {
return;
}
let _ = Regex::new(&pattern)
.map(|re| {
regexes
.write()
.map(|mut f| f.insert(key, re.into()))
.inspect_err(|e| tracing::error!("Unable to lock RegexFinder for writing: {e}"))
})
.inspect_err(|e| {
tracing::error!(
{ regex = pattern.to_string() },
"Unable to create Regex: {e}"
);
});
}
#[roto_method(runtime, RegexFinderMap)]
fn get(regexes: Val<RegexFinderMap>, key: Arc<str>) -> Val<Optional<RegexFinder>> {
(*regexes)
.read()
.map_or_else(
|e| {
tracing::error!("Unable to lock RegexFinder for reading: {e}");
Optional::None
},
|r| r.get(&key).cloned().into(),
)
.into()
}
#[roto_method(runtime, Optional<RegexFinder>)]
fn is_match(rf: Val<Optional<RegexFinder>>, s: Arc<str>) -> bool {
if let Optional::Some(re) = rf.0 {
re.is_match(&s)
} else {
false
}
}
Ok(())
}
fn register_env(runtime: &mut Runtime) -> Result<(), Error> {
#[derive(Debug, Clone, Copy)]
struct Env;
runtime.register_copy_type::<Env>("The environment of iocaine")?;
#[roto_static_method(runtime, Env)]
fn get(var: Arc<str>) -> Arc<str> {
std::env::var(var.as_ref())
.unwrap_or_else(|_| String::default())
.into()
}
Ok(())
}
pub fn register_context(runtime: &mut Runtime) -> Result<(), Error> {
register_pattern_finder(runtime)?;
register_regex_finder(runtime)?;
register_env(runtime)?;
runtime.register_context_type::<IocaineContext>()?;
Ok(())
}