feat(isolated-declarations): improve inferring the return type from function (#3750)

crates/oxc_isolated_declarations/src/class.rs

@@ -265,7 +265,7 @@ impl<'a> IsolatedDeclarations<'a> {
             }
         }
 
-        let mut inferred_accessor_type: FxHashMap<Atom<'a>, Box<'a, TSTypeAnnotation<'a>>> =
+        let mut inferred_accessor_types: FxHashMap<Atom<'a>, Box<'a, TSTypeAnnotation<'a>>> =
             FxHashMap::default();
 
         // Infer get accessor return type from set accessor
@@ -282,7 +282,7 @@ impl<'a> IsolatedDeclarations<'a> {
                     continue;
                 };
                 let name = self.ast.new_atom(&name);
-                if inferred_accessor_type.contains_key(&name) {
+                if inferred_accessor_types.contains_key(&name) {
                     // We've inferred that accessor type already
                     continue;
                 }
@@ -291,7 +291,7 @@ impl<'a> IsolatedDeclarations<'a> {
                     MethodDefinitionKind::Get => {
                         let return_type = self.infer_function_return_type(function);
                         if let Some(return_type) = return_type {
-                            inferred_accessor_type.insert(name, self.ast.copy(&return_type));
+                            inferred_accessor_types.insert(name, self.ast.copy(&return_type));
                         }
                     }
                     MethodDefinitionKind::Set => {
@@ -305,7 +305,7 @@ impl<'a> IsolatedDeclarations<'a> {
                                     |t| Some(self.ast.copy(t)),
                                 );
                             if let Some(type_annotation) = type_annotation {
-                                inferred_accessor_type.insert(name, type_annotation);
+                                inferred_accessor_types.insert(name, type_annotation);
                             }
                         }
                     }
@@ -338,7 +338,7 @@ impl<'a> IsolatedDeclarations<'a> {
                             |n| {
                                 self.transform_set_accessor_params(
                                     &function.params,
-                                    inferred_accessor_type
+                                    inferred_accessor_types
                                         .get(&self.ast.new_atom(&n))
                                         .map(|t| self.ast.copy(t)),
                                 )
@@ -368,7 +368,7 @@ impl<'a> IsolatedDeclarations<'a> {
                         }
                         MethodDefinitionKind::Get => {
                             let rt = method.key.static_name().and_then(|name| {
-                                inferred_accessor_type
+                                inferred_accessor_types
                                     .get(&self.ast.new_atom(&name))
                                     .map(|t| self.ast.copy(t))
                             });

crates/oxc_isolated_declarations/src/declaration.rs

@@ -75,10 +75,12 @@ impl<'a> IsolatedDeclarations<'a> {
             }
             if init.is_none() && binding_type.is_none() {
                 binding_type = Some(self.ast.ts_unknown_keyword(SPAN));
-                self.error(
+                if !decl.init.as_ref().is_some_and(Expression::is_function) {
-                  OxcDiagnostic::error("Variable must have an explicit type annotation with --isolatedDeclarations.")
+                    self.error(
-                      .with_label(decl.id.span()),
+                      OxcDiagnostic::error("Variable must have an explicit type annotation with --isolatedDeclarations.")
-              );
+                          .with_label(decl.id.span()),
+                  );
+                }
             }
         }
         let id = binding_type.map_or_else(

crates/oxc_isolated_declarations/src/return_type.rs

@@ -5,15 +5,38 @@ use oxc_ast::{
     },
     AstBuilder, Visit,
 };
-use oxc_span::{Atom, GetSpan};
+use oxc_span::{Atom, GetSpan, SPAN};
 use oxc_syntax::scope::ScopeFlags;
 
 use crate::{diagnostics::type_containing_private_name, IsolatedDeclarations};
 
-/// Infer return type from return statement. Does not support multiple return statements.
+/// Infer return type from return statement.
+/// ```ts
+/// function foo() {
+///    return 1;
+/// }
+/// // inferred type is number
+///
+/// function bar() {
+///   if (true) {
+///    return;
+///   }
+///   return 1;
+/// }
+/// // inferred type is number | undefined
+///
+/// function baz() {
+///  if (true) {
+///   return null;
+///  }
+///  return 1;
+/// }
+/// // We can't infer return type if there are multiple return statements with different types
+/// ```
+#[allow(clippy::option_option)]
 pub struct FunctionReturnType<'a> {
     ast: AstBuilder<'a>,
-    return_expression: Option<Expression<'a>>,
+    return_expression: Option<Option<Expression<'a>>>,
     value_bindings: Vec<Atom<'a>>,
     type_bindings: Vec<Atom<'a>>,
     return_statement_count: u8,
@@ -36,48 +59,57 @@ impl<'a> FunctionReturnType<'a> {
 
         visitor.visit_function_body(body);
 
-        if visitor.return_statement_count > 1 {
+        let expr = visitor.return_expression??;
-            return None;
+        let Some(mut expr_type) = transformer.infer_type_from_expression(&expr) else {
-        }
+            // Avoid report error in parent function
+            return if expr.is_function() {
+                Some(transformer.ast.ts_unknown_keyword(SPAN))
+            } else {
+                None
+            };
+        };
 
-        visitor.return_expression.and_then(|expr| {
+        if let Some((reference_name, is_value)) = match &expr_type {
-            let expr_type = transformer.infer_type_from_expression(&expr)?;
+            TSType::TSTypeReference(type_reference) => {
-
+                if let TSTypeName::IdentifierReference(ident) = &type_reference.type_name {
-            if let Some((reference_name, is_value)) = match &expr_type {
+                    Some((ident.name.clone(), false))
-                TSType::TSTypeReference(type_reference) => {
-                    if let TSTypeName::IdentifierReference(ident) = &type_reference.type_name {
-                        Some((ident.name.clone(), false))
-                    } else {
-                        None
-                    }
-                }
-                TSType::TSTypeQuery(query) => {
-                    if let TSTypeQueryExprName::IdentifierReference(ident) = &query.expr_name {
-                        Some((ident.name.clone(), true))
-                    } else {
-                        None
-                    }
-                }
-                _ => None,
-            } {
-                let is_defined_in_current_scope = if is_value {
-                    visitor.value_bindings.contains(&reference_name)
                 } else {
-                    visitor.type_bindings.contains(&reference_name)
+                    None
-                };
-
-                if is_defined_in_current_scope {
-                    transformer.error(type_containing_private_name(
-                        &reference_name,
-                        expr_type
-                            .get_identifier_reference()
-                            .map_or_else(|| expr_type.span(), |ident| ident.span),
-                    ));
                 }
             }
+            TSType::TSTypeQuery(query) => {
+                if let TSTypeQueryExprName::IdentifierReference(ident) = &query.expr_name {
+                    Some((ident.name.clone(), true))
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        } {
+            let is_defined_in_current_scope = if is_value {
+                visitor.value_bindings.contains(&reference_name)
+            } else {
+                visitor.type_bindings.contains(&reference_name)
+            };
 
-            Some(expr_type)
+            if is_defined_in_current_scope {
-        })
+                transformer.error(type_containing_private_name(
+                    &reference_name,
+                    expr_type
+                        .get_identifier_reference()
+                        .map_or_else(|| expr_type.span(), |ident| ident.span),
+                ));
+            }
+        }
+
+        //
+        if visitor.return_statement_count > 1 {
+            let types = transformer
+                .ast
+                .new_vec_from_iter([expr_type, transformer.ast.ts_undefined_keyword(SPAN)]);
+            expr_type = transformer.ast.ts_union_type(SPAN, types);
+        }
+        Some(expr_type)
     }
 }
 
@@ -109,8 +141,16 @@ impl<'a> Visit<'a> for FunctionReturnType<'a> {
     fn visit_return_statement(&mut self, stmt: &ReturnStatement<'a>) {
         self.return_statement_count += 1;
         if self.return_statement_count > 1 {
-            return;
+            if let Some(expr) = &self.return_expression {
+                // if last return statement is not empty, we can't infer return type
+                if expr.is_some() {
+                    self.return_expression = None;
+                    return;
+                }
+            } else {
+                return;
+            }
         }
-        self.return_expression = self.ast.copy(&stmt.argument);
+        self.return_expression = Some(self.ast.copy(&stmt.argument));
     }
 }

crates/oxc_isolated_declarations/tests/fixtures/infer-return-type.ts

@@ -0,0 +1,20 @@
+function foo() {
+   return 1;
+}
+// inferred type is number
+
+function bar() {
+  if (true) {
+   return;
+  }
+  return 1;
+}
+// inferred type is number | undefined
+
+function baz() {
+ if (true) {
+  return null;
+ }
+ return 1;
+}
+// We can't infer return type if there are multiple return statements with different types

crates/oxc_isolated_declarations/tests/snapshots/infer-return-type.snap

@@ -0,0 +1,21 @@
+---
+source: crates/oxc_isolated_declarations/tests/mod.rs
+input_file: crates/oxc_isolated_declarations/tests/fixtures/infer-return-type.ts
+---
+==================== .D.TS ====================
+
+declare function foo(): number;
+declare function bar(): ((number) | (undefined));
+declare function baz();
+
+
+==================== Errors ====================
+
+  x Function must have an explicit return type annotation with
+  | --isolatedDeclarations.
+    ,-[14:10]
+ 13 | 
+ 14 | function baz() {
+    :          ^^^
+ 15 |  if (true) {
+    `----