Ionite
1 year ago
13 changed files with 302 additions and 30 deletions
@ -0,0 +1,70 @@
|
||||
using System.Collections.Generic; |
||||
using System.Data.Common; |
||||
using System.Globalization; |
||||
using System.IO; |
||||
using System.Threading.Tasks; |
||||
using StabilityMatrix.Avalonia.Models.TagCompletion; |
||||
using Sylvan.Data.Csv; |
||||
using Sylvan; |
||||
using Sylvan.Data; |
||||
|
||||
namespace StabilityMatrix.Avalonia.Helpers; |
||||
|
||||
public class TagCsvParser |
||||
{ |
||||
private readonly Stream stream; |
||||
|
||||
public TagCsvParser(Stream stream) |
||||
{ |
||||
this.stream = stream; |
||||
} |
||||
|
||||
public async IAsyncEnumerable<TagCsvEntry> ParseAsync() |
||||
{ |
||||
var pool = new StringPool(); |
||||
var options = new CsvDataReaderOptions |
||||
{ |
||||
StringFactory = pool.GetString, |
||||
HasHeaders = false, |
||||
}; |
||||
|
||||
using var textReader = new StreamReader(stream); |
||||
await using var dataReader = await CsvDataReader.CreateAsync(textReader, options); |
||||
|
||||
while (await dataReader.ReadAsync()) |
||||
{ |
||||
var entry = new TagCsvEntry |
||||
{ |
||||
Name = dataReader.GetString(0), |
||||
Type = dataReader.GetInt32(1), |
||||
Count = dataReader.GetInt32(2), |
||||
Aliases = dataReader.GetString(3), |
||||
}; |
||||
yield return entry; |
||||
} |
||||
|
||||
/*var dataBinderOptions = new DataBinderOptions |
||||
{ |
||||
BindingMode = DataBindingMode.Any |
||||
};*/ |
||||
/*var results = dataReader.GetRecordsAsync<TagCsvEntry>(dataBinderOptions); |
||||
return results;*/ |
||||
} |
||||
|
||||
public async Task<Dictionary<string, TagCsvEntry>> GetDictionaryAsync() |
||||
{ |
||||
var dict = new Dictionary<string, TagCsvEntry>(); |
||||
|
||||
await foreach (var entry in ParseAsync()) |
||||
{ |
||||
if (entry.Name is null || entry.Type is null) |
||||
{ |
||||
continue; |
||||
} |
||||
|
||||
dict.Add(entry.Name, entry); |
||||
} |
||||
|
||||
return dict; |
||||
} |
||||
} |
||||
@ -0,0 +1,119 @@
|
||||
using System; |
||||
using System.Collections.Generic; |
||||
using System.Threading.Tasks; |
||||
using AutoComplete.Builders; |
||||
using AutoComplete.Clients.IndexSearchers; |
||||
using AutoComplete.DataStructure; |
||||
using AutoComplete.Domain; |
||||
using NLog; |
||||
using StabilityMatrix.Avalonia.Controls.CodeCompletion; |
||||
using StabilityMatrix.Avalonia.Helpers; |
||||
using StabilityMatrix.Core.Helper; |
||||
using StabilityMatrix.Core.Models; |
||||
using StabilityMatrix.Core.Models.FileInterfaces; |
||||
|
||||
namespace StabilityMatrix.Avalonia.Models.TagCompletion; |
||||
|
||||
public class CompletionProvider : ICompletionProvider |
||||
{ |
||||
private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); |
||||
|
||||
private readonly Dictionary<string, TagCsvEntry> entries = new(); |
||||
|
||||
private InMemoryIndexSearcher? searcher; |
||||
|
||||
public bool IsLoaded => searcher is not null; |
||||
|
||||
public async Task LoadFromFile(FilePath path, bool recreate = false) |
||||
{ |
||||
// Get Blake3 hash of file |
||||
var hash = await FileHash.GetBlake3Async(path); |
||||
|
||||
Logger.Trace("Loading tags from {Path} with Blake3 hash {Hash}", path, hash); |
||||
|
||||
// Check for AppData/StabilityMatrix/Temp/Tags/<hash>/*.bin |
||||
var tempTagsDir = GlobalConfig.HomeDir.JoinDir("Temp", "Tags"); |
||||
tempTagsDir.Create(); |
||||
var hashDir = tempTagsDir.JoinDir(hash); |
||||
|
||||
var headerFile = hashDir.JoinFile("header.bin"); |
||||
var indexFile = hashDir.JoinFile("index.bin"); |
||||
var tailFile = hashDir.JoinFile("tail.bin"); |
||||
|
||||
entries.Clear(); |
||||
|
||||
// If directory or any file is missing, rebuild the index |
||||
if (recreate || !(hashDir.Exists && headerFile.Exists && indexFile.Exists && tailFile.Exists)) |
||||
{ |
||||
Logger.Trace("Creating new index for {Path}", hashDir); |
||||
hashDir.Create(); |
||||
|
||||
await using var headerStream = headerFile.Info.OpenWrite(); |
||||
await using var indexStream = indexFile.Info.OpenWrite(); |
||||
await using var tailStream = tailFile.Info.OpenWrite(); |
||||
|
||||
var builder = new IndexBuilder(headerStream, indexStream, tailStream); |
||||
|
||||
// Parse csv |
||||
var csvStream = path.Info.OpenRead(); |
||||
var parser = new TagCsvParser(csvStream); |
||||
|
||||
await foreach (var entry in parser.ParseAsync()) |
||||
{ |
||||
if (string.IsNullOrWhiteSpace(entry.Name)) |
||||
{ |
||||
continue; |
||||
} |
||||
// Add to index |
||||
builder.Add(entry.Name); |
||||
// Add to local dictionary |
||||
entries.Add(entry.Name, entry); |
||||
} |
||||
|
||||
builder.Build(); |
||||
} |
||||
|
||||
searcher = new InMemoryIndexSearcher(headerFile, indexFile, tailFile); |
||||
searcher.Init(); |
||||
} |
||||
|
||||
/// <inheritdoc /> |
||||
public IEnumerable<ICompletionData> GetCompletions(string searchTerm, int itemsCount, bool suggest) |
||||
{ |
||||
if (searcher is null) |
||||
{ |
||||
throw new InvalidOperationException("Index is not loaded"); |
||||
} |
||||
|
||||
var searchOptions = new SearchOptions |
||||
{ |
||||
Term = searchTerm, |
||||
MaxItemCount = itemsCount, |
||||
SuggestWhenFoundStartsWith = suggest |
||||
}; |
||||
|
||||
var result = searcher.Search(searchOptions); |
||||
|
||||
// No results |
||||
if (result.ResultType == TrieNodeSearchResultType.NotFound) |
||||
{ |
||||
Logger.Trace("No results for {Term}", searchTerm); |
||||
return Array.Empty<ICompletionData>(); |
||||
} |
||||
|
||||
Logger.Trace("Got {Count} results for {Term}", result.Items.Length, searchTerm); |
||||
|
||||
// Get entry for each result |
||||
var completions = new List<ICompletionData>(); |
||||
foreach (var item in result.Items) |
||||
{ |
||||
if (entries.TryGetValue(item, out var entry)) |
||||
{ |
||||
var entryType = TagTypeExtensions.FromE621(entry.Type.GetValueOrDefault(-1)); |
||||
completions.Add(new TagCompletionData(entry.Name!, entryType)); |
||||
} |
||||
} |
||||
|
||||
return completions; |
||||
} |
||||
} |
||||
@ -0,0 +1,28 @@
|
||||
using System.Collections.Generic; |
||||
using System.Threading.Tasks; |
||||
using StabilityMatrix.Avalonia.Controls.CodeCompletion; |
||||
using StabilityMatrix.Core.Models.FileInterfaces; |
||||
|
||||
namespace StabilityMatrix.Avalonia.Models.TagCompletion; |
||||
|
||||
public interface ICompletionProvider |
||||
{ |
||||
/// <summary> |
||||
/// Whether the completion provider is loaded. |
||||
/// </summary> |
||||
bool IsLoaded { get; } |
||||
|
||||
/// <summary> |
||||
/// Load the completion provider from a file. |
||||
/// </summary> |
||||
Task LoadFromFile(FilePath path, bool recreate = false); |
||||
|
||||
/// <summary> |
||||
/// Returns a list of completion items for the given text. |
||||
/// </summary> |
||||
public IEnumerable<ICompletionData> GetCompletions( |
||||
string searchTerm, |
||||
int itemsCount, |
||||
bool suggest |
||||
); |
||||
} |
||||
Loading…
Reference in new issue