flow-mumu 0.2.0-rc.9

// src/throttle_with.rs
//
// flow:throttle_with(period_fn, source)
// Non-blocking throttle like `flow:throttle`, but the **period** is supplied
// by a *function* that is called **after each emission** to schedule the
// next allowed emit time. This avoids draining finite providers while we
// are waiting and guarantees retention/reuse of the same function.
//
// - `period_fn` must be a 0-arg function returning an int/long/float
//   (milliseconds, ≥ 0). It may itself be stateful (e.g., a transform
//   yielding a sequence of delays); we only advance it when we actually emit.
// - `source` is an Iterator handle or a zero-arg transform Function.
// - If the next allowed time hasn’t arrived yet, returns **"AGAIN"**
//   (temporary no-data) without sleeping.
//
// Partial application forms supported:
//   • flow:throttle_with(_)                    -> waits for (period_fn, source)
//   • flow:throttle_with(period_fn)            -> returns (source) => transform
//   • flow:throttle_with(source)               -> returns (period_fn) => transform
//   • flow:throttle_with(period_fn, source)    -> ready transform
//
// First item is allowed **immediately** (no initial delay).

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use mumu::parser::interpreter::{Interpreter, DynamicFnInfo};
use mumu::parser::interpreter::apply::apply_n_ary_function_value;
use mumu::parser::types::{FunctionValue, IteratorHandle, IteratorKind, Value};

use crate::common::is_placeholder;

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

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

fn flow_throttle_with_bridge_fn(
    interp: &mut Interpreter,
    args: Vec<Value>,
) -> Result<Value, String> {
    match args.len() {
        0 => Ok(make_partial(None, None)),

        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_partial(None, None))
            } else if let Value::Function(fb) = a {
                // period function provided first; wait for source
                Ok(make_partial(Some(fb.clone()), None))
            } else {
                // assume it's the source; wait for period function
                Ok(make_partial(None, Some(a.clone())))
            }
        }

        2 => {
            let pf = &args[0];
            let src = &args[1];

            // Handle placeholders via partial
            if is_placeholder(pf) || is_placeholder(src) {
                let pf_opt = if is_placeholder(pf) {
                    None
                } else {
                    match pf {
                        Value::Function(fb) => Some(fb.clone()),
                        other => {
                            return Err(format!(
                                "flow:throttle_with => first argument must be a Function (period provider) or '_', got {:?}",
                                other
                            ))
                        }
                    }
                };
                let src_opt = if is_placeholder(src) { None } else { Some(src.clone()) };
                return Ok(make_partial(pf_opt, src_opt));
            }

            // Finalize only when period is a Function
            let pf_fn = match pf {
                Value::Function(fb) => fb.clone(),
                other => {
                    return Err(format!(
                        "flow:throttle_with => first argument must be a Function (period provider), got {:?}",
                        other
                    ))
                }
            };
            finalize_throttle_with(interp, pf_fn, src.clone())
        }

        n => Err(format!(
            "flow:throttle_with => expected 0, 1 or 2 arguments, got {}",
            n
        )),
    }
}

fn make_partial(period_fn_opt: Option<Box<FunctionValue>>, src_opt: Option<Value>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    Value::Function(Box::new(RustClosure(
        "flow:throttle_with-partial".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, new_args: Vec<Value>| {
            let mut pf = period_fn_opt.clone();
            let mut src = src_opt.clone();

            for a in new_args {
                if pf.is_none() {
                    if is_placeholder(&a) {
                        // keep None
                    } else if let Value::Function(fb) = a {
                        pf = Some(fb);
                    } else {
                        return Err(
                            "flow:throttle_with => first param must be a Function or '_'"
                                .to_string(),
                        );
                    }
                    continue;
                }
                if src.is_none() {
                    if is_placeholder(&a) {
                        // keep None
                    } else {
                        src = Some(a);
                    }
                    continue;
                }
                return Err("flow:throttle_with => partial => too many arguments".to_string());
            }

            match (pf.clone(), src.clone()) {
                (Some(pf_fn), Some(sv)) => finalize_throttle_with(interp, pf_fn, sv),
                _ => Ok(make_partial(pf, src)),
            }
        })),
        0,
    )))
}

fn finalize_throttle_with(
    _interp: &mut Interpreter,
    period_fn: Box<FunctionValue>,
    source: Value,
) -> Result<Value, String> {
    match source {
        Value::Iterator(h) => Ok(build_throttle_with_iter(period_fn, h)),
        Value::Function(tf) => Ok(build_throttle_with_chain(period_fn, tf)),
        other => Err(format!(
            "flow:throttle_with => second argument must be Iterator or transform Function, got {:?}",
            other
        )),
    }
}

/* ──────────────────────────────────────────────────────────────────────────
   Dynamic-period throttle implementations
   * We only advance the period_fn **after a successful emission**:
   *   next_ready = now + ms(period_fn())
   * While waiting, we return "AGAIN" and DO NOT call period_fn at all.
   * First emission has no delay.
   ───────────────────────────────────────────────────────────────────────── */

