/// Path to lexicon file

#[arg(short = 'l', long, default_value = "./index/storage_vbyte.txt")]

lexicon: String,

/// Path to document info file

#[arg(short = 'd', long, default_value = "./index/docs.txt")]

docs: String,

/// Index IDs file

#[arg(short = 'i', long, default_value = "./index/index.vbyte")]

ids: String,

/// Index freqs file

#[arg(short = 'f', long, default_value = "./index/freqs.vbyte")]

freqs: String,

/// Index format

#[arg(short = 't', long, default_value = "vbyte")]

format: String,

/// Collection file (for reading document content)

#[arg(short = 'c', long, default_value = "./data/collection.tsv")]

collection: String,

/// Web server port (0 means CLI mode)

#[arg(short = 'p', long, default_value = "0")]

port: u16,

/// Snippet length (number of characters)

#[arg(short = 's', long, default_value = "200")]

snippet_len: usize,

rank: usize,

doc_id: u32,

score: f64,

url: String,

snippet: Option<String>,

#[serde(skip_serializing_if = "Option::is_none")]

term_frequencies: Option<HashMap<String, u32>>, // Term frequency for each query term in the document

fn new(collection_path: String) -> io::Result<Self> {

let collection_file = File::open(&collection_path)?;

Ok(Self {

collection_file,

doc_contents: HashMap::new(),

})

}

/// Get document content (using seek for direct positioning)

fn get_content(&mut self, _doc_id: u32, doc_info: &DocInfo) -> io::Result<String> {

// Check cache first

if let Some(content) = self.doc_contents.get(&doc_info.doc_id) {

return Ok(content.clone());

}

// Check if there's a valid offset

if doc_info.offset == 0 && doc_info.len == 0 {

// Old format data, snippet not supported

return Ok(String::new());

}

// Use seek to jump directly to document position - O(1) time complexity!

self.collection_file.seek(SeekFrom::Start(doc_info.offset))?;

// Read content of specified length

let mut buffer = vec![0u8; doc_info.len as usize];

self.collection_file.read_exact(&mut buffer)?;

// Convert to string

let content = String::from_utf8_lossy(&buffer).to_string();

// Cache (limit cache size)

if self.doc_contents.len() < 1000 {

self.doc_contents.insert(doc_info.doc_id, content.clone());

}

Ok(content)

}

/// Generate query-related snippet

fn generate_snippet(&mut self, doc_id: u32, doc_info: &DocInfo, query_terms: &[String], max_len: usize) -> io::Result<String> {

let doc = self.get_content(doc_id, doc_info)?;

if doc.is_empty() {

return Ok(String::new());

}

// split content using \t

let content_parts: Vec<&str> = doc.split('\t').collect();

let content = match content_parts.len() {

2 => content_parts[1].to_string(),

3 => content_parts[2].to_string(),

4 => content_parts[2..].join(" "),

_ => doc, // fallback to full line

};

let content_lower = content.to_lowercase();

// Strategy: Find the first occurrence of a query term (consistent with C++ version)

let mut q_pos = None;

for term in query_terms {

if let Some(pos) = content_lower.find(term) {

q_pos = Some(pos);

break;

}

}

let center_pos = q_pos.unwrap_or(0);

// Calculate window range: centered around the found term

let start = center_pos.saturating_sub(max_len / 2);

let end = (center_pos + max_len / 2).min(content.len());

// UTF-8 boundary handling: ensure start is on a character boundary

let mut safe_start = start;

while safe_start > 0 && !content.is_char_boundary(safe_start) {

safe_start -= 1;

}

// UTF-8 boundary handling: ensure end is on a character boundary

let mut safe_end = end;

while safe_end < content.len() && !content.is_char_boundary(safe_end) {

safe_end += 1;

}

safe_end = safe_end.min(content.len());

let mut snippet = content[safe_start..safe_end].to_string();

// Add ellipsis

if safe_start > 0 {

snippet = format!("...{}", snippet);

}

if safe_end < content.len() {

snippet = format!("{}...", snippet);

}

Ok(snippet)

}

capacity: usize,

cache: HashMap<String, Arc<(InvertedList, u32)>>, // (list, doc_count) wrapped in Arc

order: Vec<String>,

fn new(capacity: usize) -> Self {

Self {

capacity,

cache: HashMap::new(),

order: Vec::new(),

}

}

