Add CodeLens annotations for Server/Client Components

package.json

@@ -74,6 +74,21 @@
           "default": true,
           "description": "Highlight TypeScript interface and type alias declaration names."
         },
+        "reactComponentLens.codelens.enabled": {
+          "type": "boolean",
+          "default": false,
+          "description": "Show CodeLens annotations indicating Server/Client Component kind."
+        },
+        "reactComponentLens.codelens.clientComponent": {
+          "type": "boolean",
+          "default": true,
+          "description": "Show CodeLens for Client Components."
+        },
+        "reactComponentLens.codelens.serverComponent": {
+          "type": "boolean",
+          "default": true,
+          "description": "Show CodeLens for Server Components."
+        },
         "reactComponentLens.highlightColors": {
           "type": "object",
           "default": {

src/analyzer.ts

@@ -96,6 +96,49 @@ export class ComponentLensAnalyzer {
     this.directiveCache.delete(filePath)
   }
 
+  public async findComponentDeclaration(
+    filePath: string,
+    componentName: string,
+  ): Promise<{ character: number; line: number } | undefined> {
+    const sourceText = this.host.readFileAsync
+      ? await this.host.readFileAsync(filePath)
+      : this.host.readFile(filePath)
+
+    if (sourceText === undefined) {
+      return undefined
+    }
+
+    const signature = this.host.getSignatureAsync
+      ? await this.host.getSignatureAsync(filePath)
+      : this.host.getSignature(filePath)
+
+    if (signature === undefined) {
+      return undefined
+    }
+
+    const analysis = this.getAnalysis(filePath, sourceText, signature)
+    if (!analysis) {
+      return undefined
+    }
+
+    const component = analysis.localComponents.get(componentName)
+    if (!component || component.ranges.length === 0) {
+      return undefined
+    }
+
+    const offset = component.ranges[0]!.start
+    let line = 0
+    let lastNewline = -1
+    for (let i = 0; i < offset; i++) {
+      if (sourceText.charCodeAt(i) === 10) {
+        line++
+        lastNewline = i
+      }
+    }
+
+    return { character: offset - lastNewline - 1, line }
+  }
+
   public async analyzeDocument(
     filePath: string,
     sourceText: string,

src/codelens.ts

@@ -0,0 +1,222 @@
+import * as vscode from 'vscode'
+
+import type {
+  ComponentLensAnalyzer,
+  ComponentUsage,
+  ScopeConfig,
+} from './analyzer'
+import { createOpenSignature } from './resolver'
+
+export interface CodeLensConfig {
+  clientComponent: boolean
+  enabled: boolean
+  globalEnabled: boolean
+  serverComponent: boolean
+}
+
+const CODELENS_SCOPE: ScopeConfig = {
+  declaration: true,
+  element: true,
+  export: true,
+  import: true,
+  type: true,
+}
+
+interface LineGroup {
+  clients: Set<string>
+  components: Map<string, string>
+  servers: Set<string>
+}
+
+export class ComponentCodeLensProvider
+  implements vscode.CodeLensProvider, vscode.Disposable
+{
+  private readonly changeEmitter = new vscode.EventEmitter<void>()
+  public readonly onDidChangeCodeLenses = this.changeEmitter.event
+
+  private config: CodeLensConfig
+
+  public constructor(
+    private readonly analyzer: ComponentLensAnalyzer,
+    config: CodeLensConfig,
+  ) {
+    this.config = config
+  }
+
+  public updateConfig(config: CodeLensConfig): void {
+    this.config = config
+    this.changeEmitter.fire()
+  }
+
+  public refresh(): void {
+    this.changeEmitter.fire()
+  }
+
+  public dispose(): void {
+    this.changeEmitter.dispose()
+  }
+
+  public async provideCodeLenses(
+    document: vscode.TextDocument,
+  ): Promise<vscode.CodeLens[]> {
+    if (!this.config.globalEnabled || !this.config.enabled) {
+      return []
+    }
+
+    if (!this.config.clientComponent && !this.config.serverComponent) {
+      return []
+    }
+
+    const signature = createOpenSignature(document.version)
+    const usages = await this.analyzer.analyzeDocument(
+      document.fileName,
+      document.getText(),
+      signature,
+      CODELENS_SCOPE,
+    )
+
+    return this.buildCodeLenses(document, usages)
+  }
+
+  private async buildCodeLenses(
+    document: vscode.TextDocument,
+    usages: ComponentUsage[],
+  ): Promise<vscode.CodeLens[]> {
+    const lineMap = new Map<number, LineGroup>()
+
+    for (let i = 0; i < usages.length; i++) {
+      const usage = usages[i]!
+      if (usage.kind === 'client' && !this.config.clientComponent) {
+        continue
+      }
+      if (usage.kind === 'server' && !this.config.serverComponent) {
+        continue
+      }
+
+      if (usage.ranges.length === 0) {
+        continue
+      }
+
+      const line = document.positionAt(usage.ranges[0]!.start).line
+      let entry = lineMap.get(line)
+      if (!entry) {
+        entry = {
+          clients: new Set(),
+          components: new Map(),
+          servers: new Set(),
+        }
+        lineMap.set(line, entry)
+      }
+
+      if (usage.kind === 'client') {
+        entry.clients.add(usage.tagName)
+      } else {
+        entry.servers.add(usage.tagName)
+      }
+
+      if (!entry.components.has(usage.tagName)) {
+        entry.components.set(usage.tagName, usage.sourceFilePath)
+      }
+    }
+
+    const positions = await this.resolveDeclarationPositions(lineMap)
+    const codeLenses: vscode.CodeLens[] = []
+
+    for (const [line, { clients, servers, components }] of lineMap) {
+      const parts: string[] = []
+      if (clients.size > 0) {
+        parts.push('Client Component')
+      }
+      if (servers.size > 0) {
+        parts.push('Server Component')
+      }
+
+      if (parts.length === 0) {
+        continue
+      }
+
+      const locations = this.buildLocations(components, positions)
+      const position = new vscode.Position(line, 0)
+
+      if (locations.length > 0) {
+        codeLenses.push(
+          new vscode.CodeLens(new vscode.Range(position, position), {
+            arguments: [document.uri, position, locations, 'peek'],
+            command: 'editor.action.peekLocations',
+            title: parts.join(' · '),
+          }),
+        )
+      } else {
+        codeLenses.push(
+          new vscode.CodeLens(new vscode.Range(position, position), {
+            command: '',
+            title: parts.join(' · '),
+          }),
+        )
+      }
+    }
+
+    return codeLenses
+  }
+
+  private async resolveDeclarationPositions(
+    lineMap: Map<number, LineGroup>,
+  ): Promise<Map<string, vscode.Position>> {
+    const filesToResolve = new Map<string, Set<string>>()
+
+    for (const [, { components }] of lineMap) {
+      for (const [tagName, sourceFilePath] of components) {
+        let names = filesToResolve.get(sourceFilePath)
+        if (!names) {
+          names = new Set()
+          filesToResolve.set(sourceFilePath, names)
+        }
+        names.add(tagName)
+      }
+    }
+
+    const positions = new Map<string, vscode.Position>()
+
+    await Promise.all(
+      Array.from(filesToResolve, async ([filePath, names]) => {
+        for (const name of names) {
+          const pos = await this.analyzer.findComponentDeclaration(
+            filePath,
+            name,
+          )
+          if (pos) {
+            positions.set(
+              filePath + ':' + name,
+              new vscode.Position(pos.line, pos.character),
+            )
+          }
+        }
+      }),
+    )
+
+    return positions
+  }
+
+  private buildLocations(
+    components: Map<string, string>,
+    positions: Map<string, vscode.Position>,
+  ): vscode.Location[] {
+    const seen = new Set<string>()
+    const locations: vscode.Location[] = []
+
+    for (const [tagName, sourceFilePath] of components) {
+      if (seen.has(sourceFilePath)) {
+        continue
+      }
+      seen.add(sourceFilePath)
+
+      const uri = vscode.Uri.file(sourceFilePath)
+      const pos =
+        positions.get(sourceFilePath + ':' + tagName) ??
+        new vscode.Position(0, 0)
+      locations.push(new vscode.Location(uri, pos))
+    }
+
+    return locations
+  }
+}

src/extension.ts

@@ -4,8 +4,14 @@ import * as path from 'node:path'
 import * as vscode from 'vscode'
 
 import { ComponentLensAnalyzer, type ScopeConfig } from './analyzer'
+import { type CodeLensConfig, ComponentCodeLensProvider } from './codelens'
 import { type HighlightColors, LensDecorations } from './decorations'
-import { ImportResolver, type SourceHost } from './resolver'
+import {
+  createDiskSignature,
+  createOpenSignature,
+  ImportResolver,
+  type SourceHost,
+} from './resolver'
 
 const LANG_JSX = 'javascriptreact'
 const LANG_TSX = 'typescriptreact'
@@ -22,8 +28,19 @@ export function activate(context: vscode.ExtensionContext): void {
   const resolver = new ImportResolver(sourceHost)
   const analyzer = new ComponentLensAnalyzer(sourceHost, resolver)
   const decorations = new LensDecorations(config.highlightColors)
+  const codeLensProvider = new ComponentCodeLensProvider(
+    analyzer,
+    config.codeLens,
+  )
+
+  context.subscriptions.push(decorations, codeLensProvider)
 
-  context.subscriptions.push(decorations)
+  context.subscriptions.push(
+    vscode.languages.registerCodeLensProvider(
+      [{ language: LANG_TSX }, { language: LANG_JSX }],
+      codeLensProvider,
+    ),
+  )
 
   let refreshTimer: NodeJS.Timeout | undefined
   let watcherDisposables: vscode.Disposable[] = []
@@ -47,6 +64,7 @@ export function activate(context: vscode.ExtensionContext): void {
 
     refreshTimer = setTimeout(() => {
       refreshTimer = undefined
+      codeLensProvider.refresh()
       void refreshVisibleEditors()
     }, delay)
   }
@@ -153,6 +171,10 @@ export function activate(context: vscode.ExtensionContext): void {
         decorations.updateColors(config.highlightColors)
       }
 
+      if (event.affectsConfiguration('reactComponentLens.codelens')) {
+        codeLensProvider.updateConfig(config.codeLens)
+      }
+
       scheduleRefresh(0)
     }),
     new vscode.Disposable(() => {
@@ -171,6 +193,7 @@ export function deactivate(): void {
 }
 
 function getConfiguration(): {
+  codeLens: CodeLensConfig
   debounceMs: number
   enabled: boolean
   highlightColors: HighlightColors
@@ -184,6 +207,18 @@ function getConfiguration(): {
   )
 
   return {
+    codeLens: {
+      clientComponent: configuration.get<boolean>(
+        'codelens.clientComponent',
+        true,
+      ),
+      enabled: configuration.get<boolean>('codelens.enabled', false),
+      globalEnabled: configuration.get<boolean>('enabled', true),
+      serverComponent: configuration.get<boolean>(
+        'codelens.serverComponent',
+        true,
+      ),
+    },
     debounceMs: Math.max(0, Math.min(2000, debounceMs)),
     enabled: configuration.get<boolean>('enabled', true),
     highlightColors: {
@@ -300,14 +335,6 @@ class WorkspaceSourceHost implements SourceHost {
   }
 }
 
-function createOpenSignature(version: number): string {
-  return 'open:' + version
-}
-
-function createDiskSignature(mtimeMs: number, size: number): string {
-  return 'disk:' + mtimeMs + ':' + size
-}
-
 function normalizeColor(
   color: string | undefined,
   fallbackColor: string,