From a5c3f3da0439450cf4a3b98b2832e54e0a99c5a1 Mon Sep 17 00:00:00 2001 From: Brummel Date: Mon, 23 Mar 2026 10:51:06 +0100 Subject: [PATCH] feat: Add bidirectional type inference for lambdas This commit introduces bidirectional type inference, enabling the compiler to infer lambda parameter types based on the context of their usage. This is particularly useful for functions like `pipe`, where a lambda's signature can be deduced from the types of the streams it operates on. Key changes include: - **`extract_lambda_param_hints` function:** This new helper function in `type_checker.rs` is responsible for extracting expected parameter types for a lambda based on the callee's signature and already known argument types. - **`check_lambda_with_param_hints` function:** This function allows type-checking a lambda node using provided parameter type hints, preserving upvalue types from the enclosing scope. - **Call expression type checking:** The `check_call` function now phases its argument type checking. Non-lambda arguments are typed first, and then lambda arguments are typed using hints derived from `extract_lambda_param_hints`. - **`pipe_arg_hint_resolver`:** This new resolver for the `pipe` function's `PolymorphicFn` type allows it to provide specific type hints for its lambda argument based on the types of its stream inputs (single stream, vector of streams, or tuple of streams). - **`StaticType::PolymorphicFn` update:** The `PolymorphicFn` enum variant has been extended to include `resolve_arg_hints`, enabling functions to provide lambda parameter hints during bidirectional type inference. - **New tests:** Added tests to verify the correct inference of lambda parameters for `pipe` with vector and tuple inputs, ensuring that inner stream types are correctly propagated. Also, a test to confirm `pipe` with an optional return still produces `Stream(Float)`. --- src/ast/compiler/type_checker.rs | 164 +++++++++++++++++++++++++++++-- src/ast/rtl/streams.rs | 39 +++++++- src/ast/types.rs | 18 +++- tests/pipeline.rs | 43 ++++++++ 4 files changed, 252 insertions(+), 12 deletions(-) diff --git a/src/ast/compiler/type_checker.rs b/src/ast/compiler/type_checker.rs index b4c9434..c747d44 100644 --- a/src/ast/compiler/type_checker.rs +++ b/src/ast/compiler/type_checker.rs @@ -73,6 +73,62 @@ impl<'a> TypeContext<'a> { } } +/// Extracts expected lambda parameter types from the callee type for a specific argument position. +/// Used for bidirectional type inference: when a Call's callee expects a function at `arg_index`, +/// this returns the expected parameter types for that function, derived from the callee's signature +/// and the already-known types of non-lambda arguments. +fn extract_lambda_param_hints( + callee_ty: &StaticType, + arg_index: usize, + known_arg_types: &[Option], +) -> Option> { + /// Helper: given the expected type at a specific parameter position in a signature, + /// extract the lambda parameter types if it expects a function. + fn hints_from_param_type(param_ty: &StaticType) -> Option> { + if let StaticType::Function(sig) = param_ty { + if let StaticType::Tuple(params) = &sig.params { + return Some(params.clone()); + } + return Some(vec![sig.params.clone()]); + } + None + } + + /// Helper: extract the expected type at `arg_index` from a signature's params tuple. + fn param_at(sig: &Signature, arg_index: usize) -> Option<&StaticType> { + if let StaticType::Tuple(params) = &sig.params { + params.get(arg_index) + } else if arg_index == 0 { + Some(&sig.params) + } else { + None + } + } + + match callee_ty { + StaticType::Function(sig) => { + let expected = param_at(sig, arg_index)?; + hints_from_param_type(expected) + } + StaticType::FunctionOverloads(sigs) => { + // Try each overload — return hints from the first one that has a function at this position + for sig in sigs { + if let Some(expected) = param_at(sig, arg_index) + && let Some(hints) = hints_from_param_type(expected) + { + return Some(hints); + } + } + None + } + StaticType::PolymorphicFn { + resolve_arg_hints: Some(resolver), + .. + } => resolver(arg_index, known_arg_types), + _ => None, + } +} + pub struct TypeChecker { root_types: Rc>>, } @@ -159,6 +215,63 @@ impl TypeChecker { } } + /// Types a lambda node using externally provided parameter type hints, + /// while preserving the current scope's upvalue types. + /// Unlike `check_node_as_bound`, this keeps the enclosing `TypeContext` as parent, + /// so captured variables retain their inferred types. + fn check_lambda_with_param_hints( + &self, + node: &Node