fn get(&mut self, key: &str) -> Option<Arc<(InvertedList, u32)>> {

if let Some(value) = self.cache.get(key) {

// Update LRU order

if let Some(pos) = self.order.iter().position(|k| k == key) {

self.order.remove(pos);

}

self.order.push(key.to_string());

return Some(Arc::clone(value)); // Only clones the Arc pointer, not the data!

}

None

}

fn put(&mut self, key: String, value: (InvertedList, u32)) {

// If it already exists, update it

if self.cache.contains_key(&key) {

if let Some(pos) = self.order.iter().position(|k| k == &key) {

self.order.remove(pos);

}

self.cache.insert(key.clone(), Arc::new(value));

self.order.push(key);

return;

}

// If cache is full, remove the oldest

if self.cache.len() >= self.capacity {

if let Some(oldest) = self.order.first().cloned() {

self.cache.remove(&oldest);

self.order.remove(0);

}

}

self.cache.insert(key.clone(), Arc::new(value));

self.order.push(key);

}

fn len(&self) -> usize {

self.cache.len()

}

lexicon: HashMap<String, LexiconEntry>,

doc_infos: Vec<DocInfo>,

total_docs: u32,

avg_doc_len: f64,

ids_file: File,

freqs_file: File,

format: IndexFormat,

doc_cache: DocumentCache,

snippet_len: usize,

list_cache: InvertedListCache, // InvertedList cache

fn load(

lexicon_path: &str,

docs_path: &str,

ids_path: &str,

freqs_path: &str,

collection_path: &str,

format: IndexFormat,

snippet_len: usize,

) -> io::Result<Self> {

println!("{}", "Loading index...".bright_green().bold());

print!(" 📖 Loading lexicon... ");

io::stdout().flush()?;

let lexicon = Self::load_lexicon(lexicon_path)?;

println!("{} ({} terms)", "✓".bright_green(), lexicon.len().to_string().bright_white());

print!(" 📄 Loading document info... ");

io::stdout().flush()?;

let doc_infos = Self::load_doc_infos(docs_path)?;

let total_docs = doc_infos.len() as u32;

let avg_doc_len = doc_infos.iter().map(|d| d.term_count as f64).sum::<f64>() / total_docs as f64;

println!(

"{} ({} documents, avg length: {:.2})",

"✓".bright_green(),

total_docs.to_string().bright_white(),

avg_doc_len.to_string().bright_white()

);

print!(" 🗂️ Opening index files... ");

io::stdout().flush()?;

let ids_file = File::open(ids_path)?;

let freqs_file = File::open(freqs_path)?;

println!("{}", "✓".bright_green());

println!();

Ok(Self {

lexicon,

doc_infos,

total_docs,

avg_doc_len,

ids_file,

freqs_file,

format,

doc_cache: DocumentCache::new(collection_path.to_string())?,

snippet_len,

list_cache: InvertedListCache::new(1000), // Cache 1000 InvertedLists

})

}

fn load_lexicon(path: &str) -> io::Result<HashMap<String, LexiconEntry>> {

let file = File::open(path)?;

let reader = BufReader::new(file);

let mut lexicon = HashMap::new();

for line in reader.lines() {

let line = line?;

let parts: Vec<&str> = line.split_whitespace().collect();

if parts.len() >= 4 {

let term = parts[0].to_string();

let ids_offset = parts[1].parse().unwrap_or(0);

let freqs_offset = parts[2].parse().unwrap_or(0);

let doc_count = parts[3].parse().unwrap_or(0);

lexicon.insert(

term.clone(),

LexiconEntry {

term,

ids_offset,

freqs_offset,

doc_count,

},

);

}

}

Ok(lexicon)

}

fn load_doc_infos(path: &str) -> io::Result<Vec<DocInfo>> {

let file = File::open(path)?;

let reader = BufReader::new(file);

let mut doc_infos = Vec::new();

for line in reader.lines() {

let line = line?;

let parts: Vec<&str> = line.split('\t').collect();

if parts.len() >= 5 {

// New format: doc_id\turl\tterm_count\toffset\tlen

doc_infos.push(DocInfo {

doc_id: parts[0].parse().unwrap_or(0),

url: parts[1].to_string(),

term_count: parts[2].parse().unwrap_or(0),

offset: parts[3].parse().unwrap_or(0),

len: parts[4].parse().unwrap_or(0),

});

} else if parts.len() >= 3 {

// Old format (backward compatibility)

doc_infos.push(DocInfo {

doc_id: parts[0].parse().unwrap_or(0),

url: parts[1].to_string(),

term_count: parts[2].parse().unwrap_or(0),

offset: 0,

len: 0,

});

}

}

Ok(doc_infos)

}

