register_allocator/src/main.rs

   1 use arbitrary::{Arbitrary, Unstructured};
   2 use bigint_presentation_code_register_allocator::{
   3     function::{FnFields, Function},
   4     interned::GlobalState,
   5 };
   6 use clap::Parser;
   7 use eyre::Result;
   8 use serde::Serialize;
   9 use serde_json::{Deserializer, Value};
  10 use std::{
  11     io::{BufWriter, Read, Write},
  12     process::ExitCode,
  13 };
  14
  15 #[derive(Parser, Debug)]
  16 #[command(version, about, long_about)]
  17 struct Args {
  18     /// dump the input function before attempting register allocation
  19     #[arg(long)]
  20     dump_input: bool,
  21     /// write JSON outputs in pretty format
  22     #[arg(long)]
  23     pretty: bool,
  24     /// generate an input function from the random bytes read from stdin using
  25     /// the same functionality used for fuzzing
  26     #[arg(long)]
  27     arbitrary: bool,
  28 }
  29
  30 #[derive(Serialize, Debug)]
  31 pub enum Info<'a> {
  32     DumpInput(&'a Function),
  33 }
  34
  35 #[derive(Serialize, Debug)]
  36 pub enum Output<'a> {
  37     Error(String),
  38     Info(Info<'a>),
  39     NeedDifferentInput,
  40     Success(Value),
  41 }
  42
  43 impl<'a> Output<'a> {
  44     fn write_to_stdout(&self, pretty: bool) -> Result<()> {
  45         let mut stdout = BufWriter::new(std::io::stdout().lock());
  46         if pretty {
  47             serde_json::to_writer_pretty(&mut stdout, self)?;
  48         } else {
  49             serde_json::to_writer(&mut stdout, self)?;
  50         }
  51         writeln!(stdout)?;
  52         stdout.flush()?;
  53         Ok(())
  54     }
  55     fn from_err(e: impl Into<eyre::Report>) -> Self {
  56         Self::Error(format!("{:?}", e.into()))
  57     }
  58 }
  59
  60 fn process_input(args: &Args, input: Value, exit_code: &mut ExitCode) -> Result<()> {
  61     GlobalState::scope(|| -> Result<()> {
  62         let function = match serde_json::from_value::<Function>(input) {
  63             Ok(v) => v,
  64             Err(e) => {
  65                 Output::from_err(e).write_to_stdout(args.pretty)?;
  66                 *exit_code = ExitCode::FAILURE;
  67                 return Ok(());
  68             }
  69         };
  70         if args.dump_input {
  71             Output::Info(Info::DumpInput(&function)).write_to_stdout(args.pretty)?;
  72         }
  73         todo!()
  74     })
  75 }
  76
  77 fn arbitrary_input(input_bytes: &[u8]) -> arbitrary::Result<Value> {
  78     GlobalState::scope(|| -> arbitrary::Result<Value> {
  79         Ok(
  80             serde_json::to_value(FnFields::arbitrary_take_rest(Unstructured::new(
  81                 &input_bytes,
  82             ))?)
  83             .expect("serialization shouldn't ever fail"),
  84         )
  85     })
  86 }
  87
  88 fn main() -> Result<ExitCode> {
  89     let args = Args::parse();
  90     let Args {
  91         dump_input: _,
  92         pretty,
  93         arbitrary,
  94     } = args;
  95     let stdin = std::io::stdin().lock();
  96     let mut exit_code = ExitCode::SUCCESS;
  97     if arbitrary {
  98         let mut input_bytes = vec![];
  99         stdin.take(0x10000).read_to_end(&mut input_bytes)?;
 100         if let Ok(input) = arbitrary_input(&input_bytes) {
 101             process_input(&args, input, &mut exit_code)?;
 102         } else {
 103             Output::NeedDifferentInput.write_to_stdout(pretty)?;
 104             exit_code = ExitCode::FAILURE;
 105         };
 106     } else {
 107         for input in Deserializer::from_reader(stdin).into_iter::<Value>() {
 108             process_input(&args, input?, &mut exit_code)?;
 109         }
 110     }
 111     Ok(exit_code)
 112 }