<commit_before>#include "../aq_base.h"

#include <sstream>
#include <string>

#include "gtest/gtest.h"
#include "../../test_helpers/bart_test_helper.h"

class AQBaseTest : public ::testing::Test {
 protected:
  void SetUp() override;

  template<int dim>
  void OutputAQ();

  template<int dim>
  void InitRefBCAndCheck();

  dealii::ParameterHandler prm;
};

void AQBaseTest::SetUp() {
  prm.declare_entry("have reflective boundary", "false",
                     dealii::Patterns::Bool(), "");
  prm.declare_entry("angular quadrature order", "4",
                     dealii::Patterns::Integer(), "");
  prm.declare_entry("angular quadrature name", "gl",
                     dealii::Patterns::Selection("gl"), "");
  prm.declare_entry("number of groups", "1", dealii::Patterns::Integer(), "");
  prm.declare_entry("transport model", "regular",
                     dealii::Patterns::Selection("regular|ep"), "");
}

template<int dim>
void AQBaseTest::OutputAQ() {
  // Outputs AQData generated by MakeAQ to the deallog
  std::unique_ptr<AQBase<dim>> gl_ptr = std::make_unique<AQBase<dim>>(prm);

  gl_ptr->ProduceAQ();
  auto wi = gl_ptr->GetAQWeights();
  auto omega_i = gl_ptr->GetAQDirs();

  for (size_t i=0; i<wi.size(); ++i) {
    dealii::deallog << "Weight: " << wi[i] << "; Omega: ";
    for (int j=0; j<dim; ++j)
      dealii::deallog << omega_i[i][j] << " ";
    dealii::deallog << std::endl;
  }
}

template<int dim>
void AQBaseTest::InitRefBCAndCheck() {
  std::unique_ptr<AQBase<dim>> aq_base_ptr = std::make_unique<AQBase<dim>>(prm);
  aq_base_ptr->MakeAQ();
  // Get omegas and reflection map
  std::vector<dealii::Tensor<1, dim>> omegas = aq_base_ptr->GetAQDirs();
  std::map<std::pair<int, int>, int > reflection_map =
      aq_base_ptr->GetRefDirInd();

  for (auto const& mapping : reflection_map) {
    // Unroll this data structure
    int direction = mapping.first.second;
    int reflection = mapping.second;

    //x-direction
    EXPECT_FLOAT_EQ(omegas[direction][0], -omegas[reflection][0]);
  }
}

TEST_F(AQBaseTest, AQBase1DProduceAQ) {
  std::string filename = "aq_base_1d";
  btest::GoldTestInit(filename);
  OutputAQ<1>();
  btest::GoldTestRun(filename);
}

TEST_F(AQBaseTest, AQBase1DEpProduceAQ) {
  std::string filename = "aq_base_1d_ep";
  prm.set("transport model", "ep");
  btest::GoldTestInit(filename); // Opens deal log
  OutputAQ<1>();
  btest::GoldTestRun(filename); // Closes deal log
}

TEST_F(AQBaseTest, AQBaseBadDimProduceAQ) {
  AQBase<2> test_AQ(prm);
  ASSERT_THROW(test_AQ.ProduceAQ(), dealii::ExcMessage);
}

TEST_F(AQBaseTest, AQBaseBadDimReflectiveBC) {
  prm.set("have reflective boundary", "true");
  InitRefBCAndCheck<1>();
}

TEST_F(AQBaseTest, AQBaseSNOrderGetter) {
  AQBase<1> test_AQ(prm);
  ASSERT_EQ(test_AQ.GetSnOrder(), 4);
}

TEST_F(AQBaseTest, AQBaseGetNDir) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNDir(), 4);
}

TEST_F(AQBaseTest, AQBaseGetNDirEp) {
  prm.set("transport model", "ep");
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNDir(), 2);
}

TEST_F(AQBaseTest, AQBaseGetTotalHOVars) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNTotalHOVars(), 4);
}

TEST_F(AQBaseTest, AQBaseGetTotalHOVarsEq) {
  prm.set("transport model", "ep");
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNTotalHOVars(), 2);
}

TEST_F(AQBaseTest, PrintAQ) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();

  std::ostringstream output_string_stream;
  test_AQ.PrintAQ(&output_string_stream);

  std::string output = "transport model: regular; output_streamature name: None\n"
                       "Dim = 1, SN order = 4\n"
                       "Weights | Omega_x | Omega_y | mu\n"
                       "0.347854845137454  -0.861136311594053  \n"
                       "0.652145154862546  -0.339981043584856  \n"
                       "0.652145154862546  0.339981043584856  \n"
                       "0.347854845137454  0.861136311594052  \n";
  EXPECT_EQ(output, output_string_stream.str());
}

TEST_F(AQBaseTest, RefDirInt) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  std::map<std::pair<int, int>, int> component_index_map(test_AQ.GetCompInd());

  for (auto const& mapping : component_index_map) {
    EXPECT_EQ(mapping.first.first, 0);
    EXPECT_EQ(mapping.first.second, mapping.second);
  }
}

TEST_F(AQBaseTest, InvDirInd) {
  AQBase<1> test_AQ(prm);
  std::unordered_map<int, std::pair<int, int>>
      inv_component_map(test_AQ.GetInvCompInd());

  for (size_t i = 0; i < inv_component_map.size(); ++i) {
    EXPECT_EQ(inv_component_map[i].first, 0);
    EXPECT_EQ(inv_component_map[i].second, i);
  }
}
<commit_msg>updated PrintAQ test to use a regex for float comparison<commit_after>#include "../aq_base.h"

#include <sstream>
#include <string>

#include "gtest/gtest.h"
#include "../../test_helpers/gmock_wrapper.h"
#include "../../test_helpers/bart_test_helper.h"

class AQBaseTest : public ::testing::Test {
 protected:
  void SetUp() override;

  template<int dim>
  void OutputAQ();

  template<int dim>
  void InitRefBCAndCheck();

  dealii::ParameterHandler prm;
};

void AQBaseTest::SetUp() {
  prm.declare_entry("have reflective boundary", "false",
                     dealii::Patterns::Bool(), "");
  prm.declare_entry("angular quadrature order", "4",
                     dealii::Patterns::Integer(), "");
  prm.declare_entry("angular quadrature name", "gl",
                     dealii::Patterns::Selection("gl"), "");
  prm.declare_entry("number of groups", "1", dealii::Patterns::Integer(), "");
  prm.declare_entry("transport model", "regular",
                     dealii::Patterns::Selection("regular|ep"), "");
}

template<int dim>
void AQBaseTest::OutputAQ() {
  // Outputs AQData generated by MakeAQ to the deallog
  std::unique_ptr<AQBase<dim>> gl_ptr = std::make_unique<AQBase<dim>>(prm);

  gl_ptr->ProduceAQ();
  auto wi = gl_ptr->GetAQWeights();
  auto omega_i = gl_ptr->GetAQDirs();

  for (size_t i=0; i<wi.size(); ++i) {
    dealii::deallog << "Weight: " << wi[i] << "; Omega: ";
    for (int j=0; j<dim; ++j)
      dealii::deallog << omega_i[i][j] << " ";
    dealii::deallog << std::endl;
  }
}

template<int dim>
void AQBaseTest::InitRefBCAndCheck() {
  std::unique_ptr<AQBase<dim>> aq_base_ptr = std::make_unique<AQBase<dim>>(prm);
  aq_base_ptr->MakeAQ();
  // Get omegas and reflection map
  std::vector<dealii::Tensor<1, dim>> omegas = aq_base_ptr->GetAQDirs();
  std::map<std::pair<int, int>, int > reflection_map =
      aq_base_ptr->GetRefDirInd();

  for (auto const& mapping : reflection_map) {
    // Unroll this data structure
    int direction = mapping.first.second;
    int reflection = mapping.second;

    //x-direction
    EXPECT_FLOAT_EQ(omegas[direction][0], -omegas[reflection][0]);
  }
}

TEST_F(AQBaseTest, AQBase1DProduceAQ) {
  std::string filename = "aq_base_1d";
  btest::GoldTestInit(filename);
  OutputAQ<1>();
  btest::GoldTestRun(filename);
}

TEST_F(AQBaseTest, AQBase1DEpProduceAQ) {
  std::string filename = "aq_base_1d_ep";
  prm.set("transport model", "ep");
  btest::GoldTestInit(filename); // Opens deal log
  OutputAQ<1>();
  btest::GoldTestRun(filename); // Closes deal log
}

TEST_F(AQBaseTest, AQBaseBadDimProduceAQ) {
  AQBase<2> test_AQ(prm);
  ASSERT_THROW(test_AQ.ProduceAQ(), dealii::ExcMessage);
}

TEST_F(AQBaseTest, AQBaseBadDimReflectiveBC) {
  prm.set("have reflective boundary", "true");
  InitRefBCAndCheck<1>();
}

TEST_F(AQBaseTest, AQBaseSNOrderGetter) {
  AQBase<1> test_AQ(prm);
  ASSERT_EQ(test_AQ.GetSnOrder(), 4);
}

TEST_F(AQBaseTest, AQBaseGetNDir) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNDir(), 4);
}

TEST_F(AQBaseTest, AQBaseGetNDirEp) {
  prm.set("transport model", "ep");
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNDir(), 2);
}

TEST_F(AQBaseTest, AQBaseGetTotalHOVars) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNTotalHOVars(), 4);
}

TEST_F(AQBaseTest, AQBaseGetTotalHOVarsEq) {
  prm.set("transport model", "ep");
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  ASSERT_EQ(test_AQ.GetNTotalHOVars(), 2);
}

TEST_F(AQBaseTest, PrintAQ) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();

  std::ostringstream output_string_stream;
  test_AQ.PrintAQ(&output_string_stream);

  std::string output_regex = "transport model: regular; output_streamature name: None\n"
                             "Dim = 1, SN order = 4\n"
                             "Weights | Omega_x | Omega_y | mu\n"
                             "0.34785//d*  -0.86113//d*  \n"
                             "0.65214//d*  -0.33998//d*  \n"
                             "0.65214//d*  0.33998//d*  \n"
                             "0.34785//d*  0.86113//d*  \n";
  EXPECT_THAT(output_string_stream.str(), ::testing::ContainsRegex(output_regex));
}

TEST_F(AQBaseTest, RefDirInt) {
  AQBase<1> test_AQ(prm);
  test_AQ.MakeAQ();
  std::map<std::pair<int, int>, int> component_index_map(test_AQ.GetCompInd());

  for (auto const& mapping : component_index_map) {
    EXPECT_EQ(mapping.first.first, 0);
    EXPECT_EQ(mapping.first.second, mapping.second);
  }
}

TEST_F(AQBaseTest, InvDirInd) {
  AQBase<1> test_AQ(prm);
  std::unordered_map<int, std::pair<int, int>>
      inv_component_map(test_AQ.GetInvCompInd());

  for (size_t i = 0; i < inv_component_map.size(); ++i) {
    EXPECT_EQ(inv_component_map[i].first, 0);
    EXPECT_EQ(inv_component_map[i].second, i);
  }
}
<|endoftext|>


