<commit_before>/**
 * @file neighbor_search_impl.hpp
 * @author Ryan Curtin
 *
 * Implementation of Neighbor-Search class to perform all-nearest-neighbors on
 * two specified data sets.
 */
#ifndef __MLPACK_METHODS_NEIGHBOR_SEARCH_NEIGHBOR_SEARCH_IMPL_HPP
#define __MLPACK_METHODS_NEIGHBOR_SEARCH_NEIGHBOR_SEARCH_IMPL_HPP

#include <mlpack/core.hpp>

#include "neighbor_search_rules.hpp"

namespace mlpack {
namespace neighbor {

//! Call the tree constructor that does mapping.
template<typename TreeType>
TreeType* BuildTree(
    typename TreeType::Mat& dataset,
    std::vector<size_t>& oldFromNew,
    typename boost::enable_if_c<
        tree::TreeTraits<TreeType>::RearrangesDataset == true, TreeType*
    >::type = 0)
{
  return new TreeType(dataset, oldFromNew);
}

//! Call the tree constructor that does not do mapping.
template<typename TreeType>
TreeType* BuildTree(
    const typename TreeType::Mat& dataset,
    const std::vector<size_t>& /* oldFromNew */,
    const typename boost::enable_if_c<
        tree::TreeTraits<TreeType>::RearrangesDataset == false, TreeType*
    >::type = 0)
{
  return new TreeType(dataset);
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::
NeighborSearch(const typename TreeType::Mat& referenceSetIn,
               const typename TreeType::Mat& querySetIn,
               const bool naive,
               const bool singleMode,
               const MetricType metric) :
    referenceSet(tree::TreeTraits<TreeType>::RearrangesDataset ? referenceCopy
        : referenceSetIn),
    querySet(tree::TreeTraits<TreeType>::RearrangesDataset ? queryCopy
        : querySetIn),
    referenceTree(NULL),
    queryTree(NULL),
    treeOwner(!naive), // False if a tree was passed.  If naive, then no trees.
    hasQuerySet(true),
    naive(naive),
    singleMode(!naive && singleMode), // No single mode if naive.
    metric(metric),
    baseCases(0),
    scores(0)
{
  // C++11 will allow us to call out to other constructors so we can avoid this
  // copypasta problem.

  // We'll time tree building, but only if we are building trees.
  Timer::Start("tree_building");

  // Copy the datasets, if they will be modified during tree building.
  if (tree::TreeTraits<TreeType>::RearrangesDataset)
  {
    referenceCopy = referenceSetIn;
    queryCopy = querySetIn;
  }

  // If not in naive mode, then we need to build trees.
  if (!naive)
  {
    // The const_cast is safe; if RearrangesDataset == false, then it'll be
    // casted back to const anyway, and if not, referenceSet points to
    // referenceCopy, which isn't const.
    referenceTree = BuildTree<TreeType>(
        const_cast<typename TreeType::Mat&>(referenceSet),
        oldFromNewReferences);

    if (!singleMode)
      queryTree = BuildTree<TreeType>(
          const_cast<typename TreeType::Mat&>(querySet), oldFromNewQueries);
  }

  // Stop the timer we started above (if we need to).
  Timer::Stop("tree_building");
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::
NeighborSearch(const typename TreeType::Mat& referenceSetIn,
               const bool naive,
               const bool singleMode,
               const MetricType metric) :
    referenceSet(tree::TreeTraits<TreeType>::RearrangesDataset ? referenceCopy
        : referenceSetIn),
    querySet(tree::TreeTraits<TreeType>::RearrangesDataset ? referenceCopy
        : referenceSetIn),
    referenceTree(NULL),
    queryTree(NULL),
    treeOwner(!naive), // If naive, then we are not building any trees.
    hasQuerySet(false),
    naive(naive),
    singleMode(!naive && singleMode), // No single mode if naive.
    metric(metric),
    baseCases(0),
    scores(0)
{
  // We'll time tree building, but only if we are building trees.
  Timer::Start("tree_building");

  // Copy the dataset, if it will be modified during tree building.
  if (tree::TreeTraits<TreeType>::RearrangesDataset)
    referenceCopy = referenceSetIn;

  // If not in naive mode, then we may need to construct trees.
  if (!naive)
  {
    // The const_cast is safe; if RearrangesDataset == false, then it'll be
    // casted back to const anyway, and if not, referenceSet points to
    // referenceCopy, which isn't const.
    referenceTree = BuildTree<TreeType>(
        const_cast<typename TreeType::Mat&>(referenceSet),
        oldFromNewReferences);

    if (!singleMode)
      queryTree = new TreeType(*referenceTree);
  }

  // Stop the timer we started above.
  Timer::Stop("tree_building");
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::NeighborSearch(
    TreeType* referenceTree,
    TreeType* queryTree,
    const typename TreeType::Mat& referenceSet,
    const typename TreeType::Mat& querySet,
    const bool singleMode,
    const MetricType metric) :
    referenceSet(referenceSet),
    querySet(querySet),
    referenceTree(referenceTree),
    queryTree(queryTree),
    treeOwner(false),
    hasQuerySet(true),
    naive(false),
    singleMode(singleMode),
    metric(metric),
    baseCases(0),
    scores(0)
{
  // Nothing else to initialize.
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::NeighborSearch(
    TreeType* referenceTree,
    const typename TreeType::Mat& referenceSet,
    const bool singleMode,
    const MetricType metric) :
    referenceSet(referenceSet),
    querySet(referenceSet),
    referenceTree(referenceTree),
    queryTree(NULL),
    treeOwner(false),
    hasQuerySet(false), // In this case we will own a tree, if singleMode.
    naive(false),
    singleMode(singleMode),
    metric(metric),
    baseCases(0),
    scores(0)
{
  Timer::Start("tree_building");

  // The query tree cannot be the same as the reference tree.
  if (referenceTree && !singleMode)
    queryTree = new TreeType(*referenceTree);

  Timer::Stop("tree_building");
}

/**
 * The tree is the only member we may be responsible for deleting.  The others
 * will take care of themselves.
 */
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::~NeighborSearch()
{
  if (treeOwner)
  {
    if (referenceTree)
      delete referenceTree;
    if (queryTree)
      delete queryTree;
  }
  else if (!treeOwner && !hasQuerySet && !(singleMode || naive))
  {
    // We replicated the reference tree to create a query tree.
    delete queryTree;
  }
}

/**
 * Computes the best neighbors and stores them in resultingNeighbors and
 * distances.
 */
template<typename SortPolicy, typename MetricType, typename TreeType>
void NeighborSearch<SortPolicy, MetricType, TreeType>::Search(
    const size_t k,
    arma::Mat<size_t>& resultingNeighbors,
    arma::mat& distances)
{
  Timer::Start("computing_neighbors");

  // If we have built the trees ourselves, then we will have to map all the
  // indices back to their original indices when this computation is finished.
  // To avoid an extra copy, we will store the neighbors and distances in a
  // separate matrix.
  arma::Mat<size_t>* neighborPtr = &resultingNeighbors;
  arma::mat* distancePtr = &distances;

  // Mapping is only necessary if the tree rearranges points.
  if (tree::TreeTraits<TreeType>::RearrangesDataset)
  {
    if (treeOwner && !(singleMode && hasQuerySet))
      distancePtr = new arma::mat; // Query indices need to be mapped.

    if (treeOwner)
      neighborPtr = new arma::Mat<size_t>; // All indices need mapping.
  }

  // Set the size of the neighbor and distance matrices.
  neighborPtr->set_size(k, querySet.n_cols);
  neighborPtr->fill(size_t() - 1);
  distancePtr->set_size(k, querySet.n_cols);
  distancePtr->fill(SortPolicy::WorstDistance());

  // Create the helper object for the tree traversal.
  typedef NeighborSearchRules<SortPolicy, MetricType, TreeType> RuleType;
  RuleType rules(referenceSet, querySet, *neighborPtr, *distancePtr, metric);

  if (naive)
  {
    // The naive brute-force traversal.
    for (size_t i = 0; i < querySet.n_cols; ++i)
      for (size_t j = 0; j < referenceSet.n_cols; ++j)
        rules.BaseCase(i, j);

    baseCases += querySet.n_cols * referenceSet.n_cols;
  }
  else if (singleMode)
  {
    // The search doesn't work if the root node is also a leaf node.
    // If this is the case, it is suggested that you use the naive method.
    Log::Assert(!(referenceTree->IsLeaf()));

    // Create the traverser.
    typename TreeType::template SingleTreeTraverser<RuleType> traverser(rules);

    // Now have it traverse for each point.
    for (size_t i = 0; i < querySet.n_cols; ++i)
      traverser.Traverse(i, *referenceTree);

    scores += rules.Scores();
    baseCases += rules.BaseCases();

    Log::Info << rules.Scores() << " node combinations were scored.\n";
    Log::Info << rules.BaseCases() << " base cases were calculated.\n";
  }
  else // Dual-tree recursion.
  {
    // Create the traverser.
    typename TreeType::template DualTreeTraverser<RuleType> traverser(rules);

    traverser.Traverse(*queryTree, *referenceTree);

    scores += rules.Scores();
    baseCases += rules.BaseCases();

    Log::Info << rules.Scores() << " node combinations were scored.\n";
    Log::Info << rules.BaseCases() << " base cases were calculated.\n";
  }

  Timer::Stop("computing_neighbors");

  // Now, do we need to do mapping of indices?
  if (!treeOwner || !tree::TreeTraits<TreeType>::RearrangesDataset)
  {
    // No mapping needed.  We are done.
    return;
  }
  else if (treeOwner && hasQuerySet && !singleMode) // Map both sets.
  {
    // Set size of output matrices correctly.
    resultingNeighbors.set_size(k, querySet.n_cols);
    distances.set_size(k, querySet.n_cols);

    for (size_t i = 0; i < distances.n_cols; i++)
    {
      // Map distances (copy a column).
      distances.col(oldFromNewQueries[i]) = distancePtr->col(i);

      // Map indices of neighbors.
      for (size_t j = 0; j < distances.n_rows; j++)
      {
        resultingNeighbors(j, oldFromNewQueries[i]) =
            oldFromNewReferences[(*neighborPtr)(j, i)];
      }
    }

    // Finished with temporary matrices.
    delete neighborPtr;
    delete distancePtr;
  }
  else if (treeOwner && !hasQuerySet)
  {
    resultingNeighbors.set_size(k, querySet.n_cols);
    distances.set_size(k, querySet.n_cols);

    for (size_t i = 0; i < distances.n_cols; i++)
    {
      // Map distances (copy a column).
      distances.col(oldFromNewReferences[i]) = distancePtr->col(i);

      // Map indices of neighbors.
      for (size_t j = 0; j < distances.n_rows; j++)
      {
        resultingNeighbors(j, oldFromNewReferences[i]) =
            oldFromNewReferences[(*neighborPtr)(j, i)];
      }
    }

    // Finished with temporary matrices.
    delete neighborPtr;
    delete distancePtr;
  }
  else if (treeOwner && hasQuerySet && singleMode) // Map only references.
  {
    // Set size of neighbor indices matrix correctly.
    resultingNeighbors.set_size(k, querySet.n_cols);

    // Map indices of neighbors.
    for (size_t i = 0; i < resultingNeighbors.n_cols; i++)
    {
      for (size_t j = 0; j < resultingNeighbors.n_rows; j++)
      {
        resultingNeighbors(j, i) = oldFromNewReferences[(*neighborPtr)(j, i)];
      }
    }

    // Finished with temporary matrix.
    delete neighborPtr;
  }
} // Search

//Return a String of the Object.
template<typename SortPolicy, typename MetricType, typename TreeType>
std::string NeighborSearch<SortPolicy, MetricType, TreeType>::ToString() const
{
  std::ostringstream convert;
  convert << "NearestNeighborSearch [" << this << "]" << std::endl;
  convert << "  Reference Set: " << referenceSet.n_rows << "x" ;
  convert <<  referenceSet.n_cols << std::endl;
  if (&referenceSet != &querySet)
    convert << "  QuerySet: " << querySet.n_rows << "x" << querySet.n_cols
        << std::endl;
  convert << "  Reference Tree: " << referenceTree << std::endl;
  if (&referenceTree != &queryTree)
    convert << "  QueryTree: " << queryTree << std::endl;
  convert << "  Tree Owner: " << treeOwner << std::endl;
  convert << "  Naive: " << naive << std::endl;
  convert << "  Metric: " << std::endl;
  convert << mlpack::util::Indent(metric.ToString(),2);
  return convert.str();
}

}; // namespace neighbor
}; // namespace mlpack

#endif
<commit_msg>Fix incorrect class name.<commit_after>/**
 * @file neighbor_search_impl.hpp
 * @author Ryan Curtin
 *
 * Implementation of Neighbor-Search class to perform all-nearest-neighbors on
 * two specified data sets.
 */
#ifndef __MLPACK_METHODS_NEIGHBOR_SEARCH_NEIGHBOR_SEARCH_IMPL_HPP
#define __MLPACK_METHODS_NEIGHBOR_SEARCH_NEIGHBOR_SEARCH_IMPL_HPP

#include <mlpack/core.hpp>

#include "neighbor_search_rules.hpp"

namespace mlpack {
namespace neighbor {

//! Call the tree constructor that does mapping.
template<typename TreeType>
TreeType* BuildTree(
    typename TreeType::Mat& dataset,
    std::vector<size_t>& oldFromNew,
    typename boost::enable_if_c<
        tree::TreeTraits<TreeType>::RearrangesDataset == true, TreeType*
    >::type = 0)
{
  return new TreeType(dataset, oldFromNew);
}

//! Call the tree constructor that does not do mapping.
template<typename TreeType>
TreeType* BuildTree(
    const typename TreeType::Mat& dataset,
    const std::vector<size_t>& /* oldFromNew */,
    const typename boost::enable_if_c<
        tree::TreeTraits<TreeType>::RearrangesDataset == false, TreeType*
    >::type = 0)
{
  return new TreeType(dataset);
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::
NeighborSearch(const typename TreeType::Mat& referenceSetIn,
               const typename TreeType::Mat& querySetIn,
               const bool naive,
               const bool singleMode,
               const MetricType metric) :
    referenceSet(tree::TreeTraits<TreeType>::RearrangesDataset ? referenceCopy
        : referenceSetIn),
    querySet(tree::TreeTraits<TreeType>::RearrangesDataset ? queryCopy
        : querySetIn),
    referenceTree(NULL),
    queryTree(NULL),
    treeOwner(!naive), // False if a tree was passed.  If naive, then no trees.
    hasQuerySet(true),
    naive(naive),
    singleMode(!naive && singleMode), // No single mode if naive.
    metric(metric),
    baseCases(0),
    scores(0)
{
  // C++11 will allow us to call out to other constructors so we can avoid this
  // copypasta problem.

  // We'll time tree building, but only if we are building trees.
  Timer::Start("tree_building");

  // Copy the datasets, if they will be modified during tree building.
  if (tree::TreeTraits<TreeType>::RearrangesDataset)
  {
    referenceCopy = referenceSetIn;
    queryCopy = querySetIn;
  }

  // If not in naive mode, then we need to build trees.
  if (!naive)
  {
    // The const_cast is safe; if RearrangesDataset == false, then it'll be
    // casted back to const anyway, and if not, referenceSet points to
    // referenceCopy, which isn't const.
    referenceTree = BuildTree<TreeType>(
        const_cast<typename TreeType::Mat&>(referenceSet),
        oldFromNewReferences);

    if (!singleMode)
      queryTree = BuildTree<TreeType>(
          const_cast<typename TreeType::Mat&>(querySet), oldFromNewQueries);
  }

  // Stop the timer we started above (if we need to).
  Timer::Stop("tree_building");
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::
NeighborSearch(const typename TreeType::Mat& referenceSetIn,
               const bool naive,
               const bool singleMode,
               const MetricType metric) :
    referenceSet(tree::TreeTraits<TreeType>::RearrangesDataset ? referenceCopy
        : referenceSetIn),
    querySet(tree::TreeTraits<TreeType>::RearrangesDataset ? referenceCopy
        : referenceSetIn),
    referenceTree(NULL),
    queryTree(NULL),
    treeOwner(!naive), // If naive, then we are not building any trees.
    hasQuerySet(false),
    naive(naive),
    singleMode(!naive && singleMode), // No single mode if naive.
    metric(metric),
    baseCases(0),
    scores(0)
{
  // We'll time tree building, but only if we are building trees.
  Timer::Start("tree_building");

  // Copy the dataset, if it will be modified during tree building.
  if (tree::TreeTraits<TreeType>::RearrangesDataset)
    referenceCopy = referenceSetIn;

  // If not in naive mode, then we may need to construct trees.
  if (!naive)
  {
    // The const_cast is safe; if RearrangesDataset == false, then it'll be
    // casted back to const anyway, and if not, referenceSet points to
    // referenceCopy, which isn't const.
    referenceTree = BuildTree<TreeType>(
        const_cast<typename TreeType::Mat&>(referenceSet),
        oldFromNewReferences);

    if (!singleMode)
      queryTree = new TreeType(*referenceTree);
  }

  // Stop the timer we started above.
  Timer::Stop("tree_building");
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::NeighborSearch(
    TreeType* referenceTree,
    TreeType* queryTree,
    const typename TreeType::Mat& referenceSet,
    const typename TreeType::Mat& querySet,
    const bool singleMode,
    const MetricType metric) :
    referenceSet(referenceSet),
    querySet(querySet),
    referenceTree(referenceTree),
    queryTree(queryTree),
    treeOwner(false),
    hasQuerySet(true),
    naive(false),
    singleMode(singleMode),
    metric(metric),
    baseCases(0),
    scores(0)
{
  // Nothing else to initialize.
}

// Construct the object.
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::NeighborSearch(
    TreeType* referenceTree,
    const typename TreeType::Mat& referenceSet,
    const bool singleMode,
    const MetricType metric) :
    referenceSet(referenceSet),
    querySet(referenceSet),
    referenceTree(referenceTree),
    queryTree(NULL),
    treeOwner(false),
    hasQuerySet(false), // In this case we will own a tree, if singleMode.
    naive(false),
    singleMode(singleMode),
    metric(metric),
    baseCases(0),
    scores(0)
{
  Timer::Start("tree_building");

  // The query tree cannot be the same as the reference tree.
  if (referenceTree && !singleMode)
    queryTree = new TreeType(*referenceTree);

  Timer::Stop("tree_building");
}

/**
 * The tree is the only member we may be responsible for deleting.  The others
 * will take care of themselves.
 */
template<typename SortPolicy, typename MetricType, typename TreeType>
NeighborSearch<SortPolicy, MetricType, TreeType>::~NeighborSearch()
{
  if (treeOwner)
  {
    if (referenceTree)
      delete referenceTree;
    if (queryTree)
      delete queryTree;
  }
  else if (!treeOwner && !hasQuerySet && !(singleMode || naive))
  {
    // We replicated the reference tree to create a query tree.
    delete queryTree;
  }
}

/**
 * Computes the best neighbors and stores them in resultingNeighbors and
 * distances.
 */
template<typename SortPolicy, typename MetricType, typename TreeType>
void NeighborSearch<SortPolicy, MetricType, TreeType>::Search(
    const size_t k,
    arma::Mat<size_t>& resultingNeighbors,
    arma::mat& distances)
{
  Timer::Start("computing_neighbors");

  // If we have built the trees ourselves, then we will have to map all the
  // indices back to their original indices when this computation is finished.
  // To avoid an extra copy, we will store the neighbors and distances in a
  // separate matrix.
  arma::Mat<size_t>* neighborPtr = &resultingNeighbors;
  arma::mat* distancePtr = &distances;

  // Mapping is only necessary if the tree rearranges points.
  if (tree::TreeTraits<TreeType>::RearrangesDataset)
  {
    if (treeOwner && !(singleMode && hasQuerySet))
      distancePtr = new arma::mat; // Query indices need to be mapped.

    if (treeOwner)
      neighborPtr = new arma::Mat<size_t>; // All indices need mapping.
  }

  // Set the size of the neighbor and distance matrices.
  neighborPtr->set_size(k, querySet.n_cols);
  neighborPtr->fill(size_t() - 1);
  distancePtr->set_size(k, querySet.n_cols);
  distancePtr->fill(SortPolicy::WorstDistance());

  // Create the helper object for the tree traversal.
  typedef NeighborSearchRules<SortPolicy, MetricType, TreeType> RuleType;
  RuleType rules(referenceSet, querySet, *neighborPtr, *distancePtr, metric);

  if (naive)
  {
    // The naive brute-force traversal.
    for (size_t i = 0; i < querySet.n_cols; ++i)
      for (size_t j = 0; j < referenceSet.n_cols; ++j)
        rules.BaseCase(i, j);

    baseCases += querySet.n_cols * referenceSet.n_cols;
  }
  else if (singleMode)
  {
    // The search doesn't work if the root node is also a leaf node.
    // If this is the case, it is suggested that you use the naive method.
    Log::Assert(!(referenceTree->IsLeaf()));

    // Create the traverser.
    typename TreeType::template SingleTreeTraverser<RuleType> traverser(rules);

    // Now have it traverse for each point.
    for (size_t i = 0; i < querySet.n_cols; ++i)
      traverser.Traverse(i, *referenceTree);

    scores += rules.Scores();
    baseCases += rules.BaseCases();

    Log::Info << rules.Scores() << " node combinations were scored.\n";
    Log::Info << rules.BaseCases() << " base cases were calculated.\n";
  }
  else // Dual-tree recursion.
  {
    // Create the traverser.
    typename TreeType::template DualTreeTraverser<RuleType> traverser(rules);

    traverser.Traverse(*queryTree, *referenceTree);

    scores += rules.Scores();
    baseCases += rules.BaseCases();

    Log::Info << rules.Scores() << " node combinations were scored.\n";
    Log::Info << rules.BaseCases() << " base cases were calculated.\n";
  }

  Timer::Stop("computing_neighbors");

  // Now, do we need to do mapping of indices?
  if (!treeOwner || !tree::TreeTraits<TreeType>::RearrangesDataset)
  {
    // No mapping needed.  We are done.
    return;
  }
  else if (treeOwner && hasQuerySet && !singleMode) // Map both sets.
  {
    // Set size of output matrices correctly.
    resultingNeighbors.set_size(k, querySet.n_cols);
    distances.set_size(k, querySet.n_cols);

    for (size_t i = 0; i < distances.n_cols; i++)
    {
      // Map distances (copy a column).
      distances.col(oldFromNewQueries[i]) = distancePtr->col(i);

      // Map indices of neighbors.
      for (size_t j = 0; j < distances.n_rows; j++)
      {
        resultingNeighbors(j, oldFromNewQueries[i]) =
            oldFromNewReferences[(*neighborPtr)(j, i)];
      }
    }

    // Finished with temporary matrices.
    delete neighborPtr;
    delete distancePtr;
  }
  else if (treeOwner && !hasQuerySet)
  {
    resultingNeighbors.set_size(k, querySet.n_cols);
    distances.set_size(k, querySet.n_cols);

    for (size_t i = 0; i < distances.n_cols; i++)
    {
      // Map distances (copy a column).
      distances.col(oldFromNewReferences[i]) = distancePtr->col(i);

      // Map indices of neighbors.
      for (size_t j = 0; j < distances.n_rows; j++)
      {
        resultingNeighbors(j, oldFromNewReferences[i]) =
            oldFromNewReferences[(*neighborPtr)(j, i)];
      }
    }

    // Finished with temporary matrices.
    delete neighborPtr;
    delete distancePtr;
  }
  else if (treeOwner && hasQuerySet && singleMode) // Map only references.
  {
    // Set size of neighbor indices matrix correctly.
    resultingNeighbors.set_size(k, querySet.n_cols);

    // Map indices of neighbors.
    for (size_t i = 0; i < resultingNeighbors.n_cols; i++)
    {
      for (size_t j = 0; j < resultingNeighbors.n_rows; j++)
      {
        resultingNeighbors(j, i) = oldFromNewReferences[(*neighborPtr)(j, i)];
      }
    }

    // Finished with temporary matrix.
    delete neighborPtr;
  }
} // Search

//Return a String of the Object.
template<typename SortPolicy, typename MetricType, typename TreeType>
std::string NeighborSearch<SortPolicy, MetricType, TreeType>::ToString() const
{
  std::ostringstream convert;
  convert << "NeighborSearch [" << this << "]" << std::endl;
  convert << "  Reference Set: " << referenceSet.n_rows << "x" ;
  convert <<  referenceSet.n_cols << std::endl;
  if (&referenceSet != &querySet)
    convert << "  QuerySet: " << querySet.n_rows << "x" << querySet.n_cols
        << std::endl;
  convert << "  Reference Tree: " << referenceTree << std::endl;
  if (&referenceTree != &queryTree)
    convert << "  QueryTree: " << queryTree << std::endl;
  convert << "  Tree Owner: " << treeOwner << std::endl;
  convert << "  Naive: " << naive << std::endl;
  convert << "  Metric: " << std::endl;
  convert << mlpack::util::Indent(metric.ToString(),2);
  return convert.str();
}

}; // namespace neighbor
}; // namespace mlpack

#endif
<|endoftext|>


<commit_before>// Copyright (C) 2019 Egor Pugin <egor.pugin@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#define SW_MAIN_IMPL
#include <primitives/sw/main.h>

#include <primitives/sw/settings.h>
#include <primitives/sw/cl.h>
#include <primitives/sw/settings_program_name.h>
#include <primitives/main.h>

#include <boost/dll.hpp>

#include <random>
#include <thread>

#define CRASH_SERVER_OPTION "internal-with-crash-server"
#define CRASH_SERVER_OPTION_START "internal-with-crash-server-start-server"

static cl::opt<int> with_crash_server(CRASH_SERVER_OPTION,
    cl::desc("Start crash handler server on startup. Use this option to record a crash dump."),
    cl::Hidden, cl::ValueOptional);

static cl::opt<String> crash_server_pipe_name(CRASH_SERVER_OPTION_START,
    cl::desc("Start crash handler server."),
    cl::ReallyHidden);

extern bool gUseStackTrace;
static ::cl::opt<bool, true> bt("bt", ::cl::desc("Print backtrace"), ::cl::location(gUseStackTrace), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));
extern bool gDebugOnException;
static ::cl::opt<bool, true> doe("doe", ::cl::desc("Debug on Exception"), ::cl::location(gDebugOnException), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));
static ::cl::alias bt2("st", ::cl::desc("Alias for -bt"), ::cl::aliasopt(bt));

extern std::string gSymbolPath;
static ::cl::opt<std::string, true> symbol_path("symbol-path", ::cl::desc("Set symbol path for backtrace"), ::cl::location(gSymbolPath), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));

static ::cl::opt<int> sleep_seconds("sleep", ::cl::desc("Sleep on startup"), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));
extern bool gPauseOnError;
static ::cl::opt<bool, true> pause_on_error("pause-on-error", ::cl::desc("Pause on program error"), ::cl::Hidden, ::cl::location(gPauseOnError), ::cl::sub(::cl::getAllSubCommands()));

static path temp_directory_path(const path &subdir)
{
    return fs::temp_directory_path() / "sw" / subdir;
}

static path get_crash_dir()
{
    auto p = temp_directory_path(::sw::getProgramName()) / "dump";
    fs::create_directories(p);
    return p;
}

static std::string get_cl_option_base()
{
    return "-" PACKAGE_NAME_CLEAN ".internal.start-minidump-server";
}

#ifdef _WIN32
#include <exception_handler.h>
#include <crash_generation_server.h>
#include <common/windows/guid_string.h>

#include <Windows.h>

static google_breakpad::CrashGenerationServer* crash_server;
static std::atomic_bool client_connected;

static google_breakpad::ExceptionHandler* handler;
static std::string client_pipe_name;
static std::wstring client_pipe_name_prepared;
static const wchar_t *client_pipe_name_prepared_c_str;
static std::string filter_message;
static std::string after_minidump_write_message;

static std::string get_object_prefix()
{
    return "sw_main_";
}

static std::string get_event_name_base()
{
    return "Local\\" + get_object_prefix();
}

static std::string get_pipe_name_base()
{
    return "\\\\.\\pipe\\" + get_object_prefix();
}

static int StartCrashGenerationServer(const std::string &win_pipe_id)
{
    auto dump_path = get_crash_dir().wstring();

    auto on_client_connect = [](void* context, const google_breakpad::ClientInfo* client_info)
    {
        client_connected = true;
    };

    auto on_client_disconnect = [](void* context, const google_breakpad::ClientInfo* client_info)
    {
        exit(0);
    };

    crash_server = new google_breakpad::CrashGenerationServer(to_wstring(get_pipe_name_base() + win_pipe_id), 0,
        on_client_connect, 0,
        0, 0,
        on_client_disconnect, 0,
        0, 0,
        true, &dump_path
    );

    auto r = crash_server->Start();
    if (r)
    {
        auto ev = OpenEvent(EVENT_MODIFY_STATE, 0, to_wstring(get_event_name_base() + win_pipe_id).c_str());
        if (!ev)
        {
            printf("OpenEvent failed: %d\n", GetLastError());
            return 1;
        }
        if (!SetEvent(ev))
        {
            printf("SetEvent failed: %d\n", GetLastError());
            return 1;
        }
        Sleep(10000);
        if (client_connected)
            Sleep(INFINITE);
        else
            r = false;
    }
    return !r;
}

static int LaunchCrashServerHelper()
{
    auto prog = boost::dll::program_location().wstring();
    auto cmd = L'\"' + prog + L'\"';
    cmd += to_wstring(" " + get_cl_option_base() + " " + client_pipe_name);

    auto pipe_name = to_wstring(get_pipe_name_base() + client_pipe_name);
    auto event_name = to_wstring(get_event_name_base() + client_pipe_name);

    STARTUPINFO si{ 0 };
    PROCESS_INFORMATION pi{ 0 };

    auto ev = CreateEvent(0, 0, 0, event_name.c_str());
    if (!ev)
    {
        printf("CreateEvent failed: %d\n", GetLastError());
        return false;
    }
    if (!CreateProcess(prog.c_str(), cmd.data(), 0, 0, 0, 0, 0, 0, &si, &pi))
    {
        printf("CreateProcess failed: %d\n", GetLastError());
        return false;
    }
    if (WaitForSingleObject(ev, 5000))
    {
        printf("Cannot connect to crash server\n");
        return false;
    }

    return true;
}

static void print_string_safe(const std::string &s)
{
    WriteFile(GetStdHandle(STD_ERROR_HANDLE), s.data(), s.size(), 0, 0);
}

static bool filter_cb(void *context, EXCEPTION_POINTERS *exinfo, MDRawAssertionInfo *assertion)
{
    print_string_safe(filter_message);
    return true;
}

static bool after_minidump_write_cb(const wchar_t *dump_path,
    const wchar_t *minidump_id,
    void *context,
    EXCEPTION_POINTERS *exinfo,
    MDRawAssertionInfo *assertion,
    bool succeeded)
{
    print_string_safe(after_minidump_write_message);
    return succeeded;
}
#endif

void sw_append_symbol_path(const path &in)
{
    // bt is not available at this time
    //if (!bt)
        //return;

    // not working atm
    return;

#ifdef _WIN32
    const auto nt_symbol_path = L"_NT_SYMBOL_PATH";
    const auto delim = L";";
    const auto max_env_var_size = 32'767;

    WCHAR buf[max_env_var_size];
    if (!GetEnvironmentVariableW(nt_symbol_path, buf, max_env_var_size))
    {
        auto e = GetLastError();
        if (e != ERROR_ENVVAR_NOT_FOUND)
        {
            printf("GetEnvironmentVariableW failed: %d", e);
            return;
        }
    }

    std::wstring p = buf;
    p += delim + in.wstring();
    SetEnvironmentVariableW(nt_symbol_path, p.c_str());
#endif
}

static int startup(int argc, char *argv[])
{
    for (int i = 1; i < argc - 1; i++)
    {
        if (argv[i] == "-sleep"s || argv[i] == "--sleep"s)
        {
            std::this_thread::sleep_for(std::chrono::seconds(std::stoi(argv[i + 1])));
            break;
        }
    }

#ifdef _WIN32
    // prepare strings
    filter_message = "Unhandled exception.\nWriting minidump to "s + to_string(get_crash_dir()) + "\n";
    after_minidump_write_message = "Minidump saved.\n";

    //
    auto option_name = get_cl_option_base();
    if (argc > 1 && argv[1] == option_name)
    {
        // out-of-process server
        cl::opt<std::string> id(cl::arg{ option_name.substr(1) }, cl::Required);
        cl::ParseCommandLineOptions(argc, argv);

        return StartCrashGenerationServer(id);
    }
    else if (argc > 1 && String(argv[1]).find(CRASH_SERVER_OPTION) != String::npos)
    {
        // with early out-of-process server
        int t = MiniDumpNormal;
        auto p = String(argv[1]).find('=');
        if (p != String::npos)
            t = std::stoi(String(argv[1]).substr(p + 1));
        // start server and handler
        sw_enable_crash_server(t);
    }
    else if (!IsDebuggerPresent())
    {
        // in-process handler
        handler = new google_breakpad::ExceptionHandler(get_crash_dir(),
            filter_cb, after_minidump_write_cb, nullptr, google_breakpad::ExceptionHandler::HANDLER_ALL);
    }
#endif

#ifdef _WIN32
    SetConsoleOutputCP(CP_UTF8);
#endif

    // fix current path:
    //  make it canonical (big drive letter on Windows)
    //  make it lexically normal
    //  remove trailing slashes

    // we also use fs::u8path() here, because normalize_path() converts it to utf-8

    auto cp = fs::current_path();
    auto s = to_string(normalize_path(cp.lexically_normal()));
    while (!s.empty() && s.back() == '/')
        s.resize(s.size() - 1);
#ifdef _WIN32
    // windows does not change case of the disk letter if other parts unchanged
    cp = s;
    cp = normalize_path_windows(cp);
    // so we take parent dir first, set cd there
    cp = cp.parent_path();
    fs::current_path(cp);
    // then we restore our cwd but with our changes
#endif
    cp = s;
    fs::current_path(cp);

    //
    sw_append_symbol_path(boost::dll::program_location().parent_path().wstring());
    sw_append_symbol_path(fs::current_path());

    // init settings early
    // e.g. before parsing command line
    sw::getSettingStorage();

    return 0;
}

int main(int argc, char *argv[])
{
    if (int r = startup(argc, argv))
        return r;
    return SW_MAIN(argc, argv);
}

void sw_enable_crash_server(int level)
{
#ifdef _WIN32
    delete handler; // delete previous always

    if (level < 0)
        return;

    // with early explicit server
    client_pipe_name = std::to_string(std::random_device()());
    client_pipe_name_prepared = to_wstring(get_pipe_name_base() + client_pipe_name);
    client_pipe_name_prepared_c_str = client_pipe_name_prepared.c_str();
    if (!LaunchCrashServerHelper())
    {
        printf("Cannot start crash server.\n");
        exit(1); // exit immediately? yes for now
        return;
    }
    handler = new google_breakpad::ExceptionHandler(get_crash_dir(),
        filter_cb, after_minidump_write_cb, nullptr, google_breakpad::ExceptionHandler::HANDLER_ALL,
        (MINIDUMP_TYPE)level, client_pipe_name_prepared_c_str, nullptr);
#endif
}
<commit_msg>[sw.main] Fix work in localized dirs.<commit_after>// Copyright (C) 2019 Egor Pugin <egor.pugin@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#define SW_MAIN_IMPL
#include <primitives/sw/main.h>

#include <primitives/sw/settings.h>
#include <primitives/sw/cl.h>
#include <primitives/sw/settings_program_name.h>
#include <primitives/main.h>

#include <boost/dll.hpp>

#include <random>
#include <thread>

#define CRASH_SERVER_OPTION "internal-with-crash-server"
#define CRASH_SERVER_OPTION_START "internal-with-crash-server-start-server"

static cl::opt<int> with_crash_server(CRASH_SERVER_OPTION,
    cl::desc("Start crash handler server on startup. Use this option to record a crash dump."),
    cl::Hidden, cl::ValueOptional);

static cl::opt<String> crash_server_pipe_name(CRASH_SERVER_OPTION_START,
    cl::desc("Start crash handler server."),
    cl::ReallyHidden);

extern bool gUseStackTrace;
static ::cl::opt<bool, true> bt("bt", ::cl::desc("Print backtrace"), ::cl::location(gUseStackTrace), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));
extern bool gDebugOnException;
static ::cl::opt<bool, true> doe("doe", ::cl::desc("Debug on Exception"), ::cl::location(gDebugOnException), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));
static ::cl::alias bt2("st", ::cl::desc("Alias for -bt"), ::cl::aliasopt(bt));

extern std::string gSymbolPath;
static ::cl::opt<std::string, true> symbol_path("symbol-path", ::cl::desc("Set symbol path for backtrace"), ::cl::location(gSymbolPath), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));

