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 <stack>
#include <vector>
#include <string>
#include <boost/shared_ptr.hpp>
#include <boost/make_shared.hpp>
#include <boost/foreach.hpp>
#include <boost/bind.hpp>

namespace gtsam {

  namespace treeTraversal {

    /* ************************************************************************* */
    namespace {
      // Internal node used in DFS preorder stack
      template<typename NODE, typename DATA>
      struct TraversalNode {
        bool expanded;
        const boost::shared_ptr<NODE>& treeNode;
        DATA& parentData;
        typename FastList<DATA>::iterator dataPointer;
        TraversalNode(const boost::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 boost::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 boost::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();
        BOOST_FOREACH(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();
          BOOST_FOREACH(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;
      typedef boost::shared_ptr<Node> sharedNode;

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


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

    template<class FOREST>
    FastVector<boost::shared_ptr<typename FOREST::Node> > CloneForest(const FOREST& forest)
    {
      typedef typename FOREST::Node Node;
      boost::shared_ptr<Node> rootContainer = boost::make_shared<Node>();
      DepthFirstForest(forest, rootContainer, CloneForestVisitorPre<Node>);
      return FastVector<boost::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 boost::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);
    }
  }

}