webgraph_cli/perm/
rand.rs

1/*
2 * SPDX-FileCopyrightText: 2023 Inria
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_progress_logger::prelude::*;
12use epserde::ser::Serialize;
13use rand::prelude::SliceRandom;
14use std::io::prelude::*;
15use std::path::PathBuf;
16
17#[derive(Parser, Debug)]
18#[command(name = "rand", about = "Creates a random permutation.", long_about = None)]
19pub struct CliArgs {
20    /// The number of elements in the permutation.
21    pub len: usize,
22    /// The random permutation in binary big-endian format.
23    pub dst: PathBuf,
24
25    #[arg(short = 'e', long)]
26    /// Store the permutation in ε-serde format.
27    pub epserde: bool,
28}
29
30pub fn main(global_args: GlobalArgs, args: CliArgs) -> Result<()> {
31    create_parent_dir(&args.dst)?;
32
33    let mut rng = rand::rng();
34    let mut perm = (0..args.len).collect::<Vec<_>>();
35    perm.shuffle(&mut rng);
36
37    if args.epserde {
38        unsafe {
39            perm.store(&args.dst)
40                .with_context(|| format!("Could not store permutation to {}", args.dst.display()))
41        }?;
42    } else {
43        let mut pl = ProgressLogger::default();
44        pl.display_memory(true).item_name("index");
45        if let Some(duration) = global_args.log_interval {
46            pl.log_interval(duration);
47        }
48
49        let mut file =
50            std::io::BufWriter::new(std::fs::File::create(&args.dst).with_context(|| {
51                format!("Could not create permutation at {}", args.dst.display())
52            })?);
53        pl.start("Writing permutation indices...");
54        for perm in perm {
55            file.write_all(&perm.to_be_bytes()).with_context(|| {
56                format!("Could not write permutation to {}", args.dst.display())
57            })?;
58            pl.light_update();
59        }
60        pl.done();
61    }
62
63    Ok(())
64}