Complete function detection logic

Author: mohebifar

.vscode/settings.json

@@ -1,4 +1,9 @@
 {
   "typescript.tsdk": "node_modules/typescript/lib",
-  "rust-analyzer.linkedProjects": ["./Cargo.toml"]
+  "rust-analyzer.linkedProjects": ["./Cargo.toml"],
+  "rust-analyzer.check.command": "clippy",
+  "[rust]": {
+    "editor.defaultFormatter": "rust-lang.rust-analyzer",
+    "editor.formatOnSave": true,
+  }
 }

packages/swc-plugin/transform/src/lib.rs

@@ -1,8 +1,8 @@
-use std::{cell::RefCell, ptr::null, rc::Rc};
+use std::{cell::RefCell, rc::Rc};
 
 use serde::Deserialize;
-use swc_common::{chain, comments::Comments, FileName};
-use swc_ecma_visit::{Fold, VisitMut};
+use swc_common::{comments::Comments, FileName};
+use swc_ecma_visit::Fold;
 
 pub use crate::utils::{component_finder, State};
 
@@ -28,15 +28,12 @@ fn default_skip_components() -> Vec<String> {
     vec![]
 }
 
-pub fn react_unforget<C>(_file_name: FileName, config: Config, _comments: C) -> impl Fold + VisitMut
+pub fn react_unforget<C>(_file_name: FileName, config: Config, _comments: C) -> impl Fold
 where
     C: Comments,
 {
     let state: Rc<RefCell<State>> = Default::default();
     let config = Rc::new(config);
 
-    println!("Calling component_finder");
-    let finder = crate::utils::component_finder(config, state);
-
-    return finder
+    crate::utils::component_finder(config, state)
 }

packages/swc-plugin/transform/src/models/component.rs

@@ -1 +1,58 @@
-pub struct Component {}
+#[allow(clippy::all)]
+use std::borrow::Borrow;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
+
+use swc_ecma_ast::{BlockStmt, BlockStmtOrExpr, Stmt};
+
+use crate::models::component_segment::ComponentSegment;
+use crate::utils::babel_shims::GenericFn;
+
+#[derive(Debug)]
+pub struct Component {
+    pub name: String,
+    pub is_hook: bool,
+    pub function: Rc<RefCell<GenericFn>>,
+    segments: HashMap<Stmt, ComponentSegment>,
+}
+
+impl Component {
+    pub fn new(name: String, is_hook: bool, function: Rc<RefCell<GenericFn>>) -> Self {
+        function.borrow_mut().ensure_body_block_stmt();
+        println!("Creating component: {}", name);
+
+        Self {
+            name,
+            is_hook,
+            function,
+            segments: Default::default(),
+        }
+    }
+
+    pub fn add_segment(&mut self, stmt: Stmt, segment: ComponentSegment) {
+        self.segments.insert(stmt, segment);
+    }
+
+    pub fn get_segments(&self) -> &HashMap<Stmt, ComponentSegment> {
+        self.segments.borrow()
+    }
+
+    pub fn get_name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn prepare(&self) {}
+
+    fn get_body_block_stmt(&self) -> BlockStmt {
+        let body = self.function.as_ref().borrow().get_body();
+
+        match body {
+            Some(body) => match body {
+                BlockStmtOrExpr::BlockStmt(block_stmt) => block_stmt,
+                BlockStmtOrExpr::Expr(_) => panic!("Expected block statement"),
+            },
+            None => panic!("Function has no body"),
+        }
+    }
+}

packages/swc-plugin/transform/src/models/component_segment.rs

@@ -0,0 +1,2 @@
+#[derive(Debug)]
+pub struct ComponentSegment {}

packages/swc-plugin/transform/src/models/mod.rs

@@ -1 +1,2 @@
 pub mod component;
+pub mod component_segment;

packages/swc-plugin/transform/src/utils/ast_tools/mod.rs

@@ -1 +1,2 @@
-pub mod generic;
+pub mod generic;
+pub mod returns_like_component;

packages/swc-plugin/transform/src/utils/ast_tools/returns_like_component.rs

