Track source locations when parsing and serializing

wip

packages/tailwindcss/src/ast.bench.ts

@@ -0,0 +1,29 @@
+import { bench } from 'vitest'
+import { toCss } from './ast'
+import * as CSS from './css-parser'
+import { TrackLocations } from './track-locations'
+
+const css = String.raw
+const input = css`
+  @theme {
+    --color-primary: #333;
+  }
+  @tailwind utilities;
+  .foo {
+    color: red;
+    /* comment */
+    &:hover {
+      color: blue;
+      @apply font-bold;
+    }
+  }
+`
+const ast = CSS.parse(input)
+
+bench('toCss', () => {
+  toCss(ast)
+})
+
+bench('toCss with source maps', () => {
+  toCss(ast, new TrackLocations())
+})

packages/tailwindcss/src/ast.ts

@@ -1,3 +1,5 @@
+import type { Location, TrackLocations } from './track-locations'
+
 export type Rule = {
   kind: 'rule'
   selector: string
@@ -87,27 +89,70 @@ export function walk(
   }
 }
 
-export function toCss(ast: AstNode[]) {
+function span(value: string, location: Location) {
-  let atRoots: string = ''
+  let line = location.line
-  let seenAtProperties = new Set<string>()
+  let column = location.column
 
-  function stringify(node: AstNode, depth = 0): string {
+  let start = { line, column }
+
+  for (let i = 0; i < value.length; ++i) {
+    if (value.charCodeAt(i) === 0x0a) {
+      // Add the number of lines the comment spans
+      line += 1
+      column = 0
+    } else {
+      // Keep track of the column for accurate end locations
+      column += 1
+    }
+  }
+
+  let end = { line, column }
+
+  location.line = line
+  location.column = column
+
+  return { start, end }
+}
+
+export function toCss(ast: AstNode[], track?: TrackLocations) {
+  let atRoots: AstNode[] = []
+  let seenAtProperties = new Set<string>()
+  let location = { line: 1, column: 0 }
+
+  function stringify(node: AstNode, depth = 0, location: Location): string {
     let css = ''
     let indent = '  '.repeat(depth)
 
+    function write(node: AstNode, str: string) {
+      // Make sure the CSS is indented correctly
+      css += indent
+      location.column += indent.length
+
+      // Write the line of CSS (comment, declaration, selector, etc…)
+      css += str
+
+      // Track the start and end of the CSS
+      track?.dst(node, span(str, location))
+
+      // Add another line after the end
+      css += `\n`
+      location.line += 1
+      location.column = 0
+    }
+
     // Rule
     if (node.kind === 'rule') {
       // Pull out `@at-root` rules to append later
       if (node.selector === '@at-root') {
         for (let child of node.nodes) {
-          atRoots += stringify(child, 0)
+          atRoots.push(child)
         }
         return css
       }
 
       if (node.selector === '@tailwind utilities') {
         for (let child of node.nodes) {
-          css += stringify(child, depth)
+          css += stringify(child, depth, location)
         }
         return css
       }
@@ -120,6 +165,16 @@ export function toCss(ast: AstNode[]) {
       // @layer base, components, utilities;
       // ```
       if (node.selector[0] === '@' && node.nodes.length === 0) {
+        // Make sure the CSS is indented correctly
+        location.column += indent.length
+
+        // Track the start and end of the CSS
+        track?.dst(node, span(`${node.selector};`, location))
+
+        // Track new line at the end
+        location.line += 1
+        location.column = 0
+
         return `${indent}${node.selector};\n`
       }
 
@@ -132,33 +187,44 @@ export function toCss(ast: AstNode[]) {
         seenAtProperties.add(node.selector)
       }
 
-      css += `${indent}${node.selector} {\n`
+      write(node, `${node.selector} {`)
+
       for (let child of node.nodes) {
-        css += stringify(child, depth + 1)
+        css += stringify(child, depth + 1, location)
       }
-      css += `${indent}}\n`
+
+      write(node, `}`)
     }
 
     // Comment
     else if (node.kind === 'comment') {
-      css += `${indent}/*${node.value}*/\n`
+      write(node, `/*${node.value}*/`)
     }
 
     // Declaration
     else if (node.property !== '--tw-sort' && node.value !== undefined && node.value !== null) {
-      css += `${indent}${node.property}: ${node.value}${node.important ? '!important' : ''};\n`
+      write(node, `${node.property}: ${node.value}${node.important ? '!important' : ''};`)
+    }
+
+    return css
+  }
+
+  function stringifyAst(nodes: AstNode[], location: Location) {
+    let css = ''
+    for (let node of nodes) {
+      let result = stringify(node, 0, location)
+      if (result !== '') {
+        css += result
+      }
     }
 
     return css
   }
 
   let css = ''
-  for (let node of ast) {
-    let result = stringify(node)
-    if (result !== '') {
-      css += result
-    }
-  }
 
-  return `${css}${atRoots}`
+  css += stringifyAst(ast, location)
+  css += stringifyAst(atRoots, location)
+
+  return css
 }

packages/tailwindcss/src/css-parser.bench.ts

@@ -0,0 +1,27 @@
+import { bench } from 'vitest'
+import * as CSS from './css-parser'
+import { TrackLocations } from './track-locations'
+
+const css = String.raw
+const input = css`
+  @theme {
+    --color-primary: #333;
+  }
+  @tailwind utilities;
+  .foo {
+    color: red;
+    /* comment */
+    &:hover {
+      color: blue;
+      @apply font-bold;
+    }
+  }
+`
+
+bench('CSS', () => {
+  CSS.parse(input)
+})
+
+bench('CSS with sourcemaps', () => {
+  CSS.parse(input, new TrackLocations())
+})

packages/tailwindcss/src/css-parser.ts

@@ -1,4 +1,5 @@
 import { comment, rule, type AstNode, type Comment, type Declaration, type Rule } from './ast'
+import type { TrackLocations } from './track-locations'
 
 const BACKSLASH = 0x5c
 const SLASH = 0x2f
@@ -20,7 +21,7 @@ const DASH = 0x2d
 const AT_SIGN = 0x40
 const EXCLAMATION_MARK = 0x21
 
-export function parse(input: string) {
+export function parse(input: string, track?: TrackLocations) {
   input = input.replaceAll('\r\n', '\n')
 
   let ast: AstNode[] = []
@@ -36,9 +37,48 @@ export function parse(input: string) {
 
   let peekChar
 
+  // The current line number
+  let line = 1
+
+  // The index of the first non-whitespace character on the current line
+  let lineStart = 0
+
+  // Source location tracking
+  let sourceStartLine = 1
+  let sourceStartColumn = 0
+  let sourceEndLine = 1
+  let sourceEndColumn = 0
+
+  function sourceRange() {
+    if (!track) return null
+
+    return {
+      start: {
+        line: sourceStartLine,
+        column: sourceStartColumn,
+      },
+      end: {
+        line: sourceEndLine,
+        column: sourceEndColumn,
+      },
+    }
+  }
+
   for (let i = 0; i < input.length; i++) {
     let currentChar = input.charCodeAt(i)
 
+    if (currentChar === LINE_BREAK) {
+      line += 1
+      lineStart = i + 1
+
+      if (buffer.length === 0) {
+        sourceStartLine = line
+        sourceStartColumn = 0
+        sourceEndLine = line
+        sourceEndColumn = 0
+      }
+    }
+
     // Current character is a `\` therefore the next character is escaped,
     // consume it together with the next character and continue.
     //
@@ -81,6 +121,19 @@ export function parse(input: string) {
           j += 1
         }
 
+        // Count newline within comments
+        else if (peekChar === LINE_BREAK) {
+          line += 1
+          lineStart = j + 1
+
+          if (buffer.length === 0) {
+            sourceStartLine = line
+            sourceStartColumn = 0
+            sourceEndLine = line
+            sourceEndColumn = 0
+          }
+        }
+
         // End of the comment
         else if (peekChar === ASTERISK && input.charCodeAt(j + 1) === SLASH) {
           i = j + 1
@@ -93,7 +146,9 @@ export function parse(input: string) {
       // Collect all license comments so that we can hoist them to the top of
       // the AST.
       if (commentString.charCodeAt(2) === EXCLAMATION_MARK) {
-        licenseComments.push(comment(commentString.slice(2, -2)))
+        let node = comment(commentString.slice(2, -2))
+        licenseComments.push(node)
+        track?.src(node, sourceRange()!)
       }
     }
 
@@ -211,6 +266,19 @@ export function parse(input: string) {
               k += 1
             }
 
+            // Count newline within comments
+            else if (peekChar === LINE_BREAK) {
+              line += 1
+              lineStart = j + 1
+
+              if (buffer.length === 0) {
+                sourceStartLine = line
+                sourceStartColumn = 0
+                sourceEndLine = line
+                sourceEndColumn = 0
+              }
+            }
+
             // End of the comment
             else if (peekChar === ASTERISK && input.charCodeAt(k + 1) === SLASH) {
               j = k + 1
@@ -273,6 +341,19 @@ export function parse(input: string) {
             closingBracketStack = closingBracketStack.slice(0, -1)
           }
         }
+
+        // Count newlines
+        else if (peekChar === LINE_BREAK) {
+          line += 1
+          lineStart = j + 1
+
+          if (buffer.length === 0) {
+            sourceStartLine = line
+            sourceStartColumn = 0
+            sourceEndLine = line
+            sourceEndColumn = 0
+          }
+        }
       }
 
       let declaration = parseDeclaration(buffer, colonIdx)
@@ -281,8 +362,13 @@ export function parse(input: string) {
       } else {
         ast.push(declaration)
       }
+      track?.src(declaration, sourceRange()!)
 
       buffer = ''
+      sourceStartLine = line
+      sourceStartColumn = i - lineStart
+      sourceEndLine = line
+      sourceEndColumn = i - lineStart
     }
 
     // End of a body-less at-rule.
@@ -306,9 +392,16 @@ export function parse(input: string) {
         ast.push(node)
       }
 
+      // Track the source location for source maps
+      track?.src(node, sourceRange()!)
+
       // Reset the state for the next node.
       buffer = ''
       node = null
+      sourceStartLine = line
+      sourceStartColumn = i - lineStart
+      sourceEndLine = line
+      sourceEndColumn = i - lineStart
     }
 
     // End of a declaration.
@@ -330,7 +423,14 @@ export function parse(input: string) {
         ast.push(declaration)
       }
 
+      // Track the source location for source maps
+      track?.src(declaration, sourceRange()!)
+
       buffer = ''
+      sourceStartLine = line
+      sourceStartColumn = i - lineStart
+      sourceEndLine = line
+      sourceEndColumn = i - lineStart
     }
 
     // Start of a block.
@@ -353,9 +453,16 @@ export function parse(input: string) {
       // attached to it.
       parent = node
 
+      // Track the source location for source maps
+      track?.src(node, sourceRange()!)
+
       // Reset the state for the next node.
       buffer = ''
       node = null
+      sourceStartLine = line
+      sourceStartColumn = i - lineStart
+      sourceEndLine = line
+      sourceEndColumn = i - lineStart
     }
 
     // End of a block.
@@ -393,9 +500,16 @@ export function parse(input: string) {
             ast.push(node)
           }
 
+          // Track the source location for source maps
+          track?.src(node, sourceRange()!)
+
           // Reset the state for the next node.
           buffer = ''
           node = null
+          sourceStartLine = line
+          sourceStartColumn = i - lineStart
+          sourceEndLine = line
+          sourceEndColumn = i - lineStart
         }
 
         // But it can also happen for declarations.
@@ -417,14 +531,16 @@ export function parse(input: string) {
           // Attach the declaration to the parent.
           if (parent) {
             let importantIdx = buffer.indexOf('!important', colonIdx + 1)
-            parent.nodes.push({
+            let node = {
               kind: 'declaration',
               property: buffer.slice(0, colonIdx).trim(),
               value: buffer
                 .slice(colonIdx + 1, importantIdx === -1 ? buffer.length : importantIdx)
                 .trim(),
               important: importantIdx !== -1,
-            } satisfies Declaration)
+            } satisfies Declaration
+            parent.nodes.push(node)
+            track?.src(node, sourceRange()!)
           }
         }
       }
@@ -445,6 +561,10 @@ export function parse(input: string) {
       // Reset the state for the next node.
       buffer = ''
       node = null
+      sourceStartLine = line
+      sourceStartColumn = i - lineStart
+      sourceEndLine = line
+      sourceEndColumn = i - lineStart
     }
 
     // Any other character is part of the current node.
@@ -454,6 +574,10 @@ export function parse(input: string) {
         buffer.length === 0 &&
         (currentChar === SPACE || currentChar === LINE_BREAK || currentChar === TAB)
       ) {
+        sourceStartLine = line
+        sourceStartColumn = i + 1 - lineStart
+        sourceEndLine = line
+        sourceEndColumn = i + 1 - lineStart
         continue
       }
 

packages/tailwindcss/src/track-locations.ts

@@ -0,0 +1,28 @@
+import type { AstNode } from './ast'
+
+export type Location = {
+  line: number
+  column: number
+}
+
+export type Range = {
+  start: Location
+  end: Location
+}
+
+export class TrackLocations {
+  sources: Map<AstNode, Range[]> = new Map()
+  destinations: Map<AstNode, Range[]> = new Map()
+
+  src(node: AstNode, range: Range) {
+    let ranges = this.sources.get(node) ?? []
+    this.sources.set(node, ranges)
+    ranges.push(range)
+  }
+
+  dst(node: AstNode, range: Range) {
+    let ranges = this.destinations.get(node) ?? []
+    this.destinations.set(node, ranges)
+    ranges.push(range)
+  }
+}