From: Jacob Lifshay <programmerjake@gmail.com>
Date: Tue, 28 Feb 2023 06:37:07 +0000 (-0800)
Subject: add --arbitrary flag
X-Git-Url: https://git.libre-soc.org/?a=commitdiff_plain;p=bigint-presentation-code.git

add --arbitrary flag
---

diff --git a/register_allocator/src/fuzzing.rs b/register_allocator/src/fuzzing.rs
index 6ba5ace..c8c3c4f 100644
--- a/register_allocator/src/fuzzing.rs
+++ b/register_allocator/src/fuzzing.rs
@@ -1,4 +1,3 @@
-#![cfg(feature = "fuzzing")]
 use crate::{
     function::{
         Block, BlockTermInstKind, Constraint, CopyInstKind, FnFields, Inst, InstKind, InstStage,
diff --git a/register_allocator/src/main.rs b/register_allocator/src/main.rs
index 06cf8bf..9e97c32 100644
--- a/register_allocator/src/main.rs
+++ b/register_allocator/src/main.rs
@@ -1,17 +1,30 @@
-use bigint_presentation_code_register_allocator::{function::Function, interned::GlobalState};
+use arbitrary::{Arbitrary, Unstructured};
+use bigint_presentation_code_register_allocator::{
+    function::{FnFields, Function},
+    interned::GlobalState,
+};
 use clap::Parser;
 use eyre::Result;
 use serde::Serialize;
 use serde_json::{Deserializer, Value};
-use std::io::{BufWriter, Write};
+use std::{
+    io::{BufWriter, Read, Write},
+    process::ExitCode,
+};
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about)]
 struct Args {
+    /// dump the input function before attempting register allocation
     #[arg(long)]
     dump_input: bool,
+    /// write JSON outputs in pretty format
     #[arg(long)]
     pretty: bool,
+    /// generate an input function from the random bytes read from stdin using
+    /// the same functionality used for fuzzing
+    #[arg(long)]
+    arbitrary: bool,
 }
 
 #[derive(Serialize, Debug)]
@@ -23,6 +36,7 @@ pub enum Info<'a> {
 pub enum Output<'a> {
     Error(String),
     Info(Info<'a>),
+    NeedDifferentInput,
     Success(Value),
 }
 
@@ -43,23 +57,56 @@ impl<'a> Output<'a> {
     }
 }
 
-fn main() -> Result<()> {
-    let Args { dump_input, pretty } = Args::parse();
-    let stdin = std::io::stdin().lock();
-    for input in Deserializer::from_reader(stdin).into_iter::<Value>() {
-        GlobalState::scope(|| -> Result<()> {
-            let function = match serde_json::from_value::<Function>(input?) {
-                Ok(v) => v,
-                Err(e) => {
-                    Output::from_err(e).write_to_stdout(pretty)?;
-                    return Ok(());
-                }
-            };
-            if dump_input {
-                Output::Info(Info::DumpInput(&function)).write_to_stdout(pretty)?;
+fn process_input(args: &Args, input: Value, exit_code: &mut ExitCode) -> Result<()> {
+    GlobalState::scope(|| -> Result<()> {
+        let function = match serde_json::from_value::<Function>(input) {
+            Ok(v) => v,
+            Err(e) => {
+                Output::from_err(e).write_to_stdout(args.pretty)?;
+                *exit_code = ExitCode::FAILURE;
+                return Ok(());
             }
-            todo!()
-        })?;
+        };
+        if args.dump_input {
+            Output::Info(Info::DumpInput(&function)).write_to_stdout(args.pretty)?;
+        }
+        todo!()
+    })
+}
+
+fn arbitrary_input(input_bytes: &[u8]) -> arbitrary::Result<Value> {
+    GlobalState::scope(|| -> arbitrary::Result<Value> {
+        Ok(
+            serde_json::to_value(FnFields::arbitrary_take_rest(Unstructured::new(
+                &input_bytes,
+            ))?)
+            .expect("serialization shouldn't ever fail"),
+        )
+    })
+}
+
+fn main() -> Result<ExitCode> {
+    let args = Args::parse();
+    let Args {
+        dump_input: _,
+        pretty,
+        arbitrary,
+    } = args;
+    let stdin = std::io::stdin().lock();
+    let mut exit_code = ExitCode::SUCCESS;
+    if arbitrary {
+        let mut input_bytes = vec![];
+        stdin.take(0x10000).read_to_end(&mut input_bytes)?;
+        if let Ok(input) = arbitrary_input(&input_bytes) {
+            process_input(&args, input, &mut exit_code)?;
+        } else {
+            Output::NeedDifferentInput.write_to_stdout(pretty)?;
+            exit_code = ExitCode::FAILURE;
+        };
+    } else {
+        for input in Deserializer::from_reader(stdin).into_iter::<Value>() {
+            process_input(&args, input?, &mut exit_code)?;
+        }
     }
-    Ok(())
+    Ok(exit_code)
 }