From e5128d3db9b75885933024c80054b306a1c3d3b0 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sat, 3 Feb 2018 16:20:14 -0800
Subject: [PATCH] Wrap Query* references with WithGen

---
 src/indexer.cc                             |  2 +-
 src/messages/cquery_base.cc                |  4 +-
 src/messages/cquery_derived.cc             |  2 +-
 src/messages/cquery_member_hierarchy.cc    | 15 ++---
 src/messages/cquery_type_hierarchy_tree.cc | 72 ++++++++++------------
 src/messages/cquery_vars.cc                |  2 +-
 src/messages/text_document_code_action.cc  | 15 +++--
 src/messages/text_document_code_lens.cc    |  7 ++-
 src/messages/text_document_definition.cc   |  3 +-
 src/query.cc                               | 67 +++++++++++++++++---
 src/query.h                                | 66 +++++++++++++++++---
 src/query_utils.cc                         | 54 +++++++++++-----
 src/query_utils.h                          | 29 ++++++++-
 src/serializer.h                           | 21 -------
 14 files changed, 234 insertions(+), 125 deletions(-)

diff --git a/src/indexer.cc b/src/indexer.cc
index 56d78211..774ef61b 100644
--- a/src/indexer.cc
+++ b/src/indexer.cc
@@ -1777,7 +1777,7 @@ void OnIndexDeclaration(CXClientData client_data, const CXIdxDeclInfo* decl) {
           if (parent_type_id) {
             IndexType* parent_type_def = db->Resolve(parent_type_id.value());
             parent_type_def->derived.push_back(type_id);
-            type->def.parents.push_back(parent_type_id.value());
+            type->def.parents.push_back(*parent_type_id);
           }
         }
       }
diff --git a/src/messages/cquery_base.cc b/src/messages/cquery_base.cc
index d380230f..3563b453 100644
--- a/src/messages/cquery_base.cc
+++ b/src/messages/cquery_base.cc
@@ -39,13 +39,13 @@ struct CqueryBaseHandler : BaseMessageHandler<Ipc_CqueryBase> {
         if (!type.def)
           continue;
         std::vector<QueryLocation> locations =
-            ToQueryLocation(db, type.def->parents);
+            ToQueryLocation(db, &type.def->parents);
         out.result = GetLsLocations(db, working_files, locations);
         break;
       } else if (ref.idx.kind == SymbolKind::Func) {
         QueryFunc& func = db->funcs[ref.idx.idx];
         std::vector<QueryLocation> locations =
-            ToQueryLocation(db, func.def->base);
+            ToQueryLocation(db, &func.def->base);
         out.result = GetLsLocations(db, working_files, locations);
         break;
       }
diff --git a/src/messages/cquery_derived.cc b/src/messages/cquery_derived.cc
index 123476e4..d119ee17 100644
--- a/src/messages/cquery_derived.cc
+++ b/src/messages/cquery_derived.cc
@@ -43,7 +43,7 @@ struct CqueryDerivedHandler : BaseMessageHandler<Ipc_CqueryDerived> {
       } else if (ref.idx.kind == SymbolKind::Func) {
         QueryFunc& func = db->funcs[ref.idx.idx];
         std::vector<QueryLocation> locations =
-            ToQueryLocation(db, func.derived);
+            ToQueryLocation(db, &func.derived);
         out.result = GetLsLocations(db, working_files, locations);
         break;
       }
