<commit_before>// Copyright 2021 TF.Text Authors.
//
// 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 <limits>
#include <memory>
#include <string>
#include <vector>

#include "tensorflow/core/framework/lookup_interface.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow_text/core/kernels/wordpiece_tokenizer.h"

namespace tensorflow {
namespace text {

namespace {
string GetWordSplitChar(OpKernelConstruction* ctx) {
  string suffix_indicator;
  ([=](string* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("suffix_indicator", c));
  })(&suffix_indicator);
  return suffix_indicator;
}

int32 GetMaxCharsPerWord(OpKernelConstruction* ctx) {
  int32 max_chars_per_word;
  ([=](int32* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("max_bytes_per_word", c));
  })(&max_chars_per_word);
  return max_chars_per_word;
}

int32 GetMaxCharsPerToken(OpKernelConstruction* ctx) {
  int32 max_chars_per_token;
  ([=](int32* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("max_chars_per_token", c));
  })(&max_chars_per_token);
  return max_chars_per_token;
}

bool GetShouldUseUnknownToken(OpKernelConstruction* ctx) {
  bool use_unknown_token;
  ([=](bool* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("use_unknown_token", c));
  })(&use_unknown_token);
  return use_unknown_token;
}

string GetUnknownToken(OpKernelConstruction* ctx) {
  string unknown_token;
  ([=](string* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("unknown_token", c));
  })(&unknown_token);
  return unknown_token;
}

bool GetSplitUnknownCharacters(OpKernelConstruction* ctx) {
  bool split_unknown_characters;
  ([=](bool* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("split_unknown_characters", c));
  })(&split_unknown_characters);
  return split_unknown_characters;
}

Status GetTableHandle(const string& input_name, OpKernelContext* ctx,
                      string* container, string* table_handle) {
  {
    mutex* mu;
    TF_RETURN_IF_ERROR(ctx->input_ref_mutex(input_name, &mu));
    mutex_lock l(*mu);
    Tensor tensor;
    TF_RETURN_IF_ERROR(ctx->mutable_input(input_name, &tensor, true));
    if (tensor.NumElements() != 2) {
      return errors::InvalidArgument(
          "Lookup table handle must be scalar, but had shape: ",
          tensor.shape().DebugString());
    }
    auto h = tensor.flat<tstring>();
    *container = h(0);
    *table_handle = h(1);
  }
  return Status::OK();
}

// Gets the LookupTable stored in the ctx->resource_manager() with key
// passed by attribute with name input_name, returns null if the table
// doesn't exist.
Status GetLookupTable(const string& input_name, OpKernelContext* ctx,
                      lookup::LookupInterface** table) {
  string container;
  string table_handle;
  DataType handle_dtype;
  TF_RETURN_IF_ERROR(ctx->input_dtype(input_name, &handle_dtype));
  if (handle_dtype == DT_RESOURCE) {
    ResourceHandle handle;
    TF_RETURN_IF_ERROR(HandleFromInput(ctx, input_name, &handle));
    return LookupResource(ctx, handle, table);
  } else {
    TF_RETURN_IF_ERROR(
        GetTableHandle(input_name, ctx, &container, &table_handle));
    return ctx->resource_manager()->Lookup(container, table_handle, table);
  }
}

class LookupTableVocab : public WordpieceVocab {
 public:
  LookupTableVocab(lookup::LookupInterface* table, OpKernelContext* ctx);

  virtual LookupStatus Contains(const absl::string_view key, bool* value) const;

 private:
  // not owned
  mutable lookup::LookupInterface* table_;
  OpKernelContext* ctx_;
  Tensor default_value_;
};

Status ToStatus(const LookupStatus& status) {
  if (status.success) {
    return Status::OK();
  }

  return errors::InvalidArgument(status.error_msg);
}

constexpr int64 kOutOfVocabValue = -1;

LookupTableVocab::LookupTableVocab(lookup::LookupInterface* table,
                                   OpKernelContext* ctx)
    : table_(table), ctx_(ctx), default_value_(DT_INT64, TensorShape({1})) {
  default_value_.flat<int64>()(0) = kOutOfVocabValue;
}

LookupStatus LookupTableVocab::Contains(const absl::string_view key,
                                        bool* value) const {
  if (value == nullptr) {
    return LookupStatus("Bad 'value' param.");
  }
  Tensor keys(DT_STRING, TensorShape({1}));
  keys.flat<tstring>()(0) = tstring(key.data(), key.size());
  Tensor values(DT_INT64, TensorShape({1}));
  auto status = table_->Find(ctx_, keys, &values, default_value_);
  if (!status.ok()) return LookupStatus(status.error_message());

  if (static_cast<int64>(values.flat<int64>()(0)) != kOutOfVocabValue) {
    *value = true;
    return LookupStatus::OK();
  }
  *value = false;
  return LookupStatus::OK();
}

}  // namespace

class WordpieceTokenizeWithOffsetsOp : public OpKernel {
 public:
  explicit WordpieceTokenizeWithOffsetsOp(OpKernelConstruction* ctx)
      : OpKernel(ctx),
        suffix_indicator_(GetWordSplitChar(ctx)),
        max_bytes_per_word_(GetMaxCharsPerWord(ctx)),
        max_chars_per_token_(GetMaxCharsPerToken(ctx)),
        use_unknown_token_(GetShouldUseUnknownToken(ctx)),
        unknown_token_(GetUnknownToken(ctx)),
        split_unknown_characters_(GetSplitUnknownCharacters(ctx)) {
    string output_row_partition_type;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_row_partition_type",
                                     &output_row_partition_type));
    if (output_row_partition_type == "row_lengths") {
      row_partition_type_ = ROW_LENGTHS;
    } else if (output_row_partition_type == "row_splits") {
      row_partition_type_ = ROW_SPLITS;
    } else {
      OP_REQUIRES(
          ctx, false,
          errors::Internal("Unexpected value for output_row_partition_type"));
    }
  }

  void Compute(OpKernelContext* ctx) override {
    const Tensor* input_values;
    OP_REQUIRES_OK(ctx, ctx->input("input_values", &input_values));
    const auto& values_vec = input_values->flat<tstring>();

    lookup::LookupInterface* lookup_table;
    OP_REQUIRES_OK(ctx,
                   GetLookupTable("vocab_lookup_table", ctx, &lookup_table));
    core::ScopedUnref unref_me(lookup_table);
    LookupTableVocab vocab_map(lookup_table, ctx);

    std::vector<string> subwords;
    std::vector<int> begin_offset;
    std::vector<int> end_offset;
    std::vector<int> row_partition;

    if (row_partition_type_ == ROW_SPLITS) {
      row_partition.push_back(0);
    }

    // Iterate through all the values and wordpiece tokenize them.
    for (int i = 0; i < values_vec.size(); ++i) {
      // Tokenize into subwords and record the offset locations.
      int num_wordpieces = 0;
      OP_REQUIRES_OK(
          ctx, ToStatus(WordpieceTokenize(
                   values_vec(i), max_bytes_per_word_, max_chars_per_token_,
                   suffix_indicator_, use_unknown_token_, unknown_token_,
                   split_unknown_characters_, &vocab_map, &subwords,
                   &begin_offset, &end_offset, &num_wordpieces)));

      // Record the row splits.
      switch (row_partition_type_) {
        case ROW_LENGTHS:
          row_partition.push_back(num_wordpieces);
          break;
        case ROW_SPLITS:
          row_partition.push_back(num_wordpieces + row_partition.back());
          break;
      }
    }

    std::vector<int64> output_subwords_shape;
    output_subwords_shape.push_back(subwords.size());

    std::vector<int64> output_row_partition_shape;
    output_row_partition_shape.push_back(row_partition.size());

    Tensor* output_values;
    OP_REQUIRES_OK(ctx, ctx->allocate_output("output_values",
                                             TensorShape(output_subwords_shape),
                                             &output_values));
    auto output_values_vec = output_values->vec<tstring>();

    Tensor* output_row_partition;
    OP_REQUIRES_OK(ctx,
                   ctx->allocate_output("output_row_lengths",
                                        TensorShape(output_row_partition_shape),
                                        &output_row_partition));
    auto output_row_partition_vec = output_row_partition->vec<int64>();

    Tensor* start_values;
    OP_REQUIRES_OK(ctx, ctx->allocate_output("start_values",
                                             TensorShape(output_subwords_shape),
                                             &start_values));
    auto start_values_vec = start_values->vec<int64>();

    Tensor* limit_values;
    OP_REQUIRES_OK(ctx, ctx->allocate_output("limit_values",
                                             TensorShape(output_subwords_shape),
                                             &limit_values));
    auto limit_values_vec = limit_values->vec<int64>();

    for (int i = 0; i < subwords.size(); ++i) {
      output_values_vec(i) = subwords[i];
    }

    for (int i = 0; i < row_partition.size(); ++i) {
      output_row_partition_vec(i) = row_partition[i];
    }

    for (int i = 0; i < begin_offset.size(); ++i) {
      start_values_vec(i) = begin_offset[i];
    }

    for (int i = 0; i < end_offset.size(); ++i) {
      limit_values_vec(i) = end_offset[i];
    }
  }

 private:
  enum RowPartitionType { ROW_LENGTHS, ROW_SPLITS };

  const string suffix_indicator_;
  const int max_bytes_per_word_;
  const int max_chars_per_token_;
  const bool use_unknown_token_;
  const string unknown_token_;
  const bool split_unknown_characters_;
  RowPartitionType row_partition_type_;

  TF_DISALLOW_COPY_AND_ASSIGN(WordpieceTokenizeWithOffsetsOp);
};

REGISTER_KERNEL_BUILDER(Name("WordpieceTokenizeWithOffsets").Device(DEVICE_CPU),
                        WordpieceTokenizeWithOffsetsOp);

}  // namespace text
}  // namespace tensorflow
<commit_msg>Add WordpieceTokenizeWithOffsets with ALLOW_STATEFUL_OP_FOR_DATASET_FUNCTIONS for tf.data<commit_after>// Copyright 2021 TF.Text Authors.
//
// 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 <limits>
#include <memory>
#include <string>
#include <vector>

