Fix for timeouts

PiperOrigin-RevId: 863274693

Author: Sean Purser-Haskell

xls/contrib/xlscc/continuations.cc

@@ -278,6 +278,15 @@ OptimizationContext::GetSourcesSetTreeNodeInfoForFunction(
   return sources_set_tree_node_infos_by_function_.at(in_function).get();
 }
 
+const xls::LeafTypeTree<std::monostate>&
+OptimizationContext::GetBlankTypeTreeForType(xls::Type* type) {
+  CHECK_NE(type, nullptr);
+  auto [param_tree_it, _] = param_tree_cache_.try_emplace(type, type);
+  xls::LeafTypeTree<std::monostate>& found = param_tree_it->second;
+  CHECK(found.type()->IsEqualTo(type));
+  return found;
+}
+
 absl::StatusOr<bool> OptimizationContext::CheckNodeSourcesInSet(
     xls::FunctionBase* in_function, xls::Node* node,
     absl::flat_hash_set<const xls::Param*> sources_set,
@@ -1270,12 +1279,10 @@ absl::Status RemoveUnusedContinuationInputs(GeneratedFunction& func,
       }
 
       CHECK_EQ(continuation_in.input_node->function_base(), slice.function);
-
       if (return_value_from_params.contains(continuation_in.input_node)) {
         ++cont_in_it;
         continue;
       }
-
       // There may still be uses like forming a tuple, the element of which is
       // never indexed
       XLS_RETURN_IF_ERROR(
@@ -1349,8 +1356,8 @@ FindPassThroughs(GeneratedFunction& func, OptimizationContext& context) {
 
               // Check that param has the same number of elements as
               // continuation_out. This avoids marking slices as pass-throughs.
-              xls::LeafTypeTree<std::monostate> param_tree(
-                  source_node->GetType());
+              const xls::LeafTypeTree<std::monostate>& param_tree =
+                  context.GetBlankTypeTreeForType(source_node->GetType());
               if (param_tree.elements().size() != sources.elements().size()) {
                 disallowed = true;
                 break;

xls/contrib/xlscc/translator_types.h

@@ -1302,6 +1302,9 @@ class OptimizationContext {
   absl::StatusOr<xls::PartialInfoQueryEngine*> GetQueryEngineForFunction(
       xls::FunctionBase* in_function);
 
+  const xls::LeafTypeTree<std::monostate>& GetBlankTypeTreeForType(
+      xls::Type* type);
+
   absl::StatusOr<bool> CheckNodeSourcesInSet(
       xls::FunctionBase* in_function, xls::Node* node,
       absl::flat_hash_set<const xls::Param*> sources_set,
@@ -1315,6 +1318,9 @@ class OptimizationContext {
                       std::unique_ptr<SourcesSetTreeNodeInfo>>
       sources_set_tree_node_infos_by_function_;
 
+  absl::flat_hash_map<xls::Type*, xls::LeafTypeTree<std::monostate>>
+      param_tree_cache_;
+
   absl::flat_hash_map<xls::FunctionBase*,
                       std::unique_ptr<xls::PartialInfoQueryEngine>>
       query_engines_by_function_;