Add Generation gen; to Query{Func,Type,Var} and corresponding ToQueryLocation

src/messages/cquery_derived.cc

@@ -37,7 +37,7 @@ struct CqueryDerivedHandler : BaseMessageHandler<Ipc_CqueryDerived> {
       if (ref.idx.kind == SymbolKind::Type) {
         QueryType& type = db->types[ref.idx.idx];
         std::vector<QueryLocation> locations =
-            ToQueryLocation(db, type.derived);
+            ToQueryLocation(db, &type.derived);
         out.result = GetLsLocations(db, working_files, locations);
         break;
       } else if (ref.idx.kind == SymbolKind::Func) {

src/messages/cquery_type_hierarchy_tree.cc

@@ -85,12 +85,15 @@ BuildInheritanceHierarchyForType(QueryDatabase* db,
     entry.children.push_back(base);
 
   // Add derived.
-  for (QueryTypeId derived : root_type.derived) {
+  for (WithGen<QueryTypeId> derived : root_type.derived) {
+    QueryType& type = db->types[derived.value.id];
+    if (derived.gen == type.gen) {
       auto derived_entry =
-        BuildInheritanceHierarchyForType(db, working_files, derived);
+        BuildInheritanceHierarchyForType(db, working_files, derived.value);
       if (derived_entry)
         entry.children.push_back(*derived_entry);
     }
+  }
 
   return entry;
 }

src/messages/cquery_vars.cc

@@ -39,7 +39,7 @@ struct CqueryVarsHandler : BaseMessageHandler<Ipc_CqueryVars> {
         case SymbolKind::Type: {
           QueryType& type = db->types[id];
           std::vector<QueryLocation> locations =
-              ToQueryLocation(db, type.instances);
+              ToQueryLocation(db, &type.instances);
           out.result = GetLsLocations(db, working_files, locations);
           break;
         }

src/messages/text_document_code_lens.cc

@@ -159,10 +159,10 @@ struct TextDocumentCodeLensHandler
                       true /*force_display*/);
           AddCodeLens("derived", "derived", &common,
                       ref.loc.OffsetStartColumn(1),
-                      ToQueryLocation(db, type.derived), nullopt,
+                      ToQueryLocation(db, &type.derived), nullopt,
                       false /*force_display*/);
           AddCodeLens("var", "vars", &common, ref.loc.OffsetStartColumn(2),
-                      ToQueryLocation(db, type.instances), nullopt,
+                      ToQueryLocation(db, &type.instances), nullopt,
                       false /*force_display*/);
           break;
         }

src/query.cc

@@ -355,6 +355,24 @@ Maybe<QueryVarId> GetQueryVarIdFromUsr(QueryDatabase* query_db,
   return QueryVarId(idx);
 }
 
+template <typename T>
+void AddRangeWithGen(std::vector<WithGen<T>>* dest, const std::vector<T>& to_add, Generation gen) {
+  for (auto& x : to_add)
+    dest->push_back(WithGen<T>{gen, x});
+}
+
+template <typename T>
+void RemoveRangeWithGen(std::vector<WithGen<T>>* dest, const std::vector<T>& to_remove) {
+  dest->erase(std::remove_if(dest->begin(), dest->end(),
+                             [&](const WithGen<T>& t) {
+                               // TODO: make to_remove a set?
+                               return std::find(to_remove.begin(),
+                                                to_remove.end(),
+                                                t.value) != to_remove.end();
+                             }),
+              dest->end());
+}
+
 }  // namespace
 
 template <>
@@ -797,6 +815,7 @@ void QueryDatabase::RemoveUsrs(SymbolKind usr_kind,
         QueryType& type = types[usr_to_type[usr].id];
         if (type.symbol_idx)
           symbols[type.symbol_idx->id].kind = SymbolKind::Invalid;
+        type.gen++;
         type.def = nullopt;
       }
       break;
@@ -806,6 +825,7 @@ void QueryDatabase::RemoveUsrs(SymbolKind usr_kind,
         QueryFunc& func = funcs[usr_to_func[usr].id];
         if (func.symbol_idx)
           symbols[func.symbol_idx->id].kind = SymbolKind::Invalid;
+        func.gen++;
         func.def = nullopt;
       }
       break;
@@ -815,6 +835,7 @@ void QueryDatabase::RemoveUsrs(SymbolKind usr_kind,
         QueryVar& var = vars[usr_to_var[usr].id];
         if (var.symbol_idx)
           symbols[var.symbol_idx->id].kind = SymbolKind::Invalid;
+        var.gen++;
         var.def = nullopt;
       }
       break;
@@ -840,6 +861,13 @@ void QueryDatabase::ApplyIndexUpdate(IndexUpdate* update) {
     RemoveRange(&def.def_var_name, merge_update.to_remove);           \
     VerifyUnique(def.def_var_name);                                   \
   }
+#define HANDLE_MERGEABLE_WITH_GEN(update_var_name, def_var_name, storage_name) \
+  for (auto merge_update : update->update_var_name) {                          \
+    auto& def = storage_name[merge_update.id.id];                              \
+    AddRangeWithGen(&def.def_var_name, merge_update.to_add, def.gen);          \
+    RemoveRangeWithGen(&def.def_var_name, merge_update.to_remove);             \
+    VerifyUnique(def.def_var_name);                                            \
+  }
 
   for (const std::string& filename : update->files_removed)
     files[usr_to_file[NormalizedPath(filename)].id].def = nullopt;
@@ -847,8 +875,8 @@ void QueryDatabase::ApplyIndexUpdate(IndexUpdate* update) {
 
   RemoveUsrs(SymbolKind::Type, update->types_removed);
   ImportOrUpdate(update->types_def_update);
-  HANDLE_MERGEABLE(types_derived, derived, types);
+  HANDLE_MERGEABLE_WITH_GEN(types_derived, derived, types);
-  HANDLE_MERGEABLE(types_instances, instances, types);
+  HANDLE_MERGEABLE_WITH_GEN(types_instances, instances, types);
   HANDLE_MERGEABLE(types_uses, uses, types);
 
   RemoveUsrs(SymbolKind::Func, update->funcs_removed);

src/query.h

@@ -18,6 +18,14 @@ using QueryTypeId = Id<QueryType>;
 using QueryFuncId = Id<QueryFunc>;
 using QueryVarId = Id<QueryVar>;
 
+using Generation = uint32_t;
+
+template <typename T>
+struct WithGen {
+  Generation gen;
+  T value;
+};
+
 struct IdMap;
 
 struct QueryLocation {
@@ -239,13 +247,14 @@ struct QueryType {
   using UsesUpdate = MergeableUpdate<QueryTypeId, QueryLocation>;
 
   Usr usr;
-  optional<Def> def;
+  Generation gen;
-  std::vector<QueryTypeId> derived;
-  std::vector<QueryVarId> instances;
-  std::vector<QueryLocation> uses;
   Maybe<Id<void>> symbol_idx;
+  optional<Def> def;
+  std::vector<WithGen<QueryTypeId>> derived;
+  std::vector<WithGen<QueryVarId>> instances;
+  std::vector<QueryLocation> uses;
 
-  explicit QueryType(const Usr& usr) : usr(usr) {}
+  explicit QueryType(const Usr& usr) : usr(usr), gen(0) {}
 };
 
 struct QueryFunc {
@@ -260,13 +269,14 @@ struct QueryFunc {
   using CallersUpdate = MergeableUpdate<QueryFuncId, QueryFuncRef>;
 
   Usr usr;
+  Generation gen;
+  Maybe<Id<void>> symbol_idx;
   optional<Def> def;
   std::vector<QueryLocation> declarations;
   std::vector<QueryFuncId> derived;
   std::vector<QueryFuncRef> callers;
-  Maybe<Id<void>> symbol_idx;
 
-  explicit QueryFunc(const Usr& usr) : usr(usr) {}
+  explicit QueryFunc(const Usr& usr) : usr(usr), gen(0) {}
 };
 
 struct QueryVar {
@@ -277,12 +287,13 @@ struct QueryVar {
   using UsesUpdate = MergeableUpdate<QueryVarId, QueryLocation>;
 
   Usr usr;
+  Generation gen;
+  Maybe<Id<void>> symbol_idx;
   optional<Def> def;
   std::vector<QueryLocation> declarations;
   std::vector<QueryLocation> uses;
-  Maybe<Id<void>> symbol_idx;
 
-  explicit QueryVar(const Usr& usr) : usr(usr) {}
+  explicit QueryVar(const Usr& usr) : usr(usr), gen(0) {}
 };
 
 struct IndexUpdate {

src/query_utils.cc

@@ -13,6 +13,29 @@ int ComputeRangeSize(const Range& range) {
   return range.end.column - range.start.column;
 }
 
+template <typename Q>
+std::vector<QueryLocation> ToQueryLocation(
+    QueryDatabase* db,
+    std::vector<Q> QueryDatabase::*collection,
+    std::vector<WithGen<Id<Q>>>* ids_) {
+  auto& ids = *ids_;
+  std::vector<QueryLocation> locs;
+  locs.reserve(ids.size());
+  size_t j = 0;
+  for (size_t i = 0; i < ids.size(); i++) {
+    Q& obj = (db->*collection)[ids[i].value.id];
+    if (obj.gen == ids[i].gen) {
+      optional<QueryLocation> loc =
+          GetDefinitionSpellingOfSymbol(db, ids[i].value);
+      if (loc)
+        locs.push_back(*loc);
+      ids[j++] = ids[i];
+    }
+  }
+  ids.resize(j);
+  return locs;
+}
+
 }  // namespace
 
 optional<QueryLocation> GetDefinitionSpellingOfSymbol(QueryDatabase* db,
@@ -169,16 +192,16 @@ std::vector<QueryLocation> ToQueryLocation(
   }
   return locs;
 }
-std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
+
-                                           const std::vector<QueryVarId>& ids) {
+std::vector<QueryLocation> ToQueryLocation(
-  std::vector<QueryLocation> locs;
+    QueryDatabase* db,
-  locs.reserve(ids.size());
+    std::vector<WithGen<QueryTypeId>>* ids_) {
-  for (const QueryVarId& id : ids) {
+  return ToQueryLocation(db, &QueryDatabase::types, ids_);
-    optional<QueryLocation> loc = GetDefinitionSpellingOfSymbol(db, id);
+}
-    if (loc)
+std::vector<QueryLocation> ToQueryLocation(
-      locs.push_back(loc.value());
+    QueryDatabase* db,
-  }
+    std::vector<WithGen<QueryVarId>>* ids_) {
-  return locs;
+  return ToQueryLocation(db, &QueryDatabase::vars, ids_);
 }
 
 std::vector<QueryLocation> GetUsesOfSymbol(QueryDatabase* db,

src/query_utils.h

@@ -22,12 +22,16 @@ optional<QueryFileId> GetDeclarationFileForSymbol(QueryDatabase* db,
 std::vector<QueryLocation> ToQueryLocation(
     QueryDatabase* db,
     const std::vector<QueryFuncRef>& refs);
+std::vector<QueryLocation> ToQueryLocation(
+    QueryDatabase* db,
+    const std::vector<QueryTypeId>& refs);
 std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
-                                           const std::vector<QueryTypeId>& ids);
+                                           std::vector<WithGen<QueryTypeId>>*);
 std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
                                            const std::vector<QueryFuncId>& ids);
-std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
+std::vector<QueryLocation> ToQueryLocation(
-                                           const std::vector<QueryVarId>& ids);
+    QueryDatabase* db,
+    std::vector<WithGen<QueryVarId>>*);
 std::vector<QueryLocation> GetUsesOfSymbol(QueryDatabase* db,
                                            const SymbolIdx& symbol,
                                            bool include_decl);