fix select function

Author: CppCXY

CHANGELOG.md

@@ -17,7 +17,7 @@
 
 `FIX` custom operator infer
 
-
+`FIX` Fix select function and add std.Select type 
 
 # 0.5.1 
 

crates/emmylua_code_analysis/resources/std/builtin.lua

@@ -106,12 +106,12 @@
 
 ---@class function
 
----@alias NotNull<T> T - ?
+---@alias std.NotNull<T> T - ?
 
----@alias Nullable<T> T + ?
+---@alias std.Nullable<T> T + ?
 
 --- built-in type for Select function
----@alias Select<T, StartOrLen> unknown
+---@alias std.Select<T, StartOrLen> unknown
 ---
 --- built-in type for Unpack function
----@alias Unpack<T, Start, End> unknown
+---@alias std.Unpack<T, Start, End> unknown

crates/emmylua_code_analysis/resources/std/global.lua

@@ -21,7 +21,7 @@
 ---@generic T
 ---@param v T
 ---@param message? string
----@return NotNull<T>
+---@return std.NotNull<T>
 function assert(v, message) end
 
 ---
@@ -105,7 +105,7 @@ function getmetatable(object) end
 --- the first absent index.
 ---@generic V
 ---@param t table<any, V> | V[]
----@return fun(tbl: any):int, NotNull<V>
+---@return fun(tbl: any):int, std.NotNull<V>
 function ipairs(t) end
 
 ---
@@ -216,7 +216,7 @@ function next(table, index) end
 --- traversal.
 ---@generic K, V
 ---@param t table<K, V> | V[]
----@return fun(tbl: any):K, NotNull<V>
+---@return fun(tbl: any):K, std.NotNull<V>
 function pairs(t) end
 ---
 --- Calls function `f` with the given arguments in *protected mode*. This
@@ -311,8 +311,8 @@ function require(modname) end
 --- the total number of extra arguments it received.
 ---@generic T, Num: integer | '#'
 ---@param index Num
----@param ... T
----@return Select<T, Num>
+---@param ... T...
+---@return std.Select<T..., Num>
 function select(index, ...) end
 
 ---
@@ -343,15 +343,12 @@ function setmetatable(table, metatable) end
 --- represents 10, 'B' represents 11, and so forth, with 'Z' representing 35. If
 --- the string `e` is not a valid numeral in the given base, the function
 --- returns **nil**.
----@overload fun(e:number):number
----@overload fun(e:number, base: int):number
----@overload fun(e:string):number|nil
----@overload fun(e:string, base: int):number|nil
----@overload fun(e:any):nil
----@overload fun(e:any, base:int):nil
+---@overload fun(e:number, base?: int):number
+---@overload fun(e:string, base?: int):number?
+---@overload fun(e:any, base?:int):nil
 ---@param e any
 ---@param base? int
----@return number|nil
+---@return number?
 function tonumber(e, base) end
 
 ---

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

@@ -756,10 +756,49 @@ impl LuaMultiReturn {
         }
     }
 