<commit_before>///
/// @file  D_mpi.cpp
/// @brief Implementation of the D formula (from Xavier Gourdon's
///        algorithm) that has been distributed using MPI (Message
///        Passing Interface) and multi-threaded using OpenMP.
///
/// Copyright (C) 2020 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include <primecount-internal.hpp>
#include <DFactorTable.hpp>
#include <PiTable.hpp>
#include <Sieve.hpp>
#include <LoadBalancer.hpp>
#include <fast_div.hpp>
#include <generate.hpp>
#include <generate_phi.hpp>
#include <imath.hpp>
#include <int128_t.hpp>
#include <min.hpp>
#include <print.hpp>
#include <mpi_reduce_sum.hpp>
#include <MpiLoadBalancer.hpp>

#include <stdint.h>
#include <vector>

using namespace std;
using namespace primecount;

namespace {

/// Compute the contribution of the hard special leaves using a
/// segmented sieve. Each thread processes the interval
/// [low, low + segments * segment_size[.
///
template <typename T, typename Primes, typename DFactorTable>
T D_thread(T x,
           int64_t x_star,
           int64_t xz,
           int64_t y,
           int64_t z,
           int64_t k,
           const Primes& primes,
           const PiTable& pi,
           const DFactorTable& factor,
           ThreadSettings& thread)
{
  T sum = 0;

  int64_t low = thread.low;
  int64_t segments = thread.segments;
  int64_t segment_size = thread.segment_size;
  int64_t pi_sqrtz = pi[isqrt(z)];
  int64_t low1 = max(low, 1);
  int64_t limit = min(low + segments * segment_size, xz);
  int64_t max_b = pi[min3(isqrt(x / low1), isqrt(limit), x_star)];
  int64_t min_b = pi[min(xz / limit, x_star)];
  min_b = max(k, min_b) + 1;

  if (min_b > max_b)
    return 0;

  Sieve sieve(low, segment_size, max_b);
  auto phi = generate_phi(low, max_b, primes, pi);
  thread.init_finished();

  // Segmented sieve of Eratosthenes
  for (; low < limit; low += segment_size)
  {
    // current segment [low, high[
    int64_t high = min(low + segment_size, limit);
    low1 = max(low, 1);

    // For i < min_b there are no special leaves:
    // low <= x / (primes[i] * m) < high
    sieve.pre_sieve(primes, min_b - 1, low, high);
    int64_t b = min_b;

    // For k + 1 <= b <= pi_sqrtz
    // Find all special leaves in the current segment that are
    // composed of a prime and a square free number:
    // low <= x / (primes[b] * m) < high
    for (int64_t end = min(pi_sqrtz, max_b); b <= end; b++)
    {
      int64_t prime = primes[b];
      T xp = x / prime;
      int64_t xp_low = min(fast_div(xp, low1), z);
      int64_t xp_high = min(fast_div(xp, high), z);
      int64_t min_m = max(xp_high, z / prime);
      int64_t max_m = min(fast_div(xp, prime * prime), xp_low);

      if (prime >= max_m)
        goto next_segment;

      min_m = factor.to_index(min_m);
      max_m = factor.to_index(max_m);

      for (int64_t m = max_m; m > min_m; m--)
      {
        // mu[m] != 0 &&
        // lpf[m] > prime &&
        // mpf[m] <= y
        if (prime < factor.is_leaf(m))
        {
          int64_t xpm = fast_div64(xp, factor.to_number(m));
          int64_t stop = xpm - low;
          int64_t phi_xpm = phi[b] + sieve.count(stop);
          int64_t mu_m = factor.mu(m);
          sum -= mu_m * phi_xpm;
        }
      }

      phi[b] += sieve.get_total_count();
      sieve.cross_off_count(prime, b);
    }

    // For pi_sqrtz < b <= pi_x_star
    // Find all special leaves in the current segment
    // that are composed of 2 primes:
    // low <= x / (primes[b] * primes[l]) < high
    for (; b <= max_b; b++)
    {
      int64_t prime = primes[b];
      T xp = x / prime;
      int64_t xp_low = min(fast_div(xp, low1), y);
      int64_t xp_high = min(fast_div(xp, high), y);
      int64_t min_m = max(xp_high, prime);
      int64_t max_m = min(fast_div(xp, prime * prime), xp_low);
      int64_t l = pi[max_m];

      if (prime >= primes[l])
        goto next_segment;

      for (; primes[l] > min_m; l--)
      {
        int64_t xpq = fast_div64(xp, primes[l]);
        int64_t stop = xpq - low;
        int64_t phi_xpq = phi[b] + sieve.count(stop);
        sum += phi_xpq;
      }

      phi[b] += sieve.get_total_count();
      sieve.cross_off_count(prime, b);
    }

    next_segment:;
  }

  return sum;
}

/// D MPI worker process.
/// Asks MPI main process for new work and reports
/// partial results to MPI main process.
///
template <typename T, typename Primes, typename DFactorTable>
void D_mpi_worker(T x,
                  int64_t y,
                  int64_t z,
                  int64_t k,
                  const Primes& primes,
                  const DFactorTable& factor,
                  int threads)
{
  PiTable pi(y, threads);
  int64_t xz = x / z;
  int64_t x_star = get_x_star_gourdon(x, y);
  int64_t thread_threshold = 1 << 20;
  threads = ideal_num_threads(threads, xz, thread_threshold);
  int main_proc_id = mpi_main_proc_id();
  int proc_id = mpi_proc_id();
  MpiMsg msg;

  #pragma omp parallel for num_threads(threads)
  for (int i = 0; i < threads; i++)
  {
    ThreadSettings thread;

    while (true)
    {
      #pragma omp critical (mpi_sync)
      {
        // send result to main process
        msg.set(proc_id, i, thread.low, thread.segments, thread.segment_size, thread.sum, thread.init_secs, thread.secs);
        msg.send(main_proc_id);

        // receive new work todo
        msg.recv(proc_id);
        thread.low = msg.low();
        thread.segments = msg.segments();
        thread.segment_size = msg.segment_size();
      }

      if (thread.low >= xz)
        break;

      // Unsigned integer division is usually slightly
      // faster than signed integer division
      using UT = typename make_unsigned<T>::type;

      thread.start_time();
      UT sum = D_thread((UT) x, x_star, xz, y, z, k, primes, pi, factor, thread);
      thread.sum = (T) sum;
      thread.stop_time();
    }
  }

  msg.set_finished();
  msg.send(main_proc_id);
}

/// D MPI main process.
/// Assigns work to the MPI worker processes.
///
template <typename T>
T D_mpi_main(T x,
             int64_t z,
             T d_approx)
{
  T sum = 0;
  int64_t xz = x / z;
  int workers = mpi_num_procs() - 1;

  MpiMsg msg;
  MpiLoadBalancer loadBalancer(x, xz, d_approx);
  Status status(x);

  while (workers > 0)
  {
    // wait for results from worker process
    msg.recv_any();

    if (msg.finished())
      workers--;
    else
    {
      sum += (T) msg.sum();
      int64_t high = msg.low() + msg.segments() * msg.segment_size();

      // update msg with new work
      loadBalancer.get_work(&msg);

      // send new work to worker process
      msg.send(msg.proc_id());

      if (is_print())
        status.print(high, xz, sum, d_approx);
    }
  }

  return sum;
}

} // namespace

namespace primecount {

int64_t D_mpi(int64_t x,
              int64_t y,
              int64_t z,
              int64_t k,
              int64_t d_approx,
              int threads)
{
  print("");
  print("=== D_mpi(x, y) ===");
  print_gourdon_vars(x, y, z, k, threads);

  int64_t sum = 0;
  double time = get_time();

  if (is_mpi_main_proc())
    sum = D_mpi_main(x, z, d_approx);
  else
  {
    DFactorTable<uint16_t> factor(y, z, threads);
    auto primes = generate_primes<int32_t>(y);
    D_mpi_worker(x, y, z, k, primes, factor, threads);
  }

  print("D", sum, time);
  return sum;
}

#ifdef HAVE_INT128_T

int128_t D_mpi(int128_t x,
               int64_t y,
               int64_t z,
               int64_t k,
               int128_t d_approx,
               int threads)
{
  print("");
  print("=== D_mpi(x, y) ===");
  print_gourdon_vars(x, y, z, k, threads);

  int128_t sum = 0;
  double time = get_time();

  if (is_mpi_main_proc())
    sum = D_mpi_main(x, z, d_approx);
  else
  {
    // uses less memory
    if (y <= FactorTable<uint16_t>::max())
    {
      DFactorTable<uint16_t> factor(y, z, threads);
      auto primes = generate_primes<uint32_t>(y);
      D_mpi_worker(x, y, z, k, primes, factor, threads);
    }
    else
    {
      DFactorTable<uint32_t> factor(y, z, threads);
      auto primes = generate_primes<int64_t>(y);
      D_mpi_worker(x, y, z, k, primes, factor, threads);
    }
  }

  print("D", sum, time);
  return sum;
}

#endif

} // namespace
<commit_msg>Use OmpLock instead of critical section<commit_after>///
/// @file  D_mpi.cpp
/// @brief Implementation of the D formula (from Xavier Gourdon's
///        algorithm) that has been distributed using MPI (Message
///        Passing Interface) and multi-threaded using OpenMP.
///
/// Copyright (C) 2020 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include <primecount-internal.hpp>
#include <DFactorTable.hpp>
#include <PiTable.hpp>
#include <Sieve.hpp>
#include <LoadBalancer.hpp>
#include <fast_div.hpp>
#include <generate.hpp>
#include <generate_phi.hpp>
#include <imath.hpp>
#include <int128_t.hpp>
#include <min.hpp>
#include <OmpLock.hpp>
#include <print.hpp>
#include <mpi_reduce_sum.hpp>
#include <MpiLoadBalancer.hpp>

#include <stdint.h>
#include <vector>

using namespace std;
using namespace primecount;

