optimize auto infer return and type narrow

Author: CppCXY

crates/emmylua_code_analysis/src/compilation/analyzer/lua/closure.rs

@@ -7,7 +7,7 @@ use crate::{
         UnResolveClosureParams, UnResolveClosureReturn, UnResolveReturn,
     },
     db_index::{LuaDocReturnInfo, LuaSignatureId},
-    SignatureReturnStatus,
+    infer_expr, DbIndex, LuaInferCache, LuaMultiReturn, LuaType, SignatureReturnStatus, TypeOps,
 };
 
 use super::{func_body::analyze_func_body_returns, LuaAnalyzer, LuaReturnPoint};
@@ -86,19 +86,20 @@ fn analyze_return(
 
     let block = closure.get_block()?;
     let return_points = analyze_func_body_returns(block);
-    let returns = match analyze_return_point(analyzer, &return_points) {
-        Some(returns) => returns,
-        None => {
-            let unresolve = UnResolveReturn {
-                file_id: analyzer.file_id,
-                signature_id: signature_id.clone(),
-                return_points,
-            };
-
-            analyzer.add_unresolved(unresolve.into());
-            return None;
-        }
-    };
+    let returns =
+        match analyze_return_point(&analyzer.db, &mut analyzer.infer_cache, &return_points) {
+            Some(returns) => returns,
+            None => {
+                let unresolve = UnResolveReturn {
+                    file_id: analyzer.file_id,
+                    signature_id: signature_id.clone(),
+                    return_points,
+                };
+
+                analyzer.add_unresolved(unresolve.into());
+                return None;
+            }
+        };
     let signature = analyzer
         .db
         .get_signature_index_mut()
@@ -134,43 +135,37 @@ fn analyze_lambda_returns(
     Some(())
 }
 
-fn analyze_return_point(
-    analyzer: &mut LuaAnalyzer,
+pub fn analyze_return_point(
+    db: &DbIndex,
+    cache: &mut LuaInferCache,
     return_points: &Vec<LuaReturnPoint>,
 ) -> Option<Vec<LuaDocReturnInfo>> {
-    let mut return_infos = Vec::new();
+    let mut return_type = LuaType::Unknown;
     for point in return_points {
         match point {
             LuaReturnPoint::Expr(expr) => {
-                let expr_type = analyzer.infer_expr(&expr)?;
-                if return_infos.is_empty() {
-                    return_infos.push(LuaDocReturnInfo {
-                        name: None,
-                        type_ref: expr_type,
-                        description: None,
-                    });
-                } else {
-                    // todo merge two type
-                    // let has_return = return_infos[0].type_ref.clone();
-                    return_infos[0].type_ref = expr_type;
-                }
+                let expr_type = infer_expr(db, cache, expr.clone())?;
+                return_type = TypeOps::Union.apply(&return_type, &expr_type);
             }
             LuaReturnPoint::MuliExpr(exprs) => {
-                // todo merge type
-                if return_infos.is_empty() {
-                    for expr in exprs {
-                        let expr_type = analyzer.infer_expr(&expr)?;
-                        return_infos.push(LuaDocReturnInfo {
-                            name: None,
-                            type_ref: expr_type,
-                            description: None,
-                        });
-                    }
+                let mut multi_return = vec![];
+                for expr in exprs {
+                    let expr_type = infer_expr(db, cache, expr.clone())?;
+                    multi_return.push(expr_type);
                 }
+                let typ = LuaType::MuliReturn(LuaMultiReturn::Multi(multi_return).into());
+                return_type = TypeOps::Union.apply(&return_type, &typ);
+            }
+            LuaReturnPoint::Nil => {
+                return_type = TypeOps::Union.apply(&return_type, &LuaType::Nil);
             }
             _ => {}
         }
     }
 
-    Some(return_infos)
+    Some(vec![LuaDocReturnInfo {
+        type_ref: return_type,
+        description: None,
+        name: None,
+    }])
 }

crates/emmylua_code_analysis/src/compilation/analyzer/lua/mod.rs

@@ -6,6 +6,7 @@ mod stats;
 use std::collections::HashMap;
 
 use closure::analyze_closure;
+pub use closure::analyze_return_point;
 use emmylua_parser::{LuaAst, LuaAstNode, LuaExpr};
 pub use func_body::LuaReturnPoint;
 use module::analyze_chunk_return;

crates/emmylua_code_analysis/src/compilation/analyzer/unresolve/resolve.rs

@@ -1,6 +1,6 @@
 use crate::{
-    compilation::analyzer::lua::LuaReturnPoint,
-    db_index::{DbIndex, LuaDocReturnInfo, LuaMemberOwner, LuaType},
+    compilation::analyzer::lua::analyze_return_point,
+    db_index::{DbIndex, LuaMemberOwner, LuaType},
     semantic::{infer_expr, LuaInferCache},
     LuaPropertyOwnerId, SignatureReturnStatus,
 };
@@ -101,51 +101,7 @@ pub fn try_resolve_return_point(
     cache: &mut LuaInferCache,
     return_: &UnResolveReturn,
 ) -> Option<bool> {
-    let mut is_nullable = false;
-    let mut return_docs = Vec::new();
-    for return_point in &return_.return_points {
-        match return_point {
-            LuaReturnPoint::Expr(expr) => {
-                let expr_type = infer_expr(db, cache, expr.clone())?;
-                if return_docs.is_empty() {
-                    return_docs.push(LuaDocReturnInfo {
-                        name: None,
-                        type_ref: expr_type,
-                        description: None,
-                    });
-                } else {
-                    let last = return_docs.first_mut()?;
-                    last.type_ref = expr_type;
-                }
-            }
-            LuaReturnPoint::MuliExpr(exprs) => {
-                if return_docs.is_empty() {
-                    for expr in exprs {
-                        let expr_type = infer_expr(db, cache, expr.clone())?;
-
-                        return_docs.push(LuaDocReturnInfo {
-                            name: None,
-                            type_ref: expr_type,
-                            description: None,
-                        });
-                    }
-                }
-            }
-            LuaReturnPoint::Nil => {
-                is_nullable = true;
-            }
-            LuaReturnPoint::Error => {}
-        }
-    }
-
-    if is_nullable {
-        for doc in &mut return_docs {
-            if !doc.type_ref.is_nullable() {
-                doc.type_ref = LuaType::Nullable(doc.type_ref.clone().into());
-            }
-        }
-    }
-
+    let return_docs = analyze_return_point(&db, cache, &return_.return_points)?;
     let signature = db
         .get_signature_index_mut()
         .get_mut(&return_.signature_id)?;

crates/emmylua_code_analysis/src/db_index/type/type_ops/narrow_type.rs

@@ -88,6 +88,7 @@ pub fn narrow_down_type(source: LuaType, target: LuaType) -> Option<LuaType> {
             | LuaType::TableGeneric(_) => return Some(source),
             _ => {}
         },
+        LuaType::Instance(base) => return narrow_down_type(source, base.get_base().clone()),
         LuaType::Unknown => return Some(source),
         _ => {
             if target.is_unknown() {
@@ -100,24 +101,26 @@ pub fn narrow_down_type(source: LuaType, target: LuaType) -> Option<LuaType> {
 
     match &source {
         LuaType::Union(union) => {
-            let union_types = union
+            let mut union_types = union
                 .get_types()
                 .iter()
                 .filter_map(|t| narrow_down_type(t.clone(), target.clone()))
                 .collect::<Vec<_>>();
 
+            union_types.dedup();
             return match union_types.len() {
                 0 => Some(target),
                 1 => Some(union_types[0].clone()),
                 _ => Some(LuaType::Union(LuaUnionType::new(union_types).into())),
             };
         }
         LuaType::Nullable(inner) => {
-            let union_types = vec![LuaType::Nil, (**inner).clone()]
+            let mut union_types = vec![LuaType::Nil, (**inner).clone()]
                 .iter()
                 .filter_map(|t| narrow_down_type(t.clone(), target.clone()))
                 .collect::<Vec<_>>();
 
+            union_types.dedup();
             return match union_types.len() {
                 0 => Some(target),
                 1 => Some(union_types[0].clone()),

crates/emmylua_code_analysis/src/db_index/type/type_ops/union_type.rs

@@ -24,6 +24,13 @@ pub fn union_type(source: LuaType, target: LuaType) -> LuaType {
         // boolean | boolean const
         (LuaType::Boolean, LuaType::BooleanConst(_)) => LuaType::Boolean,
         (LuaType::BooleanConst(_), LuaType::Boolean) => LuaType::Boolean,
+        (LuaType::BooleanConst(left), LuaType::BooleanConst(right)) => {
+            if left == right {
+                LuaType::BooleanConst(left.clone())
+            } else {
+                LuaType::Boolean
+            }
+        }
         // table | table const
         (LuaType::Table, LuaType::TableConst(_)) => LuaType::Table,
         (LuaType::TableConst(_), LuaType::Table) => LuaType::Table,