From d7ab3c77b3259e25e15d77b6939af4ba8d3be4c0 Mon Sep 17 00:00:00 2001
From: hyd-dev <yd-huang@outlook.com>
Date: Mon, 15 Mar 2021 17:57:53 +0800
Subject: [PATCH] Add `rustc_interface::interface::Config::parse_sess_created`

---
 compiler/rustc_driver/src/lib.rs          | 2 ++
 compiler/rustc_interface/src/interface.rs | 9 ++++++++-
 src/librustdoc/core.rs                    | 1 +
 src/librustdoc/doctest.rs                 | 1 +
 4 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_driver/src/lib.rs b/compiler/rustc_driver/src/lib.rs
index 79bb21b29fc..25a78041c00 100644
--- a/compiler/rustc_driver/src/lib.rs
+++ b/compiler/rustc_driver/src/lib.rs
@@ -215,6 +215,7 @@ fn run_compiler(
             diagnostic_output,
             stderr: None,
             lint_caps: Default::default(),
+            parse_sess_created: None,
             register_lints: None,
             override_queries: None,
             make_codegen_backend: make_codegen_backend.take().unwrap(),
@@ -298,6 +299,7 @@ fn run_compiler(
         diagnostic_output,
         stderr: None,
         lint_caps: Default::default(),
+        parse_sess_created: None,
         register_lints: None,
         override_queries: None,
         make_codegen_backend: make_codegen_backend.unwrap(),
diff --git a/compiler/rustc_interface/src/interface.rs b/compiler/rustc_interface/src/interface.rs
index 502e7155c2e..14bffb54e7a 100644
--- a/compiler/rustc_interface/src/interface.rs
+++ b/compiler/rustc_interface/src/interface.rs
@@ -142,6 +142,9 @@ pub struct Config {
 
     pub lint_caps: FxHashMap<lint::LintId, lint::Level>,
 
+    /// This is a callback from the driver that is called when [`ParseSess`] is created.
+    pub parse_sess_created: Option<Box<dyn FnOnce(&mut ParseSess) + Send>>,
+
     /// This is a callback from the driver that is called when we're registering lints;
     /// it is called during plugin registration when we have the LintStore in a non-shared state.
     ///
@@ -166,7 +169,7 @@ pub struct Config {
 
 pub fn create_compiler_and_run<R>(config: Config, f: impl FnOnce(&Compiler) -> R) -> R {
     let registry = &config.registry;
-    let (sess, codegen_backend) = util::create_session(
+    let (mut sess, codegen_backend) = util::create_session(
         config.opts,
         config.crate_cfg,
         config.diagnostic_output,
@@ -177,6 +180,10 @@ pub fn create_compiler_and_run<R>(config: Config, f: impl FnOnce(&Compiler) -> R
         registry.clone(),
     );
 
+    if let Some(parse_sess_created) = config.parse_sess_created {
+        parse_sess_created(&mut Lrc::get_mut(&mut sess).unwrap().parse_sess);
+    }
+
     let compiler = Compiler {
         sess,
         codegen_backend,
diff --git a/src/librustdoc/core.rs b/src/librustdoc/core.rs
index e9c32396e84..c5b5ab0f3d0 100644
--- a/src/librustdoc/core.rs
+++ b/src/librustdoc/core.rs
@@ -311,6 +311,7 @@ crate fn create_config(
         diagnostic_output: DiagnosticOutput::Default,
         stderr: None,
         lint_caps,
+        parse_sess_created: None,
         register_lints: Some(box crate::lint::register_lints),
         override_queries: Some(|_sess, providers, _external_providers| {
             // Most lints will require typechecking, so just don't run them.
diff --git a/src/librustdoc/doctest.rs b/src/librustdoc/doctest.rs
index 64375964e70..80e2e23eadd 100644
--- a/src/librustdoc/doctest.rs
+++ b/src/librustdoc/doctest.rs
@@ -95,6 +95,7 @@ crate fn run(options: Options) -> Result<(), ErrorReported> {
         diagnostic_output: DiagnosticOutput::Default,
         stderr: None,
         lint_caps,
+        parse_sess_created: None,
         register_lints: Some(box crate::lint::register_lints),
         override_queries: None,
         make_codegen_backend: None,