namespace {

/// Compute the contribution of the hard special leaves using a
/// segmented sieve. Each thread processes the interval
/// [low, low + segments * segment_size[.
///
template <typename T, typename Primes, typename DFactorTable>
T D_thread(T x,
           int64_t x_star,
           int64_t xz,
           int64_t y,
           int64_t z,
           int64_t k,
           const Primes& primes,
           const PiTable& pi,
           const DFactorTable& factor,
           ThreadSettings& thread)
{
  T sum = 0;

  int64_t low = thread.low;
  int64_t segments = thread.segments;
  int64_t segment_size = thread.segment_size;
  int64_t pi_sqrtz = pi[isqrt(z)];
  int64_t low1 = max(low, 1);
  int64_t limit = min(low + segments * segment_size, xz);
  int64_t max_b = pi[min3(isqrt(x / low1), isqrt(limit), x_star)];
  int64_t min_b = pi[min(xz / limit, x_star)];
  min_b = max(k, min_b) + 1;

  if (min_b > max_b)
    return 0;

  Sieve sieve(low, segment_size, max_b);
  auto phi = generate_phi(low, max_b, primes, pi);
  thread.init_finished();

  // Segmented sieve of Eratosthenes
  for (; low < limit; low += segment_size)
  {
    // current segment [low, high[
    int64_t high = min(low + segment_size, limit);
    low1 = max(low, 1);

    // For i < min_b there are no special leaves:
    // low <= x / (primes[i] * m) < high
    sieve.pre_sieve(primes, min_b - 1, low, high);
    int64_t b = min_b;

    // For k + 1 <= b <= pi_sqrtz
    // Find all special leaves in the current segment that are
    // composed of a prime and a square free number:
    // low <= x / (primes[b] * m) < high
    for (int64_t end = min(pi_sqrtz, max_b); b <= end; b++)
    {
      int64_t prime = primes[b];
      T xp = x / prime;
      int64_t xp_low = min(fast_div(xp, low1), z);
      int64_t xp_high = min(fast_div(xp, high), z);
      int64_t min_m = max(xp_high, z / prime);
      int64_t max_m = min(fast_div(xp, prime * prime), xp_low);

      if (prime >= max_m)
        goto next_segment;

      min_m = factor.to_index(min_m);
      max_m = factor.to_index(max_m);

      for (int64_t m = max_m; m > min_m; m--)
      {
        // mu[m] != 0 &&
        // lpf[m] > prime &&
        // mpf[m] <= y
        if (prime < factor.is_leaf(m))
        {
          int64_t xpm = fast_div64(xp, factor.to_number(m));
          int64_t stop = xpm - low;
          int64_t phi_xpm = phi[b] + sieve.count(stop);
          int64_t mu_m = factor.mu(m);
          sum -= mu_m * phi_xpm;
        }
      }

      phi[b] += sieve.get_total_count();
      sieve.cross_off_count(prime, b);
    }

    // For pi_sqrtz < b <= pi_x_star
    // Find all special leaves in the current segment
    // that are composed of 2 primes:
    // low <= x / (primes[b] * primes[l]) < high
    for (; b <= max_b; b++)
    {
      int64_t prime = primes[b];
      T xp = x / prime;
      int64_t xp_low = min(fast_div(xp, low1), y);
      int64_t xp_high = min(fast_div(xp, high), y);
      int64_t min_m = max(xp_high, prime);
      int64_t max_m = min(fast_div(xp, prime * prime), xp_low);
      int64_t l = pi[max_m];

      if (prime >= primes[l])
        goto next_segment;

      for (; primes[l] > min_m; l--)
      {
        int64_t xpq = fast_div64(xp, primes[l]);
        int64_t stop = xpq - low;
        int64_t phi_xpq = phi[b] + sieve.count(stop);
        sum += phi_xpq;
      }

      phi[b] += sieve.get_total_count();
      sieve.cross_off_count(prime, b);
    }

    next_segment:;
  }

  return sum;
}

/// D MPI worker process.
/// Asks MPI main process for new work and reports
/// partial results to MPI main process.
///
template <typename T, typename Primes, typename DFactorTable>
void D_mpi_worker(T x,
                  int64_t y,
                  int64_t z,
                  int64_t k,
                  const Primes& primes,
                  const DFactorTable& factor,
                  int threads)
{
  PiTable pi(y, threads);
  int64_t xz = x / z;
  int64_t x_star = get_x_star_gourdon(x, y);
  int64_t thread_threshold = 1 << 20;
  threads = ideal_num_threads(threads, xz, thread_threshold);
  int main_proc_id = mpi_main_proc_id();
  int proc_id = mpi_proc_id();
  MpiMsg msg;
  OmpLock ompLock;

  #pragma omp parallel for num_threads(threads)
  for (int i = 0; i < threads; i++)
  {
    ThreadSettings thread;

    while (true)
    {
      {
        LockGuard lockGuard(ompLock);

        // send result to main process
        msg.set(proc_id, i, thread.low, thread.segments, thread.segment_size, thread.sum, thread.init_secs, thread.secs);
        msg.send(main_proc_id);

        // receive new work todo
        msg.recv(proc_id);
        thread.low = msg.low();
        thread.segments = msg.segments();
        thread.segment_size = msg.segment_size();
      }

      if (thread.low >= xz)
        break;

      // Unsigned integer division is usually slightly
      // faster than signed integer division
      using UT = typename make_unsigned<T>::type;

      thread.start_time();
      UT sum = D_thread((UT) x, x_star, xz, y, z, k, primes, pi, factor, thread);
      thread.sum = (T) sum;
      thread.stop_time();
    }
  }

  msg.set_finished();
  msg.send(main_proc_id);
}

/// D MPI main process.
/// Assigns work to the MPI worker processes.
///
template <typename T>
T D_mpi_main(T x,
             int64_t z,
             T d_approx)
{
  T sum = 0;
  int64_t xz = x / z;
  int workers = mpi_num_procs() - 1;

  MpiMsg msg;
  MpiLoadBalancer loadBalancer(x, xz, d_approx);
  Status status(x);

  while (workers > 0)
  {
    // wait for results from worker process
    msg.recv_any();

    if (msg.finished())
      workers--;
    else
    {
      sum += (T) msg.sum();
      int64_t high = msg.low() + msg.segments() * msg.segment_size();

      // update msg with new work
      loadBalancer.get_work(&msg);

      // send new work to worker process
      msg.send(msg.proc_id());

      if (is_print())
        status.print(high, xz, sum, d_approx);
    }
  }

  return sum;
}

} // namespace

namespace primecount {

int64_t D_mpi(int64_t x,
              int64_t y,
              int64_t z,
              int64_t k,
              int64_t d_approx,
              int threads)
{
  print("");
  print("=== D_mpi(x, y) ===");
  print_gourdon_vars(x, y, z, k, threads);

  int64_t sum = 0;
  double time = get_time();

  if (is_mpi_main_proc())
    sum = D_mpi_main(x, z, d_approx);
  else
  {
    DFactorTable<uint16_t> factor(y, z, threads);
    auto primes = generate_primes<int32_t>(y);
    D_mpi_worker(x, y, z, k, primes, factor, threads);
  }

  print("D", sum, time);
  return sum;
}

#ifdef HAVE_INT128_T

int128_t D_mpi(int128_t x,
               int64_t y,
               int64_t z,
               int64_t k,
               int128_t d_approx,
               int threads)
{
  print("");
  print("=== D_mpi(x, y) ===");
  print_gourdon_vars(x, y, z, k, threads);

  int128_t sum = 0;
  double time = get_time();

  if (is_mpi_main_proc())
    sum = D_mpi_main(x, z, d_approx);
  else
  {
    // uses less memory
    if (y <= FactorTable<uint16_t>::max())
    {
      DFactorTable<uint16_t> factor(y, z, threads);
      auto primes = generate_primes<uint32_t>(y);
      D_mpi_worker(x, y, z, k, primes, factor, threads);
    }
    else
    {
      DFactorTable<uint32_t> factor(y, z, threads);
      auto primes = generate_primes<int64_t>(y);
      D_mpi_worker(x, y, z, k, primes, factor, threads);
    }
  }

  print("D", sum, time);
  return sum;
}

#endif

} // namespace
<|endoftext|>


<commit_before>#ifndef ARCH_RUNTIME_THREAD_POOL_HPP_
#define ARCH_RUNTIME_THREAD_POOL_HPP_

#include <pthread.h>

#include <map>
#include <string>

#include "config/args.hpp"
#include "arch/runtime/event_queue.hpp"
#include "arch/runtime/system_event.hpp"
#include "arch/runtime/message_hub.hpp"
#include "arch/runtime/coroutines.hpp"
#include "arch/io/blocker_pool.hpp"
#include "arch/timer.hpp"

class linux_thread_t;

/* coro_runtime_t is borrowed from coroutines.hpp.  Please only
construct one coro_runtime_t per thread. Coroutines can only be used
when a coro_runtime_t exists. It exists to take advantage of RAII. */

struct coro_runtime_t {
    coro_runtime_t();
    ~coro_runtime_t();

#ifndef NDEBUG
    void get_coroutine_counts(std::map<std::string, size_t> *dest);
#endif
};

/* A thread pool represents a group of threads, each of which is associated with an
event queue. There is one thread pool per server. It is responsible for starting up
and shutting down the threads and event queues. */

class linux_thread_pool_t {
public:
    linux_thread_pool_t(int worker_threads, bool do_set_affinity);

    // When the process receives a SIGINT or SIGTERM, interrupt_message will be delivered to the
    // same thread that initial_message was delivered to, and interrupt_message will be set to
    // NULL. If you want to receive notification of further SIGINTs or SIGTERMs, you must call
    // set_interrupt_message() again. Returns the previous value of interrupt_message.
    static linux_thread_message_t *set_interrupt_message(linux_thread_message_t *interrupt_message);

    // Blocks while threads are working. Only returns after shutdown() is called. initial_message
    // is a thread message that will be delivered to thread zero after all of the event queues
    // have been started; it is used to start the server's activity.
    void run_thread_pool(linux_thread_message_t *initial_message);

#ifndef NDEBUG
    void enable_coroutine_summary();
#endif

    // Shut down all the threads. Can be called from any thread.
    void shutdown_thread_pool();

    ~linux_thread_pool_t();

private:
#ifndef NDEBUG
    bool coroutine_summary;
#endif

    static void *start_thread(void*);

    static void interrupt_handler(int);
    static void sigsegv_handler(int, siginfo_t *, void *) NORETURN;
    pthread_spinlock_t interrupt_message_lock;
    linux_thread_message_t *interrupt_message;

    // Used to signal the main thread for shutdown
    volatile bool do_shutdown;
    pthread_cond_t shutdown_cond;
    pthread_mutex_t shutdown_cond_mutex;

    // The number of threads to allocate for handling blocking calls
    static const int GENERIC_BLOCKER_THREAD_COUNT = 2;
    blocker_pool_t* generic_blocker_pool;

public:
    pthread_t pthreads[MAX_THREADS];
    linux_thread_t *threads[MAX_THREADS];

    // Cooperatively run a blocking function call using the generic_blocker_pool
    template <class Callable>
    static void run_in_blocker_pool(const Callable &);

    int n_threads;
    bool do_set_affinity;
    // The thread_pool that started the thread we are currently in
    static __thread linux_thread_pool_t *thread_pool;
    // The ID of the thread we are currently in
    static __thread int thread_id;
    // The event queue for the thread we are currently in (same as &thread_pool->threads[thread_id])
    static __thread linux_thread_t *thread;

private:
    DISABLE_COPYING(linux_thread_pool_t);
};

template <class Callable>
struct generic_job_t :
    public blocker_pool_t::job_t
{
    void run() {
        (*fn)();
    }

    void done() {
        // Now that the function is done, resume execution of the suspended task
        suspended->notify_sometime();
    }

    const Callable *fn;
    coro_t* suspended;
};

// Function to handle blocking calls in a separate thread pool
// This should be used for any calls that cannot otherwise be made non-blocking
template <class Callable>
void linux_thread_pool_t::run_in_blocker_pool(const Callable &fn)
{
    generic_job_t<Callable> job;
    job.fn = &fn;
    job.suspended = coro_t::self();

    rassert(thread_pool->generic_blocker_pool != NULL,
            "thread_pool_t::run_in_blocker_pool called while generic_thread_pool uninitialized");
    thread_pool->generic_blocker_pool->do_job(&job);

    // Give up execution, to be resumed when the done callback is made
    coro_t::wait();
}