#include "tensorflow/core/framework/dataset_stateful_op_allowlist.h"
#include "tensorflow/core/framework/lookup_interface.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow_text/core/kernels/wordpiece_tokenizer.h"

namespace tensorflow {
namespace text {

namespace {
string GetWordSplitChar(OpKernelConstruction* ctx) {
  string suffix_indicator;
  ([=](string* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("suffix_indicator", c));
  })(&suffix_indicator);
  return suffix_indicator;
}

int32 GetMaxCharsPerWord(OpKernelConstruction* ctx) {
  int32 max_chars_per_word;
  ([=](int32* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("max_bytes_per_word", c));
  })(&max_chars_per_word);
  return max_chars_per_word;
}

int32 GetMaxCharsPerToken(OpKernelConstruction* ctx) {
  int32 max_chars_per_token;
  ([=](int32* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("max_chars_per_token", c));
  })(&max_chars_per_token);
  return max_chars_per_token;
}

bool GetShouldUseUnknownToken(OpKernelConstruction* ctx) {
  bool use_unknown_token;
  ([=](bool* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("use_unknown_token", c));
  })(&use_unknown_token);
  return use_unknown_token;
}

string GetUnknownToken(OpKernelConstruction* ctx) {
  string unknown_token;
  ([=](string* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("unknown_token", c));
  })(&unknown_token);
  return unknown_token;
}

bool GetSplitUnknownCharacters(OpKernelConstruction* ctx) {
  bool split_unknown_characters;
  ([=](bool* c) -> void {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("split_unknown_characters", c));
  })(&split_unknown_characters);
  return split_unknown_characters;
}

Status GetTableHandle(const string& input_name, OpKernelContext* ctx,
                      string* container, string* table_handle) {
  {
    mutex* mu;
    TF_RETURN_IF_ERROR(ctx->input_ref_mutex(input_name, &mu));
    mutex_lock l(*mu);
    Tensor tensor;
    TF_RETURN_IF_ERROR(ctx->mutable_input(input_name, &tensor, true));
    if (tensor.NumElements() != 2) {
      return errors::InvalidArgument(
          "Lookup table handle must be scalar, but had shape: ",
          tensor.shape().DebugString());
    }
    auto h = tensor.flat<tstring>();
    *container = h(0);
    *table_handle = h(1);
  }
  return Status::OK();
}

// Gets the LookupTable stored in the ctx->resource_manager() with key
// passed by attribute with name input_name, returns null if the table
// doesn't exist.
Status GetLookupTable(const string& input_name, OpKernelContext* ctx,
                      lookup::LookupInterface** table) {
  string container;
  string table_handle;
  DataType handle_dtype;
  TF_RETURN_IF_ERROR(ctx->input_dtype(input_name, &handle_dtype));
  if (handle_dtype == DT_RESOURCE) {
    ResourceHandle handle;
    TF_RETURN_IF_ERROR(HandleFromInput(ctx, input_name, &handle));
    return LookupResource(ctx, handle, table);
  } else {
    TF_RETURN_IF_ERROR(
        GetTableHandle(input_name, ctx, &container, &table_handle));
    return ctx->resource_manager()->Lookup(container, table_handle, table);
  }
}

class LookupTableVocab : public WordpieceVocab {
 public:
  LookupTableVocab(lookup::LookupInterface* table, OpKernelContext* ctx);

  virtual LookupStatus Contains(const absl::string_view key, bool* value) const;

 private:
  // not owned
  mutable lookup::LookupInterface* table_;
  OpKernelContext* ctx_;
  Tensor default_value_;
};

Status ToStatus(const LookupStatus& status) {
  if (status.success) {
    return Status::OK();
  }

  return errors::InvalidArgument(status.error_msg);
}

constexpr int64 kOutOfVocabValue = -1;

LookupTableVocab::LookupTableVocab(lookup::LookupInterface* table,
                                   OpKernelContext* ctx)
    : table_(table), ctx_(ctx), default_value_(DT_INT64, TensorShape({1})) {
  default_value_.flat<int64>()(0) = kOutOfVocabValue;
}

LookupStatus LookupTableVocab::Contains(const absl::string_view key,
                                        bool* value) const {
  if (value == nullptr) {
    return LookupStatus("Bad 'value' param.");
  }
  Tensor keys(DT_STRING, TensorShape({1}));
  keys.flat<tstring>()(0) = tstring(key.data(), key.size());
  Tensor values(DT_INT64, TensorShape({1}));
  auto status = table_->Find(ctx_, keys, &values, default_value_);
  if (!status.ok()) return LookupStatus(status.error_message());

  if (static_cast<int64>(values.flat<int64>()(0)) != kOutOfVocabValue) {
    *value = true;
    return LookupStatus::OK();
  }
  *value = false;
  return LookupStatus::OK();
}

}  // namespace

class WordpieceTokenizeWithOffsetsOp : public OpKernel {
 public:
  explicit WordpieceTokenizeWithOffsetsOp(OpKernelConstruction* ctx)
      : OpKernel(ctx),
        suffix_indicator_(GetWordSplitChar(ctx)),
        max_bytes_per_word_(GetMaxCharsPerWord(ctx)),
        max_chars_per_token_(GetMaxCharsPerToken(ctx)),
        use_unknown_token_(GetShouldUseUnknownToken(ctx)),
        unknown_token_(GetUnknownToken(ctx)),
        split_unknown_characters_(GetSplitUnknownCharacters(ctx)) {
    string output_row_partition_type;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_row_partition_type",
                                     &output_row_partition_type));
    if (output_row_partition_type == "row_lengths") {
      row_partition_type_ = ROW_LENGTHS;
    } else if (output_row_partition_type == "row_splits") {
      row_partition_type_ = ROW_SPLITS;
    } else {
      OP_REQUIRES(
          ctx, false,
          errors::Internal("Unexpected value for output_row_partition_type"));
    }
  }

  void Compute(OpKernelContext* ctx) override {
    const Tensor* input_values;
    OP_REQUIRES_OK(ctx, ctx->input("input_values", &input_values));
    const auto& values_vec = input_values->flat<tstring>();

    lookup::LookupInterface* lookup_table;
    OP_REQUIRES_OK(ctx,
                   GetLookupTable("vocab_lookup_table", ctx, &lookup_table));
    core::ScopedUnref unref_me(lookup_table);
    LookupTableVocab vocab_map(lookup_table, ctx);

    std::vector<string> subwords;
    std::vector<int> begin_offset;
    std::vector<int> end_offset;
    std::vector<int> row_partition;

    if (row_partition_type_ == ROW_SPLITS) {
      row_partition.push_back(0);
    }

    // Iterate through all the values and wordpiece tokenize them.
    for (int i = 0; i < values_vec.size(); ++i) {
      // Tokenize into subwords and record the offset locations.
      int num_wordpieces = 0;
      OP_REQUIRES_OK(
          ctx, ToStatus(WordpieceTokenize(
                   values_vec(i), max_bytes_per_word_, max_chars_per_token_,
                   suffix_indicator_, use_unknown_token_, unknown_token_,
                   split_unknown_characters_, &vocab_map, &subwords,
                   &begin_offset, &end_offset, &num_wordpieces)));

      // Record the row splits.
      switch (row_partition_type_) {
        case ROW_LENGTHS:
          row_partition.push_back(num_wordpieces);
          break;
        case ROW_SPLITS:
          row_partition.push_back(num_wordpieces + row_partition.back());
          break;
      }
    }

    std::vector<int64> output_subwords_shape;
    output_subwords_shape.push_back(subwords.size());

    std::vector<int64> output_row_partition_shape;
    output_row_partition_shape.push_back(row_partition.size());

    Tensor* output_values;
    OP_REQUIRES_OK(ctx, ctx->allocate_output("output_values",
                                             TensorShape(output_subwords_shape),
                                             &output_values));
    auto output_values_vec = output_values->vec<tstring>();

    Tensor* output_row_partition;
    OP_REQUIRES_OK(ctx,
                   ctx->allocate_output("output_row_lengths",
                                        TensorShape(output_row_partition_shape),
                                        &output_row_partition));
    auto output_row_partition_vec = output_row_partition->vec<int64>();

    Tensor* start_values;
    OP_REQUIRES_OK(ctx, ctx->allocate_output("start_values",
                                             TensorShape(output_subwords_shape),
                                             &start_values));
    auto start_values_vec = start_values->vec<int64>();

    Tensor* limit_values;
    OP_REQUIRES_OK(ctx, ctx->allocate_output("limit_values",
                                             TensorShape(output_subwords_shape),
                                             &limit_values));
    auto limit_values_vec = limit_values->vec<int64>();

    for (int i = 0; i < subwords.size(); ++i) {
      output_values_vec(i) = subwords[i];
    }

    for (int i = 0; i < row_partition.size(); ++i) {
      output_row_partition_vec(i) = row_partition[i];
    }

    for (int i = 0; i < begin_offset.size(); ++i) {
      start_values_vec(i) = begin_offset[i];
    }

    for (int i = 0; i < end_offset.size(); ++i) {
      limit_values_vec(i) = end_offset[i];
    }
  }

 private:
  enum RowPartitionType { ROW_LENGTHS, ROW_SPLITS };

  const string suffix_indicator_;
  const int max_bytes_per_word_;
  const int max_chars_per_token_;
  const bool use_unknown_token_;
  const string unknown_token_;
  const bool split_unknown_characters_;
  RowPartitionType row_partition_type_;

  TF_DISALLOW_COPY_AND_ASSIGN(WordpieceTokenizeWithOffsetsOp);
};

REGISTER_KERNEL_BUILDER(Name("WordpieceTokenizeWithOffsets").Device(DEVICE_CPU),
                        WordpieceTokenizeWithOffsetsOp);
ALLOW_STATEFUL_OP_FOR_DATASET_FUNCTIONS("WordpieceTokenizeWithOffsets");

}  // namespace text
}  // namespace tensorflow
<|endoftext|>


<commit_before>// Copyright (c) 2017-2019 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 <wallet/walletutil.h>

#include <logging.h>
#include <util/system.h>

