Refactor argument UTF-8 checking into `rustc_driver::args::raw_args()`

3 files changed, 28 insertions, 23 deletions

diff --git a/compiler/rustc_driver_impl/src/args.rs b/compiler/rustc_driver_impl/src/args.rs
index 8c03f54bb59..28574457389 100644
--- a/compiler/rustc_driver_impl/src/args.rs
+++ b/compiler/rustc_driver_impl/src/args.rs
@@ -1,7 +1,4 @@
-use std::error;
-use std::fmt;
-use std::fs;
-use std::io;
+use std::{env, error, fmt, fs, io};
 
 use rustc_session::EarlyDiagCtxt;
 use rustc_span::ErrorGuaranteed;
@@ -116,6 +113,29 @@ pub fn arg_expand_all(
     result.map(|()| expander.finish())
 }
 
+/// Gets the raw unprocessed command-line arguments as Unicode strings, without doing any further
+/// processing (e.g., without `@file` expansion).
+///
+/// This function is identical to [`env::args()`] except that it emits an error when it encounters
+/// non-Unicode arguments instead of panicking.
+pub fn raw_args(early_dcx: &EarlyDiagCtxt) -> Result<Vec<String>, ErrorGuaranteed> {
+    let mut res = Ok(Vec::new());
+    for (i, arg) in env::args_os().enumerate() {
+        match arg.into_string() {
+            Ok(arg) => {
+                if let Ok(args) = &mut res {
+                    args.push(arg);
+                }
+            }
+            Err(arg) => {
+                res =
+                    Err(early_dcx.early_err(format!("argument {i} is not valid Unicode: {arg:?}")))
+            }
+        }
+    }
+    res
+}
+
 #[derive(Debug)]
 enum Error {
     Utf8Error(String),
diff --git a/compiler/rustc_driver_impl/src/lib.rs b/compiler/rustc_driver_impl/src/lib.rs
index 3831bbc1e15..cd296c9662a 100644
--- a/compiler/rustc_driver_impl/src/lib.rs
+++ b/compiler/rustc_driver_impl/src/lib.rs
@@ -1515,15 +1515,7 @@ pub fn main() -> ! {
     let mut callbacks = TimePassesCallbacks::default();
     let using_internal_features = install_ice_hook(DEFAULT_BUG_REPORT_URL, |_| ());
     let exit_code = catch_with_exit_code(|| {
-        let args = env::args_os()
-            .enumerate()
-            .map(|(i, arg)| {
-                arg.into_string().unwrap_or_else(|arg| {
-                    early_dcx.early_fatal(format!("argument {i} is not valid Unicode: {arg:?}"))
-                })
-            })
-            .collect::<Vec<_>>();
-        RunCompiler::new(&args, &mut callbacks)
+        RunCompiler::new(&args::raw_args(&early_dcx)?, &mut callbacks)
             .set_using_internal_features(using_internal_features)
             .run()
     });
diff --git a/src/librustdoc/lib.rs b/src/librustdoc/lib.rs
index faa92c9d15a..39d27b104cd 100644
--- a/src/librustdoc/lib.rs
+++ b/src/librustdoc/lib.rs
@@ -179,21 +179,14 @@ pub fn main() {
     rustc_driver::init_logger(&early_dcx, rustc_log::LoggerConfig::from_env("RUSTDOC_LOG"));
 
     let exit_code = rustc_driver::catch_with_exit_code(|| {
-        let args = env::args_os()
-            .enumerate()
-            .map(|(i, arg)| {
-                arg.into_string().unwrap_or_else(|arg| {
-                    early_dcx.early_fatal(format!("argument {i} is not valid Unicode: {arg:?}"))
-                })
-            })
-            .collect::<Vec<_>>();
-        main_args(&mut early_dcx, &args, using_internal_features)
+        let at_args = rustc_driver::args::raw_args(&early_dcx)?;
+        main_args(&mut early_dcx, &at_args, using_internal_features)
     });
     process::exit(exit_code);
 }
 
 fn init_logging(early_dcx: &EarlyDiagCtxt) {
-    let color_logs = match std::env::var("RUSTDOC_LOG_COLOR").as_deref() {
+    let color_logs = match env::var("RUSTDOC_LOG_COLOR").as_deref() {
         Ok("always") => true,
         Ok("never") => false,
         Ok("auto") | Err(VarError::NotPresent) => io::stdout().is_terminal(),