Merge pull request #1182 from lewis6991/codex/issue-1180-self-assignment-flow

fix(flow): preserve self-assignment antecedent types

Author: CppCXY

crates/emmylua_code_analysis/src/compilation/test/flow.rs

@@ -3599,6 +3599,48 @@ _2 = a[1]
         assert_eq!(ws.humanize_type(after_assign), "Counter");
     }
 
+    #[test]
+    fn test_issue_1180_self_assignment_replays_antecedent_type() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def(
+            r#"
+        local index = 1
+        index = index - 1
+        after_assign = index
+        "#,
+        );
+
+        assert_eq!(ws.expr_ty("after_assign"), ws.ty("integer"));
+    }
+
+    #[test]
+    fn test_issue_1180_while_loop_preserves_integer_self_assignments() {
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+        let file_id = ws.def(
+            r#"
+        function test()
+            local arr = {1, 2, 3, 4, "hello", 6, false, 7}
+            local index = 1
+
+            while index <= #arr do
+                local v = arr[index]
+
+                if type(v) == "boolean" then
+                    table.remove(arr, index)
+                    index = index - 1
+                    print(index)
+                end
+
+                index = index + 1
+            end
+        end
+        "#,
+        );
+        let index_type = last_name_expr_type(&ws, file_id, "index");
+
+        assert_eq!(ws.humanize_type(index_type), "integer");
+    }
+
     #[test]
     #[timeout(5000)]
     fn test_issue_1114_repeated_self_dependent_assignments_build_semantic_model() {

crates/emmylua_code_analysis/src/semantic/infer/narrow/get_type_at_flow.rs

@@ -1,7 +1,7 @@
 use emmylua_parser::{
-    BinaryOperator, LuaAssignStat, LuaAstNode, LuaCallExprStat, LuaChunk, LuaDocOpType, LuaExpr,
-    LuaForStat, LuaIndexKey, LuaIndexMemberExpr, LuaSyntaxId, LuaTableExpr, LuaUnaryExpr,
-    LuaVarExpr, UnaryOperator,
+    LuaAssignStat, LuaAstNode, LuaCallExprStat, LuaChunk, LuaDocOpType, LuaExpr, LuaForStat,
+    LuaIndexKey, LuaIndexMemberExpr, LuaSyntaxId, LuaTableExpr, LuaUnaryExpr, LuaVarExpr,
+    UnaryOperator,
 };
 use hashbrown::HashSet;
 use std::{rc::Rc, sync::Arc};
@@ -100,6 +100,14 @@ enum Continuation {
         expr_type: LuaType,
         reuse_antecedent_narrowing: bool,
     },
+    // Resume RHS replay once the assigned ref's pre-assignment type is known.
+    // This resolves self-reads before earlier RHS dependencies can repeatedly
+    // walk the same assignment chain.
+    AssignmentSelfAntecedent {
+        walk: QueryWalk,
+        replay: FlowExprReplay,
+        replay_query: FlowReplayQuery,
+    },
     // Resume an assignment result after proving the branch can reach the
     // assignment. An impossible branch must not contribute this result to a merge.
     AssignmentReachability {
@@ -218,9 +226,39 @@ impl FlowReplayQuery {
         Ok(())
     }
 
-    fn resolve_dependencies(&mut self, var_ref_id: &VarRefId, typ: LuaType) {
+    fn has_unskippable_dependency(&self, var_ref_id: &VarRefId, flow_id: FlowId) -> bool {
+        // A resolved index can skip its prefix and key dependencies. Only hoist
+        // a self-read when every earlier success path still reaches it.
+        let mut furthest_success_jump = self.next_dependency_idx;
+        for (idx, query) in self
+            .dependency_queries
+            .iter()
+            .enumerate()
+            .skip(self.next_dependency_idx)
+        {
+            if idx >= furthest_success_jump
+                && query.resolved_type.is_none()
+                && query.var_ref_id == *var_ref_id
+                && query.flow_id == flow_id
+            {
+                return true;
+            }
+            if let Some(next_idx) = query.next_dependency_idx_on_success {
+                furthest_success_jump = furthest_success_jump.max(next_idx);
+            }
+        }
+
+        false
+    }
+
+    fn resolve_dependencies_at_flow(
+        &mut self,
+        var_ref_id: &VarRefId,
+        flow_id: FlowId,
+        typ: LuaType,
+    ) {
         for query in &mut self.dependency_queries {
-            if query.var_ref_id == *var_ref_id {
+            if query.var_ref_id == *var_ref_id && query.flow_id == flow_id {
                 query.resolved_type = Some(typ.clone());
             }
         }
@@ -496,6 +534,16 @@ impl<'a> FlowTypeEngine<'a> {
                         reuse_antecedent_narrowing,
                         query_result,
                     ),
+                    Some(Continuation::AssignmentSelfAntecedent {
+                        walk,
+                        replay,
+                        replay_query,
+                    }) => self.resume_assignment_self_antecedent(
+                        walk,
+                        replay,
+                        replay_query,
+                        query_result,
+                    ),
                     Some(Continuation::AssignmentReachability { walk, result_type }) => {
                         match query_result {
                             Ok(FlowQueryResult::Unreachable) => {
@@ -692,6 +740,30 @@ impl<'a> FlowTypeEngine<'a> {
         Ok(self.finish_walk(walk, result_type))
     }
 
+    fn resume_assignment_self_antecedent(
+        &mut self,
+        walk: QueryWalk,
+        replay: FlowExprReplay,
+        mut replay_query: FlowReplayQuery,
+        antecedent_result: FlowResult,
+    ) -> Result<SchedulerStep, InferFailReason> {
+        let antecedent_type = match antecedent_result {
+            Ok(FlowQueryResult::Type(antecedent_type)) => antecedent_type,
+            Ok(FlowQueryResult::Unreachable) => {
+                return Ok(self.finish_unreachable_branch(walk));
+            }
+            Err(err) => return self.fail_query(&walk.query, err),
+        };
+
+        let antecedent_flow_id = replay_query.flow_id;
+        replay_query.resolve_dependencies_at_flow(
+            &walk.query.var_ref_id,
+            antecedent_flow_id,
+            antecedent_type,
+        );
+        self.start_expr_replay(walk, replay, replay_query)
+    }
+
     fn start_expr_replay(
         &mut self,
         walk: QueryWalk,
@@ -1106,29 +1178,31 @@ impl<'a> FlowTypeEngine<'a> {
             let expr_idx = i.min(last_expr_idx);
             let result_slot = i.saturating_sub(last_expr_idx);
             let expr = assignment_info.exprs[expr_idx].clone();
-            let mut replay_query = FlowReplayQuery::new(
+            let replay_query = FlowReplayQuery::new(
                 self.db,
                 Some(self.tree),
                 self.cache,
                 antecedent_flow_id,
-                expr.clone(),
+                expr,
                 true,
             );
-            // A plain self-dependent RHS would replay this assignment while
-            // trying to type itself. Treat that self read as unknown; `and`/`or`
-            // assignments still need the antecedent value for their semantics.
-            if explicit_var_type.is_none() && !contains_short_circuit_binary_expr(&expr) {
-                replay_query.resolve_dependencies(&var_ref_id, LuaType::Unknown);
+            let replay = FlowExprReplay::Assignment {
+                antecedent_flow_id,
+                explicit_var_type,
+                result_slot,
+            };
+            if replay_query.has_unskippable_dependency(&var_ref_id, antecedent_flow_id) {
+                let query = FlowQuery::new(self.cache, &var_ref_id, antecedent_flow_id);
+                return Ok(SchedulerStep::StartQuery {
+                    query,
+                    continuation: Some(Continuation::AssignmentSelfAntecedent {
+                        walk,
+                        replay,
+                        replay_query,
+                    }),
+                });
             }
-            return self.start_expr_replay(
-                walk,
-                FlowExprReplay::Assignment {
-                    antecedent_flow_id,
-                    explicit_var_type,
-                    result_slot,
-                },
-                replay_query,
-            );
+            return self.start_expr_replay(walk, replay, replay_query);
         }
 
         self.finish_assignment_expr_type(walk, antecedent_flow_id, explicit_var_type, LuaType::Nil)
@@ -1971,17 +2045,6 @@ fn integer_enum_assignment_declared_type(
         .then_some(declared_type)
 }
 
-fn contains_short_circuit_binary_expr(expr: &LuaExpr) -> bool {
-    expr.descendants::<LuaExpr>().any(|expr| {
-        let LuaExpr::BinaryExpr(binary_expr) = expr else {
-            return false;
-        };
-        binary_expr.get_op_token().is_some_and(|token| {
-            matches!(token.get_op(), BinaryOperator::OpAnd | BinaryOperator::OpOr)
-        })
-    })
-}
-
 fn is_partial_assignment_expr_compatible(
     db: &DbIndex,
     source_type: &LuaType,