fs::path GetWalletDir()
{
    fs::path path;

    if (gArgs.IsArgSet("-walletdir")) {
        path = gArgs.GetArg("-walletdir", "");
        if (!fs::is_directory(path)) {
            // If the path specified doesn't exist, we return the deliberately
            // invalid empty string.
            path = "";
        }
    } else {
        path = GetDataDir();
        // If a wallets directory exists, use that, otherwise default to GetDataDir
        if (fs::is_directory(path / "wallets")) {
            path /= "wallets";
        }
    }

    return path;
}

static bool IsBerkeleyBtree(const fs::path& path)
{
    if (!fs::exists(path)) return false;

    // A Berkeley DB Btree file has at least 4K.
    // This check also prevents opening lock files.
    boost::system::error_code ec;
    auto size = fs::file_size(path, ec);
    if (ec) LogPrintf("%s: %s %s\n", __func__, ec.message(), path.string());
    if (size < 4096) return false;

    fsbridge::ifstream file(path, std::ios::binary);
    if (!file.is_open()) return false;

    file.seekg(12, std::ios::beg); // Magic bytes start at offset 12
    uint32_t data = 0;
    file.read((char*) &data, sizeof(data)); // Read 4 bytes of file to compare against magic

    // Berkeley DB Btree magic bytes, from:
    //  https://github.com/file/file/blob/5824af38469ec1ca9ac3ffd251e7afe9dc11e227/magic/Magdir/database#L74-L75
    //  - big endian systems - 00 05 31 62
    //  - little endian systems - 62 31 05 00
    return data == 0x00053162 || data == 0x62310500;
}

std::vector<fs::path> ListWalletDir()
{
    const fs::path wallet_dir = GetWalletDir();
    const size_t offset = wallet_dir.string().size() + 1;
    std::vector<fs::path> paths;
    boost::system::error_code ec;

    for (auto it = fs::recursive_directory_iterator(wallet_dir, ec); it != fs::recursive_directory_iterator(); it.increment(ec)) {
        if (ec) {
            LogPrintf("%s: %s %s\n", __func__, ec.message(), it->path().string());
            continue;
        }

        // Get wallet path relative to walletdir by removing walletdir from the wallet path.
        // This can be replaced by boost::filesystem::lexically_relative once boost is bumped to 1.60.
        const fs::path path = it->path().string().substr(offset);

        if (it->status().type() == fs::directory_file && IsBerkeleyBtree(it->path() / "wallet.dat")) {
            // Found a directory which contains wallet.dat btree file, add it as a wallet.
            paths.emplace_back(path);
        } else if (it.level() == 0 && it->symlink_status().type() == fs::regular_file && IsBerkeleyBtree(it->path())) {
            if (it->path().filename() == "wallet.dat") {
                // Found top-level wallet.dat btree file, add top level directory ""
                // as a wallet.
                paths.emplace_back();
            } else {
                // Found top-level btree file not called wallet.dat. Current bitcoin
                // software will never create these files but will allow them to be
                // opened in a shared database environment for backwards compatibility.
                // Add it to the list of available wallets.
                paths.emplace_back(path);
            }
        }
    }

    return paths;
}

WalletLocation::WalletLocation(const std::string& name)
    : m_name(name)
    , m_path(fs::absolute(name, GetWalletDir()))
{
}

bool WalletLocation::Exists() const
{
    return fs::symlink_status(m_path).type() != fs::file_not_found;
}
<commit_msg>wallet: Do not iterate a directory if having an error while accessing it<commit_after>// Copyright (c) 2017-2019 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 <wallet/walletutil.h>

#include <logging.h>
#include <util/system.h>

fs::path GetWalletDir()
{
    fs::path path;

    if (gArgs.IsArgSet("-walletdir")) {
        path = gArgs.GetArg("-walletdir", "");
        if (!fs::is_directory(path)) {
            // If the path specified doesn't exist, we return the deliberately
            // invalid empty string.
            path = "";
        }
    } else {
        path = GetDataDir();
        // If a wallets directory exists, use that, otherwise default to GetDataDir
        if (fs::is_directory(path / "wallets")) {
            path /= "wallets";
        }
    }

    return path;
}

static bool IsBerkeleyBtree(const fs::path& path)
{
    if (!fs::exists(path)) return false;

    // A Berkeley DB Btree file has at least 4K.
    // This check also prevents opening lock files.
    boost::system::error_code ec;
    auto size = fs::file_size(path, ec);
    if (ec) LogPrintf("%s: %s %s\n", __func__, ec.message(), path.string());
    if (size < 4096) return false;

    fsbridge::ifstream file(path, std::ios::binary);
    if (!file.is_open()) return false;

    file.seekg(12, std::ios::beg); // Magic bytes start at offset 12
    uint32_t data = 0;
    file.read((char*) &data, sizeof(data)); // Read 4 bytes of file to compare against magic

    // Berkeley DB Btree magic bytes, from:
    //  https://github.com/file/file/blob/5824af38469ec1ca9ac3ffd251e7afe9dc11e227/magic/Magdir/database#L74-L75
    //  - big endian systems - 00 05 31 62
    //  - little endian systems - 62 31 05 00
    return data == 0x00053162 || data == 0x62310500;
}

std::vector<fs::path> ListWalletDir()
{
    const fs::path wallet_dir = GetWalletDir();
    const size_t offset = wallet_dir.string().size() + 1;
    std::vector<fs::path> paths;
    boost::system::error_code ec;

    for (auto it = fs::recursive_directory_iterator(wallet_dir, ec); it != fs::recursive_directory_iterator(); it.increment(ec)) {
        if (ec) {
            if (fs::is_directory(*it)) {
                it.no_push();
                LogPrintf("%s: %s %s -- skipping.\n", __func__, ec.message(), it->path().string());
            } else {
                LogPrintf("%s: %s %s\n", __func__, ec.message(), it->path().string());
            }
            continue;
        }

        // Get wallet path relative to walletdir by removing walletdir from the wallet path.
        // This can be replaced by boost::filesystem::lexically_relative once boost is bumped to 1.60.
        const fs::path path = it->path().string().substr(offset);

        if (it->status().type() == fs::directory_file && IsBerkeleyBtree(it->path() / "wallet.dat")) {
            // Found a directory which contains wallet.dat btree file, add it as a wallet.
            paths.emplace_back(path);
        } else if (it.level() == 0 && it->symlink_status().type() == fs::regular_file && IsBerkeleyBtree(it->path())) {
            if (it->path().filename() == "wallet.dat") {
                // Found top-level wallet.dat btree file, add top level directory ""
                // as a wallet.
                paths.emplace_back();
            } else {
                // Found top-level btree file not called wallet.dat. Current bitcoin
                // software will never create these files but will allow them to be
                // opened in a shared database environment for backwards compatibility.
                // Add it to the list of available wallets.
                paths.emplace_back(path);
            }
        }
    }

    return paths;
}

WalletLocation::WalletLocation(const std::string& name)
    : m_name(name)
    , m_path(fs::absolute(name, GetWalletDir()))
{
}

bool WalletLocation::Exists() const
{
    return fs::symlink_status(m_path).type() != fs::file_not_found;
}
<|endoftext|>


<commit_before>/**
 * @file /home/ryan/programming/atl/test/test.cpp
 * @author Ryan Domigan <ryan_domigan@sutdents@uml.edu>
 * Created on Feb 21, 2015
 */

#include <gtest/gtest.h>

#include <ffi.hpp>
#include <vm.hpp>

#include "./testing_utils.hpp"

#include "./utility.cpp"
#include "./test_type.cpp"
#include "./vm.cpp"
#include "./parser.cpp"
#include "./compile.cpp"
#include "./lists.cpp"
#include "./macros.cpp"
#include "./gc.cpp"
#include "./type_deduction.cpp"

#include "./atl.cpp"

int main(int argc, char *argv[]) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}
<commit_msg>Disable 'atl' tests<commit_after>/**
 * @file /home/ryan/programming/atl/test/test.cpp
 * @author Ryan Domigan <ryan_domigan@sutdents@uml.edu>
 * Created on Feb 21, 2015
 */

#include <gtest/gtest.h>

#include <ffi.hpp>
#include <vm.hpp>

#include "./testing_utils.hpp"

#include "./utility.cpp"
#include "./test_type.cpp"
#include "./vm.cpp"
#include "./parser.cpp"
#include "./compile.cpp"
#include "./lists.cpp"
#include "./macros.cpp"
#include "./gc.cpp"
#include "./type_deduction.cpp"

// #include "./atl.cpp"

int main(int argc, char *argv[]) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}
<|endoftext|>


<commit_before>/*
    Copyright (c) 2016-2019 Xavier Leclercq

    Permission is hereby granted, free of charge, to any person obtaining a
    copy of this software and associated documentation files (the "Software"),
    to deal in the Software without restriction, including without limitation
    the rights to use, copy, modify, merge, publish, distribute, sublicense,
    and/or sell copies of the Software, and to permit persons to whom the
    Software is furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in
    all copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
    THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
    IN THE SOFTWARE.
*/

/*
    Part of this file were copied from the Chart.js project (http://chartjs.org/)
    and translated into C++.

    The files of the Chart.js project have the following copyright and license.

    Copyright (c) 2013-2016 Nick Downie
    Released under the MIT license
    https://github.com/nnnick/Chart.js/blob/master/LICENSE.md
*/

#include "wxchartsrectangle.h"
#include <wx/pen.h>
#include <wx/brush.h>

wxChartsRectangle::wxChartsRectangle(wxDouble x,
                                     wxDouble y,
                                     const wxChartTooltipProvider::ptr tooltipProvider,
                                     const wxChartsRectangleOptions &options)
    : wxChartsElement(tooltipProvider),
    m_position(x, y), m_width(0), m_height(0), m_options(options)
{
}

