use std::collections::HashMap;
use std::time::Instant;
use ryo_analysis::SymbolId;
use crate::id::SuggestId;
use crate::suggest::{compute_priority, OpportunityId, SafetyLevel, SuggestOpportunity};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PrecheckStatus {
#[default]
NotChecked,
Passed,
Failed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SuggestIndex(pub(crate) usize);
impl SuggestIndex {
pub fn as_usize(self) -> usize {
self.0
}
}
#[derive(Debug)]
pub struct StoredSuggestion {
pub opportunity: SuggestOpportunity,
pub suggest_idx: SuggestIndex,
pub safety: SafetyLevel,
pub priority: u8,
pub precheck_status: PrecheckStatus,
pub generation: u32,
pub closed: bool,
pub close_reason: Option<String>,
pub created_at: Instant,
pub closed_at: Option<Instant>,
}
impl StoredSuggestion {
pub fn new(
opportunity: SuggestOpportunity,
suggest_idx: SuggestIndex,
safety: SafetyLevel,
pattern_weight: f32,
) -> Self {
let priority = compute_priority(opportunity.confidence, safety, pattern_weight);
Self {
opportunity,
suggest_idx,
safety,
priority,
precheck_status: PrecheckStatus::NotChecked,
generation: 0,
closed: false,
close_reason: None,
created_at: Instant::now(),
closed_at: None,
}
}
pub fn new_with_priority(
opportunity: SuggestOpportunity,
suggest_idx: SuggestIndex,
safety: SafetyLevel,
priority: u8,
) -> Self {
Self {
opportunity,
suggest_idx,
safety,
priority,
precheck_status: PrecheckStatus::NotChecked,
generation: 0,
closed: false,
close_reason: None,
created_at: Instant::now(),
closed_at: None,
}
}
pub fn close(&mut self, reason: impl Into<String>) {
self.closed = true;
self.close_reason = Some(reason.into());
self.closed_at = Some(Instant::now());
}
pub fn bump_generation(&mut self) {
self.generation += 1;
}
}
pub struct SuggestStore {
suggestions: HashMap<u32, StoredSuggestion>,
symbol_to_suggests: HashMap<SymbolId, Vec<u32>>,
dedup_index: HashMap<(SuggestIndex, OpportunityId), u32>,
next_index: u32,
}
impl Default for SuggestStore {
fn default() -> Self {
Self::new()
}
}
impl SuggestStore {
pub fn new() -> Self {
Self {
suggestions: HashMap::new(),
symbol_to_suggests: HashMap::new(),
dedup_index: HashMap::new(),
next_index: 1, }
}
pub fn insert(&mut self, suggestion: StoredSuggestion) -> Option<SuggestId> {
let dedup_key = (suggestion.suggest_idx, suggestion.opportunity.id);
if let Some(&existing_idx) = self.dedup_index.get(&dedup_key) {
if self
.suggestions
.get(&existing_idx)
.is_some_and(|s| !s.closed)
{
return None; }
}
let index = self.next_index;
self.next_index += 1;
let generation = suggestion.generation;
let targets = suggestion.opportunity.targets.clone();
for target in targets {
self.symbol_to_suggests
.entry(target)
.or_default()
.push(index);
}
self.dedup_index.insert(dedup_key, index);
self.suggestions.insert(index, suggestion);
Some(SuggestId::new(index, generation))
}
pub fn get(&self, id: SuggestId) -> Option<&StoredSuggestion> {
self.suggestions
.get(&id.index())
.filter(|sug| sug.generation == id.generation() && !sug.closed)
}
pub fn get_mut(&mut self, id: SuggestId) -> Option<&mut StoredSuggestion> {
self.suggestions
.get_mut(&id.index())
.filter(|sug| sug.generation == id.generation() && !sug.closed)
}
pub fn remove_for_symbol(&mut self, symbol: &SymbolId) {
if let Some(indices) = self.symbol_to_suggests.remove(symbol) {
for index in indices {
if let Some(removed) = self.suggestions.remove(&index) {
let dedup_key = (removed.suggest_idx, removed.opportunity.id);
self.dedup_index.remove(&dedup_key);
}
}
}
}
pub fn invalidate_for_symbol(&mut self, symbol: &SymbolId) {
if let Some(indices) = self.symbol_to_suggests.get(symbol) {
for &index in indices {
if let Some(sug) = self.suggestions.get_mut(&index) {
sug.bump_generation();
}
}
}
}
pub fn is_valid(&self, id: SuggestId) -> bool {
self.get(id).is_some()
}
pub fn current_generation(&self, id: SuggestId) -> Option<u32> {
self.suggestions.get(&id.index()).map(|s| s.generation)
}
pub fn close(&mut self, id: SuggestId, reason: impl Into<String>) -> bool {
if let Some(sug) = self.get_mut(id) {
sug.close(reason);
true
} else {
false
}
}
pub fn iter(&self) -> impl Iterator<Item = (SuggestId, &StoredSuggestion)> {
self.suggestions
.iter()
.filter(|(_, s)| !s.closed)
.map(|(&index, sug)| (SuggestId::new(index, sug.generation), sug))
}
pub fn len(&self) -> usize {
self.suggestions.iter().filter(|(_, s)| !s.closed).count()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn total_count(&self) -> usize {
self.suggestions.len()
}
pub fn clear(&mut self) {
self.suggestions.clear();
self.symbol_to_suggests.clear();
self.dedup_index.clear();
self.next_index = 1;
}
}
#[derive(Debug, Clone)]
pub struct GcConfig {
pub max_closed_age: std::time::Duration,
pub max_suggestions: usize,
pub gc_interval: std::time::Duration,
}
impl Default for GcConfig {
fn default() -> Self {
Self {
max_closed_age: std::time::Duration::from_secs(300), max_suggestions: 1000,
gc_interval: std::time::Duration::from_secs(60), }
}
}
impl SuggestStore {
pub fn gc(&mut self, config: &GcConfig, valid_symbols: &impl Fn(&SymbolId) -> bool) {
let now = Instant::now();
let mut to_remove = Vec::new();
for (&index, sug) in self.suggestions.iter() {
if sug.closed {
if let Some(closed_at) = sug.closed_at {
if now.duration_since(closed_at) > config.max_closed_age {
to_remove.push(index);
continue;
}
}
}
let any_valid = sug.opportunity.targets.iter().any(valid_symbols);
if !any_valid {
to_remove.push(index);
}
}
for index in to_remove {
if let Some(sug) = self.suggestions.remove(&index) {
let dedup_key = (sug.suggest_idx, sug.opportunity.id);
self.dedup_index.remove(&dedup_key);
for target in &sug.opportunity.targets {
if let Some(indices) = self.symbol_to_suggests.get_mut(target) {
indices.retain(|&i| i != index);
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::suggest::{OpportunityContext, OpportunityId, SuggestLocation};
fn make_opportunity(id: u32, targets: Vec<SymbolId>) -> SuggestOpportunity {
SuggestOpportunity::new(
OpportunityId::new(id),
targets,
SuggestLocation::for_test("test.rs", "Test"),
"Test suggestion",
0.9,
OpportunityContext::Derive {
derive_name: "Default".into(),
missing_impls: vec![],
},
)
}
#[test]
fn test_suggest_id_format() {
let id = SuggestId::new(1, 0);
assert_eq!(id.to_string(), "S001g0");
let id2 = SuggestId::new(42, 3);
assert_eq!(id2.to_string(), "S042g3");
}
#[test]
fn test_suggest_id_parse() {
let id: SuggestId = "S001g0".parse().unwrap();
assert_eq!(id.index(), 1);
assert_eq!(id.generation(), 0);
let id2: SuggestId = "S042g3".parse().unwrap();
assert_eq!(id2.index(), 42);
assert_eq!(id2.generation(), 3);
}
#[test]
fn test_store_insert_and_get() {
let mut store = SuggestStore::new();
let sym = SymbolId::parse("100v1").unwrap();
let opp = make_opportunity(1, vec![sym]);
let sug = StoredSuggestion::new(opp, SuggestIndex(0), SafetyLevel::Auto, 1.0);
let id = store.insert(sug).expect("first insert should succeed");
assert_eq!(id.index(), 1);
assert_eq!(id.generation(), 0);
let retrieved = store.get(id).unwrap();
assert_eq!(retrieved.safety, SafetyLevel::Auto);
}
#[test]
fn test_store_invalidation() {
let mut store = SuggestStore::new();
let sym = SymbolId::parse("100v1").unwrap();
let opp = make_opportunity(1, vec![sym]);
let sug = StoredSuggestion::new(opp, SuggestIndex(0), SafetyLevel::Auto, 1.0);
let id = store.insert(sug).expect("insert should succeed");
assert!(store.is_valid(id));
store.invalidate_for_symbol(&sym);
assert!(!store.is_valid(id));
let new_gen = store.current_generation(id).unwrap();
assert_eq!(new_gen, 1);
}
#[test]
fn test_store_close() {
let mut store = SuggestStore::new();
let sym = SymbolId::parse("100v1").unwrap();
let opp = make_opportunity(1, vec![sym]);
let sug = StoredSuggestion::new(opp, SuggestIndex(0), SafetyLevel::Auto, 1.0);
let id = store.insert(sug).expect("insert should succeed");
assert!(store.is_valid(id));
assert_eq!(store.len(), 1);
store.close(id, "Applied");
assert!(!store.is_valid(id));
assert_eq!(store.len(), 0);
assert_eq!(store.total_count(), 1); }
#[test]
fn test_store_remove_for_symbol() {
let mut store = SuggestStore::new();
let sym1 = SymbolId::parse("100v1").unwrap();
let sym2 = SymbolId::parse("200v1").unwrap();
let opp1 = make_opportunity(1, vec![sym1]);
let opp2 = make_opportunity(2, vec![sym2]);
let sug1 = StoredSuggestion::new(opp1, SuggestIndex(0), SafetyLevel::Auto, 1.0);
let sug2 = StoredSuggestion::new(opp2, SuggestIndex(0), SafetyLevel::Auto, 1.0);
let id1 = store.insert(sug1).expect("insert sug1");
let id2 = store.insert(sug2).expect("insert sug2");
assert_eq!(store.len(), 2);
store.remove_for_symbol(&sym1);
assert!(!store.is_valid(id1));
assert!(store.is_valid(id2));
assert_eq!(store.len(), 1);
}
#[test]
fn test_store_iter() {
let mut store = SuggestStore::new();
let sym = SymbolId::parse("100v1").unwrap();
for i in 0..5 {
let opp = make_opportunity(i, vec![sym]);
let sug = StoredSuggestion::new(opp, SuggestIndex(0), SafetyLevel::Auto, 1.0);
store.insert(sug);
}
let ids: Vec<_> = store.iter().map(|(id, _)| id).collect();
assert_eq!(ids.len(), 5);
}
}