fix for for range var

CppCXY · CppCXY · commit 0c22c209e6ef · 2025-04-04T20:04:58.000+08:00
close #321 close #261
diff --git a/crates/emmylua_code_analysis/resources/std/global.lua b/crates/emmylua_code_analysis/resources/std/global.lua
@@ -113,7 +113,7 @@ function getmetatable(object) end
 --- the first absent index.
 ---@generic V
 ---@param t V[] | table<any, V> | {[any]: V}
----@return fun(tbl: any):int, std.NotNull<V>
+---@return fun(tbl: any):int, V
 function ipairs(t) end
 
 ---@alias std.loadmode
@@ -231,7 +231,7 @@ function next(table, index) end
 --- traversal.
 ---@generic K, V
 ---@param t table<K, V> | V[] | {[K]: V}
----@return fun(tbl: any):K, std.NotNull<V>
+---@return fun(tbl: any):K, V
 function pairs(t) end
 ---
 --- Calls function `f` with the given arguments in *protected mode*. This
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/decl/stats.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/decl/stats.rs
@@ -4,8 +4,9 @@ use emmylua_parser::{
 };
 
 use crate::{
+    compilation::analyzer::bind_type::bind_type,
     db_index::{LocalAttribute, LuaDecl, LuaMember, LuaMemberKey},
-    LuaDeclExtra, LuaMemberFeature, LuaMemberId, LuaSemanticDeclId, LuaSignatureId,
+    LuaDeclExtra, LuaMemberFeature, LuaMemberId, LuaSemanticDeclId, LuaSignatureId, LuaType,
 };
 
 use super::{members::find_index_owner, DeclAnalyzer};
@@ -176,8 +177,13 @@ pub fn analyze_for_stat(analyzer: &mut DeclAnalyzer, stat: LuaForStat) -> Option
         },
         None,
     );
-
+    let decl_id = decl.get_id();
     analyzer.add_decl(decl);
