Fix return cast usage

Fix #720

Author: CppCXY

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

@@ -1284,4 +1284,35 @@ end
         let a = ws.expr_ty("A");
         assert_eq!(ws.humanize_type(a), "Node");
     }
+
+    #[test]
+    fn test_return_cast_multi_file() {
+        let mut ws = VirtualWorkspace::new();
+        ws.def_file(
+            "test.lua",
+            r#"
+            local M = {}
+
+            --- @return boolean
+            --- @return_cast _obj function
+            function M.is_callable(_obj) end
+
+            return M
+            "#,
+        );
+        ws.def(
+            r#"
+            local test = require("test")
+
+            local obj
+
+            if test.is_callable(obj) then
+                o = obj
+            end
+            "#,
+        );
+        let a = ws.expr_ty("o");
+        let expected = LuaType::Function;
+        assert_eq!(a, expected);
+    }
 }

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

@@ -34,12 +34,10 @@ impl LuaFlowIndex {
         self.file_flow_tree.get(file_id)
     }
 
-    pub fn get_signature_cast(
-        &self,
-        file_id: &FileId,
-        signature_id: &LuaSignatureId,
-    ) -> Option<&LuaSignatureCast> {
-        self.signature_cast_cache.get(file_id)?.get(signature_id)
+    pub fn get_signature_cast(&self, signature_id: &LuaSignatureId) -> Option<&LuaSignatureCast> {
+        self.signature_cast_cache
+            .get(&signature_id.get_file_id())?
+            .get(signature_id)
     }
 
     pub fn add_signature_cast(

crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/call_flow.rs

@@ -88,10 +88,7 @@ pub fn get_type_at_call_expr(
                 _ => {}
             }
 
-            let Some(signature_cast) = db
-                .get_flow_index()
-                .get_signature_cast(&cache.get_file_id(), &signature_id)
-            else {
+            let Some(signature_cast) = db.get_flow_index().get_signature_cast(&signature_id) else {
                 return Ok(ResultTypeOrContinue::Continue);
             };
 
@@ -105,6 +102,7 @@ pub fn get_type_at_call_expr(
                     flow_node,
                     prefix_expr,
                     signature_cast,
+                    signature_id,
                     condition_flow,
                 ),
                 name => get_type_at_call_expr_by_signature_param_name(
@@ -200,6 +198,7 @@ fn get_type_at_call_expr_by_signature_self(
     flow_node: &FlowNode,
     call_prefix: LuaExpr,
     signature_cast: &LuaSignatureCast,
+    signature_id: LuaSignatureId,
     condition_flow: InferConditionFlow,
 ) -> Result<ResultTypeOrContinue, InferFailReason> {
     let LuaExpr::IndexExpr(call_prefix_index) = call_prefix else {
@@ -221,13 +220,18 @@ fn get_type_at_call_expr_by_signature_self(
     let antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
     let antecedent_type = get_type_at_flow(db, tree, cache, root, var_ref_id, antecedent_flow_id)?;
 
-    let Some(cast_op_type) = signature_cast.cast.to_node(root) else {
+    let Some(syntax_tree) = db.get_vfs().get_syntax_tree(&signature_id.get_file_id()) else {
+        return Ok(ResultTypeOrContinue::Continue);
+    };
+
+    let signature_root = syntax_tree.get_chunk_node();
+    let Some(cast_op_type) = signature_cast.cast.to_node(&signature_root) else {
         return Ok(ResultTypeOrContinue::Continue);
     };
 
     let result_type = cast_type(
         db,
-        cache.get_file_id(),
+        signature_id.get_file_id(),
         cast_op_type,
         antecedent_type,
         condition_flow,
@@ -291,13 +295,18 @@ fn get_type_at_call_expr_by_signature_param_name(
     let antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
     let antecedent_type = get_type_at_flow(db, tree, cache, root, var_ref_id, antecedent_flow_id)?;
 
-    let Some(cast_op_type) = signature_cast.cast.to_node(root) else {
+    let Some(syntax_tree) = db.get_vfs().get_syntax_tree(&signature_id.get_file_id()) else {
+        return Ok(ResultTypeOrContinue::Continue);
+    };
+
+    let signature_root = syntax_tree.get_chunk_node();
+    let Some(cast_op_type) = signature_cast.cast.to_node(&signature_root) else {
         return Ok(ResultTypeOrContinue::Continue);
     };
 
     let result_type = cast_type(
         db,
-        cache.get_file_id(),
+        signature_id.get_file_id(),
         cast_op_type,
         antecedent_type,
         condition_flow,