Add a ReserveNodes method to DfsVisitor.

xls/ir/dfs_visitor.h

@@ -30,6 +30,12 @@ class DfsVisitor {
  public:
   virtual ~DfsVisitor() = default;
 
+  // Let the visitor know that 'count' nodes are likely to be visited in the
+  // near future. This is purely advisory to allow the visitor to reserve space
+  // to avoid reallocations. The visitor is free to ignore this hint. The user
+  // is free to call this method or not.
+  virtual void ReserveNodes(int64_t count) { visited_.reserve(count); }
+
   virtual absl::Status HandleAdd(BinOp* add) = 0;
   virtual absl::Status HandleAfterAll(AfterAll* after_all) = 0;
   virtual absl::Status HandleMinDelay(MinDelay* min_delay) = 0;

xls/passes/dataflow_simplification_pass.cc

@@ -86,6 +86,7 @@ absl::StatusOr<bool> DataflowSimplificationPass::RunOnFunctionBaseInternal(
     FunctionBase* func, const OptimizationPassOptions& options,
     PassResults* results, OptimizationContext& context) const {
   NodeSourceDataflowVisitor visitor;
+  visitor.ReserveNodes(func->node_count());
   for (Node* node : context.TopoSort(func)) {
     XLS_RETURN_IF_ERROR(node->VisitSingleNode(&visitor));
   }

xls/passes/dataflow_visitor.h

@@ -84,6 +84,10 @@ namespace xls {
 template <typename LeafT>
 class DataflowVisitor : public DfsVisitorWithDefault {
  public:
+  void ReserveNodes(int64_t count) override {
+    DfsVisitorWithDefault::ReserveNodes(count);
+    map_.reserve(count);
+  }
   absl::Status HandleArray(Array* array) override {
     std::vector<LeafTypeTreeView<LeafT>> elements;
     for (Node* operand : array->operands()) {