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        perm.store(&args.dst)
39            .with_context(|| format!("Could not store permutation to {}", args.dst.display()))?;
40    } else {
41        let mut pl = ProgressLogger::default();
42        pl.display_memory(true).item_name("index");
43        if let Some(duration) = global_args.log_interval {
44            pl.log_interval(duration);
45        }
46
47        let mut file =
48            std::io::BufWriter::new(std::fs::File::create(&args.dst).with_context(|| {
49                format!("Could not create permutation at {}", args.dst.display())
50            })?);
51        pl.start("Writing permutation indices...");
52        for perm in perm {
53            file.write_all(&perm.to_be_bytes()).with_context(|| {
54                format!("Could not write permutation to {}", args.dst.display())
55            })?;
56            pl.light_update();
57        }
58        pl.done();
59    }
60
61    Ok(())
62}