5 files changed, 50 insertions, 68 deletions

diff --git a/crates/ra_hir/src/adt.rs b/crates/ra_hir/src/adt.rs
index e0f515642..dc936e826 100644
--- a/crates/ra_hir/src/adt.rs
+++ b/crates/ra_hir/src/adt.rs
@@ -3,6 +3,7 @@
 
 use std::sync::Arc;
 
+use ra_arena::{RawId, Arena, impl_arena_id};
 use ra_syntax::{
     TreeArc,
     ast::{self, NameOwner, StructFlavor}
@@ -12,7 +13,6 @@ use crate::{
     Name, AsName, Struct, Enum, EnumVariant, Crate,
     HirDatabase, HirFileId,
     type_ref::TypeRef,
-    ids::LocationCtx,
 };
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@@ -66,43 +66,47 @@ fn variants(enum_def: &ast::EnumDef) -> impl Iterator<Item = &ast::EnumVariant>
 }
 
 impl EnumVariant {
-    pub fn source_impl(&self, db: &impl HirDatabase) -> (HirFileId, TreeArc<ast::EnumVariant>) {
+    pub(crate) fn source_impl(
+        &self,
+        db: &impl HirDatabase,
+    ) -> (HirFileId, TreeArc<ast::EnumVariant>) {
         let (file_id, enum_def) = self.parent.source(db);
         let var = variants(&*enum_def)
-            .nth(self.idx as usize)
+            .zip(db.enum_data(self.parent).variants.iter())
+            .find(|(_syntax, (id, _))| *id == self.id)
             .unwrap()
+            .0
             .to_owned();
         (file_id, var)
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct EnumVariantId(RawId);
+impl_arena_id!(EnumVariantId);
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct EnumData {
     pub(crate) name: Option<Name>,
-    pub(crate) variants: Vec<(Name, EnumVariant)>,
+    pub(crate) variants: Arena<EnumVariantId, EnumVariantData>,
 }
 
 impl EnumData {
-    fn new(enum_def: &ast::EnumDef, variants: Vec<(Name, EnumVariant)>) -> Self {
-        let name = enum_def.name().map(|n| n.as_name());
-        EnumData { name, variants }
-    }
-
     pub(crate) fn enum_data_query(db: &impl HirDatabase, e: Enum) -> Arc<EnumData> {
         let (_file_id, enum_def) = e.source(db);
-        let variants = variants(&*enum_def)
-            .enumerate()
-            .filter_map(|(idx, variant_def)| {
-                let name = variant_def.name()?.as_name();
-                let var = EnumVariant {
-                    parent: e,
-                    idx: idx as u32,
-                };
-                Some((name, var))
-            })
-            .collect();
-
-        Arc::new(EnumData::new(&*enum_def, variants))
+        let mut res = EnumData {
+            name: enum_def.name().map(|n| n.as_name()),
+            variants: Arena::default(),
+        };
+        for var in variants(&*enum_def) {
+            let data = EnumVariantData {
+                name: var.name().map(|it| it.as_name()),
+                variant_data: Arc::new(VariantData::new(var.flavor())),
+            };
+            res.variants.alloc(data);
+        }
+
+        Arc::new(res)
     }
 }
 
@@ -110,33 +114,6 @@ impl EnumData {
 pub struct EnumVariantData {
     pub(crate) name: Option<Name>,
     pub(crate) variant_data: Arc<VariantData>,
-    pub(crate) parent_enum: Enum,
-}
-
-impl EnumVariantData {
-    fn new(variant_def: &ast::EnumVariant, parent_enum: Enum) -> EnumVariantData {
-        let name = variant_def.name().map(|n| n.as_name());
-        let variant_data = VariantData::new(variant_def.flavor());
-        let variant_data = Arc::new(variant_data);
-        EnumVariantData {
-            name,
-            variant_data,
-            parent_enum,
-        }
-    }
-
-    pub(crate) fn enum_variant_data_query(
-        db: &impl HirDatabase,
-        var: EnumVariant,
-    ) -> Arc<EnumVariantData> {
-        let (file_id, variant_def) = var.source(db);
-        let enum_def = variant_def.parent_enum();
-        let ctx = LocationCtx::new(db, var.module(db), file_id);
-        let e = Enum {
-            id: ctx.to_def(enum_def),
-        };
-        Arc::new(EnumVariantData::new(&*variant_def, e))
-    }
 }
 
 /// A single field of an enum variant or struct
diff --git a/crates/ra_hir/src/code_model_api.rs b/crates/ra_hir/src/code_model_api.rs
index 4fc47aba5..249a4aba9 100644
--- a/crates/ra_hir/src/code_model_api.rs
+++ b/crates/ra_hir/src/code_model_api.rs
@@ -11,7 +11,7 @@ use crate::{
     db::HirDatabase,
     expr::BodySyntaxMapping,
     ty::{InferenceResult, VariantDef},
-    adt::VariantData,
+    adt::{VariantData, EnumVariantId},
     generics::GenericParams,
     docs::{Documentation, Docs, docs_from_ast},
     module_tree::ModuleId,
@@ -252,8 +252,20 @@ impl Enum {
         db.enum_data(*self).name.clone()
     }
 
-    pub fn variants(&self, db: &impl HirDatabase) -> Vec<(Name, EnumVariant)> {
-        db.enum_data(*self).variants.clone()
+    pub fn variants(&self, db: &impl HirDatabase) -> Vec<EnumVariant> {
+        db.enum_data(*self)
+            .variants
+            .iter()
+            .map(|(id, _)| EnumVariant { parent: *self, id })
+            .collect()
+    }
+
+    pub fn variant(&self, db: &impl HirDatabase, name: &Name) -> Option<EnumVariant> {
+        db.enum_data(*self)
+            .variants
+            .iter()
+            .find(|(_id, data)| data.name.as_ref() == Some(name))
+            .map(|(id, _)| EnumVariant { parent: *self, id })
     }
 
     pub fn generic_params(&self, db: &impl HirDatabase) -> Arc<GenericParams> {
@@ -270,7 +282,7 @@ impl Docs for Enum {
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct EnumVariant {
     pub(crate) parent: Enum,
-    pub(crate) idx: u32,
+    pub(crate) id: EnumVariantId,
 }
 
 impl EnumVariant {
@@ -285,11 +297,13 @@ impl EnumVariant {
     }
 
     pub fn name(&self, db: &impl HirDatabase) -> Option<Name> {
-        db.enum_variant_data(*self).name.clone()
+        db.enum_data(self.parent).variants[self.id].name.clone()
     }
 
     pub fn variant_data(&self, db: &impl HirDatabase) -> Arc<VariantData> {
-        db.enum_variant_data(*self).variant_data.clone()
+        db.enum_data(self.parent).variants[self.id]
+            .variant_data
+            .clone()
     }
 
     pub fn fields(&self, db: &impl HirDatabase) -> Vec<StructField> {
diff --git a/crates/ra_hir/src/db.rs b/crates/ra_hir/src/db.rs
index 5304f5d31..5a29e54d6 100644
--- a/crates/ra_hir/src/db.rs
+++ b/crates/ra_hir/src/db.rs
@@ -8,12 +8,12 @@ use crate::{
     SourceFileItems, SourceItemId, Crate, Module, HirInterner,
     query_definitions,
     Function, FnSignature, FnScopes,
-    Struct, Enum, EnumVariant,
+    Struct, Enum,
     macros::MacroExpansion,
     module_tree::ModuleTree,
     nameres::{ItemMap, lower::{LoweredModule, ImportSourceMap}},
     ty::{InferenceResult, Ty, method_resolution::CrateImplBlocks, TypableDef, VariantDef},
-    adt::{StructData, EnumData, EnumVariantData},
+    adt::{StructData, EnumData},
     impl_block::ModuleImplBlocks,
     generics::{GenericParams, GenericDef},
 };
@@ -35,9 +35,6 @@ pub trait HirDatabase: SyntaxDatabase + AsRef<HirInterner> {
     #[salsa::invoke(crate::adt::EnumData::enum_data_query)]
     fn enum_data(&self, e: Enum) -> Arc<EnumData>;
 
-    #[salsa::invoke(crate::adt::EnumVariantData::enum_variant_data_query)]
-    fn enum_variant_data(&self, var: EnumVariant) -> Arc<EnumVariantData>;
-
     #[salsa::invoke(crate::ty::infer)]
     fn infer(&self, func: Function) -> Arc<InferenceResult>;
 
diff --git a/crates/ra_hir/src/mock.rs b/crates/ra_hir/src/mock.rs
index 361366f6a..1a83a5c87 100644
--- a/crates/ra_hir/src/mock.rs
+++ b/crates/ra_hir/src/mock.rs
@@ -231,7 +231,6 @@ salsa::database_storage! {
             fn type_for_field() for db::TypeForFieldQuery;
             fn struct_data() for db::StructDataQuery;
             fn enum_data() for db::EnumDataQuery;
-            fn enum_variant_data() for db::EnumVariantDataQuery;
             fn impls_in_module() for db::ImplsInModuleQuery;
             fn impls_in_crate() for db::ImplsInCrateQuery;
             fn body_hir() for db::BodyHirQuery;
diff --git a/crates/ra_hir/src/nameres.rs b/crates/ra_hir/src/nameres.rs
index 639726b5e..028c1882f 100644
--- a/crates/ra_hir/src/nameres.rs
+++ b/crates/ra_hir/src/nameres.rs
@@ -397,13 +397,8 @@ impl ItemMap {
                 ModuleDef::Enum(e) => {
                     // enum variant
                     tested_by!(item_map_enum_importing);
-                    let matching_variant = e
-                        .variants(db)
-                        .into_iter()
-                        .find(|(n, _variant)| n == &segment.name);
-
-                    match matching_variant {
-                        Some((_n, variant)) => PerNs::both(variant.into(), (*e).into()),
+                    match e.variant(db, &segment.name) {
+                        Some(variant) => PerNs::both(variant.into(), (*e).into()),
                         None => PerNs::none(),
                     }
                 }