Remove multimap usage from memoizeDeclsByName (p4lang#4821)

* Ensure we use iterator traits to query for iterators' value type. This way we can wrap over std::vector<T*> iterator Signed-off-by: Anton Korobeynikov <[email protected]> * Use better data structures to store memoized declarations Signed-off-by: Anton Korobeynikov <[email protected]> --------- Signed-off-by: Anton Korobeynikov <[email protected]>
rupesh-chiluka-marvell · Jul 24, 2024 · c0bbf1b · c0bbf1b
1 parent 132122d
commit c0bbf1b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 17 deletions.
diff --git a/frontends/common/resolveReferences/resolveReferences.cpp b/frontends/common/resolveReferences/resolveReferences.cpp
@@ -38,10 +38,10 @@ const std::vector<const IR::IDeclaration *> &ResolutionContext::memoizeDeclarati
     return (namespaceDecls[ns] = std::move(decls));
 }
 
-std::unordered_multimap<cstring, const IR::IDeclaration *> &ResolutionContext::memoizeDeclsByName(
+ResolutionContext::NamespaceDeclsByName &ResolutionContext::memoizeDeclsByName(
     const IR::INamespace *ns) const {
     auto &namesToDecls = namespaceDeclNames[ns];
-    for (const auto *d : getDeclarations(ns)) namesToDecls.emplace(d->getName().name, d);
+    for (const auto *d : getDeclarations(ns)) namesToDecls[d->getName().name].push_back(d);
     return namesToDecls;
 }
 
@@ -63,7 +63,7 @@ std::vector<const IR::IDeclaration *> ResolutionContext::lookup(const IR::INames
 
     if (const auto *gen = current->to<IR::IGeneralNamespace>()) {
         // FIXME: implement range filtering without enumerator wrappers
-        auto *decls = Util::enumerate(getDeclsByName(gen, name));
+        auto *decls = getDeclsByName(gen, name);
         switch (type) {
             case P4::ResolutionType::Any:
                 break;

diff --git a/frontends/common/resolveReferences/resolveReferences.h b/frontends/common/resolveReferences/resolveReferences.h
@@ -18,6 +18,7 @@ limitations under the License.
 #define COMMON_RESOLVEREFERENCES_RESOLVEREFERENCES_H_
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/inlined_vector.h"
 #include "ir/ir.h"
 #include "lib/cstring.h"
 #include "lib/iterator_range.h"
@@ -36,17 +37,17 @@ class ResolutionContext : virtual public Visitor, public DeclarationLookup {
     const std::vector<const IR::IDeclaration *> &memoizeDeclarations(
         const IR::INamespace *ns) const;
 
+    using DeclsVector = absl::InlinedVector<const IR::IDeclaration *, 2>;
+    using NamespaceDeclsByName = absl::flat_hash_map<cstring, DeclsVector, Util::Hash>;
+
     // Returns a mapping from name -> decl for the given namespace, and caches the result for
     // future lookups.
-    std::unordered_multimap<cstring, const IR::IDeclaration *> &memoizeDeclsByName(
-        const IR::INamespace *ns) const;
+    NamespaceDeclsByName &memoizeDeclsByName(const IR::INamespace *ns) const;
 
     mutable absl::flat_hash_map<const IR::INamespace *, std::vector<const IR::IDeclaration *>,
                                 Util::Hash>
         namespaceDecls;
-    mutable absl::flat_hash_map<const IR::INamespace *,
-                                std::unordered_multimap<cstring, const IR::IDeclaration *>,
-                                Util::Hash>
+    mutable absl::flat_hash_map<const IR::INamespace *, NamespaceDeclsByName, Util::Hash>
         namespaceDeclNames;
 
  protected:
@@ -100,11 +101,13 @@ class ResolutionContext : virtual public Visitor, public DeclarationLookup {
     /// Returns the set of decls with the given name that exist in the given namespace.
     auto getDeclsByName(const IR::INamespace *ns, cstring name) const {
         auto nsIt = namespaceDeclNames.find(ns);
-        auto &namesToDecls =
+        const auto &namesToDecls =
             nsIt != namespaceDeclNames.end() ? nsIt->second : memoizeDeclsByName(ns);
 
-        auto decls = Values(namesToDecls.equal_range(name));
-        return Util::iterator_range(decls.begin(), decls.end());
+        auto decls = namesToDecls.find(name);
+        if (decls == namesToDecls.end())
+            return Util::Enumerator<const IR::IDeclaration *>::emptyEnumerator();
+        return Util::enumerate(decls->second);
     }
 };
 

diff --git a/lib/enumerator.h b/lib/enumerator.h
@@ -196,17 +196,17 @@ class Enumerator {
 
 /// A generic iterator returning elements of type T.
 template <typename Iter>
-class IteratorEnumerator : public Enumerator<typename Iter::value_type> {
+class IteratorEnumerator : public Enumerator<typename std::iterator_traits<Iter>::value_type> {
  protected:
     Iter begin;
     Iter end;
     Iter current;
     const char *name;
-    friend class Enumerator<typename Iter::value_type>;
+    friend class Enumerator<typename std::iterator_traits<Iter>::value_type>;
 
  public:
     IteratorEnumerator(Iter begin, Iter end, const char *name)
-        : Enumerator<typename Iter::value_type>(),
+        : Enumerator<typename std::iterator_traits<Iter>::value_type>(),
           begin(begin),
           end(end),
           current(begin),
@@ -241,7 +241,7 @@ class IteratorEnumerator : public Enumerator<typename Iter::value_type> {
         throw std::runtime_error("Unexpected enumerator state");
     }
 
-    typename Iter::value_type getCurrent() const {
+    typename std::iterator_traits<Iter>::value_type getCurrent() const {
         switch (this->state) {
             case EnumeratorState::NotStarted:
                 throw std::logic_error("You cannot call 'getCurrent' before 'moveNext'");
@@ -638,12 +638,12 @@ bool EnumeratorHandle<T>::operator!=(const EnumeratorHandle<T> &other) const {
 }
 
 template <typename Iter>
-Enumerator<typename Iter::value_type> *enumerate(Iter begin, Iter end) {
+Enumerator<typename std::iterator_traits<Iter>::value_type> *enumerate(Iter begin, Iter end) {
     return new IteratorEnumerator(begin, end, "iterator");
 }
 
 template <typename Iter>
-Enumerator<typename Iter::value_type> *enumerate(iterator_range<Iter> range) {
+Enumerator<typename std::iterator_traits<Iter>::value_type> *enumerate(iterator_range<Iter> range) {
     return new IteratorEnumerator(range.begin(), range.end(), "range");
 }