static ::cl::opt<int> sleep_seconds("sleep", ::cl::desc("Sleep on startup"), ::cl::Hidden, ::cl::sub(::cl::getAllSubCommands()));
extern bool gPauseOnError;
static ::cl::opt<bool, true> pause_on_error("pause-on-error", ::cl::desc("Pause on program error"), ::cl::Hidden, ::cl::location(gPauseOnError), ::cl::sub(::cl::getAllSubCommands()));

static path temp_directory_path(const path &subdir)
{
    return fs::temp_directory_path() / "sw" / subdir;
}

static path get_crash_dir()
{
    auto p = temp_directory_path(::sw::getProgramName()) / "dump";
    fs::create_directories(p);
    return p;
}

static std::string get_cl_option_base()
{
    return "-" PACKAGE_NAME_CLEAN ".internal.start-minidump-server";
}

#ifdef _WIN32
#include <exception_handler.h>
#include <crash_generation_server.h>
#include <common/windows/guid_string.h>

#include <Windows.h>

static google_breakpad::CrashGenerationServer* crash_server;
static std::atomic_bool client_connected;

static google_breakpad::ExceptionHandler* handler;
static std::string client_pipe_name;
static std::wstring client_pipe_name_prepared;
static const wchar_t *client_pipe_name_prepared_c_str;
static std::string filter_message;
static std::string after_minidump_write_message;

