mumu-flow 0.1.0

Stream tranform tools plugin for the Lava language
Documentation
use mumu::parser::interpreter::{Interpreter, DynamicFnInfo};
use mumu::parser::types::{
    Value,
    FunctionValue,
    InkIteratorHandle,
    InkIteratorKind,
};
use mumu::parser::interpreter::apply_n_ary_function_value;
use std::sync::{Arc, Mutex};
use std::time::{Instant, Duration};
use std::thread;

pub fn register_flow_throttle(interp: &mut Interpreter) {
    let f = Arc::new(Mutex::new(flow_throttle_bridge_fn));
    let info = DynamicFnInfo::new(f, true);
    interp.register_dynamic_function_ex("flow:throttle", info);

    interp.set_variable(
        "flow:throttle",
        Value::Function(Box::new(FunctionValue::Named("flow:throttle".to_string())))
    );
}

fn flow_throttle_bridge_fn(_interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_partial(None, None)),
        1 => {
            let one_val = &args[0];
            if is_placeholder(one_val) {
                Ok(make_partial(None, None))
            } else if let Value::Int(ms_i) = one_val {
                Ok(make_partial(Some(*ms_i), None))
            } else {
                Err(format!("flow:throttle => first param must be an int or '_', got {:?}", one_val))
            }
        }
        2 => {
            let a_ms = &args[0];
            let a_src = &args[1];

            if is_placeholder(a_ms) || is_placeholder(a_src) {
                let ms_opt = match a_ms {
                    Value::Int(n) => Some(*n),
                    _ if is_placeholder(a_ms) => None,
                    other => return Err(format!("flow:throttle => first param must be int or '_', got {:?}", other)),
                };
                let src_opt = if is_placeholder(a_src) { None } else { Some(a_src.clone()) };

                if ms_opt.is_some() && src_opt.is_some() {
                    let ms_val = ms_opt.unwrap();
                    let src_val = src_opt.unwrap();
                    finalize_throttle(ms_val, src_val)
                } else {
                    Ok(make_partial(ms_opt, src_opt))
                }
            } else {
                if let Value::Int(m) = a_ms {
                    finalize_throttle(*m, a_src.clone())
                } else {
                    Err(format!("flow:throttle => first param must be int, got {:?}", a_ms))
                }
            }
        }
        n => Err(format!("flow:throttle => up to 2 arguments, got {}", n)),
    }
}

fn finalize_throttle(ms: i32, source_val: Value) -> Result<Value, String> {
    if ms < 0 {
        return Err(format!("flow:throttle => negative ms not allowed: {}", ms));
    }
    match source_val {
        Value::InkIterator(h) => Ok(build_throttle_iter(ms, h)),
        Value::InkTransform(tf) => Ok(build_throttle_chain(ms, tf)),
        other => Err(format!("flow:throttle => second param must be InkIterator or InkTransform, got {:?}", other)),
    }
}

fn make_partial(init_ms: Option<i32>, init_src: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    let partial = ThrottlePartial {
        ms: init_ms,
        source: init_src,
    };
    let shared_state = Arc::new(Mutex::new(partial));

    let closure = RustClosure(
        "flow:throttle-partial".to_string(),
        Arc::new(Mutex::new(move |_pinterp: &mut Interpreter, new_args: Vec<Value>| {
            let mut guard = shared_state.lock().map_err(|_| "flow:throttle => partial lock error".to_string())?;

            for val in &new_args {
                if guard.ms.is_none() {
                    if is_placeholder(val) {
                        // remain none
                    } else if let Value::Int(i) = val {
                        guard.ms = Some(*i);
                    } else {
                        return Err(format!("flow:throttle => first param must be int or '_', got {:?}", val));
                    }
                } else if guard.source.is_none() {
                    if is_placeholder(val) {
                        // remain none
                    } else {
                        guard.source = Some(val.clone());
                    }
                } else {
                    return Err("flow:throttle => partial => too many arguments".to_string());
                }
            }

            if guard.ms.is_some() && guard.source.is_some() {
                let final_ms = guard.ms.unwrap();
                let final_src = guard.source.as_ref().unwrap().clone();

                drop(guard);
                finalize_throttle(final_ms, final_src)
            } else {
                let ms_c = guard.ms;
                let src_c = guard.source.clone();
                drop(guard);

                Ok(make_partial(ms_c, src_c))
            }
        })),
        0
    );

    Value::Function(Box::new(closure))
}

