webgraph_cli/perm/
bfs.rs

1/*
2 * SPDX-FileCopyrightText: 2023 Sebastiano Vigna
3 * SPDX-FileCopyrightText: 2023 Tommaso Fontana
4 *
5 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
6 */
7
8use 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    /// The basename of the graph.
23    pub src: PathBuf,
24
25    /// The filename of the permutation in binary big-endian format.
26    pub perm: PathBuf,
27
28    #[arg(short, long)]
29    /// Save the permutation in ε-serde format.
30    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    // load the graph
54    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    // create the permutation
69    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}