static std::string get_object_prefix()
{
    return "sw_main_";
}

static std::string get_event_name_base()
{
    return "Local\\" + get_object_prefix();
}

static std::string get_pipe_name_base()
{
    return "\\\\.\\pipe\\" + get_object_prefix();
}

static int StartCrashGenerationServer(const std::string &win_pipe_id)
{
    auto dump_path = get_crash_dir().wstring();

    auto on_client_connect = [](void* context, const google_breakpad::ClientInfo* client_info)
    {
        client_connected = true;
    };

    auto on_client_disconnect = [](void* context, const google_breakpad::ClientInfo* client_info)
    {
        exit(0);
    };

    crash_server = new google_breakpad::CrashGenerationServer(to_wstring(get_pipe_name_base() + win_pipe_id), 0,
        on_client_connect, 0,
        0, 0,
        on_client_disconnect, 0,
        0, 0,
        true, &dump_path
    );

    auto r = crash_server->Start();
    if (r)
    {
        auto ev = OpenEvent(EVENT_MODIFY_STATE, 0, to_wstring(get_event_name_base() + win_pipe_id).c_str());
        if (!ev)
        {
            printf("OpenEvent failed: %d\n", GetLastError());
            return 1;
        }
        if (!SetEvent(ev))
        {
            printf("SetEvent failed: %d\n", GetLastError());
            return 1;
        }
        Sleep(10000);
        if (client_connected)
            Sleep(INFINITE);
        else
            r = false;
    }
    return !r;
}

