Attempting Large Code Refactor using LLMs

Gajus Kuizinas

Prompt Engineer

AI Chatbot Developer

ChatGPT

Let's set the context. We have a React component called Menu. It has a few sub-components. We are using the "Compound Components Pattern" to export these components. The Menu component is the parent component. Then we have Menu.Item, Menu.ItemSeparator, Menu.Panel, and Menu.Reference components. This is how the export looks like:

Menu.Reference = MenuReference;
Menu.Panel = MenuPanel;
Menu.Item = MenuItem;
Menu.ItemSeparator = MenuItemSeparator;

export { Menu };

The Menu component is used in a few places. For example:

<Menu>
  <Menu.Reference>
    <Button>Open Menu</Button>
  </Menu.Reference>
  <Menu.Panel>
    <Menu.Item>Item 1</Menu.Item>
    <Menu.Item>Item 2</Menu.Item>
    <Menu.ItemSeparator />
    <Menu.Item>Item 3</Menu.Item>
  </Menu.Panel>
</Menu>

However, we want to change the export to be like this:

export const Menu = {
  Item: MenuItem,
  ItemSeparator: MenuItemSeparator,
  Menu: TopMenu,
  Panel: MenuPanel,
  Reference: MenuReference,
};

We are making this change as part of making our entire codebase easier to tree-shake. The pattern of assigning components as properties to a component function is being interpreted as a side-effect by our bundler.

This requires updating all the places where Menu is used. For example:

<Menu>
  <Menu.Reference>
    <Button>Open Menu</Button>
  </Menu.Reference>
  <Menu.Panel>
    <Menu.Item>Item 1</Menu.Item>
    <Menu.Item>Item 2</Menu.Item>
    <Menu.ItemSeparator />
    <Menu.Item>Item 3</Menu.Item>
  </Menu.Panel>
</Menu>

becomes:

<Menu.Menu>
  <Menu.Reference>
    <Button>Open Menu</Button>
  </Menu.Reference>
  <Menu.Panel>
    <Menu.Item>Item 1</Menu.Item>
    <Menu.Item>Item 2</Menu.Item>
    <Menu.ItemSeparator />
    <Menu.Item>Item 3</Menu.Item>
  </Menu.Panel>
</Menu.Menu>

In addition, we need to consider other cases, e.g.:

const MyMenu = styled(Menu)``;

<Menu
  isOpen={isActionsOpen}
  offset={[20, 0]}
  placement="bottom-start"
  toggleMenu={toggleMenu}
>

import { Menu as StandardMenu } from '@/components/Menu';

// Not our Menu component
const Menu = () => {};

<Menu></Menu>

In short, a simple search-and-replace would not work here.

The <Menu> component is referenced 325 times in 83 files. This is a large code refactor. It is tedious and error prone. Can we use LLMs to help us with this?

The Approach

The first thing I did was to update the Menu component to export the components using the new pattern. Then I ran tsc to produce a list of files with errors. Then I narrowed down to the files that mention Menu. Then for every file, I applied the ChatGPT generated suggestion. Let's go through this step by step.

Step 1: Update the `Menu` component

This is the diff of the changes I made to the Menu component:

-Menu.Reference = MenuReference;
-Menu.Panel = MenuPanel;
-Menu.Item = MenuItem;
-Menu.ItemSeparator = MenuItemSeparator;
-
-export { Menu };
+export const Menu = {
+  Item: MenuItem,
+  ItemSeparator: MenuItemSeparator,
+  Menu: TopMenu,
+  Panel: MenuPanel,
+  Reference: MenuReference,
+};

Step 2: Run `tsc` to produce a list of files with errors

This gave me a list of files with errors:

src/components/Messaging/components/Channel/components/ChannelHeader/components/ChannelHeaderActions.tsx:197:6 - error TS2604: JSX element type 'Menu' does not have any construct or call signatures.

197     <Menu
         ~~~~

