Repo created

libs/search/bookmarks/processor.cpp

@@ -0,0 +1,281 @@
+#include "search/bookmarks/processor.hpp"
+
+#include "search/emitter.hpp"
+
+#include "base/assert.hpp"
+#include "base/checked_cast.hpp"
+#include "base/dfa_helpers.hpp"
+#include "base/levenshtein_dfa.hpp"
+#include "base/stl_helpers.hpp"
+#include "base/string_utils.hpp"
+
+#include <algorithm>
+
+namespace search
+{
+namespace bookmarks
+{
+namespace
+{
+struct DocVecWrapper
+{
+  explicit DocVecWrapper(DocVec const & dv) : m_dv(dv) {}
+
+  template <typename Fn>
+  void ForEachToken(Fn && fn) const
+  {
+    for (size_t i = 0; i < m_dv.GetNumTokens(); ++i)
+      fn(StringUtf8Multilang::kDefaultCode, m_dv.GetToken(i));
+  }
+
+  DocVec const & m_dv;
+};
+
+struct RankingInfo
+{
+  bool operator<(RankingInfo const & rhs) const { return m_cosineSimilarity > rhs.m_cosineSimilarity; }
+
+  bool operator>(RankingInfo const & rhs) const { return rhs < *this; }
+
+  bool operator==(RankingInfo const & rhs) const { return !(*this < rhs) && !(*this > rhs); }
+  bool operator!=(RankingInfo const & rhs) const { return !(*this == rhs); }
+
+  double m_cosineSimilarity = 0.0;
+};
+
+struct IdInfoPair
+{
+  IdInfoPair(Id const & id, RankingInfo const & info) : m_id(id), m_info(info) {}
+
+  bool operator<(IdInfoPair const & rhs) const
+  {
+    if (m_info != rhs.m_info)
+      return m_info < rhs.m_info;
+    return m_id < rhs.m_id;
+  }
+
+  Id m_id;
+  RankingInfo m_info;
+};
+
+void FillRankingInfo(QueryVec & qv, IdfMap & idfs, DocVec const & dv, RankingInfo & info)
+{
+  info.m_cosineSimilarity = qv.Similarity(idfs, dv);
+}
+}  // namespace
+
+Processor::Processor(Emitter & emitter, base::Cancellable const & cancellable)
+  : m_emitter(emitter)
+  , m_cancellable(cancellable)
+{}
+
+void Processor::Reset()
+{
+  m_index = {};
+  m_docs.clear();
+  m_indexDescriptions = false;
+  m_indexableGroups.clear();
+  m_idToGroup.clear();
+  m_bookmarksInGroup.clear();
+}
+
+void Processor::EnableIndexingOfDescriptions(bool enable)
+{
+  m_indexDescriptions = enable;
+}
+
+void Processor::EnableIndexingOfBookmarkGroup(GroupId const & groupId, bool enable)
+{
+  bool const wasIndexable = m_indexableGroups.count(groupId) > 0;
+  if (enable)
+    m_indexableGroups.insert(groupId);
+  else
+    m_indexableGroups.erase(groupId);
+  bool const nowIndexable = m_indexableGroups.count(groupId) > 0;
+
+  if (wasIndexable == nowIndexable)
+    return;
+
+  for (auto const & id : m_bookmarksInGroup[groupId])
+    if (nowIndexable)
+      AddToIndex(id);
+    else
+      EraseFromIndex(id);
+}
+
+void Processor::Add(Id const & id, Doc const & doc)
+{
+  ASSERT_EQUAL(m_docs.count(id), 0, ());
+
+  DocVec::Builder builder;
+  doc.ForEachNameToken([&](int8_t /* lang */, strings::UniString const & token) { builder.Add(token); });
+
+  if (m_indexDescriptions)
+    doc.ForEachDescriptionToken([&](int8_t /* lang */, strings::UniString const & token) { builder.Add(token); });
+
+  DocVec const docVec(builder);
+
+  m_docs[id] = docVec;
+}
+
+void Processor::AddToIndex(Id const & id)
+{
+  ASSERT_EQUAL(m_docs.count(id), 1, ());
+
+  m_index.Add(id, DocVecWrapper(m_docs[id]));
+}
+
+void Processor::Update(Id const & id, Doc const & doc)
+{
+  auto group = kInvalidGroupId;
+  auto const groupIt = m_idToGroup.find(id);
+  if (groupIt != m_idToGroup.end())
+  {
+    // A copy to avoid use-after-free.
+    group = groupIt->second;
+    DetachFromGroup(id, group);
+  }
+
+  Erase(id);
+  Add(id, doc);
+
+  if (group != kInvalidGroupId)
+    AttachToGroup(id, group);
+}
+
+void Processor::Erase(Id const & id)
+{
+  ASSERT_EQUAL(m_docs.count(id), 1, ());
+
+  ASSERT(m_idToGroup.find(id) == m_idToGroup.end(),
+         ("A bookmark must be detached from all groups before being deleted."));
+
+  m_docs.erase(id);
+}
+
+void Processor::EraseFromIndex(Id const & id)
+{
+  ASSERT_EQUAL(m_docs.count(id), 1, ());
+
+  auto const & docVec = m_docs[id];
+  m_index.Erase(id, DocVecWrapper(docVec));
+}
+
+void Processor::AttachToGroup(Id const & id, GroupId const & group)
+{
+  auto const it = m_idToGroup.find(id);
+  if (it != m_idToGroup.end())
+    LOG(LWARNING, ("Tried to attach bookmark", id, "to group", group, "but it already belongs to group", it->second));
+
+  m_idToGroup[id] = group;
+  m_bookmarksInGroup[group].insert(id);
+  if (m_indexableGroups.count(group) > 0)
+    AddToIndex(id);
+}
+
+void Processor::DetachFromGroup(Id const & id, GroupId const & group)
+{
+  auto const it = m_idToGroup.find(id);
+  if (it == m_idToGroup.end())
+  {
+    LOG(LWARNING, ("Tried to detach bookmark", id, "from group", group, "but it does not belong to any group"));
+    return;
+  }
+
+  if (it->second != group)
+  {
+    LOG(LWARNING, ("Tried to detach bookmark", id, "from group", group, "but it only belongs to group", it->second));
+    return;
+  }
+
+  m_idToGroup.erase(it);
+  m_bookmarksInGroup[group].erase(id);
+
+  if (m_indexableGroups.count(group) > 0)
+    EraseFromIndex(id);
+
+  auto const groupIt = m_bookmarksInGroup.find(group);
+  CHECK(groupIt != m_bookmarksInGroup.end(), (group, m_bookmarksInGroup));
+  if (groupIt->second.size() == 0)
+    m_bookmarksInGroup.erase(groupIt);
+}
+
+void Processor::Search(Params const & params) const
+{
+  std::set<Id> ids;
+  auto insertId = [&ids](Id const & id, bool /* exactMatch */) { ids.insert(id); };
+
+  for (size_t i = 0; i < params.GetNumTokens(); ++i)
+  {
+    BailIfCancelled();
+
+    auto const & token = params.GetToken(i);
+    if (params.IsPrefixToken(i))
+      Retrieve<strings::PrefixDFAModifier<strings::LevenshteinDFA>>(token, insertId);
+    else
+      Retrieve<strings::LevenshteinDFA>(token, insertId);
+  }
+
+  IdfMap idfs(*this, 1.0 /* unknownIdf */);
+  auto qv = GetQueryVec(idfs, params);
+
+  std::vector<IdInfoPair> idInfos;
+  for (auto const & id : ids)
+  {
+    BailIfCancelled();
+
+    if (params.m_groupId != kInvalidGroupId)
+    {
+      auto const it = m_idToGroup.find(id);
+      if (it == m_idToGroup.end() || it->second != params.m_groupId)
+        continue;
+    }
+
+    auto it = m_docs.find(id);
+    CHECK(it != m_docs.end(), ("Can't find retrieved doc:", id));
+    auto const & doc = it->second;
+
+    RankingInfo info;
+    FillRankingInfo(qv, idfs, doc, info);
+
+    idInfos.emplace_back(id, info);
+  }
+
+  BailIfCancelled();
+  sort(idInfos.begin(), idInfos.end());
+
+  size_t numEmitted = 0;
+  for (auto const & idInfo : idInfos)
+  {
+    if (numEmitted >= params.m_maxNumResults)
+      break;
+    m_emitter.AddBookmarkResult(bookmarks::Result(idInfo.m_id));
+    ++numEmitted;
+  }
+}
+
+void Processor::Finish(bool cancelled)
+{
+  m_emitter.Finish(cancelled);
+}
+
+uint64_t Processor::GetNumDocs(strings::UniString const & token, bool isPrefix) const
+{
+  return base::asserted_cast<uint64_t>(m_index.GetNumDocs(StringUtf8Multilang::kDefaultCode, token, isPrefix));
+}
+
+QueryVec Processor::GetQueryVec(IdfMap & idfs, QueryParams const & params) const
+{
+  QueryVec::Builder builder;
+  for (size_t i = 0; i < params.GetNumTokens(); ++i)
+  {
+    auto const & token = params.GetToken(i).GetOriginal();
+    if (params.IsPrefixToken(i))
+      builder.SetPrefix(token);
+    else
+      builder.AddFull(token);
+  }
+  return {idfs, builder};
+}
+}  // namespace bookmarks
+}  // namespace search