1 // Copyright 2019-2020 Georges Racinet <georges.racinet@octobus.net>

     2 //

     3 // This software may be used and distributed according to the terms of the

     4 // GNU General Public License version 2 or any later version.

     5 

     6 use clap::*;

     7 use hg::revlog::node::*;

     8 use hg::revlog::nodemap::*;

     9 use hg::revlog::*;

    10 use memmap::MmapOptions;

    11 use rand::Rng;

    12 use std::fs::File;

    13 use std::io;

    14 use std::io::Write;

    15 use std::path::{Path, PathBuf};

    16 use std::str::FromStr;

    17 use std::time::Instant;

    18 

    19 mod index;

    20 use index::Index;

    21 

    22 fn mmap_index(repo_path: &Path) -> Index {

    23     let mut path = PathBuf::from(repo_path);

    24     path.extend([".hg", "store", "00changelog.i"].iter());

    25     Index::load_mmap(path)

    26 }

    27 

    28 fn mmap_nodemap(path: &Path) -> NodeTree {

    29     let file = File::open(path).unwrap();

    30     let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };

    31     let len = mmap.len();

    32     NodeTree::load_bytes(Box::new(mmap), len)

    33 }

    34 

    35 /// Scan the whole index and create the corresponding nodemap file at `path`

    36 fn create(index: &Index, path: &Path) -> io::Result<()> {

    37     let mut file = File::create(path)?;

    38     let start = Instant::now();

    39     let mut nm = NodeTree::default();

    40     for rev in 0..index.len() {

    41         let rev = rev as Revision;

    42         nm.insert(index, index.node(rev).unwrap(), rev).unwrap();

    43     }

    44     eprintln!("Nodemap constructed in RAM in {:?}", start.elapsed());

    45     file.write(&nm.into_readonly_and_added_bytes().1)?;

    46     eprintln!("Nodemap written to disk");

    47     Ok(())

    48 }

    49 

    50 fn query(index: &Index, nm: &NodeTree, prefix: &str) {

    51     let start = Instant::now();

    52     let res = nm.find_hex(index, prefix);

    53     println!("Result found in {:?}: {:?}", start.elapsed(), res);

    54 }

    55 

    56 fn bench(index: &Index, nm: &NodeTree, queries: usize) {

    57     let len = index.len() as u32;

    58     let mut rng = rand::thread_rng();

    59     let nodes: Vec<Node> = (0..queries)

    60         .map(|_| {

    61             index

    62                 .node((rng.gen::<u32>() % len) as Revision)

    63                 .unwrap()

    64                 .clone()

    65         })

    66         .collect();

    67     if queries < 10 {

    68         let nodes_hex: Vec<String> =

    69             nodes.iter().map(|n| n.encode_hex()).collect();

    70         println!("Nodes: {:?}", nodes_hex);

    71     }

    72     let mut last: Option<Revision> = None;

    73     let start = Instant::now();

    74     for node in nodes.iter() {

    75         last = nm.find_bin(index, node.into()).unwrap();

    76     }

    77     let elapsed = start.elapsed();

    78     println!(

    79         "Did {} queries in {:?} (mean {:?}), last was {:?} with result {:?}",

    80         queries,

    81         elapsed,

    82         elapsed / (queries as u32),

    83         nodes.last().unwrap().encode_hex(),

    84         last

    85     );

    86 }

    87 

    88 fn main() {

    89     let matches = App::new("Nodemap pure Rust example")

    90         .arg(

    91             Arg::with_name("REPOSITORY")

    92                 .help("Path to the repository, always necessary for its index")

    93                 .required(true),

    94         )

    95         .arg(

    96             Arg::with_name("NODEMAP_FILE")

    97                 .help("Path to the nodemap file, independent of REPOSITORY")

    98                 .required(true),

    99         )

   100         .subcommand(

   101             SubCommand::with_name("create")

   102                 .about("Create NODEMAP_FILE by scanning repository index"),

   103         )

   104         .subcommand(

   105             SubCommand::with_name("query")

   106                 .about("Query NODEMAP_FILE for PREFIX")

   107                 .arg(Arg::with_name("PREFIX").required(true)),

   108         )

   109         .subcommand(

   110             SubCommand::with_name("bench")

   111                 .about(

   112                     "Perform #QUERIES random successful queries on NODEMAP_FILE")

   113                 .arg(Arg::with_name("QUERIES").required(true)),

   114         )

   115         .get_matches();

   116 

   117     let repo = matches.value_of("REPOSITORY").unwrap();

   118     let nm_path = matches.value_of("NODEMAP_FILE").unwrap();

   119 

   120     let index = mmap_index(&Path::new(repo));

   121 

   122     if let Some(_) = matches.subcommand_matches("create") {

   123         println!("Creating nodemap file {} for repository {}", nm_path, repo);

   124         create(&index, &Path::new(nm_path)).unwrap();

   125         return;

   126     }

   127 

   128     let nm = mmap_nodemap(&Path::new(nm_path));

   129     if let Some(matches) = matches.subcommand_matches("query") {

   130         let prefix = matches.value_of("PREFIX").unwrap();

   131         println!(

   132             "Querying {} in nodemap file {} of repository {}",

   133             prefix, nm_path, repo

   134         );

   135         query(&index, &nm, prefix);

   136     }

   137     if let Some(matches) = matches.subcommand_matches("bench") {

   138         let queries =

   139             usize::from_str(matches.value_of("QUERIES").unwrap()).unwrap();

   140         println!(

   141             "Doing {} random queries in nodemap file {} of repository {}",

   142             queries, nm_path, repo

   143         );

   144         bench(&index, &nm, queries);

   145     }

   146 }