-    pub fn get_len(&self) -> usize {
+    pub fn get_len(&self) -> Option<i64> {
         match self {
-            LuaMultiReturn::Multi(types) => types.len(),
-            LuaMultiReturn::Base(_) => 1,
+            LuaMultiReturn::Base(_) => None,
+            LuaMultiReturn::Multi(types) => {
+                let basic_len = types.len() as i64;
+                if basic_len == 0 {
+                    return Some(0);
+                }
+
+                if let Some(LuaType::MuliReturn(last_multi)) = types.last() {
+                    let last_element_len = last_multi.get_len();
+                    return match last_element_len {
+                        Some(len) => Some(basic_len - 1 + len),
+                        None => None,
+                    };
+                }
+
+                Some(basic_len)
+            }
+        }
+    }
+
+    pub fn get_new_multi_from(&self, idx: usize) -> LuaMultiReturn {
+        match self {
+            LuaMultiReturn::Multi(types) => {
+                if types.len() == 0 {
+                    return LuaMultiReturn::Multi(Vec::new());
+                }
+
+                let mut new_types = Vec::new();
+                if idx < types.len() {
+                    new_types.extend_from_slice(&types[idx..]);
+                } else {
+                    let last = types.len() - 1;
+                    if let LuaType::MuliReturn(multi) = &types[last] {
+                        let rest_offset = idx - last;
+                        return multi.get_new_multi_from(rest_offset);
+                    }
+                }
+
+                LuaMultiReturn::Multi(new_types)
+            }
+            LuaMultiReturn::Base(t) => LuaMultiReturn::Base(t.clone()),
         }
     }
 

crates/emmylua_code_analysis/src/semantic/instantiate/instantiate_class_generic.rs

@@ -443,25 +443,30 @@ fn instantiate_select_call(source: &LuaType, index: &LuaType) -> LuaType {
     match num_or_len {
         NumOrLen::Num(i) => match multi_return {
             LuaMultiReturn::Base(_) => LuaType::MuliReturn(multi_return.clone().into()),
-            LuaMultiReturn::Multi(types) => {
-                let start = if i < 0 { types.len() as i64 + i } else { i - 1 };
-                if start < 0 || start as usize >= types.len() {
-                    return LuaType::Unknown;
+            LuaMultiReturn::Multi(_) => {
+                let total_len = multi_return.get_len();
+                if total_len.is_none() {
+                    return source.clone();
                 }
 
-                let new_types = types
-                    .iter()
-                    .skip(start as usize)
-                    .cloned()
-                    .collect::<Vec<_>>();
+                let total_len = total_len.unwrap();
+                let start = if i < 0 { total_len as i64 + i } else { i - 1 };
+                if start < 0 || start >= total_len {
+                    return source.clone();
+                }
 
-                LuaType::MuliReturn(LuaMultiReturn::Multi(new_types).into())
+                let multi = multi_return.get_new_multi_from(start as usize);
+                LuaType::MuliReturn(multi.into())
             }
         },
-        NumOrLen::Len => match multi_return {
-            LuaMultiReturn::Base(_) => LuaType::Integer,
-            LuaMultiReturn::Multi(types) => LuaType::DocIntegerConst(types.len() as i64),
-        },
+        NumOrLen::Len => {
+            let len = multi_return.get_len();
+            if let Some(len) = len {
+                LuaType::IntegerConst(len)
+            } else {
+                LuaType::Integer
+            }
+        }
         NumOrLen::LenUnknown => LuaType::Integer,
     }
 }

crates/emmylua_code_analysis/src/semantic/instantiate/instantiate_func_generic.rs

@@ -9,7 +9,7 @@ use crate::{
 };
 
 use super::{
-    instantiate_class_generic::instantiate_doc_function, tpl_pattern::tpl_pattern_match,
+    instantiate_class_generic::instantiate_doc_function, tpl_pattern::{tpl_pattern_match, variadic_tpl_pattern_match},
     type_substitutor::TypeSubstitutor,
 };
 
@@ -86,6 +86,12 @@ fn match_tpl_args(
             continue;
         };
 
+        if let LuaType::Variadic(inner) = func_param_type {
+            let rest_arg_types = &arg_types[i..];
+            variadic_tpl_pattern_match(&inner, rest_arg_types, &mut substitutor);
+            break;
+        }
+
         tpl_pattern_match(
             db,
             infer_config,

crates/emmylua_code_analysis/src/semantic/instantiate/test.rs

@@ -25,4 +25,42 @@ mod test {
         let expected = ws.ty("async fun(a: number, b: string, c:boolean): number");
         assert_eq!(ty, expected);
     }
+
+    #[test]
+    fn test_select_type() {
+        let mut ws = crate::VirtualWorkspace::new_with_init_std_lib();
+        ws.def(r#"
+        ---@param ... string
+        function ffff(...)
+            a, b, c = select(2, ...)
+        end
+        "#);
+
+        let a_ty = ws.expr_ty("a");
+        let b_ty = ws.expr_ty("b");
+        let c_ty = ws.expr_ty("c");
+        let expected = ws.ty("string");
+        assert_eq!(a_ty, expected);
+        assert_eq!(b_ty, expected);
+        assert_eq!(c_ty, expected);
+
+        ws.def(r#"
+        e, f = select(2, "a", "b", "c")
+        "#);
+
+        let e = ws.expr_ty("e");
+        let expected = ws.expr_ty("'b'");
+        let f = ws.expr_ty("f");
+        let expected_f = ws.expr_ty("'c'");
+        assert_eq!(e, expected);
+        assert_eq!(f, expected_f);
+
+        ws.def(r#"
+        h = select('#', "a", "b")
+        "#);
+
+        let h = ws.expr_ty("h");
+        let expected = ws.expr_ty("2");
+        assert_eq!(h, expected);
+    }
 }

crates/emmylua_code_analysis/src/semantic/instantiate/tpl_pattern.rs

@@ -303,7 +303,7 @@ fn func_varargs_tpl_pattern_match(
     Some(())
 }
 
-fn variadic_tpl_pattern_match(
+pub fn variadic_tpl_pattern_match(
     tpl: &LuaType,
     target_rest_types: &[LuaType],
     substitutor: &mut TypeSubstitutor,