From 37809def7cde195ce8295e3a19d2e8dedbb01f09 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Fri, 26 Jan 2018 17:31:50 -0800
Subject: [PATCH] Add `bool include_decl` parameter to GetUsesOfSymbol

Fixes #350

Function declarations/definitions are not counted if include_decl is
false.
We should do similar thing to Var and Type.
---
 src/messages/text_document_highlight.cc  |  2 +-
 src/messages/text_document_references.cc | 12 ++----------
 src/messages/text_document_rename.cc     |  2 +-
 src/query_utils.cc                       | 11 +++++++----
 src/query_utils.h                        |  5 +++--
 5 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/src/messages/text_document_highlight.cc b/src/messages/text_document_highlight.cc
index c417c82a..89229660 100644
--- a/src/messages/text_document_highlight.cc
+++ b/src/messages/text_document_highlight.cc
@@ -38,7 +38,7 @@ struct TextDocumentDocumentHighlightHandler
     for (const SymbolRef& ref :
          FindSymbolsAtLocation(working_file, file, request->params.position)) {
       // Found symbol. Return references to highlight.
-      std::vector<QueryLocation> uses = GetUsesOfSymbol(db, ref.idx);
+      std::vector<QueryLocation> uses = GetUsesOfSymbol(db, ref.idx, true);
       out.result.reserve(uses.size());
       for (const QueryLocation& use : uses) {
         if (use.path != file_id)
diff --git a/src/messages/text_document_references.cc b/src/messages/text_document_references.cc
index badf6a78..2617a1d6 100644
--- a/src/messages/text_document_references.cc
+++ b/src/messages/text_document_references.cc
@@ -53,19 +53,11 @@ struct TextDocumentReferencesHandler
 
     for (const SymbolRef& ref :
          FindSymbolsAtLocation(working_file, file, request->params.position)) {
-      optional<QueryLocation> excluded_declaration;
-      if (!request->params.context.includeDeclaration) {
-        LOG_S(INFO) << "Excluding declaration in references";
-        excluded_declaration = GetDefinitionSpellingOfSymbol(db, ref.idx);
-      }
-
       // Found symbol. Return references.
-      std::vector<QueryLocation> uses = GetUsesOfSymbol(db, ref.idx);
+      std::vector<QueryLocation> uses = GetUsesOfSymbol(
+          db, ref.idx, request->params.context.includeDeclaration);
       out.result.reserve(uses.size());
       for (const QueryLocation& use : uses) {
-        if (excluded_declaration.has_value() && use == *excluded_declaration)
-          continue;
-
         optional<lsLocation> ls_location =
             GetLsLocation(db, working_files, use);
         if (ls_location)
diff --git a/src/messages/text_document_rename.cc b/src/messages/text_document_rename.cc
index c6076280..7534c114 100644
--- a/src/messages/text_document_rename.cc
+++ b/src/messages/text_document_rename.cc
@@ -97,7 +97,7 @@ struct TextDocumentRenameHandler : BaseMessageHandler<Ipc_TextDocumentRename> {
     for (const SymbolRef& ref :
          FindSymbolsAtLocation(working_file, file, request->params.position)) {
       // Found symbol. Return references to rename.
-      std::vector<QueryLocation> uses = GetUsesOfSymbol(db, ref.idx);
+      std::vector<QueryLocation> uses = GetUsesOfSymbol(db, ref.idx, true);
       out.result =
           BuildWorkspaceEdit(db, working_files, uses, request->params.newName);
       break;
diff --git a/src/query_utils.cc b/src/query_utils.cc
index 910ab84b..112c16a0 100644
--- a/src/query_utils.cc
+++ b/src/query_utils.cc
@@ -182,7 +182,8 @@ std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
 }
 
 std::vector<QueryLocation> GetUsesOfSymbol(QueryDatabase* db,
-                                           const SymbolIdx& symbol) {
+                                           const SymbolIdx& symbol,
+                                           bool include_decl) {
   switch (symbol.kind) {
     case SymbolKind::Type: {
       QueryType& type = db->types[symbol.idx];
@@ -192,9 +193,11 @@ std::vector<QueryLocation> GetUsesOfSymbol(QueryDatabase* db,
       // TODO: the vector allocation could be avoided.
       QueryFunc& func = db->funcs[symbol.idx];
       std::vector<QueryLocation> result = ToQueryLocation(db, func.callers);
-      AddRange(&result, func.declarations);
-      if (func.def && func.def->definition_spelling)
-        result.push_back(*func.def->definition_spelling);
+      if (include_decl) {
+        AddRange(&result, func.declarations);
+        if (func.def && func.def->definition_spelling)
+          result.push_back(*func.def->definition_spelling);
+      }
       return result;
     }
     case SymbolKind::Var: {
diff --git a/src/query_utils.h b/src/query_utils.h
index c4871ada..4e7f994c 100644
--- a/src/query_utils.h
+++ b/src/query_utils.h
@@ -29,7 +29,8 @@ std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
 std::vector<QueryLocation> ToQueryLocation(QueryDatabase* db,
                                            const std::vector<QueryVarId>& ids);
 std::vector<QueryLocation> GetUsesOfSymbol(QueryDatabase* db,
-                                           const SymbolIdx& symbol);
+                                           const SymbolIdx& symbol,
+                                           bool include_decl);
 std::vector<QueryLocation> GetDeclarationsOfSymbolForGotoDefinition(
     QueryDatabase* db,
     const SymbolIdx& symbol);
@@ -65,4 +66,4 @@ std::vector<SymbolRef> FindSymbolsAtLocation(WorkingFile* working_file,
 
 void EmitDiagnostics(WorkingFiles* working_files,
                      std::string path,
-                     std::vector<lsDiagnostic> diagnostics);
\ No newline at end of file
+                     std::vector<lsDiagnostic> diagnostics);