diff --git a/frontends/p4/def_use.cpp b/frontends/p4/def_use.cpp
index 06b4e52b59b..a05414c4829 100644
--- a/frontends/p4/def_use.cpp
+++ b/frontends/p4/def_use.cpp
@@ -763,10 +763,16 @@ void ComputeWriteSet::visitVirtualMethods(const IR::IndexedVector<IR::Declaratio
     }
 }
 
+std::size_t P4::loc_t::hash() const {
+    if (!parent) return Util::Hash{}(node->id);
+
+    return Util::Hash{}(node->id, parent->hash());
+}
+
 // Returns program location of n, given the program location of n's direct parent.
 // Use to get loc if n is indirect child (e.g. grandchild) of currently being visited node.
 // In this case parentLoc is the loc of n's direct parent.
-const P4::ComputeWriteSet::loc_t *ComputeWriteSet::getLoc(const IR::Node *n,
+const P4::loc_t *ComputeWriteSet::getLoc(const IR::Node *n,
                                                           const loc_t *parentLoc) {
     loc_t tmp{n, parentLoc};
     return &*cached_locs.insert(tmp).first;
@@ -774,7 +780,7 @@ const P4::ComputeWriteSet::loc_t *ComputeWriteSet::getLoc(const IR::Node *n,
 
 // Returns program location given the context of the currently being visited node.
 // Use to get loc of currently being visited node.
-const P4::ComputeWriteSet::loc_t *ComputeWriteSet::getLoc(const Visitor::Context *ctxt) {
+const P4::loc_t *ComputeWriteSet::getLoc(const Visitor::Context *ctxt) {
     if (!ctxt) return nullptr;
     loc_t tmp{ctxt->node, getLoc(ctxt->parent)};
     return &*cached_locs.insert(tmp).first;
@@ -783,7 +789,7 @@ const P4::ComputeWriteSet::loc_t *ComputeWriteSet::getLoc(const Visitor::Context
 // Returns program location of a child node n, given the context of the
 // currently being visited node.
 // Use to get loc if n is direct child of currently being visited node.
-const P4::ComputeWriteSet::loc_t *ComputeWriteSet::getLoc(const IR::Node *n,
+const P4::loc_t *ComputeWriteSet::getLoc(const IR::Node *n,
                                                           const Visitor::Context *ctxt) {
     for (auto *p = ctxt; p; p = p->parent)
         if (p->node == n) return getLoc(p);
diff --git a/frontends/p4/def_use.h b/frontends/p4/def_use.h
index fb80923a932..fd06b35dacc 100644
--- a/frontends/p4/def_use.h
+++ b/frontends/p4/def_use.h
@@ -24,15 +24,29 @@ limitations under the License.
 #include "lib/alloc_trace.h"
 #include "lib/hash.h"
 #include "lib/hvec_map.h"
-#include "lib/ordered_map.h"
 #include "lib/ordered_set.h"
 #include "typeMap.h"
 
 namespace P4 {
 
+class ComputeWriteSet;
 class StorageFactory;
 class LocationSet;
 
+// A location in the program. Includes the context from the visitor, which needs to
+// be copied out of the Visitor::Context objects, as they are allocated on the stack and
+// will become invalid as the IR traversal continues
+struct loc_t {
+    const IR::Node *node;
+    const loc_t *parent;
+    bool operator==(const loc_t &a) const {
+        if (node != a.node) return false;
+        if (parent && a.parent) return *parent == *a.parent;
+        return parent == a.parent;
+    }
+    std::size_t hash() const;
+};
+
 /// Abstraction for something that is has a left value (variable, parameter)
 class StorageLocation : public IHasDbPrint, public ICastable {
     static unsigned crtid;
@@ -336,6 +350,14 @@ struct hash<P4::ProgramPoint> {
     typedef std::size_t result_type;
     result_type operator()(argument_type const &s) const { return s.hash(); }
 };
+
+template <>
+struct hash<P4::loc_t> {
+    typedef P4::loc_t argument_type;
+    typedef std::size_t result_type;
+    result_type operator()(argument_type const &loc) const { return loc.hash(); }
+};
+
 }  // namespace std
 
 namespace Util {
@@ -478,18 +500,6 @@ class AllDefinitions : public IHasDbPrint {
 
 class ComputeWriteSet : public Inspector, public IHasDbPrint {
  public:
-    // A location in the program. Includes the context from the visitor, which needs to
-    // be copied out of the Visitor::Context objects, as they are allocated on the stack and
-    // will become invalid as the IR traversal continues
-    struct loc_t {
-        const IR::Node *node;
-        const loc_t *parent;
-        bool operator<(const loc_t &a) const {
-            if (node != a.node) return node->id < a.node->id;
-            if (!parent || !a.parent) return parent != nullptr;
-            return *parent < *a.parent;
-        }
-    };
 
     explicit ComputeWriteSet(AllDefinitions *allDefinitions)
         : allDefinitions(allDefinitions),
@@ -499,7 +509,7 @@ class ComputeWriteSet : public Inspector, public IHasDbPrint {
           storageMap(allDefinitions->storageMap),
           lhs(false),
           virtualMethod(false),
-          cached_locs(*new std::set<loc_t>) {
+          cached_locs(*new std::unordered_set<loc_t>) {
         CHECK_NULL(allDefinitions);
         visitDagOnce = false;
     }
@@ -561,7 +571,7 @@ class ComputeWriteSet : public Inspector, public IHasDbPrint {
     /// if true we are processing an expression on the lhs of an assignment
     bool lhs;
     /// For each program location the location set it writes
-    ordered_map<loc_t, const LocationSet *> writes;
+    hvec_map<loc_t, const LocationSet *> writes;
     bool virtualMethod;  /// True if we are analyzing a virtual method
     AllocTrace memuse;
     alloc_trace_cb_t nested_trace;
@@ -570,7 +580,7 @@ class ComputeWriteSet : public Inspector, public IHasDbPrint {
     /// Creates new visitor, but with same underlying data structures.
     /// Needed to visit some program fragments repeatedly.
     ComputeWriteSet(const ComputeWriteSet *source, ProgramPoint context, Definitions *definitions,
-                    std::set<loc_t> &cached_locs)
+                    std::unordered_set<loc_t> &cached_locs)
         : allDefinitions(source->allDefinitions),
           currentDefinitions(definitions),
           returnedDefinitions(nullptr),
@@ -648,7 +658,8 @@ class ComputeWriteSet : public Inspector, public IHasDbPrint {
     }
 
  private:
-    std::set<loc_t> &cached_locs;
+    // TODO: Make absl::flat_hash_set instead?
+    std::unordered_set<loc_t> &cached_locs;
 };
 
 }  // namespace P4