class linux_thread_t :
    public linux_event_callback_t,
    public linux_queue_parent_t
{
    timer_token_t *perfmon_stats_timer;

public:
    linux_thread_t(linux_thread_pool_t *parent_pool, int thread_id);
    ~linux_thread_t();

    linux_event_queue_t queue;
    linux_message_hub_t message_hub;
    timer_handler_t timer_handler;

    /* Never accessed; its constructor and destructor set up and tear down thread-local variables
    for coroutines. */
    coro_runtime_t coro_runtime;

    void pump();   // Called by the event queue
    bool should_shut_down();   // Called by the event queue
#ifndef NDEBUG
    void initiate_shut_down(std::map<std::string, size_t> *coroutine_counts); // Can be called from any thread
#else
    void initiate_shut_down(); // Can be called from any thread
#endif
    void on_event(int events);

private:
    volatile bool do_shutdown;
    pthread_mutex_t do_shutdown_mutex;
    system_event_t shutdown_notify_event;

#ifndef NDEBUG
    std::map<std::string, size_t> *coroutine_counts_at_shutdown;
#endif
};

#endif /* ARCH_RUNTIME_THREAD_POOL_HPP_ */
<commit_msg>Made set_interrupt_message have a MUST_USE return value.<commit_after>#ifndef ARCH_RUNTIME_THREAD_POOL_HPP_
#define ARCH_RUNTIME_THREAD_POOL_HPP_

#include <pthread.h>

#include <map>
#include <string>

#include "config/args.hpp"
#include "arch/runtime/event_queue.hpp"
#include "arch/runtime/system_event.hpp"
#include "arch/runtime/message_hub.hpp"
#include "arch/runtime/coroutines.hpp"
#include "arch/io/blocker_pool.hpp"
#include "arch/timer.hpp"

class linux_thread_t;

/* coro_runtime_t is borrowed from coroutines.hpp.  Please only
construct one coro_runtime_t per thread. Coroutines can only be used
when a coro_runtime_t exists. It exists to take advantage of RAII. */

struct coro_runtime_t {
    coro_runtime_t();
    ~coro_runtime_t();

#ifndef NDEBUG
    void get_coroutine_counts(std::map<std::string, size_t> *dest);
#endif
};

/* A thread pool represents a group of threads, each of which is associated with an
event queue. There is one thread pool per server. It is responsible for starting up
and shutting down the threads and event queues. */

class linux_thread_pool_t {
public:
    linux_thread_pool_t(int worker_threads, bool do_set_affinity);

    // When the process receives a SIGINT or SIGTERM, interrupt_message will be delivered to the
    // same thread that initial_message was delivered to, and interrupt_message will be set to
    // NULL. If you want to receive notification of further SIGINTs or SIGTERMs, you must call
    // set_interrupt_message() again. Returns the previous value of interrupt_message.
    static MUST_USE linux_thread_message_t *set_interrupt_message(linux_thread_message_t *interrupt_message);

    // Blocks while threads are working. Only returns after shutdown() is called. initial_message
    // is a thread message that will be delivered to thread zero after all of the event queues
    // have been started; it is used to start the server's activity.
    void run_thread_pool(linux_thread_message_t *initial_message);

#ifndef NDEBUG
    void enable_coroutine_summary();
#endif

    // Shut down all the threads. Can be called from any thread.
    void shutdown_thread_pool();

    ~linux_thread_pool_t();

private:
#ifndef NDEBUG
    bool coroutine_summary;
#endif

    static void *start_thread(void*);

    static void interrupt_handler(int);
    static void sigsegv_handler(int, siginfo_t *, void *) NORETURN;
    pthread_spinlock_t interrupt_message_lock;
    linux_thread_message_t *interrupt_message;

    // Used to signal the main thread for shutdown
    volatile bool do_shutdown;
    pthread_cond_t shutdown_cond;
    pthread_mutex_t shutdown_cond_mutex;

    // The number of threads to allocate for handling blocking calls
    static const int GENERIC_BLOCKER_THREAD_COUNT = 2;
    blocker_pool_t* generic_blocker_pool;

public:
    pthread_t pthreads[MAX_THREADS];
    linux_thread_t *threads[MAX_THREADS];

    // Cooperatively run a blocking function call using the generic_blocker_pool
    template <class Callable>
    static void run_in_blocker_pool(const Callable &);

    int n_threads;
    bool do_set_affinity;
    // The thread_pool that started the thread we are currently in
    static __thread linux_thread_pool_t *thread_pool;
    // The ID of the thread we are currently in
    static __thread int thread_id;
    // The event queue for the thread we are currently in (same as &thread_pool->threads[thread_id])
    static __thread linux_thread_t *thread;

private:
    DISABLE_COPYING(linux_thread_pool_t);
};

template <class Callable>
struct generic_job_t :
    public blocker_pool_t::job_t
{
    void run() {
        (*fn)();
    }

    void done() {
        // Now that the function is done, resume execution of the suspended task
        suspended->notify_sometime();
    }

    const Callable *fn;
    coro_t* suspended;
};

// Function to handle blocking calls in a separate thread pool
// This should be used for any calls that cannot otherwise be made non-blocking
template <class Callable>
void linux_thread_pool_t::run_in_blocker_pool(const Callable &fn)
{
    generic_job_t<Callable> job;
    job.fn = &fn;
    job.suspended = coro_t::self();

    rassert(thread_pool->generic_blocker_pool != NULL,
            "thread_pool_t::run_in_blocker_pool called while generic_thread_pool uninitialized");
    thread_pool->generic_blocker_pool->do_job(&job);

    // Give up execution, to be resumed when the done callback is made
    coro_t::wait();
}

class linux_thread_t :
    public linux_event_callback_t,
    public linux_queue_parent_t
{
    timer_token_t *perfmon_stats_timer;

public:
    linux_thread_t(linux_thread_pool_t *parent_pool, int thread_id);
    ~linux_thread_t();

    linux_event_queue_t queue;
    linux_message_hub_t message_hub;
    timer_handler_t timer_handler;

    /* Never accessed; its constructor and destructor set up and tear down thread-local variables
    for coroutines. */
    coro_runtime_t coro_runtime;

    void pump();   // Called by the event queue
    bool should_shut_down();   // Called by the event queue
#ifndef NDEBUG
    void initiate_shut_down(std::map<std::string, size_t> *coroutine_counts); // Can be called from any thread
#else
    void initiate_shut_down(); // Can be called from any thread
#endif
    void on_event(int events);

private:
    volatile bool do_shutdown;
    pthread_mutex_t do_shutdown_mutex;
    system_event_t shutdown_notify_event;

#ifndef NDEBUG
    std::map<std::string, size_t> *coroutine_counts_at_shutdown;
#endif
};

#endif /* ARCH_RUNTIME_THREAD_POOL_HPP_ */
<|endoftext|>


<commit_before>// init: The initial user-level program

#include "types.h"
#include "user.h"
#include <fcntl.h>
#include "lib.h"
#include "major.h"

static const char *sh_argv[] = { "sh", 0 };
static const char *app_argv[][MAXARG] = {
#ifdef LWIP
  { "telnetd", 0 },
  { "httpd", 0 },
#endif
};

static struct {
  const char* name;
  int major;
} dev[] = {
  { "/dev/netif",     MAJ_NETIF },
  { "/dev/sampler",   MAJ_SAMPLER },
  { "/dev/lockstat",  MAJ_LOCKSTAT },
  { "/dev/stat",      MAJ_STAT },
  { "/dev/cmdline",   MAJ_CMDLINE},
  { "/dev/gc",   MAJ_GC},
};

static int
startone(const char **argv)
{
  int pid;

  pid = fork(0);
  if(pid < 0){
    fprintf(1, "init: fork failed\n");
    exit();
  }
  if(pid == 0){
    exec(argv[0], argv);
    fprintf(1, "init: exec %s failed\n", argv[0]);
    exit();
  }
  return pid;
}

static void
runcmdline(void)
{
  const char* argv[3] = { "sh", 0, 0 };
  char buf[256];
  char* b;
  long r;
  int fd;

  fd = open("/dev/cmdline", O_RDONLY);
  if (fd < 0)
    return;

  r = read(fd, buf, sizeof(buf)-1);
  if (r < 0)
    return;
  buf[r] = 0;

  if ((b = strchr(buf, '$'))) {
    argv[1] = b+1;
    startone(argv);
  }
}

int
main(void)
{
  int pid, wpid;

  if(open("console", O_RDWR) < 0){
    mknod("console", 1, 1);
    open("console", O_RDWR);
  }
  dup(0);  // stdout
  dup(0);  // stderr

  mkdir("dev", 0777);
  for (int i = 0; i < NELEM(dev); i++)
    if (mknod(dev[i].name, dev[i].major, 1) < 0)
      fprintf(2, "init: mknod %s failed\n", dev[i].name);

  for (u32 i = 0; i < NELEM(app_argv); i++)
    startone(app_argv[i]);

  runcmdline();

  for(;;){
    pid = startone(sh_argv);
    while((wpid=wait()) >= 0 && wpid != pid)
      fprintf(1, "zombie!\n");
  }
}
<commit_msg>init: Pass -c to sh when running kernel command line<commit_after>// init: The initial user-level program

#include "types.h"
#include "user.h"
#include <fcntl.h>
#include "lib.h"
#include "major.h"

static const char *sh_argv[] = { "sh", 0 };
static const char *app_argv[][MAXARG] = {
#ifdef LWIP
  { "telnetd", 0 },
  { "httpd", 0 },
#endif
};

static struct {
  const char* name;
  int major;
} dev[] = {
  { "/dev/netif",     MAJ_NETIF },
  { "/dev/sampler",   MAJ_SAMPLER },
  { "/dev/lockstat",  MAJ_LOCKSTAT },
  { "/dev/stat",      MAJ_STAT },
  { "/dev/cmdline",   MAJ_CMDLINE},
  { "/dev/gc",   MAJ_GC},
};

static int
startone(const char **argv)
{
  int pid;

  pid = fork(0);
  if(pid < 0){
    fprintf(1, "init: fork failed\n");
    exit();
  }
  if(pid == 0){
    exec(argv[0], argv);
    fprintf(1, "init: exec %s failed\n", argv[0]);
    exit();
  }
  return pid;
}

static void
runcmdline(void)
{
  const char* argv[4] = { "sh", "-c", 0 };
  char buf[256];
  char* b;
  long r;
  int fd;

  fd = open("/dev/cmdline", O_RDONLY);
  if (fd < 0)
    return;

  r = read(fd, buf, sizeof(buf)-1);
  if (r < 0)
    return;
  buf[r] = 0;

  if ((b = strchr(buf, '$'))) {
    argv[2] = b+1;
    startone(argv);
  }
}