fn search(&mut self, query: &str, mode: QueryMode, top_k: usize, with_snippet: bool) -> io::Result<(Vec<SearchResult>, f64, f64, f64, f64, usize)> {

let query_terms = process_text(query);

if query_terms.is_empty() {

return Ok((Vec::new(), 0.0, 0.0, 0.0, 0.0, 0));

}

let mut cache_hits = 0;

let mut inverted_lists: Vec<(String, Arc<(InvertedList, u32)>)> = Vec::new();

// Load inverted lists, using cache (timed)

let load_start = std::time::Instant::now();

for term in &query_terms {

if let Some(lexicon_entry) = self.lexicon.get(term) {

// Try to get from cache first

if let Some(list_arc) = self.list_cache.get(term) {

cache_hits += 1;

inverted_lists.push((term.clone(), list_arc));

} else {

// Cache miss, read from file

let mut ids_reader = self.ids_file.try_clone()?;

let mut freqs_reader = self.freqs_file.try_clone()?;

let list = InvertedList::read_from_merged_index(&mut ids_reader, &mut freqs_reader, lexicon_entry, self.format)?;

// Add to cache (wrapped in Arc)

self.list_cache.put(term.clone(), (list.clone(), lexicon_entry.doc_count));

// Get the Arc back from cache

if let Some(list_arc) = self.list_cache.get(term) {

inverted_lists.push((term.clone(), list_arc));

}

}

}

}

let load_time = load_start.elapsed().as_secs_f64() * 1000.0;

if inverted_lists.is_empty() {

return Ok((Vec::new(), 0.0, 0.0, 0.0, 0.0, 0));

}

// Calculate BM25 scores (collect term frequencies at the same time)

let search_start = std::time::Instant::now();

let (scores, term_frequencies_map) = match mode {

QueryMode::Conjunctive => self.conjunctive_search(&inverted_lists),

QueryMode::Disjunctive => self.disjunctive_search(&inverted_lists),

};

let search_time = search_start.elapsed().as_secs_f64() * 1000.0;

// Generate results and sort (timed)

let sort_start = std::time::Instant::now();

let mut results: Vec<_> = scores

.into_iter()

.map(|(doc_id, score)| {

let doc_info = self.doc_infos.get(doc_id as usize);

let url = doc_info.map(|d| d.url.clone()).unwrap_or_default();

let term_frequencies = term_frequencies_map.get(&doc_id).cloned();

SearchResult {

rank: 0,

doc_id,

score,

url,

snippet: None,

term_frequencies,

}

})

.collect();

results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap().then_with(|| a.doc_id.cmp(&b.doc_id)));

results.truncate(top_k);

// Update rankings

for (i, result) in results.iter_mut().enumerate() {

result.rank = i + 1;

}

let sort_time = sort_start.elapsed().as_secs_f64() * 1000.0;

// Generate snippets (timed separately)

let snippet_start = std::time::Instant::now();

if with_snippet {

for result in results.iter_mut() {

let doc_info = self.doc_infos.get(result.doc_id as usize);

if let Some(doc_info) = doc_info {

result.snippet = self.doc_cache.generate_snippet(result.doc_id, doc_info, &query_terms, self.snippet_len).ok();

}

}

}

let snippet_time = snippet_start.elapsed().as_secs_f64() * 1000.0;

Ok((results, load_time, search_time, sort_time, snippet_time, cache_hits))

}

