metrics have been added

1 files changed, 111 insertions, 0 deletions

diff --git a/library/cpp/tvmauth/client/misc/roles/entities_index.cpp b/library/cpp/tvmauth/client/misc/roles/entities_index.cpp
new file mode 100644
index 0000000000..6c3cd5c192
--- /dev/null
+++ b/library/cpp/tvmauth/client/misc/roles/entities_index.cpp
@@ -0,0 +1,111 @@
+#include "entities_index.h"
+
+#include <util/stream/str.h>
+
+#include <set>
+
+namespace NTvmAuth::NRoles {
+    TEntitiesIndex::TStage::TStage(const std::set<TString>& k)
+        : Keys_(k.begin(), k.end())
+    {
+    }
+
+    // TODO TStringBuf
+    bool TEntitiesIndex::TStage::GetNextKeySet(std::vector<TString>& out) {
+        out.clear();
+        out.reserve(Keys_.size());
+
+        ++Id_;
+        for (size_t idx = 0; idx < Keys_.size(); ++idx) {
+            bool need = (Id_ >> idx) & 0x01;
+
+            if (need) {
+                out.push_back(Keys_[idx]);
+            }
+        }
+
+        return !out.empty();
+    }
+
+    TEntitiesIndex::TEntitiesIndex(const std::vector<TEntityPtr>& entities) {
+        const std::set<TString> uniqueKeys = GetUniqueSortedKeys(entities);
+        Idx_.Entities = entities;
+        Idx_.SubTree.reserve(uniqueKeys.size() * entities.size());
+
+        TStage stage(uniqueKeys);
+        std::vector<TString> keyset;
+        while (stage.GetNextKeySet(keyset)) {
+            for (const TEntityPtr& e : entities) {
+                TSubTree* currentBranch = &Idx_;
+
+                for (const TString& key : keyset) {
+                    auto it = e->find(key);
+                    if (it == e->end()) {
+                        continue;
+                    }
+
+                    auto [i, ok] = currentBranch->SubTree.emplace(
+                        TKeyValue{it->first, it->second},
+                        TSubTree());
+
+                    currentBranch = &i->second;
+                    currentBranch->Entities.push_back(e);
+                }
+            }
+        }
+
+        MakeUnique(Idx_);
+    }
+
+    std::set<TString> TEntitiesIndex::GetUniqueSortedKeys(const std::vector<TEntityPtr>& entities) {
+        std::set<TString> res;
+
+        for (const TEntityPtr& e : entities) {
+            for (const auto& [key, value] : *e) {
+                res.insert(key);
+            }
+        }
+
+        return res;
+    }
+
+    void TEntitiesIndex::MakeUnique(TSubTree& branch) {
+        auto& vec = branch.Entities;
+        std::sort(vec.begin(), vec.end());
+        vec.erase(std::unique(vec.begin(), vec.end()), vec.end());
+
+        for (auto& [_, restPart] : branch.SubTree) {
+            MakeUnique(restPart);
+        }
+    }
+
+    static void Print(const TEntitiesIndex::TSubTree& part, IOutputStream& out, size_t offset = 0) {
+        std::vector<std::pair<TKeyValue, const TEntitiesIndex::TSubTree*>> vec;
+        vec.reserve(part.SubTree.size());
+
+        for (const auto& [key, value] : part.SubTree) {
+            vec.push_back({key, &value});
+        }
+
+        std::sort(vec.begin(), vec.end(), [](const auto& l, const auto& r) {
+            if (l.first.Key != r.first.Key) {
+                return l.first.Key < r.first.Key;
+            }
+            return l.first.Value < r.first.Value;
+        });
+
+        for (const auto& [key, value] : vec) {
+            out << TString(offset, ' ') << "\"" << key.Key << "/" << key.Value << "\"" << Endl;
+            Print(*value, out, offset + 4);
+        }
+    }
+
+    TString TEntitiesIndex::PrintDebugString() const {
+        TStringStream res;
+        res << Endl;
+
+        Print(Idx_, res);
+
+        return res.Str();
+    }
+}