iocaine 2.2.0

The deadliest poison known to AI
Documentation
// SPDX-FileCopyrightText: 2025 Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

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(())
}