int
main(void)
{
  int pid, wpid;

  if(open("console", O_RDWR) < 0){
    mknod("console", 1, 1);
    open("console", O_RDWR);
  }
  dup(0);  // stdout
  dup(0);  // stderr

  mkdir("dev", 0777);
  for (int i = 0; i < NELEM(dev); i++)
    if (mknod(dev[i].name, dev[i].major, 1) < 0)
      fprintf(2, "init: mknod %s failed\n", dev[i].name);

  for (u32 i = 0; i < NELEM(app_argv); i++)
    startone(app_argv[i]);

  runcmdline();

  for(;;){
    pid = startone(sh_argv);
    while((wpid=wait()) >= 0 && wpid != pid)
      fprintf(1, "zombie!\n");
  }
}
<|endoftext|>


<commit_before>/*****************************************************************************
*  Copyright 2011 Sergey Shekyan
*
*  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.
* *****************************************************************************/

/*****
 * Author: Sergey Shekyan sshekyan@qualys.com
 *
 * Slow HTTP attack  vulnerability test tool
 *  http://code.google.com/p/slowhttptest/
 *****/

#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <stdio.h>
#include <string.h>
#include <cmath>
#include <string>

#include <openssl/ssl.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>

#include "slowlog.h"
#include "slowsocket.h"
#include "slowurl.h"

namespace slowhttptest {
SlowSocket::SlowSocket()
    : sockfd_(-1), requests_to_send_(0),
      followups_to_send_(0), last_followup_timing_(0),
      offset_(0), ssl_(0), buf_(0),
      start_in_millisecs_(0), connected_in_millisecs_(0),
      stop_in_millisecs_(0), state_(eInit) {
}

SlowSocket::~SlowSocket() {
  close();
}

int SlowSocket::set_nonblocking() {
  int flags;

  if(-1 == (flags = fcntl(sockfd_, F_GETFL, 0))) {
    flags = 0;
  }
  return fcntl(sockfd_, F_SETFL, flags | O_NONBLOCK);
}

bool SlowSocket::init(addrinfo* addr, const Url* url, int& maxfd,
                      int followups_to_send) {
  addrinfo* res;
  bool connect_initiated_ = false;
  for (res = addr; !connect_initiated_ && res; res = res->ai_next) {
    sockfd_ = socket(res->ai_family, res->ai_socktype,
     res->ai_protocol);
    if(-1 == sockfd_) {
      slowlog(LOG_ERROR, "failed to create socket: %s\n", strerror(errno));
      return false;
    }

    if(-1 == set_nonblocking()) {
      slowlog(LOG_ERROR, "failed to set socket %d to non-blocking \n", sockfd_);
      return false;
    }

    slowlog(LOG_DEBUG, "non-blocking socket fd = %d created\n", sockfd_);
    if(connect_initiated_ = url->isSSL() ? connect_ssl(addr) : connect_plain(addr)) {
      break; //found right addrinfo
    }

  }

  followups_to_send_ = followups_to_send;
  requests_to_send_ = 1;

  if(sockfd_ > maxfd) {
    maxfd = sockfd_;
  }
  return true;
}

bool SlowSocket::connect_plain(addrinfo* addr) {
  errno = 0;

  if (connect(sockfd_, addr->ai_addr, addr->ai_addrlen) < 0
      && EINPROGRESS != errno) {
    slowlog(LOG_ERROR, "cannot connect socket %d: %s\n", sockfd_, strerror(errno));
    close();
    return false;
  }
  return true;
}

bool SlowSocket::connect_ssl(addrinfo* addr) {
  // Establish regular connection.
  if(!connect_plain(addr))  return false;

  // Init SSL related stuff.
  // TODO(vagababov): this is not thread safe of pretty.
  static bool ssl_is_initialized = false;
  if (!ssl_is_initialized) {
    SSL_library_init();
    ssl_is_initialized = true;
  }
  SSL_METHOD* method = NULL;
  SSL_CTX* ssl_ctx = NULL;
  method = (SSL_METHOD*)SSLv23_client_method();
  ssl_ctx = SSL_CTX_new(method);
  if(!ssl_ctx) {
    slowlog(LOG_ERROR, "cannot create new SSL context\n");
    close();
    return false;
  }
  ssl_ = SSL_new(ssl_ctx);
  if(!ssl_) {
    slowlog(LOG_ERROR, "cannot create SSL structure for a connection\n");
    close();
    return false;
  }
  SSL_set_fd(ssl_, sockfd_);
  int ret = SSL_connect(ssl_);
  if(ret <= 0) {
    int err = SSL_get_error(ssl_, ret);
    //slowlog(LOG_ERROR, "socket %d: SSL connect error: %d\n", sockfd_, err);
    if(SSL_ERROR_WANT_READ != err && SSL_ERROR_WANT_WRITE != err) {
      close();
      return false;
    }
  }
  return true;
}

int SlowSocket::recv_slow(void *buf, size_t len) {
  int ret=ssl_ ? SSL_read(ssl_, buf, len)
    : recv(sockfd_, buf, len, 0);
  if(ssl_) {
    if(ret <0) {
      int err = SSL_get_error(ssl_, ret);
      if(err == SSL_ERROR_WANT_WRITE) {
        requests_to_send_ = 1;
      }
    }
    if(SSL_is_init_finished(ssl_) && (state_ == eConnecting)) {
      requests_to_send_ = 1;
    }
  }
  return ret;
}

int SlowSocket::send_slow(const void* buf, size_t len, const SendType type) {
  int ret;
  if(ssl_) {
    if(!SSL_is_init_finished(ssl_)) {
      ret = SSL_do_handshake(ssl_);
      if(ret <= 0) {
        int err = SSL_get_error(ssl_, ret);
        if(SSL_ERROR_WANT_READ != err && SSL_ERROR_WANT_WRITE != err) {
          slowlog(LOG_ERROR, "socket %d: SSL connect error: %d\n", sockfd_, err);
          close();
          return -1;
        } else {
          if(SSL_ERROR_WANT_READ == err) {
            requests_to_send_=0;
          }
          else {
            requests_to_send_=1;
          }
          errno = EAGAIN;
          return -1;
        }
      } else { //connected and handhsake finished
        requests_to_send_=1;
      }
    } else {
      if(requests_to_send_ > 0) { //report for initial data only
        slowlog(LOG_DEBUG, "SSL connection is using %s\n", SSL_get_cipher(ssl_));
      }
    }
  }
  // VA: this is not good. create a "prepare" method.
  // initial send
  if(buf_ == 0) {
    buf_ = buf;
    offset_ = len;
  }

  ret = ssl_ ? SSL_write(ssl_, buf_, offset_)
                 : send(sockfd_, buf_, offset_, 0);

  // entire data was sent
  if(ret > 0 && ret == offset_) {
    if(eInitialSend == type) {
      --requests_to_send_;
    } else if(eFollowUpSend == type) {
      --followups_to_send_;
    }
    buf_ = 0;
    offset_ = 0;
  } else if(ret > 0 && ret < offset_) {
    buf_ = static_cast<const char*>(buf_) + ret;
    offset_ -= ret;
  }
  return ret;
}

void SlowSocket::close() {
  if (-1 == sockfd_) return;

  slowlog(LOG_DEBUG, "closing slow, sock is %d\n", sockfd_);
  if(ssl_) {
    SSL_free(ssl_);
    ssl_ = NULL;
  }
  requests_to_send_ = 0;
  followups_to_send_ = 0;
  ::close(sockfd_);
  sockfd_ = -1;
}

void SlowSocket::set_state(SocketState state) {
  timeval t;
  gettimeofday(&t, 0);
  switch(state) {
    case eInit:
      break;
    case eConnecting:
      set_start(&t);
      break;
    case eConnected:
      set_connected(&t);
      break;
    case eError:
      break;
    case eClosed:
      set_stop(&t);
      break;
    default:
      break;
  }
  state_ = state;
}

}  // namespace slowhttptest
<commit_msg>Formatting fixes<commit_after>/*****************************************************************************
*  Copyright 2011 Sergey Shekyan
*
*  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.
* *****************************************************************************/

/*****
 * Author: Sergey Shekyan sshekyan@qualys.com
 *
 * Slow HTTP attack  vulnerability test tool
 *  http://code.google.com/p/slowhttptest/
 *****/

#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <stdio.h>
#include <string.h>
#include <cmath>
#include <string>

#include <openssl/ssl.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>

#include "slowlog.h"
#include "slowsocket.h"
#include "slowurl.h"