static int LaunchCrashServerHelper()
{
    auto prog = boost::dll::program_location().wstring();
    auto cmd = L'\"' + prog + L'\"';
    cmd += to_wstring(" " + get_cl_option_base() + " " + client_pipe_name);

    auto pipe_name = to_wstring(get_pipe_name_base() + client_pipe_name);
    auto event_name = to_wstring(get_event_name_base() + client_pipe_name);

    STARTUPINFO si{ 0 };
    PROCESS_INFORMATION pi{ 0 };

    auto ev = CreateEvent(0, 0, 0, event_name.c_str());
    if (!ev)
    {
        printf("CreateEvent failed: %d\n", GetLastError());
        return false;
    }
    if (!CreateProcess(prog.c_str(), cmd.data(), 0, 0, 0, 0, 0, 0, &si, &pi))
    {
        printf("CreateProcess failed: %d\n", GetLastError());
        return false;
    }
    if (WaitForSingleObject(ev, 5000))
    {
        printf("Cannot connect to crash server\n");
        return false;
    }

    return true;
}

static void print_string_safe(const std::string &s)
{
    WriteFile(GetStdHandle(STD_ERROR_HANDLE), s.data(), s.size(), 0, 0);
}

static bool filter_cb(void *context, EXCEPTION_POINTERS *exinfo, MDRawAssertionInfo *assertion)
{
    print_string_safe(filter_message);
    return true;
}

static bool after_minidump_write_cb(const wchar_t *dump_path,
    const wchar_t *minidump_id,
    void *context,
    EXCEPTION_POINTERS *exinfo,
    MDRawAssertionInfo *assertion,
    bool succeeded)
{
    print_string_safe(after_minidump_write_message);
    return succeeded;
}
#endif

void sw_append_symbol_path(const path &in)
{
    // bt is not available at this time
    //if (!bt)
        //return;

    // not working atm
    return;

#ifdef _WIN32
    const auto nt_symbol_path = L"_NT_SYMBOL_PATH";
    const auto delim = L";";
    const auto max_env_var_size = 32'767;

    WCHAR buf[max_env_var_size];
    if (!GetEnvironmentVariableW(nt_symbol_path, buf, max_env_var_size))
    {
        auto e = GetLastError();
        if (e != ERROR_ENVVAR_NOT_FOUND)
        {
            printf("GetEnvironmentVariableW failed: %d", e);
            return;
        }
    }

    std::wstring p = buf;
    p += delim + in.wstring();
    SetEnvironmentVariableW(nt_symbol_path, p.c_str());
#endif
}

static int startup(int argc, char *argv[])
{
    for (int i = 1; i < argc - 1; i++)
    {
        if (argv[i] == "-sleep"s || argv[i] == "--sleep"s)
        {
            std::this_thread::sleep_for(std::chrono::seconds(std::stoi(argv[i + 1])));
            break;
        }
    }

#ifdef _WIN32
    // prepare strings
    filter_message = "Unhandled exception.\nWriting minidump to "s + to_string(get_crash_dir()) + "\n";
    after_minidump_write_message = "Minidump saved.\n";

    //
    auto option_name = get_cl_option_base();
    if (argc > 1 && argv[1] == option_name)
    {
        // out-of-process server
        cl::opt<std::string> id(cl::arg{ option_name.substr(1) }, cl::Required);
        cl::ParseCommandLineOptions(argc, argv);

        return StartCrashGenerationServer(id);
    }
    else if (argc > 1 && String(argv[1]).find(CRASH_SERVER_OPTION) != String::npos)
    {
        // with early out-of-process server
        int t = MiniDumpNormal;
        auto p = String(argv[1]).find('=');
        if (p != String::npos)
            t = std::stoi(String(argv[1]).substr(p + 1));
        // start server and handler
        sw_enable_crash_server(t);
    }
    else if (!IsDebuggerPresent())
    {
        // in-process handler
        handler = new google_breakpad::ExceptionHandler(get_crash_dir(),
            filter_cb, after_minidump_write_cb, nullptr, google_breakpad::ExceptionHandler::HANDLER_ALL);
    }
#endif

#ifdef _WIN32
    SetConsoleOutputCP(CP_UTF8);
#endif

    // fix current path:
    //  make it canonical (big drive letter on Windows)
    //  make it lexically normal
    //  remove trailing slashes

    // we also use fs::u8path() here, because normalize_path() converts it to utf-8

    auto cp = fs::current_path();
    auto s = to_path_string(normalize_path(cp.lexically_normal()));
    while (!s.empty() && s.back() == '/')
        s.resize(s.size() - 1);
#ifdef _WIN32
    // windows does not change case of the disk letter if other parts unchanged
    cp = s;
    cp = normalize_path_windows(cp);
    // so we take parent dir first, set cd there
    cp = cp.parent_path();
    fs::current_path(cp);
    // then we restore our cwd but with our changes
#endif
    cp = s;
    fs::current_path(cp);

    //
    sw_append_symbol_path(boost::dll::program_location().parent_path().wstring());
    sw_append_symbol_path(fs::current_path());

    // init settings early
    // e.g. before parsing command line
    sw::getSettingStorage();

    return 0;
}

int main(int argc, char *argv[])
{
    if (int r = startup(argc, argv))
        return r;
    return SW_MAIN(argc, argv);
}

