1use crate::{create_parent_dir, GlobalArgs};
9use anyhow::{Context, Result};
10use clap::Parser;
11use dsi_bitstream::dispatch::factory::CodesReaderFactoryHelper;
12use dsi_bitstream::prelude::*;
13use dsi_progress_logger::prelude::*;
14use epserde::prelude::Serialize;
15use std::io::{BufWriter, Write};
16use std::path::PathBuf;
17use webgraph::prelude::*;
18
19#[derive(Parser, Debug)]
20#[command(name = "bfs", about = "Computes the permutation induced by a breadth-first visit.", long_about = None)]
21pub struct CliArgs {
22 pub src: PathBuf,
24
25 pub perm: PathBuf,
27
28 #[arg(short, long)]
29 pub epserde: bool,
31}
32
33pub fn main(global_args: GlobalArgs, args: CliArgs) -> Result<()> {
34 create_parent_dir(&args.perm)?;
35
36 match get_endianness(&args.src)?.as_str() {
37 #[cfg(feature = "be_bins")]
38 BE::NAME => bfs::<BE>(global_args, args),
39 #[cfg(feature = "le_bins")]
40 LE::NAME => bfs::<LE>(global_args, args),
41 e => panic!("Unknown endianness: {}", e),
42 }
43}
44
45pub fn bfs<E: Endianness + 'static + Send + Sync>(
46 global_args: GlobalArgs,
47 args: CliArgs,
48) -> Result<()>
49where
50 MemoryFactory<E, MmapHelper<u32>>: CodesReaderFactoryHelper<E>,
51 for<'a> LoadModeCodesReader<'a, E, LoadMmap>: BitSeek,
52{
53 let graph = BvGraph::with_basename(&args.src)
55 .mode::<LoadMmap>()
56 .flags(MemoryFlags::TRANSPARENT_HUGE_PAGES | MemoryFlags::RANDOM_ACCESS)
57 .endianness::<E>()
58 .load()?;
59
60 let mut pl = ProgressLogger::default();
61 pl.display_memory(true)
62 .item_name("nodes")
63 .expected_updates(Some(graph.num_nodes()));
64 if let Some(duration) = global_args.log_interval {
65 pl.log_interval(duration);
66 }
67
68 let mut perm = vec![0; graph.num_nodes()];
70 pl.start("Computing BFS permutation...");
71 let mut visit = webgraph::visits::breadth_first::Seq::new(&graph);
72 for (i, event) in visit.into_iter().enumerate() {
73 perm[event.node] = i;
74 pl.light_update();
75 }
76 pl.done();
77
78 if args.epserde {
79 unsafe {
80 perm.store(&args.perm)
81 .with_context(|| format!("Could not write permutation to {}", args.perm.display()))
82 }?;
83 } else {
84 let mut file = std::fs::File::create(&args.perm)
85 .with_context(|| format!("Could not create permutation at {}", args.perm.display()))?;
86 let mut buf = BufWriter::new(&mut file);
87 pl.start(format!("Storing the nodes to {}", args.perm.display()));
88 for word in perm.iter() {
89 buf.write_all(&word.to_be_bytes()).with_context(|| {
90 format!("Could not write permutation to {}", args.perm.display())
91 })?;
92 pl.light_update();
93 }
94 pl.done();
95 }
96 Ok(())
97}