void wxChartsRectangle::Draw(wxGraphicsContext &gc) const
{
    wxGraphicsPath path = gc.CreatePath();

    path.AddRectangle(m_position.m_x, m_position.m_y, m_width, m_height);

    wxBrush brush(m_options.GetBrushOptions().GetColor());
    gc.SetBrush(brush);
    gc.FillPath(path);

    if (m_options.GetBorders() == wxALL)
    {
        wxPen pen(m_options.GetPenOptions().GetColor(), 2);
        gc.SetPen(pen);
        gc.StrokePath(path);
    }
    else
    {
        wxPen pen(m_options.GetPenOptions().GetColor(), 2);
        gc.SetPen(pen);

        int borders = m_options.GetBorders();
        if (borders & wxTOP)
        {
            gc.StrokeLine(m_position.m_x, m_position.m_y,
                m_position.m_x + m_width, m_position.m_y);
        }
        if (borders & wxRIGHT)
        {
            gc.StrokeLine(m_position.m_x + m_width, m_position.m_y,
                m_position.m_x + m_width, m_position.m_y + m_height);
        }
        if (borders & wxBOTTOM)
        {
            gc.StrokeLine(m_position.m_x, m_position.m_y + m_height,
                m_position.m_x + m_width, m_position.m_y + m_height);
        }
        if (borders & wxLEFT)
        {
            gc.StrokeLine(m_position.m_x, m_position.m_y,
                m_position.m_x, m_position.m_y + m_height);
        }
    }
}

bool wxChartsRectangle::HitTest(const wxPoint &point) const
{
    bool x = ((m_position.m_x <= point.x) && (point.x <= (m_position.m_x + m_width)));
    bool y = ((m_position.m_y <= point.y) && (point.y <= (m_position.m_y + m_height)));
    return (x && y);
}

wxPoint2DDouble wxChartsRectangle::GetTooltipPosition() const
{
    return wxPoint2DDouble(m_position.m_x + (m_width / 2),
        m_position.m_y + (m_height / 2));
}

const wxPoint2DDouble& wxChartsRectangle::GetPosition() const
{
    return m_position;
}

void wxChartsRectangle::SetPosition(wxDouble x, wxDouble y)
{
    m_position.m_x = x;
    m_position.m_y = y;
}

void wxChartsRectangle::SetPosition(wxPoint2DDouble position)
{
    m_position = position;
}

wxDouble wxChartsRectangle::GetWidth() const
{
    return m_width;
}

wxDouble wxChartsRectangle::GetHeight() const
{
    return m_height;
}

void wxChartsRectangle::SetSize(wxDouble width, wxDouble height)
{
    m_width = width;
    m_height = height;
}
<commit_msg>Use pen width as defined by options in wxChartsRectangle<commit_after>/*
    Copyright (c) 2016-2019 Xavier Leclercq

    Permission is hereby granted, free of charge, to any person obtaining a
    copy of this software and associated documentation files (the "Software"),
    to deal in the Software without restriction, including without limitation
    the rights to use, copy, modify, merge, publish, distribute, sublicense,
    and/or sell copies of the Software, and to permit persons to whom the
    Software is furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in
    all copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
    THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
    IN THE SOFTWARE.
*/

/*
    Part of this file were copied from the Chart.js project (http://chartjs.org/)
    and translated into C++.

    The files of the Chart.js project have the following copyright and license.

    Copyright (c) 2013-2016 Nick Downie
    Released under the MIT license
    https://github.com/nnnick/Chart.js/blob/master/LICENSE.md
*/

#include "wxchartsrectangle.h"
#include "wxchartsutilities.h"
#include <wx/pen.h>
#include <wx/brush.h>

wxChartsRectangle::wxChartsRectangle(wxDouble x,
                                     wxDouble y,
                                     const wxChartTooltipProvider::ptr tooltipProvider,
                                     const wxChartsRectangleOptions &options)
    : wxChartsElement(tooltipProvider),
    m_position(x, y), m_width(0), m_height(0), m_options(options)
{
}

void wxChartsRectangle::Draw(wxGraphicsContext &gc) const
{
    wxGraphicsPath path = gc.CreatePath();

    path.AddRectangle(m_position.m_x, m_position.m_y, m_width, m_height);

    wxBrush brush(m_options.GetBrushOptions().GetColor());
    gc.SetBrush(brush);
    gc.FillPath(path);

    if (m_options.GetBorders() == wxALL)
    {
        wxPen pen = wxChartsUtilities::CreatePen(m_options.GetPenOptions());
        gc.SetPen(pen);
        gc.StrokePath(path);
    }
    else
    {
        wxPen pen = wxChartsUtilities::CreatePen(m_options.GetPenOptions());
        gc.SetPen(pen);

        int borders = m_options.GetBorders();
        if (borders & wxTOP)
        {
            gc.StrokeLine(m_position.m_x, m_position.m_y,
                m_position.m_x + m_width, m_position.m_y);
        }
        if (borders & wxRIGHT)
        {
            gc.StrokeLine(m_position.m_x + m_width, m_position.m_y,
                m_position.m_x + m_width, m_position.m_y + m_height);
        }
        if (borders & wxBOTTOM)
        {
            gc.StrokeLine(m_position.m_x, m_position.m_y + m_height,
                m_position.m_x + m_width, m_position.m_y + m_height);
        }
        if (borders & wxLEFT)
        {
            gc.StrokeLine(m_position.m_x, m_position.m_y,
                m_position.m_x, m_position.m_y + m_height);
        }
    }
}

bool wxChartsRectangle::HitTest(const wxPoint &point) const
{
    bool x = ((m_position.m_x <= point.x) && (point.x <= (m_position.m_x + m_width)));
    bool y = ((m_position.m_y <= point.y) && (point.y <= (m_position.m_y + m_height)));
    return (x && y);
}

wxPoint2DDouble wxChartsRectangle::GetTooltipPosition() const
{
    return wxPoint2DDouble(m_position.m_x + (m_width / 2),
        m_position.m_y + (m_height / 2));
}

const wxPoint2DDouble& wxChartsRectangle::GetPosition() const
{
    return m_position;
}

void wxChartsRectangle::SetPosition(wxDouble x, wxDouble y)
{
    m_position.m_x = x;
    m_position.m_y = y;
}

void wxChartsRectangle::SetPosition(wxPoint2DDouble position)
{
    m_position = position;
}

wxDouble wxChartsRectangle::GetWidth() const
{
    return m_width;
}

wxDouble wxChartsRectangle::GetHeight() const
{
    return m_height;
}

void wxChartsRectangle::SetSize(wxDouble width, wxDouble height)
{
    m_width = width;
    m_height = height;
}
<|endoftext|>


<commit_before>#line 2 "togo/core/threading/task_manager.cpp"
/**
@copyright MIT license; see @ref index or the accompanying LICENSE file.
*/

#include <togo/core/config.hpp>
#include <togo/core/error/assert.hpp>
#include <togo/core/utility/utility.hpp>
#include <togo/core/log/log.hpp>
#include <togo/core/collection/array.hpp>
#include <togo/core/collection/priority_queue.hpp>
#include <togo/core/threading/condvar.hpp>
#include <togo/core/threading/mutex.hpp>
#include <togo/core/threading/thread.hpp>
#include <togo/core/threading/task_manager.hpp>

#include <cstring>
#include <cstdio>

#undef TOGO_TEST_LOG_ENABLE
#if defined(TOGO_TEST_TASK_MANAGER)
	#define TOGO_TEST_LOG_ENABLE
#endif
#include <togo/core/log/test.hpp>

namespace togo {

namespace {

static constexpr TaskID const ID_NULL{0};

inline bool operator==(TaskID const& x, TaskID const& y) {
	return x._value == y._value;
}

inline bool operator!=(TaskID const& x, TaskID const& y) {
	return x._value != y._value;
}

inline unsigned task_sort_key(Task const& task) {
	// Sort by num_incomplete, then by priority
	return (~unsigned{task.num_incomplete} << 16 & 0xFFFF0000) | unsigned{task.priority};
}

bool task_less(Task* const& x, Task* const& y) {
	return task_sort_key(*x) < task_sort_key(*y);
}

enum : unsigned {
	// 128 tasks
	ID_SHIFT = 7,
	NUM_TASKS = 1 << ID_SHIFT,
	ID_ADD = 1 << ID_SHIFT,
	INDEX_MASK = NUM_TASKS - 1,