void sw_enable_crash_server(int level)
{
#ifdef _WIN32
    delete handler; // delete previous always

    if (level < 0)
        return;

    // with early explicit server
    client_pipe_name = std::to_string(std::random_device()());
    client_pipe_name_prepared = to_wstring(get_pipe_name_base() + client_pipe_name);
    client_pipe_name_prepared_c_str = client_pipe_name_prepared.c_str();
    if (!LaunchCrashServerHelper())
    {
        printf("Cannot start crash server.\n");
        exit(1); // exit immediately? yes for now
        return;
    }
    handler = new google_breakpad::ExceptionHandler(get_crash_dir(),
        filter_cb, after_minidump_write_cb, nullptr, google_breakpad::ExceptionHandler::HANDLER_ALL,
        (MINIDUMP_TYPE)level, client_pipe_name_prepared_c_str, nullptr);
#endif
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "coins.h"
#include "random.h"
#include "uint256.h"
#include "test/test_bitcoin.h"

#include <vector>
#include <map>

#include <boost/test/unit_test.hpp>
#include "zcash/IncrementalMerkleTree.hpp"

namespace
{
class CCoinsViewTest : public CCoinsView
{
    uint256 hashBestBlock_;
    uint256 hashBestAnchor_;
    std::map<uint256, CCoins> map_;
    std::map<uint256, ZCIncrementalMerkleTree> mapAnchors_;
    std::map<uint256, bool> mapSerials_;

public:
    bool GetAnchorAt(const uint256& rt, ZCIncrementalMerkleTree &tree) const {
        if (rt.IsNull()) {
            ZCIncrementalMerkleTree new_tree;
            tree = new_tree;
            return true;
        }

        std::map<uint256, ZCIncrementalMerkleTree>::const_iterator it = mapAnchors_.find(rt);
        if (it == mapAnchors_.end()) {
            return false;
        } else {
            tree = it->second;
            return true;
        }
    }

    bool GetSerial(const uint256 &serial) const
    {
        std::map<uint256, bool>::const_iterator it = mapSerials_.find(serial);

        if (it == mapSerials_.end()) {
            return false;
        } else {
            // The map shouldn't contain any false entries.
            assert(it->second);
            return true;
        }
    }

    uint256 GetBestAnchor() const { return hashBestAnchor_; }

    bool GetCoins(const uint256& txid, CCoins& coins) const
    {
        std::map<uint256, CCoins>::const_iterator it = map_.find(txid);
        if (it == map_.end()) {
            return false;
        }
        coins = it->second;
        if (coins.IsPruned() && insecure_rand() % 2 == 0) {
            // Randomly return false in case of an empty entry.
            return false;
        }
        return true;
    }

    bool HaveCoins(const uint256& txid) const
    {
        CCoins coins;
        return GetCoins(txid, coins);
    }

    uint256 GetBestBlock() const { return hashBestBlock_; }

    bool BatchWrite(CCoinsMap& mapCoins,
                    const uint256& hashBlock,
                    const uint256& hashAnchor,
                    CAnchorsMap& mapAnchors,
                    CSerialsMap& mapSerials)
    {
        for (CCoinsMap::iterator it = mapCoins.begin(); it != mapCoins.end(); ) {
            map_[it->first] = it->second.coins;
            if (it->second.coins.IsPruned() && insecure_rand() % 3 == 0) {
                // Randomly delete empty entries on write.
                map_.erase(it->first);
            }
            mapCoins.erase(it++);
        }
        for (CAnchorsMap::iterator it = mapAnchors.begin(); it != mapAnchors.end(); ) {
            if (it->second.entered) {
                std::map<uint256, ZCIncrementalMerkleTree>::iterator ret =
                    mapAnchors_.insert(std::make_pair(it->first, ZCIncrementalMerkleTree())).first;

                ret->second = it->second.tree;
            } else {
                mapAnchors_.erase(it->first);
            }
            mapAnchors.erase(it++);
        }
        for (CSerialsMap::iterator it = mapSerials.begin(); it != mapSerials.end(); ) {
            if (it->second.entered) {
                mapSerials_[it->first] = true;
            } else {
                mapSerials_.erase(it->first);
            }
            mapSerials.erase(it++);
        }
        mapCoins.clear();
        mapAnchors.clear();
        mapSerials.clear();
        hashBestBlock_ = hashBlock;
        hashBestAnchor_ = hashAnchor;
        return true;
    }

    bool GetStats(CCoinsStats& stats) const { return false; }
};

class CCoinsViewCacheTest : public CCoinsViewCache
{
public:
    CCoinsViewCacheTest(CCoinsView* base) : CCoinsViewCache(base) {}

    void SelfTest() const
    {
        // Manually recompute the dynamic usage of the whole data, and compare it.
        size_t ret = memusage::DynamicUsage(cacheCoins);
        for (CCoinsMap::iterator it = cacheCoins.begin(); it != cacheCoins.end(); it++) {
            ret += memusage::DynamicUsage(it->second.coins);
        }
        BOOST_CHECK_EQUAL(memusage::DynamicUsage(*this), ret);
    }

};

}

BOOST_AUTO_TEST_CASE(serials_test)
{
    CCoinsViewTest base;
    CCoinsViewCacheTest cache(&base);

    uint256 myserial = GetRandHash();

    BOOST_CHECK(!cache.GetSerial(myserial));
    cache.SetSerial(myserial, true);
    BOOST_CHECK(cache.GetSerial(myserial));
    cache.Flush();

    CCoinsViewCacheTest cache2(&base);

    BOOST_CHECK(cache2.GetSerial(myserial));
    cache2.SetSerial(myserial, false);
    BOOST_CHECK(!cache2.GetSerial(myserial));
    cache2.Flush();

    CCoinsViewCacheTest cache3(&base);

    BOOST_CHECK(!cache3.GetSerial(myserial));
}

void appendRandomCommitment(ZCIncrementalMerkleTree &tree)
{
    libzcash::SpendingKey k = libzcash::SpendingKey::random();
    libzcash::PaymentAddress addr = k.address();

    libzcash::Note note(addr.a_pk, 0, uint256(), uint256());

    tree.append(note.cm());
}

BOOST_AUTO_TEST_CASE(anchors_flush_test)
{
    CCoinsViewTest base;
    uint256 newrt;
    {
        CCoinsViewCacheTest cache(&base);
        ZCIncrementalMerkleTree tree;
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));
        appendRandomCommitment(tree);

        newrt = tree.root();

        cache.PushAnchor(tree);
        cache.Flush();
    }

    {
        CCoinsViewCacheTest cache(&base);
        ZCIncrementalMerkleTree tree;
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));

        // Get the cached entry.
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));

        uint256 check_rt = tree.root();

        BOOST_CHECK(check_rt == newrt);
    }
}

BOOST_AUTO_TEST_CASE(anchors_test)
{
    // TODO: These tests should be more methodical.
    //       Or, integrate with Bitcoin's tests later.

    CCoinsViewTest base;
    CCoinsViewCacheTest cache(&base);

    BOOST_CHECK(cache.GetBestAnchor() == uint256());

    {
        ZCIncrementalMerkleTree tree;

        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);

        ZCIncrementalMerkleTree save_tree_for_later;
        save_tree_for_later = tree;

        uint256 newrt = tree.root();
        uint256 newrt2;

        cache.PushAnchor(tree);
        BOOST_CHECK(cache.GetBestAnchor() == newrt);

        {
            ZCIncrementalMerkleTree confirm_same;
            BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), confirm_same));

            BOOST_CHECK(confirm_same.root() == newrt);
        }

        appendRandomCommitment(tree);
        appendRandomCommitment(tree);

        newrt2 = tree.root();

        cache.PushAnchor(tree);
        BOOST_CHECK(cache.GetBestAnchor() == newrt2);

        ZCIncrementalMerkleTree test_tree;
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), test_tree));

        BOOST_CHECK(tree.root() == test_tree.root());

        {
            ZCIncrementalMerkleTree test_tree2;
            cache.GetAnchorAt(newrt, test_tree2);

            BOOST_CHECK(test_tree2.root() == newrt);
        }

        {
            cache.PopAnchor(newrt);
            ZCIncrementalMerkleTree obtain_tree;
            assert(!cache.GetAnchorAt(newrt2, obtain_tree)); // should have been popped off
            assert(cache.GetAnchorAt(newrt, obtain_tree));

            assert(obtain_tree.root() == newrt);
        }
    }
}

BOOST_FIXTURE_TEST_SUITE(coins_tests, BasicTestingSetup)

static const unsigned int NUM_SIMULATION_ITERATIONS = 40000;

// This is a large randomized insert/remove simulation test on a variable-size
// stack of caches on top of CCoinsViewTest.
//
// It will randomly create/update/delete CCoins entries to a tip of caches, with
// txids picked from a limited list of random 256-bit hashes. Occasionally, a
// new tip is added to the stack of caches, or the tip is flushed and removed.
//
// During the process, booleans are kept to make sure that the randomized
// operation hits all branches.
BOOST_AUTO_TEST_CASE(coins_cache_simulation_test)
{
    // Various coverage trackers.
    bool removed_all_caches = false;
    bool reached_4_caches = false;
    bool added_an_entry = false;
    bool removed_an_entry = false;
    bool updated_an_entry = false;
    bool found_an_entry = false;
    bool missed_an_entry = false;

    // A simple map to track what we expect the cache stack to represent.
    std::map<uint256, CCoins> result;

    // The cache stack.
    CCoinsViewTest base; // A CCoinsViewTest at the bottom.
    std::vector<CCoinsViewCacheTest*> stack; // A stack of CCoinsViewCaches on top.
    stack.push_back(new CCoinsViewCacheTest(&base)); // Start with one cache.

    // Use a limited set of random transaction ids, so we do test overwriting entries.
    std::vector<uint256> txids;
    txids.resize(NUM_SIMULATION_ITERATIONS / 8);
    for (unsigned int i = 0; i < txids.size(); i++) {
        txids[i] = GetRandHash();
    }

    for (unsigned int i = 0; i < NUM_SIMULATION_ITERATIONS; i++) {
        // Do a random modification.
        {
            uint256 txid = txids[insecure_rand() % txids.size()]; // txid we're going to modify in this iteration.
            CCoins& coins = result[txid];
            CCoinsModifier entry = stack.back()->ModifyCoins(txid);
            BOOST_CHECK(coins == *entry);
            if (insecure_rand() % 5 == 0 || coins.IsPruned()) {
                if (coins.IsPruned()) {
                    added_an_entry = true;
                } else {
                    updated_an_entry = true;
                }
                coins.nVersion = insecure_rand();
                coins.vout.resize(1);
                coins.vout[0].nValue = insecure_rand();
                *entry = coins;
            } else {
                coins.Clear();
                entry->Clear();
                removed_an_entry = true;
            }
        }

        // Once every 1000 iterations and at the end, verify the full cache.
        if (insecure_rand() % 1000 == 1 || i == NUM_SIMULATION_ITERATIONS - 1) {
            for (std::map<uint256, CCoins>::iterator it = result.begin(); it != result.end(); it++) {
                const CCoins* coins = stack.back()->AccessCoins(it->first);
                if (coins) {
                    BOOST_CHECK(*coins == it->second);
                    found_an_entry = true;
                } else {
                    BOOST_CHECK(it->second.IsPruned());
                    missed_an_entry = true;
                }
            }
            BOOST_FOREACH(const CCoinsViewCacheTest *test, stack) {
                test->SelfTest();
            }
        }

        if (insecure_rand() % 100 == 0) {
            // Every 100 iterations, change the cache stack.
            if (stack.size() > 0 && insecure_rand() % 2 == 0) {
                stack.back()->Flush();
                delete stack.back();
                stack.pop_back();
            }
            if (stack.size() == 0 || (stack.size() < 4 && insecure_rand() % 2)) {
                CCoinsView* tip = &base;
                if (stack.size() > 0) {
                    tip = stack.back();
                } else {
                    removed_all_caches = true;
                }
                stack.push_back(new CCoinsViewCacheTest(tip));
                if (stack.size() == 4) {
                    reached_4_caches = true;
                }
            }
        }
    }

    // Clean up the stack.
    while (stack.size() > 0) {
        delete stack.back();
        stack.pop_back();
    }

    // Verify coverage.
    BOOST_CHECK(removed_all_caches);
    BOOST_CHECK(reached_4_caches);
    BOOST_CHECK(added_an_entry);
    BOOST_CHECK(removed_an_entry);
    BOOST_CHECK(updated_an_entry);
    BOOST_CHECK(found_an_entry);
    BOOST_CHECK(missed_an_entry);
}