namespace slowhttptest {
SlowSocket::SlowSocket()
    : sockfd_(-1), requests_to_send_(0),
      followups_to_send_(0), last_followup_timing_(0),
      offset_(0), ssl_(0), buf_(0),
      start_in_millisecs_(0), connected_in_millisecs_(0),
      stop_in_millisecs_(0), state_(eInit) {
}

SlowSocket::~SlowSocket() {
  close();
}

int SlowSocket::set_nonblocking() {
  int flags;

  if(-1 == (flags = fcntl(sockfd_, F_GETFL, 0))) {
    flags = 0;
  }
  return fcntl(sockfd_, F_SETFL, flags | O_NONBLOCK);
}

bool SlowSocket::init(addrinfo* addr, const Url* url, int& maxfd,
                      int followups_to_send) {
  addrinfo* res;
  bool connect_initiated_ = false;
  for (res = addr; !connect_initiated_ && res; res = res->ai_next) {
    sockfd_ = socket(res->ai_family, res->ai_socktype,
                     res->ai_protocol);
    if(-1 == sockfd_) {
      slowlog(LOG_ERROR, "failed to create socket: %s\n", strerror(errno));
      return false;
    }

    if(-1 == set_nonblocking()) {
      slowlog(LOG_ERROR, "failed to set socket %d to non-blocking \n", sockfd_);
      return false;
    }

    slowlog(LOG_DEBUG, "non-blocking socket fd = %d created\n", sockfd_);
    if(connect_initiated_ = url->isSSL() ? connect_ssl(addr) : connect_plain(addr)) {
      break; //found right addrinfo
    }
  }

  followups_to_send_ = followups_to_send;
  requests_to_send_ = 1;

  maxdf = max(sockfd_, maxfd);
  return true;
}

bool SlowSocket::connect_plain(addrinfo* addr) {
  errno = 0;

  if (connect(sockfd_, addr->ai_addr, addr->ai_addrlen) < 0
      && EINPROGRESS != errno) {
    slowlog(LOG_ERROR, "cannot connect socket %d: %s\n", sockfd_, strerror(errno));
    close();
    return false;
  }
  return true;
}

bool SlowSocket::connect_ssl(addrinfo* addr) {
  // Establish regular connection.
  if(!connect_plain(addr))  return false;

  // Init SSL related stuff.
  // TODO(vagababov): this is not thread safe of pretty.
  static bool ssl_is_initialized = false;
  if (!ssl_is_initialized) {
    SSL_library_init();
    ssl_is_initialized = true;
  }
  SSL_METHOD* method = NULL;
  SSL_CTX* ssl_ctx = NULL;
  method = (SSL_METHOD*)SSLv23_client_method();
  ssl_ctx = SSL_CTX_new(method);
  if(!ssl_ctx) {
    slowlog(LOG_ERROR, "cannot create new SSL context\n");
    close();
    return false;
  }
  ssl_ = SSL_new(ssl_ctx);
  if(!ssl_) {
    slowlog(LOG_ERROR, "cannot create SSL structure for a connection\n");
    close();
    return false;
  }
  SSL_set_fd(ssl_, sockfd_);
  int ret = SSL_connect(ssl_);
  if(ret <= 0) {
    int err = SSL_get_error(ssl_, ret);
    //slowlog(LOG_ERROR, "socket %d: SSL connect error: %d\n", sockfd_, err);
    if(SSL_ERROR_WANT_READ != err && SSL_ERROR_WANT_WRITE != err) {
      close();
      return false;
    }
  }
  return true;
}

int SlowSocket::recv_slow(void *buf, size_t len) {
  int ret = ssl_ ? SSL_read(ssl_, buf, len)
                 : recv(sockfd_, buf, len, 0);
  if(ssl_) {
    if(ret < 0) {
      int err = SSL_get_error(ssl_, ret);
      if(err == SSL_ERROR_WANT_WRITE) {
        requests_to_send_ = 1;
      }
    }
    if(SSL_is_init_finished(ssl_) && (state_ == eConnecting)) {
      requests_to_send_ = 1;
    }
  }
  return ret;
}

int SlowSocket::send_slow(const void* buf, size_t len, const SendType type) {
  int ret;
  if(ssl_) {
    if(!SSL_is_init_finished(ssl_)) {
      ret = SSL_do_handshake(ssl_);
      if(ret <= 0) {
        int err = SSL_get_error(ssl_, ret);
        if(SSL_ERROR_WANT_READ != err && SSL_ERROR_WANT_WRITE != err) {
          slowlog(LOG_ERROR, "socket %d: SSL connect error: %d\n", sockfd_, err);
          close();
          return -1;
        } else {
          if(SSL_ERROR_WANT_READ == err) {
            requests_to_send_=0;
          }
          else {
            requests_to_send_=1;
          }
          errno = EAGAIN;
          return -1;
        }
      } else { //connected and handhsake finished
        requests_to_send_=1;
      }
    } else {
      if(requests_to_send_ > 0) { //report for initial data only
        slowlog(LOG_DEBUG, "SSL connection is using %s\n", SSL_get_cipher(ssl_));
      }
    }
  }
  // VA: this is not good. create a "prepare" method.
  // initial send
  if(buf_ == 0) {
    buf_ = buf;
    offset_ = len;
  }

  ret = ssl_ ? SSL_write(ssl_, buf_, offset_)
                 : send(sockfd_, buf_, offset_, 0);

  // entire data was sent
  if(ret > 0 && ret == offset_) {
    if(eInitialSend == type) {
      --requests_to_send_;
    } else if(eFollowUpSend == type) {
      --followups_to_send_;
    }
    buf_ = 0;
    offset_ = 0;
  } else if(ret > 0 && ret < offset_) {
    buf_ = static_cast<const char*>(buf_) + ret;
    offset_ -= ret;
  }
  return ret;
}

void SlowSocket::close() {
  if (-1 == sockfd_) return;

  slowlog(LOG_DEBUG, "closing slow, sock is %d\n", sockfd_);
  if(ssl_) {
    SSL_free(ssl_);
    ssl_ = NULL;
  }
  requests_to_send_ = 0;
  followups_to_send_ = 0;
  ::close(sockfd_);
  sockfd_ = -1;
}

void SlowSocket::set_state(SocketState state) {
  timeval t;
  gettimeofday(&t, 0);
  switch(state) {
    case eInit:
      break;
    case eConnecting:
      set_start(&t);
      break;
    case eConnected:
      set_connected(&t);
      break;
    case eError:
      break;
    case eClosed:
      set_stop(&t);
      break;
    default:
      break;
  }
  state_ = state;
}

}  // namespace slowhttptest
<|endoftext|>


<commit_before>/*
 * Copyright 2004-2015 Cray Inc.
 * Other additional copyright holders may be indicated within.
 *
 * The entirety of this work is 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 "ScopeBase.h"

#include "AstDumpToNode.h"
#include "ScopeModule.h"
#include "stmt.h"
#include "symbol.h"
#include "VisibleSymbol.h"

/************************************ | *************************************
*                                                                           *
*                                                                           *
*                                                                           *
************************************* | ************************************/

ScopeBase::ScopeBase(const ScopeBase* parent)
{
  mParent = parent;
}

ScopeBase::~ScopeBase()
{

}

void ScopeBase::extend(BlockStmt* block)
{
  std::vector<ModuleSymbol*> visited;

  extend(block, visited);
}

void ScopeBase::extend(BlockStmt*                  block,
                       std::vector<ModuleSymbol*>& visited)
{
  for_alist(stmt, block->body)
  {
    if (DefExpr* defExpr = toDefExpr(stmt))
    {
      addSym(defExpr->sym);
    }
    else if (CallExpr* callExpr = toCallExpr(stmt))
    {
      if (callExpr->isPrimitive(PRIM_USE) == true)
        addUse(callExpr->get(1));
    }
  }
}

void ScopeBase::addSym(Symbol* symbol)
{
  if (isConflicted(symbol) == false)
    mSymbols.push_back(symbol);
}

void ScopeBase::addUse(Expr* expr)
{
  if (ModuleSymbol* module = moduleByName(expr))
  {
    addUse(module);
  }
  else
  {
    AstDumpToNode logger(stdout, 3);

    printf("ScopeBase::extend: Failed to find module for\n");
    printf("   ");
    expr->accept(&logger);
    printf("\n\n");
  }
}

void ScopeBase::addUse(ModuleSymbol* module)
{
  std::vector<ModuleSymbol*> visited;

  addUse(module, visited);
}

void ScopeBase::addUse(ModuleSymbol*               module,
                       std::vector<ModuleSymbol*>& visited)
{
  if (isConflicted(module)       == false &&
      isVisited(module, visited) == false)
  {
    ScopeModule* usedScope = new ScopeModule(module);

    mUsedScopes.push_back(usedScope);

    visited.push_back(module);
  }
}

ModuleSymbol* ScopeBase::moduleByName(Expr* expr) const
{
  ModuleSymbol* retval = 0;

  if (UnresolvedSymExpr* unres = toUnresolvedSymExpr(expr))
  {
    std::vector<VisibleSymbol> symbols;

    visibleSymbols(unres, symbols);

    if (symbols.size() == 0)
    {
      const char* name = unres->unresolved;

      forv_Vec(ModuleSymbol, mod, gModuleSymbols)
      {
        if (strcmp(mod->name, name) == 0)
        {
          retval = mod;
          break;
        }
      }
    }
    else
      retval = toModuleSymbol(symbols[0].symbol());
  }

  return retval;
}

bool ScopeBase::isVisited(ModuleSymbol*               module,
                          std::vector<ModuleSymbol*>& visited)
{
  bool retval = false;

  for (size_t i = 0; i < visited.size() && retval == false; i++)
    retval = (visited[i] == module) ? true : false;

  return retval;
}

// A scope can have many definitions for a function but
// only one definition for other symbols
bool ScopeBase::isConflicted(Symbol* sym) const
{
  bool retval = false;

  if (isFnSymbol(sym) == false)
  {
    for (size_t i = 0; i < mSymbols.size() && retval == false; i++)
    {
      if (mSymbols[i] == sym)
      {
        AstDumpToNode logger(stdout, 3);

        printf("Error: Symbol is redefined\n");
        printf("   ");
        sym->accept(&logger);
        printf("\n");

        retval = true;
      }
    }
  }

  return retval;
}

// It would be surprising if a module were 'used' more than
// once in a single scope
bool ScopeBase::isConflicted(ModuleSymbol* mod) const
{
  bool retval = false;

  for (size_t i = 0; i < mUsedScopes.size() && retval == false; i++)
  {
    if (mUsedScopes[i]->isScopeForModule(mod) == true)
    {
      AstDumpToNode logger(stdout, 3);

      printf("Warning: Module is 'used' more than once in one scope\n");
      printf("   ");
      mod->accept(&logger);
      printf("\n");

      retval = true;
    }
  }

  return retval;
}

/************************************ | *************************************
*                                                                           *
*                                                                           *
*                                                                           *
************************************* | ************************************/

void ScopeBase::visibleSymbols(UnresolvedSymExpr*          unres,
                               std::vector<VisibleSymbol>& symbols) const
{
  visibleSymbols(unres, symbols, 0);
}

void ScopeBase::visibleSymbols(UnresolvedSymExpr*          unres,
                               std::vector<VisibleSymbol>& symbols,
                               int                         distance)  const
{
  const char* needle = unres->unresolved;

  for (size_t i = 0; i < mSymbols.size(); i++)
  {
    if (mSymbols[i]->name == needle)
    {
      VisibleSymbol vis(mSymbols[i], this, distance);

      symbols.push_back(vis);
    }
  }

  for (size_t i = 0; i < mUsedScopes.size(); i++)
  {
    mUsedScopes[i]->visibleSymbols(unres, symbols, distance + 1);
  }

  if (mParent != 0)
  {
    mParent->visibleSymbols(unres, symbols, distance + 1);
  }
}

/************************************ | *************************************
*                                                                           *
*                                                                           *
*                                                                           *
************************************* | ************************************/

void ScopeBase::describe(bool recursive) const
{
  printf("#<Scope ");
  describeHeader(stdout);
  printf("\n");

  printf("   symbols:\n");

  for (size_t i = 0; i < mSymbols.size(); i++)
    printf("       %3ld: %s\n", i, mSymbols[i]->name);

  if (mUsedScopes.size() > 0)
  {
    printf("\n");
    printf("   modUses:\n");

    for (size_t i = 0; i < mUsedScopes.size(); i++)
      printf("       %3ld: %s\n", i, mUsedScopes[i]->name());
  }

  printf(">\n\n");

  if (recursive == true)
  {
    for (size_t i = 0; i < mUsedScopes.size(); i++)
      mUsedScopes[i]->describe(true);

    if (mParent != 0)
      mParent->describe(true);
  }
}
<commit_msg>Fix size_t vs int issue for linux32 smoke test<commit_after>/*
 * Copyright 2004-2015 Cray Inc.
 * Other additional copyright holders may be indicated within.
 *
 * The entirety of this work is 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 "ScopeBase.h"

#include "AstDumpToNode.h"
#include "ScopeModule.h"
#include "stmt.h"
#include "symbol.h"
#include "VisibleSymbol.h"

/************************************ | *************************************
*                                                                           *
*                                                                           *
*                                                                           *
************************************* | ************************************/

ScopeBase::ScopeBase(const ScopeBase* parent)
{
  mParent = parent;
}

ScopeBase::~ScopeBase()
{

}

void ScopeBase::extend(BlockStmt* block)
{
  std::vector<ModuleSymbol*> visited;

  extend(block, visited);
}

