Add a hashing function for state and symbolic variables.

Signed-off-by: fruffy <[email protected]>
p4lang · Aug 1, 2024 · 252c182 · 252c182
1 parent a7d6e35
commit 252c182
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 6 deletions.
diff --git a/backends/p4tools/common/lib/ir_compare.h b/backends/p4tools/common/lib/ir_compare.h
@@ -25,5 +25,16 @@ struct StateVariableLess {
     }
 };
 
+/// Hash for StateVariable pointers. We only hash the label.
+struct StateVariableHash {
+    size_t operator()(const IR::StateVariable *s1) const {
+        return Util::Hasher<uint64_t>()(s1->hash());
+    }
+
+    size_t operator()(const IR::StateVariable &s1) const {
+        return Util::Hasher<uint64_t>()(s1.hash());
+    }
+};
+
 }  // namespace IR
 #endif /* BACKENDS_P4TOOLS_COMMON_LIB_IR_COMPARE_H_ */
diff --git a/backends/p4tools/common/lib/model.h b/backends/p4tools/common/lib/model.h
@@ -4,15 +4,16 @@
 #include <map>
 #include <utility>
 
-#include "absl/container/btree_map.h"
+#include "absl/container/flat_hash_map.h"
 #include "ir/ir.h"
 #include "ir/solver.h"
 #include "ir/visitor.h"
 
 namespace P4Tools {
 
 /// Symbolic maps map a state variable to a IR::Expression.
-using SymbolicMapType = absl::btree_map<IR::StateVariable, const IR::Expression *>;
+using SymbolicMapType = absl::flat_hash_map<IR::StateVariable, const IR::Expression *,
+                                            IR::StateVariableHash, IR::StateVariableEqual>;
 
 /// Represents a solution found by the solver. A model is a concretized form of a symbolic
 /// environment. All the expressions in a Model must be of type IR::Literal.

diff --git a/backends/p4tools/p4tools.def b/backends/p4tools/p4tools.def
@@ -126,6 +126,45 @@ class StateVariable : Expression {
         return 0;
     }
 
+    uint64_t hash() const {
+        return hash(0, ref);
+    }
+
+    uint64_t hash(uint64_t seed, const IR::Expression *expression) const {
+        // expression is a Member.
+        if (const auto *member = expression->to<IR::Member>()) {
+            return hash(seed, member);
+        }
+        // expression is a PathExpression.
+        if (const auto *pathExpression = expression->to<IR::PathExpression>()) {
+            return hash(seed, pathExpression);
+        }
+        // expression is a ArrayIndex.
+        if (const auto *arrayIndex = expression->to<IR::ArrayIndex>()) {
+            return hash(seed, arrayIndex);
+        }
+        BUG("Either %1% of type %2% is not a valid StateVariable", expression,
+            expression->node_type_name());
+    }
+
+    uint64_t hash(uint64_t seed, const IR::Member *member) const {
+        return hash(Util::hash_combine(seed, std::hash<cstring>()(member->member)), member->expr);
+    }
+
+    uint64_t hash(uint64_t seed, const IR::PathExpression *pathExpression) const {
+        return Util::hash_combine(seed, std::hash<cstring>()(pathExpression->path->name));
+    }
+
+    uint64_t hash(uint64_t seed, const IR::ArrayIndex *arrayIndex) const {
+        const auto *arrayIndexVal = arrayIndex->right->to<IR::Constant>();
+        BUG_CHECK(
+            arrayIndexVal != nullptr,
+            "Value %1% is not a constant. Only constants are supported as part of a state variable.",
+            arrayIndex->right);
+        return hash(Util::hash_combine(seed, std::hash<big_int>()(arrayIndexVal->value)),
+                                 arrayIndex->left);
+    }
+
     toString { return ref->toString(); }
 
     dbprint { ref->dbprint(out); }

diff --git a/ir/compare.h b/ir/compare.h
@@ -25,5 +25,16 @@ struct SymbolicVariableLess {
     }
 };
 
+/// Hash for SymbolicVariable pointers. We only hash the label.
+struct SymbolicVariableHash {
+    size_t operator()(const IR::SymbolicVariable *s1) const {
+        return std::hash<cstring>()(s1->label);
+    }
+
+    size_t operator()(const IR::SymbolicVariable &s1) const {
+        return std::hash<cstring>()(s1.label);
+    }
+};
+
 }  // namespace IR
 #endif /* IR_COMPARE_H_ */
diff --git a/ir/solver.h b/ir/solver.h
@@ -1,11 +1,12 @@
 #ifndef IR_SOLVER_H_
 #define IR_SOLVER_H_
 
+#include <functional>
 #include <optional>
 #include <vector>
 
-#include "absl/container/btree_map.h"
-#include "ir/compare.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "ir/ir.h"
 #include "lib/castable.h"
 #include "lib/cstring.h"
@@ -15,8 +16,10 @@
 using Constraint = IR::Expression;
 
 /// This type maps symbolic variables to their value assigned by the solver.
-using SymbolicMapping =
-    absl::btree_map<const IR::SymbolicVariable *, const IR::Expression *, IR::SymbolicVariableLess>;
+using SymbolicMapping = absl::flat_hash_map<const IR::SymbolicVariable *, const IR::Expression *,
+                                            IR::SymbolicVariableHash, IR::SymbolicVariableEqual>;
+using SymbolicSet = absl::flat_hash_set<const IR::SymbolicVariable *, IR::SymbolicVariableHash,
+                                        IR::SymbolicVariableEqual>;
 
 /// Provides a higher-level interface for an SMT solver.
 class AbstractSolver : public ICastable {