BOOST_AUTO_TEST_SUITE_END()
<commit_msg>Move new coins tests to within coins_tests test suite.<commit_after>// Copyright (c) 2014 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "coins.h"
#include "random.h"
#include "uint256.h"
#include "test/test_bitcoin.h"

#include <vector>
#include <map>

#include <boost/test/unit_test.hpp>
#include "zcash/IncrementalMerkleTree.hpp"

namespace
{
class CCoinsViewTest : public CCoinsView
{
    uint256 hashBestBlock_;
    uint256 hashBestAnchor_;
    std::map<uint256, CCoins> map_;
    std::map<uint256, ZCIncrementalMerkleTree> mapAnchors_;
    std::map<uint256, bool> mapSerials_;

public:
    bool GetAnchorAt(const uint256& rt, ZCIncrementalMerkleTree &tree) const {
        if (rt.IsNull()) {
            ZCIncrementalMerkleTree new_tree;
            tree = new_tree;
            return true;
        }

        std::map<uint256, ZCIncrementalMerkleTree>::const_iterator it = mapAnchors_.find(rt);
        if (it == mapAnchors_.end()) {
            return false;
        } else {
            tree = it->second;
            return true;
        }
    }

    bool GetSerial(const uint256 &serial) const
    {
        std::map<uint256, bool>::const_iterator it = mapSerials_.find(serial);

        if (it == mapSerials_.end()) {
            return false;
        } else {
            // The map shouldn't contain any false entries.
            assert(it->second);
            return true;
        }
    }

    uint256 GetBestAnchor() const { return hashBestAnchor_; }

    bool GetCoins(const uint256& txid, CCoins& coins) const
    {
        std::map<uint256, CCoins>::const_iterator it = map_.find(txid);
        if (it == map_.end()) {
            return false;
        }
        coins = it->second;
        if (coins.IsPruned() && insecure_rand() % 2 == 0) {
            // Randomly return false in case of an empty entry.
            return false;
        }
        return true;
    }

    bool HaveCoins(const uint256& txid) const
    {
        CCoins coins;
        return GetCoins(txid, coins);
    }

    uint256 GetBestBlock() const { return hashBestBlock_; }

    bool BatchWrite(CCoinsMap& mapCoins,
                    const uint256& hashBlock,
                    const uint256& hashAnchor,
                    CAnchorsMap& mapAnchors,
                    CSerialsMap& mapSerials)
    {
        for (CCoinsMap::iterator it = mapCoins.begin(); it != mapCoins.end(); ) {
            map_[it->first] = it->second.coins;
            if (it->second.coins.IsPruned() && insecure_rand() % 3 == 0) {
                // Randomly delete empty entries on write.
                map_.erase(it->first);
            }
            mapCoins.erase(it++);
        }
        for (CAnchorsMap::iterator it = mapAnchors.begin(); it != mapAnchors.end(); ) {
            if (it->second.entered) {
                std::map<uint256, ZCIncrementalMerkleTree>::iterator ret =
                    mapAnchors_.insert(std::make_pair(it->first, ZCIncrementalMerkleTree())).first;

                ret->second = it->second.tree;
            } else {
                mapAnchors_.erase(it->first);
            }
            mapAnchors.erase(it++);
        }
        for (CSerialsMap::iterator it = mapSerials.begin(); it != mapSerials.end(); ) {
            if (it->second.entered) {
                mapSerials_[it->first] = true;
            } else {
                mapSerials_.erase(it->first);
            }
            mapSerials.erase(it++);
        }
        mapCoins.clear();
        mapAnchors.clear();
        mapSerials.clear();
        hashBestBlock_ = hashBlock;
        hashBestAnchor_ = hashAnchor;
        return true;
    }

    bool GetStats(CCoinsStats& stats) const { return false; }
};

class CCoinsViewCacheTest : public CCoinsViewCache
{
public:
    CCoinsViewCacheTest(CCoinsView* base) : CCoinsViewCache(base) {}

    void SelfTest() const
    {
        // Manually recompute the dynamic usage of the whole data, and compare it.
        size_t ret = memusage::DynamicUsage(cacheCoins);
        for (CCoinsMap::iterator it = cacheCoins.begin(); it != cacheCoins.end(); it++) {
            ret += memusage::DynamicUsage(it->second.coins);
        }
        BOOST_CHECK_EQUAL(memusage::DynamicUsage(*this), ret);
    }

};

}

void appendRandomCommitment(ZCIncrementalMerkleTree &tree)
{
    libzcash::SpendingKey k = libzcash::SpendingKey::random();
    libzcash::PaymentAddress addr = k.address();

    libzcash::Note note(addr.a_pk, 0, uint256(), uint256());

    tree.append(note.cm());
}

BOOST_FIXTURE_TEST_SUITE(coins_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(serials_test)
{
    CCoinsViewTest base;
    CCoinsViewCacheTest cache(&base);

    uint256 myserial = GetRandHash();

    BOOST_CHECK(!cache.GetSerial(myserial));
    cache.SetSerial(myserial, true);
    BOOST_CHECK(cache.GetSerial(myserial));
    cache.Flush();

    CCoinsViewCacheTest cache2(&base);

    BOOST_CHECK(cache2.GetSerial(myserial));
    cache2.SetSerial(myserial, false);
    BOOST_CHECK(!cache2.GetSerial(myserial));
    cache2.Flush();

    CCoinsViewCacheTest cache3(&base);

    BOOST_CHECK(!cache3.GetSerial(myserial));
}

BOOST_AUTO_TEST_CASE(anchors_flush_test)
{
    CCoinsViewTest base;
    uint256 newrt;
    {
        CCoinsViewCacheTest cache(&base);
        ZCIncrementalMerkleTree tree;
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));
        appendRandomCommitment(tree);

        newrt = tree.root();

        cache.PushAnchor(tree);
        cache.Flush();
    }

    {
        CCoinsViewCacheTest cache(&base);
        ZCIncrementalMerkleTree tree;
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));

        // Get the cached entry.
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));

        uint256 check_rt = tree.root();

        BOOST_CHECK(check_rt == newrt);
    }
}

BOOST_AUTO_TEST_CASE(anchors_test)
{
    // TODO: These tests should be more methodical.
    //       Or, integrate with Bitcoin's tests later.

    CCoinsViewTest base;
    CCoinsViewCacheTest cache(&base);

    BOOST_CHECK(cache.GetBestAnchor() == uint256());

    {
        ZCIncrementalMerkleTree tree;

        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), tree));
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);
        appendRandomCommitment(tree);

        ZCIncrementalMerkleTree save_tree_for_later;
        save_tree_for_later = tree;

        uint256 newrt = tree.root();
        uint256 newrt2;

        cache.PushAnchor(tree);
        BOOST_CHECK(cache.GetBestAnchor() == newrt);

        {
            ZCIncrementalMerkleTree confirm_same;
            BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), confirm_same));

            BOOST_CHECK(confirm_same.root() == newrt);
        }

        appendRandomCommitment(tree);
        appendRandomCommitment(tree);

        newrt2 = tree.root();

        cache.PushAnchor(tree);
        BOOST_CHECK(cache.GetBestAnchor() == newrt2);

        ZCIncrementalMerkleTree test_tree;
        BOOST_CHECK(cache.GetAnchorAt(cache.GetBestAnchor(), test_tree));

        BOOST_CHECK(tree.root() == test_tree.root());

        {
            ZCIncrementalMerkleTree test_tree2;
            cache.GetAnchorAt(newrt, test_tree2);

            BOOST_CHECK(test_tree2.root() == newrt);
        }

        {
            cache.PopAnchor(newrt);
            ZCIncrementalMerkleTree obtain_tree;
            assert(!cache.GetAnchorAt(newrt2, obtain_tree)); // should have been popped off
            assert(cache.GetAnchorAt(newrt, obtain_tree));

            assert(obtain_tree.root() == newrt);
        }
    }
}

static const unsigned int NUM_SIMULATION_ITERATIONS = 40000;

// This is a large randomized insert/remove simulation test on a variable-size
// stack of caches on top of CCoinsViewTest.
//
// It will randomly create/update/delete CCoins entries to a tip of caches, with
// txids picked from a limited list of random 256-bit hashes. Occasionally, a
// new tip is added to the stack of caches, or the tip is flushed and removed.
//
// During the process, booleans are kept to make sure that the randomized
// operation hits all branches.
BOOST_AUTO_TEST_CASE(coins_cache_simulation_test)
{
    // Various coverage trackers.
    bool removed_all_caches = false;
    bool reached_4_caches = false;
    bool added_an_entry = false;
    bool removed_an_entry = false;
    bool updated_an_entry = false;
    bool found_an_entry = false;
    bool missed_an_entry = false;

    // A simple map to track what we expect the cache stack to represent.
    std::map<uint256, CCoins> result;

    // The cache stack.
    CCoinsViewTest base; // A CCoinsViewTest at the bottom.
    std::vector<CCoinsViewCacheTest*> stack; // A stack of CCoinsViewCaches on top.
    stack.push_back(new CCoinsViewCacheTest(&base)); // Start with one cache.

    // Use a limited set of random transaction ids, so we do test overwriting entries.
    std::vector<uint256> txids;
    txids.resize(NUM_SIMULATION_ITERATIONS / 8);
    for (unsigned int i = 0; i < txids.size(); i++) {
        txids[i] = GetRandHash();
    }

    for (unsigned int i = 0; i < NUM_SIMULATION_ITERATIONS; i++) {
        // Do a random modification.
        {
            uint256 txid = txids[insecure_rand() % txids.size()]; // txid we're going to modify in this iteration.
            CCoins& coins = result[txid];
            CCoinsModifier entry = stack.back()->ModifyCoins(txid);
            BOOST_CHECK(coins == *entry);
            if (insecure_rand() % 5 == 0 || coins.IsPruned()) {
                if (coins.IsPruned()) {
                    added_an_entry = true;
                } else {
                    updated_an_entry = true;
                }
                coins.nVersion = insecure_rand();
                coins.vout.resize(1);
                coins.vout[0].nValue = insecure_rand();
                *entry = coins;
            } else {
                coins.Clear();
                entry->Clear();
                removed_an_entry = true;
            }
        }

        // Once every 1000 iterations and at the end, verify the full cache.
        if (insecure_rand() % 1000 == 1 || i == NUM_SIMULATION_ITERATIONS - 1) {
            for (std::map<uint256, CCoins>::iterator it = result.begin(); it != result.end(); it++) {
                const CCoins* coins = stack.back()->AccessCoins(it->first);
                if (coins) {
                    BOOST_CHECK(*coins == it->second);
                    found_an_entry = true;
                } else {
                    BOOST_CHECK(it->second.IsPruned());
                    missed_an_entry = true;
                }
            }
            BOOST_FOREACH(const CCoinsViewCacheTest *test, stack) {
                test->SelfTest();
            }
        }

        if (insecure_rand() % 100 == 0) {
            // Every 100 iterations, change the cache stack.
            if (stack.size() > 0 && insecure_rand() % 2 == 0) {
                stack.back()->Flush();
                delete stack.back();
                stack.pop_back();
            }
            if (stack.size() == 0 || (stack.size() < 4 && insecure_rand() % 2)) {
                CCoinsView* tip = &base;
                if (stack.size() > 0) {
                    tip = stack.back();
                } else {
                    removed_all_caches = true;
                }
                stack.push_back(new CCoinsViewCacheTest(tip));
                if (stack.size() == 4) {
                    reached_4_caches = true;
                }
            }
        }
    }

    // Clean up the stack.
    while (stack.size() > 0) {
        delete stack.back();
        stack.pop_back();
    }

    // Verify coverage.
    BOOST_CHECK(removed_all_caches);
    BOOST_CHECK(reached_4_caches);
    BOOST_CHECK(added_an_entry);
    BOOST_CHECK(removed_an_entry);
    BOOST_CHECK(updated_an_entry);
    BOOST_CHECK(found_an_entry);
    BOOST_CHECK(missed_an_entry);
}

