From 99ce43f042ba68eb31d91fdac7b4951af2c76044 Mon Sep 17 00:00:00 2001
From: Tianyu Yao <skyyao@meta.com>
Date: Wed, 21 Jun 2023 11:30:16 -0700
Subject: [PATCH] Inject changed schema files when loading from saved state

Reviewed By: kassens

Differential Revision: D46753491

fbshipit-source-id: 381f9c601cbdcd11b336cc45eb664476a3752542
---
 compiler/crates/relay-compiler/src/config.rs             | 8 ++++++++
 compiler/crates/relay-compiler/src/file_source/mod.rs    | 2 +-
 .../src/file_source/watchman_file_source.rs              | 9 +++++++++
 compiler/crates/relay-compiler/src/lib.rs                | 2 ++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/compiler/crates/relay-compiler/src/config.rs b/compiler/crates/relay-compiler/src/config.rs
index 478dc4097a590..2188c406e218e 100644
--- a/compiler/crates/relay-compiler/src/config.rs
+++ b/compiler/crates/relay-compiler/src/config.rs
@@ -57,6 +57,7 @@ use crate::build_project::artifact_writer::ArtifactFileWriter;
 use crate::build_project::artifact_writer::ArtifactWriter;
 use crate::build_project::generate_extra_artifacts::GenerateExtraArtifactsFn;
 use crate::build_project::AdditionalValidations;
+use crate::compiler_state::CompilerState;
 use crate::compiler_state::ProjectName;
 use crate::compiler_state::ProjectSet;
 use crate::errors::ConfigValidationError;
@@ -77,6 +78,9 @@ type PostArtifactsWriter = Box<
 type OperationPersisterCreator =
     Box<dyn Fn(&ProjectConfig) -> Option<Box<dyn OperationPersister + Send + Sync>> + Send + Sync>;
 
+type UpdateCompilerStateFromSavedState =
+    Option<Box<dyn Fn(&mut CompilerState, &Config) + Send + Sync>>;
+
 /// The full compiler config. This is a combination of:
 /// - the configuration file
 /// - the absolute path to the root of the compiled projects
@@ -149,6 +153,9 @@ pub struct Config {
     /// The async function is called before the compiler connects to the file
     /// source.
     pub initialize_resources: Option<Box<dyn Fn() + Send + Sync>>,
+
+    /// Runs in `try_saved_state` when the compiler state is initialized from saved state.
+    pub update_compiler_state_from_saved_state: UpdateCompilerStateFromSavedState,
 }
 
 pub enum FileSourceKind {
@@ -410,6 +417,7 @@ Example file:
             custom_transforms: None,
             export_persisted_query_ids_to_file: None,
             initialize_resources: None,
+            update_compiler_state_from_saved_state: None,
         };
 
         let mut validation_errors = Vec::new();
diff --git a/compiler/crates/relay-compiler/src/file_source/mod.rs b/compiler/crates/relay-compiler/src/file_source/mod.rs
index 7d5b0f2b659be..6dd295cea8cb3 100644
--- a/compiler/crates/relay-compiler/src/file_source/mod.rs
+++ b/compiler/crates/relay-compiler/src/file_source/mod.rs
@@ -35,7 +35,7 @@ pub use source_control_update_status::SourceControlUpdateStatus;
 pub use watchman_client::prelude::Clock;
 use watchman_file_source::WatchmanFileSource;
 
-use self::external_file_source::ExternalFileSourceResult;
+pub use self::external_file_source::ExternalFileSourceResult;
 pub use self::extract_graphql::extract_javascript_features_from_file;
 pub use self::extract_graphql::source_for_location;
 pub use self::extract_graphql::FsSourceReader;
diff --git a/compiler/crates/relay-compiler/src/file_source/watchman_file_source.rs b/compiler/crates/relay-compiler/src/file_source/watchman_file_source.rs
index 6954e4958d00e..685cf95d80ab8 100644
--- a/compiler/crates/relay-compiler/src/file_source/watchman_file_source.rs
+++ b/compiler/crates/relay-compiler/src/file_source/watchman_file_source.rs
@@ -101,6 +101,7 @@ impl<'config> WatchmanFileSource<'config> {
                     saved_state_config.clone(),
                     saved_state_loader.as_ref(),
                     saved_state_version,
+                    self.config,
                 )
                 .await
             {
@@ -241,6 +242,7 @@ impl<'config> WatchmanFileSource<'config> {
         saved_state_config: ScmAwareClockData,
         saved_state_loader: &'_ (dyn SavedStateLoader + Send + Sync),
         saved_state_version: &str,
+        config: &Config,
     ) -> std::result::Result<Result<CompilerState>, &'static str> {
         let scm_since = Clock::ScmAware(FatClockData {
             clock: ClockSpec::null(),
@@ -296,6 +298,13 @@ impl<'config> WatchmanFileSource<'config> {
             .write()
             .unwrap()
             .push(file_source_result);
+
+        if let Some(update_compiler_state_from_saved_state) =
+            &config.update_compiler_state_from_saved_state
+        {
+            update_compiler_state_from_saved_state(&mut compiler_state, config);
+        }
+
         if let Err(parse_error) = perf_logger_event.time("merge_file_source_changes", || {
             compiler_state.merge_file_source_changes(self.config, perf_logger, true)
         }) {
diff --git a/compiler/crates/relay-compiler/src/lib.rs b/compiler/crates/relay-compiler/src/lib.rs
index 040e0d3ff0bd7..e91730b8afcc6 100644
--- a/compiler/crates/relay-compiler/src/lib.rs
+++ b/compiler/crates/relay-compiler/src/lib.rs
@@ -55,6 +55,8 @@ pub use config::ProjectConfig;
 pub use config::RemotePersistConfig;
 pub use config::SchemaLocation;
 pub use file_source::source_for_location;
+pub use file_source::ExternalFileSourceResult;
+pub use file_source::File;
 pub use file_source::FileCategorizer;
 pub use file_source::FileGroup;
 pub use file_source::FileSource;