void ScopeBase::extend(BlockStmt*                  block,
                       std::vector<ModuleSymbol*>& visited)
{
  for_alist(stmt, block->body)
  {
    if (DefExpr* defExpr = toDefExpr(stmt))
    {
      addSym(defExpr->sym);
    }
    else if (CallExpr* callExpr = toCallExpr(stmt))
    {
      if (callExpr->isPrimitive(PRIM_USE) == true)
        addUse(callExpr->get(1));
    }
  }
}

void ScopeBase::addSym(Symbol* symbol)
{
  if (isConflicted(symbol) == false)
    mSymbols.push_back(symbol);
}

void ScopeBase::addUse(Expr* expr)
{
  if (ModuleSymbol* module = moduleByName(expr))
  {
    addUse(module);
  }
  else
  {
    AstDumpToNode logger(stdout, 3);

    printf("ScopeBase::extend: Failed to find module for\n");
    printf("   ");
    expr->accept(&logger);
    printf("\n\n");
  }
}

void ScopeBase::addUse(ModuleSymbol* module)
{
  std::vector<ModuleSymbol*> visited;

  addUse(module, visited);
}

void ScopeBase::addUse(ModuleSymbol*               module,
                       std::vector<ModuleSymbol*>& visited)
{
  if (isConflicted(module)       == false &&
      isVisited(module, visited) == false)
  {
    ScopeModule* usedScope = new ScopeModule(module);

    mUsedScopes.push_back(usedScope);

    visited.push_back(module);
  }
}

ModuleSymbol* ScopeBase::moduleByName(Expr* expr) const
{
  ModuleSymbol* retval = 0;

  if (UnresolvedSymExpr* unres = toUnresolvedSymExpr(expr))
  {
    std::vector<VisibleSymbol> symbols;

    visibleSymbols(unres, symbols);

    if (symbols.size() == 0)
    {
      const char* name = unres->unresolved;

      forv_Vec(ModuleSymbol, mod, gModuleSymbols)
      {
        if (strcmp(mod->name, name) == 0)
        {
          retval = mod;
          break;
        }
      }
    }
    else
      retval = toModuleSymbol(symbols[0].symbol());
  }

  return retval;
}

bool ScopeBase::isVisited(ModuleSymbol*               module,
                          std::vector<ModuleSymbol*>& visited)
{
  bool retval = false;

  for (size_t i = 0; i < visited.size() && retval == false; i++)
    retval = (visited[i] == module) ? true : false;

  return retval;
}

// A scope can have many definitions for a function but
// only one definition for other symbols
bool ScopeBase::isConflicted(Symbol* sym) const
{
  bool retval = false;

  if (isFnSymbol(sym) == false)
  {
    for (size_t i = 0; i < mSymbols.size() && retval == false; i++)
    {
      if (mSymbols[i] == sym)
      {
        AstDumpToNode logger(stdout, 3);

        printf("Error: Symbol is redefined\n");
        printf("   ");
        sym->accept(&logger);
        printf("\n");

        retval = true;
      }
    }
  }

  return retval;
}

// It would be surprising if a module were 'used' more than
// once in a single scope
bool ScopeBase::isConflicted(ModuleSymbol* mod) const
{
  bool retval = false;

  for (size_t i = 0; i < mUsedScopes.size() && retval == false; i++)
  {
    if (mUsedScopes[i]->isScopeForModule(mod) == true)
    {
      AstDumpToNode logger(stdout, 3);

      printf("Warning: Module is 'used' more than once in one scope\n");
      printf("   ");
      mod->accept(&logger);
      printf("\n");

      retval = true;
    }
  }

  return retval;
}

/************************************ | *************************************
*                                                                           *
*                                                                           *
*                                                                           *
************************************* | ************************************/

void ScopeBase::visibleSymbols(UnresolvedSymExpr*          unres,
                               std::vector<VisibleSymbol>& symbols) const
{
  visibleSymbols(unres, symbols, 0);
}

void ScopeBase::visibleSymbols(UnresolvedSymExpr*          unres,
                               std::vector<VisibleSymbol>& symbols,
                               int                         distance)  const
{
  const char* needle = unres->unresolved;

  for (size_t i = 0; i < mSymbols.size(); i++)
  {
    if (mSymbols[i]->name == needle)
    {
      VisibleSymbol vis(mSymbols[i], this, distance);

      symbols.push_back(vis);
    }
  }

  for (size_t i = 0; i < mUsedScopes.size(); i++)
  {
    mUsedScopes[i]->visibleSymbols(unres, symbols, distance + 1);
  }

  if (mParent != 0)
  {
    mParent->visibleSymbols(unres, symbols, distance + 1);
  }
}

/************************************ | *************************************
*                                                                           *
*                                                                           *
*                                                                           *
************************************* | ************************************/

void ScopeBase::describe(bool recursive) const
{
  printf("#<Scope ");
  describeHeader(stdout);
  printf("\n");

  printf("   symbols:\n");

  for (size_t i = 0; i < mSymbols.size(); i++)
    printf("       %3d: %s\n", (int) i, mSymbols[i]->name);

  if (mUsedScopes.size() > 0)
  {
    printf("\n");
    printf("   modUses:\n");

    for (size_t i = 0; i < mUsedScopes.size(); i++)
      printf("       %3d: %s\n", (int) i, mUsedScopes[i]->name());
  }

  printf(">\n\n");

  if (recursive == true)
  {
    for (size_t i = 0; i < mUsedScopes.size(); i++)
      mUsedScopes[i]->describe(true);

    if (mParent != 0)
      mParent->describe(true);
  }
}
<|endoftext|>


<commit_before>#ifndef ALEPH_TOPOLOGY_MESH_HH__
#define ALEPH_TOPOLOGY_MESH_HH__

#include <cassert>

#include <algorithm>
#include <iterator>
#include <memory>
#include <vector>

namespace aleph
{

namespace topology
{

/**
  @class Mesh
  @brief Half-edge mesh data structure

  This data structure is capable of representing two-dimensional piecewise
  linear manifolds. In order to speed up standard queries, this class uses
  a standard half-edge data structure.
*/

template <class Position = float, class Data = float> class Mesh
{
public:
  struct Face;
  struct HalfEdge;
  struct Vertex;

  using Index           = std::size_t;

  using FacePointer     = std::shared_ptr<Face>;
  using HalfEdgePointer = std::shared_ptr<HalfEdge>;
  using VertexPointer   = std::shared_ptr<Vertex>;

  struct HalfEdge
  {
    FacePointer   face;
    VertexPointer vertex;

    HalfEdgePointer next; // Next half-edge (counter-clockwise)
    HalfEdgePointer prev; // Previous half-edge
    HalfEdgePointer pair; // Opposite half-edge

    VertexPointer source() const noexcept
    {
      return pair->vertex;
    }

    VertexPointer target() const noexcept
    {
      return vertex;
    }
  };

  struct Face
  {
    HalfEdgePointer edge;
  };

  struct Vertex
  {
    Position x = Position();
    Position y = Position();
    Position z = Position();
    Data     d = Data();

    HalfEdgePointer edge;
  };

  // Mesh attributes ---------------------------------------------------

  std::size_t vertices() const noexcept
  {
    return _vertices.size();
  }

  std::size_t faces() const noexcept
  {
    // TODO: Not yet implemented
    return 0;
  }

  // Mesh modification -------------------------------------------------

  /** Adds a new vertex to the mesh */
  void addVertex( Position x, Position y, Position z, Data d = Data() )
  {
    Vertex v;

    v.x = x;
    v.y = y;
    v.z = z;
    v.d = d;

    _vertices.push_back( std::make_shared<Vertex>( v ) );
  }

  /**
    Adds a new face to the mesh. This function expects a range of vertex IDs
    that make up the face. The vertices of the face need to sorted correctly
    in order for the orientation to be consistent.
  */

  template <class InputIterator> void addFace( InputIterator begin, InputIterator end )
  {
    FacePointer face = std::make_shared<Face>();

    // Stores all half-edges created (or found) by this function in the
    // order in which they belong to the face.
    std::vector<HalfEdgePointer> edges;
    edges.reserve( std::distance( begin, end ) );

    for( InputIterator it = begin; it != end; ++it )
    {
      auto curr = it;
      auto next = std::next( it );

      if( next == end )
        next = begin;

      auto source = _vertices.at( *curr );   // Edge source vertex
      auto target = _vertices.at( *next );   // Edge target vertex
      auto edge   = getEdge( *curr, *next ); // Edge

      if( !edge )
      {
        edge      = std::make_shared<HalfEdge>();
        auto pair = std::make_shared<HalfEdge>();

        edge->face   = face;
        edge->pair   = pair;
        edge->vertex = target;
        pair->vertex = source;
        pair->pair   = edge;

        if( !source->edge )
          source->edge = edge;

        if( !target->edge )
          target->edge = pair;
      }
      else
      {
        assert( !edge->face );
        edge->face = face;
      }

      edges.push_back( edge );
    }

    // Set 'next' and 'prev' pointers correctly ------------------------
    //
    // We first traverse all edges that bound the current face. Here, it
    // should be possible to traverse the face directly, so we require a
    // proper pointer in both directions.

    for( auto itEdge = edges.begin(); itEdge != edges.end(); ++itEdge )
    {
      auto curr = itEdge;
      auto prev = std::prev( curr );
      auto next = std::next( curr );

      if( curr == edges.begin() )
        prev = std::prev( edges.end() );

      if( next == edges.end() )
        next = edges.begin();

      auto&& edge = *itEdge;

      assert( !edge->next );
      assert( !edge->prev );

      edge->next = *next;
      edge->prev = *prev;
    }
  }

  // Mesh queries ------------------------------------------------------

  std::vector<VertexPointer> getLowerNeighbours( const Vertex& v ) const noexcept
  {
    auto neighbours = this->getNeighbours( v );
    auto d          = v.d;

    neighbours.erase( std::remove_if( neighbours.begin(), neighbours.end(),
                                      [&d] ( const VertexPointer& neighbour )
                                      {
                                        return neighbour->d >= d;
                                      } ),
                      neighbours.end() );

    return neighbours;
  }

  std::vector<VertexPointer> getHigherNeighbours( const Vertex& v ) const noexcept
  {
    auto neighbours = this->getNeighbours( v );
    auto d          = v.d;

    neighbours.erase( std::remove_if( neighbours.begin(), neighbours.end(),
                                      [&d] ( const VertexPointer& neighbour )
                                      {
                                        return neighbour->d <= d;
                                      } ),
                      neighbours.end() );

    return neighbours;
  }

  VertexPointer vertex( std::size_t id ) const
  {
    return _vertices.at( id );
  }

private:

  /** Gets all vertices that are adjacent to a given vertex */
  std::vector<VertexPointer> getNeighbours( const Vertex& v ) const noexcept
  {
    std::vector<VertexPointer> neighbours;

    auto edge = v.edge;
    do
    {
      if( edge )
        neighbours.push_back( edge->target() );
      else
        break;

      edge = edge->pair->next;
    }
    while( edge != v.edge );

    return neighbours;
  }

  /** Gets all edges that are incident on a given vertex. */
  std::vector<HalfEdgePointer> getEdges( const Vertex& v ) const noexcept
  {
    std::vector<HalfEdgePointer> edges;

    auto edge = v.edge;
    do
    {
      if( edge )
        edges.push_back( edge );
      else
        break;

      edge = edge->pair->next;
    }
    while( edge != v.edge );

    return edges;
  }

  /**
    Check whether a given (directed) edge already exists. If so,
    a pointer to the edge is being returned.
  */

  HalfEdgePointer getEdge( Index u, Index v ) const noexcept
  {
    auto source = _vertices.at(u);           // Edge source vertex
    auto target = _vertices.at(v);           // Edge target vertex
    auto edges  = this->getEdges( *source ); // Incident edges

    auto itEdge = std::find_if( edges.begin(), edges.end(),
                                [&source, &target] ( const HalfEdgePointer& edge )
                                {
                                  return edge->source() == source && edge->target() == target;
                                } );

    if( itEdge != edges.end() )
      return *itEdge;
    else
      return nullptr;
  }

  /**
    Stores all vertex pointers. This is sufficient to store the
    complete mesh.
  */

  std::vector<VertexPointer> _vertices;
};

} // namespace topology

} // namespace aleph