BOOST_AUTO_TEST_SUITE_END()
<|endoftext|>


<commit_before>/* -*- mode: c++; c-basic-offset:4 -*-
    commands/changepassphrasecommand.cpp

    This file is part of Kleopatra, the KDE keymanager
    Copyright (c) 2008 Klarälvdalens Datakonsult AB

    Kleopatra is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    Kleopatra is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

    In addition, as a special exception, the copyright holders give
    permission to link the code of this program with any edition of
    the Qt library by Trolltech AS, Norway (or with modified versions
    of Qt that use the same license as Qt), and distribute linked
    combinations including the two.  You must obey the GNU General
    Public License in all respects for all of the code used other than
    Qt.  If you modify this file, you may extend this exception to
    your version of the file, but you are not obligated to do so.  If
    you do not wish to do so, delete this exception statement from
    your version.
*/

#include <config-kleopatra.h>

#include "changepassphrasecommand.h"

#include "command_p.h"

#include <utils/gnupg-helper.h>

#include <gpgme++/key.h>

#include <KLocale>
#include <KMessageBox>

using namespace Kleo;
using namespace Kleo::Commands;
using namespace GpgME;

ChangePassphraseCommand::ChangePassphraseCommand( KeyListController * c )
    : GnuPGProcessCommand( c )
{

}

ChangePassphraseCommand::ChangePassphraseCommand( QAbstractItemView * v, KeyListController * c )
    : GnuPGProcessCommand( v, c )
{

}

ChangePassphraseCommand::ChangePassphraseCommand( const Key & key )
    : GnuPGProcessCommand( key )
{

}

ChangePassphraseCommand::~ChangePassphraseCommand() {}

QStringList ChangePassphraseCommand::arguments() const {
    const Key key = d->key();
    if ( key.protocol() == OpenPGP )
        return QStringList() << gpgPath() << "--edit-key" << key.primaryFingerprint() << "passwd";
    else
        return QStringList() << gpgSmPath() << "--passwd" << key.primaryFingerprint();
}

QString ChangePassphraseCommand::errorCaption() const {
    return i18n( "Passphrase Change Error" );
}

QString ChangePassphraseCommand::successCaption() const {
    return i18n( "Passphrase Change Finished" );
}

QString ChangePassphraseCommand::crashExitMessage( const QStringList & args ) const {
    return i18nc("@info",
                 "<para>The GPG or GpgSM process that tried to change the passphrase "
                 "ended prematurely because of an unexpected error.</para>"
                 "<para>Please check the output of <icode>%1</icode> for details.</para>", args.join( " " ) ) ;
}

QString ChangePassphraseCommand::errorExitMessage( const QStringList & args ) const {
    return i18nc("@info",
                 "<para>An error occurred while trying to change the passphrase.</para> "
                 "<para>The output from <command>%1</command> was: <message>%2</message></para>",
                 args[0], errorString() );
}

QString ChangePassphraseCommand::successMessage( const QStringList & ) const {
    return i18n( "Passphrase changed successfully." );
}

#include "moc_changepassphrasecommand.cpp"
<commit_msg>add --batch to enforce non-interactive editing<commit_after>/* -*- mode: c++; c-basic-offset:4 -*-
    commands/changepassphrasecommand.cpp

    This file is part of Kleopatra, the KDE keymanager
    Copyright (c) 2008 Klarälvdalens Datakonsult AB

    Kleopatra is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    Kleopatra is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

    In addition, as a special exception, the copyright holders give
    permission to link the code of this program with any edition of
    the Qt library by Trolltech AS, Norway (or with modified versions
    of Qt that use the same license as Qt), and distribute linked
    combinations including the two.  You must obey the GNU General
    Public License in all respects for all of the code used other than
    Qt.  If you modify this file, you may extend this exception to
    your version of the file, but you are not obligated to do so.  If
    you do not wish to do so, delete this exception statement from
    your version.
*/

#include <config-kleopatra.h>

#include "changepassphrasecommand.h"

#include "command_p.h"

#include <utils/gnupg-helper.h>

#include <gpgme++/key.h>

#include <KLocale>
#include <KMessageBox>

using namespace Kleo;
using namespace Kleo::Commands;
using namespace GpgME;

ChangePassphraseCommand::ChangePassphraseCommand( KeyListController * c )
    : GnuPGProcessCommand( c )
{

}

ChangePassphraseCommand::ChangePassphraseCommand( QAbstractItemView * v, KeyListController * c )
    : GnuPGProcessCommand( v, c )
{

}

ChangePassphraseCommand::ChangePassphraseCommand( const Key & key )
    : GnuPGProcessCommand( key )
{

}

ChangePassphraseCommand::~ChangePassphraseCommand() {}

QStringList ChangePassphraseCommand::arguments() const {
    const Key key = d->key();
    if ( key.protocol() == OpenPGP )
        return QStringList() << gpgPath() << << "--batch" << "--edit-key" << key.primaryFingerprint() << "passwd";
    else
        return QStringList() << gpgSmPath() << "--passwd" << key.primaryFingerprint();
}

QString ChangePassphraseCommand::errorCaption() const {
    return i18n( "Passphrase Change Error" );
}

QString ChangePassphraseCommand::successCaption() const {
    return i18n( "Passphrase Change Finished" );
}

QString ChangePassphraseCommand::crashExitMessage( const QStringList & args ) const {
    return i18nc("@info",
                 "<para>The GPG or GpgSM process that tried to change the passphrase "
                 "ended prematurely because of an unexpected error.</para>"
                 "<para>Please check the output of <icode>%1</icode> for details.</para>", args.join( " " ) ) ;
}

QString ChangePassphraseCommand::errorExitMessage( const QStringList & args ) const {
    return i18nc("@info",
                 "<para>An error occurred while trying to change the passphrase.</para> "
                 "<para>The output from <command>%1</command> was: <message>%2</message></para>",
                 args[0], errorString() );
}

QString ChangePassphraseCommand::successMessage( const QStringList & ) const {
    return i18n( "Passphrase changed successfully." );
}

#include "moc_changepassphrasecommand.cpp"
<|endoftext|>


<commit_before>CheckCalibStatus(const Char_t *fileName)
{

  TFile *file = TFile::Open(fileName);
  AliCDBEntry *cdbe = (AliCDBEntry *)file->Get("AliCDBEntry");
  AliTOFChannelOnlineStatusArray *array = (AliTOFChannelOnlineStatusArray *)cdbe->GetObject();

  TH1F *h = new TH1F("h", "Channel status;index;status", array->GetSize(), 0., array->GetSize(););
  for (Int_t i = 0; i <  array->GetSize(); i++)
    h->SetBinContent(i + 1, array->GetStatus(i));

  h->Draw();

}
<commit_msg>added to CheckCalibStatus macro two histograms in order to display, on the basis of detector configuration stored in OCDB, enabled and noisy channels in a user-friendly fashion<commit_after>CheckCalibStatus(const Char_t *fileName)
{

  TFile *file = TFile::Open(fileName);
  AliCDBEntry *cdbe = (AliCDBEntry *)file->Get("AliCDBEntry");
  AliTOFChannelOnlineStatusArray *array = (AliTOFChannelOnlineStatusArray *)cdbe->GetObject();

  TH1F *hStatus = new TH1F("hStatus", "Channel status;index;status", array->GetSize(), 0., array->GetSize(););
  TH2F *hNoiseMap = new TH2F("hNoiseMap", "Noise map;sector;strip", 72, 0., 18., 91, 0., 91.);
  TH2F *hEnableMap = new TH2F("hEnableMap", "Enable map;sector;strip", 72, 0., 18., 91, 0., 91.);

  AliTOFcalibHisto calib;
  calib.LoadCalibHisto();

  Int_t sector, sectorStrip, padx, fea;
  Float_t hitmapx, hitmapy;
  for (Int_t i = 0; i <  array->GetSize(); i++) {
    hStatus->SetBinContent(i + 1, array->GetStatus(i));
    sector = calib.GetCalibMap(AliTOFcalibHisto::kSector, i);
    sectorStrip = calib.GetCalibMap(AliTOFcalibHisto::kSectorStrip, i);
    padx = calib.GetCalibMap(AliTOFcalibHisto::kPadX, i);
    fea = padx / 12;
    hitmapx = sector + ((Double_t)(3 - fea) + 0.5) / 4.;
    hitmapy = sectorStrip;
    if (array->GetHWStatus(i) == AliTOFChannelOnlineStatusArray::kTOFHWOk) hEnableMap->Fill(hitmapx, hitmapy);
    if (array->GetNoiseStatus(i) == AliTOFChannelOnlineStatusArray::kTOFNoiseBad) hNoiseMap->Fill(hitmapx, hitmapy);
  }

  TFile *fout = TFile::Open("CheckCalibStatus.root", "RECREATE");
  hStatus->Write();
  hNoiseMap->Write();
  hEnableMap->Write();
  fout->Close();

}
<|endoftext|>


<commit_before>/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/kernels/sendrecv_ops.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"

namespace tensorflow {

static string GetRendezvousKeyPrefix(const string& send_device,
                                     const string& recv_device,
                                     const uint64 send_device_incarnation,
                                     const string& tensor_name) {
  return strings::StrCat(send_device, ";",
                         strings::FpToString(send_device_incarnation), ";",
                         recv_device, ";", tensor_name);
}

static void GetRendezvousKey(const string& key_prefix,
                             const FrameAndIter& frame_iter, string* key) {
  key->clear();
  strings::StrAppend(key, key_prefix, ";", frame_iter.frame_id, ":",
                     frame_iter.iter_id);
}

static FrameAndIter GetFrameAndIter(OpKernelContext* ctx,
                                    bool hostmem_sendrecv) {
  if (hostmem_sendrecv && ctx->call_frame() != nullptr) {
    // Host memory send/recv pairs are added by
    // common_runtime/memory_types.cc.  When the pair of nodes are
    // added inside a function, we need to use the function call frame
    // to formulate the unique rendezvous key.
    return FrameAndIter(reinterpret_cast<uint64>(ctx->call_frame()), 0);
  } else {
    return ctx->frame_iter();
  }
}

SendOp::SendOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
  string send_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device));
  string recv_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device));
  uint64 send_device_incarnation;
  OP_REQUIRES_OK(
      ctx, ctx->GetAttr("send_device_incarnation",
                        reinterpret_cast<int64*>(&send_device_incarnation)));
  string tensor_name;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name));
  key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device,
                                       send_device_incarnation, tensor_name);
  // The vast majority of Send nodes are outside any loop context, so
  // proactively cache the rendezvous key for the top-level.
  GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_);
  OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_));
  if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) {
    hostmem_sendrecv_ = false;
  }
}