fn build_throttle_iter(ms: i32, handle: InkIteratorHandle) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    let env = ThrottleIterEnv {
        ms_delay: ms,
        upstream: handle,
        last_item_time: Instant::now(),
    };
    let shared = Arc::new(Mutex::new(env));

    let clos = RustClosure(
        "flow:throttle-iterator".to_string(),
        Arc::new(Mutex::new(move |_itp: &mut Interpreter, _a: Vec<Value>| {
            let mut guard = shared.lock().map_err(|_| "flow:throttle => lock error (iterator)".to_string())?;

            let now = Instant::now();
            let allowed = guard.last_item_time + Duration::from_millis(guard.ms_delay as u64);
            if now < allowed {
                let wait_dur = allowed - now;
                thread::sleep(wait_dur);
            }

            let kind = &guard.upstream.kind;
            let result = match kind {
                InkIteratorKind::Core(state_arc) => {
                    let mut i_guard = state_arc.lock().map_err(|_| "flow:throttle => lock error on InkIterator".to_string())?;
                    if i_guard.done || i_guard.current >= i_guard.end {
                        i_guard.done = true;
                        Err("NO_MORE_DATA".to_string())
                    } else {
                        let val = Value::Int(i_guard.current);
                        i_guard.current += 1;
                        if i_guard.current >= i_guard.end {
                            i_guard.done = true;
                        }
                        Ok(val)
                    }
                }
                InkIteratorKind::Plugin(plugin_arc) => {
                    let mut plugin = plugin_arc.lock().map_err(|_| "flow:throttle => plugin lock error".to_string())?;
                    match plugin.next_value() {
                        Ok(val) => Ok(val),
                        Err(e) => Err(e),
                    }
                }
            };

            if let Ok(_) = result {
                guard.last_item_time = Instant::now();
            }
            result
        })),
        0
    );

    Value::InkTransform(Box::new(clos))
}

fn build_throttle_chain(ms: i32, upstream_tf: Box<FunctionValue>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    let env = ThrottleChainEnv {
        ms_delay: ms,
        upstream: upstream_tf,
        last_item_time: Instant::now(),
    };
    let shared = Arc::new(Mutex::new(env));

    let clos = RustClosure(
        "flow:throttle-transform".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            let mut guard = shared.lock().map_err(|_| "flow:throttle => lock error (chain)".to_string())?;

            let now = Instant::now();
            let allowed = guard.last_item_time + Duration::from_millis(guard.ms_delay as u64);
            if now < allowed {
                thread::sleep(allowed - now);
            }

            let res = apply_n_ary_function_value(interp, guard.upstream.clone(), vec![]);
            match res {
                Ok(item) => {
                    guard.last_item_time = Instant::now();
                    Ok(item)
                }
                Err(e) => Err(e),
            }
        })),
        0
    );

    Value::InkTransform(Box::new(clos))
}

#[derive(Clone)]
struct ThrottlePartial {
    ms: Option<i32>,
    source: Option<Value>,
}

struct ThrottleIterEnv {
    ms_delay: i32,
    upstream: InkIteratorHandle,
    last_item_time: std::time::Instant,
}

struct ThrottleChainEnv {
    ms_delay: i32,
    upstream: Box<FunctionValue>,
    last_item_time: std::time::Instant,
}

fn is_placeholder(v: &Value) -> bool {
    match v {
        Value::Placeholder => true,
        Value::SingleString(s) if s == "_" => true,
        Value::StrArray(ss) if ss.len() == 1 && ss[0] == "_" => true,
        _ => false,
    }
}