Track data flow into variables and function returns

CMakeLists.txt

@@ -211,6 +211,7 @@ target_sources(ccls PRIVATE
   src/messages/ccls_navigate.cc
   src/messages/ccls_reload.cc
   src/messages/ccls_vars.cc
+  src/messages/ccls_dataFlowInto.cc
   src/messages/initialize.cc
   src/messages/textDocument_code.cc
   src/messages/textDocument_completion.cc

src/indexer.cc

@@ -22,6 +22,7 @@ limitations under the License.
 #include "sema_manager.hh"
 
 #include <clang/AST/AST.h>
+#include <clang/AST/RecursiveASTVisitor.h>
 #include <clang/Frontend/FrontendAction.h>
 #include <clang/Index/IndexDataConsumer.h>
 #include <clang/Index/IndexingAction.h>
@@ -681,6 +682,8 @@ public:
   }
 
 public:
+  std::unordered_set<const Decl *> visited_decls;
+
   IndexDataConsumer(IndexParam &param) : param(param) {}
   void initialize(ASTContext &Ctx) override {
     this->Ctx = param.Ctx = &Ctx;
@@ -796,6 +799,194 @@ public:
     IndexParam::DeclInfo *info;
     Usr usr = GetUsr(D, &info);
 
+    if (true) {
+      class FuncVisitor : public RecursiveASTVisitor<FuncVisitor> {
+      public:
+        FuncVisitor(IndexDataConsumer *consumer, IndexFile *db, SourceManager &SM,
+                    const LangOptions &Lang, int lid)
+            : consumer(consumer), db(db), SM(SM), Lang(Lang), lid(lid) {}
+
+        std::stack<FunctionDecl*> function_stack;
+
+        bool TraverseDecl(Decl *decl) {
+          if (decl->hasBody() && !consumer->visited_decls.insert(decl).second)
+            return true;
+          return RecursiveASTVisitor::TraverseDecl(decl);
+        }
+        #define TRAVERSE_FUNCTION_DECL_CHILD(TYPE) \
+        bool Traverse##TYPE(TYPE *decl) {\
+          function_stack.push(decl);\
+          bool ret = RecursiveASTVisitor::Traverse##TYPE(decl);\
+          function_stack.pop();\
+          return ret;\
+        }
+        TRAVERSE_FUNCTION_DECL_CHILD(FunctionDecl)
+        TRAVERSE_FUNCTION_DECL_CHILD(CXXDeductionGuideDecl)
+        TRAVERSE_FUNCTION_DECL_CHILD(CXXMethodDecl)
+        TRAVERSE_FUNCTION_DECL_CHILD(CXXConstructorDecl)
+        TRAVERSE_FUNCTION_DECL_CHILD(CXXConversionDecl)
+        TRAVERSE_FUNCTION_DECL_CHILD(CXXDestructorDecl)
+
+        bool VisitUnaryOperator(UnaryOperator *unop) {
+          if (unop->isIncrementDecrementOp()) {
+            AddDataFlow(unop->getSubExpr(), unop, unop->getSourceRange());
+          }
+          return RecursiveASTVisitor::VisitUnaryOperator(unop);
+        }
+        bool VisitBinaryOperator(BinaryOperator *binop) {
+          if (binop->getOpcode() == BinaryOperator::Opcode::BO_Assign) {
+            AddDataFlow(binop->getLHS(), binop->getRHS(), binop->getSourceRange());
+          } else if (binop->isAssignmentOp()) {
+            AddDataFlow(binop->getLHS(), binop, binop->getSourceRange());
+          }
+          return RecursiveASTVisitor::VisitBinaryOperator(binop);
+        }
+        bool VisitVarDecl(VarDecl *decl) {
+          if (decl->hasInit()) {
+            AddDataFlow(decl, decl->getInit(), decl->getSourceRange());
+          }
+          return RecursiveASTVisitor::VisitVarDecl(decl);
+        }
+        bool VisitFieldDecl(FieldDecl *decl) {
+          if (decl->hasInClassInitializer()) {
+            if (decl->getInClassInitializer()) {
+              AddDataFlow(decl, decl->getInClassInitializer(), decl->getSourceRange());
+            } else {
+              // Allow re-visit of this decl.
+              consumer->visited_decls.erase(decl);
+            }
+          }
+          return RecursiveASTVisitor::VisitFieldDecl(decl);
+        }
+        bool VisitCXXConstructorDecl(CXXConstructorDecl *decl) {
+          for (const auto &init : decl->inits()) {
+            if (init->isMemberInitializer()) {
+              if (auto *defaultInit =
+                      dyn_cast<CXXDefaultInitExpr>(init->getInit())) {
+                AddDataFlow(init->getMember(), defaultInit->getExpr(), defaultInit->getSourceRange());
+              } else {
+                AddDataFlow(init->getMember(), init->getInit(), init->getSourceRange());
+              }
+            }
+            
+            if (auto *constructorInit =
+                    dyn_cast<CXXConstructExpr>(init->getInit())) {
+              if (CXXConstructorDecl *decl = constructorInit->getConstructor()) {
+                int len = std::min(constructorInit->getNumArgs(), decl->getNumParams());
+                for (int i = 0; i < len; ++i) {
+                  AddDataFlow(decl->getParamDecl(i), constructorInit->getArg(i), constructorInit->getSourceRange());
+                }
+              }
+            }
+          }
+          return RecursiveASTVisitor::VisitCXXConstructorDecl(decl);
+        }
+        bool VisitCXXDefaultInitExpr(CXXDefaultInitExpr *expr) {
+          AddDataFlow(expr->getField(), expr->getExpr(), expr->getSourceRange());
+          return RecursiveASTVisitor::VisitCXXDefaultInitExpr(expr);
+        }
+
+        bool VisitDeclStmt(DeclStmt *stmt) {
+          for (auto &decl : stmt->decls()) {
+            if (VarDecl *named = dyn_cast<VarDecl>(decl)) {
+              if (named->hasInit()) {
+                AddDataFlow(named, named->getInit(), named->getSourceRange());
+              }
+            }
+          }
+          return RecursiveASTVisitor::VisitDeclStmt(stmt);
+        }
+
+        bool VisitCallExpr(CallExpr *expr) {
+          if (FunctionDecl *decl = expr->getDirectCallee()) {
+            int len = std::min(expr->getNumArgs(), decl->getNumParams());
+            for (int i = 0; i < len; ++i) {
+              AddDataFlow(decl->getParamDecl(i), expr->getArg(i), expr->getSourceRange());
+            }
+          }
+          return RecursiveASTVisitor::VisitCallExpr(expr);
+        }
+        bool VisitCXXConstructExpr(CXXConstructExpr *expr) {
+          if (CXXConstructorDecl *decl = expr->getConstructor()) {
+            int len = std::min(expr->getNumArgs(), decl->getNumParams());
+            for (int i = 0; i < len; ++i) {
+              AddDataFlow(decl->getParamDecl(i), expr->getArg(i), expr->getSourceRange());
+            }
+          }
+          return RecursiveASTVisitor::VisitCXXConstructExpr(expr);
+        }
+
+        bool VisitReturnStmt(ReturnStmt *stmt) {
+          if (!function_stack.empty()) {
+            if (Expr *retVal = stmt->getRetValue()) {
+              if (const FunctionDecl *func = dyn_cast<FunctionDecl>(function_stack.top()))
+                AddDataFlow(func, retVal, stmt->getSourceRange());
+            }
+          }
+          return RecursiveASTVisitor::VisitReturnStmt(stmt);
+        }
+
+      private:
+        void AddDataFlow(const Expr *target, const Expr *source, SourceRange flowSourceRange) {
+          if (!source)
+            target->dumpColor();
+          if (const auto *targetDeclRef =
+                  dyn_cast<DeclRefExpr>(target->IgnoreParenCasts())) {
+            AddDataFlow(targetDeclRef->getDecl(), source, flowSourceRange);
+          } else if (const auto *targetFieldRef =
+                         dyn_cast<MemberExpr>(target->IgnoreParenCasts())) {
+            AddDataFlow(targetFieldRef->getMemberDecl(), source, flowSourceRange);
+          }
+        }
+        void AddDataFlow(const ValueDecl *target, const Expr *source, SourceRange flowSourceRange) {
+          if (!source)
+            target->dumpColor();
+          auto& var = db->ToVar(consumer->GetUsr(target));
+          var.data_flow_into.push_back(GetDataFlow(source, flowSourceRange));
+        }
+        void AddDataFlow(const FunctionDecl *target, const Expr *source, SourceRange flowSourceRange) {
+          if (!source)
+            target->dumpColor();
+          auto& func = db->ToFunc(consumer->GetUsr(target));
+          func.data_flow_into_return.push_back(GetDataFlow(source, flowSourceRange));
+        }
+        DataFlow GetDataFlow(const Expr *source, SourceRange flowSourceRange) {
+          auto data_flowRange = FromCharSourceRange(
+              SM, Lang,
+              CharSourceRange::getTokenRange(flowSourceRange));
+          Use write{{data_flowRange, Role::None}, lid};
+
+          if (const auto *data_flowExpr =
+                  dyn_cast<DeclRefExpr>(source->IgnoreParenCasts())) {
+            auto *data_flow_decl = data_flowExpr->getDecl();
+            write.role = Role::Read;
+            return DataFlow{consumer->GetUsr(data_flow_decl), write};
+          } else if (const auto *data_flowExpr =
+                         dyn_cast<MemberExpr>(source->IgnoreParenCasts())) {
+            auto *data_flow_decl = data_flowExpr->getMemberDecl();
+            write.role = Role::Read;
+            return DataFlow{consumer->GetUsr(data_flow_decl), write};
+          } else if (const auto *data_flowExpr =
+                         dyn_cast<CallExpr>(source->IgnoreParenCasts())) {
+            if (auto *data_flow_callee = data_flowExpr->getDirectCallee()) {
+              write.role = Role::Call;
+              return DataFlow{consumer->GetUsr(data_flow_callee), write};
+            }
+          }
+
+          return DataFlow{0, write};
+        }
+
+        IndexDataConsumer *consumer;
+        IndexFile *db;
+        SourceManager &SM;
+        const LangOptions &Lang;
+        int lid;
+      };
+      FuncVisitor funcVisitor(this, db, SM, Lang, lid);
+      funcVisitor.TraverseDecl(const_cast<Decl *>(ASTNode.OrigD));
+    }
+
     auto do_def_decl = [&](auto *entity) {
       Use use{{loc, role}, lid};
       if (is_def) {
@@ -1183,8 +1374,8 @@ public:
 };
 } // namespace
 
