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)`.
This commit is contained in:
2026-03-23 10:51:06 +01:00
parent 1875cfdc9b
commit a5c3f3da04
4 changed files with 252 additions and 12 deletions
+156 -6
View File
@@ -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<StaticType>],
) -> Option<Vec<StaticType>> {
/// 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<Vec<StaticType>> {
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<std::cell::RefCell<Vec<StaticType>>>,
}
@@ -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<P: BoundLike>(
&self,
node: &Node<P>,
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<P: BoundLike>(
&self,
node: &Node<P>,
@@ -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 arg_count = elements.len();
let mut typed_elements: Vec<Option<Rc<TypedNode>>> = vec![None; arg_count];
let mut known_types: Vec<Option<StaticType>> = 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);
elem_types.push(t.ty.clone());
typed_elements.push(Rc::new(t));
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],
&param_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<Rc<TypedNode>> = typed_elements
.into_iter()
.map(|e| e.expect("all args should be typed"))
.collect();
let elem_types: Vec<StaticType> =
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),
}
+38 -1
View File
@@ -342,6 +342,40 @@ fn pipe_type_resolver(args_ty: &StaticType) -> Option<StaticType> {
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<StaticType>],
) -> Option<Vec<StaticType>> {
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<T>, 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)");
+14 -4
View File
@@ -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<StaticType>]) -> Option<Vec<StaticType>>;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[allow(unpredictable_function_pointer_comparisons)]
pub enum StaticType {
@@ -338,8 +343,13 @@ pub enum StaticType {
FunctionOverloads(Vec<Signature>),
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<StaticType>),
/// `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<StaticType>,
resolve_arg_hints: Option<ArgHintResolver>,
},
/// 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, "<polymorphic-fn>"),
StaticType::PolymorphicFn { .. } => write!(f, "<polymorphic-fn>"),
StaticType::Error => write!(f, "<error>"),
}
}
@@ -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,
}
}
+43
View File
@@ -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> → 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();