Add pipe-buffered function
Implements `pipe-buffered`, a new primitive that allows accumulating values from multiple input streams into internal series buffers before executing a user-provided lambda. This is useful for implementing lookback logic and ensuring that the lambda only executes when sufficient data is available in all input streams. The function takes a `lookback` integer, input streams, and a lambda. The lambda receives `Series` objects for each input, allowing indexed access to past values. The `pipe-buffered` function manages the buffering and fill-gate logic.
This commit is contained in:
+253
-2
@@ -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<TypedNode>,
|
||||
args: Rc<TypedNode>,
|
||||
node_ty: &StaticType,
|
||||
_subst: &HashMap<u32, StaticType>,
|
||||
) -> Option<NodeKind<TypedPhase>> {
|
||||
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<PipeFn>,
|
||||
) -> Box<PipeFn> {
|
||||
let series: Vec<Rc<dyn SeriesStorage>> = 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<Value> = 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<StaticType> {
|
||||
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<StaticType> {
|
||||
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<U>`, unwrapping `Optional<U>` to `U` (filter pattern).
|
||||
fn buffered_pipe_type_resolver(args_ty: &StaticType) -> Option<StaticType> {
|
||||
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<T>` instead of raw `T`.
|
||||
fn buffered_pipe_arg_hint_resolver(
|
||||
arg_index: usize,
|
||||
known_args: &[Option<StaticType>],
|
||||
) -> Option<Vec<StaticType>> {
|
||||
// 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<T> — 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<StaticType> = 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)]
|
||||
|
||||
+3
-2
@@ -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<Record>, not just Record or Any
|
||||
// Lambda param should be Series<Record>, 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}"
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user