+    bind_type(
+        analyzer.db,
+        decl_id.into(),
+        crate::LuaTypeCache::DocType(LuaType::Integer),
+    );
 
     Some(())
 }
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/lua/for_range_stat.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/lua/for_range_stat.rs
@@ -0,0 +1,123 @@
+use std::sync::Arc;
+
+use emmylua_parser::{LuaAstToken, LuaForRangeStat};
+
+use crate::{
+    compilation::analyzer::unresolve::UnResolveIterVar, DbIndex, InferFailReason, LuaDeclId,
+    LuaFunctionType, LuaOperatorMetaMethod, LuaType, LuaTypeCache, TypeOps,
+};
+
+use super::LuaAnalyzer;
+
+pub fn analyze_for_range_stat(
+    analyzer: &mut LuaAnalyzer,
+    for_range_stat: LuaForRangeStat,
+) -> Option<()> {
+    let var_name_list = for_range_stat.get_var_name_list();
+    let first_iter_expr = for_range_stat.get_expr_list().next()?;
+    let first_iter_type = analyzer.infer_expr(&first_iter_expr);
+
+    match first_iter_type {
+        Ok(first_iter_type) => {
+            let iter_doc_func = infer_for_range_iter_expr_func(analyzer.db, first_iter_type);
+
+            if let Some(doc_func) = iter_doc_func {
+                let multi_return = doc_func.get_multi_return();
+                let mut idx = 0;
+                for var_name in var_name_list {
+                    let position = var_name.get_position();
+                    let decl_id = LuaDeclId::new(analyzer.file_id, position);
+                    let ret_type = multi_return
+                        .get_type(idx)
+                        .cloned()
+                        .unwrap_or(LuaType::Unknown);
+                    let ret_type = TypeOps::Remove.apply(&ret_type, &LuaType::Nil);
+                    analyzer
+                        .db
+                        .get_type_index_mut()
+                        .bind_type(decl_id.into(), LuaTypeCache::InferType(ret_type));
+                    idx += 1;
+                }
+                return Some(());
+            } else {
+                for var_name in var_name_list {
+                    let position = var_name.get_position();
+                    let decl_id = LuaDeclId::new(analyzer.file_id, position);
+                    analyzer
+                        .db
+                        .get_type_index_mut()
+                        .bind_type(decl_id.into(), LuaTypeCache::InferType(LuaType::Unknown));
+                }
+                return Some(());
+            }
+        }
+        Err(InferFailReason::None) => {
+            for var_name in var_name_list {
+                let position = var_name.get_position();
+                let decl_id = LuaDeclId::new(analyzer.file_id, position);
+                analyzer
+                    .db
+                    .get_type_index_mut()
+                    .bind_type(decl_id.into(), LuaTypeCache::InferType(LuaType::Unknown));
+            }
+            return Some(());
+        }
+        Err(reason) => {
+            let mut idx = 0;
+            for var_name in var_name_list {
+                let position = var_name.get_position();
+                let decl_id = LuaDeclId::new(analyzer.file_id, position);
+                let unresolved = UnResolveIterVar {
+                    file_id: analyzer.file_id,
+                    decl_id,
+                    iter_expr: first_iter_expr.clone(),
+                    ret_idx: idx,
+                    reason: reason.clone(),
+                };
+                analyzer.add_unresolved(unresolved.into());
+                idx += 1;
+            }
+        }
+    }
+
+    Some(())
+}
+
+pub fn infer_for_range_iter_expr_func(
+    db: &mut DbIndex,
+    iter_expr_type: LuaType,
+) -> Option<Arc<LuaFunctionType>> {
+    match iter_expr_type {
+        LuaType::DocFunction(func) => Some(func),
+        LuaType::Ref(type_decl_id) => {
+            let type_decl = db.get_type_index().get_type_decl(&type_decl_id)?;
+            if type_decl.is_alias() {
+                let alias_origin = type_decl.get_alias_origin(db, None)?;
+                match alias_origin {
+                    LuaType::DocFunction(doc_func) => Some(doc_func),
+                    _ => None,
+                }
+            } else if type_decl.is_class() {
+                let operator_index = db.get_operator_index();
+                let operator_ids = operator_index
+                    .get_operators(&type_decl_id.into(), LuaOperatorMetaMethod::Call)?;
+                operator_ids
+                    .iter()
+                    .filter_map(|overload_id| {
+                        let operator = operator_index.get_operator(overload_id)?;
+                        let func = operator.get_operator_func();
+                        match func {
+                            LuaType::DocFunction(f) => {
+                                return Some(f.clone());
+                            }
+                            _ => None,
+                        }
+                    })
+                    .nth(0)
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/lua/mod.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/lua/mod.rs
@@ -1,4 +1,5 @@
 mod closure;
+mod for_range_stat;
 mod func_body;
 mod metatable;
 mod module;
@@ -9,12 +10,14 @@ use std::collections::HashMap;
 use closure::analyze_closure;
 pub use closure::analyze_return_point;
 use emmylua_parser::{LuaAst, LuaAstNode, LuaExpr};
+use for_range_stat::analyze_for_range_stat;
+pub use for_range_stat::infer_for_range_iter_expr_func;
 pub use func_body::LuaReturnPoint;
 use metatable::analyze_setmetatable;
 use module::analyze_chunk_return;
 use stats::{
-    analyze_assign_stat, analyze_for_range_stat, analyze_func_stat, analyze_local_func_stat,
-    analyze_local_stat, analyze_table_field,
+    analyze_assign_stat, analyze_func_stat, analyze_local_func_stat, analyze_local_stat,
+    analyze_table_field,
 };
 
 use crate::{
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/lua/stats.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/lua/stats.rs
@@ -1,14 +1,14 @@
 use emmylua_parser::{
-    BinaryOperator, LuaAssignStat, LuaAstNode, LuaAstToken, LuaExpr, LuaForRangeStat, LuaFuncStat,
-    LuaIndexExpr, LuaLocalFuncStat, LuaLocalStat, LuaTableField, LuaVarExpr, PathTrait,
+    BinaryOperator, LuaAssignStat, LuaAstNode, LuaExpr, LuaFuncStat, LuaIndexExpr,
+    LuaLocalFuncStat, LuaLocalStat, LuaTableField, LuaVarExpr, PathTrait,
 };
 
 use crate::{
     compilation::analyzer::{
         bind_type::{add_member, bind_type},
-        unresolve::{UnResolveDecl, UnResolveIterVar, UnResolveMember},
+        unresolve::{UnResolveDecl, UnResolveMember},
     },
-    db_index::{LuaDeclId, LuaMemberId, LuaMemberOwner, LuaOperatorMetaMethod, LuaType},
+    db_index::{LuaDeclId, LuaMemberId, LuaMemberOwner, LuaType},
     InferFailReason, LuaTypeCache, LuaTypeOwner,
 };
 
@@ -393,111 +393,6 @@ fn merge_type_owner_and_unresolve_expr(
     Some(())
 }
 
-pub fn analyze_for_range_stat(
-    analyzer: &mut LuaAnalyzer,
-    for_range_stat: LuaForRangeStat,
-) -> Option<()> {
-    let var_name_list = for_range_stat.get_var_name_list();
-    let first_iter_expr = for_range_stat.get_expr_list().next()?;
-    let first_iter_type = analyzer.infer_expr(&first_iter_expr);
-
-    match first_iter_type {
-        Ok(first_iter_type) => {
-            let iter_doc_func = match first_iter_type {
-                LuaType::DocFunction(doc_func) => Some(doc_func),
-                LuaType::Ref(type_decl_id) => {
-                    let type_decl = analyzer.db.get_type_index().get_type_decl(&type_decl_id)?;
-                    if type_decl.is_alias() {
-                        let alias_origin = type_decl.get_alias_origin(analyzer.db, None)?;
-                        match alias_origin {
-                            LuaType::DocFunction(doc_func) => Some(doc_func),
-                            _ => None,
-                        }
-                    } else if type_decl.is_class() {
-                        let operator_index = analyzer.db.get_operator_index();
-                        let operator_ids = operator_index
-                            .get_operators(&type_decl_id.into(), LuaOperatorMetaMethod::Call)?;
-                        operator_ids
-                            .iter()
-                            .filter_map(|overload_id| {
-                                let operator = operator_index.get_operator(overload_id)?;
-                                let func = operator.get_operator_func();
-                                match func {
-                                    LuaType::DocFunction(f) => {
-                                        return Some(f.clone());
-                                    }
-                                    _ => None,
-                                }
-                            })
-                            .nth(0)
-                    } else {
-                        None
-                    }
-                }
-                _ => None,
-            };
-
-            if let Some(doc_func) = iter_doc_func {
-                let multi_return = doc_func.get_multi_return();
-                let mut idx = 0;
-                for var_name in var_name_list {
-                    let position = var_name.get_position();
-                    let decl_id = LuaDeclId::new(analyzer.file_id, position);
-                    let ret_type = multi_return
-                        .get_type(idx)
-                        .unwrap_or(&LuaType::Unknown)
-                        .clone();
-                    analyzer
-                        .db
-                        .get_type_index_mut()
-                        .bind_type(decl_id.into(), LuaTypeCache::InferType(ret_type));
-                    idx += 1;
-                }
-                return Some(());
-            } else {
-                for var_name in var_name_list {
-                    let position = var_name.get_position();
-                    let decl_id = LuaDeclId::new(analyzer.file_id, position);
-                    analyzer
-                        .db
-                        .get_type_index_mut()
-                        .bind_type(decl_id.into(), LuaTypeCache::InferType(LuaType::Unknown));
-                }
-                return Some(());
-            }
-        }
-        Err(InferFailReason::None) => {
-            for var_name in var_name_list {
-                let position = var_name.get_position();
-                let decl_id = LuaDeclId::new(analyzer.file_id, position);
-                analyzer
-                    .db
-                    .get_type_index_mut()
-                    .bind_type(decl_id.into(), LuaTypeCache::InferType(LuaType::Unknown));
-            }
-            return Some(());
-        }
-        Err(reason) => {
-            let mut idx = 0;
-            for var_name in var_name_list {
-                let position = var_name.get_position();
-                let decl_id = LuaDeclId::new(analyzer.file_id, position);
-                let unresolved = UnResolveIterVar {
-                    file_id: analyzer.file_id,
-                    decl_id,
-                    iter_expr: first_iter_expr.clone(),
-                    ret_idx: idx,
-                    reason: reason.clone(),
-                };
-                analyzer.add_unresolved(unresolved.into());
-                idx += 1;
-            }
-        }
-    }
-
-    Some(())
-}
-
 pub fn analyze_func_stat(analyzer: &mut LuaAnalyzer, func_stat: LuaFuncStat) -> Option<()> {
     let closure = func_stat.get_closure()?;
     let func_name = func_stat.get_func_name()?;
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/unresolve/resolve.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/unresolve/resolve.rs
@@ -5,12 +5,12 @@ use emmylua_parser::{LuaAstNode, LuaExpr, LuaLocalStat, LuaTableExpr};
 use crate::{
     compilation::analyzer::{
         bind_type::{add_member, bind_type},
-        lua::analyze_return_point,
+        lua::{analyze_return_point, infer_for_range_iter_expr_func},
     },
     db_index::{DbIndex, LuaMemberOwner, LuaType},
     semantic::{infer_expr, LuaInferCache},
     InFiled, InferFailReason, LuaDeclId, LuaMember, LuaMemberId, LuaMemberKey, LuaSemanticDeclId,
-    LuaTypeCache, SignatureReturnStatus,
+    LuaTypeCache, SignatureReturnStatus, TypeOps,
 };
 
 use super::{
@@ -260,22 +260,33 @@ pub fn try_resolve_return_point(
 pub fn try_resolve_iter_var(
     db: &mut DbIndex,
     cache: &mut LuaInferCache,
-    iter_var: &UnResolveIterVar,
+    iter_var: &mut UnResolveIterVar,
 ) -> Option<bool> {
     if !check_reach_reason(db, cache, &iter_var.reason).unwrap_or(false) {
         return None;
     }
 
-    let expr_type = infer_expr(db, cache, iter_var.iter_expr.clone()).ok()?;
-    let func = match expr_type {
-        LuaType::DocFunction(func) => func,
-        _ => return Some(true),
+    let expr_type = match infer_expr(db, cache, iter_var.iter_expr.clone()) {
+        Ok(t) => t,
+        Err(InferFailReason::None) => return Some(true),
+        Err(reason) => {
+            iter_var.reason = reason;
+            return None;
+        }
+    };
+    let func = match infer_for_range_iter_expr_func(db, expr_type.clone()) {
+        Some(func) => func,
+        None => {
+            return Some(true);
+        }
     };
 
     let multi_return = func.get_multi_return();
-    let iter_type = multi_return
+    let mut iter_type = multi_return
         .get_type(iter_var.ret_idx)
-        .unwrap_or(&LuaType::Nil);
+        .cloned()
+        .unwrap_or(LuaType::Unknown);
+    iter_type = TypeOps::Remove.apply(&iter_type, &LuaType::Nil);
     let decl_id = iter_var.decl_id;
     bind_type(
         db,
diff --git a/crates/emmylua_code_analysis/src/compilation/test/for_range_var_infer_test.rs b/crates/emmylua_code_analysis/src/compilation/test/for_range_var_infer_test.rs
