graph.rs

use std::{cmp::Ordering, io::SeekFrom};

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

pub struct Node {
    id: u32,
    name: String,
}

impl Node {
    pub fn new(id: u32, name: String) -> Self {
        Self { id, name }
    }
}

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 - 1;

        while max >= min {
            let cur = (max + min) / 2;

            let node = self.get_node_from_name_idx(cur).await?;

            let c = name.cmp(&node.name);
            match c {
                Ordering::Less => max = cur - 1,
                Ordering::Equal => {
                    return Ok(Some(node.id));
                }
                Ordering::Greater => min = cur + 1,
            }
        }

        Ok(None)
    }

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

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

            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)).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 - 1,
                Ordering::Equal => {
                    return Ok(Some(host));
                }
                Ordering::Greater => min = cur + 1,
            }

            if max == min {
                return Ok(None);
            }
        }
    }

    /*async fn get_nearby_matching_names(
        &self,
        cur: u64,
        name: &str,
        wildcard: bool,
    ) -> anyhow::Result<Vec<u64>> {
        let mut min = cur;
        loop {
            let node = self.get_node_from_name_idx(min).await?;

            if node.name != node {}

            min -= 1;
        }
    }*/

    async fn get_node_from_name_idx(&mut self, idx: u64) -> anyhow::Result<Node> {
        self.name_index.seek(SeekFrom::Start(idx * 8)).await?;
        let offset = self.name_index.read_u64_le().await?;

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

        Ok(Node::new(id, host))
    }
}

#[derive(Debug)]
pub struct Edge {
    pub from: u32,
    pub to: u32,
}

pub struct EdgeReader {
    edges_from: File,
    edges_to: File,
    num_edges: u64,
}

impl EdgeReader {
    pub async fn new() -> anyhow::Result<Self> {
        let edges_from = File::open("edges_from.bin").await?;
        let edges_to = File::open("edges_to.bin").await?;
        let num_edges = edges_from.metadata().await?.len() / 8; // 8 bytes per edge

        Ok(Self {
            edges_from,
            edges_to,
            num_edges,
        })
    }

    pub async fn get_nodes_from(&mut self, from: u32) -> anyhow::Result<Vec<Edge>> {
        let mut min = 0;
        let mut max = self.num_edges - 1;

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

            let edge = self.get_edge_from_from_idx(cur).await?;

            let c = from.cmp(&edge.from);
            match c {
                Ordering::Less => max = cur - 1,
                Ordering::Equal => {
                    return self.get_all_consecutive_matching_from(cur, from).await;
                }
                Ordering::Greater => min = cur + 1,
            }
        }
    }

    pub async fn get_nodes_to(&mut self, to: u32) -> anyhow::Result<Vec<Edge>> {
        let mut min = 0;
        let mut max = self.num_edges - 1;

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

            let edge = self.get_edge_from_to_idx(cur).await?;

            let c = to.cmp(&edge.to);
            match c {
                Ordering::Less => max = cur - 1,
                Ordering::Equal => {
                    return self.get_all_consecutive_matching_to(cur, to).await;
                }
                Ordering::Greater => min = cur + 1,
            }
        }
    }

    async fn get_all_consecutive_matching_from(
        &mut self,
        cur: u64,
        from: u32,
    ) -> anyhow::Result<Vec<Edge>> {
        let mut out = vec![];
        let mut min = cur;
        let mut max = cur + 1;

        loop {
            let edge = self.get_edge_from_from_idx(min).await?;
            if edge.from != from {
                break;
            }

            out.push(edge);

            min -= 1;
        }

        loop {
            let edge = self.get_edge_from_from_idx(max).await?;
            if edge.from != from {
                break;
            }

            out.push(edge);

            max += 1;
        }

        Ok(out)
    }

    async fn get_edge_from_from_idx(&mut self, idx: u64) -> anyhow::Result<Edge> {
        self.edges_from.seek(SeekFrom::Start(idx * 8)).await?;
        let from = self.edges_from.read_u32_le().await?;
        let to = self.edges_from.read_u32_le().await?;

        Ok(Edge { from, to })
    }

    async fn get_all_consecutive_matching_to(
        &mut self,
        cur: u64,
        to: u32,
    ) -> anyhow::Result<Vec<Edge>> {
        let mut out = vec![];
        let mut min = cur;
        let mut max = cur + 1;

        loop {
            let edge = self.get_edge_from_to_idx(min).await?;
            if edge.to != to {
                break;
            }

            out.push(edge);

            min -= 1;
        }

        loop {
            let edge = self.get_edge_from_from_idx(max).await?;
            if edge.to != to {
                break;
            }

            out.push(edge);

            max += 1;
        }

        Ok(out)
    }

    async fn get_edge_from_to_idx(&mut self, idx: u64) -> anyhow::Result<Edge> {
        self.edges_to.seek(SeekFrom::Start(idx * 8)).await?;
        let to = self.edges_to.read_u32_le().await?;
        let from = self.edges_to.read_u32_le().await?;

        Ok(Edge { from, to })
    }
}