src/components/Messaging/components/Channel/components/ChannelHeader/components/ChannelHeaderActions.tsx:197:6 - error TS2786: 'Menu' cannot be used as a JSX component.
  Its type '{ Item: <C extends React.ElementType = "div">({ as, children, disabled, hideDisabledLockIcon, iconFillColor, onClick, icon: Icon, elementRef, destructive, iconSize, closeMenuOnClick, withDarkUnlockedIcon, ...props }: MenuItemProps<C>) => JSX.Element; ItemSeparator: StyledComponent<"div", DefaultTheme, {}, never>; Me...' is not a valid JSX element type.

197     <Menu
         ~~~~

src/components/Messaging/components/Channel/components/Message/Message.tsx:139:34 - error TS2339: Property 'id' does not exist on type '{ readonly " $fragmentSpreads": FragmentRefs<"MessageAvatar_memberAvatars">; }'.

139     (participant) => participant.id === message?.user?.userProfileId,
                                     ~~

src/components/Messaging/components/Sidebar/components/ChannelListControls/components/ConversationJobPostings.tsx:170:8 - error TS2604: JSX element type 'Menu' does not have any construct or call signatures.

170       <Menu
           ~~~~

Step 3: Narrow down to the files that mention `Menu`

I wrote a function to parse the tsc output:

const parseTscErrorOutput = (output: string) => {
  return [...new Set(output.trim().split('\n')
    .filter((line) => line.startsWith('src'))
    .filter((line) => line.includes('Menu'))
    .map((line) => {
      return line.split(':')[0];
    })
    .sort((a, b) => {
      return a.localeCompare(b);
    }))];
};

The idea is to get a list of files that mention Menu.

Step 4: Apply ChatGPT generated suggestion to each file

I wrote a function to apply the ChatGPT generated suggestion to each file:

for (const filePath of affectedFiles) {
  console.log(`Processing ${filePath}`);

  const inputCode = await readFile(filePath, 'utf8');

  const chatCompletion = await openai.createChatCompletion({
    messages: [{ content: createPrompt(inputCode), role: 'user' }],
    model: 'gpt-4',
  });

  const response = String(chatCompletion.data.choices[0].message?.content);

  await writeFile(filePath, normalizeResponse(response));

  try {
    await formatFile(filePath);
  } catch {
    console.warn('[warn] could not apply eslint fixes');
  }

  console.log(`Updated ${filePath}`);
}

The createPrompt function is used to create the prompt for ChatGPT. It looks like this:

const createPrompt = (code: string) => {
  return `
The export of the Menu component has changed from:

\`\`\`
Menu.Reference = MenuReference;
Menu.Panel = MenuPanel;
Menu.Item = MenuItem;
Menu.ItemSeparator = MenuItemSeparator;

export { Menu };
\`\`\`

to:

\`\`\`
export const Menu = {
  Item: MenuItem,
  ItemSeparator: MenuItemSeparator,
  Menu: TopMenu,
  Panel: MenuPanel,
  Reference: MenuReference,
};
\`\`\`

This means that the React files that consume the Menu component need to be updated to reference <Menu.Menu> instead of <Menu>.

Example (1):

\`\`\`
<Menu>
  ...
</Menu>
\`\`\`

becomes:

\`\`\`
<Menu.Menu>
  ...
</Menu.Menu>
\`\`\`


Example (2):

\`\`\`
const MyMenu = styled(Menu);
\`\`\`

becomes:

\`\`\`
const MyMenu = styled(Menu.Menu);
\`\`\`

Update the following source code to use the new Menu component export:

\`\`\`
${code}
\`\`\`

Instructions:

* Do not truncate the response
* Only refactor <Menu> to <Menu.Menu>
* Only output code
* Menu.Item, Menu.ItemSeparator, Menu.Panel, Menu.Reference must not be refactored
* Do not refactor already correct references, e.g. <Menu.Menu> must not be refactored to <Menu.Menu.Menu>
* Do not wrap the output
* Do not refactor references if they are locally redefined
* Do not add comments
`;
};

The normalizeResponse function is used to strip the ``` from the response. It looks like this:

const normalizeResponse = (response: string): string => {
  const lines = response.trim().split('\n');

  const blockStart = lines.findIndex((line) => line.startsWith('```'));

  if (blockStart === -1) {
    return response;
  }

  const blockEnd = lines
    .slice(blockStart + 1)
    .findIndex((line) => line.startsWith('```'));

  if (blockEnd === -1) {
    throw new Error('Could not find end of code block');
  }

  const code = lines
    .slice(blockStart + 1, blockStart + blockEnd + 1)
    .join('\n');

  return code;
};

