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 }) } }