From 00798609b2c9b1cef61763c39e88f723be0e7a94 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Wed, 17 Feb 2021 03:55:02 +0000
Subject: [PATCH] prealloc writing done: --save-raw

---
 src/arg.rs    | 10 ++++++++--
 src/config.rs | 12 ++++++++++--
 src/main.rs   | 37 ++++++++++++++++++++++++++++++++-----
 3 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/src/arg.rs b/src/arg.rs
index ebc0c41..2f357ff 100644
--- a/src/arg.rs
+++ b/src/arg.rs
@@ -33,6 +33,7 @@ const OPTIONS_NORMAL: &'static [&'static str] = &[
     "-v			Verbose mode. Output extra information.",
     #[cfg(feature="inspect")] "--save <file>		Dump the collected data to this file for further inspection.",
     #[cfg(feature="inspect")] "-D			Dump the collected data to `stdout` (see `--save`.)",
+    #[cfg(feature="prealloc")] "--save-raw <file>	Dump the collected data to this file uncompressed.",
     "-			Stop parsing arguments, treat all the rest as paths.",
 ];
 
@@ -140,12 +141,17 @@ fn parse<I: IntoIterator<Item=String>>(args: I) -> eyre::Result<Mode>
 		    cfg.output_level = config::OutputLevel::Verbose;
 		},
 		#[cfg(feature="inspect")] "-D" => {
-		    cfg.serialise_output = Some(None);
+		    cfg.serialise_output = Some(config::OutputSerialisationMode::Stdout);
 		},
 		#[cfg(feature="inspect")] "--save" => {
 		    let file = args.next().ok_or(eyre!("`--save` expects a filename parameter"))
 			.with_suggestion(suggestion_intended_arg.clone())?;
-		    cfg.serialise_output = Some(Some(file.into()));
+		    cfg.serialise_output = Some(config::OutputSerialisationMode::File(file.into()));
+		},
+		#[cfg(feature="prealloc")] "--save-raw" => {
+		    let file = args.next().ok_or(eyre!("`--save-raw` expects a filename parameter"))
+			.with_suggestion(suggestion_intended_arg.clone())?;
+		    cfg.serialise_output = Some(config::OutputSerialisationMode::PreallocFile(file.into()));
 		},
 		"--recursive" => {
 		    let max = args.next().ok_or(eyre!("`--recursive` expects a parameter"))
diff --git a/src/config.rs b/src/config.rs
index 42bc62b..82376ef 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -96,6 +96,14 @@ impl From<usize> for Recursion
     }
 }
 
+#[cfg(feature="inspect")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum OutputSerialisationMode
+{
+    Stdout,
+    File(PathBuf),
+    #[cfg(feature="prealloc")] PreallocFile(PathBuf),
+}
 
 /// Configuration for this run
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -108,7 +116,7 @@ pub struct Config
     pub output_level: OutputLevel,
     
     #[cfg(feature="inspect")]
-    pub serialise_output: Option<Option<PathBuf>>, // Some(None) means dump to `stdout`
+    pub serialise_output: Option<OutputSerialisationMode>,
 }
 
 impl Config
@@ -131,7 +139,7 @@ impl Config
     #[inline] pub fn is_using_stdout(&self) -> bool
     {
 	#[cfg(feature="inspect")] {
-	    return if let Some(None) = self.serialise_output {
+	    return if let Some(OutputSerialisationMode::Stdout) = self.serialise_output {
 		true
 	    } else {
 		false
diff --git a/src/main.rs b/src/main.rs
index 80ad764..6b304b4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -72,10 +72,11 @@ async fn normal(cfg: config::Config) -> eyre::Result<()>
     match cfg.serialise_output.as_ref().map(|ser_out| {
 	type BoxedWrite = Box<dyn tokio::io::AsyncWrite + Unpin>;
 	use futures::FutureExt;
+	use config::OutputSerialisationMode;
 	match ser_out {
-	    Some(output_file) => {
+	    OutputSerialisationMode::File(output_file) => {
 		use tokio::fs::OpenOptions;
-		async move {
+		(async move {
 		    let stream = OpenOptions::new()
 			.write(true)
 			.truncate(true)
@@ -85,16 +86,42 @@ async fn normal(cfg: config::Config) -> eyre::Result<()>
 			.with_section(|| format!("{:?}", output_file).header("File was"))?;
 
 		    Ok::<BoxedWrite, eyre::Report>(Box::new(stream))
-		}.boxed()
+		}.boxed(), None)
+	    },
+	    OutputSerialisationMode::Stdout => (async move { Ok::<BoxedWrite, _>(Box::new(tokio::io::stdout())) }.boxed(), None),
+	    #[cfg(feature="prealloc")] OutputSerialisationMode::PreallocFile(output_file) => {
+		(async move {
+		    Ok::<BoxedWrite, _>(Box::new(tokio::io::sink())) // we use a sink as a shim stream since it will never be used when tuple item `.1` is Some()
+		}.boxed(), Some(output_file))
 	    },
-	    None => async move { Ok::<BoxedWrite, _>(Box::new(tokio::io::stdout())) }.boxed(),
 	}
     }) {
-	Some(stream_fut) => {
+	// We use tuple item `.1` here to indicate if we're in normal write mode.
+	// None -> normal
+	// Some(path) -> prealloc
+	Some((stream_fut, None)) => {
 	    let stream = stream_fut.await?; 
 	    serial::write_async(stream, &graph).await
 		.wrap_err(eyre!("Failed to serialise graph to stream"))?;
 	},
+	#[cfg(feature="prealloc")] Some((_task_fut, Some(output_file))) => {
+	    use tokio::fs::OpenOptions;
+	    let file = OpenOptions::new()
+		.write(true)
+		.read(true) //needed for map I think?
+		.truncate(true)
+		.create(true)
+		.open(&output_file).await
+		.wrap_err(eyre!("Failed to open file for mapping"))
+		.with_section(|| format!("{:?}", output_file).header("File was"))?;
+	    let mut file = file.into_std().await;
+	    tokio::task::spawn_blocking(move || {
+		serial::write_sync_map(&mut file, &graph)
+	    }).await.wrap_err(eyre!("Prealloc panicked while dumping"))
+		.with_section(|| format!("{:?}", output_file).header("File was"))?
+		.wrap_err(eyre!("Prealloc failed to dump graph to file"))
+		.with_section(|| format!("{:?}", output_file).header("File was"))?;
+	},
 	_ => (),
     }