pub mod data;
pub mod loader;
use serde::{Deserialize, Serialize};
pub use data::build_default_corpus;
pub use loader::load_and_verify_corpus;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum EvalCategory {
NlToSymbol,
ExactPartialIdentifier,
ConceptToImpl,
ErrorLogToOrigin,
CallerCalleeDataFlow,
InterfaceToImpl,
ConfigDocToCode,
SimilarAlgoDifferentName,
SameNameDifferentBehavior,
ChangedDeletedFreshness,
LargeGeneratedDistractors,
MultiLanguage,
CrossLanguage,
HardNegatives,
}
impl EvalCategory {
pub fn as_str(&self) -> &'static str {
match self {
Self::NlToSymbol => "nl_to_symbol",
Self::ExactPartialIdentifier => "exact_partial_identifier",
Self::ConceptToImpl => "concept_to_impl",
Self::ErrorLogToOrigin => "error_log_to_origin",
Self::CallerCalleeDataFlow => "caller_callee_data_flow",
Self::InterfaceToImpl => "interface_to_impl",
Self::ConfigDocToCode => "config_doc_to_code",
Self::SimilarAlgoDifferentName => "similar_algo_different_name",
Self::SameNameDifferentBehavior => "same_name_different_behavior",
Self::ChangedDeletedFreshness => "changed_deleted_freshness",
Self::LargeGeneratedDistractors => "large_generated_distractors",
Self::MultiLanguage => "multi_language",
Self::CrossLanguage => "cross_language",
Self::HardNegatives => "hard_negatives",
}
}
pub fn description(&self) -> &'static str {
match self {
Self::NlToSymbol => "Natural-language intent to symbol",
Self::ExactPartialIdentifier => "Exact and partial identifier",
Self::ConceptToImpl => "Concept to implementation",
Self::ErrorLogToOrigin => "Error/log string to origin",
Self::CallerCalleeDataFlow => "Caller/callee/data-flow retrieval",
Self::InterfaceToImpl => "Interface to implementation",
Self::ConfigDocToCode => "Configuration/docs to code",
Self::SimilarAlgoDifferentName => "Similar algorithms with different names",
Self::SameNameDifferentBehavior => "Same names with different behavior",
Self::ChangedDeletedFreshness => "Changed/deleted files and freshness",
Self::LargeGeneratedDistractors => "Large files and generated-looking distractors",
Self::MultiLanguage => "Multi-language (Rust/TS/Python/Go/Java/C/C++)",
Self::CrossLanguage => "Cross-language conceptual queries",
Self::HardNegatives => "Hard negatives sharing syntax, names, or comments",
}
}
pub fn all() -> [EvalCategory; 14] {
[
Self::NlToSymbol,
Self::ExactPartialIdentifier,
Self::ConceptToImpl,
Self::ErrorLogToOrigin,
Self::CallerCalleeDataFlow,
Self::InterfaceToImpl,
Self::ConfigDocToCode,
Self::SimilarAlgoDifferentName,
Self::SameNameDifferentBehavior,
Self::ChangedDeletedFreshness,
Self::LargeGeneratedDistractors,
Self::MultiLanguage,
Self::CrossLanguage,
Self::HardNegatives,
]
}
pub fn from_str(s: &str) -> Option<Self> {
match s {
"nl_to_symbol" => Some(Self::NlToSymbol),
"exact_partial_identifier" => Some(Self::ExactPartialIdentifier),
"concept_to_impl" => Some(Self::ConceptToImpl),
"error_log_to_origin" => Some(Self::ErrorLogToOrigin),
"caller_callee_data_flow" => Some(Self::CallerCalleeDataFlow),
"interface_to_impl" => Some(Self::InterfaceToImpl),
"config_doc_to_code" => Some(Self::ConfigDocToCode),
"similar_algo_different_name" => Some(Self::SimilarAlgoDifferentName),
"same_name_different_behavior" => Some(Self::SameNameDifferentBehavior),
"changed_deleted_freshness" => Some(Self::ChangedDeletedFreshness),
"large_generated_distractors" => Some(Self::LargeGeneratedDistractors),
"multi_language" => Some(Self::MultiLanguage),
"cross_language" => Some(Self::CrossLanguage),
"hard_negatives" => Some(Self::HardNegatives),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Split {
Train,
Eval,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Language {
Rust,
TypeScript,
JavaScript,
Python,
Go,
Java,
C,
Cpp,
}
impl Language {
pub fn as_str(&self) -> &'static str {
match self {
Self::Rust => "rust",
Self::TypeScript => "typescript",
Self::JavaScript => "javascript",
Self::Python => "python",
Self::Go => "go",
Self::Java => "java",
Self::C => "c",
Self::Cpp => "cpp",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorpusCase {
pub id: String,
pub query: String,
pub category: EvalCategory,
pub split: Split,
pub relevant_symbols: Vec<String>,
pub relevant_files: Vec<String>,
pub languages: Vec<Language>,
pub expected_position: Option<(usize, usize)>,
pub notes: Option<String>,
pub hard_negatives: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Corpus {
pub cases: Vec<CorpusCase>,
}
impl Corpus {
pub fn new() -> Self {
Self { cases: Vec::new() }
}
pub fn add_case(&mut self, case: CorpusCase) {
self.cases.push(case);
}
pub fn all(&self) -> &[CorpusCase] {
&self.cases
}
pub fn eval_cases(&self) -> impl Iterator<Item = &CorpusCase> {
self.cases.iter().filter(|c| c.split == Split::Eval)
}
pub fn train_cases(&self) -> impl Iterator<Item = &CorpusCase> {
self.cases.iter().filter(|c| c.split == Split::Train)
}
pub fn cases_for_category(&self, cat: EvalCategory) -> impl Iterator<Item = &CorpusCase> {
self.cases.iter().filter(move |c| c.category == cat)
}
pub fn verify_category_coverage(&self) -> Result<(), CorpusVerificationError> {
let mut missing: Vec<String> = Vec::new();
let mut insufficient: Vec<String> = Vec::new();
const MIN_CASES_PER_CATEGORY: usize = 2;
for cat in EvalCategory::all() {
let count = self.cases_for_category(cat).count();
if count == 0 {
missing.push(format!("{}: 0 cases", cat.as_str()));
} else if count < MIN_CASES_PER_CATEGORY {
insufficient.push(format!(
"{}: {} cases (need {})",
cat.as_str(),
count,
MIN_CASES_PER_CATEGORY
));
}
}
if !missing.is_empty() {
return Err(CorpusVerificationError::MissingCategories(missing));
}
if !insufficient.is_empty() {
return Err(CorpusVerificationError::InsufficientCases(insufficient));
}
Ok(())
}
pub fn verify_split_integrity(&self) -> Result<(), CorpusVerificationError> {
let mut empty_labels: Vec<String> = Vec::new();
let mut seen_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
let mut duplicate_ids: Vec<String> = Vec::new();
for case in &self.cases {
if case.relevant_symbols.is_empty() && case.relevant_files.is_empty() {
empty_labels.push(case.id.clone());
}
if !seen_ids.insert(case.id.clone()) {
duplicate_ids.push(case.id.clone());
}
}
let mut eval_missing: Vec<String> = Vec::new();
for cat in EvalCategory::all() {
let has_eval = self.cases_for_category(cat).any(|c| c.split == Split::Eval);
if !has_eval {
eval_missing.push(cat.as_str().to_string());
}
}
if !empty_labels.is_empty() {
return Err(CorpusVerificationError::EmptyLabels(empty_labels));
}
if !duplicate_ids.is_empty() {
return Err(CorpusVerificationError::DuplicateCaseIds(duplicate_ids));
}
if !eval_missing.is_empty() {
return Err(CorpusVerificationError::NoEvalCasesForCategory(
eval_missing,
));
}
Ok(())
}
pub fn len(&self) -> usize {
self.cases.len()
}
pub fn eval_len(&self) -> usize {
self.eval_cases().count()
}
pub fn train_len(&self) -> usize {
self.train_cases().count()
}
pub fn is_empty(&self) -> bool {
self.cases.is_empty()
}
}
impl Default for Corpus {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub enum CorpusVerificationError {
MissingCategories(Vec<String>),
InsufficientCases(Vec<String>),
EmptyLabels(Vec<String>),
DuplicateCaseIds(Vec<String>),
NoEvalCasesForCategory(Vec<String>),
}
impl std::fmt::Display for CorpusVerificationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingCategories(cats) => {
write!(f, "Missing categories: {}", cats.join(", "))
}
Self::InsufficientCases(cats) => {
write!(f, "Insufficient cases in: {}", cats.join(", "))
}
Self::EmptyLabels(ids) => write!(f, "Cases with empty labels: {}", ids.join(", ")),
Self::DuplicateCaseIds(ids) => {
write!(f, "Duplicate case IDs: {}", ids.join(", "))
}
Self::NoEvalCasesForCategory(cats) => {
write!(f, "No eval cases for categories: {}", cats.join(", "))
}
}
}
}
impl std::error::Error for CorpusVerificationError {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_category_all_returns_14() {
let all = EvalCategory::all();
assert_eq!(all.len(), 14);
}
#[test]
fn test_category_roundtrip_strings() {
for cat in EvalCategory::all() {
let s = cat.as_str();
let parsed = EvalCategory::from_str(s).expect("should parse");
assert_eq!(parsed, cat);
}
}
#[test]
fn test_category_descriptions_nonempty() {
for cat in EvalCategory::all() {
assert!(!cat.description().is_empty());
}
}
#[test]
fn test_corpus_add_and_count() {
let mut corpus = Corpus::new();
assert_eq!(corpus.len(), 0);
assert!(corpus.is_empty());
corpus.add_case(CorpusCase {
id: "test-1".to_string(),
query: "find function".to_string(),
category: EvalCategory::NlToSymbol,
split: Split::Eval,
relevant_symbols: vec!["my_func".to_string()],
relevant_files: vec!["src/main.rs".to_string()],
languages: vec![Language::Rust],
expected_position: None,
notes: None,
hard_negatives: Vec::new(),
});
assert_eq!(corpus.len(), 1);
assert!(!corpus.is_empty());
assert_eq!(corpus.eval_len(), 1);
assert_eq!(corpus.train_len(), 0);
}
#[test]
fn test_corpus_split_filtering() {
let mut corpus = Corpus::new();
for i in 0..4 {
corpus.add_case(CorpusCase {
id: format!("case-{i}"),
query: format!("query-{i}"),
category: EvalCategory::NlToSymbol,
split: if i % 2 == 0 {
Split::Eval
} else {
Split::Train
},
relevant_symbols: vec!["sym".to_string()],
relevant_files: vec!["file.rs".to_string()],
languages: vec![Language::Rust],
expected_position: None,
notes: None,
hard_negatives: Vec::new(),
});
}
assert_eq!(corpus.eval_len(), 2);
assert_eq!(corpus.train_len(), 2);
}
#[test]
fn test_corpus_case_for_category() {
let mut corpus = Corpus::new();
corpus.add_case(CorpusCase {
id: "1".to_string(),
query: "q1".to_string(),
category: EvalCategory::NlToSymbol,
split: Split::Eval,
relevant_symbols: vec!["s".to_string()],
relevant_files: vec!["f.rs".to_string()],
languages: vec![Language::Rust],
expected_position: None,
notes: None,
hard_negatives: Vec::new(),
});
corpus.add_case(CorpusCase {
id: "2".to_string(),
query: "q2".to_string(),
category: EvalCategory::ConceptToImpl,
split: Split::Eval,
relevant_symbols: vec!["s".to_string()],
relevant_files: vec!["f.rs".to_string()],
languages: vec![Language::Rust],
expected_position: None,
notes: None,
hard_negatives: Vec::new(),
});
assert_eq!(
corpus.cases_for_category(EvalCategory::NlToSymbol).count(),
1
);
assert_eq!(
corpus
.cases_for_category(EvalCategory::ConceptToImpl)
.count(),
1
);
assert_eq!(
corpus
.cases_for_category(EvalCategory::HardNegatives)
.count(),
0
);
}
}