@@ -0,0 +1,66 @@
+use swc_ecma_ast::{ArrowExpr, Expr, FnDecl, Function, Lit, ReturnStmt};
+use swc_ecma_utils::ExprExt;
+use swc_ecma_visit::{noop_visit_type, Visit};
+
+pub struct ReturnTypeVisitorState {
+    return_type_mathces_component: bool,
+    return_count: u32,
+}
+
+impl Default for ReturnTypeVisitorState {
+    fn default() -> Self {
+        Self {
+            return_type_mathces_component: true,
+            return_count: 0,
+        }
+    }
+}
+
+impl ReturnTypeVisitorState {
+    pub fn does_return_type_match_component(&self) -> bool {
+        self.return_type_mathces_component && self.return_count > 0
+    }
+}
+
+#[derive(Default)]
+pub struct VisitReturnStatements {
+    pub state: ReturnTypeVisitorState,
+}
+
+impl Visit for VisitReturnStatements {
+    noop_visit_type!();
+
+    // We don't want to visit nested functions
+    fn visit_fn_decl(&mut self, _: &FnDecl) {}
+    fn visit_function(&mut self, _: &Function) {}
+    fn visit_arrow_expr(&mut self, _: &ArrowExpr) {}
+
+    fn visit_return_stmt(&mut self, n: &ReturnStmt) {
+        self.state.return_count += 1;
+
+        if let Some(arg) = &n.arg {
+            if !is_component_return_type(arg.as_expr()) {
+                self.state.return_type_mathces_component = false;
+            }
+        }
+    }
+}
+
+/**
+ * Check if the given expression qualifies as a component return type.
+ */
+pub fn is_component_return_type(node: &Expr) -> bool {
+    if let Expr::JSXElement(_) = node {
+        return true;
+    }
+
+    if let Expr::JSXFragment(_) = node {
+        return true;
+    }
+
+    if let Expr::Lit(literal) = node {
+        return matches!(literal, Lit::Null(_));
+    }
+
+    false
+}

packages/swc-plugin/transform/src/utils/babel_shims/mod.rs

@@ -1,6 +1,70 @@
-use swc_ecma_ast::{ArrowExpr, Function};
+use std::ops::Deref;
 
-pub enum Fn {
-    Function(Function),
-    ArrowFunction(ArrowExpr),
+use swc_common::DUMMY_SP;
+use swc_ecma_ast::{ArrowExpr, BlockStmt, BlockStmtOrExpr, Function};
+use swc_ecma_utils::ExprFactory;
+use swc_ecma_visit::VisitWith;
+
+use super::ast_tools::returns_like_component::{is_component_return_type, VisitReturnStatements};
+
+#[derive(Debug, Clone)]
+pub enum GenericFn {
+    Function(Box<Function>),
+    ArrowFunction(Box<ArrowExpr>),
+}
+
+impl GenericFn {
+    pub fn get_body(&self) -> Option<BlockStmtOrExpr> {
+        match self {
+            GenericFn::Function(f) => {
+                let body = &f.body;
+                body.as_ref()
+                    .map(|body| BlockStmtOrExpr::BlockStmt(body.clone()))
+            }
+            GenericFn::ArrowFunction(f) => Some(f.body.deref().clone()),
+        }
+    }
+
+    pub fn update_body(&mut self, body: BlockStmt) {
+        match self {
+            GenericFn::Function(f) => {
+                f.body = Some(body);
+            }
+            GenericFn::ArrowFunction(f) => {
+                f.body = Box::new(body.into());
+            }
+        }
+    }
+
+    pub fn ensure_body_block_stmt(&mut self) {
+        if let GenericFn::ArrowFunction(arrow_function) = self {
+            let body = arrow_function.body.deref();
+            if let BlockStmtOrExpr::Expr(expr) = body {
+                let return_arg = expr.clone().into_return_stmt();
+                let block_stmt = BlockStmt {
+                    span: DUMMY_SP,
+                    stmts: vec![return_arg.into()],
+                };
+
+                self.update_body(block_stmt);
+            }
+        }
+    }
+
+    pub fn return_type_matches_component(&self) -> bool {
+        let fn_body = self.get_body();
+
+        if let Some(fn_body) = fn_body {
+            match fn_body {
+                BlockStmtOrExpr::BlockStmt(ref block_stmt) => {
+                    let mut visitor: VisitReturnStatements = Default::default();
+                    block_stmt.visit_children_with(&mut visitor);
+                    visitor.state.does_return_type_match_component()
+                }
+                BlockStmtOrExpr::Expr(ref expr) => is_component_return_type(expr.deref()),
+            }
+        } else {
+            false
+        }
+    }
 }