🐛 fix(check): type argument selectors on unknown input and self[...]

crates/mq-check/src/constraint.rs

@@ -1688,6 +1688,12 @@ pub(super) fn generate_symbol_constraints(
         // Attr selectors (`Selector::Attr`) return the concrete attribute type via `attr_kind_to_type`.
         // Non-Attr selectors (`.h1`, `.code`, etc.) always return `Type::markdown()`.
         SymbolKind::Selector(ref selector) => {
+            // An argument selector runs on the unknown input, so it may select from a single
+            // node or from an array of nodes (`nodes | compact(.code)`).
+            let varies_with_input = ctx.get_piped_input(symbol_id).is_none()
+                && matches!(ctx.input_type(), Type::Dynamic)
+                && !matches!(selector, mq_lang::Selector::Attr(_) | mq_lang::Selector::Recursive);
+
             // Compute this selector's own output type based on the incoming piped input.
             let own_type = if let Some(piped_ty) = ctx.get_piped_input(symbol_id).cloned() {
                 let resolved = ctx.resolve_type(&piped_ty);
@@ -1757,6 +1763,8 @@ pub(super) fn generate_symbol_constraints(
                 // an array of all descendant Markdown nodes.
                 if let mq_lang::Selector::Attr(attr_kind) = selector {
                     attr_kind_to_type(attr_kind)
+                } else if varies_with_input {
+                    Type::Var(ctx.fresh_var())
                 } else if matches!(selector, mq_lang::Selector::Recursive) {
                     Type::array(Type::markdown())
                 } else {
@@ -1779,8 +1787,14 @@ pub(super) fn generate_symbol_constraints(
             if child_selectors.is_empty() {
                 ctx.set_symbol_type(symbol_id, own_type);
             } else {
-                // Thread the output type through the chain of child selectors.
-                ctx.set_piped_input(child_selectors[0], own_type);
+                // Thread the output type through the chain of child selectors. The children of
+                // an argument selector are checked against the selected node kind.
+                let chain_input = if varies_with_input {
+                    selector_kinds(selector).map_or_else(Type::markdown, Type::Node)
+                } else {
+                    own_type
+                };
+                ctx.set_piped_input(child_selectors[0], chain_input);
                 for i in 0..child_selectors.len() {
                     if let Some(child_sym) = hir.symbol(child_selectors[i]) {
                         generate_symbol_constraints(
@@ -1797,7 +1811,11 @@ pub(super) fn generate_symbol_constraints(
                     }
                 }
                 // The parent's type is the last child's type.
-                let last_ty = ctx.get_or_create_symbol_type(*child_selectors.last().unwrap());
+                let last_ty = if varies_with_input {
+                    Type::Var(ctx.fresh_var())
+                } else {
+                    ctx.get_or_create_symbol_type(*child_selectors.last().unwrap())
+                };
                 ctx.set_symbol_type(symbol_id, last_ty);
             }
         }
@@ -1953,6 +1971,99 @@ fn is_index_assign_target(symbol_id: SymbolId, hir: &Hir, children_index: &Child
         .is_some_and(|parent| get_children(children_index, parent).first() == Some(&symbol_id))
 }
 
+/// Constrains the bracket access `v[i]` or `v[a:b]` on a value of type `container_ty`.
+///
+/// Two args in one bracket group are a range slice `v[start:end]` lowered as `v(start, end)`.
+/// One arg is an element access `v[i]` lowered as `v(i)`. Chained accesses are lowered the same
+/// way: `v[0][1]` is `v(0, 1)` and `v[0][:key]` is `v(0, :key)`. The bracket group count tells
+/// them apart from a slice. A String/Symbol literal after the first argument can never be a
+/// slice bound, so it marks such a chain. A chain mixed with slices is ambiguous and left
+/// unconstrained.
+fn constrain_bracket_access(
+    hir: &Hir,
+    ctx: &mut InferenceContext,
+    symbol_id: SymbolId,
+    children: &[SymbolId],
+    explicit_arg_tys: &[Type],
+    container_ty: Type,
+    range: Option<mq_lang::Range>,
+) {
+    let groups = hir.bracket_group_count(symbol_id);
+    let is_chained_access = children.iter().skip(1).any(|&child_id| {
+        hir.symbol(child_id)
+            .is_some_and(|s| matches!(s.kind, SymbolKind::String | SymbolKind::Symbol))
+    });
+    if is_chained_access || (groups > 1 && children.len() != groups) {
+        // The container's element type is not tracked through the chain, so leave the
+        // result unconstrained.
+        let ty_var = ctx.fresh_var();
+        ctx.set_symbol_type(symbol_id, Type::Var(ty_var));
+    } else if groups > 1 {
+        // Element access chain: each index applies to the previous result.
+        let mut container_ty = container_ty;
+        for (i, (&arg_id, arg_ty)) in children.iter().zip(explicit_arg_tys).enumerate() {
+            let result_ty = Type::Var(ctx.fresh_var());
+            constrain_element_access(hir, ctx, container_ty, arg_id, arg_ty, result_ty.clone(), range);
+            if i + 1 == children.len() {
+                ctx.set_symbol_type(symbol_id, result_ty);
+                break;
+            }
+            container_ty = result_ty;
+        }
+    } else if children.len() == 2 {
+        // Range slice: delegate to the "slice" builtin.
+        let mut slice_arg_tys = vec![container_ty];
+        slice_arg_tys.extend(explicit_arg_tys.iter().cloned());
+        let defer = might_receive_piped_input(hir, symbol_id);
+        resolve_builtin_call(ctx, symbol_id, "slice", &slice_arg_tys, range, defer);
+    } else {
+        let result_ty = Type::Var(ctx.fresh_var());
+        ctx.set_symbol_type(symbol_id, result_ty.clone());
+        if let (Some(&arg_id), Some(arg_ty)) = (children.first(), explicit_arg_tys.first()) {
+            constrain_element_access(hir, ctx, container_ty, arg_id, arg_ty, result_ty, range);
+        }
+    }
+}
+
+/// Constrains the element access `container[arg]`, whose result is `result_ty`.
+///
+/// It is deferred as a tuple access so that Array/Tuple/Dict container types are handled
+/// correctly after unification.
+fn constrain_element_access(
+    hir: &Hir,
+    ctx: &mut InferenceContext,
+    container_ty: Type,
+    arg_id: SymbolId,
+    arg_ty: &Type,
+    result_ty: Type,
+    range: Option<mq_lang::Range>,
+) {
+    let literal_index = hir
+        .symbol(arg_id)
+        .filter(|arg_sym| matches!(arg_sym.kind, SymbolKind::Number))
+        .and_then(|arg_sym| arg_sym.value.as_ref()?.parse::<usize>().ok());
+
+    // Only constrain the index to Number when it is a literal numeric index (e.g. v[0]).
+    // For variable indices (e.g. a String key used for Dict access), skip this constraint:
+    // the correct element type is resolved via resolve_deferred_tuple_accesses once the
+    // container's type is known.
+    if literal_index.is_some() {
+        ctx.add_constraint(Constraint::Equal(
+            arg_ty.clone(),
+            Type::Number,
+            range,
+            ConstraintOrigin::General,
+        ));
+    }
+
+    ctx.add_deferred_tuple_access(infer::DeferredTupleAccess {
+        result_ty,
+        container_ty,
+        index: literal_index,
+        range,
+    });
+}
+
 /// Constrain a call: a user function (with arity, defaults and piped input), or a builtin.
 fn generate_call_constraints(
     hir: &Hir,
@@ -1977,6 +2088,16 @@ fn generate_call_constraints(
                 record_attr_call(hir, symbol_id, &children, &explicit_arg_tys, func_name == "get", ctx);
             }
 
+            // `self[i]` and `self[a:b]` index or slice the piped input.
+            if func_name == "self" && hir.bracket_group_count(symbol_id) > 0 {
+                let self_ty = ctx
+                    .get_piped_input(symbol_id)
+                    .cloned()
+                    .unwrap_or_else(|| Type::Var(ctx.fresh_var()));
+                constrain_bracket_access(hir, ctx, symbol_id, &children, &explicit_arg_tys, self_ty, range);
+                return;
+            }
+
             // Try user-defined function first (via HIR reference resolution)
             if let Some(def_id) = hir.resolve_reference_symbol(symbol_id) {
                 let def_symbol = hir.symbol(def_id);
@@ -2148,69 +2269,17 @@ fn generate_call_constraints(
                                 }
                             } else {
                                 // Non-function variable with bracket access.
-                                // Two args → range slice v[start:end] lowered as v(start, end).
-                                // One arg  → element access v[i] lowered as v(i).
-                                //
-                                // Chained accesses are lowered the same way: `v[0][:key]` is
-                                // `v(0, :key)`. A String/Symbol literal after the first argument
-                                // can never be a slice bound, so it marks such a chain.
-                                let is_chained_access = children.iter().skip(1).any(|&child_id| {
-                                    hir.symbol(child_id)
-                                        .is_some_and(|s| matches!(s.kind, SymbolKind::String | SymbolKind::Symbol))
-                                });
-                                if is_chained_access {
-                                    // The container's element type is not tracked through the
-                                    // chain, so leave the result unconstrained.
-                                    let ty_var = ctx.fresh_var();
-                                    ctx.set_symbol_type(symbol_id, Type::Var(ty_var));
-                                } else if children.len() == 2 {
-                                    // Range slice: delegate to the "slice" builtin.
-                                    // Use original_func_ty (not func_ty) so the container's
-                                    // type variable stays unified with its definition site.
-                                    let mut slice_arg_tys = vec![original_func_ty.clone()];
-                                    slice_arg_tys.extend(explicit_arg_tys.iter().cloned());
-                                    let defer = might_receive_piped_input(hir, symbol_id);
-                                    resolve_builtin_call(ctx, symbol_id, "slice", &slice_arg_tys, range, defer);
-                                } else {
-                                    // Element access: defer as a tuple access so that
-                                    // Array/Tuple/Dict container types are handled correctly
-                                    // after unification.
-                                    let literal_index = children.first().and_then(|&arg_id| {
-                                        let arg_sym = hir.symbol(arg_id)?;
-                                        if matches!(arg_sym.kind, SymbolKind::Number) {
-                                            arg_sym.value.as_ref()?.parse::<usize>().ok()
-                                        } else {
-                                            None
-                                        }
-                                    });
-
-                                    // Only constrain the index to Number when it is a literal
-                                    // numeric index (e.g. v[0]). For variable indices
-                                    // (e.g. a String key used for Dict access), skip this
-                                    // constraint: the correct element type is resolved via
-                                    // resolve_deferred_tuple_accesses once the container's
-                                    // type is known.
-                                    if literal_index.is_some()
-                                        && let Some(index_ty) = explicit_arg_tys.first()
-                                    {
-                                        ctx.add_constraint(Constraint::Equal(
-                                            index_ty.clone(),
-                                            Type::Number,
-                                            range,
-                                            ConstraintOrigin::General,
-                                        ));
-                                    }
-
-                                    let ty_var = ctx.fresh_var();
-                                    ctx.set_symbol_type(symbol_id, Type::Var(ty_var));
-
-                                    ctx.add_deferred_tuple_access(infer::DeferredTupleAccess {
-                                        call_symbol_id: symbol_id,
-                                        def_id,
-                                        index: literal_index,
-                                        range: get_symbol_range(hir, symbol_id),
-                                    });
-                                }
+                                // Use original_func_ty (not func_ty) so the container's type
+                                // variable stays unified with its definition site.
+                                constrain_bracket_access(
+                                    hir,
+                                    ctx,
+                                    symbol_id,
+                                    &children,
+                                    &explicit_arg_tys,
+                                    original_func_ty.clone(),
+                                    range,
+                                );
                             }
                         }
                     }

crates/mq-check/src/deferred.rs

@@ -323,10 +323,7 @@ pub(crate) fn resolve_deferred_tuple_accesses(ctx: &mut InferenceContext) -> boo
 
     let mut resolved_any = false;
     for access in &accesses {
-        let var_ty = match ctx.get_symbol_type(access.def_id).cloned() {
-            Some(ty) => ctx.resolve_type(&ty),
-            None => continue,
-        };
+        let var_ty = ctx.resolve_type(&access.container_ty);
 
         match &var_ty {
             types::Type::Tuple(elems) => {
@@ -341,13 +338,13 @@ pub(crate) fn resolve_deferred_tuple_accesses(ctx: &mut InferenceContext) -> boo
                     // Dynamic index, return Union of all element types
                     types::Type::union(elems.clone())
                 };
-                let call_ty = ctx.get_or_create_symbol_type(access.call_symbol_id);
+                let call_ty = access.result_ty.clone();
                 ctx.add_constraint(Constraint::Equal(call_ty, result_ty, None, ConstraintOrigin::General));
                 resolved_any = true;
             }
             types::Type::Array(elem) => {
                 // Normal array, bind to element type
-                let call_ty = ctx.get_or_create_symbol_type(access.call_symbol_id);
+                let call_ty = access.result_ty.clone();
                 ctx.add_constraint(Constraint::Equal(
                     call_ty,
                     *elem.clone(),
@@ -389,15 +386,15 @@ pub(crate) fn resolve_deferred_tuple_accesses(ctx: &mut InferenceContext) -> boo
                     } else {
                         types::Type::union(elem_types)
                     };
-                    let call_ty = ctx.get_or_create_symbol_type(access.call_symbol_id);
+                    let call_ty = access.result_ty.clone();
                     ctx.add_constraint(Constraint::Equal(call_ty, result_ty, None, ConstraintOrigin::General));
                     resolved_any = true;
                 }
             }
             types::Type::Dict(_, value) => {
                 // Dict bracket access: result type is the dict's value type.
                 // e.g. `d[key]` where d: Dict(String, V) → V
-                let call_ty = ctx.get_or_create_symbol_type(access.call_symbol_id);
+                let call_ty = access.result_ty.clone();
                 ctx.add_constraint(Constraint::Equal(
                     call_ty,
                     *value.clone(),
@@ -408,7 +405,7 @@ pub(crate) fn resolve_deferred_tuple_accesses(ctx: &mut InferenceContext) -> boo
             }
             types::Type::String => {
                 // String bracket access: `s[n]` or `s[n:m]` returns a String (character/slice).
-                let call_ty = ctx.get_or_create_symbol_type(access.call_symbol_id);
+                let call_ty = access.result_ty.clone();
                 ctx.add_constraint(Constraint::Equal(
                     call_ty,
                     types::Type::String,
@@ -439,7 +436,7 @@ pub(crate) fn resolve_deferred_tuple_accesses(ctx: &mut InferenceContext) -> boo
                     access.range,
                     ConstraintOrigin::General,
                 ));
-                let call_ty = ctx.get_or_create_symbol_type(access.call_symbol_id);
+                let call_ty = access.result_ty.clone();
                 ctx.add_constraint(Constraint::Equal(call_ty, elem_ty, None, ConstraintOrigin::General));
                 resolved_any = true;
             }
@@ -448,6 +445,16 @@ pub(crate) fn resolve_deferred_tuple_accesses(ctx: &mut InferenceContext) -> boo
     resolved_any
 }
 
+/// Resolves deferred tuple accesses until no more resolve, solving constraints between rounds.
+///
+/// Chained lookups like `v[0][0][0]` use the previous result as the next container, so each
+/// one resolves in a later round. Every productive round removes at least one queued access.
+pub(crate) fn drain_deferred_tuple_accesses(ctx: &mut InferenceContext) {
+    while resolve_deferred_tuple_accesses(ctx) {
+        unify::solve_constraints(ctx);
+    }
+}
+
 /// Result of checking union member overload compatibility.
 enum UnionCheckResult {
     /// All non-None members return a consistent type, with optional None propagation.

crates/mq-check/src/infer.rs

@@ -165,10 +165,10 @@ pub struct DeferredCallReturnAccess {
 /// the variable's Tuple type is fully resolved.
 #[derive(Debug, Clone)]
 pub struct DeferredTupleAccess {
-    /// The call symbol ID (the bracket access expression)
-    pub call_symbol_id: SymbolId,
-    /// The variable definition symbol ID
-    pub def_id: SymbolId,
+    /// The type of the accessed element
+    pub result_ty: Type,
+    /// The type of the indexed value
+    pub container_ty: Type,
     /// The literal index value (None if dynamic index)
     pub index: Option<usize>,
     /// Source range for error reporting

crates/mq-check/src/lib.rs

@@ -454,9 +454,7 @@ impl TypeChecker {
         laps.lap("narrowing");
 
         // Resolve deferred tuple index accesses now that variable types are known.
-        if deferred::resolve_deferred_tuple_accesses(&mut ctx) {
-            unify::solve_constraints(&mut ctx);
-        }
+        deferred::drain_deferred_tuple_accesses(&mut ctx);
 
         // Check the value assigned through `a[i][j] = v` against the element it replaces.
         if deferred::resolve_deferred_index_assigns(&mut ctx) {
@@ -500,18 +498,14 @@ impl TypeChecker {
         // `resolve_deferred_overloads` runs. This ensures that index accesses on variables
         // with union types containing unresolved vars (e.g., `Union(None, Var)`) are
         // retried with the now-concrete member types.
-        if deferred::resolve_deferred_tuple_accesses(&mut ctx) {
-            unify::solve_constraints(&mut ctx);
-        }
+        deferred::drain_deferred_tuple_accesses(&mut ctx);
 
         laps.lap("deferred_accesses");
 
         // Resolve deferred try/catch branch-type merges, then retry accesses/operators
         // that depend on a `let v = try: ... catch: ...` variable's now-final type.
         deferred::resolve_deferred_try_catches(&mut ctx);
-        if deferred::resolve_deferred_tuple_accesses(&mut ctx) {
-            unify::solve_constraints(&mut ctx);
-        }
+        deferred::drain_deferred_tuple_accesses(&mut ctx);
         laps.lap("deferred_try_catches");
         deferred::resolve_deferred_overloads(&mut ctx);
         laps.lap("resolve_deferred_overloads");