Fix #318

Author: CppCXY

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

@@ -0,0 +1,38 @@
+#[cfg(test)]
+mod test {
+    use smol_str::SmolStr;
+
+    use crate::{LuaType, LuaUnionType, VirtualWorkspace};
+
+    #[test]
+    fn test_issue_318() {
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+
+        ws.def(
+            r#"
+        local map = {
+            a = 'a',
+            b = 'b',
+            c = 'c',
+        }
+        local key      --- @type string
+        c = map[key]   -- type should be ('a'|'b'|'c'|nil)
+
+        "#,
+        );
+
+        let c_ty = ws.expr_ty("c");
+
+        let union_type = LuaType::Union(
+            LuaUnionType::new(vec![
+                LuaType::StringConst(SmolStr::new("a").into()),
+                LuaType::StringConst(SmolStr::new("b").into()),
+                LuaType::StringConst(SmolStr::new("c").into()),
+                LuaType::Nil,
+            ])
+            .into(),
+        );
+
+        assert_eq!(c_ty, union_type);
+    }
+}

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

@@ -9,6 +9,7 @@ mod for_range_var_infer_test;
 mod infer_str_tpl_test;
 mod inherit_type;
 mod mathlib_test;
+mod memebr_infer_test;
 mod metatable_test;
 mod module_annotation;
 mod multi_return;

crates/emmylua_code_analysis/src/db_index/type/types.rs

@@ -1,4 +1,8 @@
-use std::{collections::HashMap, hash::Hash, sync::Arc};
+use std::{
+    collections::HashMap,
+    hash::{Hash, Hasher},
+    sync::Arc,
+};
 
 use internment::ArcIntern;
 use rowan::TextRange;
@@ -631,8 +635,7 @@ impl From<LuaObjectType> for LuaType {
         LuaType::Object(t.into())
     }
 }
-
-#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+#[derive(Debug, Clone)]
 pub struct LuaUnionType {
     types: Vec<LuaType>,
 }
@@ -655,6 +658,50 @@ impl LuaUnionType {
     }
 }
 
+impl PartialEq for LuaUnionType {
+    fn eq(&self, other: &Self) -> bool {
+        if self.types.len() != other.types.len() {
+            return false;
+        }
+        let mut counts = HashMap::new();
+        // Count occurrences in self.types
+        for t in &self.types {
+            *counts.entry(t).or_insert(0) += 1;
+        }
+        // Decrease counts for other.types
+        for t in &other.types {
+            match counts.get_mut(t) {
+                Some(count) if *count > 0 => *count -= 1,
+                _ => return false,
+            }
+        }
+        true
+    }
+}
+
+impl Eq for LuaUnionType {}
+
+impl std::hash::Hash for LuaUnionType {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        // To get an order-insensitive hash, combine:
+        // - the number of elements
+        // - the sum and product of the hashes of individual elements.
+        // This is a simple and fast commutative hash.
+        let mut sum: u64 = 0;
+        let mut prod: u64 = 1;
+        for t in &self.types {
+            let mut hasher = std::collections::hash_map::DefaultHasher::new();
+            t.hash(&mut hasher);
+            let h = hasher.finish();
+            sum = sum.wrapping_add(h);
+            prod = prod.wrapping_mul(h.wrapping_add(1));
+        }
+        self.types.len().hash(state);
+        sum.hash(state);
+        prod.hash(state);
+    }
+}
+
 impl From<LuaUnionType> for LuaType {
     fn from(t: LuaUnionType) -> Self {
         LuaType::Union(t.into())

crates/emmylua_code_analysis/src/semantic/infer/infer_index.rs

@@ -453,28 +453,67 @@ fn infer_member_by_index_table(
     table_range: &InFiled<TextRange>,
     index_expr: LuaIndexMemberExpr,
 ) -> InferResult {
-    let metatable = db
-        .get_metatable_index()
-        .get(table_range)
-        .ok_or(InferFailReason::FieldDotFound)?;
-
-    let meta_owner = LuaOperatorOwner::Table(metatable.clone());
-    let operator_ids = db
-        .get_operator_index()
-        .get_operators(&meta_owner, LuaOperatorMetaMethod::Index)
-        .ok_or(InferFailReason::FieldDotFound)?;
+    let metatable = db.get_metatable_index().get(table_range);
+    match metatable {
+        Some(metatable) => {
+            let meta_owner = LuaOperatorOwner::Table(metatable.clone());
+            let operator_ids = db
+                .get_operator_index()
+                .get_operators(&meta_owner, LuaOperatorMetaMethod::Index)
+                .ok_or(InferFailReason::FieldDotFound)?;
+
+            let index_key = index_expr.get_index_key().ok_or(InferFailReason::None)?;
+
+            for operator_id in operator_ids {
+                let operator = db
+                    .get_operator_index()
+                    .get_operator(operator_id)
+                    .ok_or(InferFailReason::None)?;
+                let operand = operator.get_operand(db);
+                let return_type = operator.get_result(db)?;
+                let typ = infer_index_metamethod(db, cache, &index_key, &operand, &return_type)?;
+                return Ok(typ);
+            }
+        }
+        None => {
+            let index_key = index_expr.get_index_key().ok_or(InferFailReason::None)?;
+            if let LuaIndexKey::Expr(expr) = index_key {
+                let key_type = infer_expr(db, cache, expr.clone())?;
+                let members = db
+                    .get_member_index()
+                    .get_members(&LuaMemberOwner::Element(table_range.clone()));
+                if let Some(members) = members {
+                    let mut result_type = LuaType::Unknown;
+                    for member in members {
+                        let member_key_type = match member.get_key() {
+                            LuaMemberKey::Name(s) => LuaType::StringConst(s.clone().into()),
+                            LuaMemberKey::Integer(i) => LuaType::IntegerConst(*i),
+                            _ => continue,
+                        };
+                        if check_type_compact(db, &key_type, &member_key_type).is_ok() {
+                            let member_type = db
+                                .get_type_index()
+                                .get_type_cache(&member.get_id().into())
+                                .map(|it| it.as_type())
+                                .unwrap_or(&LuaType::Unknown);
+
+                            result_type = TypeOps::Union.apply(&result_type, member_type);
+                        }
+                    }
 
-    let index_key = index_expr.get_index_key().ok_or(InferFailReason::None)?;
+                    if !result_type.is_unknown() {
+                        if matches!(
+                            key_type,
+                            LuaType::String | LuaType::Number | LuaType::Integer
+                        ) {
+                            result_type = TypeOps::Union.apply(&result_type, &LuaType::Nil);
+                        }
 
-    for operator_id in operator_ids {
-        let operator = db
-            .get_operator_index()
-            .get_operator(operator_id)
-            .ok_or(InferFailReason::None)?;
-        let operand = operator.get_operand(db);
-        let return_type = operator.get_result(db)?;
-        let typ = infer_index_metamethod(db, cache, &index_key, &operand, &return_type)?;
-        return Ok(typ);
+                        return Ok(result_type);
+                    }
+                }
+            }
+        }
     }
 
     Err(InferFailReason::FieldDotFound)