fn build_throttle_with_iter(period_fn: Box<FunctionValue>, handle: IteratorHandle) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    #[derive(Debug)]
    struct Env {
        period_fn: Box<FunctionValue>,
        upstream: IteratorHandle,
        next_ready: Option<Instant>, // None => first item allowed immediately
    }

    let shared = Arc::new(Mutex::new(Env {
        period_fn,
        upstream: handle,
        next_ready: None,
    }));

    let clos = RustClosure(
        "flow:throttle_with-iterator".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            // Gate without touching the period provider
            {
                let st = shared
                    .lock()
                    .map_err(|_| "flow:throttle_with => lock error (iterator gate)".to_string())?;
                if let Some(ready_at) = st.next_ready {
                    if Instant::now() < ready_at {
                        return Err("AGAIN".to_string());
                    }
                }
            }

            // Pull from upstream
            let res = {
                let st = shared
                    .lock()
                    .map_err(|_| "flow:throttle_with => lock error (iterator pull)".to_string())?;
                match &st.upstream.kind {
                    IteratorKind::Core(state_arc) => {
                        let mut g = state_arc
                            .lock()
                            .map_err(|_| "flow:throttle_with => Iterator lock error".to_string())?;
                        if g.done || g.current >= g.end {
                            g.done = true;
                            Err("NO_MORE_DATA".to_string())
                        } else {
                            let v = Value::Int(g.current);
                            g.current += 1;
                            if g.current >= g.end {
                                g.done = true;
                            }
                            Ok(v)
                        }
                    }
                    IteratorKind::Plugin(plugin_arc) => {
                        let mut plugin = plugin_arc
                            .lock()
                            .map_err(|_| "flow:throttle_with => plugin iterator lock error".to_string())?;
                        plugin.next_value()
                    }
                }
            };

            // On success: advance the provider and schedule next_ready
            if res.is_ok() {
                let ms = current_ms(
                    interp,
                    {
                        let st = shared.lock().map_err(|_| {
                            "flow:throttle_with => lock error (iterator get-fn)".to_string()
                        })?;
                        st.period_fn.clone()
                    },
                )?;
                let mut st = shared.lock().map_err(|_| {
                    "flow:throttle_with => lock error (iterator schedule)".to_string()
                })?;
                st.next_ready = Some(Instant::now() + Duration::from_millis(ms));
            }

            res
        })),
        0,
    );

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

fn build_throttle_with_chain(period_fn: Box<FunctionValue>, upstream_tf: Box<FunctionValue>) -> Value {
    use mumu::parser::types::FunctionValue::RustClosure;

    #[derive(Debug)]
    struct Env {
        period_fn: Box<FunctionValue>,
        upstream: Box<FunctionValue>,
        next_ready: Option<Instant>, // None => first item allowed immediately
    }

    let shared = Arc::new(Mutex::new(Env {
        period_fn,
        upstream: upstream_tf,
        next_ready: None,
    }));

    let clos = RustClosure(
        "flow:throttle_with-transform".to_string(),
        Arc::new(Mutex::new(move |interp: &mut Interpreter, _args: Vec<Value>| {
            // Gate without touching the period provider
            {
                let st = shared
                    .lock()
                    .map_err(|_| "flow:throttle_with => lock error (transform gate)".to_string())?;
                if let Some(ready_at) = st.next_ready {
                    if Instant::now() < ready_at {
                        return Err("AGAIN".to_string());
                    }
                }
            }

            // Pull one from upstream
            let v = {
                let st = shared
                    .lock()
                    .map_err(|_| "flow:throttle_with => lock error (transform pull)".to_string())?;
                apply_n_ary_function_value(interp, st.upstream.clone(), vec![])
            }?;

            // Advance the provider and schedule next_ready
            let ms = current_ms(
                interp,
                {
                    let st = shared.lock().map_err(|_| {
                        "flow:throttle_with => lock error (transform get-fn)".to_string()
                    })?;
                    st.period_fn.clone()
                },
            )?;
            {
                let mut st = shared.lock().map_err(|_| {
                    "flow:throttle_with => lock error (transform schedule)".to_string()
                })?;
                st.next_ready = Some(Instant::now() + Duration::from_millis(ms));
            }

            Ok(v)
        })),
        0,
    );

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

/* ──────────────────────────────────────────────────────────────────────────
   Helper: call the period provider (0-arg) and coerce to non-negative ms
   ───────────────────────────────────────────────────────────────────────── */

fn current_ms(interp: &mut Interpreter, pf: Box<FunctionValue>) -> Result<u64, String> {
    let v = apply_n_ary_function_value(interp, pf, vec![])?;
    match v {
        Value::Int(i) if i >= 0 => Ok(i as u64),
        Value::Long(l) if l >= 0 => Ok(l as u64),
        Value::Float(f) if f >= 0.0 => Ok(f as u64),
        Value::Int(i) => Err(format!(
            "flow:throttle_with => period function returned negative int: {}",
            i
        )),
        Value::Long(l) => Err(format!(
            "flow:throttle_with => period function returned negative long: {}",
            l
        )),
        Value::Float(f) => Err(format!(
            "flow:throttle_with => period function returned negative float: {}",
            f
        )),
        other => Err(format!(
            "flow:throttle_with => period function must return number (ms), got {:?}",
            other
        )),
    }
}