[RF] Fix quadratic scaling when building large workspaces

roofit/roofitcore/src/RooAbsCollection.cxx

@@ -48,7 +48,7 @@ implemented using the container denoted by RooAbsCollection::Storage_t.
 #include <iostream>
 #include <fstream>
 #include <memory>
-
+#include <unordered_set>
 
 namespace RooFit {
 namespace Detail {
@@ -1559,28 +1559,87 @@ void RooAbsCollection::sort(bool reverse) {
 /// that RooAbsArg in the collection.
 
 void RooAbsCollection::sortTopologically() {
-   std::unordered_set<TNamed const *> seenArgs;
-   for (std::size_t iArg = 0; iArg < _list.size(); ++iArg) {
-      RooAbsArg *arg = _list[iArg];
-      bool movedArg = false;
-      for (RooAbsArg *server : arg->servers()) {
-         if (seenArgs.find(server->namePtr()) == seenArgs.end()) {
-            auto found = std::find_if(_list.begin(), _list.end(),
-                                      [server](RooAbsArg *elem) { return elem->namePtr() == server->namePtr(); });
-            if (found != _list.end()) {
-               _list.erase(found);
-               _list.insert(_list.begin() + iArg, server);
-               movedArg = true;
-               break;
+   // Depth-first post-order traversal over the server edges, seeded with the
+   // args in collection order. Like in the original insertion-based
+   // implementation, servers are matched to collection members by name, with
+   // an encountered server instance standing in for a same-name member.
+   // Every arg and server link is visited only once.
+   std::unordered_set<TNamed const *> memberNames;
+   memberNames.reserve(_list.size());
+   bool hasDuplicates = false;
+   for (RooAbsArg *arg : _list) {
+      hasDuplicates |= !memberNames.insert(arg->namePtr()).second;
+   }
+
+   // The traversal emits each name once, which would shrink a collection
+   // with multiple same-name entries (allowed in a RooArgList). Fall back to
+   // the original algorithm in that case.
+   if (hasDuplicates) {
+      std::unordered_set<TNamed const *> seenArgsLegacy;
+      for (std::size_t iArg = 0; iArg < _list.size(); ++iArg) {
+         RooAbsArg *arg = _list[iArg];
+         bool movedArg = false;
+         for (RooAbsArg *server : arg->servers()) {
+            if (seenArgsLegacy.find(server->namePtr()) == seenArgsLegacy.end()) {
+               auto found = std::find_if(_list.begin(), _list.end(),
+                                         [server](RooAbsArg *elem) { return elem->namePtr() == server->namePtr(); });
+               if (found != _list.end()) {
+                  _list.erase(found);
+                  _list.insert(_list.begin() + iArg, server);
+                  movedArg = true;
+                  break;
+               }
             }
          }
+         if (movedArg) {
+            --iArg;
+            continue;
+         }
+         seenArgsLegacy.insert(arg->namePtr());
       }
-      if (movedArg) {
-         --iArg;
+      return;
+   }
+
+   Storage_t sorted;
+   sorted.reserve(_list.size());
+   std::unordered_set<TNamed const *> seenArgs;
+   seenArgs.reserve(_list.size());
+   // Names on the current traversal path, to skip cyclic server edges
+   // instead of looping forever.
+   std::unordered_set<TNamed const *> onStack;
+
+   // Pairs of the arg to process and the index of the next of its servers to look at
+   std::vector<std::pair<RooAbsArg *, std::size_t>> stack;
+   for (RooAbsArg *arg : _list) {
+      if (seenArgs.count(arg->namePtr()))
          continue;
+      stack.emplace_back(arg, 0);
+      onStack.insert(arg->namePtr());
+      while (!stack.empty()) {
+         RooAbsArg *node = stack.back().first;
+         bool descended = false;
+         auto const &servers = node->servers();
+         while (stack.back().second < servers.size()) {
+            RooAbsArg *server = servers[stack.back().second++];
+            TNamed const *serverName = server->namePtr();
+            if (memberNames.count(serverName) && !seenArgs.count(serverName) && !onStack.count(serverName)) {
+               stack.emplace_back(server, 0);
+               onStack.insert(serverName);
+               descended = true;
+               break;
+            }
+         }
+         if (descended)
+            continue;
+         if (seenArgs.insert(node->namePtr()).second) {
+            sorted.push_back(node);
+         }
+         onStack.erase(node->namePtr());
+         stack.pop_back();
       }
-      seenArgs.insert(arg->namePtr());
    }
+
+   _list = std::move(sorted);
 }
 
 ////////////////////////////////////////////////////////////////////////////////

roofit/roofitcore/src/RooRealIntegral.cxx

@@ -46,7 +46,7 @@ integration is performed in the various implementations of the RooAbsIntegrator
 
 #include <iostream>
 #include <memory>
-
+#include <unordered_map>
 
 namespace {
 
@@ -131,14 +131,26 @@ void addParameterToServers(RooAbsReal const &function, RooAbsArg &leaf, std::vec
 enum class MarkedState { Dependent, Independent, AlreadyAdded };
 
 /// Mark all args that recursively are value clients of "dep".
-void unmarkDepValueClients(RooAbsArg const &dep, RooArgSet const &args, std::vector<MarkedState> &marked)
+void unmarkDepValueClients(RooAbsArg const &dep, std::unordered_map<RooAbsArg const *, std::size_t> const &indexMap,
+                           std::vector<MarkedState> &marked)
 {
-   assert(args.size() == marked.size());
-   auto index = args.index(dep);
-   if (index >= 0) {
-      marked[index] = MarkedState::Dependent;
-      for (RooAbsArg *client : dep.valueClients()) {
-         unmarkDepValueClients(*client, args, marked);
+   auto found = indexMap.find(&dep);
+   if (found == indexMap.end())
+      return;
+   marked[found->second] = MarkedState::Dependent;
+
+   // Iterative depth-first traversal of the value clients within the
+   // computation graph, visiting every arg at most once.
+   std::vector<RooAbsArg const *> stack{&dep};
+   while (!stack.empty()) {
+      RooAbsArg const *arg = stack.back();
+      stack.pop_back();
+      for (RooAbsArg *client : arg->valueClients()) {
+         auto foundClient = indexMap.find(client);
+         if (foundClient != indexMap.end() && marked[foundClient->second] != MarkedState::Dependent) {
+            marked[foundClient->second] = MarkedState::Dependent;
+            stack.push_back(client);
+         }
       }
    }
 }
@@ -154,22 +166,38 @@ getValueAndShapeServers(RooAbsReal const &function, RooArgSet const &depList, co
    RooArgSet allArgs{allArgsList};
    allArgs.sortTopologically();
 
+   // Maps from arg to index in allArgs for constant-time lookups. Two maps,
+   // because the graph can contain same-name instances (e.g. cloned sub-trees
+   // from projections): dependent client chains are matched by instance,
+   // while the final server matching is done by name.
+   std::unordered_map<RooAbsArg const *, std::size_t> indexMap;
+   std::unordered_map<TNamed const *, std::size_t> indexMapByName;
+   indexMap.reserve(allArgs.size());
+   indexMapByName.reserve(allArgs.size());
+   for (std::size_t i = 0; i < allArgs.size(); ++i) {
+      indexMap.emplace(allArgs[i], i);
+      indexMapByName.emplace(allArgs[i]->namePtr(), i);
+   }
+
    // Figure out what are all the value servers only
    RooArgList allValueArgsList;
    function.treeNodeServerList(&allValueArgsList, nullptr, true, true, /*valueOnly=*/true, false);
    RooArgSet allValueArgs{allValueArgsList};
 
    // All "marked" args will be added as value servers to the integral
    std::vector<MarkedState> marked(allArgs.size(), MarkedState::Independent);
-   marked.back() = MarkedState::Dependent; // We don't want to consider the function itself
+   // We don't want to consider the function itself
+   if (auto foundFunc = indexMap.find(&function); foundFunc != indexMap.end()) {
+      marked[foundFunc->second] = MarkedState::Dependent;
+   }
 
    // Mark all args that are (indirect) value servers of the integration
    // variable or the integration variable itself. If something was marked,
    // it means the integration variable was in the compute graph and we will
    // add it to the server list.
    for (RooAbsArg *dep : depList) {
       if (RooAbsArg *depInArgs = allArgs.find(dep->GetName())) {
-         unmarkDepValueClients(*depInArgs, allArgs, marked);
+         unmarkDepValueClients(*depInArgs, indexMap, marked);
          addObservableToServers(function, *depInArgs, serversToAdd, rangeName);
       }
    }
@@ -179,10 +207,10 @@ getValueAndShapeServers(RooAbsReal const &function, RooArgSet const &depList, co
    for (std::size_t i = 0; i < allArgs.size(); ++i) {
       if (marked[i] == MarkedState::Dependent) {
          for (RooAbsArg *server : allArgs[i]->servers()) {
-            int index = allArgs.index(server->GetName());
-            if (index >= 0 && marked[index] == MarkedState::Independent) {
+            auto found = indexMapByName.find(server->namePtr());
+            if (found != indexMapByName.end() && marked[found->second] == MarkedState::Independent) {
                addParameterToServers(function, *server, serversToAdd, !allValueArgs.find(*server));
-               marked[index] = MarkedState::AlreadyAdded;
+               marked[found->second] = MarkedState::AlreadyAdded;
             }
          }
       }

roofit/roofitcore/src/RooWorkspace.cxx

@@ -88,6 +88,7 @@ and try reading again.
 #include <fstream>
 #include <cstring>
 #include <unordered_map>
+#include <unordered_set>
 
 namespace {
 
@@ -489,20 +490,32 @@ bool RooWorkspace::import(const RooAbsArg& inArg,
     }
   }
 
-  // Make list of conflicting nodes
+  // When existing nodes are recycled and no renaming is requested, nodes
+  // whose names already exist in the workspace are not cloned at all: the
+  // existing nodes are used anyway, and the clones would be discarded right
+  // after the import. One deliberate difference to the general code path:
+  // importWorkspaceHook() and the class-code import don't run for recycled
+  // nodes anymore, since there is no clone to run them on.
+  const bool pruneClones = useExistingNodes && !suffix && varMap.empty() && !noRecursion;
+
+  // Make list of conflicting nodes. With useExistingNodes, this list stays
+  // empty, and the branch set is only otherwise used by the rename-all mode
+  // that excludes pruning, so the scan is skipped when pruning is active.
   RooArgSet conflictNodes ;
   RooArgSet branchSet ;
-  if (noRecursion) {
-    branchSet.add(inArg) ;
-  } else {
-    inArg.branchNodeServerList(&branchSet) ;
-  }
+  if (!pruneClones) {
+     if (noRecursion) {
+        branchSet.add(inArg);
+     } else {
+        inArg.branchNodeServerList(&branchSet);
+     }
 
-  for (const auto branch : branchSet) {
-    RooAbsArg* wsbranch = _allOwnedNodes.find(branch->GetName()) ;
-    if (wsbranch && wsbranch!=branch && !branch->getAttribute("RooWorkspace::Recycle") && !useExistingNodes) {
-      conflictNodes.add(*branch) ;
-    }
+     for (const auto branch : branchSet) {
+        RooAbsArg *wsbranch = _allOwnedNodes.find(branch->GetName());
+        if (wsbranch && wsbranch != branch && !branch->getAttribute("RooWorkspace::Recycle") && !useExistingNodes) {
+           conflictNodes.add(*branch);
+        }
+     }
   }
 
   // Terminate here if there are conflicts and no resolution protocol
@@ -515,7 +528,73 @@ bool RooWorkspace::import(const RooAbsArg& inArg,
   // Now create a working copy of the incoming object tree
   RooArgSet cloneSet;
   cloneSet.useHashMapForFind(true); // Accelerate finding
-  RooArgSet(inArg).snapshot(cloneSet, !noRecursion);
+  // Nodes of the input graph whose names already exist in the workspace and
+  // that are hence recycled instead of cloned.
+  std::vector<RooAbsArg const *> recycledSources;
+  if (pruneClones) {
+     // Traverse the whole graph like the snapshot in the general code path
+     // does (depth-first pre-order over the server links, keeping the first
+     // encountered instance for each name), but clone only the nodes whose
+     // names are not in the workspace yet.
+     std::vector<RooAbsArg const *> toClone;
+     std::unordered_set<TNamed const *> visited;
+     visited.insert(inArg.namePtr());
+     std::vector<RooAbsArg const *> stack{&inArg};
+     while (!stack.empty()) {
+        RooAbsArg const *arg = stack.back();
+        stack.pop_back();
+        if (RooAbsArg *wsnode = _allOwnedNodes.find(arg->GetName())) {
+           if (!silence) {
+              coutI(ObjectHandling) << "RooWorkspace::import(" << GetName() << ") using existing copy of "
+                                    << wsnode->ClassName() << "::" << wsnode->GetName() << " for import of "
+                                    << inArg.ClassName() << "::" << inArg.GetName() << std::endl;
+           }
+           recycledSources.push_back(arg);
+        } else {
+           toClone.push_back(arg);
+        }
+        // Push the servers in reverse order so that they are popped in
+        // forward order, matching the traversal order of the snapshot.
+        auto const &servers = arg->servers().containedObjects();
+        for (auto it = servers.rbegin(); it != servers.rend(); ++it) {
+           if (visited.insert((*it)->namePtr()).second) {
+              stack.push_back(*it);
+           }
+        }
+     }
+     for (RooAbsArg const *arg : toClone) {
+        auto clone = std::unique_ptr<RooAbsArg>{static_cast<RooAbsArg *>(arg->Clone())};
+        clone->setAttribute("SnapShot_ExtRefClone");
+        cloneSet.addOwned(std::move(clone));
+     }
+     // Redirect the server links of the clones to the other clones or to the
+     // existing nodes in the workspace.
+     RooArgSet redirectSet;
+     redirectSet.useHashMapForFind(true);
+     redirectSet.add(cloneSet);
+     for (RooAbsArg const *arg : recycledSources) {
+        redirectSet.add(*_allOwnedNodes.find(arg->GetName()));
+     }
+     for (RooAbsArg *clone : cloneSet) {
+        clone->redirectServers(redirectSet, true);
+     }
+  } else {
+     RooArgSet(inArg).snapshot(cloneSet, !noRecursion);
+  }
+
+  // Import the expensive-object cache payloads for recycled nodes like the
+  // general code path does via the discarded clones, which always hand over
+  // the global cache instance because the cache pointer of a RooAbsArg is
+  // transient and not copied when cloning.
+  for (RooAbsArg const *arg : recycledSources) {
+     _eocache.importCacheObjects(RooExpensiveObjectCache::instance(), arg->GetName(), true);
+  }
+
+  if (pruneClones && cloneSet.empty()) {
+     // The whole computation graph is already in the workspace.
+     return false;
+  }
+
   RooAbsArg* cloneTop = cloneSet.find(inArg.GetName()) ;
 
   // Mark all nodes for renaming if we are not in conflictOnly mode
@@ -528,7 +607,10 @@ bool RooWorkspace::import(const RooAbsArg& inArg,
   // Track whether any node in cloneSet actually gets renamed below: only then
   // does the working copy have to be cloned a second time (see there).
   bool nodesRenamed = false ;
-  string topName2 = cloneTop->GetName() ;
+  // With clone pruning, the top-level node itself may have been recycled, in
+  // which case it has no clone in cloneSet. No renaming can happen then, so
+  // the original name is already final.
+  string topName2 = cloneTop ? cloneTop->GetName() : inArg.GetName();
   if (!renameConflictOrig) {
     // Mark all nodes to be imported for renaming following conflict resolution protocol
     for (const auto cnode : conflictNodes) {