Use clap for PIR cli (#244)

Cargo.toml

@@ -13,12 +13,15 @@ license = "MIT"
 rust-version = "1.73"
 
 [workspace.dependencies]
+clap = { version = "^4.4.18", features = ["derive"] }
 console = "^0.15.8"
 criterion = "^0.5.1"
 doc-comment = "^0.3.3"
+env_logger = "^0.11.1"
 ethnum = "^1.5.0"
 indicatif = "^0.17.7"
 itertools = "^0.12.1"
+log = "^0.4.20"
 ndarray = "^0.15.6"
 num-bigint = "^0.4.4"
 num-bigint-dig = "^0.8.4"

crates/fhe/Cargo.toml

@@ -33,12 +33,15 @@ thiserror.workspace = true
 prost-build.workspace = true
 
 [dev-dependencies]
+clap.workspace = true
+console.workspace = true
 criterion.workspace = true
+env_logger.workspace = true
 indicatif.workspace = true
 itertools.workspace = true
+log.workspace = true
 ndarray.workspace = true
 rand.workspace = true
-console.workspace = true
 
 [[bench]]
 name = "bfv"
@@ -58,8 +61,5 @@ name = "mulpir"
 [[example]]
 name = "sealpir"
 
-[[example]]
-name = "util"
-
 [[example]]
 name = "voting"

crates/fhe/examples/mulpir.rs

@@ -6,104 +6,40 @@
 // We use the same parameters as in the paper to enable an apple-to-apple
 // comparison.
 
+mod pir;
 mod util;
 
-use console::style;
+use clap::Parser;
 use fhe::bfv;
 use fhe_traits::{
     DeserializeParametrized, FheDecoder, FheDecrypter, FheEncoder, FheEncrypter, Serialize,
 };
 use fhe_util::{inverse, transcode_to_bytes};
 use indicatif::HumanBytes;
 use rand::{rngs::OsRng, thread_rng, RngCore};
-use std::{env, error::Error, process::exit, time::Instant};
+use std::{error::Error, time::Instant};
 use util::{
     encode_database, generate_database, number_elements_per_plaintext,
     timeit::{timeit, timeit_n},
 };
 
-fn print_notice_and_exit(max_element_size: usize, error: Option<String>) {
-    println!(
-        "{} MulPIR with fhe.rs",
-        style("  overview:").magenta().bold()
-    );
-    println!(
-        "{} mulpir- [-h] [--help] [--database_size=<value>] [--element_size=<value>]",
-        style("     usage:").magenta().bold()
-    );
-    println!(
-        "{} {} must be at least 1, and {} must be between 1 and {}",
-        style("constraints:").magenta().bold(),
-        style("database_size").blue(),
-        style("element_size").blue(),
-        max_element_size
-    );
-    if let Some(error) = error {
-        println!("{} {}", style("     error:").red().bold(), error);
-    }
-    exit(0);
-}
-
 fn main() -> Result<(), Box<dyn Error>> {
     // We use the parameters reported in Table 1 of https://eprint.iacr.org/2019/1483.pdf.
     let degree = 8192;
     let plaintext_modulus: u64 = (1 << 20) + (1 << 19) + (1 << 17) + (1 << 16) + (1 << 14) + 1;
     let moduli_sizes = [50, 55, 55];
 
+    let args = pir::Cli::parse();
+    let database_size = args.database_size;
+    let elements_size = args.element_size;
+
     // Compute what is the maximum byte-length of an element to fit within one
     // ciphertext. Each coefficient of the ciphertext polynomial can contain
     // floor(log2(plaintext_modulus)) bits.
     let max_element_size = ((plaintext_modulus.ilog2() as usize) * degree) / 8;
-
-    // This executable is a command line tool which enables to specify different
-    // database and element sizes.
-    let args: Vec<String> = env::args().skip(1).collect();
-
-    // Print the help if requested.
-    if args.contains(&"-h".to_string()) || args.contains(&"--help".to_string()) {
-        print_notice_and_exit(max_element_size, None)
-    }
-
-    // Use the default values from <https://eprint.iacr.org/2019/1483.pdf>.
-    let mut database_size = 1 << 20;
-    let mut elements_size = 288;
-
-    // Update the database size and/or element size depending on the arguments
-    // provided.
-    for arg in &args {
-        if arg.starts_with("--database_size") {
-            let a: Vec<&str> = arg.rsplit('=').collect();
-            if a.len() != 2 || a[0].parse::<usize>().is_err() {
-                print_notice_and_exit(
-                    max_element_size,
-                    Some("Invalid `--database_size` command".to_string()),
-                )
-            } else {
-                database_size = a[0].parse::<usize>()?
-            }
-        } else if arg.starts_with("--element_size") {
-            let a: Vec<&str> = arg.rsplit('=').collect();
-            if a.len() != 2 || a[0].parse::<usize>().is_err() {
-                print_notice_and_exit(
-                    max_element_size,
-                    Some("Invalid `--element_size` command".to_string()),
-                )
-            } else {
-                elements_size = a[0].parse::<usize>()?
-            }
-        } else {
-            print_notice_and_exit(
-                max_element_size,
-                Some(format!("Unrecognized command: {arg}")),
-            )
-        }
-    }
-
     if elements_size > max_element_size || elements_size == 0 || database_size == 0 {
-        print_notice_and_exit(
-            max_element_size,
-            Some("Element or database sizes out of bound".to_string()),
-        )
+        log::error!("Invalid parameters: database_size = {database_size}, elements_size = {elements_size}. The maximum element size if {max_element_size}.");
+        clap::Error::new(clap::error::ErrorKind::InvalidValue).exit();
     }
 
     // The parameters are within bound, let's go! Let's first display some

crates/fhe/examples/pir.rs

@@ -0,0 +1,21 @@
+use clap::Parser;
+
+#[derive(Parser)]
+pub struct Cli {
+    #[arg(
+        long,
+        help = "The number of elements in the database",
+        default_value = "65536"
+    )]
+    pub database_size: usize,
+
+    #[arg(
+        long,
+        help = "The size of each database element",
+        default_value = "1024"
+    )]
+    pub element_size: usize,
+}
+
+#[allow(dead_code)]
+fn main() {}