8 files changed, 357 insertions, 5 deletions

diff --git a/Cargo.lock b/Cargo.lock
index fe71b9971..a3cfe5dc4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -348,6 +348,14 @@ dependencies = [
 ]
 
 [[package]]
+name = "expect"
+version = "0.1.0"
+dependencies = [
+ "once_cell",
+ "stdx",
+]
+
+[[package]]
 name = "filetime"
 version = "0.2.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/crates/expect/Cargo.toml b/crates/expect/Cargo.toml
new file mode 100644
index 000000000..09eb57a43
--- /dev/null
+++ b/crates/expect/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "expect"
+version = "0.1.0"
+authors = ["rust-analyzer developers"]
+edition = "2018"
+
+[dependencies]
+once_cell = "1"
+stdx = { path = "../stdx" }
diff --git a/crates/expect/src/lib.rs b/crates/expect/src/lib.rs
new file mode 100644
index 000000000..08d0eafdf
--- /dev/null
+++ b/crates/expect/src/lib.rs
@@ -0,0 +1,308 @@
+//! Snapshot testing library, see
+//! https://github.com/rust-analyzer/rust-analyzer/pull/5101
+use std::{
+    collections::HashMap,
+    env, fmt, fs,
+    ops::Range,
+    path::{Path, PathBuf},
+    sync::Mutex,
+};
+
+use once_cell::sync::Lazy;
+use stdx::{lines_with_ends, trim_indent};
+
+const HELP: &str = "
+You can update all `expect![[]]` tests by:
+
+    env UPDATE_EXPECT=1 cargo test
+
+To update a single test, place the cursor on `expect` token and use `run` feature of rust-analyzer.
+";
+
+fn update_expect() -> bool {
+    env::var("UPDATE_EXPECT").is_ok()
+}
+
+/// expect![[""]]
+#[macro_export]
+macro_rules! expect {
+    [[$lit:literal]] => {$crate::Expect {
+        file: file!(),
+        line: line!(),
+        column: column!(),
+        data: $lit,
+    }};
+    [[]] => { $crate::expect![[""]] };
+}
+
+#[derive(Debug)]
+pub struct Expect {
+    pub file: &'static str,
+    pub line: u32,
+    pub column: u32,
+    pub data: &'static str,
+}
+
+impl Expect {
+    pub fn assert_eq(&self, actual: &str) {
+        let trimmed = self.trimmed();
+        if &trimmed == actual {
+            return;
+        }
+        Runtime::fail(self, &trimmed, actual);
+    }
+    pub fn assert_debug_eq(&self, actual: &impl fmt::Debug) {
+        let actual = format!("{:#?}\n", actual);
+        self.assert_eq(&actual)
+    }
+
+    fn trimmed(&self) -> String {
+        if !self.data.contains('\n') {
+            return self.data.to_string();
+        }
+        trim_indent(self.data)
+    }
+
+    fn locate(&self, file: &str) -> Location {
+        let mut target_line = None;
+        let mut line_start = 0;
+        for (i, line) in lines_with_ends(file).enumerate() {
+            if i == self.line as usize - 1 {
+                let pat = "expect![[";
+                let offset = line.find(pat).unwrap();
+                let literal_start = line_start + offset + pat.len();
+                let indent = line.chars().take_while(|&it| it == ' ').count();
+                target_line = Some((literal_start, indent));
+                break;
+            }
+            line_start += line.len();
+        }
+        let (literal_start, line_indent) = target_line.unwrap();
+        let literal_length = file[literal_start..].find("]]").unwrap();
+        let literal_range = literal_start..literal_start + literal_length;
+        Location { line_indent, literal_range }
+    }
+}
+
+#[derive(Default)]
+struct Runtime {
+    help_printed: bool,
+    per_file: HashMap<&'static str, FileRuntime>,
+}
+static RT: Lazy<Mutex<Runtime>> = Lazy::new(Default::default);
+
+impl Runtime {
+    fn fail(expect: &Expect, expected: &str, actual: &str) {
+        let mut rt = RT.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
+        let mut updated = "";
+        if update_expect() {
+            updated = " (updated)";
+            rt.per_file
+                .entry(expect.file)
+                .or_insert_with(|| FileRuntime::new(expect))
+                .update(expect, actual);
+        }
+        let print_help = !rt.help_printed && !update_expect();
+        rt.help_printed = true;
+
+        let help = if print_help { HELP } else { "" };
+        panic!(
+            "\n
+error: expect test failed{}
+  --> {}:{}:{}
+{}
+Expect:
+----
+{}
+----
+
+Actual:
+----
+{}
+----
+",
+            updated, expect.file, expect.line, expect.column, help, expected, actual
+        )
+    }
+}
+
+struct FileRuntime {
+    path: PathBuf,
+    original_text: String,
+    patchwork: Patchwork,
+}
+
+impl FileRuntime {
+    fn new(expect: &Expect) -> FileRuntime {
+        let path = workspace_root().join(expect.file);
+        let original_text = fs::read_to_string(&path).unwrap();
+        let patchwork = Patchwork::new(original_text.clone());
+        FileRuntime { path, original_text, patchwork }
+    }
+    fn update(&mut self, expect: &Expect, actual: &str) {
+        let loc = expect.locate(&self.original_text);
+        let patch = format_patch(loc.line_indent.clone(), actual);
+        self.patchwork.patch(loc.literal_range, &patch);
+        fs::write(&self.path, &self.patchwork.text).unwrap()
+    }
+}
+
+#[derive(Debug)]
+struct Location {
+    line_indent: usize,
+    literal_range: Range<usize>,
+}
+
+#[derive(Debug)]
+struct Patchwork {
+    text: String,
+    indels: Vec<(Range<usize>, usize)>,
+}
+
+impl Patchwork {
+    fn new(text: String) -> Patchwork {
+        Patchwork { text, indels: Vec::new() }
+    }
+    fn patch(&mut self, mut range: Range<usize>, patch: &str) {
+        self.indels.push((range.clone(), patch.len()));
+        self.indels.sort_by_key(|(delete, _insert)| delete.start);
+
+        let (delete, insert) = self
+            .indels
+            .iter()
+            .take_while(|(delete, _)| delete.start < range.start)
+            .map(|(delete, insert)| (delete.end - delete.start, insert))
+            .fold((0usize, 0usize), |(x1, y1), (x2, y2)| (x1 + x2, y1 + y2));
+
+        for pos in &mut [&mut range.start, &mut range.end] {
+            **pos += insert;
+            **pos -= delete
+        }
+
+        self.text.replace_range(range, &patch);
+    }
+}
+
+fn format_patch(line_indent: usize, patch: &str) -> String {
+    let mut max_hashes = 0;
+    let mut cur_hashes = 0;
+    for byte in patch.bytes() {
+        if byte != b'#' {
+            cur_hashes = 0;
+            continue;
+        }
+        cur_hashes += 1;
+        max_hashes = max_hashes.max(cur_hashes);
+    }
+    let hashes = &"#".repeat(max_hashes + 1);
+    let indent = &" ".repeat(line_indent);
+    let is_multiline = patch.contains('\n');
+
+    let mut buf = String::new();
+    buf.push('r');
+    buf.push_str(hashes);
+    buf.push('"');
+    if is_multiline {
+        buf.push('\n');
+    }
+    let mut final_newline = false;
+    for line in lines_with_ends(patch) {
+        if is_multiline {
+            buf.push_str(indent);
+            buf.push_str("    ");
+        }
+        buf.push_str(line);
+        final_newline = line.ends_with('\n');
+    }
+    if final_newline {
+        buf.push_str(indent);
+    }
+    buf.push('"');
+    buf.push_str(hashes);
+    buf
+}
+
+fn workspace_root() -> PathBuf {
+    Path::new(
+        &env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| env!("CARGO_MANIFEST_DIR").to_owned()),
+    )
+    .ancestors()
+    .nth(2)
+    .unwrap()
+    .to_path_buf()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_expect_macro() {
+        let empty = expect![[]];
+        expect![[r#"
+            Expect {
+                file: "crates/expect/src/lib.rs",
+                line: 241,
+                column: 21,
+                data: "",
+            }
+        "#]]
+        .assert_debug_eq(&empty);
+
+        let expect = expect![["
+            hello
+            world
+        "]];
+        expect![[r#"
+            Expect {
+                file: "crates/expect/src/lib.rs",
+                line: 252,
+                column: 22,
+                data: "\n            hello\n            world\n        ",
+            }
+        "#]]
+        .assert_debug_eq(&expect);
+    }
+
+    #[test]
+    fn test_format_patch() {
+        let patch = format_patch(0, "hello\nworld\n");
+        expect![[r##"
+            r#"
+                hello
+                world
+            "#"##]]
+        .assert_eq(&patch);
+
+        let patch = format_patch(4, "single line");
+        expect![[r##"r#"single line"#"##]].assert_eq(&patch);
+    }
+
+    #[test]
+    fn test_patchwork() {
+        let mut patchwork = Patchwork::new("one two three".to_string());
+        patchwork.patch(4..7, "zwei");
+        patchwork.patch(0..3, "один");
+        patchwork.patch(8..13, "3");
+        expect![[r#"
+            Patchwork {
+                text: "один zwei 3",
+                indels: [
+                    (
+                        0..3,
+                        8,
+                    ),
+                    (
+                        4..7,
+                        4,
+                    ),
+                    (
+                        8..13,
+                        1,
+                    ),
+                ],
+            }
+        "#]]
+        .assert_debug_eq(&patchwork);
+    }
+}
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 38e3c3324..4d0684b2a 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -23,7 +23,7 @@ use ra_ide::{
 };
 use ra_prof::profile;
 use ra_project_model::TargetKind;
-use ra_syntax::{AstNode, SyntaxKind, TextRange, TextSize};
+use ra_syntax::{algo, ast, AstNode, SyntaxKind, TextRange, TextSize};
 use serde::{Deserialize, Serialize};
 use serde_json::to_value;
 use stdx::{format_to, split_delim};
@@ -407,8 +407,21 @@ pub(crate) fn handle_runnables(
     let file_id = from_proto::file_id(&snap, &params.text_document.uri)?;
     let line_index = snap.analysis.file_line_index(file_id)?;
     let offset = params.position.map(|it| from_proto::offset(&line_index, it));
-    let mut res = Vec::new();
     let cargo_spec = CargoTargetSpec::for_file(&snap, file_id)?;
+
+    let expect_test = match offset {
+        Some(offset) => {
+            let source_file = snap.analysis.parse(file_id)?;
+            algo::find_node_at_offset::<ast::MacroCall>(source_file.syntax(), offset)
+                .and_then(|it| it.path())
+                .and_then(|it| it.segment())
+                .and_then(|it| it.name_ref())
+                .map_or(false, |it| it.text() == "expect")
+        }
+        None => false,
+    };
+
+    let mut res = Vec::new();
     for runnable in snap.analysis.runnables(file_id)? {
         if let Some(offset) = offset {
             if !runnable.nav.full_range().contains_inclusive(offset) {
@@ -418,8 +431,12 @@ pub(crate) fn handle_runnables(
         if should_skip_target(&runnable, cargo_spec.as_ref()) {
             continue;
         }
-
-        res.push(to_proto::runnable(&snap, file_id, runnable)?);
+        let mut runnable = to_proto::runnable(&snap, file_id, runnable)?;
+        if expect_test {
+            runnable.label = format!("{} + expect", runnable.label);
+            runnable.args.expect_test = Some(true);
+        }
+        res.push(runnable);
     }
 
     // Add `cargo check` and `cargo test` for the whole package
@@ -438,6 +455,7 @@ pub(crate) fn handle_runnables(
                             spec.package.clone(),
                         ],
                         executable_args: Vec::new(),
+                        expect_test: None,
                     },
                 })
             }
@@ -451,6 +469,7 @@ pub(crate) fn handle_runnables(
                     workspace_root: None,
                     cargo_args: vec!["check".to_string(), "--workspace".to_string()],
                     executable_args: Vec::new(),
+                    expect_test: None,
                 },
             });
         }
diff --git a/crates/rust-analyzer/src/lsp_ext.rs b/crates/rust-analyzer/src/lsp_ext.rs
index 1371f6cb4..1befe678c 100644
--- a/crates/rust-analyzer/src/lsp_ext.rs
+++ b/crates/rust-analyzer/src/lsp_ext.rs
@@ -161,6 +161,8 @@ pub struct CargoRunnable {
     pub cargo_args: Vec<String>,
     // stuff after --
     pub executable_args: Vec<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub expect_test: Option<bool>,
 }
 
 pub enum InlayHints {}
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index f6cb8e4bb..a03222ae9 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -666,6 +666,7 @@ pub(crate) fn runnable(
             workspace_root: workspace_root.map(|it| it.into()),
             cargo_args,
             executable_args,
+            expect_test: None,
         },
     })
 }
diff --git a/editors/code/src/lsp_ext.ts b/editors/code/src/lsp_ext.ts
index e16ea799c..fdb99956b 100644
--- a/editors/code/src/lsp_ext.ts
+++ b/editors/code/src/lsp_ext.ts
@@ -60,6 +60,7 @@ export interface Runnable {
         workspaceRoot?: string;
         cargoArgs: string[];
         executableArgs: string[];
+        expectTest?: boolean;
     };
 }
 export const runnables = new lc.RequestType<RunnablesParams, Runnable[], void>("experimental/runnables");
diff --git a/editors/code/src/run.ts b/editors/code/src/run.ts
index 766b05112..e1430e31f 100644
--- a/editors/code/src/run.ts
+++ b/editors/code/src/run.ts
@@ -108,12 +108,16 @@ export async function createTask(runnable: ra.Runnable, config: Config): Promise
     if (runnable.args.executableArgs.length > 0) {
         args.push('--', ...runnable.args.executableArgs);
     }
+    const env: { [key: string]: string } = { "RUST_BACKTRACE": "short" };
+    if (runnable.args.expectTest) {
+        env["UPDATE_EXPECT"] = "1";
+    }
     const definition: tasks.CargoTaskDefinition = {
         type: tasks.TASK_TYPE,
         command: args[0], // run, test, etc...
         args: args.slice(1),
         cwd: runnable.args.workspaceRoot,
-        env: Object.assign({}, process.env as { [key: string]: string }, { "RUST_BACKTRACE": "short" }),
+        env: Object.assign({}, process.env as { [key: string]: string }, env),
     };
 
     const target = vscode.workspace.workspaceFolders![0]; // safe, see main activate()