Implementing a Simplified React Forget Compiler

This challenge asks you to implement a core piece of the React Forget compiler. React Forget optimizes React components by memoizing expensive computations and preventing unnecessary re-renders. Your task is to build a simplified version of this compiler that can identify and transform specific patterns in React component code written in TypeScript.

Problem Description

You need to create a TypeScript compiler plugin that analyzes React function components and applies optimizations. Specifically, you will focus on memoizing useMemo calls within function components. The compiler should identify useMemo calls that are directly returning a value and transform them into a more optimized form if certain conditions are met.

Key Requirements:

Identify useMemo calls: The compiler should be able to parse TypeScript code and locate useMemo calls within function component bodies.
Analyze useMemo dependencies: For each useMemo call, identify its dependencies (the second argument to useMemo).
Detect simple return values: The compiler should recognize useMemo calls where the first argument is a simple arrow function that directly returns a value (e.g., () => someValue).
Transform useMemo: If a useMemo call meets the criteria (simple return value and dependencies are stable), transform it into a more direct assignment or a simplified memoization pattern. For this challenge, we'll simplify the transformation: if the dependencies are simple (e.g., literals, variables without complex logic), we will aim to remove the useMemo wrapper and directly use the value, assuming the React runtime would handle memoization internally for such simple cases. If dependencies are complex, we'll leave useMemo as is.
Preserve component structure: The transformation should not alter the overall structure or functionality of the React component.
Handle basic TypeScript AST: You'll be working with the Abstract Syntax Tree (AST) of TypeScript code. Familiarity with typescript's AST node types will be crucial.

Expected Behavior:

The compiler should take a TypeScript file as input and produce a transformed TypeScript file as output.

Example Transformation:

Before:

import React, { useMemo } from 'react';

interface Props {
  a: number;
  b: string;
}

function MyComponent({ a, b }: Props) {
  const memoizedValue = useMemo(() => {
    console.log('Calculating...');
    return a * 2 + b.length;
  }, [a, b]);

  return <div>{memoizedValue}</div>;
}

After (simplified transformation for this challenge):

import React from 'react'; // useMemo might be implicitly handled by React runtime

interface Props {
  a: number;
  b: string;
}

function MyComponent({ a, b }: Props) {
  // Simplified: Assuming React runtime or internal logic handles memoization for simple dependencies
  const memoizedValue = a * 2 + b.length; // Direct calculation, no explicit useMemo

  return <div>{memoizedValue}</div>;
}

Edge Cases:

useMemo calls with complex logic in the factory function (not just a direct return).
useMemo calls with complex dependency arrays.
useMemo calls within other hooks or functions, not directly in the component body.
useMemo calls that don't return a value directly (e.g., contain multiple statements before return).
useMemo calls with no dependencies.

For this simplified challenge, we will focus on the common case: useMemo within a component function, with a simple arrow function returning a value, and a dependency array that we can reasonably assess as "simple".

Examples

Example 1:

Input:
import React, { useMemo } from 'react';

function MyComponent(props: { count: number }) {
  const derivedData = useMemo(() => {
    return props.count * 10;
  }, [props.count]);

  return <div>{derivedData}</div>;
}

Output:
import React from 'react';

function MyComponent(props: { count: number }) {
  const derivedData = props.count * 10;

  return <div>{derivedData}</div>;
}

Explanation: The useMemo hook is used to calculate props.count * 10. The dependency props.count is a simple variable. The compiler identifies this pattern and transforms it by removing the useMemo call and directly assigning the calculated value, assuming the React runtime can manage memoization internally for such simple computations.

Example 2:

Input:
import React, { useMemo } from 'react';

function AnotherComponent(props: { name: string, age: number }) {
  const formattedString = useMemo(() => {
    if (props.age > 18) {
      return `Adult: ${props.name}`;
    } else {
      return `Minor: ${props.name}`;
    }
  }, [props.name, props.age]);

  return <p>{formattedString}</p>;
}

Output:
import React, { useMemo } from 'react';

function AnotherComponent(props: { name: string, age: number }) {
  const formattedString = useMemo(() => {
    if (props.age > 18) {
      return `Adult: ${props.name}`;
    } else {
      return `Minor: ${props.name}`;
    }
  }, [props.name, props.age]);

  return <p>{formattedString}</p>;
}

Explanation: This useMemo call contains conditional logic within its factory function. Because it's not a single, direct return statement, and the logic is more complex than a simple expression, the compiler deems it not suitable for the simplified transformation and leaves it unchanged.

Example 3:

Input:
import React, { useMemo } from 'react';

function ComponentWithObject(props: { data: { id: number } }) {
  const computedObject = useMemo(() => {
    return { ...props.data, timestamp: Date.now() };
  }, [props.data]);

  return <div>{JSON.stringify(computedObject)}</div>;
}

Output:
import React, { useMemo } from 'react';

function ComponentWithObject(props: { data: { id: number } }) {
  const computedObject = useMemo(() => {
    return { ...props.data, timestamp: Date.now() };
  }, [props.data]);

  return <div>{JSON.stringify(computedObject)}</div>;
}

Explanation: The factory function directly returns an object literal. However, this object literal includes a call to Date.now(), which is a function call and can be considered side-effect-prone or non-deterministic in a compilation context. For this simplified challenge, we will treat such cases (function calls within the return value that aren't simple variable access or arithmetic) as too complex for automatic removal and leave the useMemo intact. A more advanced compiler might analyze Date.now()'s stability.

Constraints

The input code will be valid TypeScript code.
You should only consider useMemo calls declared within the top-level body of a React function component.
The transformation should be applied only to VariableDeclaration nodes that are assigned the result of a CallExpression where the callee is useMemo.
The useMemo factory function (the first argument) must be an ArrowFunction that contains a single ReturnStatement whose expression is not a BlockStatement (i.e., it's a direct expression return).
The dependency array (second argument) should be an ArrayLiteralExpression. For simplicity, we will consider dependencies to be "simple" if they are Identifiers or PropertyAccessExpressions that resolve to stable values (e.g., props.someValue, state.count). We will not attempt to deeply analyze the stability of complex expressions or object properties. For this challenge, we will transform if all dependencies are simple identifiers or property accesses.
The compiler should output valid TypeScript code, preserving original formatting as much as possible (though exact formatting isn't critical).
Performance of the compiler itself is not a primary concern for this challenge, but the generated code should be efficient.

Notes

You will need to use the typescript package to parse the code into an AST and then traverse and transform it.
Consider using the ts-morph library for easier AST manipulation if you find the native TypeScript API verbose.
Think about how to identify a "React function component." A simple heuristic is a function declaration that returns JSX.
The core of the problem lies in accurately identifying the AST nodes corresponding to useMemo and its arguments, and then programmatically constructing the new AST nodes for the transformation.
This is a simplified implementation. Real-world compilers handle many more edge cases, optimizations, and complex AST structures. Focus on the core useMemo transformation for direct returns with simple dependencies.
You might need helper functions to determine if a node represents a "simple" dependency.