-const int IndexFile::kMajorVersion = 19;
+const int IndexFile::kMajorVersion = 20;
-const int IndexFile::kMinorVersion = 1;
+const int IndexFile::kMinorVersion = 0;
 
 IndexFile::IndexFile(llvm::sys::fs::UniqueID UniqueID, const std::string &path,
                      const std::string &contents)
@@ -1349,6 +1540,7 @@ Index(SemaManager *manager, WorkingFiles *wfiles, VFS *vfs,
       // e.g. declaration + out-of-line definition
       Uniquify(it.second.derived);
       Uniquify(it.second.uses);
+      Uniquify(it.second.data_flow_into_return);
     }
     for (auto &it : entry->usr2type) {
       Uniquify(it.second.derived);
@@ -1357,8 +1549,10 @@ Index(SemaManager *manager, WorkingFiles *wfiles, VFS *vfs,
       Uniquify(it.second.def.bases);
       Uniquify(it.second.def.funcs);
     }
-    for (auto &it : entry->usr2var)
+    for (auto &it : entry->usr2var) {
       Uniquify(it.second.uses);
+      Uniquify(it.second.data_flow_into);
+    }
 
     // Update dependencies for the file.
     for (auto &[_, file] : param.UID2File) {
@@ -1404,6 +1598,17 @@ void Reflect(JsonReader &vis, DeclRef &v) {
   v.file_id = static_cast<int>(strtol(s + 1, &s, 10));
 }
 
+void Reflect(JsonReader &vis, DataFlow &v) {
+  std::string t = vis.GetString();
+  char *s = const_cast<char *>(t.c_str());
+  v.from = strtoull(s + 1, &s, 10);
+  s = strchr(s, '|') + 1;
+  v.use.range = Range::FromString(s);
+  s = strchr(s, '|');
+  v.use.role = static_cast<Role>(strtol(s + 1, &s, 10));
+  v.use.file_id = static_cast<int>(strtol(s + 1, &s, 10));
+}
+
 void Reflect(JsonWriter &vis, SymbolRef &v) {
   char buf[99];
   snprintf(buf, sizeof buf, "%s|%" PRIu64 "|%d|%d", v.range.ToString().c_str(),
@@ -1425,6 +1630,13 @@ void Reflect(JsonWriter &vis, DeclRef &v) {
   std::string s(buf);
   Reflect(vis, s);
 }
+void Reflect(JsonWriter &vis, DataFlow &v) {
+  char buf[99];
+  snprintf(buf, sizeof buf, "%" PRIu64 "|%s|%d|%d", v.from,
+           v.use.range.ToString().c_str(), int(v.use.role), v.use.file_id);
+  std::string s(buf);
+  Reflect(vis, s);
+}
 
 void Reflect(BinaryReader &vis, SymbolRef &v) {
   Reflect(vis, v.range);
@@ -1441,6 +1653,10 @@ void Reflect(BinaryReader &vis, DeclRef &v) {
   Reflect(vis, static_cast<Use &>(v));
   Reflect(vis, v.extent);
 }
+void Reflect(BinaryReader &vis, DataFlow &v) {
+  Reflect(vis, v.from);
+  Reflect(vis, v.use);
+}
 
 void Reflect(BinaryWriter &vis, SymbolRef &v) {
   Reflect(vis, v.range);
@@ -1457,4 +1673,9 @@ void Reflect(BinaryWriter &vis, DeclRef &v) {
   Reflect(vis, static_cast<Use &>(v));
   Reflect(vis, v.extent);
 }
+void Reflect(BinaryWriter &vis, DataFlow &v) {
+  Reflect(vis, v.from);
+  Reflect(vis, v.use);
+}
+
 } // namespace ccls

src/indexer.hh

@@ -130,18 +130,34 @@ struct DeclRef : Use {
   Range extent;
 };
 
+struct DataFlow {
+  Usr from;
+  Use use;
+
+  bool operator==(const DataFlow &o) const {
+    return from == o.from && use == o.use;
+  }
+  bool operator<(const DataFlow &o) const {
+    return from != o.from ? from < o.from : use < o.use;
+  }
+};
+
 void Reflect(JsonReader &visitor, SymbolRef &value);
 void Reflect(JsonReader &visitor, Use &value);
 void Reflect(JsonReader &visitor, DeclRef &value);
+void Reflect(JsonReader &visitor, DataFlow &value);
 void Reflect(JsonWriter &visitor, SymbolRef &value);
 void Reflect(JsonWriter &visitor, Use &value);
 void Reflect(JsonWriter &visitor, DeclRef &value);
+void Reflect(JsonWriter &visitor, DataFlow &value);
 void Reflect(BinaryReader &visitor, SymbolRef &value);
 void Reflect(BinaryReader &visitor, Use &value);
 void Reflect(BinaryReader &visitor, DeclRef &value);
+void Reflect(BinaryReader &visitor, DataFlow &value);
 void Reflect(BinaryWriter &visitor, SymbolRef &value);
 void Reflect(BinaryWriter &visitor, Use &value);
 void Reflect(BinaryWriter &visitor, DeclRef &value);
+void Reflect(BinaryWriter &visitor, DataFlow &value);
 
 template <typename D> struct NameMixin {
   std::string_view Name(bool qualified) const {
@@ -191,6 +207,7 @@ struct IndexFunc : NameMixin<IndexFunc> {
   std::vector<DeclRef> declarations;
   std::vector<Usr> derived;
   std::vector<Use> uses;
+  std::vector<DataFlow> data_flow_into_return;
 };
 
 struct TypeDef : NameMixin<TypeDef> {
@@ -271,6 +288,7 @@ struct IndexVar {
   Def def;
   std::vector<DeclRef> declarations;
   std::vector<Use> uses;
+  std::vector<DataFlow> data_flow_into;
 };
 
 struct IndexInclude {
@@ -349,3 +367,4 @@ MAKE_HASHABLE(ccls::SymbolRef, t.range, t.usr, t.kind, t.role);
 MAKE_HASHABLE(ccls::ExtentRef, t.range, t.usr, t.kind, t.role, t.extent);
 MAKE_HASHABLE(ccls::Use, t.range, t.file_id)
 MAKE_HASHABLE(ccls::DeclRef, t.range, t.file_id)
+MAKE_HASHABLE(ccls::DataFlow, t.from, t.use)

src/message_handler.cc

@@ -163,6 +163,7 @@ MessageHandler::MessageHandler() {
   Bind("$ccls/navigate", &MessageHandler::ccls_navigate);
   Bind("$ccls/reload", &MessageHandler::ccls_reload);
   Bind("$ccls/vars", &MessageHandler::ccls_vars);
+  Bind("$ccls/dataFlowInto", &MessageHandler::ccls_dataFlowInto);
   Bind("exit", &MessageHandler::exit);
   Bind("initialize", &MessageHandler::initialize);
   Bind("shutdown", &MessageHandler::shutdown);
@@ -275,7 +276,7 @@ void EmitSemanticHighlight(DB *db, WorkingFile *wfile, QueryFile &file) {
 
   // Group symbols together.
   std::unordered_map<SymbolIdx, CclsSemanticHighlightSymbol> grouped_symbols;
-  for (auto &[sym, refcnt] : file.symbol2refcnt) {
+  for (auto [sym, refcnt] : file.symbol2refcnt) {
     if (refcnt <= 0) continue;
     std::string_view detailed_name;
     SymbolKind parent_kind = SymbolKind::Unknown;
@@ -285,7 +286,10 @@ void EmitSemanticHighlight(DB *db, WorkingFile *wfile, QueryFile &file) {
     // This switch statement also filters out symbols that are not highlighted.
     switch (sym.kind) {
     case Kind::Func: {
-      idx = db->func_usr[sym.usr];
+      auto func_it = db->func_usr.find(sym.usr);
+      if (func_it == db->func_usr.end())
+        continue;
+      idx = func_it->second;
       const QueryFunc &func = db->funcs[idx];
       const QueryFunc::Def *def = func.AnyDef();
       if (!def)
@@ -319,7 +323,10 @@ void EmitSemanticHighlight(DB *db, WorkingFile *wfile, QueryFile &file) {
       break;
     }
     case Kind::Type: {
-      idx = db->type_usr[sym.usr];
+      auto type_it = db->type_usr.find(sym.usr);
+      if (type_it == db->type_usr.end())
+        continue;
+      idx = type_it->second;
       const QueryType &type = db->types[idx];
       for (auto &def : type.def) {
         kind = def.kind;
@@ -332,7 +339,10 @@ void EmitSemanticHighlight(DB *db, WorkingFile *wfile, QueryFile &file) {
       break;
     }
     case Kind::Var: {
-      idx = db->var_usr[sym.usr];
+      auto var_it = db->var_usr.find(sym.usr);
+      if (var_it == db->var_usr.end())
+        continue;
+      idx = var_it->second;
       const QueryVar &var = db->vars[idx];
       for (auto &def : var.def) {
         kind = def.kind;

src/message_handler.hh

@@ -246,6 +246,7 @@ private:
   void ccls_navigate(JsonReader &, ReplyOnce &);
   void ccls_reload(JsonReader &);
   void ccls_vars(JsonReader &, ReplyOnce &);
+  void ccls_dataFlowInto(JsonReader &, ReplyOnce &);
   void exit(EmptyParam &);
   void initialize(JsonReader &, ReplyOnce &);
   void shutdown(EmptyParam &, ReplyOnce &);

src/messages/ccls_dataFlowInto.cc

@@ -0,0 +1,108 @@
+/* Copyright 2017-2018 ccls Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "message_handler.hh"
+#include "pipeline.hh"
+#include "query.hh"
+
+#include <unordered_set>
+
+namespace ccls {
+namespace {
+struct Param : TextDocumentPositionParam {
+};
+REFLECT_STRUCT(Param, textDocument, position);
+} // namespace
+
+struct Out_cclsDataFlowInto {
+  int id;
+  Location location;
+  // Empty if the |levels| limit is reached.
+  std::vector<Out_cclsDataFlowInto> children;
+};
+REFLECT_STRUCT(Out_cclsDataFlowInto, id ,location, children);
+
+Out_cclsDataFlowInto BuildDataFlow(Usr usr, Kind kind, Use use, std::unordered_set<Usr> seen, DB* db, WorkingFiles* wfiles, int& id) {
+  Out_cclsDataFlowInto result = {id++};
+  if (auto loc = GetLsLocation(db, wfiles, use)) {
+    result.location = *loc;
+  }
+  if (!seen.insert(usr).second) return result;
+
+  const std::vector<DataFlow>* data_flow_sources = nullptr;
+  if (kind == Kind::Var) {
+    auto& var = db->Var(usr);
+    data_flow_sources = &var.data_flow_into;
+  } else if (kind == Kind::Func) {
+    auto& func = db->Func(usr);
+    data_flow_sources = &func.data_flow_into_return;
+  }
+  if (data_flow_sources != nullptr) {
+    if (data_flow_sources->empty()) {
+      if (kind == Kind::Var) {
+        auto& var = db->Var(usr);
+        for (auto& def : var.def) {
+          if (def.spell) {
+            if (auto loc = GetLsLocation(db, wfiles, *def.spell)) {
+              result.children.push_back({id++, *loc});
+            }
+          }
+        }
+      } else if (kind == Kind::Func) {
+        auto& func = db->Func(usr);
+        for (auto& def : func.def) {
+          if (def.spell) {
+            if (auto loc = GetLsLocation(db, wfiles, *def.spell)) {
+              result.children.push_back({id++, *loc});
+            }
+          }
+        }
+      }
+    } else {
+      for (auto& write : *data_flow_sources) {
+        if (write.use.role == Role::Read) {
+          result.children.push_back(BuildDataFlow(write.from, Kind::Var, write.use, seen, db, wfiles, id));
+        } else if (write.use.role == Role::Call) {
+          result.children.push_back(BuildDataFlow(write.from, Kind::Func, write.use, seen, db, wfiles, id));
+        } else if (auto loc = GetLsLocation(db, wfiles, write.use)) {
+          result.children.push_back({id++, *loc});
+        }
+      }
+    }
+  }
+  return result;
+}
+
+void MessageHandler::ccls_dataFlowInto(JsonReader &reader, ReplyOnce &reply) {
+  Param param;
+  Reflect(reader, param);
+  QueryFile *file = FindFile(param.textDocument.uri.GetPath());
+  WorkingFile *wf = file ? wfiles->GetFile(file->def->path) : nullptr;
+  if (!wf) {
+    reply.NotReady(file);
+    return;
+  }
+
+  std::optional<Out_cclsDataFlowInto> result;
+  for (SymbolRef sym : FindSymbolsAtLocation(wf, file, param.position)) {
+    // Found symbol. Return references.
+    int id = 0;
+    std::unordered_set<Usr> seen = {};
+    result = BuildDataFlow(sym.usr, sym.kind, Use{{sym.range, sym.role},file->id}, seen, db, wfiles, id);
+    break;
+  }
+  reply(result);
+}
+} // namespace ccls

src/query.cc

@@ -99,6 +99,7 @@ IndexUpdate IndexUpdate::CreateDelta(IndexFile *previous, IndexFile *current) {
     r.funcs_declarations[func.usr].first = std::move(func.declarations);
     r.funcs_uses[func.usr].first = std::move(func.uses);
     r.funcs_derived[func.usr].first = std::move(func.derived);
+    r.funcs_data_flow_into_return[func.usr].first = std::move(func.data_flow_into_return);
   }
   for (auto &it : current->usr2func) {
     auto &func = it.second;
@@ -107,6 +108,7 @@ IndexUpdate IndexUpdate::CreateDelta(IndexFile *previous, IndexFile *current) {
     r.funcs_declarations[func.usr].second = std::move(func.declarations);
     r.funcs_uses[func.usr].second = std::move(func.uses);
     r.funcs_derived[func.usr].second = std::move(func.derived);
+    r.funcs_data_flow_into_return[func.usr].second = std::move(func.data_flow_into_return);
   }
 
   r.types_hint = current->usr2type.size() - previous->usr2type.size();
@@ -136,6 +138,7 @@ IndexUpdate IndexUpdate::CreateDelta(IndexFile *previous, IndexFile *current) {
       r.vars_removed.emplace_back(var.usr, var.def);
     r.vars_declarations[var.usr].first = std::move(var.declarations);
     r.vars_uses[var.usr].first = std::move(var.uses);
+    r.vars_data_flow_into[var.usr].first = std::move(var.data_flow_into);
   }
   for (auto &it : current->usr2var) {
     auto &var = it.second;
@@ -143,6 +146,7 @@ IndexUpdate IndexUpdate::CreateDelta(IndexFile *previous, IndexFile *current) {
       r.vars_def_update.emplace_back(it.first, var.def);
     r.vars_declarations[var.usr].second = std::move(var.declarations);
     r.vars_uses[var.usr].second = std::move(var.uses);
+    r.vars_data_flow_into[var.usr].second = std::move(var.data_flow_into);
   }
 
   r.files_def_update = BuildFileDefUpdate(std::move(*current));
@@ -239,6 +243,8 @@ void DB::ApplyIndexUpdate(IndexUpdate *u) {
                  Use &use, int delta) {
     use.file_id =
         use.file_id == -1 ? u->file_id : lid2fid.find(use.file_id)->second;
+    if (usr == 0) return;
+
     ExtentRef sym{{use.range, usr, kind, use.role}};
     int &v = files[use.file_id].symbol2refcnt[sym];
     v += delta;
@@ -319,6 +325,20 @@ void DB::ApplyIndexUpdate(IndexUpdate *u) {
   REMOVE_ADD(func, derived);
   for (auto &[usr, p] : u->funcs_uses)
     UpdateUses(usr, Kind::Func, func_usr, funcs, p, true);
+  for (auto &[usr, p] : u->funcs_data_flow_into_return) {
+    auto R = func_usr.try_emplace(usr, func_usr.size());
+    if (R.second)
+      funcs.emplace_back().usr = usr;
+    auto &func = funcs[R.first->second];
+    for (DataFlow &src : p.first) {
+      Ref(prev_lid2file_id, src.from, src.use.role == Role::Read ? Kind::Var : Kind::Func, src.use, -1);
+    }
+    RemoveRange(func.data_flow_into_return, p.first);
+    for (DataFlow &src : p.second) {
+      Ref(lid2file_id, src.from, src.use.role == Role::Read ? Kind::Var : Kind::Func, src.use, 1);
+    }
+    AddRange(func.data_flow_into_return, p.second);
+  }
 
   if ((t = types.size() + u->types_hint) > types.capacity()) {
     t = size_t(t * grow);
@@ -361,6 +381,20 @@ void DB::ApplyIndexUpdate(IndexUpdate *u) {
   REMOVE_ADD(var, declarations);
   for (auto &[usr, p] : u->vars_uses)
     UpdateUses(usr, Kind::Var, var_usr, vars, p, false);
+  for (auto &[usr, p] : u->vars_data_flow_into) {
+    auto R = var_usr.try_emplace(usr, var_usr.size());
+    if (R.second)
+      vars.emplace_back().usr = usr;
+    auto &var = vars[R.first->second];
+    for (DataFlow &src : p.first) {
+      Ref(prev_lid2file_id, src.from, src.use.role == Role::Read ? Kind::Var : Kind::Func, src.use, -1);
+    }
+    RemoveRange(var.data_flow_into, p.first);
+    for (DataFlow &src : p.second) {
+      Ref(lid2file_id, src.from, src.use.role == Role::Read ? Kind::Var : Kind::Func, src.use, 1);
+    }
+    AddRange(var.data_flow_into, p.second);
+  }
 
 #undef REMOVE_ADD
 }

src/query.hh

@@ -83,6 +83,7 @@ struct QueryFunc : QueryEntity<QueryFunc, FuncDef> {
   std::vector<DeclRef> declarations;
   std::vector<Usr> derived;
   std::vector<Use> uses;
+  std::vector<DataFlow> data_flow_into_return;
 };
 
 struct QueryType : QueryEntity<QueryType, TypeDef> {
@@ -99,6 +100,7 @@ struct QueryVar : QueryEntity<QueryVar, VarDef> {
   llvm::SmallVector<Def, 1> def;
   std::vector<DeclRef> declarations;
   std::vector<Use> uses;
+  std::vector<DataFlow> data_flow_into;
 };
 
 struct IndexUpdate {
@@ -125,6 +127,7 @@ struct IndexUpdate {
   Update<DeclRef> funcs_declarations;
   Update<Use> funcs_uses;
   Update<Usr> funcs_derived;
+  Update<DataFlow> funcs_data_flow_into_return;
 
   // Type updates.
   int types_hint;
@@ -141,6 +144,7 @@ struct IndexUpdate {
   std::vector<std::pair<Usr, QueryVar::Def>> vars_def_update;
   Update<DeclRef> vars_declarations;
   Update<Use> vars_uses;
+  Update<DataFlow> vars_data_flow_into;
 };
 
 struct DenseMapInfoForUsr {

src/serializer.cc

@@ -308,6 +308,7 @@ template <typename TVisitor> void Reflect1(TVisitor &vis, IndexFunc &v) {
   REFLECT_MEMBER2("declarations", v.declarations);
   REFLECT_MEMBER2("derived", v.derived);
   REFLECT_MEMBER2("uses", v.uses);
+  REFLECT_MEMBER2("data_flow_into_return", v.data_flow_into_return);
   ReflectMemberEnd(vis);
 }
 void Reflect(JsonReader &vis, IndexFunc &v) { Reflect1(vis, v); }
@@ -357,6 +358,7 @@ template <typename TVisitor> void Reflect1(TVisitor &vis, IndexVar &v) {
 
   REFLECT_MEMBER2("declarations", v.declarations);
   REFLECT_MEMBER2("uses", v.uses);
+  REFLECT_MEMBER2("data_flow_into", v.data_flow_into);
   ReflectMemberEnd(vis);
 }
 void Reflect(JsonReader &vis, IndexVar &v) { Reflect1(vis, v); }