	FLAG_SHUTDOWN = 1 << 0,
};

inline Task& get_task(TaskManager& tm, TaskID const id) {
	return tm._tasks[id._value & INDEX_MASK].task;
}

inline bool is_complete(Task const& task, TaskID const expected_id) {
	return task.id != expected_id || task.num_incomplete == 0;
}

inline bool is_ready(Task const& task) {
	return task.num_incomplete <= 1;
}

inline void free_slot(TaskManager& tm, Task& task) {
	TaskSlot& slot = *reinterpret_cast<TaskSlot*>(&task);
	unsigned index = slot.id._value & INDEX_MASK;
	slot.hole.next = nullptr;
	while (index--) {
		TaskSlot& slot_before = tm._tasks[index];
		if (slot_before.id == ID_NULL) {
			slot.hole.next = slot_before.hole.next;
			slot_before.hole.next = &slot;
			break;
		}
	}
	if (!slot.hole.next) {
		// No hole between slot and head
		slot.hole.next = tm._first_hole;
		tm._first_hole = &slot;
	}
	slot.id = ID_NULL;
}

inline Task& add_task(
	TaskManager& tm,
	TaskWork const& work,
	u16 const priority
) {
	TOGO_ASSERT(tm._num_tasks != NUM_TASKS, "cannot add task to full task manager");
	TOGO_DEBUG_ASSERTE(tm._first_hole != nullptr);
	TaskSlot& slot = *tm._first_hole;
	tm._first_hole = slot.hole.next;
	std::memset(&slot.task, 0, sizeof(Task));
	slot.task.id._value = tm._id_gen | (&slot - tm._tasks);
	slot.task.work = work;
	slot.task.priority = priority;
	slot.task.num_incomplete = 1;
	tm._id_gen = max(tm._id_gen + ID_ADD, u32{ID_ADD});
	return slot.task;
}

inline void queue_task(TaskManager& tm, Task& task) {
	/*TOGO_TEST_LOG_DEBUGF(
		"queue_task: push: %04u %03u @ %04hu / %04hu\n",
		task.id._value >> ID_SHIFT,
		task.id._value & INDEX_MASK,
		task.num_incomplete,
		task.priority
	);*/
	priority_queue::push(tm._queue, &task);
	/*#if defined(TOGO_TEST_TASK_MANAGER)
		Task* const front = priority_queue::front(tm._queue);
		TOGO_TEST_LOG_DEBUGF(
			"queue_task: front: %04u %03u @ %04hu / %04hu\n",
			front->id._value >> ID_SHIFT,
			front->id._value & INDEX_MASK,
			front->num_incomplete,
			front->priority
		);
	#endif*/
	condvar::signal(tm._work_signal, tm._mutex);
}

inline void complete_task(TaskManager& tm, Task& task) {
	// TODO: Reorder each parent in the priority queue
	TOGO_TEST_LOG_DEBUGF(
		"complete_task  : %-32s: %04u %03u @ %04hu / %04hu\n",
		thread::name(),
		task.id._value >> ID_SHIFT,
		task.id._value & INDEX_MASK,
		task.num_incomplete,
		task.priority
	);
	task.num_incomplete = 0;
	TaskID parent_id = task.parent;
	while (parent_id != ID_NULL) {
		Task& parent = get_task(tm, parent_id);
		--parent.num_incomplete;
		parent_id = parent.num_incomplete ? parent.parent : ID_NULL;
	}
	free_slot(tm, task);
	condvar::signal_all(tm._work_signal, tm._mutex);
}

void execute_pending(TaskManager& tm, TaskID const wait_id) {
	Task* task = nullptr;
	Task const* wait_task = nullptr;
	if (wait_id != ID_NULL) {
		wait_task = &get_task(tm, wait_id);
	}
	MutexLock lock{tm._mutex};
	while (true) {
		if (task) {
			complete_task(tm, *task);
			task = nullptr;
		}
		if (wait_task && is_complete(*wait_task, wait_id)) {
			return;
		} else if (tm._flags & FLAG_SHUTDOWN) {
			return;
		} else if (
			priority_queue::any(tm._queue) &&
			is_ready(*priority_queue::front(tm._queue))
		) {
			task = priority_queue::front(tm._queue);
			priority_queue::pop(tm._queue);
			TOGO_TEST_LOG_DEBUGF(
				"execute_pending: %-32s: %04u %03u @ %04hu / %04hu [take]\n",
				thread::name(),
				task->id._value >> ID_SHIFT,
				task->id._value & INDEX_MASK,
				task->num_incomplete,
				task->priority
			);
			// If !task->work.func, the task is empty
			if (task->work.func) {
				// task->id and task->work should not be modified
				// after adding and the task should not be destroyed
				// by any other function, so this is free of race
				// conditions.
				mutex::unlock(tm._mutex);
				task->work.func(task->id, task->work.data);
				mutex::lock(tm._mutex);
			}
			continue;
		}
		condvar::wait(tm._work_signal, lock);
	}
}

void* worker_func(void* const tm_void) {
	TaskManager& tm = *static_cast<TaskManager*>(tm_void);
	TOGO_TEST_LOG_DEBUGF("worker_func: start: %s\n", thread::name());
	execute_pending(tm, ID_NULL);
	TOGO_TEST_LOG_DEBUGF("worker_func: shutdown: %s\n", thread::name());
	return nullptr;
}

} // anonymous namespace

static_assert(
	NUM_TASKS == array_extent(&TaskManager::_tasks),
	"ID_SHIFT does not fit with the size of TaskManager::_tasks"
);

// class TaskManager implementation

TaskManager::~TaskManager() {
	{
		MutexLock lock{_mutex};
		_flags |= FLAG_SHUTDOWN;
		condvar::signal_all(_work_signal, lock);
	}
	for (Thread* const worker_thread : _workers) {
		thread::join(worker_thread);
	}
}

TaskManager::TaskManager(unsigned num_workers, Allocator& allocator)
	: _tasks()
	, _queue(task_less, allocator)
	, _workers(allocator)
	, _mutex(MutexType::normal)
	, _work_signal()
	, _first_hole(_tasks)
	, _num_tasks(0)
	, _flags(0)
	, _id_gen(ID_ADD)
{
	TaskSlot* const last = _tasks + (NUM_TASKS - 1);
	for (TaskSlot* slot = _tasks; slot != last; ++slot) {
		slot->hole.next = slot + 1;
	}
	last->hole.next = nullptr;
	priority_queue::reserve(_queue, NUM_TASKS);
	if (num_workers) {
		array::reserve(_workers, num_workers);
		char name[32];
		while (num_workers--) {
			std::snprintf(name, sizeof(name), "tm-%8p-worker-%u", this, num_workers);
			array::push_back(_workers, thread::create(name, this, worker_func, allocator));
		}
	}
}

// interface task_manager implementation

/// Add task with no children.
///
/// The task may start immediately.
TaskID task_manager::add(
	TaskManager& tm,
	TaskWork const& work,
	u16 const priority IGEN_DEFAULT(0)
) {
	MutexLock lock{tm._mutex};
	Task& task = add_task(tm, work, priority);
	queue_task(tm, task);
	return task.id;
}

/// Add task with a hold on its execution.
TaskID task_manager::add_hold(
	TaskManager& tm,
	TaskWork const& work,
	u16 const priority IGEN_DEFAULT(0)
) {
	MutexLock lock{tm._mutex};
	return add_task(tm, work, priority).id;
}

/// Set task parent.
///
/// To prevent race conditions, this should only be used with a parent
/// that was created with add_hold().
void task_manager::set_parent(
	TaskManager& tm,
	TaskID const child_id,
	TaskID const parent_id
) {
	TOGO_ASSERT(child_id != parent_id, "cannot make task a child of itself");
	MutexLock lock{tm._mutex};
	Task& child = get_task(tm, child_id);
	if (child.id != child_id) {
		TOGO_TEST_LOG_DEBUGF(
			"set_parent: child does not exist or has completed: %04u %03u\n",
			child_id._value >> ID_SHIFT,
			child_id._value & INDEX_MASK
		);
		return;
	}
	TOGO_ASSERT(child.parent == ID_NULL, "child task already has a parent");
	TOGO_ASSERT(parent_id != ID_NULL, "parent ID is invalid");
	Task& parent = get_task(tm, parent_id);
	TOGO_ASSERT(parent.id == parent_id, "parent task does not exist");
	++parent.num_incomplete;
	child.parent = parent_id;
}

/// End the hold on a task.
void task_manager::end_hold(TaskManager& tm, TaskID const id) {
	MutexLock lock{tm._mutex};
	Task& task = get_task(tm, id);
	TOGO_ASSERT(task.id == id, "id is not valid");
	queue_task(tm, task);
}

/// Wait for a task to complete.
///
/// This function will execute any available tasks while id is incomplete.
void task_manager::wait(TaskManager& tm, TaskID const id) {
	TOGO_DEBUG_ASSERT(id != ID_NULL, "attempted to wait on null ID");
	execute_pending(tm, id);
}

} // namespace togo
<commit_msg>lib/core/threading/task_manager: explicitly cast `this` pointer to shut up -Wformat-pedantic >_><commit_after>#line 2 "togo/core/threading/task_manager.cpp"
/**
@copyright MIT license; see @ref index or the accompanying LICENSE file.
*/

#include <togo/core/config.hpp>
#include <togo/core/error/assert.hpp>
#include <togo/core/utility/utility.hpp>
#include <togo/core/log/log.hpp>
#include <togo/core/collection/array.hpp>
#include <togo/core/collection/priority_queue.hpp>
#include <togo/core/threading/condvar.hpp>
#include <togo/core/threading/mutex.hpp>
#include <togo/core/threading/thread.hpp>
#include <togo/core/threading/task_manager.hpp>

#include <cstring>
#include <cstdio>

#undef TOGO_TEST_LOG_ENABLE
#if defined(TOGO_TEST_TASK_MANAGER)
	#define TOGO_TEST_LOG_ENABLE
#endif
#include <togo/core/log/test.hpp>

namespace togo {

namespace {

static constexpr TaskID const ID_NULL{0};

inline bool operator==(TaskID const& x, TaskID const& y) {
	return x._value == y._value;
}

inline bool operator!=(TaskID const& x, TaskID const& y) {
	return x._value != y._value;
}

inline unsigned task_sort_key(Task const& task) {
	// Sort by num_incomplete, then by priority
	return (~unsigned{task.num_incomplete} << 16 & 0xFFFF0000) | unsigned{task.priority};
}

bool task_less(Task* const& x, Task* const& y) {
	return task_sort_key(*x) < task_sort_key(*y);
}

enum : unsigned {
	// 128 tasks
	ID_SHIFT = 7,
	NUM_TASKS = 1 << ID_SHIFT,
	ID_ADD = 1 << ID_SHIFT,
	INDEX_MASK = NUM_TASKS - 1,

