Skip to content
Snippets Groups Projects
graph.rs 2.55 KiB
Newer Older
use std::{cmp::Ordering, io::SeekFrom};

use anyhow::Context;
use tokio::{
    fs::File,
    io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, BufReader},
};

pub struct NodeReader {
    nodes: File,
    name_index: File,
    id_index: File,
    num_nodes: u64,
}

impl NodeReader {
    pub async fn new() -> anyhow::Result<Self> {
        let nodes = File::open("nodes.bin").await?;
        let name_index = File::open("nodes_by_name_index.bin").await?;
        let id_index = File::open("nodes_by_id_index.bin").await?;

        let num_nodes = name_index.metadata().await?.len() / 8; // 8 bytes per node

        Ok(Self {
            nodes,
            name_index,
            id_index,
            num_nodes,
        })
    }

    pub async fn id_from_name(&mut self, name: &str) -> anyhow::Result<Option<u32>> {
        // binary search for name

        let mut min = 0;
        let mut max = self.num_nodes;

        loop {
            let cur = (max + min) / 2 + 1;

            self.name_index.seek(SeekFrom::Start(cur * 8)).await?;
            let offset = self.name_index.read_u64_le().await?;

            self.nodes.seek(SeekFrom::Start(offset.into())).await?;
            let id = self.nodes.read_u32_le().await?;
            let host = BufReader::new(&mut self.nodes)
                .lines()
                .next_line()
                .await?
                .context("Missing line")?;

            let c = name.cmp(&host);
            match c {
                Ordering::Less => max = cur,
                Ordering::Equal => {
                    return Ok(Some(id));
                }
                Ordering::Greater => min = cur,
            }
        }
    }
Stephen D's avatar
Stephen D committed

    pub async fn name_from_id(&mut self, id: u32) -> anyhow::Result<Option<String>> {
        let mut min = 0;
        let mut max = self.num_nodes;

        loop {
            let cur = (max + min) / 2 + 1;

            self.id_index.seek(SeekFrom::Start(cur * 12 + 4)).await?;
            let offset = self.id_index.read_u64_le().await?;

            self.nodes.seek(SeekFrom::Start(offset.into())).await?;
            let host_id = self.nodes.read_u32_le().await?;
            let host = BufReader::new(&mut self.nodes)
                .lines()
                .next_line()
                .await?
                .context("Missing line")?;

            let c = id.cmp(&host_id);
            match c {
                Ordering::Less => max = cur,
                Ordering::Equal => {
                    return Ok(Some(host));
                }
                Ordering::Greater => min = cur,
            }
        }
    }