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

@@ -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",

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_;