fix: stop repeated assignments from hanging

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

@@ -8,6 +8,7 @@ mod test {
     const LARGE_LINEAR_ASSIGNMENT_STEPS: usize = 2048;
     const MAXWELLHOME_ARRAY_VALUES: usize = 2048;
     const ISSUE_1100_HIGHLIGHT_GROUPS: usize = 2048;
+    const ISSUE_1114_REPEATED_BINARY_STEPS: usize = 512;
 
     #[test]
     fn test_closure_return() {
@@ -2942,6 +2943,228 @@ _2 = a[1]
         assert_eq!(ws.humanize_type(after_assign), "Pattern");
     }
 
+    #[test]
+    fn test_assignment_binary_rhs_replays_non_self_dependency() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def(
+            r#"
+        ---@class FooValue
+        ---@field kind "foo"
+        ---@field value integer
+
+        ---@class BarValue
+        ---@field kind "bar"
+        ---@field value string
+
+        local right ---@type FooValue|BarValue
+
+        if right.kind == "foo" then
+            local value
+            value = right.value + 1
+            after_assign = value
+        end
+        "#,
+        );
+
+        let after_assign = ws.expr_ty("after_assign");
+        assert_eq!(ws.humanize_type(after_assign), "integer");
+    }
+
+    #[test]
+    fn test_assignment_rhs_keeps_flow_dependent_concat_operator() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def(
+            r#"
+        ---@class Rope
+        ---@operator concat(Rope): Rope
+
+        local left ---@type Rope?
+        local right ---@type Rope?
+
+        if not left then return end
+        if not right then return end
+        left = left .. right
+        after_assign = left
+        "#,
+        );
+
+        let after_assign = ws.expr_ty("after_assign");
+        assert_eq!(ws.humanize_type(after_assign), "Rope");
+    }
+
+    #[test]
+    fn test_assignment_rhs_keeps_flow_dependent_add_operator() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def(
+            r#"
+        ---@class Counter
+        ---@operator add(Counter): Counter
+
+        local left ---@type Counter?
+        local right ---@type Counter?
+
+        if not left then return end
+        if not right then return end
+        left = left + right
+        after_assign = left
+        "#,
+        );
+
+        let after_assign = ws.expr_ty("after_assign");
+        assert_eq!(ws.humanize_type(after_assign), "Counter");
+    }
+
+    #[test]
+    #[timeout(5000)]
+    fn test_issue_1114_repeated_self_concat_builds_semantic_model() {
+        let mut ws = VirtualWorkspace::new();
+        let repeated_assignments =
+            "wnd = wnd .. config.pic[idx][index]\n".repeat(ISSUE_1114_REPEATED_BINARY_STEPS);
+        let block = format!(
+            r#"
+        function f(idx, index)
+            local wnd = ""
+            local child = ""
+            {repeated_assignments}
+            return wnd:format(child:sub(1, -2))
+        end
+        "#
+        );
+
+        let file_id = ws.def(&block);
+
+        assert!(
+            ws.analysis
+                .compilation
+                .get_semantic_model(file_id)
+                .is_some(),
+            "expected semantic model for repeated self-concat repro"
+        );
+    }
+
+    #[test]
+    #[timeout(5000)]
+    fn test_issue_1114_repeated_self_add_builds_semantic_model() {
+        let mut ws = VirtualWorkspace::new();
+        let repeated_assignments =
+            "total = total + config.pic[idx][index]\n".repeat(ISSUE_1114_REPEATED_BINARY_STEPS);
+        let block = format!(
+            r#"
+        function f(idx, index)
+            local total = 0
+            {repeated_assignments}
+            return total
+        end
+        "#
+        );
+
+        let file_id = ws.def(&block);
+
+        assert!(
+            ws.analysis
+                .compilation
+                .get_semantic_model(file_id)
+                .is_some(),
+            "expected semantic model for repeated self-add repro"
+        );
+    }
+
+    #[test]
+    #[timeout(5000)]
+    fn test_issue_1114_repeated_parenthesized_self_concat_builds_semantic_model() {
+        let mut ws = VirtualWorkspace::new();
+        let repeated_assignments =
+            "wnd = (wnd .. config.pic[idx][index])\n".repeat(ISSUE_1114_REPEATED_BINARY_STEPS);
+        let block = format!(
+            r#"
+        function f(idx, index)
+            local wnd = ""
+            {repeated_assignments}
+            return wnd
+        end
+        "#
+        );
+
+        let file_id = ws.def(&block);
+
+        assert!(
+            ws.analysis
+                .compilation
+                .get_semantic_model(file_id)
+                .is_some(),
+            "expected semantic model for repeated parenthesized self-concat repro"
+        );
+    }
+
+    #[test]
+    #[timeout(5000)]
+    fn test_issue_1114_repeated_unary_self_assignment_builds_semantic_model() {
+        let mut ws = VirtualWorkspace::new();
+        let repeated_assignments =
+            "total = -(total + config.pic[idx][index])\n".repeat(ISSUE_1114_REPEATED_BINARY_STEPS);
+        let block = format!(
+            r#"
+        function f(idx, index)
+            local total = 0
+            {repeated_assignments}
+            return total
+        end
+        "#
+        );
+
+        let file_id = ws.def(&block);
+
+        assert!(
+            ws.analysis
+                .compilation
+                .get_semantic_model(file_id)
+                .is_some(),
+            "expected semantic model for repeated unary self-assignment repro"
+        );
+    }
+
+    #[test]
+    #[timeout(5000)]
+    fn test_issue_1114_repeated_comparison_self_assignment_builds_semantic_model() {
+        let mut ws = VirtualWorkspace::new();
+        let repeated_assignments = "enabled = enabled == config.pic[idx][index]\n"
+            .repeat(ISSUE_1114_REPEATED_BINARY_STEPS);
+        let block = format!(
+            r#"
+        function f(idx, index)
+            local enabled = true
+            {repeated_assignments}
+            return enabled
+        end
+        "#
+        );
+
+        let file_id = ws.def(&block);
+
+        assert!(
+            ws.analysis
+                .compilation
+                .get_semantic_model(file_id)
+                .is_some(),
+            "expected semantic model for repeated comparison self-assignment repro"
+        );
+    }
+
+    #[test]
+    fn test_binary_assignment_infer_error_keeps_previous_type() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def(
+            r#"
+        local value = "prior"
+        value = config.pic + 1
+        after_assign = value
+        "#,
+        );
+
+        let after_assign = ws.expr_ty("after_assign");
+        assert_eq!(ws.humanize_type(after_assign), "string");
+    }
+
     #[test]
     fn test_eq_uses_branch_narrowed_rhs_ref_type() {
         let mut ws = VirtualWorkspace::new();

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

@@ -1,13 +1,13 @@
 use emmylua_parser::{
-    LuaAssignStat, LuaAstNode, LuaChunk, LuaDocOpType, LuaExpr, LuaIndexKey, LuaIndexMemberExpr,
-    LuaSyntaxId, LuaTableExpr, LuaVarExpr,
+    BinaryOperator, LuaAssignStat, LuaAstNode, LuaChunk, LuaDocOpType, LuaExpr, LuaIndexKey,
+    LuaIndexMemberExpr, LuaSyntaxId, LuaTableExpr, LuaVarExpr, UnaryOperator,
 };
 use hashbrown::HashSet;
 use std::{rc::Rc, sync::Arc};
 
 use crate::{
     CacheEntry, DbIndex, FlowId, FlowNode, FlowNodeKind, FlowTree, InferFailReason, LuaDeclId,
-    LuaInferCache, LuaMemberId, LuaSignatureId, LuaType, TypeOps, check_type_compact,
+    LuaInferCache, LuaMemberId, LuaSignatureId, LuaType, LuaTypeNode, TypeOps, check_type_compact,
     semantic::{
         cache::{FlowAssignmentInfo, FlowMode, FlowVarCache},
         infer::{
@@ -1049,9 +1049,21 @@ impl<'a> FlowTypeEngine<'a> {
                 Some(self.tree),
                 self.cache,
                 antecedent_flow_id,
-                expr,
+                expr.clone(),
                 true,
             );
+            // PERF: Replaying `x = <operator using x>` chains can walk every
+            // prior assignment; built-in operators only need their result shape.
+            if explicit_var_type.is_none()
+                && replay_query
+                    .dependency_queries
+                    .iter()
+                    .any(|query| query.var_ref_id == var_ref_id)
+                && let Some(expr_type) =
+                    self_dependent_assignment_operator_type(self.db, self.cache, &expr, result_slot)
+            {
+                return Ok(self.finish_walk(walk, expr_type));
+            }
             return self.start_expr_replay(
                 walk,
                 FlowExprReplay::Assignment {
@@ -1647,6 +1659,89 @@ fn preserves_assignment_expr_type(typ: &LuaType) -> bool {
     matches!(typ, LuaType::TableConst(_) | LuaType::Object(_)) || is_exact_assignment_expr_type(typ)
 }
 
+fn self_dependent_assignment_operator_type(
+    db: &DbIndex,
+    cache: &mut LuaInferCache,
+    expr: &LuaExpr,
+    result_slot: usize,
+) -> Option<LuaType> {
+    let fallback_type = match expr {
+        LuaExpr::ParenExpr(paren_expr) => {
+            return self_dependent_assignment_operator_type(
+                db,
+                cache,
+                &paren_expr.get_expr()?,
+                result_slot,
+            );
+        }
+        LuaExpr::BinaryExpr(binary_expr) => {
+            if binary_expr.get_exprs().is_some_and(|(left, right)| {
+                [left, right]
+                    .into_iter()
+                    .any(|operand| expr_infers_custom_type(db, cache, operand))
+            }) {
+                return None;
+            }
+
+            match binary_expr.get_op_token()?.get_op() {
+                BinaryOperator::OpAdd
+                | BinaryOperator::OpSub
+                | BinaryOperator::OpMul
+                | BinaryOperator::OpDiv
+                | BinaryOperator::OpMod
+                | BinaryOperator::OpPow => LuaType::Number,
+                BinaryOperator::OpIDiv
+                | BinaryOperator::OpBAnd
+                | BinaryOperator::OpBOr
+                | BinaryOperator::OpBXor
+                | BinaryOperator::OpShl
+                | BinaryOperator::OpShr => LuaType::Integer,
+                BinaryOperator::OpConcat => LuaType::String,
+                BinaryOperator::OpLt
+                | BinaryOperator::OpLe
+                | BinaryOperator::OpGt
+                | BinaryOperator::OpGe
+                | BinaryOperator::OpEq
+                | BinaryOperator::OpNe => LuaType::Boolean,
+                _ => return None,
+            }
+        }
+        LuaExpr::UnaryExpr(unary_expr) => {
+            if unary_expr
+                .get_expr()
+                .is_some_and(|operand| expr_infers_custom_type(db, cache, operand))
+            {
+                return None;
+            }
+
+            match unary_expr.get_op_token()?.get_op() {
+                UnaryOperator::OpNot => LuaType::Boolean,
+                UnaryOperator::OpLen | UnaryOperator::OpBNot => LuaType::Integer,
+                UnaryOperator::OpUnm => LuaType::Number,
+                UnaryOperator::OpNop => return None,
+            }
+        }
+        _ => return None,
+    };
+
+    let expr_type = match try_infer_expr_no_flow(db, cache, expr.clone()) {
+        Ok(Some(expr_type)) => expr_type
+            .get_result_slot_type(result_slot)
+            .unwrap_or(LuaType::Nil),
+        Ok(None) | Err(_) if result_slot == 0 => fallback_type,
+        Ok(None) | Err(_) => return None,
+    };
+
+    (expr_type.is_boolean() || expr_type.is_number() || expr_type.is_string()).then_some(expr_type)
+}
+
+fn expr_infers_custom_type(db: &DbIndex, cache: &mut LuaInferCache, expr: LuaExpr) -> bool {
+    try_infer_expr_no_flow(db, cache, expr)
+        .ok()
+        .flatten()
+        .is_some_and(|typ| typ.any_type(LuaType::is_custom_type))
+}
+
 fn is_partial_assignment_expr_compatible(
     db: &DbIndex,
     source_type: &LuaType,