feat: Support transforming hooks and arrow functions

ROADMAP.md

@@ -8,13 +8,17 @@ Type: `Feature` - `Bug` - `Enhancement` - `Refactor` - `Unknown`
 - [x] `Enhancement` Memoize any non-primitive value used in JSX
 - [x] `Enhancement` Support memoization of object destructuring
 - [x] `Enhancement` Unwrap JSX elements into memoizable variables
-- [ ] `Enhancement` Support side-effect for dependencies
-- [ ] `Feature` Support transforming hooks
-- [ ] `Enhancement` Support memoization of callbacks declared using function declarations
+- [x] `Enhancement` Support components with multiple returns
+- [x] `Enhancement` Support side-effect for dependencies
+- [x] `Feature` Support transforming hooks
+- [ ] `Enhancement` Support memoization of callbacks declared using function declarations after
+
+### Low Priority
+
 - [ ] `Enhancement` Skip memoization of useState setter functions
 - [ ] `Enhancement` When unwrapping array and object patterns, optimize the unwrapping by limiting it to component variables
 - [ ] `Enhancement` Support React.createElement calls
-- [ ] `Enhancement` Hot module reloading improvement utilizing a checksum to invalidate cache
+- [x] ~`Enhancement` Hot module reloading improvement utilizing a checksum to invalidate cache~ Not needed
 - [ ] `Feature` Memoize array map items
 
 ## `Feature` ESLint Plugin
@@ -30,5 +34,4 @@ The following are unknowns that need to be researched and tested.
 
 - [ ] Assignment expressions in JSX `<div>{(a = 1)} {a}</div>` - is the order of execution guaranteed?
 - [ ] Source maps - is it possible to generate source maps for the transformed code?
-- [ ] Hot reloading - how does hot reloading work with the transformed code?
-- [ ] Memoization of values declared after the first return statement
+- [x] ~Hot reloading - how does hot reloading work with the transformed code?~ It works!

packages/compiler/src/classes/tests/Component-fixtures.test.ts

@@ -15,7 +15,7 @@ const parseCodeAndRun = (fixtureName: string) => {
 
 describe("Component fixtures", () => {
   describe("applyModification", () => {
-    it.each(Array.from({ length: 12 }, (_, i) => `fixture_${i + 1}`))(
+    it.each(Array.from({ length: 13 }, (_, i) => `fixture_${i + 1}`))(
       "%s",
       (fixtureName) => {
         const [, program] = parseCodeAndRun(fixtureName);

packages/compiler/src/classes/tests/__snapshots__/Component-fixtures.test.ts.snap

@@ -1299,3 +1299,9 @@ exports[`Component fixtures applyModification fixture_12 1`] = `
   return null;
 }"
 `;
+
+exports[`Component fixtures applyModification fixture_13 1`] = `
+"const MyComponent = () => {
+  return <div />;
+};"
+`;

packages/compiler/src/fixtures/fixture_13.js

@@ -0,0 +1 @@
+const MyComponent = () => <div />;

packages/compiler/src/utils/expand-arrow-function-to-block-statement.ts

@@ -0,0 +1,19 @@
+import * as t from "@babel/types";
+
+export function expandArrowFunctionToBlockStatement(
+  path: babel.NodePath<t.Function>
+) {
+  if (!path.isArrowFunctionExpression()) {
+    return;
+  }
+
+  const pathBody = path.get("body");
+
+  if (pathBody.isBlockStatement()) {
+    return;
+  }
+
+  pathBody.replaceWith(
+    t.blockStatement([t.returnStatement(pathBody.node as t.Expression)])
+  );
+}

packages/compiler/src/utils/find-component.ts

@@ -1,6 +1,9 @@
 import type * as babel from "@babel/core";
+import * as t from "@babel/types";
 import { getReturnsOfFunction } from "./get-returns-of-function";
 import { Component } from "../classes/Component";
+import { doesMatchHookName } from "./is-hook-call";
+import { expandArrowFunctionToBlockStatement } from "./expand-arrow-function-to-block-statement";
 
 function doesIdMatchComponentName(name: string) {
   return /^_?[A-Z]/.test(name);
@@ -18,7 +21,7 @@ export function findComponents(program: babel.NodePath<babel.types.Program>) {
       }
 
       if (
-        path.isFunctionExpression() &&
+        (path.isFunctionExpression() || path.isArrowFunctionExpression()) &&
         path.parentPath.isVariableDeclarator()
       ) {
         const varIdPath = path.parentPath.get("id");
@@ -27,22 +30,46 @@ export function findComponents(program: babel.NodePath<babel.types.Program>) {
         }
       }
 
-      if (!id || !doesIdMatchComponentName(id.name)) {
+      const nameMatchesComponentName = id && doesIdMatchComponentName(id.name);
+      const nameMatchesHook = id && doesMatchHookName(id.name);
+
+      if (!nameMatchesComponentName && !nameMatchesHook) {
         return;
       }
 
+      expandArrowFunctionToBlockStatement(path);
+
       const returns = getReturnsOfFunction(path);
 
       if (returns.length === 0) {
         return;
       }
 
-      const allReturnsMatch = returns.every((ret) => {
-        return (
-          ret.get("argument").isNullLiteral() ||
-          ret.get("argument").isJSXElement() ||
-          ret.get("argument").isJSXFragment()
+      if (nameMatchesHook) {
+        components.push(
+          new Component(path as babel.NodePath<babel.types.Function>)
         );
+      }
+
+      const allReturnsMatch = returns.every((ret) => {
+        const argument = ret.get("argument");
+        if (isComponentReturnType(argument.node)) {
+          return true;
+        }
+
+        if (argument.isIdentifier()) {
+          const binding = argument.scope.getBinding(argument.node.name);
+          const variableDeclarator = binding?.path.find((bindingPath) =>
+            bindingPath.isVariableDeclarator()
+          ) as babel.NodePath<babel.types.VariableDeclarator> | undefined;
+
+          if (variableDeclarator) {
+            const init = variableDeclarator.node;
+            if (init && isComponentReturnType(init)) {
+              return true;
+            }
+          }
+        }
       });
 
       if (allReturnsMatch) {
@@ -55,3 +82,7 @@ export function findComponents(program: babel.NodePath<babel.types.Program>) {
 
   return components;
 }
+
+function isComponentReturnType(node: babel.types.Node | null | undefined) {
+  return t.isJSXElement(node) || t.isJSXFragment(node) || t.isNullLiteral(node);
+}