3 files changed, 167 insertions, 40 deletions

diff --git a/crates/ra_hir_def/src/adt.rs b/crates/ra_hir_def/src/adt.rs
index 2bd18255c..ae99afe39 100644
--- a/crates/ra_hir_def/src/adt.rs
+++ b/crates/ra_hir_def/src/adt.rs
@@ -2,13 +2,17 @@
 
 use std::sync::Arc;
 
-use hir_expand::name::{AsName, Name};
-use ra_arena::Arena;
+use hir_expand::{
+    either::Either,
+    name::{AsName, Name},
+    Source,
+};
+use ra_arena::{map::ArenaMap, Arena};
 use ra_syntax::ast::{self, NameOwner, TypeAscriptionOwner};
 
 use crate::{
-    db::DefDatabase2, type_ref::TypeRef, AstItemDef, EnumId, LocalEnumVariantId,
-    LocalStructFieldId, StructOrUnionId,
+    db::DefDatabase2, trace::Trace, type_ref::TypeRef, AstItemDef, EnumId, HasChildSource,
+    LocalEnumVariantId, LocalStructFieldId, StructOrUnionId, VariantId,
 };
 
 /// Note that we use `StructData` for unions as well!
@@ -61,17 +65,9 @@ impl EnumData {
     pub(crate) fn enum_data_query(db: &impl DefDatabase2, e: EnumId) -> Arc<EnumData> {
         let src = e.source(db);
         let name = src.value.name().map(|n| n.as_name());
-        let variants = src
-            .value
-            .variant_list()
-            .into_iter()
-            .flat_map(|it| it.variants())
-            .map(|var| EnumVariantData {
-                name: var.name().map(|it| it.as_name()),
-                variant_data: Arc::new(VariantData::new(var.kind())),
-            })
-            .collect();
-        Arc::new(EnumData { name, variants })
+        let mut trace = Trace::new_for_arena();
+        lower_enum(&mut trace, &src.value);
+        Arc::new(EnumData { name, variants: trace.into_arena() })
     }
 
     pub(crate) fn variant(&self, name: &Name) -> Option<LocalEnumVariantId> {
@@ -80,38 +76,109 @@ impl EnumData {
     }
 }
 
+impl HasChildSource for EnumId {
+    type ChildId = LocalEnumVariantId;
+    type Value = ast::EnumVariant;
+    fn child_source(&self, db: &impl DefDatabase2) -> Source<ArenaMap<Self::ChildId, Self::Value>> {
+        let src = self.source(db);
+        let mut trace = Trace::new_for_map();
+        lower_enum(&mut trace, &src.value);
+        src.with_value(trace.into_map())
+    }
+}
+
+fn lower_enum(
+    trace: &mut Trace<LocalEnumVariantId, EnumVariantData, ast::EnumVariant>,
+    ast: &ast::EnumDef,
+) {
+    for var in ast.variant_list().into_iter().flat_map(|it| it.variants()) {
+        trace.alloc(
+            || var.clone(),
+            || EnumVariantData {
+                name: var.name().map(|it| it.as_name()),
+                variant_data: Arc::new(VariantData::new(var.kind())),
+            },
+        )
+    }
+}
+
 impl VariantData {
     fn new(flavor: ast::StructKind) -> Self {
-        match flavor {
-            ast::StructKind::Tuple(fl) => {
-                let fields = fl
-                    .fields()
-                    .enumerate()
-                    .map(|(i, fd)| StructFieldData {
-                        name: Name::new_tuple_field(i),
-                        type_ref: TypeRef::from_ast_opt(fd.type_ref()),
-                    })
-                    .collect();
-                VariantData::Tuple(fields)
-            }
-            ast::StructKind::Record(fl) => {
-                let fields = fl
-                    .fields()
-                    .map(|fd| StructFieldData {
-                        name: fd.name().map(|n| n.as_name()).unwrap_or_else(Name::missing),
-                        type_ref: TypeRef::from_ast_opt(fd.ascribed_type()),
-                    })
-                    .collect();
-                VariantData::Record(fields)
-            }
-            ast::StructKind::Unit => VariantData::Unit,
+        let mut trace = Trace::new_for_arena();
+        match lower_struct(&mut trace, &flavor) {
+            StructKind::Tuple => VariantData::Tuple(trace.into_arena()),
+            StructKind::Record => VariantData::Record(trace.into_arena()),
+            StructKind::Unit => VariantData::Unit,
         }
     }
 
     pub fn fields(&self) -> Option<&Arena<LocalStructFieldId, StructFieldData>> {
-        match self {
+        match &self {
             VariantData::Record(fields) | VariantData::Tuple(fields) => Some(fields),
             _ => None,
         }
     }
 }
+
+impl HasChildSource for VariantId {
+    type ChildId = LocalStructFieldId;
+    type Value = Either<ast::TupleFieldDef, ast::RecordFieldDef>;
+
+    fn child_source(&self, db: &impl DefDatabase2) -> Source<ArenaMap<Self::ChildId, Self::Value>> {
+        let src = match self {
+            VariantId::EnumVariantId(it) => {
+                // I don't really like the fact that we call into parent source
+                // here, this might add to more queries then necessary.
+                let src = it.parent.child_source(db);
+                src.map(|map| map[it.local_id].kind())
+            }
+            VariantId::StructId(it) => it.0.source(db).map(|it| it.kind()),
+        };
+        let mut trace = Trace::new_for_map();
+        lower_struct(&mut trace, &src.value);
+        src.with_value(trace.into_map())
+    }
+}
+
+enum StructKind {
+    Tuple,
+    Record,
+    Unit,
+}
+
+fn lower_struct(
+    trace: &mut Trace<
+        LocalStructFieldId,
+        StructFieldData,
+        Either<ast::TupleFieldDef, ast::RecordFieldDef>,
+    >,
+    ast: &ast::StructKind,
+) -> StructKind {
+    match ast {
+        ast::StructKind::Tuple(fl) => {
+            for (i, fd) in fl.fields().enumerate() {
+                trace.alloc(
+                    || Either::A(fd.clone()),
+                    || StructFieldData {
+                        name: Name::new_tuple_field(i),
+                        type_ref: TypeRef::from_ast_opt(fd.type_ref()),
+                    },
+                )
+            }
+            StructKind::Tuple
+        }
+        ast::StructKind::Record(fl) => {
+            for fd in fl.fields() {
+                trace.alloc(
+                    || Either::B(fd.clone()),
+                    || StructFieldData {
+                        name: fd.name().map(|n| n.as_name()).unwrap_or_else(Name::missing),
+                        type_ref: TypeRef::from_ast_opt(fd.ascribed_type()),
+                    },
+                )
+            }
+            StructKind::Record
+        }
+        ast::StructKind::Unit => StructKind::Unit,
+    }
+}
diff --git a/crates/ra_hir_def/src/lib.rs b/crates/ra_hir_def/src/lib.rs
index 3a0420da0..2edf743ab 100644
--- a/crates/ra_hir_def/src/lib.rs
+++ b/crates/ra_hir_def/src/lib.rs
@@ -20,6 +20,8 @@ pub mod generics;
 pub mod resolver;
 pub mod data;
 
+mod trace;
+
 #[cfg(test)]
 mod test_db;
 #[cfg(test)]
@@ -31,7 +33,7 @@ pub mod nameres;
 use std::hash::{Hash, Hasher};
 
 use hir_expand::{ast_id_map::FileAstId, db::AstDatabase, AstId, HirFileId, Source};
-use ra_arena::{impl_arena_id, RawId};
+use ra_arena::{impl_arena_id, map::ArenaMap, RawId};
 use ra_db::{salsa, CrateId, FileId};
 use ra_syntax::{ast, AstNode, SyntaxNode};
 
@@ -550,3 +552,12 @@ impl HasSource for ConstLoc {
         Source::new(self.ast_id.file_id(), node)
     }
 }
+
+pub trait HasChildSource {
+    type ChildId;
+    type Value;
+    fn child_source(
+        &self,
+        db: &impl db::DefDatabase2,
+    ) -> Source<ArenaMap<Self::ChildId, Self::Value>>;
+}
diff --git a/crates/ra_hir_def/src/trace.rs b/crates/ra_hir_def/src/trace.rs
new file mode 100644
index 000000000..fc26f5a48
--- /dev/null
+++ b/crates/ra_hir_def/src/trace.rs
@@ -0,0 +1,49 @@
+//! Trace is a pretty niche data structure which is used when lowering a CST
+//! into HIR.
+//!
+//! Lowering process calculates two bits of information:
+//! * the lowered syntax itself
+//! * a mapping between lowered syntax and original syntax
+//!
+//! Due to the way salsa works, the mapping is usually hot lava, as it contains
+//! absolute offsets. The `Trace` structure (inspired, at least in name, by
+//! Kotlin's `BindingTrace`) allows use the same code to compute both
+//! projections.
+use ra_arena::{map::ArenaMap, Arena, ArenaId, RawId};
+
+pub(crate) struct Trace<ID: ArenaId, T, V> {
+    for_arena: bool,
+    arena: Arena<ID, T>,
+    map: ArenaMap<ID, V>,
+    len: u32,
+}
+
+impl<ID: ra_arena::ArenaId, T, V> Trace<ID, T, V> {
+    pub(crate) fn new_for_arena() -> Trace<ID, T, V> {
+        Trace { for_arena: true, arena: Arena::default(), map: ArenaMap::default(), len: 0 }
+    }
+
+    pub(crate) fn new_for_map() -> Trace<ID, T, V> {
+        Trace { for_arena: false, arena: Arena::default(), map: ArenaMap::default(), len: 0 }
+    }
+
+    pub(crate) fn alloc(&mut self, value: impl Fn() -> V, data: impl Fn() -> T) {
+        if self.for_arena {
+            self.arena.alloc(data());
+        } else {
+            let id = ID::from_raw(RawId::from(self.len));
+            self.len += 1;
+            self.map.insert(id, value());
+        }
+    }
+
+    pub(crate) fn into_arena(self) -> Arena<ID, T> {
+        assert!(self.for_arena);
+        self.arena
+    }
+
+    pub(crate) fn into_map(self) -> ArenaMap<ID, V> {
+        assert!(!self.for_arena);
+        self.map
+    }
+}