support TypeGuard<T>

Author: CppCXY

CHANGELOG.md

@@ -58,6 +58,24 @@ end
 
 `NEW` Support `Lua 5.5` global decl grammar
 
+`NEW` Support `TypeGuard<T>` as return type. For example:
+```lua
+
+---@return TypeGuard<string>
+local function is_string(value)
+    return type(value) == "string"
+end
+
+local a
+
+if is_string(a) then
+    print(a:sub(1, 1))
+else
+    print("a is not a string")
+end
+```
+
+
 # 0.7.2
 
 `FIX` Fix reading configuration file encoded with UTF-8 BOM

crates/emmylua_code_analysis/resources/std/builtin.lua

@@ -125,3 +125,5 @@
 ---@alias collectgarbage_opt std.collectgarbage_opt
 
 ---@alias metatable std.metatable
+
+---@alias TypeGuard<T> boolean

crates/emmylua_code_analysis/src/compilation/analyzer/doc/infer_type.rs

@@ -270,6 +270,12 @@ fn infer_special_generic_type(
             ));
         }
         "std.Unpack" => {}
+        "TypeGuard" => {
+            let first_doc_param_type = generic_type.get_generic_types()?.get_types().next()?;
+            let first_param = infer_type(analyzer, first_doc_param_type);
+
+            return Some(LuaType::TypeGuard(first_param.into()));
+        }
         _ => {}
     }
 

crates/emmylua_code_analysis/src/db_index/flow/type_assert.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{ops::Deref, sync::Arc};
 
 use crate::{infer_expr, DbIndex, InferFailReason, LuaInferCache, LuaType, TypeOps};
 use emmylua_parser::{LuaAstNode, LuaCallExpr, LuaExpr, LuaSyntaxId, LuaSyntaxNode};
@@ -182,28 +182,32 @@ fn call_assertion(
     let Some(signature) = db.get_signature_index().get(&signature_id) else {
         return Err(InferFailReason::None);
     };
-    // donot change the condition
-    if !signature.get_return_type().is_boolean() {
-        return Err(InferFailReason::None);
-    }
 
-    let Some(cast) = db.get_flow_index().get_call_cast(signature_id) else {
-        return Err(InferFailReason::None);
-    };
-
-    let param_name = if param_idx >= 0 {
-        let Some(param_name) = signature.get_param_name_by_id(param_idx as usize) else {
-            return Err(InferFailReason::None);
-        };
+    let return_type = signature.get_return_type();
+    // donot change the condition
+    match return_type {
+        LuaType::Boolean => {
+            let Some(cast) = db.get_flow_index().get_call_cast(signature_id) else {
+                return Err(InferFailReason::None);
+            };
+
+            let param_name = if param_idx >= 0 {
+                let Some(param_name) = signature.get_param_name_by_id(param_idx as usize) else {
+                    return Err(InferFailReason::None);
+                };
 
-        param_name
-    } else {
-        "self".to_string()
-    };
+                param_name
+            } else {
+                "self".to_string()
+            };
 
-    let Some(typeassert) = cast.get(&param_name) else {
-        return Err(InferFailReason::None);
-    };
+            let Some(typeassert) = cast.get(&param_name) else {
+                return Err(InferFailReason::None);
+            };
 
-    Ok(typeassert.clone())
+            Ok(typeassert.clone())
+        }
+        LuaType::TypeGuard(inner) => Ok(TypeAssertion::Force(inner.deref().clone())),
+        _ => return Err(InferFailReason::None),
+    }
 }

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

@@ -88,6 +88,10 @@ pub fn humanize_type(db: &DbIndex, ty: &LuaType, level: RenderLevel) -> String {
         LuaType::MultiLineUnion(multi_union) => {
             humanize_multi_line_union_type(db, multi_union, level)
         }
+        LuaType::TypeGuard(inner) => {
+            let type_str = humanize_type(db, inner, level.next_level());
+            format!("TypeGuard<{}>", type_str)
+        }
         _ => "unknown".to_string(),
     }
 }

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

@@ -58,6 +58,7 @@ pub enum LuaType {
     Namespace(ArcIntern<SmolStr>),
     Call(Arc<LuaAliasCallType>),
     MultiLineUnion(Arc<LuaMultiLineUnion>),
+    TypeGuard(Arc<LuaType>),
 }
 
 impl PartialEq for LuaType {
@@ -103,6 +104,7 @@ impl PartialEq for LuaType {
             (LuaType::DocIntegerConst(a), LuaType::DocIntegerConst(b)) => a == b,
             (LuaType::Namespace(a), LuaType::Namespace(b)) => a == b,
             (LuaType::MultiLineUnion(a), LuaType::MultiLineUnion(b)) => a == b,
+            (LuaType::TypeGuard(a), LuaType::TypeGuard(b)) => a == b,
             _ => false, // 不同变体之间不相等
         }
     }
@@ -174,6 +176,10 @@ impl Hash for LuaType {
                 let ptr = Arc::as_ptr(a);
                 (43, ptr).hash(state)
             }
+            LuaType::TypeGuard(a) => {
+                let ptr = Arc::as_ptr(a);
+                (44, ptr).hash(state)
+            }
         }
     }
 }
@@ -402,6 +408,10 @@ impl LuaType {
     pub fn is_member_owner(&self) -> bool {
         matches!(self, LuaType::Ref(_) | LuaType::TableConst(_))
     }
+
+    pub fn is_type_guard(&self) -> bool {
+        matches!(self, LuaType::TypeGuard(_))
+    }
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]

crates/emmylua_code_analysis/src/semantic/infer/infer_call/mod.rs

@@ -568,6 +568,7 @@ pub(crate) fn unwrapp_return_type(
                 }
             }
         }
+        LuaType::TypeGuard(_) => return Ok(LuaType::Boolean),
         _ => {}
     }
 

crates/emmylua_code_analysis/src/semantic/type_check/mod.rs

@@ -112,6 +112,12 @@ fn check_general_type_compact(
             compact_type,
             check_guard.next_level()?,
         ),
+        LuaType::TypeGuard(_) => {
+            if compact_type.is_boolean() {
+                return Ok(());
+            }
+            return Err(TypeCheckFailReason::TypeNotMatch);
+        }
         _ => Err(TypeCheckFailReason::TypeNotMatch),
     }
 }
@@ -142,6 +148,7 @@ fn escape_type(db: &DbIndex, typ: &LuaType) -> Option<LuaType> {
             let union = multi_union.to_union();
             return Some(union);
         }
+        LuaType::TypeGuard(_) => return Some(LuaType::Boolean),
         _ => {}
     }