feat: transform code for display

apps/v4/components/component-source.tsx

@@ -4,6 +4,7 @@ import * as React from "react"
 
 import { highlightCode } from "@/lib/highlight-code"
 import { getRegistryItem } from "@/lib/registry"
+import { transformForDisplay } from "@/lib/rehype"
 import { cn } from "@/lib/utils"
 import { CodeCollapsibleWrapper } from "@/components/code-collapsible-wrapper"
 import { CopyButton } from "@/components/copy-button"
@@ -53,6 +54,9 @@ export async function ComponentSource({
   code = code.replaceAll("export default", "export")
   code = code.replaceAll("/* eslint-disable react/no-children-prop */\n", "")
 
+  // Apply transforms (cn-* → Tailwind, IconPlaceholder → icons, etc.).
+  code = await transformForDisplay(code, styleName)
+
   const lang = language ?? title?.split(".").pop() ?? "tsx"
   const highlightedCode = await highlightCode(code, lang)
 

apps/v4/lib/rehype.ts

@@ -1,7 +1,17 @@
-import fs from "fs"
+import fs, { promises as fsPromises } from "fs"
 import path from "path"
+import type { configSchema } from "shadcn/schema"
+import {
+  createStyleMap,
+  transformIcons,
+  transformMenu,
+  transformRender,
+  transformStyle,
+} from "shadcn/utils"
+import { Project, ScriptKind } from "ts-morph"
 import { u } from "unist-builder"
 import { visit } from "unist-util-visit"
+import { z } from "zod"
 
 import { Index as StylesIndex } from "@/registry/__index__"
 import { getActiveStyle } from "@/registry/_legacy-styles"
@@ -18,6 +28,103 @@ function getIndexForStyle(styleName: string) {
   return { index: StylesIndex, key: styleName }
 }
 
+// Extract style name from compound style (e.g., "radix-nova" → "nova").
+function getStyleFromStyleName(styleName: string) {
+  const parts = styleName.split("-")
+  return parts.length > 1 ? parts.slice(1).join("-") : styleName
+}
+
+// Build minimal config for transforms.
+function buildDisplayConfig(styleName: string): z.infer<typeof configSchema> {
+  return {
+    $schema: "https://ui.shadcn.com/schema.json",
+    style: styleName,
+    rsc: true,
+    tsx: true,
+    tailwind: {
+      config: "",
+      css: "",
+      baseColor: "neutral",
+      cssVariables: true,
+      prefix: "",
+    },
+    iconLibrary: "lucide",
+    aliases: {
+      components: "@/components",
+      utils: "@/lib/utils",
+      ui: "@/components/ui",
+      lib: "@/lib",
+      hooks: "@/hooks",
+    },
+    resolvedPaths: {
+      cwd: "/",
+      tailwindConfig: "",
+      tailwindCss: "",
+      utils: "@/lib/utils",
+      components: "@/components",
+      lib: "@/lib",
+      hooks: "@/hooks",
+      ui: "@/components/ui",
+    },
+  }
+}
+
+// Cache for style maps to avoid repeated file reads.
+const styleMapCache = new Map<string, Record<string, string>>()
+
+async function getStyleMap(styleName: string) {
+  const style = getStyleFromStyleName(styleName)
+
+  if (styleMapCache.has(style)) {
+    return styleMapCache.get(style)!
+  }
+
+  try {
+    const cssPath = path.join(
+      process.cwd(),
+      `registry/styles/style-${style}.css`
+    )
+    const css = await fsPromises.readFile(cssPath, "utf-8")
+    const styleMap = createStyleMap(css)
+    styleMapCache.set(style, styleMap)
+    return styleMap
+  } catch {
+    // Return empty style map if file not found.
+    return {}
+  }
+}
+
+export async function transformForDisplay(content: string, styleName: string) {
+  try {
+    // 1. Apply style transformation (cn-* → Tailwind classes).
+    const styleMap = await getStyleMap(styleName)
+    let transformed = await transformStyle(content, { styleMap })
+
+    // 2. Apply icon/menu/render transforms.
+    const config = buildDisplayConfig(styleName)
+    const project = new Project({ compilerOptions: {} })
+    const sourceFile = project.createSourceFile("component.tsx", transformed, {
+      scriptKind: ScriptKind.TSX,
+    })
+
+    const transformers = [transformIcons, transformMenu, transformRender]
+    for (const transformer of transformers) {
+      await transformer({
+        filename: "component.tsx",
+        raw: transformed,
+        sourceFile,
+        config,
+      })
+    }
+
+    return sourceFile.getText()
+  } catch (error) {
+    // Return original content if transform fails.
+    console.error("Transform failed:", error)
+    return content
+  }
+}
+
 interface UnistNode {
   type: string
   name?: string
@@ -37,12 +144,23 @@ export interface UnistTree {
   children: UnistNode[]
 }
 
+// Collected node info for async processing.
+interface NodeToProcess {
+  node: UnistNode
+  type: "ComponentSource" | "ComponentPreview"
+  name: string
+  styleName: string
+  fileName?: string
+  srcPath?: string
+}
+
 export function rehypeComponent() {
   return async (tree: UnistTree) => {
     const activeStyle = await getActiveStyle()
+    const nodesToProcess: NodeToProcess[] = []
 
+    // First pass: collect all nodes that need processing.
     visit(tree, (node: UnistNode) => {
-      // src prop overrides both name and fileName.
       const { value: srcPath } =
         (getNodeAttributeByName(node, "src") as {
           name: string
@@ -59,71 +177,15 @@ export function rehypeComponent() {
           (getNodeAttributeByName(node, "styleName")?.value as string) ||
           activeStyle.name
 
-        if (!name && !srcPath) {
-          return null
-        }
-
-        try {
-          let src: string
-
-          if (srcPath) {
-            src = path.join(process.cwd(), srcPath)
-          } else {
-            const { index, key } = getIndexForStyle(styleName)
-            const component = index[key]?.[name]
-            src = fileName
-              ? component.files.find((file: unknown) => {
-                  if (typeof file === "string") {
-                    return (
-                      file.endsWith(`${fileName}.tsx`) ||
-                      file.endsWith(`${fileName}.ts`)
-                    )
-                  }
-                  return false
-                }) || component.files[0]?.path
-              : component.files[0]?.path
-          }
-
-          // Read the source file.
-          const filePath = src
-          let source = fs.readFileSync(filePath, "utf8")
-
-          // Replace imports.
-          // TODO: Use @swc/core and a visitor to replace this.
-          // For now a simple regex should do.
-          source = source.replaceAll(
-            `@/registry/${styleName}/`,
-            "@/components/"
-          )
-          source = source.replaceAll(`@/registry/bases/radix/`, "@/components/")
-          source = source.replaceAll(`@/registry/bases/base/`, "@/components/")
-          source = source.replaceAll("export default", "export")
-
-          // Add code as children so that rehype can take over at build time.
-          node.children?.push(
-            u("element", {
-              tagName: "pre",
-              properties: {
-                __src__: src,
-              },
-              children: [
-                u("element", {
-                  tagName: "code",
-                  properties: {
-                    className: ["language-tsx"],
-                  },
-                  children: [
-                    {
-                      type: "text",
-                      value: source,
-                    },
-                  ],
-                }),
-              ],
-            })
-          )
-        } catch (error) {
-          console.error(error)
+        if (name || srcPath) {
+          nodesToProcess.push({
+            node,
+            type: "ComponentSource",
+            name,
+            styleName,
+            fileName,
+            srcPath,
+          })
         }
       }
 
@@ -133,32 +195,62 @@ export function rehypeComponent() {
           (getNodeAttributeByName(node, "styleName")?.value as string) ||
           activeStyle.name
 
-        if (!name) {
-          return null
+        if (name) {
+          nodesToProcess.push({
+            node,
+            type: "ComponentPreview",
+            name,
+            styleName,
+          })
         }
+      }
+    })
 
+    // Second pass: process all collected nodes asynchronously.
+    await Promise.all(
+      nodesToProcess.map(async (item) => {
         try {
-          const { index, key } = getIndexForStyle(styleName)
-          const component = index[key]?.[name]
-          const src = component.files[0]?.path
+          let src: string
+
+          if (item.srcPath) {
+            src = path.join(process.cwd(), item.srcPath)
+          } else {
+            const { index, key } = getIndexForStyle(item.styleName)
+            const component = index[key]?.[item.name]
+
+            if (item.type === "ComponentSource" && item.fileName) {
+              src =
+                component.files.find((file: unknown) => {
+                  if (typeof file === "string") {
+                    return (
+                      file.endsWith(`${item.fileName}.tsx`) ||
+                      file.endsWith(`${item.fileName}.ts`)
+                    )
+                  }
+                  return false
+                }) || component.files[0]?.path
+            } else {
+              src = component.files[0]?.path
+            }
+          }
 
           // Read the source file.
-          const filePath = src
-          let source = fs.readFileSync(filePath, "utf8")
+          let source = fs.readFileSync(src, "utf8")
 
           // Replace imports.
-          // TODO: Use @swc/core and a visitor to replace this.
-          // For now a simple regex should do.
           source = source.replaceAll(
-            `@/registry/${styleName}/`,
+            `@/registry/${item.styleName}/`,
             "@/components/"
           )
           source = source.replaceAll(`@/registry/bases/radix/`, "@/components/")
           source = source.replaceAll(`@/registry/bases/base/`, "@/components/")
           source = source.replaceAll("export default", "export")
 
+          // Apply transforms (cn-* → Tailwind, IconPlaceholder → icons, etc.).
+          source = await transformForDisplay(source, item.styleName)
+
           // Add code as children so that rehype can take over at build time.
-          node.children?.push(
+          item.node.children?.push(
             u("element", {
               tagName: "pre",
               properties: {
@@ -183,8 +275,8 @@ export function rehypeComponent() {
         } catch (error) {
           console.error(error)
         }
-      }
-    })
+      })
+    )
   }
 }