#endif
<commit_msg>Implemented vertices & faces query for meshes<commit_after>#ifndef ALEPH_TOPOLOGY_MESH_HH__
#define ALEPH_TOPOLOGY_MESH_HH__

#include <cassert>

#include <algorithm>
#include <iterator>
#include <memory>
#include <unordered_set>
#include <vector>

namespace aleph
{

namespace topology
{

/**
  @class Mesh
  @brief Half-edge mesh data structure

  This data structure is capable of representing two-dimensional piecewise
  linear manifolds. In order to speed up standard queries, this class uses
  a standard half-edge data structure.
*/

template <class Position = float, class Data = float> class Mesh
{
public:
  struct Face;
  struct HalfEdge;
  struct Vertex;

  using Index           = std::size_t;

  using FacePointer     = std::shared_ptr<Face>;
  using HalfEdgePointer = std::shared_ptr<HalfEdge>;
  using VertexPointer   = std::shared_ptr<Vertex>;

  struct HalfEdge
  {
    FacePointer   face;
    VertexPointer vertex;

    HalfEdgePointer next; // Next half-edge (counter-clockwise)
    HalfEdgePointer prev; // Previous half-edge
    HalfEdgePointer pair; // Opposite half-edge

    VertexPointer source() const noexcept
    {
      return pair->vertex;
    }

    VertexPointer target() const noexcept
    {
      return vertex;
    }
  };

  struct Face
  {
    HalfEdgePointer edge;
  };

  struct Vertex
  {
    Position x = Position();
    Position y = Position();
    Position z = Position();
    Data     d = Data();

    HalfEdgePointer edge;
  };

  // Mesh attributes ---------------------------------------------------

  std::size_t vertices() const noexcept
  {
    return _vertices.size();
  }

  std::size_t faces() const noexcept
  {
    std::unordered_set<FacePointer> faces;

    for( auto&& vertex : _vertices )
    {
      auto&& edge = vertex->edge;
      auto&& face = edge->face;

      faces.insert( face );
    }

    return faces.size();
  }

  // Mesh modification -------------------------------------------------

  /** Adds a new vertex to the mesh */
  void addVertex( Position x, Position y, Position z, Data d = Data() )
  {
    Vertex v;

    v.x = x;
    v.y = y;
    v.z = z;
    v.d = d;

    _vertices.push_back( std::make_shared<Vertex>( v ) );
  }

  /**
    Adds a new face to the mesh. This function expects a range of vertex IDs
    that make up the face. The vertices of the face need to sorted correctly
    in order for the orientation to be consistent.
  */

  template <class InputIterator> void addFace( InputIterator begin, InputIterator end )
  {
    FacePointer face = std::make_shared<Face>();

    // Stores all half-edges created (or found) by this function in the
    // order in which they belong to the face.
    std::vector<HalfEdgePointer> edges;
    edges.reserve( std::distance( begin, end ) );

    for( InputIterator it = begin; it != end; ++it )
    {
      auto curr = it;
      auto next = std::next( it );

      if( next == end )
        next = begin;

      auto source = _vertices.at( *curr );   // Edge source vertex
      auto target = _vertices.at( *next );   // Edge target vertex
      auto edge   = getEdge( *curr, *next ); // Edge

      if( !edge )
      {
        edge      = std::make_shared<HalfEdge>();
        auto pair = std::make_shared<HalfEdge>();

        edge->face   = face;
        edge->pair   = pair;
        edge->vertex = target;
        pair->vertex = source;
        pair->pair   = edge;

        if( !source->edge )
          source->edge = edge;

        if( !target->edge )
          target->edge = pair;
      }
      else
      {
        assert( !edge->face );
        edge->face = face;
      }

      edges.push_back( edge );
    }

    // Set 'next' and 'prev' pointers correctly ------------------------
    //
    // We first traverse all edges that bound the current face. Here, it
    // should be possible to traverse the face directly, so we require a
    // proper pointer in both directions.

    for( auto itEdge = edges.begin(); itEdge != edges.end(); ++itEdge )
    {
      auto curr = itEdge;
      auto prev = std::prev( curr );
      auto next = std::next( curr );

      if( curr == edges.begin() )
        prev = std::prev( edges.end() );

      if( next == edges.end() )
        next = edges.begin();

      auto&& edge = *itEdge;

      assert( !edge->next );
      assert( !edge->prev );

      edge->next = *next;
      edge->prev = *prev;
    }
  }

  // Mesh queries ------------------------------------------------------

  std::vector<VertexPointer> getLowerNeighbours( const Vertex& v ) const noexcept
  {
    auto neighbours = this->getNeighbours( v );
    auto d          = v.d;

    neighbours.erase( std::remove_if( neighbours.begin(), neighbours.end(),
                                      [&d] ( const VertexPointer& neighbour )
                                      {
                                        return neighbour->d >= d;
                                      } ),
                      neighbours.end() );

    return neighbours;
  }

  std::vector<VertexPointer> getHigherNeighbours( const Vertex& v ) const noexcept
  {
    auto neighbours = this->getNeighbours( v );
    auto d          = v.d;

    neighbours.erase( std::remove_if( neighbours.begin(), neighbours.end(),
                                      [&d] ( const VertexPointer& neighbour )
                                      {
                                        return neighbour->d <= d;
                                      } ),
                      neighbours.end() );

    return neighbours;
  }

  VertexPointer vertex( std::size_t id ) const
  {
    return _vertices.at( id );
  }

private:

  /** Gets all vertices that are adjacent to a given vertex */
  std::vector<VertexPointer> getNeighbours( const Vertex& v ) const noexcept
  {
    std::vector<VertexPointer> neighbours;

    auto edge = v.edge;
    do
    {
      if( edge )
        neighbours.push_back( edge->target() );
      else
        break;

      edge = edge->pair->next;
    }
    while( edge != v.edge );

    return neighbours;
  }

  /** Gets all edges that are incident on a given vertex. */
  std::vector<HalfEdgePointer> getEdges( const Vertex& v ) const noexcept
  {
    std::vector<HalfEdgePointer> edges;

    auto edge = v.edge;
    do
    {
      if( edge )
        edges.push_back( edge );
      else
        break;

      edge = edge->pair->next;
    }
    while( edge != v.edge );

    return edges;
  }

  /**
    Check whether a given (directed) edge already exists. If so,
    a pointer to the edge is being returned.
  */

  HalfEdgePointer getEdge( Index u, Index v ) const noexcept
  {
    auto source = _vertices.at(u);           // Edge source vertex
    auto target = _vertices.at(v);           // Edge target vertex
    auto edges  = this->getEdges( *source ); // Incident edges

    auto itEdge = std::find_if( edges.begin(), edges.end(),
                                [&source, &target] ( const HalfEdgePointer& edge )
                                {
                                  return edge->source() == source && edge->target() == target;
                                } );

    if( itEdge != edges.end() )
      return *itEdge;
    else
      return nullptr;
  }

  /**
    Stores all vertex pointers. This is sufficient to store the
    complete mesh.
  */

  std::vector<VertexPointer> _vertices;
};

} // namespace topology

} // namespace aleph

#endif
<|endoftext|>


<commit_before>#ifndef _NODO_BINARY_TREE_HPP_
#define _NODO_BINARY_TREE_HPP_
#include <iostream>

template<class T>
class NodeBT
{
	private:
		T _key;
		NodeBT<T> *_left, *_right;

	public:
		NodeBT() : _key(0), _left(NULL), _right(NULL) { };
		NodeBT(const T in) : _key(in), _left(NULL), _right(NULL) { };
		NodeBT(const NodeBT<T>& in) : _key(in.key), _left(NULL), _right(NULL) { };

		T getKey() const { return _key; }
		NodeBT<T>* getLeft() const { return _left; }
		NodeBT<T>* getRight() const { return _right; }
		bool isLeaf() const { return _left == NULL && _right == NULL; }
		bool isFullNode() const { return _left != NULL || _right != NULL; }

		void setKey(const T);
		void setLeft(const NodeBT<T>*);
		void setRight(const NodeBT<T>*);
};

/**
 * Key setter
 * */
template<class T>
void NodeBT<T>::setKey(const T item)
{
	_key = item;
}

/**
 * Left setter
 * */
template<class T>
void NodeBT<T>::setLeft(const NodeBT<T>* l)
{
	_left = l;
}

/**
 * Right setter
 * */
template<class T>
void NodeBT<T>::setRight(const NodeBT<T>* r)
{
	_right = r;
}

#endif
<commit_msg>Constructor added<commit_after>#ifndef _NODO_BINARY_TREE_HPP_
#define _NODO_BINARY_TREE_HPP_
#include <iostream>

template<class T>
class NodeBT
{
	private:
		T _key;
		NodeBT<T> *_left, *_right;

	public:
		NodeBT() : _key(0), _left(NULL), _right(NULL) { };
		NodeBT(const T in) : _key(in), _left(NULL), _right(NULL) { };
		NodeBT(const T in, NodeBT<T>* l, NodeBT<T>* r) : _key(in), _left(l), _right(r) { };
		NodeBT(const NodeBT<T>& in) : _key(in.key), _left(NULL), _right(NULL) { };

		T getKey() const { return _key; }
		NodeBT<T>* getLeft() const { return _left; }
		NodeBT<T>* getRight() const { return _right; }
		bool isLeaf() const { return _left == NULL && _right == NULL; }
		bool isFullNode() const { return _left != NULL || _right != NULL; }

		void setKey(const T);
		void setLeft(const NodeBT<T>*);
		void setRight(const NodeBT<T>*);
};

/**
 * Key setter
 * */
template<class T>
void NodeBT<T>::setKey(const T item)
{
	_key = item;
}

/**
 * Left setter
 * */
template<class T>
void NodeBT<T>::setLeft(const NodeBT<T>* l)
{
	_left = l;
}

/**
 * Right setter
 * */
template<class T>
void NodeBT<T>::setRight(const NodeBT<T>* r)
{
	_right = r;
}

#endif
<|endoftext|>