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();