diff --git a/src/ast/rtl/streams.rs b/src/ast/rtl/streams.rs index cc023db..618a313 100644 --- a/src/ast/rtl/streams.rs +++ b/src/ast/rtl/streams.rs @@ -1,8 +1,8 @@ use crate::ast::compiler::call_hooks::{InferenceAccess, RtlCompilerHook}; use crate::ast::diagnostics::Diagnostics; use crate::ast::nodes::{Node, NodeKind, TypedNode, TypedPhase}; -use crate::ast::rtl::series::{RingBuffer, ScalarValue, SeriesMember}; -use crate::ast::types::{NativeFunction, NodeIdentity, PipeFn, SourceLocation, StreamStorage, Value}; +use crate::ast::rtl::series::{create_typed_series, RingBuffer, ScalarValue, SeriesMember}; +use crate::ast::types::{NativeFunction, NodeIdentity, PipeFn, SeriesStorage, SourceLocation, StreamStorage, Value}; use crate::ast::vm::{GlobalStore, VM}; use std::cell::RefCell; use std::collections::HashMap; @@ -528,6 +528,227 @@ fn pipe_arg_hint_resolver( } } +// ============================================================================ +// pipe-buffered: Pipe with automatic value accumulation into Series +// ============================================================================ + +/// Compiler hook for `(pipe-buffered lookback inputs lambda)`. +/// +/// - **post_call:** Validates 3 arguments (Int, Streams, Lambda) and checks that +/// the input stream count matches the lambda parameter count. +/// - **finalize:** Replaces the callee with a factory closure that builds a +/// wrapper-executor. The wrapper pushes values into internal Series, checks the +/// fill gate, and forwards Series objects to the user lambda. +struct LookbackPipeHook { + globals: GlobalStore, +} + +impl RtlCompilerHook for LookbackPipeHook { + fn post_call( + &self, + args: &TypedNode, + ret_ty: StaticType, + _ctx: &dyn InferenceAccess, + diag: &mut Diagnostics, + ) -> StaticType { + // Validate: (pipe-buffered Int Streams Lambda) — 3 elements + let NodeKind::Tuple { elements } = &args.kind else { return ret_ty }; + if elements.len() != 3 || !matches!(elements[0].ty, StaticType::Int) { + return ret_ty; + } + + let expected = match &elements[1].ty { + StaticType::Stream(_) | StaticType::Series(_) => 1, + StaticType::Vector(_, count) => *count, + StaticType::Tuple(elems) => elems.len(), + _ => return ret_ty, + }; + + if let NodeKind::Lambda { params, .. } = &elements[2].kind { + let actual = match ¶ms.kind { + NodeKind::Tuple { elements } => elements.len(), + _ => 1, + }; + if actual != expected { + diag.push_error( + format!( + "pipe-buffered: lambda expects {} parameter(s) but {} input stream(s) provided", + actual, expected + ), + Some(elements[2].identity.clone()), + ); + } + } + ret_ty + } + + fn finalize( + &self, + _callee: Rc, + args: Rc, + node_ty: &StaticType, + _subst: &HashMap, + ) -> Option> { + let StaticType::Stream(inner) = node_ty else { return None }; + if matches!(inner.as_ref(), StaticType::Any | StaticType::TypeVar(_)) { + return None; + } + let element_type = *inner.clone(); + let globals = self.globals.clone(); + let input_element_types = extract_input_element_types(&args); + + let factory: Value = Value::Function(Rc::new(NativeFunction { + func: Rc::new(move |call_args: &[Value]| { + let lookback = call_args[0].as_int().unwrap() as usize; + let obs = extract_obs_streams(&call_args[1]); + let user_exec = build_pipe_executor(&call_args[2], globals.clone()); + let wrapper = build_buffered_wrapper(lookback, &input_element_types, user_exec); + Value::Stream(build_pipeline_node(obs, wrapper, &element_type)) + }), + purity: Purity::Impure, + })); + + let factory_node = Rc::new(Node { + kind: NodeKind::Constant(factory), + ty: StaticType::Any, + identity: NodeIdentity::new(SourceLocation { line: 0, col: 0 }), + comments: Rc::from([]), + }); + Some(NodeKind::Call { callee: factory_node, args }) + } +} + +/// Builds a wrapper executor that accumulates values into internal Series +/// and only calls the user lambda once the fill gate is satisfied. +fn build_buffered_wrapper( + lookback: usize, + input_element_types: &[StaticType], + mut user_exec: Box, +) -> Box { + let series: Vec> = input_element_types + .iter() + .map(|ty| create_typed_series(ty, lookback)) + .collect(); + let mut fill_gate_open = false; + + Box::new(move |args: &[Value]| { + // 1. Push incoming values into internal series + for (i, s) in series.iter().enumerate() { + s.as_pushable().unwrap().push_value(args[i].clone()); + } + + // 2. Fill gate: wait until all series have enough data + if !fill_gate_open { + if series.iter().all(|s| s.len() >= lookback) { + fill_gate_open = true; + } else { + return Value::Void; // Filtered by PipeStream::notify + } + } + + // 3. Call user lambda with Series objects instead of raw values + let series_args: Vec = series + .iter() + .map(|s| Value::Series(Rc::clone(s))) + .collect(); + user_exec(&series_args) + }) +} + +/// Extracts input element types from the typed AST for `pipe-buffered`. +/// Args layout: `Tuple([Int, inputs, Lambda])` — inputs at index 1. +fn extract_input_element_types(args: &TypedNode) -> Vec { + let NodeKind::Tuple { elements } = &args.kind else { return vec![] }; + if elements.len() < 2 { + return vec![]; + } + let input_ty = &elements[1].ty; + + fn inner_type(ty: &StaticType) -> StaticType { + match ty { + StaticType::Stream(inner) | StaticType::Series(inner) => *inner.clone(), + _ => StaticType::Any, + } + } + + match input_ty { + StaticType::Stream(_) | StaticType::Series(_) => vec![inner_type(input_ty)], + StaticType::Vector(elem_ty, count) => vec![inner_type(elem_ty); *count], + StaticType::Tuple(elems) => elems.iter().map(inner_type).collect(), + _ => vec![StaticType::Any], + } +} + +/// Extracts element types from runtime stream values (for the native fn fallback). +fn extract_runtime_input_types(val: &Value) -> Vec { + match val { + Value::Tuple(elements) => elements + .iter() + .map(|v| match v { + Value::Stream(s) => s.element_type(), + _ => StaticType::Any, + }) + .collect(), + Value::Stream(s) => vec![s.element_type()], + _ => vec![StaticType::Any], + } +} + +/// Resolves the return type of `pipe-buffered` from its argument types. +/// Argument layout: `Tuple([Int, inputs, Lambda])`. +/// Returns `Stream`, unwrapping `Optional` to `U` (filter pattern). +fn buffered_pipe_type_resolver(args_ty: &StaticType) -> Option { + let StaticType::Tuple(elements) = args_ty else { return None }; + if elements.len() != 3 || !matches!(elements[0], StaticType::Int) { + return None; + } + let StaticType::Function(sig) = &elements[2] else { return None }; + + let inner = if let StaticType::Optional(inner) = &sig.ret { + *inner.clone() + } else { + sig.ret.clone() + }; + Some(StaticType::Stream(Box::new(inner))) +} + +/// Provides expected lambda parameter types for `pipe-buffered`. +/// Lambda is at arg index 2. Parameters are wrapped as `Series` instead of raw `T`. +fn buffered_pipe_arg_hint_resolver( + arg_index: usize, + known_args: &[Option], +) -> Option> { + // Lambda is at position 2; only provide hints for that position + if arg_index != 2 { + return None; + } + // Inputs are at position 1 + let input_ty = known_args.get(1)?.as_ref()?; + + fn extract_inner(ty: &StaticType) -> StaticType { + match ty { + StaticType::Stream(inner) | StaticType::Series(inner) => *inner.clone(), + _ => StaticType::Any, + } + } + + // Wrap each inner type in Series — the lambda sees Series, not raw values + let wrap = |t: StaticType| StaticType::Series(Box::new(t)); + + match input_ty { + StaticType::Stream(_) | StaticType::Series(_) => { + Some(vec![wrap(extract_inner(input_ty))]) + } + StaticType::Vector(elem_ty, count) => { + Some(vec![wrap(extract_inner(elem_ty)); *count]) + } + StaticType::Tuple(elements) => { + Some(elements.iter().map(|e| wrap(extract_inner(e))).collect()) + } + _ => None, + } +} + pub fn register(env: &Environment) { // (create-random-ohlc seed limit) -> StreamNode let generators = env.pipeline_generators.clone(); @@ -717,6 +938,36 @@ pub fn register(env: &Environment) { "(pipe ohlc (fn [bar] (.close bar)))", "(pipe [stream-a stream-b] (fn [a b] (+ a b)))", ]); + + // (pipe-buffered lookback inputs lambda) -> StreamNode + // Like pipe, but accumulates values into internal Series with the given lookback + // depth. The lambda receives Series objects and only fires once all series have + // at least `lookback` elements (fill gate). + let fn_globals = env.global_store(); + let hook_globals = fn_globals.clone(); + env.register_native_fn( + "pipe-buffered", + StaticType::PolymorphicFn { + resolve_return: buffered_pipe_type_resolver, + resolve_arg_hints: Some(buffered_pipe_arg_hint_resolver), + }, + Purity::Impure, + move |args: &[Value]| { + assert!(args.len() == 3, "pipe-buffered expects 3 arguments (lookback, inputs, lambda)"); + let lookback = args[0].as_int().unwrap() as usize; + let obs = extract_obs_streams(&args[1]); + let input_types: Vec = extract_runtime_input_types(&args[1]); + let user_exec = build_pipe_executor(&args[2], fn_globals.clone()); + let wrapper = build_buffered_wrapper(lookback, &input_types, user_exec); + Value::Stream(build_pipeline_node(obs, wrapper, &StaticType::Any)) + }, + ).with_compiler_hook(Rc::new(LookbackPipeHook { globals: hook_globals })) + .doc("Like pipe, but accumulates values into Series before firing the lambda.") + .description("The lambda only fires once all input series have at least N elements (fill gate). Lambda parameters are Series objects with lookback indexing (0 = newest).") + .examples(&[ + "(pipe-buffered 20 ohlc (fn [bars] (- (.close (bars 0)) (.close (bars 19)))))", + "(pipe-buffered 2 [stream-a stream-b] (fn [a b] (+ (a 0) (b 0))))", + ]); } #[cfg(test)] diff --git a/tests/pipeline.rs b/tests/pipeline.rs index a5c7ca1..85a1984 100644 --- a/tests/pipeline.rs +++ b/tests/pipeline.rs @@ -200,9 +200,10 @@ fn test_pipe_buffered_lambda_params_are_series() { )"; let dump = env.dump_ast(source).expect("dump_ast failed"); - // Lambda param should be Series, not just Record or Any + // Lambda param should be Series, not just Record or Any. + // The AST dump uses Debug format: "Series(Record(...))" assert!( - dump.contains("series<"), + dump.contains("Series(Record("), "Lambda param should be typed as Series, not raw value.\nDump:\n{dump}" ); }