Ensure next-value-opt handles of non-synth state elements

docs_src/passes_list.md

@@ -689,7 +689,7 @@ A pass which reduces the width of operations eliminating redundant or unused bit
 
 Pass which tries to optimize `next_value` nodes.
 
- Optimizations include: - removing literal predicates on `next_value` nodes (removing the   `next_value` node if dead), - splitting `next_value` nodes with `select`-based values (if small), - splitting `next_value` nodes with `priority_sel`-based values, and - splitting `next_value` nodes with `one_hot_sel`-based values (where safe).
+ Optimizations include: - removing literal predicates on `next_value` nodes (removing the   `next_value` node if dead), - splitting `next_value` nodes with `select`-based values (if small), - splitting `next_value` nodes with `priority_sel`-based values, and - splitting `next_value` nodes with `one_hot_sel`-based values (where safe). - Identifies `state-element`s which are non-synth and performing an identity   update operation and removes the read of the synth version in this case.   (This allows for later cond-spec passes to predicate the read.)
 
  For best results, first modernizes old-style values on `next (...)` lines, converting them to `next_value` nodes.
 

xls/ir/nodes.h

@@ -857,6 +857,10 @@ class Next final : public Node {
 
   bool IsDefinitelyEqualTo(const Node* other) const final;
 
+  StateElement* state_element() const {
+    return state_read()->As<StateRead>()->state_element();
+  }
+
  private:
   static constexpr int64_t kPredicateOperand = 2;
 

xls/passes/BUILD

@@ -4003,35 +4003,40 @@ xls_pass(
         ":pass_base",
         ":query_engine",
         ":stateless_query_engine",
-        "//xls/common:math_util",
+        "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/ir",
         "//xls/ir:bits",
         "//xls/ir:node_util",
         "//xls/ir:op",
+        "//xls/ir:state_element",
         "//xls/ir:value",
-        "//xls/ir:value_utils",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
+        "@cppitertools",
     ],
 )
 
 cc_test(
     name = "next_value_optimization_pass_test",
     srcs = ["next_value_optimization_pass_test.cc"],
     deps = [
+        ":dce_pass",
         ":next_value_optimization_pass",
         ":optimization_pass",
         ":pass_base",
+        ":pass_test_helpers",
         "//xls/common:xls_gunit_main",
         "//xls/common/fuzzing:fuzztest",
         "//xls/common/status:matchers",
+        "//xls/common/status:status_macros",
         "//xls/fuzzer/ir_fuzzer:ir_fuzz_domain",
         "//xls/fuzzer/ir_fuzzer:ir_fuzz_test_library",
         "//xls/ir",

xls/passes/dataflow_visitor.h

@@ -385,18 +385,13 @@ class DataflowVisitor : public DfsVisitorWithDefault {
 
   // Sets the leaf type tree value associated with `node`.
   absl::Status SetValue(Node* node, LeafTypeTreeView<LeafT> value) {
-    XLS_RET_CHECK_EQ(node->GetType(), value.type());
-    map_.insert_or_assign(
-        node, std::make_unique<SharedLeafTypeTree<LeafT>>(value.AsShared()));
-    return absl::OkStatus();
+    return SetValue(node, value.AsShared());
   }
   absl::Status SetValue(Node* node, LeafTypeTree<LeafT>&& value) {
-    XLS_RET_CHECK_EQ(node->GetType(), value.type());
-    map_.insert_or_assign(node, std::make_unique<SharedLeafTypeTree<LeafT>>(
-                                    std::move(value).AsShared()));
-    return absl::OkStatus();
+    return SetValue(node, std::move(value).AsShared());
   }
-  absl::Status SetValue(Node* node, SharedLeafTypeTree<LeafT>&& value) {
+  // Base set-value function. Any overrides must call this.
+  virtual absl::Status SetValue(Node* node, SharedLeafTypeTree<LeafT>&& value) {
     XLS_RET_CHECK_EQ(node->GetType(), value.type());
     map_.insert_or_assign(
         node, std::make_unique<SharedLeafTypeTree<LeafT>>(std::move(value)));