From 00f0821722976758cf39ab671d90395bec0bf2a0 Mon Sep 17 00:00:00 2001
From: Ziwei Wang <ziweiwzw217@gmail.com>
Date: Tue, 30 Apr 2024 16:54:17 -0400
Subject: [PATCH] Feat: Add codeLens for bitbake function references

---
 client/src/ui/BitbakeCommands.ts | 17 +++++++++++++++++
 server/src/server.ts             | 31 +++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/client/src/ui/BitbakeCommands.ts b/client/src/ui/BitbakeCommands.ts
index d8dc6877..501b19f2 100644
--- a/client/src/ui/BitbakeCommands.ts
+++ b/client/src/ui/BitbakeCommands.ts
@@ -52,6 +52,8 @@ export function registerBitbakeCommands (context: vscode.ExtensionContext, bitba
     vscode.commands.registerCommand('bitbake.stop-toaster', async () => { await stopToaster(bitBakeProjectScanner.bitbakeDriver) }),
     vscode.commands.registerCommand('bitbake.clear-workspace-state', async () => { await clearAllWorkspaceState(context) }),
     vscode.commands.registerCommand('bitbake.examine-dependency-taskexp', async (uri) => { await examineDependenciesTaskexp(bitbakeWorkspace, bitBakeProjectScanner, uri) }),
+    vscode.commands.registerCommand('bitbake.codeLens.showReferences', async (uri, position) => { await showReferences(uri, position) }),
+
     // Handles enqueued parsing requests (onSave)
     vscode.tasks.onDidEndTask((e) => {
       if (e.execution.task.name === 'Bitbake: Parse') {
@@ -70,6 +72,21 @@ export function registerBitbakeCommands (context: vscode.ExtensionContext, bitba
   )
 }
 
+async function showReferences (uri: any, position: any): Promise<void> {
+  if (typeof uri !== 'string' || position === undefined) {
+    return
+  }
+
+  let _position: vscode.Position
+  try {
+    _position = new vscode.Position(position.line, position.character)
+    const locations = await vscode.commands.executeCommand('vscode.executeReferenceProvider', vscode.Uri.parse(uri), position)
+    await vscode.commands.executeCommand('editor.action.showReferences', vscode.Uri.parse(uri), _position, locations)
+  } catch (error: any) {
+    void vscode.window.showErrorMessage('Failed to show references: ' + JSON.stringify(error))
+  }
+}
+
 async function clearAllWorkspaceState (context: vscode.ExtensionContext, key?: string): Promise<void> {
   for (const key of context.workspaceState.keys()) {
     await context.workspaceState.update(key, undefined).then(
diff --git a/server/src/server.ts b/server/src/server.ts
index 9f50d8d6..fa4151e0 100644
--- a/server/src/server.ts
+++ b/server/src/server.ts
@@ -5,6 +5,7 @@
  * ------------------------------------------------------------------------------------------ */
 
 import path from 'path'
+import * as LSP from 'vscode-languageserver/node'
 import {
   type Connection,
   type InitializeResult,
@@ -111,6 +112,9 @@ disposables.push(
         },
         renameProvider: {
           prepareProvider: true
+        },
+        codeLensProvider: {
+          resolveProvider: true
         }
       }
     }
@@ -156,6 +160,33 @@ disposables.push(
     analyzer.clearRecipeLocalFiles()
   }),
 
+  connection.onCodeLens(async (params): Promise<LSP.CodeLens[]> => {
+    const codeLenses: LSP.CodeLens[] = []
+    const uri = params.textDocument.uri
+    // TODO: check the setting to see if showing the references code lens is allowed
+    const allSymbols = analyzer.getGlobalDeclarationSymbolsForUri(uri)
+    allSymbols.forEach((symbol) => {
+      if (symbol.kind === LSP.SymbolKind.Function) {
+        const codeLens = LSP.CodeLens.create(symbol.location.range)
+
+        codeLens.command = {
+          title: 'Show References',
+          command: 'bitbake.codeLens.showReferences',
+          arguments: [uri, symbol.location.range.start]
+        }
+
+        codeLens.data = { uri, position: symbol.location.range.start }
+
+        codeLenses.push(codeLens)
+      }
+    })
+    return codeLenses
+  }),
+
+  connection.onCodeLensResolve((codeLens): LSP.CodeLens => {
+    return codeLens
+  }),
+
   connection.onRequest(
     RequestMethod.EmbeddedLanguageTypeOnPosition,
     async ({ uriString, position }: RequestParams['EmbeddedLanguageTypeOnPosition']): RequestResult['EmbeddedLanguageTypeOnPosition'] => {