From 2bc7d25b0e854f4f570f6d967cf03036821d2a41 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Thu, 19 Sep 2024 16:26:28 +0200
Subject: [PATCH] improve algorithm

A bit of a vague commit message, but this does a lot of things. I could
split it up, but not sure if it's worth it. Instead, let's talk about
it.

While working on keeping track of comment locations I was running into
some issues. Not the end of the world, but we could make things better.

Paired with Jordan on this to rework the algorithm. The idea is that we
now do multiple passes which is technically slower, but now we can work
on separate units of work.

- Step #1 is to prepare the at-rule. This means that rules with multiple
  selectors will be split in multiple nodes with the their own single
  selector.
- Step #2 is to collect all the classes we want to create an `@utility`
  for.
- Step #3 is to create a clone of the main `@layer utilities` for all
  the non-`@utility` leftover nodes (E.g.: rules with element and ID
  selectors).
- Step #4 is to create a clone of the main `@layer utilities` node for
  every single `@utility <name>` we want to create.
- Step #5 is to go over every clone, and eliminate everything that is
  not part of the `@utility` in question. So we can remove siblings
  (except for comments near it) and go up the chain.
- Step #6 is now to go over the initial `@layer utilities` clone we set
  aside, and remove everything that's not part of any of the clones.
- Step #7 is cleanup work, where empty nodes are removed, and rules with
  a selector of `&` are replaced by its children. This is done in a
  depth-first traversal instead of breadth first.

Co-authored-by: Jordan Pittman <jordan@cryptica.me>
---
 .../migrate-at-layer-utilities.test.ts        | 345 ++++++++++++--
 .../codemods/migrate-at-layer-utilities.ts    | 427 +++++++++++-------
 2 files changed, 569 insertions(+), 203 deletions(-)

diff --git a/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.test.ts b/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.test.ts
index 784c02207f8b..a493639f2ea3 100644
--- a/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.test.ts
+++ b/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.test.ts
@@ -1,6 +1,6 @@
 import dedent from 'dedent'
 import postcss from 'postcss'
-import { expect, it } from 'vitest'
+import { describe, expect, it } from 'vitest'
 import { migrateAtLayerUtilities } from './migrate-at-layer-utilities'
 
 const css = dedent
@@ -21,11 +21,138 @@ it('should migrate simple `@layer utilities` to `@utility`', async () => {
         }
       }
     `),
+  ).toMatchInlineSnapshot(`
+    "@utility foo {
+      color: red;
+    }"
+  `)
+})
+
+it('should split multiple selectors in separate utilities', async () => {
+  expect(
+    await migrate(css`
+      @layer utilities {
+        .foo,
+        .bar {
+          color: red;
+        }
+      }
+    `),
   ).toMatchInlineSnapshot(`
     "@utility foo {
       color: red;
     }
