treeTraversal-inst.h

Go to the documentation of this file.

/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

#pragma once

#include <gtsam/base/treeTraversal/parallelTraversalTasks.h>
#include <gtsam/base/treeTraversal/statistics.h>

#include <gtsam/base/FastList.h>
#include <gtsam/base/FastVector.h>
#include <gtsam/inference/Key.h>
#include <gtsam/config.h> // for GTSAM_USE_TBB

#include <stack>
#include <vector>
#include <string>
#include <memory>
#include <cassert>

namespace gtsam {

namespace treeTraversal {

/* ************************************************************************* */
namespace {
// Internal node used in DFS preorder stack
template<typename NODE, typename DATA>
struct TraversalNode {
  bool expanded;
  const std::shared_ptr<NODE>& treeNode;
  DATA& parentData;
  typename FastList<DATA>::iterator dataPointer;
  TraversalNode(const std::shared_ptr<NODE>& _treeNode, DATA& _parentData) :
      expanded(false), treeNode(_treeNode), parentData(_parentData) {
  }
};

// Do nothing - default argument for post-visitor for tree traversal
struct no_op {
  template<typename NODE, typename DATA>
  void operator()(const std::shared_ptr<NODE>& node, const DATA& data) {
  }
};

}

template<class FOREST, typename DATA, typename VISITOR_PRE,
    typename VISITOR_POST>
void DepthFirstForest(FOREST& forest, DATA& rootData, VISITOR_PRE& visitorPre,
    VISITOR_POST& visitorPost) {
  // Typedefs
  typedef typename FOREST::Node Node;
  typedef std::shared_ptr<Node> sharedNode;

  // Depth first traversal stack
  typedef TraversalNode<typename FOREST::Node, DATA> TraversalNode;
  typedef FastList<TraversalNode> Stack;
  Stack stack;
  FastList<DATA> dataList; // List to store node data as it is returned from the pre-order visitor

  // Add roots to stack (insert such that they are visited and processed in order
  {
    typename Stack::iterator insertLocation = stack.begin();
    for(const sharedNode& root: forest.roots())
      stack.insert(insertLocation, TraversalNode(root, rootData));
  }

  // Traverse
  while (!stack.empty()) {
    // Get next node
    TraversalNode& node = stack.front();

    if (node.expanded) {
      // If already expanded, then the data stored in the node is no longer needed, so visit
      // then delete it.
      (void) visitorPost(node.treeNode, *node.dataPointer);
      dataList.erase(node.dataPointer);
      stack.pop_front();
    } else {
      // If not already visited, visit the node and add its children (use reverse iterators so
      // children are processed in the order they appear)
      node.dataPointer = dataList.insert(dataList.end(),
          visitorPre(node.treeNode, node.parentData));
      typename Stack::iterator insertLocation = stack.begin();
      for(const sharedNode& child: node.treeNode->children)
        stack.insert(insertLocation, TraversalNode(child, *node.dataPointer));
      node.expanded = true;
    }
  }
  assert(dataList.empty());
}

template<class FOREST, typename DATA, typename VISITOR_PRE>
void DepthFirstForest(FOREST& forest, DATA& rootData, VISITOR_PRE& visitorPre) {
  no_op visitorPost;
  DepthFirstForest(forest, rootData, visitorPre, visitorPost);
}

template<class FOREST, typename DATA, typename VISITOR_PRE,
    typename VISITOR_POST>
void DepthFirstForestParallel(FOREST& forest, DATA& rootData,
    VISITOR_PRE& visitorPre, VISITOR_POST& visitorPost,
    int problemSizeThreshold = 10) {
#ifdef GTSAM_USE_TBB
  // Typedefs
  typedef typename FOREST::Node Node;

  internal::CreateRootTask<Node>(forest.roots(), rootData, visitorPre,
      visitorPost, problemSizeThreshold);
#else
  DepthFirstForest(forest, rootData, visitorPre, visitorPost);
#endif
}

/* ************************************************************************* */
namespace {
template<typename NODE>
std::shared_ptr<NODE> CloneForestVisitorPre(
    const std::shared_ptr<NODE>& node,
    const std::shared_ptr<NODE>& parentPointer) {
  // Clone the current node and add it to its cloned parent
  std::shared_ptr<NODE> clone = std::make_shared<NODE>(*node);
  clone->children.clear();
  parentPointer->children.push_back(clone);
  return clone;
}
}

template<class FOREST>
FastVector<std::shared_ptr<typename FOREST::Node> > CloneForest(
    const FOREST& forest) {
  typedef typename FOREST::Node Node;
  std::shared_ptr<Node> rootContainer = std::make_shared<Node>();
  DepthFirstForest(forest, rootContainer, CloneForestVisitorPre<Node>);
  return FastVector<std::shared_ptr<Node> >(rootContainer->children.begin(),
      rootContainer->children.end());
}

/* ************************************************************************* */
namespace {
struct PrintForestVisitorPre {
  const KeyFormatter& formatter;
  PrintForestVisitorPre(const KeyFormatter& formatter) :
      formatter(formatter) {
  }
  template<typename NODE> std::string operator()(
      const std::shared_ptr<NODE>& node, const std::string& parentString) {
    // Print the current node
    node->print(parentString + "-", formatter);
    // Increment the indentation
    return parentString + "| ";
  }
};
}

template<class FOREST>
void PrintForest(const FOREST& forest, std::string str,
    const KeyFormatter& keyFormatter) {
  PrintForestVisitorPre visitor(keyFormatter);
  DepthFirstForest(forest, str, visitor);
}
}  // namespace treeTraversal

}  // namespace gtsam