Auto merge of #18139 - ShoyuVanilla:issue-18109, r=Veykril

fix: Extend `type_variable_table` when modifying index is larger than the table size

Fixes #18109

Whenever we create an inference variable in r-a, we extend `type_variable_table` to matching size here;

https://github.com/rust-lang/rust-analyzer/blob/f4aca78c92e03354327c8f6c7fefaef9f45ab166/crates/hir-ty/src/infer/unify.rs#L378-L381

But sometimes, an inference variable is [created from chalk](https://github.com/rust-lang/chalk/blob/ab710e0c9b455403b138ef72a2fb90967a58eff3/chalk-solve/src/infer/unify.rs#L743) and passed to r-a as a type of an expression or a pattern.
If r-a set diverging flag to this before the table is extended to a sufficient size, it panics here;

https://github.com/rust-lang/rust-analyzer/blob/f4aca78c92e03354327c8f6c7fefaef9f45ab166/crates/hir-ty/src/infer/unify.rs#L275-L277

I think that extending table when setting diverging flag is reasonable becase we are already doing such extending to a size that covers the inference vars created from chalk and this change only covers the order-dependent random cases that this might fail

2 files changed, 81 insertions, 4 deletions

diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/infer/unify.rs b/src/tools/rust-analyzer/crates/hir-ty/src/infer/unify.rs
index d38380e4a3a..7300453ff00 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/infer/unify.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/infer/unify.rs
@@ -266,14 +266,16 @@ impl<'a> InferenceTable<'a> {
             }
             let v = InferenceVar::from(i as u32);
             let root = self.var_unification_table.inference_var_root(v);
-            if let Some(data) = self.type_variable_table.get_mut(root.index() as usize) {
-                *data |= TypeVariableFlags::DIVERGING;
-            }
+            self.modify_type_variable_flag(root, |f| {
+                *f |= TypeVariableFlags::DIVERGING;
+            });
         }
     }
 
     pub(super) fn set_diverging(&mut self, iv: InferenceVar, diverging: bool) {
-        self.type_variable_table[iv.index() as usize].set(TypeVariableFlags::DIVERGING, diverging);
+        self.modify_type_variable_flag(iv, |f| {
+            f.set(TypeVariableFlags::DIVERGING, diverging);
+        });
     }
 
     fn fallback_value(&self, iv: InferenceVar, kind: TyVariableKind) -> Ty {
@@ -370,6 +372,18 @@ impl<'a> InferenceTable<'a> {
         var
     }
 
+    fn modify_type_variable_flag<F>(&mut self, var: InferenceVar, cb: F)
+    where
+        F: FnOnce(&mut TypeVariableFlags),
+    {
+        let idx = var.index() as usize;
+        if self.type_variable_table.len() <= idx {
+            self.extend_type_variable_table(idx);
+        }
+        if let Some(f) = self.type_variable_table.get_mut(idx) {
+            cb(f);
+        }
+    }
     fn extend_type_variable_table(&mut self, to_index: usize) {
         let count = to_index - self.type_variable_table.len() + 1;
         self.type_variable_table.extend(iter::repeat(TypeVariableFlags::default()).take(count));
diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/tests/regression.rs b/src/tools/rust-analyzer/crates/hir-ty/src/tests/regression.rs
index db289f430ea..a3cf12d8a16 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/tests/regression.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/tests/regression.rs
@@ -2232,3 +2232,66 @@ async fn f<A, B, C>() -> Bar {}
         "#]],
     );
 }
+
+#[test]
+fn issue_18109() {
+    check_infer(
+        r#"
+//- minicore: option
+struct Map<T, U>(T, U);
+
+impl<T, U> Map<T, U> {
+    fn new() -> Self { loop {} }
+    fn get(&self, _: &T) -> Option<&U> { loop {} }
+}
+
+fn test(x: bool) {
+    let map = Map::new();
+    let _ = match x {
+        true => {
+            let Some(val) = map.get(&8) else { return };
+            *val
+        }
+        false => return,
+        _ => 42,
+    };
+}
+"#,
+        expect![[r#"
+            69..80 '{ loop {} }': Map<T, U>
+            71..78 'loop {}': !
+            76..78 '{}': ()
+            93..97 'self': &'? Map<T, U>
+            99..100 '_': &'? T
+            120..131 '{ loop {} }': Option<&'? U>
+            122..129 'loop {}': !
+            127..129 '{}': ()
+            143..144 'x': bool
+            152..354 '{     ...  }; }': ()
+            162..165 'map': Map<i32, i32>
+            168..176 'Map::new': fn new<i32, i32>() -> Map<i32, i32>
+            168..178 'Map::new()': Map<i32, i32>
+            188..189 '_': i32
+            192..351 'match ...     }': i32
+            198..199 'x': bool
+            210..214 'true': bool
+            210..214 'true': bool
+            218..303 '{     ...     }': i32
+            236..245 'Some(val)': Option<&'? i32>
+            241..244 'val': &'? i32
+            248..251 'map': Map<i32, i32>
+            248..259 'map.get(&8)': Option<&'? i32>
+            256..258 '&8': &'? i32
+            257..258 '8': i32
+            265..275 '{ return }': !
+            267..273 'return': !
+            289..293 '*val': i32
+            290..293 'val': &'? i32
+            312..317 'false': bool
+            312..317 'false': bool
+            321..327 'return': !
+            337..338 '_': bool
+            342..344 '42': i32
+        "#]],
+    );
+}