-    "
+
+    @utility bar {
+      color: red;
+    }"
+  `)
+})
+
+it('should merge `@utility` with the same name', async () => {
+  expect(
+    await migrate(css`
+      @layer utilities {
+        .foo {
+          color: red;
+        }
+      }
+
+      .bar {
+        color: blue;
+      }
+
+      @layer utilities {
+        .foo {
+          font-weight: bold;
+        }
+      }
+    `),
+  ).toMatchInlineSnapshot(`
+    "@utility foo {
+      color: red;
+      font-weight: bold;
+    }
+
+    .bar {
+      color: blue;
+    }"
+  `)
+})
+
+it('should leave non-class utilities alone', async () => {
+  expect(
+    await migrate(css`
+      @layer utilities {
+        /* 1. */
+        #before {
+          /* 1.1. */
+          color: red;
+          /* 1.2. */
+          .bar {
+            /* 1.2.1. */
+            font-weight: bold;
+          }
+        }
+
+        /* 2. */
+        .foo {
+          /* 2.1. */
+          color: red;
+          /* 2.2. */
+          .bar {
+            /* 2.2.1. */
+            font-weight: bold;
+          }
+        }
+
+        /* 3. */
+        #after {
+          /* 3.1. */
+          color: blue;
+          /* 3.2. */
+          .bar {
+            /* 3.2.1. */
+            font-weight: bold;
+          }
+        }
+      }
+    `),
+  ).toMatchInlineSnapshot(`
+    "@layer utilities {
+      /* 1. */
+      #before {
+        /* 1.1. */
+        color: red;
+        /* 1.2. */
+        .bar {
+          /* 1.2.1. */
+          font-weight: bold;
+        }
+      }
+
+      /* 3. */
+      #after {
+        /* 3.1. */
+        color: blue;
+        /* 3.2. */
+        .bar {
+          /* 3.2.1. */
+          font-weight: bold;
+        }
+      }
+    }
+
+    @utility foo {
+      /* 2. */
+      /* 2.1. */
+      color: red;
+      /* 2.2. */
+      .bar {
+        /* 2.2.1. */
+        font-weight: bold;
+      }
+    }"
   `)
 })
 
@@ -57,8 +184,7 @@ it('should migrate simple `@layer utilities` with nesting to `@utility`', async
       &:focus {
         color: green;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -79,10 +205,10 @@ it('should migrate multiple simple `@layer utilities` to `@utility`', async () =
     "@utility foo {
       color: red;
     }
+
     @utility bar {
       color: blue;
-    }
-    "
+    }"
   `)
 })
 
