Skip to content

Commit

Permalink
Better indexing logic (#147)
Browse files Browse the repository at this point in the history
Better Indexing
  • Loading branch information
saranshsaini authored Dec 20, 2024
1 parent c11b6dd commit 527f3c1
Show file tree
Hide file tree
Showing 3 changed files with 184 additions and 49 deletions.
200 changes: 152 additions & 48 deletions CodeiumVS/LanguageServer/LanguageServer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,10 @@ public class LanguageServer
private readonly HttpClient _httpClient;
private readonly CodeiumVSPackage _package;


public readonly LanguageServerController Controller;


public LanguageServer()
{
_package = CodeiumVSPackage.Instance;
Expand Down Expand Up @@ -158,7 +160,7 @@ public async Task SignInAsync()

// TODO: should we use timeout = Timeout.InfiniteTimeSpan? default value is 100s (1m40s)
GetAuthTokenResponse? result =
await RequestCommandAsync<GetAuthTokenResponse>("GetAuthToken", new {});
await RequestCommandAsync<GetAuthTokenResponse>("GetAuthToken", new { });

if (result == null)
{
Expand Down Expand Up @@ -302,7 +304,7 @@ await _package.LogAsync(
KnownMonikers.StatusError,
true,
null,
[..actions, ..NotificationInfoBar.SupportActions]);
[.. actions, .. NotificationInfoBar.SupportActions]);
}
else
{
Expand Down Expand Up @@ -373,7 +375,8 @@ private void ThreadDownloadLanguageServer(IVsThreadedWaitDialog4 progressDialog)
webClient.DownloadFileCompleted += (s, e) =>
{
ThreadHelper.JoinableTaskFactory
.RunAsync(async delegate {
.RunAsync(async delegate
{
await ThreadDownload_OnCompletedAsync(e, progressDialog, downloadDest);
})
.FireAndForget();
Expand Down Expand Up @@ -591,7 +594,7 @@ await _package.LogAsync(
KnownMonikers.StatusError,
true,
null,
[..actions, ..NotificationInfoBar.SupportActions]);
[.. actions, .. NotificationInfoBar.SupportActions]);

return;
}
Expand Down Expand Up @@ -739,47 +742,138 @@ private async Task InitializeTrackedWorkspaceAsync()
DTE dte = (DTE)ServiceProvider.GlobalProvider.GetService(typeof(DTE));
await _package.LogAsync($"Number of top-level projects: {dte.Solution.Projects.Count}");

List<string> processedProjects = new List<string>();
var documents = dte.Documents;
var openFilePaths = new HashSet<string>();
if (_package.SettingsPage.IndexOpenFiles)
{
foreach (EnvDTE.Document doc in documents)
{
await _package.LogAsync($"Open File: {doc.Path}");
openFilePaths.Add(doc.Path);
}
}

async Task ProcessProjectAsync(EnvDTE.Project project)
var inputProjectsToIndex = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
string projectListPath = _package.SettingsPage.IndexingFilesListPath.Trim();
try
{
try
if (!string.IsNullOrEmpty(projectListPath) && File.Exists(projectListPath))
{
string projectFullName = project.FullName;
await _package.LogAsync($"Project Full Name: {projectFullName}");
if (!string.IsNullOrEmpty(projectFullName) && !processedProjects.Contains(projectFullName))
string[] lines = File.ReadAllLines(projectListPath);
foreach (string line in lines)
{
processedProjects.Add(projectFullName);
string projectDir = Path.GetDirectoryName(projectFullName);
await _package.LogAsync($"Project Dir: {projectDir}");
AddTrackedWorkspaceResponse response = await AddTrackedWorkspaceAsync(projectDir);
if (response != null)
string trimmedLine = line.Trim();
if (!string.IsNullOrEmpty(trimmedLine))
{
_initializedWorkspace = true;
inputProjectsToIndex.Add(trimmedLine);
}
}
await _package.LogAsync($"Number of Projects loaded from {projectListPath}: {inputProjectsToIndex.Count}");
}
}
catch (Exception ex)
{
await _package.LogAsync($"Error reading project list: {ex.Message}");
}

// Process sub-projects (e.g., project references)
foreach (EnvDTE.ProjectItem item in project.ProjectItems)
List<string> projectsToIndex = await GetFilesToIndex(inputProjectsToIndex, openFilePaths, dte);
await _package.LogAsync($"Number of projects to index: {projectsToIndex.Count}");

for (int i = 0; i < projectsToIndex.Count; i++)
{
try
{
await _package.LogAsync($"Processing Project {i + 1} of {projectsToIndex.Count}: {projectsToIndex[i]}");
AddTrackedWorkspaceResponse response = await AddTrackedWorkspaceAsync(projectsToIndex[i]);
if (response != null)
{
if (item.SubProject != null)
_initializedWorkspace = true;
}
}
catch (Exception ex)
{
await _package.LogAsync($"Error processing project {i + 1} of {projectsToIndex.Count}: {ex.Message}");
}
}
}

private async Task<List<string>> GetFilesToIndex(HashSet<string> inputProjectsToIndex, HashSet<string> openFilePaths, DTE dte)
{
await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
int maxToIndex = 15;
HashSet<string> specifiedProjectsToIndexPath = new HashSet<string>();
HashSet<string> openFilesProjectsToIndexPath = new HashSet<string>();
HashSet<string> remainingProjectsToIndexPath = new HashSet<string>();
HashSet<string> processedProjects = new HashSet<string>();
async Task AddFilesToIndexLists(EnvDTE.Project project)
{
if (specifiedProjectsToIndexPath.Count == inputProjectsToIndex.Count && openFilePaths.Count == 0 && (specifiedProjectsToIndexPath.Count + remainingProjectsToIndexPath.Count + openFilesProjectsToIndexPath.Count) >= maxToIndex)
{
return;
}
string projectFullName = project.FullName;
string projectName = Path.GetFileNameWithoutExtension(projectFullName);
if (!string.IsNullOrEmpty(projectFullName) && !processedProjects.Contains(projectFullName))
{
string projectDir = Path.GetDirectoryName(projectFullName);

// There are three cases.
// 1. The project is in the list of projects to index passed in by the user. These take priority. The entire solution is searched until all are found.
// 2. Find the project the open file is a member of. Not sure if nested projects could match the same file multiple times so delete from the set when found.
// 3. Any other project. Tops it up to the max amount to index if the previous two cases didnt.
if (inputProjectsToIndex.Contains(projectName))
{
await _package.LogAsync($"Found in input list {projectName}");
specifiedProjectsToIndexPath.Add(projectDir);
}
else if (openFilePaths.Count != 0)
{
string matchingFile = null;
foreach (var filePath in openFilePaths)
{
if (filePath.StartsWith(projectDir, StringComparison.OrdinalIgnoreCase))
{
await _package.LogAsync($"Found in open files {filePath}");
matchingFile = filePath;
break;
}
}
if (!string.IsNullOrEmpty(matchingFile))
{
await ProcessProjectAsync(item.SubProject);
openFilesProjectsToIndexPath.Add(projectDir);
openFilePaths.Remove(matchingFile);
}
}
else
{
await _package.LogAsync($"Found in remaining {projectName}");
remainingProjectsToIndexPath.Add(projectDir);
}
processedProjects.Add(projectFullName);
}
catch (Exception ex)

foreach (EnvDTE.ProjectItem item in project.ProjectItems)
{
await _package.LogAsync("Error: Failed to initialize tracked workspace: " + ex.Message);
if (item.SubProject != null)
{
await AddFilesToIndexLists(item.SubProject);

}
}
}

foreach (EnvDTE.Project project in dte.Solution.Projects)
{
await ProcessProjectAsync(project);
await AddFilesToIndexLists(project);
}
List<string> result = new List<string>();
result.AddRange(specifiedProjectsToIndexPath);
result.AddRange(openFilesProjectsToIndexPath);
result.AddRange(remainingProjectsToIndexPath);
return result;
}


private async Task<T?> RequestCommandAsync<T>(string command, object data,
CancellationToken cancellationToken = default)
{
Expand All @@ -800,21 +894,28 @@ public async Task<IList<CompletionItem>?>
var uri = new System.Uri(absolutePath);
var absoluteUri = uri.AbsoluteUri;
GetCompletionsRequest data =
new() { metadata = GetMetadata(),
document = new() { text = text,
editor_language = language.Name,
language = language.Type,
cursor_offset = (ulong)cursorPosition,
line_ending = lineEnding,
absolute_path = absolutePath,
absolute_uri = absoluteUri,
relative_path = Path.GetFileName(absolutePath) },
editor_options = new() {
tab_size = (ulong)tabSize,
insert_spaces = insertSpaces,
disable_autocomplete_in_comments =
new()
{
metadata = GetMetadata(),
document = new()
{
text = text,
editor_language = language.Name,
language = language.Type,
cursor_offset = (ulong)cursorPosition,
line_ending = lineEnding,
absolute_path = absolutePath,
absolute_uri = absoluteUri,
relative_path = Path.GetFileName(absolutePath)
},
editor_options = new()
{
tab_size = (ulong)tabSize,
insert_spaces = insertSpaces,
disable_autocomplete_in_comments =
!_package.SettingsPage.EnableCommentCompletion,
} };
}
};

GetCompletionsResponse? result =
await RequestCommandAsync<GetCompletionsResponse>("GetCompletions", data, token);
Expand All @@ -831,7 +932,7 @@ public async Task AcceptCompletionAsync(string completionId)

public async Task<GetProcessesResponse?> GetProcessesAsync()
{
return await RequestCommandAsync<GetProcessesResponse>("GetProcesses", new {});
return await RequestCommandAsync<GetProcessesResponse>("GetProcesses", new { });
}

public async Task<AddTrackedWorkspaceResponse?> AddTrackedWorkspaceAsync(string workspacePath)
Expand All @@ -842,16 +943,19 @@ public async Task AcceptCompletionAsync(string completionId)

public Metadata GetMetadata()
{
return new() { request_id = _metadata.request_id++,
api_key = _metadata.api_key,
ide_name = _metadata.ide_name,
ide_version = _metadata.ide_version,

extension_name = _metadata.extension_name,
extension_version = _metadata.extension_version,
session_id = _metadata.session_id,
locale = _metadata.locale,
disable_telemetry = _metadata.disable_telemetry };
return new()
{
request_id = _metadata.request_id++,
api_key = _metadata.api_key,
ide_name = _metadata.ide_name,
ide_version = _metadata.ide_version,

extension_name = _metadata.extension_name,
extension_version = _metadata.extension_version,
session_id = _metadata.session_id,
locale = _metadata.locale,
disable_telemetry = _metadata.disable_telemetry
};
}

public async Task<IList<FunctionInfo>?>
Expand Down
31 changes: 31 additions & 0 deletions CodeiumVS/SettingsPage.cs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ public class SettingsPage : DialogPage
private bool enableIndexing = true;
private bool enableCodeLens = true;
private int indexingMaxFileCount = 5000;
private string indexingFilesListPath = "";
private bool indexOpenFiles = true;

[Category("Codeium")]
[DisplayName("Enterprise Mode")]
Expand Down Expand Up @@ -139,4 +141,33 @@ public int IndexingMaxWorkspaceSize
indexingMaxFileCount = value;
}
}

[Category("Codeium")]
[DisplayName("Files to Index List Path")]
[Description(
"Complete path to a .txt file that contains a line separated list of file paths to index. Leave blank to index full solution.")]
public string IndexingFilesListPath
{
get {
return indexingFilesListPath;
}
set {
indexingFilesListPath = value;
}
}
[Category("Codeium")]
[DisplayName("Index Open Files")]
[Description(
"Complete path to a .txt file that contains a line separated list of file paths to index. Leave blank to index full solution.")]
public bool IndexOpenFiles
{
get
{
return indexOpenFiles;
}
set
{
indexOpenFiles = value;
}
}
}
2 changes: 1 addition & 1 deletion CodeiumVS/SuggestionUI/SuggestionTagger.cs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using System.Collections.Generic;
using System.Collections.Generic;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Media;
Expand Down

0 comments on commit 527f3c1

Please sign in to comment.