From 2e896be40dc638cbb85601b8126e6e9f49be3e1c Mon Sep 17 00:00:00 2001
From: Alex Light <allight@google.com>
Date: Wed, 23 Oct 2024 15:45:36 -0700
Subject: [PATCH] Add Specialize function to proc-state range query engine.

This is used by some potential users and can be implemented better than the default do-nothing implementation.

Also fix a small typo.

PiperOrigin-RevId: 689129685
---
 xls/passes/BUILD                           | 1 +
 xls/passes/proc_state_range_query_engine.h | 9 ++++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/xls/passes/BUILD b/xls/passes/BUILD
index bc56d6a339..8e19a6bcd2 100644
--- a/xls/passes/BUILD
+++ b/xls/passes/BUILD
@@ -3420,6 +3420,7 @@ cc_library(
         ":back_propagate_range_analysis",
         ":dataflow_visitor",
         ":node_dependency_analysis",
+        ":predicate_state",
         ":query_engine",
         ":range_query_engine",
         ":ternary_query_engine",
diff --git a/xls/passes/proc_state_range_query_engine.h b/xls/passes/proc_state_range_query_engine.h
index ab7f8f046d..81091b2e12 100644
--- a/xls/passes/proc_state_range_query_engine.h
+++ b/xls/passes/proc_state_range_query_engine.h
@@ -19,6 +19,7 @@
 #include <optional>
 #include <utility>
 
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
@@ -30,6 +31,7 @@
 #include "xls/ir/node.h"
 #include "xls/ir/proc.h"
 #include "xls/ir/ternary.h"
+#include "xls/passes/predicate_state.h"
 #include "xls/passes/query_engine.h"
 #include "xls/passes/range_query_engine.h"
 #include "xls/passes/ternary_query_engine.h"
@@ -81,7 +83,7 @@ class ProcStateRangeQueryEngine final : public QueryEngine {
   }
 
   bool AtLeastOneTrue(absl::Span<TreeBitLocation const> bits) const override {
-    return AtLeastOneTrue(bits);
+    return inner_.AtLeastOneTrue(bits);
   }
 
   bool KnownEquals(const TreeBitLocation& a,
@@ -136,6 +138,11 @@ class ProcStateRangeQueryEngine final : public QueryEngine {
     return inner_.GetTernary(node);
   }
 
+  std::unique_ptr<QueryEngine> SpecializeGivenPredicate(
+      const absl::flat_hash_set<PredicateState>& state) const override {
+    return inner_.SpecializeGivenPredicate(state);
+  }
+
  private:
   // Actual range results from the proc-state aware analysis.
   std::unique_ptr<TernaryQueryEngine> ternary_;