Local Search Engine in Rust

Question

I made a simple search engine using the xkcd API in Rust which turned out better than I'd hoped for!

I decided to use tf-idf as a way to rank results, which I feel like has some room for improvement. Any ideas regarding different algorithms would be fantastic.

The way I tried to organize the data was using separate structs, where the comic data from the xkcd endpoint is represented by Comic, the index of the comic's text is represented by ComicIndex, and all ComicIndexes are held in a Document (I wasn't sure what to name this, it's the struct which holds all of the comics and all of the word frequencies). As you might be able to tell from reading my code, I struggled with finding good, readable names for some variables.

Additionally, I had originally kept everything in one single file, but later split the code up into 3 files (4 if you count `utils.rs which just holds strings) for ease of viewing.

Anyways, here's the code:

main.rs

mod comic;
mod search;
mod utils;

use comic::{download_all, download_and_append_to_document};
use search::search;
use utils::{URL, INFO, RED, RESET, GREEN, CYAN, MAGENTA};
use clap::Parser;
use std::fs::File;

#[derive(Parser, Debug)]
#[command(name = "XKCD Search Engine")]
#[command(author = "Joey M.")]
#[command(version = "1.0.0")]
#[command(about = "Simple search engine for XKCD comics.", long_about = None)]
struct Args {
    /// Query to search XKCD comics for.
    #[clap(short, long)]
    search: Option<String>,

    /// Limit the amount of search results. Defaults to 10.
    #[clap(short, long)]
    limit: Option<usize>,

    /// Download and index a specific comic.
    #[clap(short, long)]
    download: Option<usize>,

    /// Download all comics.
    #[clap(short = 'D', long)]
    download_all: bool, 
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let args = Args::parse();
    let mut limit = 10;
    let mut query = String::new();

    if let Some(l) = args.limit {
        limit = l;
    }

    if let Some(s) = args.search {
        query = s;
    }

    let client = reqwest::Client::new();

    if args.download_all {
        let latest: comic::Comic = client.get(format!("{URL}/{INFO}")).send().await?.json().await?;

        println!("Downloading and indexing comics 1..={}.", latest.num);

        download_all(latest, &client).await?;
    } else if let Some(c) = args.download {
        println!("Downloading comic number {c}");
        download_and_append_to_document(c as u16, &client).await?;
    } else {
        let file = File::open(utils::DATA_PATH)?;
        let mut doc: comic::Document  = serde_json::from_reader(file)?;

        if query.is_empty() {
            println!("Search query is required.");
            return Ok(());
        }

        let res = search(&query, &mut doc);
        let top_20: Vec<(f32, comic::ComicIndex)> = res.into_iter().take(limit).collect();

        let padding = 4;
        let max_rank_len = 6 + padding;
        let max_name_len = top_20.iter().map(|(_, c)| c.comic.title.len()).max().unwrap_or(0).max(6) + padding;
        let max_num_len = top_20.iter().map(|(_, c)| c.comic.num.to_string().len()).max().unwrap_or(0).max(13) + padding;

        println!("┌{0:─<width_rank$}┬{0:─<width_rank$}┬{0:─<width_name$}┬{0:─<width_num$}┐", "", width_rank = max_rank_len, width_name = max_name_len, width_num = max_num_len);
        println!("│{MAGENTA} {:<width_rank$}{RESET}│{CYAN} {:<width_rank$}{RESET}│{GREEN} {:<width_name$}{RESET}│{RED} {:<width_num$}{RESET}│", "Index", "Rank", "Title", "Comic Number", width_rank = max_rank_len-1, width_name = max_name_len-1, width_num = max_num_len-1);
        println!("├{0:─<width_rank$}┼{0:─<width_rank$}┼{0:─<width_name$}┼{0:─<width_num$}┤", "", width_rank = max_rank_len, width_name = max_name_len, width_num = max_num_len);

        for (index, (rank, c)) in top_20.iter().enumerate() {
            println!(
                "│ {:<width_rank$}│ {:<width_rank$.4}│ {:<width_name$}│ {:<width_num$}│",
                index + 1,
                -rank,
                c.comic.title,
                c.comic.num,
                width_rank = max_rank_len-1,
                width_name = max_name_len-1,
                width_num = max_num_len-1,
            );
        }

        println!("└{0:─<width_rank$}┴{0:─<width_rank$}┴{0:─<width_name$}┴{0:─<width_num$}┘", "", width_rank = max_rank_len, width_name = max_name_len, width_num = max_num_len);
    }

    Ok(())
}

comic.rs

use std::collections::HashMap;
use std::fs::File;
use futures::{stream, StreamExt};
use tokio::sync::Mutex;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use anyhow;
use reqwest;
use crate::utils::{DATA_PATH, URL, INFO};

const PARALLEL_REQUESTS: usize = 8;

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Comic {
    pub month: String,
    pub num: u16,
    pub year: String,
    pub day: String,
    pub title: String,
    pub alt: String,
    pub img: String,
    pub transcript: String,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ComicIndex {
    pub title_freq: HashMap<String, usize>,
    pub title_len: usize,
    pub alt_freq: HashMap<String, usize>,
    pub alt_len: usize,
    pub comic: Comic,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ComicFrequency {
    pub alt_freq: HashMap<String, usize>,
    pub title_freq: HashMap<String, usize>,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Document {
    pub comics: Vec<ComicIndex>,
    pub frequency: ComicFrequency,
}

pub async fn download_all(latest: Comic, client: &reqwest::Client) -> anyhow::Result<()> {
    let c = Mutex::new(Vec::new());

    let frequencies = Arc::new(Mutex::new(ComicFrequency {
        title_freq: HashMap::new(),
        alt_freq: HashMap::new(),
    }));

    let indices = (1..=latest.num).collect::<Vec<_>>();

    let bodies = stream::iter(indices)
        .map(|index| {
            let url = format!("{URL}/{index}/{INFO}");
            let client = client.clone();
            let frequencies = Arc::clone(&frequencies);

            tokio::spawn(async move {
                let resp = client.get(&url).send().await?;
                let mut parsed = resp.json::<Comic>().await?;

                index_comic(&mut parsed, &mut *frequencies.lock().await).await
            })
        })
        .buffer_unordered(PARALLEL_REQUESTS);
    bodies
        .for_each(|x| async {
            match x {
                Ok(Ok(x)) => {
                    c.lock().await.push(x);
                },
                Ok(Err(e)) => eprintln!("Got a reqwest::Error: {}", e),
                Err(e) => eprintln!("Got a tokio::JoinError: {}", e),
            }
        })
        .await;

    let c = c.lock().await;
    let frequencies = frequencies.lock().await.clone();

    let doc = Document {
        comics: (*c).to_vec(),
        frequency: frequencies,
    };

    let file = File::create(DATA_PATH)?;
    let _ = serde_json::to_writer(file, &doc)?;

    Ok(())
}


pub async fn download_and_append_to_document(comic_number: u16, client: &reqwest::Client) -> anyhow::Result<()> { 
    let file = File::open(DATA_PATH)?;
    let mut doc: Document  = serde_json::from_reader(file)?;

    if !doc.comics.iter().any(|com| com.comic.num == comic_number) {
        let mut comic: Comic = client.get(format!("{URL}/{comic_number}/{INFO}")).send().await?.json().await?;

        let index: ComicIndex = index_comic(&mut comic, &mut doc.frequency).await?;

        doc.comics.push(index);

        let file = File::create(DATA_PATH)?;
        let _ = serde_json::to_writer(file, &doc)?;
    } else {
        println!("Comic already exists in data");
    }

    Ok(())
}

pub async fn index_comic(c: &mut Comic, df: &mut ComicFrequency) -> anyhow::Result<ComicIndex> {
    let mut alt_frequencies: HashMap<String, usize> = HashMap::new();
    let mut title_frequencies: HashMap<String, usize> = HashMap::new();

    // index alt text
    let mut alt = c.alt.to_lowercase();
    alt.retain(|c| !r#"(),".;:--'"#.contains(c)); // strip punctuation

    let mut alen = 0;
    for word in alt.split_whitespace() {
        alen += 1;
        *alt_frequencies.entry(word.to_owned()).or_default() += 1;
    }

    let mut title = c.title.to_lowercase();
    title.retain(|c| !r#"(),".;:-'"#.contains(c)); // strip punctuation
    let mut tlen = 0;

    for word in title.split_whitespace() {
        *title_frequencies.entry(word.to_owned()).or_default() += 1;
        tlen += 1;
    }

    // for each unique word in alt text and title, add it to df

    for text in title_frequencies.keys() {
        *df.title_freq.entry(text.to_owned()).or_default() += 1;
    }

    for word in alt_frequencies.keys() {
        *df.alt_freq.entry(word.to_owned()).or_default() += 1;
    }

    let i = ComicIndex {
        title_freq: title_frequencies,
        title_len: tlen,
        alt_freq: alt_frequencies,
        alt_len: alen,
        comic: c.clone()
    };

    Ok(i)
}

search.rs

use crate::comic::{ComicIndex, ComicFrequency, Document};
use crate::utils::Field;

pub fn search(q: &str, doc: &mut Document) -> Vec<(f32, ComicIndex)> {
    // preprocessing on query
    let mut query = q.to_lowercase();
    query.retain(|c| !r#"(),".;:'"#.contains(c)); // strip punctuation

    let mut result = Vec::new();

    for comic in &doc.comics {
        let mut rank: f32 = 0.0;

        for word in query.split_whitespace() {
            // TODO: add transcript once transcript generated for all comics
            // TODO: fix partial matches not working
            // rank += (tf-idf of alt * alt weight) + (tf-idf of title * title weight); where body weight + title weight = 1 <-------------------

            rank += calculate_tf(&word, &comic, &Field::Title) * calculate_idf(&word, &doc.frequency, &Field::Title, doc.comics.len()) * 0.6f32;
            rank +=  calculate_tf(&word, &comic, &Field::Alt) * calculate_idf(&word, &doc.frequency, &Field::Alt, doc.comics.len()) * 0.4f32; 
        }

        if rank < 0.0 {
            result.push((rank, comic.clone()));
        }
    }


    result.sort_by(|(a, _), (b, _)| a.partial_cmp(b).expect(&format!("{a} and {a} are not comparable")));

    result
}

fn calculate_tf(string: &str, c: &ComicIndex, field: &Field) -> f32 {
    match field {
        Field::Title => {
            let size = c.title_len as f32;
            let elements = c.title_freq.get(string).cloned().unwrap_or(0) as f32;

            elements / size
        },
        Field::Alt => {
            let size = c.alt_len as f32;
            let elements = c.alt_freq.get(string).cloned().unwrap_or(0) as f32;

            elements / size
        },
    }
}

fn calculate_idf(string: &str, df: &ComicFrequency, field: &Field, length: usize) -> f32 {
    match field {
        Field::Title => {
            let length = length as f32;
            let num = df.title_freq.get(string).cloned().unwrap_or(1) as f32;
            (num / length).log10()
        },
        Field::Alt => {
            let length = length as f32;
            let num = df.alt_freq.get(string).cloned().unwrap_or(1) as f32;
            (num / length).log10()
        },
    }
}

utils.rs

pub const URL: &str = "https://xkcd.com";

pub const DATA_PATH: &str = "data/document.json";
pub const INFO: &str = "info.0.json";
pub const RED: &str = "\x1b[31m";
pub const GREEN: &str = "\x1b[32m";
pub const CYAN: &str = "\x1b[34m";
pub const MAGENTA: &str = "\x1b[35m";
pub const RESET: &str = "\x1b[0m";

// TODO: use Transcript as a field and start indexing transcripts
pub enum Field {
    Title,
    Alt,
}

This is the github project if anyone wants to clone the entire project: link.

Example Output:

    Finished release [optimized] target(s) in 0.04s
     Running `target/release/xkcs --search 'Google Maps'`
┌──────────┬──────────┬───────────────────────┬─────────────────┐
│ Index    │ Rank     │ Title                 │ Comic Number    │
├──────────┼──────────┼───────────────────────┼─────────────────┤
│ 1        │ 1.7424   │ Maps                  │ 2617            │
│ 2        │ 1.6927   │ Google Maps           │ 461             │
│ 3        │ 1.1075   │ Google Trends Maps    │ 2126            │
│ 4        │ 0.9261   │ Reaction Maps         │ 2260            │
│ 5        │ 0.8042   │ Google Trends         │ 522             │
│ 6        │ 0.8042   │ Google Announcement   │ 1361            │
│ 7        │ 0.5362   │ Working for Google    │ 192             │
│ 8        │ 0.5362   │ 2020 Google Trends    │ 2302            │
│ 9        │ 0.1035   │ Expedition            │ 1169            │
│ 10       │ 0.0908   │ Nostalgia             │ 318             │
└──────────┴──────────┴───────────────────────┴─────────────────┘

The main issues and areas of improvement I can see are:

• Improving the search algorithm
  • Right now, partial word matches are discarded due to the way words are indexed. Currently, I'm checking a HashMap of frequencies for the every string literal in the provided query, so no partial matches or related terms are displayed.
  • I could probably do this with something like Levenshtein distances? I would be interested to hear feedback on this.
• Improving the mass of spaghetti code in main.rs:
  • I didn't want to install yet another module, so I hacked together a table formatter. This code looks really bad.
  • I should have probably just moved it to a separate file, or install a module which does this for me.
• Improving user experience
  • Right now, to view the comic you have to manually copy and paste the number into the xkcd website. This is probably not that helpful.
  • In hindsight, I probably should have added the alt text to the final table, although that wouldn't have been super helpful either.

Finally, something which may be concerning is that my code for the tf-idf calculations only returns negative numbers. I'm not sure where I messed up in my implementation, but I think the tf-idf calculations are possibly incorrect (although they give decent results).

Additionally, I'm not sure how good this Rust code is (especially the mess which is download_all), and I'm extremely curious about what can be improved.

Advice is much appreciated.

Answer 1

Looks like nice little tool. I like that you split it up into several modules. I just started learning Rust so I'm probably missing a few things but below my take on it.

main.rs

1. This:

   let mut limit = 10;
   let mut query = String::new();
   
   if let Some(l) = args.limit {
       limit = l;
   }
   
   if let Some(s) = args.search {
       query = s;
   }
   

   can be simplified to

    let limit = args.limit.unwrap_or(10);
    let query = args.search.unwrap_or(String::new());
   

2. The download and search options are mutually exclusive (if a download option is given then search is ignored) which seems unnecessarily restrictive (and also not obvious to a user). Just replace the last else clause with if !query.is_empty() (which also avoids loading the entire dataset before checking that no query term has been specified).

comic.rs

1. Shared mutable state is an unfortunately sometimes necessary evil but it should be avoided wherever possible:

   let c = Mutex::new(Vec::new());
   
   let frequencies = Arc::new(Mutex::new(ComicFrequency {
       title_freq: HashMap::new(),
       alt_freq: HashMap::new(),
   }));
   

   Instead of doing the above I'd consider restructuring the code so each parallel process collects the data locally and then returns the data structure. The calling code can the collect the results and merge them. No locks needed. I can't give an hint right now on how to do this since I lack familiarity with Rust but I'd be pretty surprised if it couldn't be achieved.

2. DRY and SRP (Single Responsibility Principle) violations - download_all and download_and_append_document have duplicated code and also IMHO do too much. There should be a function that is download_and_index_comic which returns an indexed comic and has basically no external dependencies. This can then be used to compose download_all and the wrapper function to combine and store it.

3. I'd consider storing individual files by comic index. This way checking whether a specific comic already exists is a simple file exists check. You don't have to load the entire structure, search through it and then write it back. It might make loading the entire dataset a tad slower since you have to load many files but I'd argue that trade-off to probably be worth it. By the time filesystem slowness gets you due to large number of files you should be storing it in a database anyway.

search.rs

1. I don't see why doc needs to be mutable when searching. Also semantically speaking this seems wrong.

2. Try to encode invariants in code rather than just in comments

   // rank += (tf-idf of alt * alt weight) + (tf-idf of title * title weight); where body weight + title weight = 1 <-------------------
   
   rank += calculate_tf(&word, &comic, &Field::Title) * calculate_idf(&word, &doc.frequency, &Field::Title, doc.comics.len()) * 0.6f32;
   rank +=  calculate_tf(&word, &comic, &Field::Alt) * calculate_idf(&word, &doc.frequency, &Field::Alt, doc.comics.len()) * 0.4f32; 
   

   The above should be encapsulated in a calculate_rank(word: &str, comic: &ComicIndex, title_weight: f32) function that automatically computes alt_weight = 1.0 - title_weight. Ideally the compiler could enforce that title_weight is between 0 and 1 but unfortunately that's not really possible at the moment (generic_const_exprs are experimental and don't support floating point).

General

1. Hard-coded DATA_PATH and PARALLEL_REQUESTS - an argument parser is already in place, should be easy enough to add these. Gives you the ability to tweak without having to re-compile.

2. Apparently there is a crate for printing tables. There is even a way to output clickable links which would improve UX a lot.

3. This term r#"(),".;:'"# appears in several places - consolidate to a named constant to keep it DRY.