	FLAG_SHUTDOWN = 1 << 0,
};

inline Task& get_task(TaskManager& tm, TaskID const id) {
	return tm._tasks[id._value & INDEX_MASK].task;
}

inline bool is_complete(Task const& task, TaskID const expected_id) {
	return task.id != expected_id || task.num_incomplete == 0;
}

inline bool is_ready(Task const& task) {
	return task.num_incomplete <= 1;
}

inline void free_slot(TaskManager& tm, Task& task) {
	TaskSlot& slot = *reinterpret_cast<TaskSlot*>(&task);
	unsigned index = slot.id._value & INDEX_MASK;
	slot.hole.next = nullptr;
	while (index--) {
		TaskSlot& slot_before = tm._tasks[index];
		if (slot_before.id == ID_NULL) {
			slot.hole.next = slot_before.hole.next;
			slot_before.hole.next = &slot;
			break;
		}
	}
	if (!slot.hole.next) {
		// No hole between slot and head
		slot.hole.next = tm._first_hole;
		tm._first_hole = &slot;
	}
	slot.id = ID_NULL;
}

inline Task& add_task(
	TaskManager& tm,
	TaskWork const& work,
	u16 const priority
) {
	TOGO_ASSERT(tm._num_tasks != NUM_TASKS, "cannot add task to full task manager");
	TOGO_DEBUG_ASSERTE(tm._first_hole != nullptr);
	TaskSlot& slot = *tm._first_hole;
	tm._first_hole = slot.hole.next;
	std::memset(&slot.task, 0, sizeof(Task));
	slot.task.id._value = tm._id_gen | (&slot - tm._tasks);
	slot.task.work = work;
	slot.task.priority = priority;
	slot.task.num_incomplete = 1;
	tm._id_gen = max(tm._id_gen + ID_ADD, u32{ID_ADD});
	return slot.task;
}

inline void queue_task(TaskManager& tm, Task& task) {
	/*TOGO_TEST_LOG_DEBUGF(
		"queue_task: push: %04u %03u @ %04hu / %04hu\n",
		task.id._value >> ID_SHIFT,
		task.id._value & INDEX_MASK,
		task.num_incomplete,
		task.priority
	);*/
	priority_queue::push(tm._queue, &task);
	/*#if defined(TOGO_TEST_TASK_MANAGER)
		Task* const front = priority_queue::front(tm._queue);
		TOGO_TEST_LOG_DEBUGF(
			"queue_task: front: %04u %03u @ %04hu / %04hu\n",
			front->id._value >> ID_SHIFT,
			front->id._value & INDEX_MASK,
			front->num_incomplete,
			front->priority
		);
	#endif*/
	condvar::signal(tm._work_signal, tm._mutex);
}

inline void complete_task(TaskManager& tm, Task& task) {
	// TODO: Reorder each parent in the priority queue
	TOGO_TEST_LOG_DEBUGF(
		"complete_task  : %-32s: %04u %03u @ %04hu / %04hu\n",
		thread::name(),
		task.id._value >> ID_SHIFT,
		task.id._value & INDEX_MASK,
		task.num_incomplete,
		task.priority
	);
	task.num_incomplete = 0;
	TaskID parent_id = task.parent;
	while (parent_id != ID_NULL) {
		Task& parent = get_task(tm, parent_id);
		--parent.num_incomplete;
		parent_id = parent.num_incomplete ? parent.parent : ID_NULL;
	}
	free_slot(tm, task);
	condvar::signal_all(tm._work_signal, tm._mutex);
}

void execute_pending(TaskManager& tm, TaskID const wait_id) {
	Task* task = nullptr;
	Task const* wait_task = nullptr;
	if (wait_id != ID_NULL) {
		wait_task = &get_task(tm, wait_id);
	}
	MutexLock lock{tm._mutex};
	while (true) {
		if (task) {
			complete_task(tm, *task);
			task = nullptr;
		}
		if (wait_task && is_complete(*wait_task, wait_id)) {
			return;
		} else if (tm._flags & FLAG_SHUTDOWN) {
			return;
		} else if (
			priority_queue::any(tm._queue) &&
			is_ready(*priority_queue::front(tm._queue))
		) {
			task = priority_queue::front(tm._queue);
			priority_queue::pop(tm._queue);
			TOGO_TEST_LOG_DEBUGF(
				"execute_pending: %-32s: %04u %03u @ %04hu / %04hu [take]\n",
				thread::name(),
				task->id._value >> ID_SHIFT,
				task->id._value & INDEX_MASK,
				task->num_incomplete,
				task->priority
			);
			// If !task->work.func, the task is empty
			if (task->work.func) {
				// task->id and task->work should not be modified
				// after adding and the task should not be destroyed
				// by any other function, so this is free of race
				// conditions.
				mutex::unlock(tm._mutex);
				task->work.func(task->id, task->work.data);
				mutex::lock(tm._mutex);
			}
			continue;
		}
		condvar::wait(tm._work_signal, lock);
	}
}

void* worker_func(void* const tm_void) {
	TaskManager& tm = *static_cast<TaskManager*>(tm_void);
	TOGO_TEST_LOG_DEBUGF("worker_func: start: %s\n", thread::name());
	execute_pending(tm, ID_NULL);
	TOGO_TEST_LOG_DEBUGF("worker_func: shutdown: %s\n", thread::name());
	return nullptr;
}

} // anonymous namespace

static_assert(
	NUM_TASKS == array_extent(&TaskManager::_tasks),
	"ID_SHIFT does not fit with the size of TaskManager::_tasks"
);

// class TaskManager implementation

TaskManager::~TaskManager() {
	{
		MutexLock lock{_mutex};
		_flags |= FLAG_SHUTDOWN;
		condvar::signal_all(_work_signal, lock);
	}
	for (Thread* const worker_thread : _workers) {
		thread::join(worker_thread);
	}
}

TaskManager::TaskManager(unsigned num_workers, Allocator& allocator)
	: _tasks()
	, _queue(task_less, allocator)
	, _workers(allocator)
	, _mutex(MutexType::normal)
	, _work_signal()
	, _first_hole(_tasks)
	, _num_tasks(0)
	, _flags(0)
	, _id_gen(ID_ADD)
{
	TaskSlot* const last = _tasks + (NUM_TASKS - 1);
	for (TaskSlot* slot = _tasks; slot != last; ++slot) {
		slot->hole.next = slot + 1;
	}
	last->hole.next = nullptr;
	priority_queue::reserve(_queue, NUM_TASKS);
	if (num_workers) {
		array::reserve(_workers, num_workers);
		char name[32];
		while (num_workers--) {
			std::snprintf(name, sizeof(name), "tm-%8p-worker-%u", static_cast<void*>(this), num_workers);
			array::push_back(_workers, thread::create(name, this, worker_func, allocator));
		}
	}
}

// interface task_manager implementation

/// Add task with no children.
///
/// The task may start immediately.
TaskID task_manager::add(
	TaskManager& tm,
	TaskWork const& work,
	u16 const priority IGEN_DEFAULT(0)
) {
	MutexLock lock{tm._mutex};
	Task& task = add_task(tm, work, priority);
	queue_task(tm, task);
	return task.id;
}

/// Add task with a hold on its execution.
TaskID task_manager::add_hold(
	TaskManager& tm,
	TaskWork const& work,
	u16 const priority IGEN_DEFAULT(0)
) {
	MutexLock lock{tm._mutex};
	return add_task(tm, work, priority).id;
}

/// Set task parent.
///
/// To prevent race conditions, this should only be used with a parent
/// that was created with add_hold().
void task_manager::set_parent(
	TaskManager& tm,
	TaskID const child_id,
	TaskID const parent_id
) {
	TOGO_ASSERT(child_id != parent_id, "cannot make task a child of itself");
	MutexLock lock{tm._mutex};
	Task& child = get_task(tm, child_id);
	if (child.id != child_id) {
		TOGO_TEST_LOG_DEBUGF(
			"set_parent: child does not exist or has completed: %04u %03u\n",
			child_id._value >> ID_SHIFT,
			child_id._value & INDEX_MASK
		);
		return;
	}
	TOGO_ASSERT(child.parent == ID_NULL, "child task already has a parent");
	TOGO_ASSERT(parent_id != ID_NULL, "parent ID is invalid");
	Task& parent = get_task(tm, parent_id);
	TOGO_ASSERT(parent.id == parent_id, "parent task does not exist");
	++parent.num_incomplete;
	child.parent = parent_id;
}

/// End the hold on a task.
void task_manager::end_hold(TaskManager& tm, TaskID const id) {
	MutexLock lock{tm._mutex};
	Task& task = get_task(tm, id);
	TOGO_ASSERT(task.id == id, "id is not valid");
	queue_task(tm, task);
}

/// Wait for a task to complete.
///
/// This function will execute any available tasks while id is incomplete.
void task_manager::wait(TaskManager& tm, TaskID const id) {
	TOGO_DEBUG_ASSERT(id != ID_NULL, "attempted to wait on null ID");
	execute_pending(tm, id);
}

} // namespace togo
<|endoftext|>


<commit_before>#define CATCH_CONFIG_RUNNER
#include "catch.hpp"

#include "variant.hpp"
#include <cstdint>
#include <limits>
#include <string>
#include <ostream>
#include <memory>

TEST_CASE( "variant version", "[variant]" ) {
    unsigned int version = VARIANT_VERSION;
    REQUIRE(version == 100);
    #if VARIANT_VERSION == 100
        REQUIRE(true);
    #else
        REQUIRE(false);
    #endif
}

TEST_CASE( "variant type traits", "[variant]" ) {
    REQUIRE((util::detail::type_traits<bool, bool, int, double, std::string>::id == 3));
    REQUIRE((util::detail::type_traits<int, bool, int, double, std::string>::id == 2));
    REQUIRE((util::detail::type_traits<double, bool, int, double, std::string>::id == 1));
    REQUIRE((util::detail::type_traits<std::string, bool, int, double, std::string>::id == 0));
    REQUIRE((util::detail::type_traits<long, bool, int, double, std::string>::id == util::detail::invalid_value));
    REQUIRE((util::detail::type_traits<std::vector<int>, bool, int, double, std::string>::id == util::detail::invalid_value));
}

TEST_CASE( "variant can be moved into vector", "[variant]" ) {
    typedef util::variant<bool,std::string> variant_type;
    variant_type v(std::string("test"));
    std::vector<variant_type> vec;
    vec.emplace_back(std::move(v));
    REQUIRE(v.get<std::string>() != std::string("test"));
    REQUIRE(vec.at(0).get<std::string>() == std::string("test"));
}

TEST_CASE( "variant should support built-in types", "[variant]" ) {
    SECTION( "bool" ) {
        util::variant<bool> v(true);
        REQUIRE(v.valid());
        REQUIRE(v.is<bool>());
        REQUIRE(v.get<bool>() == true);
        v.set<bool>(false);
        REQUIRE(v.get<bool>() == false);
        v = true;
        REQUIRE(v == util::variant<bool>(true));
    }
    SECTION( "nullptr" ) {
        typedef std::nullptr_t value_type;
        util::variant<value_type> v(nullptr);
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        // TODO: commented since it breaks on windows: 'operator << is ambiguous'
        //REQUIRE(v.get<value_type>() == nullptr);
        // FIXME: does not compile: ./variant.hpp:340:14: error: use of overloaded operator '<<' is ambiguous (with operand types 'std::__1::basic_ostream<char>' and 'const nullptr_t')
        // https://github.com/mapbox/variant/issues/14
        //REQUIRE(v == util::variant<value_type>(nullptr));
    }
    SECTION( "unique_ptr" ) {
        typedef std::unique_ptr<std::string> value_type;
        util::variant<value_type> v(value_type(new std::string("hello")));
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(*v.get<value_type>().get() == *value_type(new std::string("hello")).get());
    }
    SECTION( "string" ) {
        typedef std::string value_type;
        util::variant<value_type> v(value_type("hello"));
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == value_type("hello"));
        v.set<value_type>(value_type("there"));
        REQUIRE(v.get<value_type>() == value_type("there"));
        v = value_type("variant");
        REQUIRE(v == util::variant<value_type>(value_type("variant")));
    }
    SECTION( "size_t" ) {
        typedef std::size_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(value_type(0));
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int8_t" ) {
        typedef std::int8_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int16_t" ) {
        typedef std::int16_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int32_t" ) {
        typedef std::int32_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int64_t" ) {
        typedef std::int64_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
}

