1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

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