1 files changed, 60 insertions, 0 deletions

diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs
index 5b1a4680b..1925db8b2 100644
--- a/crates/ra_assists/src/assist_context.rs
+++ b/crates/ra_assists/src/assist_context.rs
@@ -19,6 +19,7 @@ use crate::{
     assist_config::{AssistConfig, SnippetCap},
     Assist, AssistId, GroupLabel, ResolvedAssist,
 };
+use rustc_hash::FxHashMap;
 
 /// `AssistContext` allows to apply an assist or check if it could be applied.
 ///
@@ -138,6 +139,16 @@ impl Assists {
         let label = Assist::new(id, label.into(), None, target);
         self.add_impl(label, f)
     }
+    pub(crate) fn add_in_multiple_files(
+        &mut self,
+        id: AssistId,
+        label: impl Into<String>,
+        target: TextRange,
+        f: impl FnOnce(&mut AssistDirector),
+    ) -> Option<()> {
+        let label = Assist::new(id, label.into(), None, target);
+        self.add_impl_multiple_files(label, f)
+    }
     pub(crate) fn add_group(
         &mut self,
         group: &GroupLabel,
@@ -162,6 +173,31 @@ impl Assists {
         Some(())
     }
 
+    fn add_impl_multiple_files(
+        &mut self,
+        label: Assist,
+        f: impl FnOnce(&mut AssistDirector),
+    ) -> Option<()> {
+        if !self.resolve {
+            self.buf.push((label, None));
+            return None;
+        }
+        let mut director = AssistDirector::default();
+        f(&mut director);
+        let changes = director.finish();
+        let file_edits: Vec<SourceFileEdit> =
+            changes.into_iter().map(|mut change| change.source_file_edits.pop().unwrap()).collect();
+
+        let source_change = SourceChange {
+            source_file_edits: file_edits,
+            file_system_edits: vec![],
+            is_snippet: false,
+        };
+
+        self.buf.push((label, Some(source_change)));
+        Some(())
+    }
+
     fn finish(mut self) -> Vec<(Assist, Option<SourceChange>)> {
         self.buf.sort_by_key(|(label, _edit)| label.target.len());
         self.buf
@@ -255,3 +291,27 @@ impl AssistBuilder {
         res
     }
 }
+
+pub(crate) struct AssistDirector {
+    builders: FxHashMap<FileId, AssistBuilder>,
+}
+
+impl AssistDirector {
+    pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) {
+        let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id));
+        f(&mut builder);
+    }
+
+    fn finish(self) -> Vec<SourceChange> {
+        self.builders
+            .into_iter()
+            .map(|(_, builder)| builder.finish())
+            .collect::<Vec<SourceChange>>()
+    }
+}
+
+impl Default for AssistDirector {
+    fn default() -> Self {
+        AssistDirector { builders: FxHashMap::default() }
+    }
+}