struct MissionInteger
{
    typedef uint64_t value_type;
    value_type val_;
    public:
      MissionInteger(uint64_t val) :
       val_(val) {}

    bool operator==(MissionInteger const& rhs) const
    {
        return (val_ == rhs.get());
    }

    uint64_t get() const
    {
        return val_;
    }
};

// TODO - remove after https://github.com/mapbox/variant/issues/14
std::ostream& operator<< (std::ostream& os, MissionInteger const& rhs)
{
    os << rhs.get();
    return os;
}

TEST_CASE( "variant should support custom types", "[variant]" ) {
    // http://www.missionintegers.com/integer/34838300
    util::variant<MissionInteger> v(MissionInteger(34838300));
    REQUIRE(v.valid());
    REQUIRE(v.is<MissionInteger>());
    REQUIRE(v.get<MissionInteger>() == MissionInteger(34838300));
    REQUIRE(v.get<MissionInteger>().get() == MissionInteger::value_type(34838300));
    // TODO: should both of the set usages below compile?
    v.set<MissionInteger>(MissionInteger::value_type(0));
    v.set<MissionInteger>(MissionInteger(0));
    REQUIRE(v.get<MissionInteger>().get() == MissionInteger::value_type(0));
    v = MissionInteger(1);
    REQUIRE(v == util::variant<MissionInteger>(MissionInteger(1)));
}

int main (int argc, char* const argv[])
{
    int result = Catch::Session().run( argc, argv );
    if (!result) printf("\x1b[1;32m ✓ \x1b[0m\n");
    return result;
}
<commit_msg>move detail tests + add initial comments (@artemp, please review)<commit_after>#define CATCH_CONFIG_RUNNER
#include "catch.hpp"

#include "variant.hpp"
#include <cstdint>
#include <limits>
#include <string>
#include <ostream>
#include <memory>

TEST_CASE( "variant version", "[variant]" ) {
    unsigned int version = VARIANT_VERSION;
    REQUIRE(version == 100);
    #if VARIANT_VERSION == 100
        REQUIRE(true);
    #else
        REQUIRE(false);
    #endif
}

TEST_CASE( "variant can be moved into vector", "[variant]" ) {
    typedef util::variant<bool,std::string> variant_type;
    variant_type v(std::string("test"));
    std::vector<variant_type> vec;
    vec.emplace_back(std::move(v));
    REQUIRE(v.get<std::string>() != std::string("test"));
    REQUIRE(vec.at(0).get<std::string>() == std::string("test"));
}

TEST_CASE( "variant should support built-in types", "[variant]" ) {
    SECTION( "bool" ) {
        util::variant<bool> v(true);
        REQUIRE(v.valid());
        REQUIRE(v.is<bool>());
        REQUIRE(v.get<bool>() == true);
        v.set<bool>(false);
        REQUIRE(v.get<bool>() == false);
        v = true;
        REQUIRE(v == util::variant<bool>(true));
    }
    SECTION( "nullptr" ) {
        typedef std::nullptr_t value_type;
        util::variant<value_type> v(nullptr);
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        // TODO: commented since it breaks on windows: 'operator << is ambiguous'
        //REQUIRE(v.get<value_type>() == nullptr);
        // FIXME: does not compile: ./variant.hpp:340:14: error: use of overloaded operator '<<' is ambiguous (with operand types 'std::__1::basic_ostream<char>' and 'const nullptr_t')
        // https://github.com/mapbox/variant/issues/14
        //REQUIRE(v == util::variant<value_type>(nullptr));
    }
    SECTION( "unique_ptr" ) {
        typedef std::unique_ptr<std::string> value_type;
        util::variant<value_type> v(value_type(new std::string("hello")));
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(*v.get<value_type>().get() == *value_type(new std::string("hello")).get());
    }
    SECTION( "string" ) {
        typedef std::string value_type;
        util::variant<value_type> v(value_type("hello"));
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == value_type("hello"));
        v.set<value_type>(value_type("there"));
        REQUIRE(v.get<value_type>() == value_type("there"));
        v = value_type("variant");
        REQUIRE(v == util::variant<value_type>(value_type("variant")));
    }
    SECTION( "size_t" ) {
        typedef std::size_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(value_type(0));
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int8_t" ) {
        typedef std::int8_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int16_t" ) {
        typedef std::int16_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int32_t" ) {
        typedef std::int32_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
    SECTION( "int64_t" ) {
        typedef std::int64_t value_type;
        util::variant<value_type> v(std::numeric_limits<value_type>::max());
        REQUIRE(v.valid());
        REQUIRE(v.is<value_type>());
        REQUIRE(v.get<value_type>() == std::numeric_limits<value_type>::max());
        v.set<value_type>(0);
        REQUIRE(v.get<value_type>() == value_type(0));
        v = value_type(1);
        REQUIRE(v == util::variant<value_type>(value_type(1)));
    }
}

struct MissionInteger
{
    typedef uint64_t value_type;
    value_type val_;
    public:
      MissionInteger(uint64_t val) :
       val_(val) {}

    bool operator==(MissionInteger const& rhs) const
    {
        return (val_ == rhs.get());
    }

    uint64_t get() const
    {
        return val_;
    }
};

// TODO - remove after https://github.com/mapbox/variant/issues/14
std::ostream& operator<< (std::ostream& os, MissionInteger const& rhs)
{
    os << rhs.get();
    return os;
}

TEST_CASE( "variant should support custom types", "[variant]" ) {
    // http://www.missionintegers.com/integer/34838300
    util::variant<MissionInteger> v(MissionInteger(34838300));
    REQUIRE(v.valid());
    REQUIRE(v.is<MissionInteger>());
    REQUIRE(v.get<MissionInteger>() == MissionInteger(34838300));
    REQUIRE(v.get<MissionInteger>().get() == MissionInteger::value_type(34838300));
    // TODO: should both of the set usages below compile?
    v.set<MissionInteger>(MissionInteger::value_type(0));
    v.set<MissionInteger>(MissionInteger(0));
    REQUIRE(v.get<MissionInteger>().get() == MissionInteger::value_type(0));
    v = MissionInteger(1);
    REQUIRE(v == util::variant<MissionInteger>(MissionInteger(1)));
}

// Test detail code (internal api not mean for public use)
TEST_CASE( "variant type traits", "[variant::detail]" ) {
    // users should not create variants with duplicated types
    // however our type indexing should still work
    // Note index for types in reverse order
    REQUIRE((util::detail::type_traits<bool, bool, int, double, std::string>::id == 3));
    REQUIRE((util::detail::type_traits<int, bool, int, double, std::string>::id == 2));
    REQUIRE((util::detail::type_traits<double, bool, int, double, std::string>::id == 1));
    REQUIRE((util::detail::type_traits<std::string, bool, int, double, std::string>::id == 0));
    REQUIRE((util::detail::type_traits<long, bool, int, double, std::string>::id == util::detail::invalid_value));
    REQUIRE((util::detail::type_traits<std::vector<int>, bool, int, double, std::string>::id == util::detail::invalid_value));
}

int main (int argc, char* const argv[])
{
    int result = Catch::Session().run( argc, argv );
    if (!result) printf("\x1b[1;32m ✓ \x1b[0m\n");
    return result;
}
<|endoftext|>


<commit_before>#define WIN32_LEAN_AND_MEAN

#include <SDKDDKVer.h>
#include <windows.h>
#include "HookAPI.h"
#include <ws2tcpip.h>
#include <atomic>
#include <concurrent_queue.h>

#pragma comment (lib, "Ws2_32.lib")

#define BUFLEN 4096

// Socket used for TeamSpeak communication
std::atomic<SOCKET> client = NULL;
// Lua state used when requiring script files
std::atomic<lua_State *> state = NULL;
// Boolean telling the socket thread whether to keep the network loop running
std::atomic<bool> running = FALSE;
// Queue of commands to be sent to the Lua script
Concurrency::concurrent_queue<char *> queue;

// Tries to connect to TeamSpeak using ClientQuery
// Returns a new socket or NULL on failure

