flow-mumu 0.2.0-rc.9

// src/drain_then.rs
//
// flow:drain_then(source, callback)
//   • Works like flow:drain_async, but the callback is **required**.
//   • Spawns a background poller that drains `source` in micro-bursts per tick.
//   • When EOF is reached for the **entire upstream pipeline** (i.e., the composed
//     transform feeding this sink returns NO_MORE_DATA), we invoke:
//         callback(total_count)
//   • Returns immediately (Int(0)) so it can be used inside flow:compose.
//
// Tuning (same as drain_async):
//   • LAVA_FLOW_BURST     — max items processed per poll tick (default: 64)
//   • LAVA_FLOW_BUDGET_US — per-tick time budget in microseconds (default: 500)
//
// Partial application patterns:
//   • flow:drain_then(_, cb)      -> waits for source, uses cb
//   • flow:drain_then(cb)         -> returns a (source) => … function
//   • flow:drain_then(source, cb) -> starts draining immediately
//
// IMPORTANT — cooperative throttling and keep-alive (matches flow:drain_async):
//   • If upstream returns **"AGAIN"** (temporary no-data), this is **not** EOF.
//   • We stop this tick (no busy-wait), and if we processed nothing but saw
//     "AGAIN", we return **1** as a keep-alive so the interpreter does not
//     prematurely idle out the pollers just before the throttle window opens.

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

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

use crate::common::{is_placeholder, to_transform};

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

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

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

        // 1-arg form is ambiguous: treat a Function as the *callback* (preferred for compose),
        // a Placeholder as “still waiting”, and any other value as the *source*.
        1 => {
            let a = &args[0];
            if is_placeholder(a) {
                Ok(make_partial(None, None))
            } else if let Value::Function(fb) = a {
                // User supplied callback first: return ctor waiting for source
                Ok(make_partial(None, Some(Value::Function(fb.clone()))))
            } else {
                // User supplied source first: return partial waiting for callback
                Ok(make_partial(Some(a.clone()), None))
            }
        }

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

            // Handle placeholders via partial
            if is_placeholder(src) || is_placeholder(cb) {
                let s_opt = if is_placeholder(src) { None } else { Some(src.clone()) };
                let c_opt = if is_placeholder(cb)  { None } else { Some(cb.clone())  };
                return Ok(make_partial(s_opt, c_opt));
            }

            // Finalize only when callback is a Function
            let cb_fn = match cb {
                Value::Function(fb) => fb.clone(),
                other => {
                    return Err(format!(
                        "flow:drain_then => callback must be Function, got {:?}",
                        other
                    ))
                }
            };
            spawn_drain_then(interp, src.clone(), cb_fn)
        }

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

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

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

            for a in new_args {
                if s.is_none() {
                    if !is_placeholder(&a) {
                        s = Some(a);
                    }
                    continue;
                }
                if c.is_none() {
                    if !is_placeholder(&a) {
                        c = Some(a);
                    }
                    continue;
                }
                return Err("flow:drain_then => partial => too many arguments".to_string());
            }

            match (s.clone(), c.clone()) {
                (Some(source), Some(cb_val)) => {
                    let cb_fn = match cb_val {
                        Value::Function(fb) => fb,
                        other => {
                            return Err(format!(
                                "flow:drain_then => callback must be Function, got {:?}",
                                other
                            ))
                        }
                    };
                    spawn_drain_then(interp, source, cb_fn)
                }
                _ => Ok(make_partial(s, c)),
            }
        })),
        0,
    )))
}

fn spawn_drain_then(
    interp: &mut Interpreter,
    source: Value,
    cb_fn: Box<FunctionValue>,
) -> Result<Value, String> {
    // Normalize Iterator or zero-arg Function into a pullable transform.
    // This transform represents the **entire upstream compose chain** feeding us.
    // When it returns NO_MORE_DATA, the whole compose is done for this run.
    let tf = to_transform(source)?;

    // Shared state with the poller
    struct DrainState {
        tf: Box<FunctionValue>,
        count: i32,
        done: bool,
        callback: Box<FunctionValue>,
        callback_fired: bool,
    }

    let state = Arc::new(Mutex::new(DrainState {
        tf,
        count: 0,
        done: false,
        callback: cb_fn,
        callback_fired: false,
    }));

    // Micro-burst tuning (identical to flow:drain_async)
    let burst_per_tick: usize = std::env::var("LAVA_FLOW_BURST")
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(64);
    let budget_us: u64 = std::env::var("LAVA_FLOW_BUDGET_US")
        .ok()
        .and_then(|s| s.parse::<u64>().ok())
        .unwrap_or(500);
    let time_budget = Duration::from_micros(budget_us);

    // Poller: drain up to `burst_per_tick` items or until `time_budget` elapses.
    // The callback is fired only once, when the upstream transform signals EOF.
    let st_for_poller = Arc::clone(&state);
    let poller: PollerFn = Arc::new(Mutex::new(move |interp: &mut Interpreter| {
        let tick_start = Instant::now();
        let mut processed = 0usize;
        let mut saw_again = false;

        // Safety: if already done and not fired yet (edge case), fire 0 now.
        {
            let mut guard = st_for_poller.lock().unwrap();
            if guard.done && !guard.callback_fired {
                let total = guard.count;
                let cb    = guard.callback.clone();
                guard.callback_fired = true;
                drop(guard);
                let _ = apply_n_ary_function_value(interp, cb, vec![Value::Int(total)]);
                return 0;
            }
        }

        loop {
            // Pull one without holding the lock during the call.
            let tf_clone = {
                let guard = st_for_poller.lock().unwrap();
                if guard.done {
                    break;
                }
                guard.tf.clone()
            };

            match apply_n_ary_function_value(interp, tf_clone, vec![]) {
                Ok(_v) => {
                    let mut guard = st_for_poller.lock().unwrap();
                    guard.count = guard.count.saturating_add(1);
                    drop(guard);
                    processed += 1;
                }
                Err(e) if e == "AGAIN" => {
                    // Temporary no-data; not EOF. Note for keep-alive and stop this tick.
                    saw_again = true;
                    break;
                }
                Err(e) if e == "NO_MORE_DATA" => {
                    let mut guard = st_for_poller.lock().unwrap();
                    guard.done = true;
                    let total = guard.count;
                    let cb    = guard.callback.clone();
                    guard.callback_fired = true;
                    drop(guard);

                    // Fire completion callback with total drain count.
                    let _ = apply_n_ary_function_value(interp, cb, vec![Value::Int(total)]);
                    break;
                }
                Err(_e) => {
                    // Treat any other error as terminal for safety (no callback).
                    let mut guard = st_for_poller.lock().unwrap();
                    guard.done = true;
                    drop(guard);
                    break;
                }
            }

            if processed >= burst_per_tick || tick_start.elapsed() >= time_budget {
                break;
            }
        }

        // Keep-alive: if this tick saw only "AGAIN", report 1 so the outer
        // idle-round logic does not cut us off just before the throttle window.
        if processed == 0 && saw_again {
            return 1;
        }

        processed
    }));

    interp.add_poller(poller);
    // Non-blocking return; friendly for compose/pipe usage.
    Ok(Value::Int(0))
}