From 59e30080a0baccb10e5ece0e5392cb94289ddc4b Mon Sep 17 00:00:00 2001
From: Shoyu Vanilla <modulo641@gmail.com>
Date: Wed, 18 Sep 2024 23:31:42 +0900
Subject: [PATCH] fix: Extend `type_variable_table` when modifying index is
 larger than table size

---
 .../crates/hir-ty/src/infer/unify.rs          | 22 +++++--
 .../crates/hir-ty/src/tests/regression.rs     | 63 +++++++++++++++++++
 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 @@ pub(super) fn propagate_diverging_flag(&mut self) {
             }
             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 @@ pub(crate) fn normalize_projection_ty(&mut self, proj_ty: ProjectionTy) -> Ty {
         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
+        "#]],
+    );
+}