redis-lua 0.4.0

use crate::types::ScriptArg;
use futures::prelude::*;

/// Script information which is generated by proc-macro.
#[derive(Clone, Debug)]
pub struct Info {
    /// The entire script including arguments initialization.
    script: &'static str,
    /// The script excluding arguments initialization.
    body: &'static str,
    /// The list of arguments.
    args: &'static [&'static str],
}

impl Info {
    /// Create the new script information.
    pub fn new(script: &'static str, body: &'static str, args: &'static [&'static str]) -> Self {
        Self { script, body, args }
    }
}

/// To make sure `Script` be object safe.
fn _object_safe(_: &dyn Script) {}

/// Represents a complete invocable script which has a complete set of arguments.
pub trait Script {
    /// Retrieve all the script information.
    fn info(&self, _: &mut Vec<Info>, _: &mut Vec<ScriptArg>);

    /// Join another script making self as inner.
    fn join<T: Script>(self, other: T) -> ScriptJoin<Self, T>
    where
        Self: Sized,
    {
        ScriptJoin(self, other)
    }

    /// Invoke the script.
    fn invoke<T>(self, con: &mut dyn redis::ConnectionLike) -> redis::RedisResult<T>
    where
        T: redis::FromRedisValue,
        Self: Sized,
    {
        let mut info = vec![];
        let mut args = vec![];
        self.info(&mut info, &mut args);
        let script = gen_script(&info, &args);
        let mut invoke = script.prepare_invoke();
        for wr in args {
            invoke.arg(wr);
        }
        invoke.invoke(con)
    }

    /// Invoke the script asynchronously.
    fn invoke_async<'a, C, T>(self, con: &'a mut C) -> redis::RedisFuture<'a, T>
    where
        C: redis::aio::ConnectionLike + Send,
        T: redis::FromRedisValue + Send,
        Self: Sized + Send + 'a,
    {
        async move {
            let mut info = vec![];
            let mut args = vec![];
            self.info(&mut info, &mut args);
            let script = gen_script(&info, &args);
            let mut invoke = script.prepare_invoke();
            for wr in args {
                invoke.arg(wr);
            }
            invoke.invoke_async(con).await
        }
        .boxed()
    }
}

impl<S: Script + ?Sized> Script for Box<S> {
    fn info(&self, infos: &mut Vec<Info>, args: &mut Vec<ScriptArg>) {
        (**self).info(infos, args);
    }
}

impl Script for () {
    fn info(&self, _: &mut Vec<Info>, _: &mut Vec<ScriptArg>) {}
}

/// Represents the set of two scripts which are joined.
pub struct ScriptJoin<S, T>(S, T);

impl<S, T> Script for ScriptJoin<S, T>
where
    S: Script,
    T: Script,
{
    fn info(&self, info: &mut Vec<Info>, args: &mut Vec<ScriptArg>) {
        self.0.info(info, args);
        self.1.info(info, args);
    }
}

/// Take another script as the inner of the script.
pub trait TakeScript<I> {
    type Item;

    /// Take the inner script.
    fn take(self, inner: I) -> Self::Item;
}

/// Generate a script from a list of script information.
pub fn gen_script(info: &[Info], args: &[ScriptArg]) -> redis::Script {
    assert!(info.len() > 0, "No script information");

    // Generate the joined script.
    let mut arg_index = 0;
    let mut script = String::new();
    let last = info.len() - 1;
    for (index, info) in info.iter().enumerate() {
        let prefix = if index == last { "return " } else { "" };
        let mut init = String::new();

        for arg in info.args {
            let pack = args[arg_index].pack();

            arg_index += 1;

            if pack {
                init += &format!("local {} = cmsgpack.unpack(ARGV[{}]) ", arg, arg_index);
            } else {
                init += &format!("local {} = ARGV[{}] ", arg, arg_index);
            }
        }

        script += &format!("{}(function() {} {} end)();\n", prefix, init, info.body);
    }
    redis::Script::new(&script)
}