@@ -107,14 +233,14 @@ it('should not migrate Rules inside of Rules to a `@utility`', async () => {
     "@utility foo {
       color: red;
     }
+
     @utility bar {
       color: blue;
 
       .baz {
         color: green;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -134,8 +260,7 @@ it('should invert at-rules to make them migrate-able', async () => {
       @media (min-width: 640px) {
         color: red;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -165,13 +290,11 @@ it('should migrate at-rules with multiple utilities and invert them', async () =
       }
     }
 
-
     @utility bar {
       @media (min-width: 640px) {
         color: blue;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -208,11 +331,13 @@ it('should migrate deeply nested at-rules with multiple utilities and invert the
         color: red;
       }
     }
+
     @utility bar {
       @media (min-width: 640px) {
         color: blue;
       }
     }
+
     @utility baz {
       @media (min-width: 640px) {
         @media (min-width: 1024px) {
@@ -220,6 +345,7 @@ it('should migrate deeply nested at-rules with multiple utilities and invert the
         }
       }
     }
+
     @utility qux {
       @media (min-width: 640px) {
         @media (min-width: 1024px) {
@@ -228,8 +354,7 @@ it('should migrate deeply nested at-rules with multiple utilities and invert the
           }
         }
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -247,8 +372,7 @@ it('should migrate classes with pseudo elements', async () => {
       &::-webkit-scrollbar {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -266,8 +390,7 @@ it('should migrate classes with attribute selectors', async () => {
       &[data-checked=""] {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -285,8 +408,7 @@ it('should migrate classes with element selectors', async () => {
       & main {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -304,8 +426,7 @@ it('should migrate classes attached to an element selector', async () => {
       &main {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -323,8 +444,7 @@ it('should migrate classes with id selectors', async () => {
       &#main {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -343,12 +463,12 @@ it('should migrate classes with another attached class', async () => {
         display: none;
       }
     }
+
     @utility main {
       &.no-scrollbar {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -367,17 +487,18 @@ it('should migrate a selector with multiple classes to multiple @utility definit
         display: none;
       }
     }
+
     @utility bar {
       .foo &:hover .baz:focus {
         display: none;
       }
     }
+
     @utility baz {
       .foo .bar:hover &:focus {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -405,9 +526,7 @@ it('should merge `@utility` definitions with the same name', async () => {
         @apply ml-[-41px];
         content: counter(step);
       }
-    }
-
-    "
+    }"
   `)
 })
 
@@ -426,17 +545,18 @@ it('should not migrate nested classes inside a `:not(…)`', async () => {
         display: none;
       }
     }
+
     @utility bar {
       .foo &:not(.qux):has(.baz) {
         display: none;
       }
     }
+
     @utility baz {
       .foo .bar:not(.qux):has(&) {
         display: none;
       }
-    }
-    "
+    }"
   `)
 })
 
@@ -489,6 +609,7 @@ it('should migrate advanced combinations', async () => {
         }
       }
     }
+
     @utility bar {
       @media (width >= 100px) {
         @supports (display: none) {
@@ -496,12 +617,10 @@ it('should migrate advanced combinations', async () => {
             display: none;
           }
         }
-      }
-
-      @media (width >= 100px) {
         color: red;
       }
     }
+
     @utility baz {
       @media (width >= 100px) {
         @supports (display: none) {
@@ -510,9 +629,155 @@ it('should migrate advanced combinations', async () => {
           }
         }
       }
-    }
+    }"
+  `)
+})
 
+describe('comments', () => {
+  it('should preserve comment location for a simple utility', async () => {
+    expect(
+      await migrate(css`
+        /* Start of utilities: */
+        @layer utilities {
+          /* Utility #1 */
+          .foo {
+            /* Declarations: */
+            color: red;
+          }
+        }
+      `),
+    ).toMatchInlineSnapshot(`
+      "/* Start of utilities: */
+      @utility foo {
+        /* Utility #1 */
+        /* Declarations: */
+        color: red;
+      }"
+    `)
+  })
+
+  it('should copy comments when creating multiple utilities from a single selector', async () => {
+    expect(
+      await migrate(css`
+        /* Start of utilities: */
+        @layer utilities {
+          /* Foo & Bar */
+          .foo .bar {
+            /* Declarations: */
+            color: red;
+          }
+        }
+      `),
+    ).toMatchInlineSnapshot(`
+      "/* Start of utilities: */
+      @utility foo {
+        /* Foo & Bar */
+        & .bar {
+          /* Declarations: */
+          color: red;
+        }
+      }
+      @utility bar {
+        /* Foo & Bar */
+        .foo & {
+          /* Declarations: */
+          color: red;
+        }
+      }"
+    `)
+  })
+
+  it('should preserve comments for utilities wrapped in at-rules', async () => {
+    expect(
+      await migrate(css`
+        /* Start of utilities: */
+        @layer utilities {
+          /* Mobile only */
+          @media (width <= 640px) {
+            /* Utility #1 */
+            .foo {
+              /* Declarations: */
+              color: red;
+            }
+          }
+        }
+      `),
+    ).toMatchInlineSnapshot(`
+      "/* Start of utilities: */
+      @utility foo {
+        /* Mobile only */
+        @media (width <= 640px) {
+          /* Utility #1 */
+          /* Declarations: */
+          color: red;
+        }
+      }"
+    `)
+  })
+
+  it('should preserve comment locations as best as possible', async () => {
+    expect(
+      await migrate(css`
+        /* Above */
+        .before {
+          /* Inside */
+        }
+        /* After */
+
+        /* Tailwind Utilities: */
+        @layer utilities {
+          /* Chrome, Safari and Opera */
+          /* Second comment */
+          @media (min-width: 640px) {
+            /* Foobar */
+            .no-scrollbar::-webkit-scrollbar {
+              display: none;
+            }
+          }
 
-    "
-  `)
+          /* Firefox, IE and Edge */
+          /* Second comment */
+          .no-scrollbar {
+            -ms-overflow-style: none; /* IE and Edge */
+            scrollbar-width: none; /* Firefox */
+          }
+        }
+
+        /* Above */
+        .after {
+          /* Inside */
+        }
+        /* After */
+      `),
+    ).toMatchInlineSnapshot(`
+      "/* Above */
+      .before {
+        /* Inside */
+      }
+      /* After */
+
+      /* Tailwind Utilities: */
+      @utility no-scrollbar {
+        /* Chrome, Safari and Opera */
+        /* Second comment */
+        @media (min-width: 640px) {
+          /* Foobar */
+          &::-webkit-scrollbar {
+            display: none;
+          }
+        }
+
+        /* Firefox, IE and Edge */
+        /* Second comment */
+        -ms-overflow-style: none; /* IE and Edge */
+        scrollbar-width: none; /* Firefox */
+      }
+
+      /* Above */
+      .after {
+        /* Inside */
+      }
+      /* After */"
+    `)
+  })
 })
diff --git a/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.ts b/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.ts
index 7d3f9737e997..590196056acd 100644
--- a/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.ts
+++ b/packages/@tailwindcss-upgrade/src/codemods/migrate-at-layer-utilities.ts
@@ -1,6 +1,7 @@
-import { AtRule, Container, parse, Rule, type Plugin } from 'postcss'
+import { AtRule, parse, Rule, type ChildNode, type Comment, type Plugin } from 'postcss'
 import SelectorParser from 'postcss-selector-parser'
 import { format } from 'prettier'
+import { segment } from '../../../tailwindcss/src/utils/segment'
 
 enum WalkAction {
   // Continue walking the tree. Default behavior.
@@ -37,181 +38,260 @@ function walk<T>(rule: Walkable<T>, cb: (rule: T) => void | WalkAction): undefin
   return result
 }
 
+// Depth first walk reversal implementation.
+function walkDepth<T>(rule: Walkable<T>, cb: (rule: T) => void) {
+  rule?.each?.((node) => {
+    walkDepth(node as Walkable<T>, cb)
+    cb(node)
+  })
+}
+
 export function migrateAtLayerUtilities(): Plugin {
   async function migrate(atRule: AtRule) {
+    // Only migrate `@layer utilities` and `@layer components`.
     if (atRule.params !== 'utilities' && atRule.params !== 'components') return
 
-    // Upgrade every Rule in `@layer utilities` to an `@utility` at-rule.
+    // If the `@layer utilities` contains CSS that should not be turned into an
+    // `@utility` at-rule, then we have to keep it around (including the
+    // `@layer utilities` wrapper). To prevent this from being processed over
+    // and over again, we mark it as seen and bail early.
+    if (atRule.raws.seen) return
+
+    // Keep rules that should not be turned into utilities as is. This will
+    // include rules with element or ID selectors.
+    let defaultsAtRule = atRule.clone({ raws: { seen: true } })
+
+    // Clone each rule with multiple selectors into their own rule with a single
+    // selector.
     walk(atRule, (node) => {
-      if (!(node instanceof Rule)) return
-
-      // Fan out each utility into its own rule.
-      //
-      // E.g.:
-      // ```css
-      // .foo .bar:hover .baz {
-      //   color: red;
-      // }
-      // ```
-      //
-      // Becomes:
-      // ```css
-      // @utility foo {
-      //   & .bar:hover .baz {
-      //     color: red;
-      //   }
-      // }
-      //
-      // @utility bar {
-      //   .foo &:hover .baz {
-      //     color: red;
-      //   }
-      // }
-      //
-      // @utility baz {
-      //   .foo .bar:hover & {
-      //     color: red;
-      //   }
-      // }
-      // ```
-      let utilitySelectors: [name: string, selector: string][] = []
+      if (node.type !== 'rule') return
+
+      // Clone the node for each selector
+      let selectors = segment(node.selector, ',')
+      if (selectors.length > 1) {
+        let clonedNodes: Rule[] = []
+        for (let selector of selectors) {
+          let clone = node.clone({ selector })
+          clonedNodes.push(clone)
+        }
+        node.replaceWith(clonedNodes)
+      }
+
+      return WalkAction.Skip
+    })
+
+    // Track all the classes that we want to create an `@utility` for.
+    let classes = new Set<string>()
+
+    walk(atRule, (node) => {
+      if (node.type !== 'rule') return
+
+      // Find all the classes in the selector
       SelectorParser((selectors) => {
         selectors.each((selector) => {
-          walk(selector, (node) => {
+          walk(selector, (selectorNode) => {
             // Ignore everything in `:not(…)`
-            if (node.type === 'pseudo' && node.value === ':not') {
+            if (selectorNode.type === 'pseudo' && selectorNode.value === ':not') {
               return WalkAction.Skip
             }
 
-            // Replace the class with `&` and track the new selector
-            if (node.type === 'class') {
-              // Work on a clone of the selector, so we can safely manipulate
-              // it without affecting the original.
-              let clone = selector.clone()
-
-              // Find the node in the clone based on the position of the
-              // original node.
-              let target = clone.atPosition(node.source!.start!.line, node.source!.start!.column)
-
-              // Keep moving the target to the front until we hit the start or
-              // find a combinator. This is to prevent `.foo.bar` from becoming
-              // `.bar&`. Instead we want `&.bar`.
-              let parent = target.parent!
-              let idx = (target.parent?.index(target) ?? 0) - 1
-              while (idx >= 0 && parent.at(idx)?.type !== 'combinator') {
-                let current = parent.at(idx + 1)
-                let previous = parent.at(idx)
-                parent.at(idx + 1).replaceWith(previous)
-                parent.at(idx).replaceWith(current)
-
-                idx--
-              }
+            if (selectorNode.type === 'class') {
+              classes.add(selectorNode.value)
+            }
+          })
+        })
+      }).processSync(node.selector, { updateSelector: false })
 
-              // Replace the class with `&`
-              target.replaceWith(SelectorParser.nesting())
+      return WalkAction.Skip
+    })
+
+    // Remove all the nodes from the default `@layer utilities` that we know
+    // should be turned into `@utility` at-rules.
+    walk(defaultsAtRule, (node) => {
+      if (node.type !== 'rule') return
 
-              // Track the new selector
-              utilitySelectors.push([node.value, clone.toString()])
+      SelectorParser((selectors) => {
+        selectors.each((selector) => {
+          walk(selector, (selectorNode) => {
+            // Ignore everything in `:not(…)`
+            if (selectorNode.type === 'pseudo' && selectorNode.value === ':not') {
+              return WalkAction.Skip
+            }
+
+            // Remove the node if the class is in the list
+            if (selectorNode.type === 'class' && classes.has(selectorNode.value)) {
+              node.remove()
+              return WalkAction.Stop
             }
           })
+          node.selector = selector.toString()
         })
       }).processSync(node.selector, { updateSelector: false })
+    })
 
-      // Wrap the new rules in `@utility` at-rules
-      let newRules: AtRule[] = []
-      for (let [name, selector] of utilitySelectors) {
-        if (selector === '&') {
-          newRules.push(new AtRule({ name: 'utility', params: name, nodes: node.nodes }))
-        } else {
-          newRules.push(
-            new AtRule({
-              name: 'utility',
-              params: name,
-              nodes: [new Rule({ selector, nodes: node.nodes })],
-            }),
-          )
+    // Upgrade every Rule in `@layer utilities` to an `@utility` at-rule.
+    let clones: AtRule[] = [defaultsAtRule]
+    for (let cls of classes) {
+      let clone = atRule.clone()
+      clones.push(clone)
+
+      walk(clone, (node) => {
+        if (node.type !== 'rule') return
+
+        // Fan out each utility into its own rule.
+        //
+        // E.g.:
+        // ```css
+        // .foo .bar:hover .baz {
+        //   color: red;
+        // }
+        // ```
+        //
+        // Becomes:
+        // ```css
+        // @utility foo {
+        //   & .bar:hover .baz {
+        //     color: red;
+        //   }
+        // }
+        //
+        // @utility bar {
+        //   .foo &:hover .baz {
+        //     color: red;
+        //   }
+        // }
+        //
+        // @utility baz {
+        //   .foo .bar:hover & {
+        //     color: red;
+        //   }
+        // }
+        // ```
+        let containsClass = false
+        SelectorParser((selectors) => {
+          selectors.each((selector) => {
+            walk(selector, (selectorNode) => {
+              // Ignore everything in `:not(…)`
+              if (selectorNode.type === 'pseudo' && selectorNode.value === ':not') {
+                return WalkAction.Skip
+              }
+
+              // Replace the class with `&` and track the new selector
+              if (selectorNode.type === 'class' && selectorNode.value === cls) {
+                containsClass = true
+
+                // Find the node in the clone based on the position of the
+                // original node.
+                let target = selector.atPosition(
+                  selectorNode.source!.start!.line,
+                  selectorNode.source!.start!.column,
+                )
+
+                // Keep moving the target to the front until we hit the start or
+                // find a combinator. This is to prevent `.foo.bar` from
+                // becoming `.bar&`. Instead we want `&.bar`.
+                let parent = target.parent!
+                let idx = (target.parent?.index(target) ?? 0) - 1
+                while (idx >= 0 && parent.at(idx)?.type !== 'combinator') {
+                  let current = parent.at(idx + 1)
+                  let previous = parent.at(idx)
+                  parent.at(idx + 1).replaceWith(previous)
+                  parent.at(idx).replaceWith(current)
+
+                  idx--
+                }
+
+                // Replace the class with `&`
+                target.replaceWith(SelectorParser.nesting())
+              }
+            })
+          })
+
+          // Update the selector
+          node.selector = selectors.toString()
+        }).processSync(node.selector)
+
+        // Cleanup all the nodes that should not be part of the `@utility` rule.
+        if (!containsClass) {
+          let toRemove: (Comment | Rule)[] = [node]
+          let idx = node.parent?.index(node) ?? null
+          if (idx !== null) {
+            for (let i = idx - 1; i >= 0; i--) {
+              if (node.parent?.nodes.at(i)?.type === 'rule') {
+                break
+              }
+              if (node.parent?.nodes.at(i)?.type === 'comment') {
+                toRemove.push(node.parent?.nodes.at(i) as Comment)
+              }
+            }
+          }
+          for (let node of toRemove) {
+            node.remove()
+          }
         }
-      }
 
-      node.replaceWith(newRules)
+        return WalkAction.Skip
+      })
 
-      return WalkAction.Skip
-    })
+      // Migrate the `@layer utilities` to `@utility <name>`
+      clone.name = 'utility'
+      clone.params = cls
 
-    // Hoist all `@utility` at-rules to the top. It could be that these were
-    // (deeply) nested in other at-rules like `@media`.
-    //
-    // ```css
-    // @media (min-width: 640px) {
-    //   @utility foo {
-    //     color: red;
-    //   }
-    // }
-    // ```
-    //
-    // Becomes:
-    // ```css
-    // @utility foo {
-    //   @media (min-width: 640px) {
-    //     color: red;
-    //   }
-    // }
-    // ```
-    let trees: AtRule[] = []
-    walk(atRule, (node) => {
-      if (node.type !== 'atrule' || node.name !== 'utility') return
-
-      // Track the parents of the node, so we can reconstruct the tree later.
-      let parents = []
-      let parent: Container | null = node.parent ?? null
-
-      while (
-        parent &&
-        !(
-          parent instanceof AtRule &&
-          parent.name === 'layer' &&
-          (parent.params === 'utilities' || parent.params === 'components')
-        )
-      ) {
-        parents.push(parent.clone({ nodes: [] }))
+      // Mark the node as pretty so that it gets formatted by Prettier later.
+      clone.raws.tailwind_pretty = true
+      clone.raws.before += '\n\n'
+    }
 
-        // Move up the tree
-        // @ts-expect-error
-        parent = parent.parent ?? null
-      }
+    // Cleanup
+    for (let idx = clones.length - 1; idx >= 0; idx--) {
+      let clone = clones[idx]
 
-      // Work on a clone of the node, so we can safely manipulate it.
-      let nodeClone = node.clone()
+      walkDepth(clone, (node) => {
+        // Remove comments from the main `@layer utilities` we want to keep,
+        // that are part of any of the other clones.
+        if (clone === defaultsAtRule) {
+          if (node.type === 'comment') {
+            let found = false
+            for (let other of clones) {
+              if (other === defaultsAtRule) continue
 
-      // Reconstruct the tree
-      for (let parent of parents) {
-        let children = nodeClone.nodes ?? []
-        // Inject _my_ children into the parent
-        parent.append(children)
+              walk(other, (child) => {
+                if (
+                  child.type === 'comment' &&
+                  child.source?.start?.offset === node.source?.start?.offset
+                ) {
+                  node.remove()
+                  found = true
+                  return WalkAction.Stop
+                }
+              })
 
-        // Inject the parent into the node
-        nodeClone.removeAll()
-        nodeClone.append(parent)
-      }
+              if (found) {
+                return WalkAction.Skip
+              }
+            }
+          }
+        }
 
-      // Collect the newly new tree (which is the manipulated node)
-      trees.push(nodeClone)
+        // Remove empty rules
+        if ((node.type === 'rule' || node.type === 'atrule') && node.nodes?.length === 0) {
+          node.remove()
+        }
 
-      // Already handled the `@utility` at-rule, no need to go deeper.
-      return WalkAction.Skip
-    })
+        // Replace `&` selectors with its children
+        else if (node.type === 'rule' && node.selector === '&') {
+          node.replaceWith(node.nodes)
+        }
+      })
 
-    // Replace `@layer utilities` with the newly constructed trees.
-    //
-    // Prettier is used to generate cleaner output, but it's only used on the
-    // nodes that we migrated.
-    atRule.replaceWith(
-      await Promise.all(
-        trees.map(async (tree) => {
-          return parse(await format(tree.toString(), { parser: 'css', semi: true }))
-        }),
-      ),
-    )
+      // Remove empty clones entirely
+      if (clone.nodes?.length === 0) {
+        clones.splice(idx, 1)
+      }
+    }
+
+    // Finally, replace the original `@layer utilities` with the new rules.
+    atRule.replaceWith(clones)
   }
 
   return {
@@ -219,23 +299,44 @@ export function migrateAtLayerUtilities(): Plugin {
     AtRule: {
       layer: migrate,
     },
-    OnceExit: (root) => {
-      // Merge `@utility` definitions with the same name.
-      let nameToAtRule = new Map<string, AtRule>()
-
-      root.walkAtRules('utility', (node) => {
-        let existing = nameToAtRule.get(node.params)
-        if (existing) {
-          // Add a newline between each `@utility` at-rule
-          if (node.first) {
-            node.first.raws.before = `\n${node.first?.raws.before ?? ''}`
+    OnceExit: async (root) => {
+      // Prettier is used to generate cleaner output, but it's only used on the
+      // nodes that were marked as `pretty` during the migration.
+      {
+        // Find the nodes to format
+        let nodesToFormat: ChildNode[] = []
+        walk(root, (child) => {
+          if (child.raws.tailwind_pretty) {
+            nodesToFormat.push(child)
+            return WalkAction.Skip
           }
-          existing.append(node.nodes ?? [])
-          node.remove()
-        } else {
-          nameToAtRule.set(node.params, node)
-        }
-      })
+        })
+
+        // Format the nodes
+        await Promise.all(
+          nodesToFormat.map(async (node) => {
+            node.replaceWith(parse(await format(node.toString(), { parser: 'css', semi: true })))
+          }),
+        )
+      }
+
+      // Merge `@utility <name>` with the same name into a single rule. This can
+      // happen when the same classes is used in multiple `@layer utilities`
+      // blocks.
+      {
+        let utilities = new Map<string, AtRule>()
+        walk(root, (child) => {
+          if (child.type === 'atrule' && child.name === 'utility') {
+            let existing = utilities.get(child.params)
+            if (existing) {
+              existing.append(child.nodes!)
+              child.remove()
+            } else {
+              utilities.set(child.params, child)
+            }
+          }
+        })
+      }
     },
   }
 }