fn conjunctive_search(&self, lists: &[(String, Arc<(InvertedList, u32)>)]) -> (HashMap<u32, f64>, HashMap<u32, HashMap<String, u32>>) {

let mut scores = HashMap::new();

let mut term_frequencies_map: HashMap<u32, HashMap<String, u32>> = HashMap::new();

if lists.is_empty() {

return (scores, term_frequencies_map);

}

// Sort by posting list length (shorter first, to speed up intersection calculation)

let mut sorted_lists: Vec<&(String, Arc<(InvertedList, u32)>)> = lists.iter().collect();

sorted_lists.sort_by_key(|(_, list_arc)| list_arc.0.doc_ids.len());

// Use C++ approach: compute intersection first, then calculate scores

// 1. Start from the shortest list

let mut intersection = sorted_lists[0].1.0.doc_ids.clone();

// 2. Intersect with other lists sequentially

for i in 1..sorted_lists.len() {

if intersection.is_empty() {

break; // Early exit

}

intersection = intersect_sorted(&intersection, &sorted_lists[i].1.0.doc_ids);

}

// 3. Calculate BM25 scores for each document in the intersection, and collect term frequencies

for doc_id in intersection {

let mut total_score = 0.0;

let mut doc_term_frequencies = HashMap::new();

for (term, list_arc) in lists {

let (list, doc_count) = &**list_arc;

if let Some(freq) = list.get_freq(doc_id) {

let score = self.bm25_score(freq, *doc_count, doc_id);

total_score += score;

doc_term_frequencies.insert(term.clone(), freq);

}

}

scores.insert(doc_id, total_score);

term_frequencies_map.insert(doc_id, doc_term_frequencies);

}

(scores, term_frequencies_map)

}

fn disjunctive_search(&self, lists: &[(String, Arc<(InvertedList, u32)>)]) -> (HashMap<u32, f64>, HashMap<u32, HashMap<String, u32>>) {

let mut scores: HashMap<u32, f64> = HashMap::new();

let mut term_frequencies_map: HashMap<u32, HashMap<String, u32>> = HashMap::new();

for (term, list_arc) in lists {

let (list, doc_count) = &**list_arc;

for i in 0..list.doc_ids.len() {

let doc_id = list.doc_ids[i];

let freq = list.freqs[i];

let score = self.bm25_score(freq, *doc_count, doc_id);

*scores.entry(doc_id).or_insert(0.0) += score;

// Collect term frequencies

term_frequencies_map.entry(doc_id).or_insert_with(HashMap::new).insert(term.clone(), freq);

}

}

(scores, term_frequencies_map)

}

fn bm25_score(&self, freq: u32, doc_freq: u32, doc_id: u32) -> f64 {

let doc_len = self.doc_infos.get(doc_id as usize).map(|d| d.term_count as f64).unwrap_or(self.avg_doc_len);

let idf = ((self.total_docs as f64 - doc_freq as f64 + 0.5) / (doc_freq as f64 + 0.5)).ln();

let tf_component = (freq as f64 * (K + 1.0)) / (freq as f64 + K * (1.0 - B + B * doc_len / self.avg_doc_len));

idf * tf_component

}

println!("{}", "Commands:".bright_green().bold());

println!(" /mode conj - Switch to Conjunctive (AND) mode");

println!(" /mode disj - Switch to Disjunctive (OR) mode");

println!(" /snippet on - Enable snippet generation");

println!(" /snippet off - Disable snippet generation");

println!(" /quit - Exit the program");

println!(" <query> - Search for documents");

println!();

let mut current_mode = QueryMode::Disjunctive;

let mut show_snippet = false;

let stdin = io::stdin();