SOCKET Connect(PCSTR hostname, PCSTR port)
{
	// Creates a structure used for containing the socket information
	WSADATA wsaData;
	int result = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (result != 0) return NULL;

	// Creates a structure used for containing the connection information
	struct addrinfo *info, *ptr, hints;
	ZeroMemory(&hints, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;
	result = getaddrinfo(hostname, port, &hints, &info);
	if (result != 0)
	{
		WSACleanup();
		return NULL;
	}

	// Tries to connect to the specified address
	SOCKET client = INVALID_SOCKET;
	for (ptr = info; ptr != NULL; ptr = ptr->ai_next)
	{
		client = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
		if (client == INVALID_SOCKET)
		{
			WSACleanup();
			return NULL;
		}
		result = connect(client, ptr->ai_addr, (int)ptr->ai_addrlen);
		if (result == SOCKET_ERROR)
		{
			closesocket(client);
			client = INVALID_SOCKET;
			continue;
		}
		break;
	}

	// Releases the address information strcuture from memory
	freeaddrinfo(info);

	// Return NULL if the connection could not be established
	if (client == INVALID_SOCKET)
	{
		WSACleanup();
		return NULL;
	}

	// Sets keep alive to true for the socket
	char value = 1;
	setsockopt(client, SOL_SOCKET, SO_KEEPALIVE, &value, sizeof(value));

	return client;
}

// Socket thread
// Executes a loop reading messages and passing them to the Lua script

DWORD WINAPI Main(LPVOID param)
{
	int length;
	char buffer[BUFLEN];
	char *ptr, *context;
	const char *separator = "\n\r";
	const char *command = "clientnotifyregister schandlerid=0 event=any\n";

	// Starts the network loop
	for (running = true; running; )
	{
		// Connects to the local TeamSpeak ClientQuery and retries on failure
		client = Connect("127.0.0.1", "25639");
		if (client == NULL) continue;

		// Receives the initial ClientQuery headers
		for (int header = 182; header > 0; header -= length)
			length = recv(client, buffer, BUFLEN, 0);

		// Sends a "listen to all events" command and receives it's response
		send(client, command, strlen(command), 0);
		recv(client, buffer, BUFLEN, 0);

		// Reads messages from the ClientQuery input stream
		do
		{
			// Receives a message and null terminates it
			buffer[length = recv(client, buffer, BUFLEN, 0)] = 0;
			// Splits up different messages that might be in the same buffer
			ptr = strtok_s(buffer, separator, &context);
			while (ptr != NULL)
			{
				// Copies the messages to heap memory and saves them to a queue
				int length = strlen(ptr) + 1;
				char *queued = new char[length];
				strcpy_s(queued, length, ptr);
				queue.push(queued);
				ptr = strtok_s(NULL, separator, &context);
			}

		} while (length > 0 && running);

		// Once the connection is closed cleans up all of its resources
		{
			SOCKET socket = client;
			client = NULL;
			closesocket(socket);
			WSACleanup();
		}
	}

	return 0;
}

// Called by the Lua script
// Sends a message using the ClientQuery socket

int SendCommand(lua_State *L)
{
	// Stops if the socket is not connected or no message has been sent
	if (client == NULL || lua_type(L, 1) == -1) return 0;

	// Reads the message and sends it
	const char *buffer = lua_tostring(L, 1);
	send(client, buffer, strlen(buffer), 0);
	send(client, "\n", 1, 0);
	return 0;
}

// Requires the TeamSpeak Lua script and loads it into the game

void WINAPI OnRequire(lua_State *L, LPCSTR file, LPVOID param)
{
	// If the required file matches any we need to override
	if (strcmp(file, "lib/managers/chatmanager") == 0)
	{
		// Load that file into the game
		if (luaL_loadfile(L, "TeamSpeak/TeamSpeak.lua") == 0)
		{
			// And execute it
			lua_pcall(L, 0, LUA_MULTRET, 0);

			// Indexes the global TeamSpeak variable
			lua_getglobal(L, "TeamSpeak");
			int index = lua_gettop(L);

			// Creates a function called Send that maps to a C++ function
			lua_pushcfunction(L, &SendCommand);
			lua_setfield(L, index, "Send");

			// Saves the current Lua state and creates the network thread
			state = L;
			if (!running) CreateThread(NULL, 0, Main, NULL, 0, NULL);
		}
	}
}

// Runs on game ticks and sends messages from TeamSpeak to the Lua script

void WINAPI OnGameTick(lua_State *L, LPCSTR type, LPVOID param)
{
	// If the Lua state matches out initial script load state,
	// the tick type is an update tick and out message queue is not empty
	if (L == state && strcmp(type, "update") == 0 && !queue.empty())
	{
		char *message;

		// Indexes the global TeamSpeak variable
		lua_getglobal(L, "TeamSpeak");
		int index = lua_gettop(L);

		// And sends each message over to a Lua function
		while (queue.try_pop(message))
		{
			lua_getfield(L, index, "OnReceive");
			lua_pushlstring(L, message, strlen(message));
			lua_pcall(L, 1, 0, 0);
			delete[] message;
		}
	}
}

// DLL entry point

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
		// Registers a callback for when the game requires an internal file
		RegisterCallback(REQUIRE_CALLBACK, &OnRequire, NULL);
		// Registers a callback for game ticks
		RegisterCallback(GAMETICK_CALLBACK, &OnGameTick, NULL);
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
		break;
	case DLL_PROCESS_DETACH:
		// Stops the network loop
		running = FALSE;
		break;
	}
	return TRUE;
}
<commit_msg>improve comments<commit_after>#define WIN32_LEAN_AND_MEAN

#include <SDKDDKVer.h>
#include <windows.h>
#include "HookAPI.h"
#include <ws2tcpip.h>
#include <atomic>
#include <concurrent_queue.h>

#pragma comment (lib, "Ws2_32.lib")

// Default buffer length
#define BUFLEN 4096

// Socket used for TeamSpeak communication
std::atomic<SOCKET> client = NULL;
// Lua state used when requiring script files
std::atomic<lua_State *> state = NULL;
// Boolean telling the socket thread whether to keep the network loop running
std::atomic<bool> running = FALSE;
// Queue of commands to be sent to the Lua script
Concurrency::concurrent_queue<char *> queue;

// Tries to connect to TeamSpeak using ClientQuery
// Returns a new socket or NULL on failure

SOCKET Connect(PCSTR hostname, PCSTR port)
{
	// Creates a structure used for containing the socket information
	WSADATA wsaData;
	int result = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (result != 0) return NULL;

	// Creates a structure used for containing the connection information
	struct addrinfo *info, *ptr, hints;
	ZeroMemory(&hints, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;
	result = getaddrinfo(hostname, port, &hints, &info);
	if (result != 0)
	{
		WSACleanup();
		return NULL;
	}

	// Tries to connect to the specified address
	SOCKET client = INVALID_SOCKET;
	for (ptr = info; ptr != NULL; ptr = ptr->ai_next)
	{
		client = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
		if (client == INVALID_SOCKET)
		{
			WSACleanup();
			return NULL;
		}
		result = connect(client, ptr->ai_addr, (int)ptr->ai_addrlen);
		if (result == SOCKET_ERROR)
		{
			closesocket(client);
			client = INVALID_SOCKET;
			continue;
		}
		break;
	}

	// Releases the address information strcuture from memory
	freeaddrinfo(info);

	// Return NULL if the connection could not be established
	if (client == INVALID_SOCKET)
	{
		WSACleanup();
		return NULL;
	}

	// Sets keep alive to true for the socket
	char value = 1;
	setsockopt(client, SOL_SOCKET, SO_KEEPALIVE, &value, sizeof(value));

	return client;
}

// Socket thread
// Executes a loop reading messages and passing them to the Lua script

DWORD WINAPI Main(LPVOID param)
{
	int length;
	char buffer[BUFLEN];
	char *ptr, *context;
	const char *separator = "\n\r";
	const char *command = "clientnotifyregister schandlerid=0 event=any\n";

	// Starts the network loop
	for (running = true; running;)
	{
		// Connects to the local TeamSpeak ClientQuery and retries on failure
		client = Connect("127.0.0.1", "25639");
		if (client == NULL) continue;

		// Receives the initial ClientQuery headers
		for (int header = 182; header > 0; header -= length)
			length = recv(client, buffer, BUFLEN, 0);

		// Sends a "listen to all events" command and receives its response
		send(client, command, strlen(command), 0);
		recv(client, buffer, BUFLEN, 0);

		// Reads messages from the ClientQuery input stream
		do
		{
			// Receives a message and null terminates it
			buffer[length = recv(client, buffer, BUFLEN, 0)] = 0;
			// Splits up different messages that might be in the same buffer
			ptr = strtok_s(buffer, separator, &context);
			while (ptr != NULL)
			{
				// Copies the messages to heap memory and saves them to a queue
				int length = strlen(ptr) + 1;
				char *queued = new char[length];
				strcpy_s(queued, length, ptr);
				queue.push(queued);
				ptr = strtok_s(NULL, separator, &context);
			}

		} while (length > 0 && running);

		// Once the connection is closed cleans up all of its resources
		{
			SOCKET socket = client;
			client = NULL;
			closesocket(socket);
			WSACleanup();
		}
	}

	return 0;
}

// Called by the Lua script when a TeamSpeak command is used
// Sends a message using the ClientQuery socket

int SendCommand(lua_State *L)
{
	// Stops if the socket is not connected or no message has been sent
	if (client == NULL || lua_type(L, 1) == -1) return 0;

	// Reads the message and sends it
	const char *buffer = lua_tostring(L, 1);
	send(client, buffer, strlen(buffer), 0);
	send(client, "\n", 1, 0);
	return 0;
}

// Requires the TeamSpeak Lua script and loads it into the game

void WINAPI OnRequire(lua_State *L, LPCSTR file, LPVOID param)
{
	// If the required file matches any we need to override
	if (strcmp(file, "lib/managers/chatmanager") == 0)
	{
		// Load that file into the game
		if (luaL_loadfile(L, "TeamSpeak/TeamSpeak.lua") == 0)
		{
			// And execute it
			lua_pcall(L, 0, LUA_MULTRET, 0);

			// Indexes the global TeamSpeak variable
			lua_getglobal(L, "TeamSpeak");
			int index = lua_gettop(L);

			// Maps C++ functions to Lua variables inside the TeamSpeak object
			lua_pushcfunction(L, &SendCommand);
			lua_setfield(L, index, "Send");

			// Saves the current Lua state and creates the network thread
			state = L;
			if (!running) CreateThread(NULL, 0, Main, NULL, 0, NULL);
		}
	}
}

// Runs on game ticks and sends messages from TeamSpeak to the Lua script

void WINAPI OnGameTick(lua_State *L, LPCSTR type, LPVOID param)
{
	// If the Lua state matches out initial script load state,
	// the tick type is an update tick and out message queue is not empty
	if (L == state && strcmp(type, "update") == 0 && !queue.empty())
	{
		char *message;

		// Indexes the global TeamSpeak variable
		lua_getglobal(L, "TeamSpeak");
		int index = lua_gettop(L);

		// And sends each message over to a Lua function
		while (queue.try_pop(message))
		{
			lua_getfield(L, index, "OnReceive");
			lua_pushlstring(L, message, strlen(message));
			lua_pcall(L, 1, 0, 0);
			delete[] message;
		}
	}
}

// DLL entry point

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
		// Registers a callback for when the game requires an internal file
		RegisterCallback(REQUIRE_CALLBACK, &OnRequire, NULL);
		// Registers a callback for game ticks
		RegisterCallback(GAMETICK_CALLBACK, &OnGameTick, NULL);
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
		break;
	case DLL_PROCESS_DETACH:
		// Stops the network loop
		running = FALSE;
		break;
	}
	return TRUE;
}
<|endoftext|>