void SendOp::Compute(OpKernelContext* ctx) {
  OP_REQUIRES(
      ctx, ctx->rendezvous() != nullptr,
      errors::Internal("Op kernel context needs to provide a rendezvous."));

  // The device context may be passed between the Send/Recv
  // boundary, so that the device context used to produce the Tensor
  // is used when performing the copy on the recv side (which may be
  // a different device).
  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->input_alloc_attr(0);

  FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
  if (frame_iter == FrameAndIter(0, 0)) {
    // Use the cached rendezvous key.
    VLOG(2) << "Send " << parsed_key_.buf_;
    ctx->SetStatus(ctx->rendezvous()->Send(parsed_key_, args, ctx->input(0),
                                           ctx->is_input_dead()));
    return;
  } else {
    Rendezvous::ParsedKey in_loop_parsed;
    GetRendezvousKey(key_prefix_, frame_iter, &in_loop_parsed.buf_);
    VLOG(2) << "Send " << in_loop_parsed.buf_;
    OP_REQUIRES_OK(ctx,
                   Rendezvous::ParseKey(in_loop_parsed.buf_, &in_loop_parsed));

    ctx->SetStatus(ctx->rendezvous()->Send(in_loop_parsed, args, ctx->input(0),
                                           ctx->is_input_dead()));
    return;
  }
}

REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_CPU), SendOp);
REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_GPU), SendOp);

#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_SYCL), SendOp);
REGISTER_KERNEL_BUILDER(
    Name("_HostSend").Device(DEVICE_SYCL).HostMemory("tensor"), SendOp);
#endif  // TENSORFLOW_USE_SYCL

REGISTER_KERNEL_BUILDER(Name("_HostSend").Device(DEVICE_CPU), SendOp);
REGISTER_KERNEL_BUILDER(
    Name("_HostSend").Device(DEVICE_GPU).HostMemory("tensor"), SendOp);

RecvOp::RecvOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {
  string send_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device));
  string recv_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device));
  uint64 send_device_incarnation;
  OP_REQUIRES_OK(
      ctx, ctx->GetAttr("send_device_incarnation",
                        reinterpret_cast<int64*>(&send_device_incarnation)));
  string tensor_name;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name));
  key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device,
                                       send_device_incarnation, tensor_name);
  // The vast majority of Recv nodes are outside any loop context, so
  // proactively cache the rendezvous key for the top-level.
  GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_);
  OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_));
  if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) {
    hostmem_sendrecv_ = false;
  }
}

namespace {
Rendezvous::DoneCallback make_recv_callback(OpKernelContext* ctx,
                                            AsyncOpKernel::DoneCallback done) {
  using namespace std::placeholders;
  return std::bind(
      [ctx](AsyncOpKernel::DoneCallback done,
            // Begin unbound arguments.
            const Status& s, const Rendezvous::Args& send_args,
            const Rendezvous::Args& recv_args, const Tensor& val,
            bool is_dead) {
        ctx->SetStatus(s);
        if (s.ok()) {
          // 'ctx' allocates the output tensor of the expected type.
          // The runtime checks whether the tensor received here is
          // the same type.
          if (!is_dead) {
            ctx->set_output(0, val);
          }
        }
        done();
      },
      std::move(done), _1, _2, _3, _4, _5);
}
}  // namespace

void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  OP_REQUIRES_ASYNC(
      ctx, ctx->rendezvous() != nullptr,
      errors::Internal("Op kernel context needs to provide a rendezvous."),
      done);

  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->output_alloc_attr(0);

  FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
  if (frame_iter == FrameAndIter(0, 0)) {
    VLOG(2) << "Recv " << parsed_key_.buf_;
    ctx->rendezvous()->RecvAsync(parsed_key_, args,
                                 make_recv_callback(ctx, std::move(done)));
  } else {
    Rendezvous::ParsedKey in_loop_parsed;
    GetRendezvousKey(key_prefix_, frame_iter, &in_loop_parsed.buf_);
    VLOG(2) << "Recv " << in_loop_parsed.buf_;
    OP_REQUIRES_OK_ASYNC(
        ctx, Rendezvous::ParseKey(in_loop_parsed.buf_, &in_loop_parsed), done);
    ctx->rendezvous()->RecvAsync(in_loop_parsed, args,
                                 make_recv_callback(ctx, std::move(done)));
  }
}

REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_CPU), RecvOp);
REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_GPU), RecvOp);

#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_SYCL), RecvOp);
#endif  // TENSORFLOW_USE_SYCL

REGISTER_KERNEL_BUILDER(Name("_HostRecv").Device(DEVICE_CPU), RecvOp);
REGISTER_KERNEL_BUILDER(
    Name("_HostRecv").Device(DEVICE_GPU).HostMemory("tensor"), RecvOp);

#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(
    Name("_HostRecv").Device(DEVICE_SYCL).HostMemory("tensor"), RecvOp);
#endif  // TENSORFLOW_USE_SYCL

}  // end namespace tensorflow
<commit_msg>Code refactor.<commit_after>/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/kernels/sendrecv_ops.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"

namespace tensorflow {

static string GetRendezvousKeyPrefix(const string& send_device,
                                     const string& recv_device,
                                     const uint64 send_device_incarnation,
                                     const string& tensor_name) {
  return strings::StrCat(send_device, ";",
                         strings::FpToString(send_device_incarnation), ";",
                         recv_device, ";", tensor_name);
}

static void GetRendezvousKey(const string& key_prefix,
                             const FrameAndIter& frame_iter, string* key) {
  key->clear();
  strings::StrAppend(key, key_prefix, ";", frame_iter.frame_id, ":",
                     frame_iter.iter_id);
}

static FrameAndIter GetFrameAndIter(OpKernelContext* ctx,
                                    bool hostmem_sendrecv) {
  if (hostmem_sendrecv && ctx->call_frame() != nullptr) {
    // Host memory send/recv pairs are added by
    // common_runtime/memory_types.cc.  When the pair of nodes are
    // added inside a function, we need to use the function call frame
    // to formulate the unique rendezvous key.
    return FrameAndIter(reinterpret_cast<uint64>(ctx->call_frame()), 0);
  } else {
    return ctx->frame_iter();
  }
}

SendOp::SendOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
  string send_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device));
  string recv_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device));
  uint64 send_device_incarnation;
  OP_REQUIRES_OK(
      ctx, ctx->GetAttr("send_device_incarnation",
                        reinterpret_cast<int64*>(&send_device_incarnation)));
  string tensor_name;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name));
  key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device,
                                       send_device_incarnation, tensor_name);
  // The vast majority of Send nodes are outside any loop context, so
  // proactively cache the rendezvous key for the top-level.
  GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_);
  OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_));
  if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) {
    hostmem_sendrecv_ = false;
  }
}

void SendOp::Compute(OpKernelContext* ctx) {
  OP_REQUIRES(
      ctx, ctx->rendezvous() != nullptr,
      errors::Internal("Op kernel context needs to provide a rendezvous."));

  // The device context may be passed between the Send/Recv
  // boundary, so that the device context used to produce the Tensor
  // is used when performing the copy on the recv side (which may be
  // a different device).
  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->input_alloc_attr(0);

  FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
  if (frame_iter == FrameAndIter(0, 0)) {
    // Use the cached rendezvous key.
    VLOG(2) << "Send " << parsed_key_.buf_;
    ctx->SetStatus(ctx->rendezvous()->Send(parsed_key_, args, ctx->input(0),
                                           ctx->is_input_dead()));
    return;
  } else {
    Rendezvous::ParsedKey in_loop_parsed;
    GetRendezvousKey(key_prefix_, frame_iter, &in_loop_parsed.buf_);
    VLOG(2) << "Send " << in_loop_parsed.buf_;
    OP_REQUIRES_OK(ctx,
                   Rendezvous::ParseKey(in_loop_parsed.buf_, &in_loop_parsed));

    ctx->SetStatus(ctx->rendezvous()->Send(in_loop_parsed, args, ctx->input(0),
                                           ctx->is_input_dead()));
    return;
  }
}

REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_CPU), SendOp);
REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_GPU), SendOp);

#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_SYCL), SendOp);
REGISTER_KERNEL_BUILDER(
    Name("_HostSend").Device(DEVICE_SYCL).HostMemory("tensor"), SendOp);
#endif  // TENSORFLOW_USE_SYCL

REGISTER_KERNEL_BUILDER(
    Name("_HostSend").Device(DEVICE_GPU).HostMemory("tensor"), SendOp);

RecvOp::RecvOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {
  string send_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device));
  string recv_device;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device));
  uint64 send_device_incarnation;
  OP_REQUIRES_OK(
      ctx, ctx->GetAttr("send_device_incarnation",
                        reinterpret_cast<int64*>(&send_device_incarnation)));
  string tensor_name;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name));
  key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device,
                                       send_device_incarnation, tensor_name);
  // The vast majority of Recv nodes are outside any loop context, so
  // proactively cache the rendezvous key for the top-level.
  GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_);
  OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_));
  if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) {
    hostmem_sendrecv_ = false;
  }
}

namespace {
Rendezvous::DoneCallback make_recv_callback(OpKernelContext* ctx,
                                            AsyncOpKernel::DoneCallback done) {
  using namespace std::placeholders;
  return std::bind(
      [ctx](AsyncOpKernel::DoneCallback done,
            // Begin unbound arguments.
            const Status& s, const Rendezvous::Args& send_args,
            const Rendezvous::Args& recv_args, const Tensor& val,
            bool is_dead) {
        ctx->SetStatus(s);
        if (s.ok()) {
          // 'ctx' allocates the output tensor of the expected type.
          // The runtime checks whether the tensor received here is
          // the same type.
          if (!is_dead) {
            ctx->set_output(0, val);
          }
        }
        done();
      },
      std::move(done), _1, _2, _3, _4, _5);
}
}  // namespace

void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  OP_REQUIRES_ASYNC(
      ctx, ctx->rendezvous() != nullptr,
      errors::Internal("Op kernel context needs to provide a rendezvous."),
      done);

  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->output_alloc_attr(0);

  FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
  if (frame_iter == FrameAndIter(0, 0)) {
    VLOG(2) << "Recv " << parsed_key_.buf_;
    ctx->rendezvous()->RecvAsync(parsed_key_, args,
                                 make_recv_callback(ctx, std::move(done)));
  } else {
    Rendezvous::ParsedKey in_loop_parsed;
    GetRendezvousKey(key_prefix_, frame_iter, &in_loop_parsed.buf_);
    VLOG(2) << "Recv " << in_loop_parsed.buf_;
    OP_REQUIRES_OK_ASYNC(
        ctx, Rendezvous::ParseKey(in_loop_parsed.buf_, &in_loop_parsed), done);
    ctx->rendezvous()->RecvAsync(in_loop_parsed, args,
                                 make_recv_callback(ctx, std::move(done)));
  }
}

REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_CPU), RecvOp);
REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_GPU), RecvOp);

#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_SYCL), RecvOp);
#endif  // TENSORFLOW_USE_SYCL

REGISTER_KERNEL_BUILDER(
    Name("_HostRecv").Device(DEVICE_GPU).HostMemory("tensor"), RecvOp);

#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(
    Name("_HostRecv").Device(DEVICE_SYCL).HostMemory("tensor"), RecvOp);
#endif  // TENSORFLOW_USE_SYCL

// Environment variable `DISABLE_HOST_SEND_RECV_REGISTRATION` is used to disable
// hostSend and hostRecv registration on CPU device in the mock environment.
static bool InitModule() {
  if (!std::getenv("DISABLE_HOST_SEND_RECV_REGISTRATION")) {
    REGISTER_KERNEL_BUILDER(Name("_HostRecv").Device(DEVICE_CPU), RecvOp);
    REGISTER_KERNEL_BUILDER(Name("_HostSend").Device(DEVICE_CPU), SendOp);
  }
  return true;
}

static bool module_initialized = InitModule();

}  // end namespace tensorflow
<|endoftext|>