diff --git a/src/messages/cquery_member_hierarchy.cc b/src/messages/cquery_member_hierarchy.cc
index ef72c070..b4d000b6 100644
--- a/src/messages/cquery_member_hierarchy.cc
+++ b/src/messages/cquery_member_hierarchy.cc
@@ -27,7 +27,8 @@ struct Out_CqueryMemberHierarchy
     : public lsOutMessage<Out_CqueryMemberHierarchy> {
   struct Entry {
     std::string_view name;
-    size_t type_id;
+    // FIXME Usr
+    RawId type_id;
     lsLocation location;
   };
   lsRequestId id;
@@ -47,8 +48,8 @@ BuildInitial(QueryDatabase* db, WorkingFiles* working_files, QueryTypeId root) {
     return {};
 
   Out_CqueryMemberHierarchy::Entry entry;
-  entry.name = root_type.def->ShortName();
   entry.type_id = root.id;
+  entry.name = root_type.def->ShortName();
   entry.location = *def_loc;
   return {entry};
 }
@@ -60,12 +61,12 @@ ExpandNode(QueryDatabase* db, WorkingFiles* working_files, QueryTypeId root) {
     return {};
 
   std::vector<Out_CqueryMemberHierarchy::Entry> ret;
-  for (auto& var_id : root_type.def->vars) {
-    QueryVar& var = db->vars[var_id.id];
+  EachWithGen<QueryVar>(db->vars, root_type.def->vars, [&](QueryVar& var) {
     Out_CqueryMemberHierarchy::Entry entry;
     entry.name = var.def->ShortName();
+    // FIXME WithGen
     entry.type_id =
-        var.def->variable_type ? var.def->variable_type->id : size_t(-1);
+        var.def->variable_type ? var.def->variable_type->value.id : RawId(-1);
     if (var.def->definition_spelling) {
       optional<lsLocation> loc =
           GetLsLocation(db, working_files, *var.def->definition_spelling);
@@ -74,7 +75,7 @@ ExpandNode(QueryDatabase* db, WorkingFiles* working_files, QueryTypeId root) {
         entry.location = *loc;
     }
     ret.push_back(std::move(entry));
-  }
+  });
   return ret;
 }
 
@@ -100,7 +101,7 @@ struct CqueryMemberHierarchyInitialHandler
       if (ref.idx.kind == SymbolKind::Var) {
         QueryVar& var = db->vars[ref.idx.idx];
         if (var.def && var.def->variable_type)
-          out.result = BuildInitial(db, working_files, *var.def->variable_type);
+          out.result = BuildInitial(db, working_files, var.def->variable_type->value);
         break;
       }
     }
diff --git a/src/messages/cquery_type_hierarchy_tree.cc b/src/messages/cquery_type_hierarchy_tree.cc
index 0510d386..1048e349 100644
--- a/src/messages/cquery_type_hierarchy_tree.cc
+++ b/src/messages/cquery_type_hierarchy_tree.cc
@@ -30,29 +30,22 @@ MAKE_REFLECT_STRUCT(Out_CqueryTypeHierarchyTree, jsonrpc, id, result);
 std::vector<Out_CqueryTypeHierarchyTree::TypeEntry>
 BuildParentInheritanceHierarchyForType(QueryDatabase* db,
                                        WorkingFiles* working_files,
-                                       QueryTypeId root) {
-  QueryType& root_type = db->types[root.id];
-  if (!root_type.def)
-    return {};
-
+                                       QueryType& root_type) {
   std::vector<Out_CqueryTypeHierarchyTree::TypeEntry> parent_entries;
   parent_entries.reserve(root_type.def->parents.size());
 
-  for (QueryTypeId parent_id : root_type.def->parents) {
-    QueryType& parent_type = db->types[parent_id.id];
-    if (!parent_type.def)
-      continue;
+  EachWithGen<QueryType>(
+      db->types, root_type.def->parents, [&](QueryType& parent_type) {
+        Out_CqueryTypeHierarchyTree::TypeEntry parent_entry;
+        parent_entry.name = parent_type.def->detailed_name;
+        if (parent_type.def->definition_spelling)
+          parent_entry.location = GetLsLocation(
+              db, working_files, *parent_type.def->definition_spelling);
+        parent_entry.children = BuildParentInheritanceHierarchyForType(
+            db, working_files, parent_type);
 
-    Out_CqueryTypeHierarchyTree::TypeEntry parent_entry;
-    parent_entry.name = parent_type.def->detailed_name;
-    if (parent_type.def->definition_spelling)
-      parent_entry.location = GetLsLocation(
-          db, working_files, *parent_type.def->definition_spelling);
-    parent_entry.children =
-        BuildParentInheritanceHierarchyForType(db, working_files, parent_id);
-
-    parent_entries.push_back(parent_entry);
-  }
+        parent_entries.push_back(parent_entry);
+      });
 
   return parent_entries;
 }
@@ -60,11 +53,7 @@ BuildParentInheritanceHierarchyForType(QueryDatabase* db,
 optional<Out_CqueryTypeHierarchyTree::TypeEntry>
 BuildInheritanceHierarchyForType(QueryDatabase* db,
                                  WorkingFiles* working_files,
-                                 QueryTypeId root_id) {
-  QueryType& root_type = db->types[root_id.id];
-  if (!root_type.def)
-    return nullopt;
-
+                                 QueryType& root_type) {
   Out_CqueryTypeHierarchyTree::TypeEntry entry;
 
   // Name and location.
@@ -80,20 +69,17 @@ BuildInheritanceHierarchyForType(QueryDatabase* db,
   base.name = "[[Base]]";
   base.location = entry.location;
   base.children =
-      BuildParentInheritanceHierarchyForType(db, working_files, root_id);
+    BuildParentInheritanceHierarchyForType(db, working_files, root_type);
   if (!base.children.empty())
     entry.children.push_back(base);
 
   // Add 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.value);
-      if (derived_entry)
-        entry.children.push_back(*derived_entry);
-    }
-  }
+  EachWithGen<QueryType>(db->types, root_type.derived, [&](QueryType& type) {
+    auto derived_entry =
+        BuildInheritanceHierarchyForType(db, working_files, type);
+    if (derived_entry)
+      entry.children.push_back(*derived_entry);
+  });
 
   return entry;
 }
@@ -108,8 +94,9 @@ BuildParentInheritanceHierarchyForFunc(QueryDatabase* db,
   if (!root_func.def || root_func.def->base.empty())
     return {};
 
-  for (QueryFuncId parent_id : root_func.def->base) {
-    QueryFunc& parent_func = db->funcs[parent_id.id];
+  // FIXME WithGen
+  for (auto parent_id : root_func.def->base) {
+    QueryFunc& parent_func = db->funcs[parent_id.value.id];
     if (!parent_func.def)
       continue;
 
@@ -119,7 +106,7 @@ BuildParentInheritanceHierarchyForFunc(QueryDatabase* db,
       parent_entry.location = GetLsLocation(
           db, working_files, *parent_func.def->definition_spelling);
     parent_entry.children =
-        BuildParentInheritanceHierarchyForFunc(db, working_files, parent_id);
+        BuildParentInheritanceHierarchyForFunc(db, working_files, parent_id.value);
 
     entries.push_back(parent_entry);
   }
@@ -155,9 +142,10 @@ BuildInheritanceHierarchyForFunc(QueryDatabase* db,
     entry.children.push_back(base);
 
   // Add derived.
-  for (QueryFuncId derived : root_func.derived) {
+  // FIXME WithGen
+  for (auto derived : root_func.derived) {
     auto derived_entry =
-        BuildInheritanceHierarchyForFunc(db, working_files, derived);
+        BuildInheritanceHierarchyForFunc(db, working_files, derived.value);
     if (derived_entry)
       entry.children.push_back(*derived_entry);
   }
@@ -182,8 +170,10 @@ struct CqueryTypeHierarchyTreeHandler
     for (const SymbolRef& ref :
          FindSymbolsAtLocation(working_file, file, request->params.position)) {
       if (ref.idx.kind == SymbolKind::Type) {
-        out.result = BuildInheritanceHierarchyForType(db, working_files,
-                                                      QueryTypeId(ref.idx.idx));
+        QueryType& type = db->types[ref.idx.idx];
+        if (type.def)
+          out.result =
+              BuildInheritanceHierarchyForType(db, working_files, type);
         break;
       }
       if (ref.idx.kind == SymbolKind::Func) {
diff --git a/src/messages/cquery_vars.cc b/src/messages/cquery_vars.cc
index cba61321..facb9e4b 100644
--- a/src/messages/cquery_vars.cc
+++ b/src/messages/cquery_vars.cc
@@ -33,7 +33,7 @@ struct CqueryVarsHandler : BaseMessageHandler<Ipc_CqueryVars> {
           QueryVar& var = db->vars[id];
           if (!var.def || !var.def->variable_type)
             continue;
-          id = var.def->variable_type->id;
+          id = var.def->variable_type->value.id;
         }
         // fallthrough
         case SymbolKind::Type: {
diff --git a/src/messages/text_document_code_action.cc b/src/messages/text_document_code_action.cc
index 6de17dec..be8bd0d4 100644
--- a/src/messages/text_document_code_action.cc
+++ b/src/messages/text_document_code_action.cc
@@ -159,7 +159,7 @@ optional<lsTextEdit> BuildAutoImplementForFunction(QueryDatabase* db,
     optional<std::string> type_name;
     optional<lsPosition> same_file_insert_end;
     if (func.def->declaring_type) {
-      QueryType& declaring_type = db->types[func.def->declaring_type->id];
+      QueryType& declaring_type = db->types[func.def->declaring_type->value.id];
       if (declaring_type.def) {
         type_name = std::string(declaring_type.def->ShortName());
         optional<lsRange> ls_type_def_extent = GetLsRange(
@@ -353,17 +353,16 @@ struct TextDocumentCodeActionHandler
           // Get implementation file.
           Out_TextDocumentCodeAction::Command command;
 
-          for (QueryFuncId func_id : type.def->funcs) {
-            QueryFunc& func_def = db->funcs[func_id.id];
-            if (!func_def.def || func_def.def->definition_extent)
-              continue;
-
+          EachWithGen<QueryFunc>(db->funcs, type.def->funcs, [&](QueryFunc&
+                                                                     func_def) {
+            if (func_def.def->definition_extent)
+              return;
             EnsureImplFile(db, file_id, impl_uri /*out*/, impl_file_id /*out*/);
             optional<lsTextEdit> edit = BuildAutoImplementForFunction(
                 db, working_files, working_file, default_line, file_id,
                 *impl_file_id, func_def);
             if (!edit)
-              continue;
+              return;
 
             ++num_edits;
 
@@ -379,7 +378,7 @@ struct TextDocumentCodeActionHandler
             } else {
               command.arguments.edits.push_back(*edit);
             }
-          }
+          });
           if (command.arguments.edits.empty())
             break;
 
diff --git a/src/messages/text_document_code_lens.cc b/src/messages/text_document_code_lens.cc
index 5b78c8e7..e90fe4f5 100644
--- a/src/messages/text_document_code_lens.cc
+++ b/src/messages/text_document_code_lens.cc
@@ -216,13 +216,14 @@ struct TextDocumentCodeLensHandler
 
           AddCodeLens("derived", "derived", &common,
                       ref.loc.OffsetStartColumn(offset++),
-                      ToQueryLocation(db, func.derived), nullopt,
+                      ToQueryLocation(db, &func.derived), nullopt,
                       false /*force_display*/);
 
           // "Base"
           if (func.def->base.size() == 1) {
+            // FIXME WithGen
             optional<QueryLocation> base_loc =
-                GetDefinitionSpellingOfSymbol(db, func.def->base[0]);
+                GetDefinitionSpellingOfSymbol(db, func.def->base[0].value);
             if (base_loc) {
               optional<lsLocation> ls_base =
                   GetLsLocation(db, working_files, *base_loc);
@@ -244,7 +245,7 @@ struct TextDocumentCodeLensHandler
             }
           } else {
             AddCodeLens("base", "base", &common, ref.loc.OffsetStartColumn(1),
-                        ToQueryLocation(db, func.def->base), nullopt,
+                        ToQueryLocation(db, &func.def->base), nullopt,
                         false /*force_display*/);
           }
 
diff --git a/src/messages/text_document_definition.cc b/src/messages/text_document_definition.cc
index 488ee8cc..973912c7 100644
--- a/src/messages/text_document_definition.cc
+++ b/src/messages/text_document_definition.cc
@@ -32,10 +32,11 @@ std::vector<QueryLocation> GetGotoDefinitionTargets(QueryDatabase* db,
       std::vector<QueryLocation> ret =
           GetDeclarationsOfSymbolForGotoDefinition(db, symbol);
       QueryVar& var = db->vars[symbol.idx];
+      // FIXME WithGen
       if (var.def && var.def->variable_type) {
         std::vector<QueryLocation> types =
             GetDeclarationsOfSymbolForGotoDefinition(
-                db, SymbolIdx(SymbolKind::Type, var.def->variable_type->id));
+                db, SymbolIdx(SymbolKind::Type, var.def->variable_type->value.id));
         ret.insert(ret.end(), types.begin(), types.end());
       }
       return ret;
diff --git a/src/query.cc b/src/query.cc
index 78ea66a8..c970307f 100644
--- a/src/query.cc
+++ b/src/query.cc
@@ -44,11 +44,11 @@ optional<QueryType::Def> ToQuery(const IdMap& id_map,
   result.comments = type.comments;
   result.definition_spelling = id_map.ToQuery(type.definition_spelling);
   result.definition_extent = id_map.ToQuery(type.definition_extent);
-  result.alias_of = id_map.ToQuery(type.alias_of);
-  result.parents = id_map.ToQuery(type.parents);
-  result.types = id_map.ToQuery(type.types);
-  result.funcs = id_map.ToQuery(type.funcs);
-  result.vars = id_map.ToQuery(type.vars);
+  result.alias_of = id_map.ToQuery(type.alias_of,0);
+  result.parents = id_map.ToQuery(type.parents,0);
+  result.types = id_map.ToQuery(type.types,0);
+  result.funcs = id_map.ToQuery(type.funcs,0);
+  result.vars = id_map.ToQuery(type.vars,0);
   return result;
 }
 
@@ -67,9 +67,9 @@ optional<QueryFunc::Def> ToQuery(const IdMap& id_map,
   result.comments = func.comments;
   result.definition_spelling = id_map.ToQuery(func.definition_spelling);
   result.definition_extent = id_map.ToQuery(func.definition_extent);
-  result.declaring_type = id_map.ToQuery(func.declaring_type);
-  result.base = id_map.ToQuery(func.base);
-  result.locals = id_map.ToQuery(func.locals);
+  result.declaring_type = id_map.ToQuery(func.declaring_type,0);
+  result.base = id_map.ToQuery(func.base,0);
+  result.locals = id_map.ToQuery(func.locals,0);
   result.callees = id_map.ToQuery(func.callees);
   return result;
 }
@@ -86,7 +86,7 @@ optional<QueryVar::Def> ToQuery(const IdMap& id_map, const IndexVar::Def& var) {
   result.comments = var.comments;
   result.definition_spelling = id_map.ToQuery(var.definition_spelling);
   result.definition_extent = id_map.ToQuery(var.definition_extent);
-  result.variable_type = id_map.ToQuery(var.variable_type);
+  result.variable_type = id_map.ToQuery(var.variable_type,0);
   result.parent_id = var.parent_id;
   result.parent_kind = var.parent_kind;
   result.kind = var.kind;
@@ -436,6 +436,20 @@ QueryVarId IdMap::ToQuery(IndexVarId id) const {
   assert(cached_var_ids_.find(id) != cached_var_ids_.end());
   return QueryVarId(cached_var_ids_.find(id)->second);
 }
+WithGen<QueryTypeId> IdMap::ToQuery(IndexTypeId id,int) const {
+  assert(cached_type_ids_.find(id) != cached_type_ids_.end());
+  return QueryTypeId(cached_type_ids_.find(id)->second);
+}
+WithGen<QueryFuncId> IdMap::ToQuery(IndexFuncId id,int) const {
+  if (id == IndexFuncId())
+    return QueryFuncId();
+  assert(cached_func_ids_.find(id) != cached_func_ids_.end());
+  return QueryFuncId(cached_func_ids_.find(id)->second);
+}
+WithGen<QueryVarId> IdMap::ToQuery(IndexVarId id,int) const {
+  assert(cached_var_ids_.find(id) != cached_var_ids_.end());
+  return QueryVarId(cached_var_ids_.find(id)->second);
+}
 QueryFuncRef IdMap::ToQuery(IndexFuncRef ref) const {
   return QueryFuncRef(ToQuery(ref.id), ToQuery(ref.loc), ref.is_implicit);
 }
@@ -464,6 +478,21 @@ optional<QueryVarId> IdMap::ToQuery(optional<IndexVarId> id) const {
     return nullopt;
   return ToQuery(id.value());
 }
+optional<WithGen<QueryTypeId>> IdMap::ToQuery(optional<IndexTypeId> id, int) const {
+  if (!id)
+    return nullopt;
+  return ToQuery(id.value(), 0);
+}
+optional<WithGen<QueryFuncId>> IdMap::ToQuery(optional<IndexFuncId> id, int) const {
+  if (!id)
+    return nullopt;
+  return ToQuery(id.value(), 0);
+}
+optional<WithGen<QueryVarId>> IdMap::ToQuery(optional<IndexVarId> id, int) const {
+  if (!id)
+    return nullopt;
+  return ToQuery(id.value(), 0);
+}
 optional<QueryFuncRef> IdMap::ToQuery(optional<IndexFuncRef> ref) const {
   if (!ref)
     return nullopt;
@@ -485,6 +514,15 @@ std::vector<Out> ToQueryTransform(const IdMap& id_map,
     result.push_back(id_map.ToQuery(in));
   return result;
 }
+template <typename In, typename Out>
+std::vector<WithGen<Out>> ToQueryTransformG(const IdMap& id_map,
+                                  const std::vector<In>& input) {
+  std::vector<WithGen<Out>> result;
+  result.reserve(input.size());
+  for (const In& in : input)
+    result.push_back(id_map.ToQuery(in, 0));
+  return result;
+}
 std::vector<QueryLocation> IdMap::ToQuery(std::vector<Range> ranges) const {
   return ToQueryTransform<Range, QueryLocation>(*this, ranges);
 }
@@ -497,6 +535,15 @@ std::vector<QueryFuncId> IdMap::ToQuery(std::vector<IndexFuncId> ids) const {
 std::vector<QueryVarId> IdMap::ToQuery(std::vector<IndexVarId> ids) const {
   return ToQueryTransform<IndexVarId, QueryVarId>(*this, ids);
 }
+std::vector<WithGen<QueryTypeId>> IdMap::ToQuery(std::vector<IndexTypeId> ids, int) const {
+  return ToQueryTransformG<IndexTypeId, QueryTypeId>(*this, ids);
+}
+std::vector<WithGen<QueryFuncId>> IdMap::ToQuery(std::vector<IndexFuncId> ids, int) const {
+  return ToQueryTransformG<IndexFuncId, QueryFuncId>(*this, ids);
+}
+std::vector<WithGen<QueryVarId>> IdMap::ToQuery(std::vector<IndexVarId> ids, int) const {
+  return ToQueryTransformG<IndexVarId, QueryVarId>(*this, ids);
+}
 std::vector<QueryFuncRef> IdMap::ToQuery(std::vector<IndexFuncRef> refs) const {
   return ToQueryTransform<IndexFuncRef, QueryFuncRef>(*this, refs);
 }
@@ -882,7 +929,7 @@ void QueryDatabase::ApplyIndexUpdate(IndexUpdate* update) {
   RemoveUsrs(SymbolKind::Func, update->funcs_removed);
   ImportOrUpdate(update->funcs_def_update);
   HANDLE_MERGEABLE(funcs_declarations, declarations, funcs);
-  HANDLE_MERGEABLE(funcs_derived, derived, funcs);
+  HANDLE_MERGEABLE_WITH_GEN(funcs_derived, derived, funcs);
   HANDLE_MERGEABLE(funcs_callers, callers, funcs);
 
   RemoveUsrs(SymbolKind::Var, update->vars_removed);
diff --git a/src/query.h b/src/query.h
index 8d35fbde..7aafef07 100644
--- a/src/query.h
+++ b/src/query.h
@@ -18,15 +18,34 @@ using QueryTypeId = Id<QueryType>;
 using QueryFuncId = Id<QueryFunc>;
 using QueryVarId = Id<QueryVar>;
 
+struct IdMap;
+
 using Generation = uint32_t;
 
+// Example use: |WithGen<Id<QueryType>>|, to mark an |Id| reference with
+// generation, so that by comparising the generation with that stored in the
+// referenced Query object, we can tell if the reference is stale (the
+// referenced object has been deleted or reused).
 template <typename T>
 struct WithGen {
   Generation gen;
   T value;
+  WithGen() : gen(-1) {}
+  WithGen(const T& value) : gen(-1), value(value) {}
+  WithGen(Generation gen, const T& value) : gen(gen), value(value) {}
+
+  bool HasValue() const { return value.HasValue(); }
+  explicit operator bool() const { return HasValue(); }
+
+  bool operator==(const WithGen& o) const {
+    return gen == o.gen && value == o.value;
+  }
 };
 
-struct IdMap;
+template <typename TVisitor, typename T>
+void Reflect(TVisitor& visitor, WithGen<T>& value) {
+  Reflect(visitor, value.value);
+}
 
 struct QueryLocation {
   QueryFileId path;
@@ -157,10 +176,23 @@ struct MergeableUpdate {
 
   MergeableUpdate(TId id, const std::vector<TValue>& to_add)
       : id(id), to_add(to_add) {}
+  MergeableUpdate(TId id, const std::vector<WithGen<TValue>>& to_add) : id(id) {
+    for (auto& x : to_add)
+      this->to_add.push_back(x.value);
+  }
   MergeableUpdate(TId id,
                   const std::vector<TValue>& to_add,
                   const std::vector<TValue>& to_remove)
       : id(id), to_add(to_add), to_remove(to_remove) {}
+  MergeableUpdate(TId id,
+                  const std::vector<WithGen<TValue>>& to_add,
+                  const std::vector<WithGen<TValue>>& to_remove)
+    : id(id) {
+    for (auto& x : to_add)
+      this->to_add.push_back(x.value);
+    for (auto& x : to_remove)
+      this->to_remove.push_back(x.value);
+  }
 };
 template <typename TVisitor, typename TId, typename TValue>
 void Reflect(TVisitor& visitor, MergeableUpdate<TId, TValue>& value) {
@@ -237,9 +269,9 @@ MAKE_REFLECT_STRUCT(QueryFile::Def,
                     dependencies);
 
 struct QueryType {
-  using Def = TypeDefDefinitionData<QueryTypeId,
-                                    QueryFuncId,
-                                    QueryVarId,
+  using Def = TypeDefDefinitionData<WithGen<QueryTypeId>,
+                                    WithGen<QueryFuncId>,
+                                    WithGen<QueryVarId>,
                                     QueryLocation>;
   using DefUpdate = WithUsr<Def>;
   using DerivedUpdate = MergeableUpdate<QueryTypeId, QueryTypeId>;
@@ -258,9 +290,9 @@ struct QueryType {
 };
 
 struct QueryFunc {
-  using Def = FuncDefDefinitionData<QueryTypeId,
-                                    QueryFuncId,
-                                    QueryVarId,
+  using Def = FuncDefDefinitionData<WithGen<QueryTypeId>,
+                                    WithGen<QueryFuncId>,
+                                    WithGen<QueryVarId>,
                                     QueryFuncRef,
                                     QueryLocation>;
   using DefUpdate = WithUsr<Def>;
@@ -273,15 +305,17 @@ struct QueryFunc {
   Maybe<Id<void>> symbol_idx;
   optional<Def> def;
   std::vector<QueryLocation> declarations;
-  std::vector<QueryFuncId> derived;
+  std::vector<WithGen<QueryFuncId>> derived;
   std::vector<QueryFuncRef> callers;
 
   explicit QueryFunc(const Usr& usr) : usr(usr), gen(0) {}
 };
 
 struct QueryVar {
-  using Def =
-      VarDefDefinitionData<QueryTypeId, QueryFuncId, QueryVarId, QueryLocation>;
+  using Def = VarDefDefinitionData<WithGen<QueryTypeId>,
+                                   WithGen<QueryFuncId>,
+                                   WithGen<QueryVarId>,
+                                   QueryLocation>;
   using DefUpdate = WithUsr<Def>;
   using DeclarationsUpdate = MergeableUpdate<QueryVarId, QueryLocation>;
   using UsesUpdate = MergeableUpdate<QueryVarId, QueryLocation>;
@@ -414,25 +448,37 @@ struct IdMap {
 
   IdMap(QueryDatabase* query_db, const IdCache& local_ids);
 
+  // FIXME Too verbose
+  // clang-format off
   QueryLocation ToQuery(Range range) const;
   QueryTypeId ToQuery(IndexTypeId id) const;
   QueryFuncId ToQuery(IndexFuncId id) const;
   QueryVarId ToQuery(IndexVarId id) const;
+  WithGen<QueryTypeId> ToQuery(IndexTypeId id, int) const;
+  WithGen<QueryFuncId> ToQuery(IndexFuncId id, int) const;
+  WithGen<QueryVarId> ToQuery(IndexVarId id, int) const;
   QueryFuncRef ToQuery(IndexFuncRef ref) const;
   QueryLocation ToQuery(IndexFunc::Declaration decl) const;
   optional<QueryLocation> ToQuery(optional<Range> range) const;
   optional<QueryTypeId> ToQuery(optional<IndexTypeId> id) const;
   optional<QueryFuncId> ToQuery(optional<IndexFuncId> id) const;
   optional<QueryVarId> ToQuery(optional<IndexVarId> id) const;
+  optional<WithGen<QueryTypeId>> ToQuery(optional<IndexTypeId> id,int) const;
+  optional<WithGen<QueryFuncId>> ToQuery(optional<IndexFuncId> id,int) const;
+  optional<WithGen<QueryVarId>> ToQuery(optional<IndexVarId> id,int) const;
   optional<QueryFuncRef> ToQuery(optional<IndexFuncRef> ref) const;
   optional<QueryLocation> ToQuery(optional<IndexFunc::Declaration> decl) const;
   std::vector<QueryLocation> ToQuery(std::vector<Range> ranges) const;
   std::vector<QueryTypeId> ToQuery(std::vector<IndexTypeId> ids) const;
   std::vector<QueryFuncId> ToQuery(std::vector<IndexFuncId> ids) const;
   std::vector<QueryVarId> ToQuery(std::vector<IndexVarId> ids) const;
+  std::vector<WithGen<QueryTypeId>> ToQuery(std::vector<IndexTypeId> ids,int) const;
+  std::vector<WithGen<QueryFuncId>> ToQuery(std::vector<IndexFuncId> ids,int) const;
+  std::vector<WithGen<QueryVarId>> ToQuery(std::vector<IndexVarId> ids,int) const;
   std::vector<QueryFuncRef> ToQuery(std::vector<IndexFuncRef> refs) const;
   std::vector<QueryLocation> ToQuery(
       std::vector<IndexFunc::Declaration> decls) const;
+  // clang-format on
 
   SymbolIdx ToSymbol(IndexTypeId id) const;
   SymbolIdx ToSymbol(IndexFuncId id) const;
diff --git a/src/query_utils.cc b/src/query_utils.cc
index ae54abc3..3b59f735 100644
--- a/src/query_utils.cc
+++ b/src/query_utils.cc
@@ -193,6 +193,11 @@ std::vector<QueryLocation> ToQueryLocation(
   return locs;
 }
 
+std::vector<QueryLocation> ToQueryLocation(
+    QueryDatabase* db,
+    std::vector<WithGen<QueryFuncId>>* ids_) {
+  return ToQueryLocation(db, &QueryDatabase::funcs, ids_);
+}
 std::vector<QueryLocation> ToQueryLocation(
     QueryDatabase* db,
     std::vector<WithGen<QueryTypeId>>* ids_) {
@@ -282,25 +287,33 @@ bool HasCallersOnSelfOrBaseOrDerived(QueryDatabase* db, QueryFunc& root) {
     return true;
 
   // Check for base calls.
-  std::queue<QueryFuncId> queue;
-  PushRange(&queue, root.def->base);
+  std::queue<QueryFunc*> queue;
+  EachWithGen<QueryFunc>(db->funcs, root.def->base, [&](QueryFunc& func) {
+    queue.push(&func);
+  });
   while (!queue.empty()) {
-    QueryFunc& func = db->funcs[queue.front().id];
+    QueryFunc& func = *queue.front();
     queue.pop();
     if (!func.callers.empty())
       return true;
     if (func.def)
-      PushRange(&queue, func.def->base);
+      EachWithGen<QueryFunc>(db->funcs, func.def->base, [&](QueryFunc& func1) {
+        queue.push(&func1);
+      });
   }
 
   // Check for derived calls.
-  PushRange(&queue, root.derived);
+  EachWithGen<QueryFunc>(db->funcs, root.derived, [&](QueryFunc& func1) {
+    queue.push(&func1);
+  });
   while (!queue.empty()) {
-    QueryFunc& func = db->funcs[queue.front().id];
+    QueryFunc& func = *queue.front();
     queue.pop();
     if (!func.callers.empty())
       return true;
-    PushRange(&queue, func.derived);
+    EachWithGen<QueryFunc>(db->funcs, func.derived, [&](QueryFunc& func1) {
+      queue.push(&func1);
+    });
   }
 
   return false;
@@ -312,15 +325,19 @@ std::vector<QueryFuncRef> GetCallersForAllBaseFunctions(QueryDatabase* db,
   if (!root.def)
     return callers;
 
-  std::queue<QueryFuncId> queue;
-  PushRange(&queue, root.def->base);
+  std::queue<QueryFunc*> queue;
+  EachWithGen<QueryFunc>(db->funcs, root.def->base, [&](QueryFunc& func1) {
+    queue.push(&func1);
+  });
   while (!queue.empty()) {
-    QueryFunc& func = db->funcs[queue.front().id];
+    QueryFunc& func = *queue.front();
     queue.pop();
 
     AddRange(&callers, func.callers);
     if (func.def)
-      PushRange(&queue, func.def->base);
+      EachWithGen<QueryFunc>(db->funcs, func.def->base, [&](QueryFunc& func1) {
+        queue.push(&func1);
+      });
   }
 
   return callers;
@@ -330,14 +347,18 @@ std::vector<QueryFuncRef> GetCallersForAllDerivedFunctions(QueryDatabase* db,
                                                            QueryFunc& root) {
   std::vector<QueryFuncRef> callers;
 
-  std::queue<QueryFuncId> queue;
-  PushRange(&queue, root.derived);
+  std::queue<QueryFunc*> queue;
+  EachWithGen<QueryFunc>(db->funcs, root.derived, [&](QueryFunc& func) {
+    queue.push(&func);
+  });
 
   while (!queue.empty()) {
-    QueryFunc& func = db->funcs[queue.front().id];
+    QueryFunc& func = *queue.front();
     queue.pop();
 
-    PushRange(&queue, func.derived);
+    EachWithGen<QueryFunc>(db->funcs, func.derived, [&](QueryFunc& func1) {
+      queue.push(&func1);
+    });
     AddRange(&callers, func.callers);
   }
 
@@ -493,7 +514,8 @@ optional<lsSymbolInformation> GetSymbolInfo(QueryDatabase* db,
       info.kind = lsSymbolKind::Function;
 
       if (func.def->declaring_type.has_value()) {
-        QueryType& container = db->types[func.def->declaring_type->id];
+        // FIXME WithGen
+        QueryType& container = db->types[func.def->declaring_type->value.id];
         if (container.def)
           info.kind = lsSymbolKind::Method;
       }
diff --git a/src/query_utils.h b/src/query_utils.h
index ac2bb2d3..6df1dca1 100644
--- a/src/query_utils.h
+++ b/src/query_utils.h
@@ -25,13 +25,14 @@ std::vector<QueryLocation> ToQueryLocation(
 std::vector<QueryLocation> ToQueryLocation(
     QueryDatabase* db,
     const std::vector<QueryTypeId>& refs);
+std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
+                                           std::vector<WithGen<QueryFuncId>>*);
 std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
                                            std::vector<WithGen<QueryTypeId>>*);
+std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
+                                           std::vector<WithGen<QueryVarId>>*);
 std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
                                            const std::vector<QueryFuncId>& ids);
-std::vector<QueryLocation> ToQueryLocation(
-    QueryDatabase* db,
-    std::vector<WithGen<QueryVarId>>*);
 std::vector<QueryLocation> GetUsesOfSymbol(QueryDatabase* db,
                                            const SymbolIdx& symbol,
                                            bool include_decl);
@@ -71,3 +72,25 @@ std::vector<SymbolRef> FindSymbolsAtLocation(WorkingFile* working_file,
 void EmitDiagnostics(WorkingFiles* working_files,
                      std::string path,
                      std::vector<lsDiagnostic> diagnostics);
+
+template <typename Q>
+void EachWithGen(std::vector<Q>& collection, WithGen<Id<Q>> x, std::function<void(Q&)> fn) {
+  Q& obj = collection[x.value.id];
+  // FIXME Deprecate optional<Def> def
+  if (obj.gen == x.gen && obj.def)
+    fn(obj);
+}
+
+template <typename Q>
+void EachWithGen(std::vector<Q>& collection, std::vector<WithGen<Id<Q>>>& ids, std::function<void(Q&)> fn) {
+  size_t j = 0;
+  for (WithGen<Id<Q>> x : ids) {
+    Q& obj = collection[x.value.id];
+    if (obj.gen == x.gen) {
+      if (obj.def) // FIXME Deprecate optional<Def> def
+        fn(obj);
+      ids[j++] = x;
+    }
+  }
+  ids.resize(j);
+}
diff --git a/src/serializer.h b/src/serializer.h
index c3439f56..b53aeeb5 100644
--- a/src/serializer.h
+++ b/src/serializer.h
@@ -122,27 +122,6 @@ struct IndexFile;
     visitor.EndArray();                                \
   }
 
-// API:
-/*
-template<typename TVisitor, typename T>
-void Reflect(TVisitor& visitor, T& value) {
-  static_assert(false, "Missing implementation");
-}
-template<typename TVisitor>
-void DefaultReflectMemberStart(TVisitor& visitor) {
-  static_assert(false, "Missing implementation");
-}
-template<typename TVisitor, typename T>
-bool ReflectMemberStart(TVisitor& visitor, T& value) {
-  static_assert(false, "Missing implementation");
-  return true;
-}
-template<typename TVisitor, typename T>
-void ReflectMemberEnd(TVisitor& visitor, T& value) {
-  static_assert(false, "Missing implementation");
-}
-*/
-
 //// Elementary types
 
 void Reflect(Reader& visitor, uint8_t& value);