(More about why it is even necessary in the Gotchas section.)

The formatFile function is used to format the file using eslint. It is intentionally called after the writeFile function. This is because the LLM generated code sometimes is not syntactically correct (see Gotchas). However, the suggestions it generates might still be valuable. So we first write the file, then we format it. If it fails, we just ignore it.

Step 5: Manually verify the changes

The next step is to manually verify the changes. I did this by using a GUI git client (Tower).

I will summarize the results later, but in short, I had to unstage a few files and manually fix them.

Step 6: Run `tsc` again

At the end, I ran tsc again to see if there are any errors. There were none. We were done.

Results

The total time it took to run the script was 1 hour and 15 minutes.

The total time it took to manually verify the changes was ~5 minutes.

Out of the 83 files, 75 were updated successfully, and 8 had issues.

I had to make changes to 8 files manually. The changes were mostly minor. For example, in one file, the LLM generated code was:

-export const MenuItem = styled(Menu.Item)`
+export const MenuItem = styled(Menu.Menu.Item)`

However, thanks to the GUI git client, I was able to discard this and similar mistakes with a click of a button. Overall, fixing the issues didn't take more than a few minutes.

Overall, it was very satisfying to just sit back and watch the changes appear one by one in the GUI git client.

Watching changes come into the codebase

Gotchas

Let's talk about the gotchas.

Gotcha 1: could not force LLM to output only code consistently

The first gotcha is that I couldn't stop ChatGPT from wrapping the output in ```. I tried adding Do not wrap the output instruction and similar, but it didn't help. Sometimes it would even prepend text, e.g.

import { Menu, MenuItem } from '@/components/Menu';
Here is the updated code:

```
import { Menu as MenuExport, MenuItem } from '@/components/Menu';

So I had to write normalizeResponse function to strip it.

A better approach here would be to use functions and pass the code as parameter. That's what I will try next, as I iterate on this approach.

Gotcha 2: LLM generated output is truncated

The other gotcha is that sometimes ChatGPT will replace parts of the response with:

{/* ... rest of the code ... */}

// Remainder of the code is the same until the return statement

For what it is worth, thanks to GUI Git again, it was simple to correct these files.

Later I added Do not truncate the response instruction and it seems to have dramatically reduced instances like this.

Gotcha 3: could not stop LLM from adding comments

The other gotcha is that sometimes ChatGPT will add comments to the code. For example:

-    <Menu
+    <Menu.Menu // new export of Menu component
      isOpen={isActionsOpen}
      offset={[20, 0]}
      placement="bottom-start"
      toggleMenu={toggleMenu}

Adding Do not add comments instructions reduced this, but didn't eliminate it.

Also, adding Do not add comments instruction sometimes caused ChatGPT to strip existing comments, e.g.

import { type FormEvent, useCallback, useReducer, useState } from 'react';
-// disable useContraMutation eslint error
-// eslint-disable-next-line no-restricted-imports
import { graphql, useMutation } from 'react-relay';

Overall, this isn't a big deal since (once again) it can be easily excluded from suggestions using a GUI Git client.

Other notes

Despite the cost different, gpt-4 is worth it. Earlier attempts using gpt-3 produced far less impressive outputs.

Including examples of how to handle edge cases reduces overall error rate.

In future iterations, I will try to produce git patches rather than the entire output.

Conclusion

It's important to note that even though some manual work was required, LLMs helped to almost entirely eliminate laborious and repetitive job I had to as an engineer. Overall, I probably spent more time setting up everything and experimenting with prompts than it would have taken me to do this refactor. However, this is just one of many refactors that we need to do. I am hoping that the next refactor will be faster. I will keep you posted.

Like this project

Posted Jun 29, 2023

I'v successfully utilized a LLM to automate and streamline the refactoring process, enhancing its efficiency and reliability.

Likes

Views

Clients