5 files changed, 193 insertions, 83 deletions

diff --git a/crates/ra_hir/src/code_model/src.rs b/crates/ra_hir/src/code_model/src.rs
index fc466c1f0..a45c062bf 100644
--- a/crates/ra_hir/src/code_model/src.rs
+++ b/crates/ra_hir/src/code_model/src.rs
@@ -1,13 +1,13 @@
 //! FIXME: write short doc here
 
-use hir_def::{HasSource as _, Lookup};
+use hir_def::{HasChildSource, HasSource as _, Lookup, VariantId};
 use ra_syntax::ast::{self, AstNode};
 
 use crate::{
     db::{AstDatabase, DefDatabase, HirDatabase},
     ids::AstItemDef,
     Const, Either, Enum, EnumVariant, FieldSource, Function, HasBody, HirFileId, MacroDef, Module,
-    ModuleSource, Static, Struct, StructField, Trait, TypeAlias, Union, VariantDef,
+    ModuleSource, Static, Struct, StructField, Trait, TypeAlias, Union,
 };
 
 pub use hir_expand::Source;
@@ -46,33 +46,12 @@ impl Module {
 impl HasSource for StructField {
     type Ast = FieldSource;
     fn source(self, db: &(impl DefDatabase + AstDatabase)) -> Source<FieldSource> {
-        let var_data = self.parent.variant_data(db);
-        let fields = var_data.fields().unwrap();
-        let ss;
-        let es;
-        let (file_id, struct_kind) = match self.parent {
-            VariantDef::Struct(s) => {
-                ss = s.source(db);
-                (ss.file_id, ss.value.kind())
-            }
-            VariantDef::EnumVariant(e) => {
-                es = e.source(db);
-                (es.file_id, es.value.kind())
-            }
-        };
-
-        let field_sources = match struct_kind {
-            ast::StructKind::Tuple(fl) => fl.fields().map(|it| FieldSource::Pos(it)).collect(),
-            ast::StructKind::Record(fl) => fl.fields().map(|it| FieldSource::Named(it)).collect(),
-            ast::StructKind::Unit => Vec::new(),
-        };
-        let value = field_sources
-            .into_iter()
-            .zip(fields.iter())
-            .find(|(_syntax, (id, _))| *id == self.id)
-            .unwrap()
-            .0;
-        Source { file_id, value }
+        let var = VariantId::from(self.parent);
+        let src = var.child_source(db);
+        src.map(|it| match it[self.id].clone() {
+            Either::A(it) => FieldSource::Pos(it),
+            Either::B(it) => FieldSource::Named(it),
+        })
     }
 }
 impl HasSource for Struct {
@@ -96,18 +75,7 @@ impl HasSource for Enum {
 impl HasSource for EnumVariant {
     type Ast = ast::EnumVariant;
     fn source(self, db: &(impl DefDatabase + AstDatabase)) -> Source<ast::EnumVariant> {
-        let enum_data = db.enum_data(self.parent.id);
-        let src = self.parent.id.source(db);
-        let value = src
-            .value
-            .variant_list()
-            .into_iter()
-            .flat_map(|it| it.variants())
-            .zip(enum_data.variants.iter())
-            .find(|(_syntax, (id, _))| *id == self.id)
-            .unwrap()
-            .0;
-        Source { file_id: src.file_id, value }
+        self.parent.id.child_source(db).map(|map| map[self.id].clone())
     }
 }
 impl HasSource for Function {
diff --git a/crates/ra_hir/src/from_id.rs b/crates/ra_hir/src/from_id.rs
index e8ed04056..8900fc1f2 100644
--- a/crates/ra_hir/src/from_id.rs
+++ b/crates/ra_hir/src/from_id.rs
@@ -5,13 +5,13 @@
 
 use hir_def::{
     AdtId, AssocItemId, ConstId, DefWithBodyId, EnumId, EnumVariantId, FunctionId, GenericDefId,
-    ModuleDefId, StaticId, StructId, TypeAliasId, UnionId,
+    ModuleDefId, StaticId, StructId, TypeAliasId, UnionId, VariantId,
 };
 
 use crate::{
     ty::{CallableDef, TypableDef},
     Adt, AssocItem, Const, Crate, DefWithBody, EnumVariant, Function, GenericDef, ModuleDef,
-    Static, TypeAlias,
+    Static, TypeAlias, VariantDef,
 };
 
 impl From<ra_db::CrateId> for Crate {
@@ -70,6 +70,12 @@ impl From<EnumVariantId> for EnumVariant {
     }
 }
 
+impl From<EnumVariant> for EnumVariantId {
+    fn from(def: EnumVariant) -> Self {
+        EnumVariantId { parent: def.parent.id, local_id: def.id }
+    }
+}
+
 impl From<ModuleDefId> for ModuleDef {
     fn from(id: ModuleDefId) -> Self {
         match id {
@@ -219,3 +225,12 @@ impl From<CallableDef> for GenericDefId {
         }
     }
 }
+
+impl From<VariantDef> for VariantId {
+    fn from(def: VariantDef) -> Self {
+        match def {
+            VariantDef::Struct(it) => VariantId::StructId(it.id),
+            VariantDef::EnumVariant(it) => VariantId::EnumVariantId(it.into()),
+        }
+    }
+}
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
+    }
+}