loop {

let mode_str = match current_mode {

QueryMode::Conjunctive => "AND".bright_yellow(),

QueryMode::Disjunctive => "OR".bright_green(),

};

print!("Query [{}{}]> ", mode_str, if show_snippet { " +snippet".bright_cyan() } else { "".normal() });

io::stdout().flush()?;

let mut query = String::new();

stdin.read_line(&mut query)?;

let query = query.trim();

if query.is_empty() {

continue;

}

if query.starts_with("/") {

match query {

"/quit" | "/exit" | "/q" => {

println!("Goodbye! 👋");

break;

}

"/mode conj" | "/mode conjunctive" | "/mode and" => {

current_mode = QueryMode::Conjunctive;

println!("{} Switched to Conjunctive (AND) mode", "✓".bright_green());

}

"/mode disj" | "/mode disjunctive" | "/mode or" => {

current_mode = QueryMode::Disjunctive;

println!("{} Switched to Disjunctive (OR) mode", "✓".bright_green());

}

"/snippet on" => {

show_snippet = true;

println!("{} Snippet generation enabled", "✓".bright_green());

}

"/snippet off" => {

show_snippet = false;

println!("{} Snippet generation disabled", "✓".bright_green());

}

_ => {

println!("{} Unknown command: {}", "✗".bright_red(), query);

}

}

continue;

}

let start = std::time::Instant::now();

let (search_outcome, cache_len) = {

let mut processor = processor_state.blocking_lock();

let outcome = processor.search(query, current_mode, 10, show_snippet);

let cache_len = processor.list_cache.len();

(outcome, cache_len)

};

let (results, load_time, search_time, sort_time, snippet_time, cache_hits) = match search_outcome {

Ok(data) => data,

Err(err) => {

eprintln!("{} Failed to execute search: {}", "✗".bright_red(), err);

continue;

}

};

let total_time = start.elapsed().as_secs_f64() * 1000.0;

println!();

println!(

"{} Found {} results in {} (load: {}, search: {}, sort: {}, snippet: {}, cache: {}/{})",

"→".bright_green(),

results.len().to_string().bright_white().bold(),

format!("{:.2}ms", total_time).bright_yellow(),

format!("{:.2}ms", load_time).bright_cyan(),

format!("{:.2}ms", search_time).bright_cyan(),

format!("{:.2}ms", sort_time).bright_cyan(),

format!("{:.2}ms", snippet_time).bright_cyan(),

cache_hits.to_string().bright_magenta(),

cache_len.to_string().bright_magenta()

);

println!("{}", "─".repeat(80).bright_black());

for result in results {

println!(

"{:>2}. {} Doc#{}",

result.rank.to_string().bright_cyan(),

format!("[{:.4}]", result.score).bright_yellow(),

result.doc_id.to_string().bright_white()

);

// Display URL (if available)

if !result.url.is_empty() {

println!(" 🔗 {}", result.url.bright_blue());

}

// Display snippet (if available)

if let Some(snippet) = result.snippet {

println!(" 📄 {}", snippet.dimmed());

}

println!(); // Empty line to separate results

}

println!("{}", "─".repeat(80).bright_black());

println!();

}

Ok(())

q: String,

#[serde(default = "default_mode")]

mode: String,

#[serde(default = "default_top_k")]

top_k: usize,

#[serde(default)]

snippet: bool,

query: String,

mode: String,

results: Vec<SearchResult>,

count: usize,

time_ms: f64,

load_time_ms: f64,

search_time_ms: f64,

sort_time_ms: f64,

snippet_time_ms: f64,

cache_hits: usize,

let mode = if params.mode == "conjunctive" || params.mode == "and" {

QueryMode::Conjunctive

} else {

QueryMode::Disjunctive

};

let start = std::time::Instant::now();

let mut processor = state.lock().await;

let (results, load_time, search_time, sort_time, snippet_time, cache_hits) = processor

.search(&params.q, mode, params.top_k, params.snippet)

.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

let elapsed = start.elapsed();

Ok(Json(SearchResponse {

query: params.q,

mode: params.mode,

count: results.len(),

results,

time_ms: elapsed.as_secs_f64() * 1000.0,

load_time_ms: load_time,

search_time_ms: search_time,

sort_time_ms: sort_time,

snippet_time_ms: snippet_time,

cache_hits,

}))

let static_service = get_service(ServeDir::new("static").append_index_html_on_directories(true));

let app = Router::new()

.route("/api/search", get(search_handler))

.nest_service("/", static_service)

.layer(CorsLayer::permissive())

.with_state(state);

let addr = format!("0.0.0.0:{}", port);

let listener = tokio::net::TcpListener::bind(&addr)

.await

.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

axum::serve(listener, app).await.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

Ok(())

let args = Args::parse();

let format = IndexFormat::from_str(&args.format).unwrap_or(IndexFormat::VarByte);

let processor = QueryProcessor::load(&args.lexicon, &args.docs, &args.ids, &args.freqs, &args.collection, format, args.snippet_len)?;

let shared_processor = Arc::new(tokio::sync::Mutex::new(processor));

if args.port > 0 {

let server_state = shared_processor.clone();

let port = args.port;

std::thread::spawn(move || {

let runtime = tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime");

if let Err(err) = runtime.block_on(run_web_server(server_state, port)) {

eprintln!("{} Web server exited: {}", "✗".bright_red(), err);

}

});

println!("{} Web server running on http://localhost:{}", "✓".bright_green(), args.port);

run_cli_mode(shared_processor)

} else {

run_cli_mode(shared_processor)

}