embed default corpus

Cargo.toml

@@ -10,9 +10,20 @@ anyhow = "1.0.71"
 clap = "4.3.3"
 glob = "0.3.1"
 log = "0.4.19"
+lz4_flex = { version = "0.11.1", default-features = false }
+postcard = { version = "1.0.4", features = ["use-std"], default-features = false }
+serde = { version = "1.0.164", features = ["derive"] }
 simple_logger = "4.1.0"
 tablestream = "0.1.3"
 
 [dev-dependencies]
 assert_approx_eq = "1.1.0"
 rand = "0.8.5"
+
+[build-dependencies]
+anyhow = "1.0.71"
+glob = "0.3.1"
+log = "0.4.19"
+lz4_flex = { version = "0.11.1", default-features = false }
+postcard = { version = "1.0.4", features = ["use-std"], default-features = false }
+serde = { version = "1.0.164", features = ["derive"] }

build.rs

@@ -0,0 +1,23 @@
+#[allow(dead_code)]
+mod corpus {
+    include!("src/corpus.rs");
+}
+
+fn main() {
+    // load default corpus
+    let default = corpus::load_corpus("cpu_rec_corpus/*.corpus").unwrap();
+    println!("cargo:rerun-if-changed=cpu_rec_corpus");
+
+    // serialize to bytes
+    let bytes = postcard::to_stdvec(&default).unwrap();
+
+    // compress
+    let bytes = lz4_flex::compress(&bytes);
+
+    // output path to target build folder
+    let mut outfile = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
+    outfile.push("default.pc");
+
+    // write to file
+    std::fs::write(outfile, bytes).unwrap();
+}

src/corpus.rs

@@ -16,12 +16,13 @@
 use anyhow::{Context, Error, Ok, Result};
 use glob::glob;
 use log::debug;
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::str::FromStr;
 use std::string::String;
 
-#[derive(Debug)]
+#[derive(Debug, Deserialize, Serialize)]
 pub struct CorpusStats {
     pub arch: String,
     bigrams_freq: HashMap<(u8, u8), f32>,

src/main.rs

@@ -202,7 +202,7 @@ fn main() -> Result<()> {
         .propagate_version(true)
         .author("Raphaël Rigo <devel@syscall.eu>")
         .about("Identifies CPU architectures in binaries")
-        .arg(arg!(--corpus <corpus_dir>).default_value("cpu_rec_corpus"))
+        .arg(arg!(--corpus <corpus_dir>))
         .arg(arg!(-d - -debug))
         .arg(arg!(-v - -verbose))
         .arg(
@@ -223,17 +223,33 @@ fn main() -> Result<()> {
     };
     simple_logger::init_with_level(level)?;
 
-    let corpus_dir = args.get_one::<String>("corpus").unwrap().to_owned();
-    if !Path::new(&corpus_dir).is_dir() {
-        return Err(Error::msg(format!(
-            "{} is not a valid directory",
-            corpus_dir
-        )));
-    }
-    let corpus_files: String = args.get_one::<String>("corpus").unwrap().to_owned() + "/*.corpus";
-    println!("Loading corpus from {}", corpus_files);
+    let default_corpus: Vec<CorpusStats> = {
+        // serialized bytes embedded from build.rs
+        let bytes = include_bytes!(concat!(env!("OUT_DIR"), "/default.pc"));
+
+        // decompress
+        let bytes = lz4_flex::decompress(bytes, 39364573).unwrap();
+
+        // deserialize
+        postcard::from_bytes(&bytes).unwrap()
+    };
 
-    let corpus_stats = load_corpus(&corpus_files)?;
+    let corpus_stats = match args.get_one::<&str>("corpus") {
+        // if no arg given, use embedded corpus
+        None => default_corpus,
+        // attempt to load the given corpus folder
+        Some(corpus_dir) => {
+            if !Path::new(corpus_dir).is_dir() {
+                return Err(Error::msg(format!(
+                    "{} is not a valid directory",
+                    corpus_dir
+                )));
+            }
+            let corpus_files = format!("{corpus_dir}/*.corpus");
+            println!("Loading corpus from {}", corpus_files);
+            load_corpus(&corpus_files)?
+        }
+    };
 
     info!("Corpus size: {}", corpus_stats.len());