, + param_hints: &[StaticType], + ctx: &mut TypeContext, + diag: &mut Diagnostics, + ) -> TypedNode { + let NodeKind::Lambda { + params, + body, + info, + } = &node.kind + else { + return self.check_node(node, ctx, diag); + }; + + let upvalues = &info.upvalues; + let positional_count = info.positional_count; + + let mut upvalue_types = Vec::with_capacity(upvalues.len()); + for &addr in upvalues { + upvalue_types.push(ctx.get_type(addr)); + } + + let mut lambda_ctx = TypeContext::new(64, upvalue_types, ctx.root_types, Some(ctx)); + + let hint_ty = StaticType::Tuple(param_hints.to_vec()); + let params_typed = self.check_params(params.as_ref(), &hint_ty, &mut lambda_ctx, diag); + + lambda_ctx.current_params_ty = Some(params_typed.ty.clone()); + + let body_typed = self.check_node(body, &mut lambda_ctx, diag); + let ret_ty = body_typed.ty.clone(); + + let fn_ty = StaticType::Function(Box::new(Signature { + params: params_typed.ty.clone(), + ret: ret_ty, + })); + + Node { + identity: node.identity.clone(), + kind: NodeKind::Lambda { + params: Rc::new(params_typed), + body: Rc::new(body_typed), + info: LambdaBinding { + upvalues: upvalues.clone(), + positional_count, + }, + }, + ty: fn_ty, + } + } + fn check_params( &self, node: &Node

, @@ -613,17 +726,54 @@ impl TypeChecker { let callee_typed = self.check_node(callee, ctx, diag); let args_typed = if let NodeKind::Tuple { elements } = &args.kind { - let mut typed_elements = Vec::new(); - let mut elem_types = Vec::new(); - for e in elements { - let t = self.check_node(e, ctx, diag); - elem_types.push(t.ty.clone()); - typed_elements.push(Rc::new(t)); + let arg_count = elements.len(); + let mut typed_elements: Vec>> = vec![None; arg_count]; + let mut known_types: Vec> = vec![None; arg_count]; + let mut lambda_indices = Vec::new(); + + // Phase 1: Type non-lambda arguments first + for (i, e) in elements.iter().enumerate() { + if matches!(e.kind, NodeKind::Lambda { .. }) { + lambda_indices.push(i); + } else { + let t = self.check_node(e, ctx, diag); + known_types[i] = Some(t.ty.clone()); + typed_elements[i] = Some(Rc::new(t)); + } } + + // Phase 2: Type lambda arguments with parameter hints (if available) + for i in lambda_indices { + let hints = extract_lambda_param_hints( + &callee_typed.ty, + i, + &known_types, + ); + let t = if let Some(param_types) = hints { + self.check_lambda_with_param_hints( + &elements[i], + ¶m_types, + ctx, + diag, + ) + } else { + self.check_node(&elements[i], ctx, diag) + }; + known_types[i] = Some(t.ty.clone()); + typed_elements[i] = Some(Rc::new(t)); + } + + let final_elements: Vec> = typed_elements + .into_iter() + .map(|e| e.expect("all args should be typed")) + .collect(); + let elem_types: Vec = + final_elements.iter().map(|e| e.ty.clone()).collect(); + Node { identity: args.identity.clone(), kind: NodeKind::Tuple { - elements: typed_elements, + elements: final_elements, }, ty: StaticType::Tuple(elem_types), } diff --git a/src/ast/rtl/streams.rs b/src/ast/rtl/streams.rs index 9d3c774..af7cece 100644 --- a/src/ast/rtl/streams.rs +++ b/src/ast/rtl/streams.rs @@ -342,6 +342,40 @@ fn pipe_type_resolver(args_ty: &StaticType) -> Option { None } +/// Provides expected lambda parameter types for bidirectional type inference. +/// For `pipe`, the lambda at position 1 receives the inner type of the input stream(s). +fn pipe_arg_hint_resolver( + arg_index: usize, + known_args: &[Option], +) -> Option> { + if arg_index != 1 { + return None; + } + let input_ty = known_args.first()?.as_ref()?; + /// Extract the inner type from a single stream or series. + fn extract_inner(ty: &StaticType) -> StaticType { + match ty { + StaticType::Stream(inner) | StaticType::Series(inner) => *inner.clone(), + _ => StaticType::Any, + } + } + + match input_ty { + StaticType::Stream(_) | StaticType::Series(_) => { + Some(vec![extract_inner(input_ty)]) + } + // Vector: homogeneous fixed-size array, e.g. `[src]` → Vector(Stream, 1) + StaticType::Vector(elem_ty, count) => { + Some(vec![extract_inner(elem_ty); *count]) + } + // Tuple: heterogeneous, e.g. `[src1 src2]` with different stream types + StaticType::Tuple(elements) => { + Some(elements.iter().map(extract_inner).collect()) + } + _ => None, + } +} + pub fn register(env: &Environment) { // (create-random-ohlc seed limit) -> StreamNode let generators = env.pipeline_generators.clone(); @@ -499,7 +533,10 @@ pub fn register(env: &Environment) { let globals = env.root_values.clone(); env.register_native_fn( "pipe", - StaticType::PolymorphicFn(pipe_type_resolver), + StaticType::PolymorphicFn { + resolve_return: pipe_type_resolver, + resolve_arg_hints: Some(pipe_arg_hint_resolver), + }, Purity::Impure, move |args: &[Value]| { assert!(args.len() == 2, "pipe expects exactly 2 arguments (inputs, lambda)"); diff --git a/src/ast/types.rs b/src/ast/types.rs index ae77722..4155a49 100644 --- a/src/ast/types.rs +++ b/src/ast/types.rs @@ -314,6 +314,11 @@ pub struct Signature { pub ret: StaticType, } +/// Callback that provides expected parameter types for a lambda argument +/// during bidirectional type inference. Receives `(arg_index, known_arg_types)` +/// and returns the expected lambda parameter types. +pub type ArgHintResolver = fn(usize, &[Option]) -> Option>; + #[derive(Debug, Clone, PartialEq, Eq, Hash)] #[allow(unpredictable_function_pointer_comparisons)] pub enum StaticType { @@ -338,8 +343,13 @@ pub enum StaticType { FunctionOverloads(Vec), Object(&'static str), /// A polymorphic native function whose return type is computed from its argument types. - /// The function pointer receives the full argument type and returns the resolved return type. - PolymorphicFn(fn(&StaticType) -> Option), + /// `resolve_return` receives the full argument type and returns the resolved return type. + /// `resolve_arg_hints` optionally provides expected parameter types for lambda arguments + /// during bidirectional type inference in the Call handler. + PolymorphicFn { + resolve_return: fn(&StaticType) -> Option, + resolve_arg_hints: Option, + }, /// A diagnostic poison type, allowing type-checking to continue after an error. Error, } @@ -398,7 +408,7 @@ impl fmt::Display for StaticType { write!(f, "overloads({} variants)", sigs.len()) } StaticType::Object(name) => write!(f, "{}", name), - StaticType::PolymorphicFn(_) => write!(f, ""), + StaticType::PolymorphicFn { .. } => write!(f, ""), StaticType::Error => write!(f, ""), } } @@ -521,7 +531,7 @@ impl StaticType { .find(|sig| sig.params == *args_ty) // 1. Try exact match first .or_else(|| sigs.iter().find(|sig| sig.params.is_assignable_from(args_ty))) // 2. Fallback to implicit coercion .map(|sig| sig.ret.clone()), - StaticType::PolymorphicFn(resolver) => resolver(args_ty), + StaticType::PolymorphicFn { resolve_return, .. } => resolve_return(args_ty), _ => None, } } diff --git a/tests/pipeline.rs b/tests/pipeline.rs index ec0aac2..8da5557 100644 --- a/tests/pipeline.rs +++ b/tests/pipeline.rs @@ -1,6 +1,49 @@ use myc::ast::environment::Environment; use myc::ast::types::Value; +/// Verifies that bidirectional type inference correctly propagates stream inner types +/// into lambda parameters for `pipe` calls with a single stream in a vector `[src]`. +/// The lambda parameter should be inferred as the stream's Record type, not `Any`. +#[test] +fn test_pipe_infers_lambda_param_from_vector_input() { + let env = Environment::new(); + let source = "(do + (def src (create-random-ohlc 42 5)) + (pipe [src] (fn [tick] (.close tick))) + )"; + let dump = env.dump_ast(source).expect("dump_ast failed"); + + // The lambda parameter `tick` should be inferred as the OHLC Record type, + // NOT as `Any`. This validates the Vector → Stream → T hint resolution. + assert!( + !dump.contains("params: Tuple([Any])"), + "Lambda params should NOT be Any — bidirectional inference failed.\nDump:\n{dump}" + ); + // The pipe call should return Stream(Float), not Stream(Any) + assert!( + dump.contains("Stream(Float)"), + "Pipe should produce Stream(Float), not Stream(Any).\nDump:\n{dump}" + ); + + // Also verify it runs correctly + let res = env.run_script(source); + assert!(res.is_ok(), "Script failed: {:?}", res.err()); +} + +/// Verifies bidirectional inference for pipe with multiple stream inputs in a tuple. +#[test] +fn test_pipe_infers_lambda_params_from_tuple_inputs() { + let env = Environment::new(); + let source = "(do + (def src1 (create-random-ohlc 42 5)) + (def src2 (create-random-ohlc 99 5)) + (pipe [src1 src2] (fn [t1 t2] (+ (.close t1) (.close t2)))) + )"; + let res = env.run_script(source); + assert!(res.is_ok(), "Script failed: {:?}", res.err()); +} + +/// Verifies that pipe with Optional return (filter pattern) still produces Stream(Float). #[test] fn test_pipeline_optional_type() { let env = Environment::new();