Track utilities per-candidate in @apply

packages/tailwindcss/src/apply.ts

@@ -2,6 +2,7 @@ import { Features } from '.'
 import { rule, toCss, walk, WalkAction, type AstNode } from './ast'
 import { compileCandidates } from './compile'
 import type { DesignSystem } from './design-system'
+import type { SourceLocation } from './source-maps/source'
 import { DefaultMap } from './utils/default-map'
 
 export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
@@ -155,12 +156,20 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
       let node = parent.nodes[i]
       if (node.kind !== 'at-rule' || node.name !== '@apply') continue
 
-      let candidates = node.params.split(/\s+/g)
+      let parts = node.params.split(/(\s+)/g)
+      let candidateOffsets: Record<string, number> = {}
+
+      let offset = 0
+      for (let [idx, part] of parts.entries()) {
+        if (idx % 2 === 0) candidateOffsets[part] = offset
+        offset += part.length
+      }
 
       // Replace the `@apply` rule with the actual utility classes
       {
         // Parse the candidates to an AST that we can replace the `@apply` rule
         // with.
+        let candidates = Object.keys(candidateOffsets)
         let compiled = compileCandidates(candidates, designSystem, {
           onInvalidCandidate: (candidate) => {
             throw new Error(`Cannot apply unknown utility class: ${candidate}`)
@@ -168,12 +177,36 @@ export function substituteAtApply(ast: AstNode[], designSystem: DesignSystem) {
         })
 
         let src = node.src
-        let candidateAst = compiled.astNodes
 
-        walk(candidateAst, (node) => {
-          node.src = src
+        let details = compiled.astNodes.map((node) => {
+          let candidate = compiled.nodeSorting.get(node)?.candidate
+          let candidateOffset = candidate ? candidateOffsets[candidate] : undefined
+
+          node = structuredClone(node)
+
+          if (!src || !candidate || candidateOffset === undefined) {
+            return { node, src }
+          }
+
+          let candidateSrc: SourceLocation = [src[0], src[1], src[2]]
+
+          candidateSrc[1] += 7 + candidateOffset
+          candidateSrc[2] = candidateSrc[1] + candidate.length
+
+          return { node: structuredClone(node), src: candidateSrc }
         })
 
+        for (let { node, src } of details) {
+          // While the original nodes may have come from an `@utility` we still
+          // want to replace the source because the `@apply` is ultimately the
+          // reason the node was emitted into the AST.
+          walk([node], (node) => {
+            node.src = src
+          })
+        }
+
+        let candidateAst = details.map((d) => d.node)
+
         // Collect the nodes to insert in place of the `@apply` rule. When a rule
         // was used, we want to insert its children instead of the rule because we
         // don't want the wrapping selector.