1 files changed, 41 insertions, 4 deletions

diff --git a/crates/ra_hir/src/code_model/src.rs b/crates/ra_hir/src/code_model/src.rs
index 32bd9c661..b9ffb0c7a 100644
--- a/crates/ra_hir/src/code_model/src.rs
+++ b/crates/ra_hir/src/code_model/src.rs
@@ -1,11 +1,15 @@
-use ra_syntax::ast;
+use ra_syntax::{
+    ast::{self, AstNode},
+    SyntaxNode,
+};
 
 use crate::{
-    ids::AstItemDef, AstDatabase, Const, DefDatabase, Enum, EnumVariant, FieldSource, Function,
-    HirFileId, MacroDef, Module, ModuleSource, Static, Struct, StructField, Trait, TypeAlias,
-    Union,
+    ids::AstItemDef, AstDatabase, Const, DefDatabase, Either, Enum, EnumVariant, FieldSource,
+    Function, HasBody, HirDatabase, HirFileId, MacroDef, Module, ModuleSource, Static, Struct,
+    StructField, Trait, TypeAlias, Union,
 };
 
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct Source<T> {
     pub file_id: HirFileId,
     pub ast: T,
@@ -16,6 +20,15 @@ pub trait HasSource {
     fn source(self, db: &(impl DefDatabase + AstDatabase)) -> Source<Self::Ast>;
 }
 
+impl<T> Source<T> {
+    pub(crate) fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
+        Source { file_id: self.file_id, ast: f(self.ast) }
+    }
+    pub(crate) fn file_syntax(&self, db: &impl AstDatabase) -> SyntaxNode {
+        db.parse_or_expand(self.file_id).expect("source created from invalid file")
+    }
+}
+
 /// NB: Module is !HasSource, because it has two source nodes at the same time:
 /// definition and declaration.
 impl Module {
@@ -108,3 +121,27 @@ impl HasSource for MacroDef {
         Source { file_id: self.id.0.file_id(), ast: self.id.0.to_node(db) }
     }
 }
+
+pub trait HasBodySource: HasBody + HasSource
+where
+    Self::Ast: AstNode,
+{
+    fn expr_source(
+        self,
+        db: &impl HirDatabase,
+        expr_id: crate::expr::ExprId,
+    ) -> Option<Source<Either<ast::Expr, ast::RecordField>>> {
+        let source_map = self.body_source_map(db);
+        let source_ptr = source_map.expr_syntax(expr_id)?;
+        let root = source_ptr.file_syntax(db);
+        let source = source_ptr.map(|ast| ast.map(|it| it.to_node(&root), |it| it.to_node(&root)));
+        Some(source)
+    }
+}
+
+impl<T> HasBodySource for T
+where
+    T: HasBody + HasSource,
+    T::Ast: AstNode,
+{
+}