<commit_before>#if defined(TLANG_WITH_LLVM)
// A helper for the llvm backend

#include <llvm/Transforms/Utils/Cloning.h>
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/TargetRegistry.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Transforms/InstCombine/InstCombine.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Scalar/GVN.h"
#include "llvm/Transforms/Utils.h"
#include "llvm/Bitcode/BitcodeReader.h"
#include <llvm/Linker/Linker.h>
#include <llvm/Demangle/Demangle.h>

#include "util.h"
#include "taichi_llvm_context.h"
#include "backends/llvm_jit.h"

TLANG_NAMESPACE_BEGIN

static llvm::ExitOnError exit_on_err;

TaichiLLVMContext::TaichiLLVMContext(Arch arch) : arch(arch) {
  llvm::InitializeAllTargets();
  llvm::remove_fatal_error_handler();
  llvm::install_fatal_error_handler(
      [](void *user_data, const std::string &reason, bool gen_crash_diag) {
        TC_ERROR("LLVM Fatal Error: {}", reason);
      },
      nullptr);

  if (arch == Arch::x86_64) {
    llvm::InitializeNativeTarget();
    llvm::InitializeNativeTargetAsmPrinter();
    llvm::InitializeNativeTargetAsmParser();
  } else {
    LLVMInitializeNVPTXTarget();
    LLVMInitializeNVPTXTargetMC();
    LLVMInitializeNVPTXTargetInfo();
    LLVMInitializeNVPTXAsmPrinter();
  }
  ctx = std::make_unique<llvm::LLVMContext>();
  TC_INFO("Creating llvm context for arch: {}", arch_name(arch));
  jit = exit_on_err(TaichiLLVMJIT::create(arch));
}

TaichiLLVMContext::~TaichiLLVMContext() {
}

llvm::Type *TaichiLLVMContext::get_data_type(DataType dt) {
  if (dt == DataType::i32) {
    return llvm::Type::getInt32Ty(*ctx);
  } else if (dt == DataType::f32) {
    return llvm::Type::getFloatTy(*ctx);
  } else if (dt == DataType::f64) {
    return llvm::Type::getDoubleTy(*ctx);
  } else {
    TC_INFO(data_type_name(dt));
    TC_NOT_IMPLEMENTED
  }
  return nullptr;
}

std::string find_existing_command(const std::vector<std::string> &commands) {
  for (auto &cmd : commands) {
    if (command_exist(cmd)) {
      return cmd;
    }
  }
  TC_P(commands);
  TC_ERROR("No command found.");
}

std::string get_runtime_fn(Arch arch) {
  return fmt::format("runtime_{}.bc", arch_name(arch));
}

void compile_runtime_bitcode(Arch arch) {
  static std::set<int> runtime_compiled;
  if (runtime_compiled.find((int)arch) == runtime_compiled.end()) {
    auto clang = find_existing_command({"clang-7", "clang"});
    TC_ASSERT(command_exist("llvm-as"));
    TC_TRACE("Compiling runtime module bitcode...");
    auto runtime_folder = get_repo_dir() + "/src/runtime/";
    std::string macro = fmt::format(" -D ARCH_{} ", arch_name(arch));
    int ret = std::system(
        fmt::format(
            "{} -S {}runtime.cpp -o {}runtime.ll -emit-llvm -std=c++17 {}",
            clang, runtime_folder, runtime_folder, macro)
            .c_str());
    if (ret) {
      TC_ERROR("Runtime compilation failed.");
    }
    std::system(fmt::format("llvm-as {}runtime.ll -o {}{}", runtime_folder,
                            runtime_folder, get_runtime_fn(arch))
                    .c_str());
    runtime_compiled.insert((int)arch);
  }
}

void compile_runtimes() {
  compile_runtime_bitcode(Arch::x86_64);
#if defined(TLANG_WITH_CUDA)
  compile_runtime_bitcode(Arch::gpu);
#endif
}

std::string libdevice_path() {
#if defined(TLANG_WITH_CUDA)
  auto folder =
      fmt::format("/usr/local/cuda-{}/nvvm/libdevice/", TLANG_CUDA_VERSION);
  if (is_release()) {
    folder = compiled_lib_dir;
  }
  auto cuda_version_major = int(std::atof(TLANG_CUDA_VERSION));
  return fmt::format("{}/libdevice.{}.bc", folder, cuda_version_major);
#else
  TC_NOT_IMPLEMENTED;
  return "";
#endif
}

std::unique_ptr<llvm::Module> TaichiLLVMContext::get_init_module() {
  return clone_runtime_module();
}

std::unique_ptr<llvm::Module> module_from_bitcode_file(std::string bitcode_path,
                                                       llvm::LLVMContext *ctx) {
  TI_AUTO_PROF
  std::ifstream ifs(bitcode_path);
  std::string bitcode(std::istreambuf_iterator<char>(ifs),
                      (std::istreambuf_iterator<char>()));
  auto runtime =
      parseBitcodeFile(MemoryBufferRef(bitcode, "runtime_bitcode"), *ctx);
  if (!runtime) {
    TC_ERROR("Runtime bitcode load failure.");
  }
  bool module_broken = llvm::verifyModule(*runtime.get(), &llvm::errs());
  TC_ERROR_IF(module_broken, "Module broken");
  return std::move(runtime.get());
}

std::unique_ptr<llvm::Module> TaichiLLVMContext::clone_runtime_module() {
  TI_AUTO_PROF
  if (!runtime_module) {
    if (is_release()) {
      runtime_module = module_from_bitcode_file(
          fmt::format("{}/{}", compiled_lib_dir, get_runtime_fn(arch)),
          ctx.get());
    } else {
      compile_runtime_bitcode(arch);
      runtime_module = module_from_bitcode_file(
          get_repo_dir() + fmt::format("/src/runtime/{}", get_runtime_fn(arch)),
          ctx.get());
    }
    if (arch == Arch::gpu) {
      runtime_module->setTargetTriple("nvptx64-nvidia-cuda");

      auto patch_intrinsic = [&](std::string name, Intrinsic::ID intrin,
                                 bool ret = true) {
        auto func = runtime_module->getFunction(name);
        func->getEntryBlock().eraseFromParent();
        auto bb = llvm::BasicBlock::Create(*ctx, "entry", func);
        IRBuilder<> builder(*ctx);
        builder.SetInsertPoint(bb);
        if (ret) {
          builder.CreateRet(builder.CreateIntrinsic(intrin, {}, {}));
        } else {
          builder.CreateIntrinsic(intrin, {}, {});
          builder.CreateRetVoid();
        }
      };

      patch_intrinsic("thread_idx", Intrinsic::nvvm_read_ptx_sreg_tid_x);
      patch_intrinsic("block_idx", Intrinsic::nvvm_read_ptx_sreg_ctaid_x);
      patch_intrinsic("block_dim", Intrinsic::nvvm_read_ptx_sreg_ntid_x);
      patch_intrinsic("grid_dim", Intrinsic::nvvm_read_ptx_sreg_nctaid_x);
      patch_intrinsic("block_barrier", Intrinsic::nvvm_barrier0, false);

      link_module_with_libdevice(runtime_module);

      // runtime_module->print(llvm::errs(), nullptr);
    }
    /*
    for (auto &f : *runtime_module) {
      TC_INFO("Loaded runtime function: {}", std::string(f.getName()));
    }
    */
  }
  return llvm::CloneModule(*runtime_module);
}

void TaichiLLVMContext::link_module_with_libdevice(
    std::unique_ptr<llvm::Module> &module) {
  TI_AUTO_PROF
  auto libdevice_module = module_from_bitcode_file(libdevice_path(), ctx.get());

  std::vector<std::string> libdevice_function_names;
  for (auto &f : *libdevice_module) {
    if (!f.isDeclaration()) {
      libdevice_function_names.push_back(f.getName());
    }
  }

  libdevice_module->setTargetTriple("nvptx64-nvidia-cuda");
  module->setDataLayout(libdevice_module->getDataLayout());

  bool failed = llvm::Linker::linkModules(*module, std::move(libdevice_module));
  if (failed) {
    TC_ERROR("CUDA libdevice linking failure.");
  }

  for (auto func_name : libdevice_function_names) {
    module->getFunction(func_name)->setLinkage(Function::InternalLinkage);
  }
}

std::unique_ptr<llvm::Module> TaichiLLVMContext::clone_struct_module() {
  TC_ASSERT(struct_module);
  return llvm::CloneModule(*struct_module);
}

void TaichiLLVMContext::set_struct_module(
    const std::unique_ptr<llvm::Module> &module) {
  TC_ASSERT(module);
  if (llvm::verifyModule(*module, &llvm::errs())) {
    module->print(llvm::errs(), nullptr);
    TC_ERROR("module broken");
  }
  struct_module = llvm::CloneModule(*module);
}

template <typename T>
llvm::Value *TaichiLLVMContext::get_constant(DataType dt, T t) {
  if (dt == DataType::f32) {
    return llvm::ConstantFP::get(*ctx, llvm::APFloat((float32)t));
  } else if (dt == DataType::f64) {
    return llvm::ConstantFP::get(*ctx, llvm::APFloat((float64)t));
  } else if (dt == DataType::i32) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(32, t, true));
  } else if (dt == DataType::u32) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(32, t, false));
  } else if (dt == DataType::i64) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(64, t, true));
  } else if (dt == DataType::u64) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(64, t, false));
  } else {
    TC_NOT_IMPLEMENTED
    return nullptr;
  }
}

template llvm::Value *TaichiLLVMContext::get_constant(DataType dt, int32 t);

template <typename T>
llvm::Value *TaichiLLVMContext::get_constant(T t) {
  using TargetType = T;
  if constexpr (std::is_same_v<TargetType, float32> ||
                std::is_same_v<TargetType, float64>) {
    return llvm::ConstantFP::get(*ctx, llvm::APFloat(t));
  } else if (std::is_same_v<TargetType, bool>) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(1, t, false));
  } else if (std::is_same_v<TargetType, int32> ||
             std::is_same_v<TargetType, uint32>) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(32, t, true));
  } else if (std::is_same_v<TargetType, int64> ||
             std::is_same_v<TargetType, uint64>) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(64, t, true));
  } else {
    TC_NOT_IMPLEMENTED
    return nullptr;
  }
}

std::string TaichiLLVMContext::type_name(llvm::Type *type) {
  std::string type_name;
  llvm::raw_string_ostream rso(type_name);
  type->print(rso);
  return rso.str();
}

std::size_t TaichiLLVMContext::get_type_size(llvm::Type *type) {
  return jit->get_type_size(type);
}

template llvm::Value *TaichiLLVMContext::get_constant(float32 t);
template llvm::Value *TaichiLLVMContext::get_constant(float64 t);

template llvm::Value *TaichiLLVMContext::get_constant(bool t);

template llvm::Value *TaichiLLVMContext::get_constant(int32 t);
template llvm::Value *TaichiLLVMContext::get_constant(uint32 t);

template llvm::Value *TaichiLLVMContext::get_constant(int64 t);
template llvm::Value *TaichiLLVMContext::get_constant(uint64 t);

TLANG_NAMESPACE_END
#endif
<commit_msg>num_instructions(llvm::Function *)<commit_after>#if defined(TLANG_WITH_LLVM)
// A helper for the llvm backend

#include <llvm/Transforms/Utils/Cloning.h>
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/TargetRegistry.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Transforms/InstCombine/InstCombine.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Scalar/GVN.h"
#include "llvm/Transforms/Utils.h"
#include "llvm/Bitcode/BitcodeReader.h"
#include <llvm/Linker/Linker.h>
#include <llvm/Demangle/Demangle.h>

#include "util.h"
#include "taichi_llvm_context.h"
#include "backends/llvm_jit.h"

TLANG_NAMESPACE_BEGIN

static llvm::ExitOnError exit_on_err;

TaichiLLVMContext::TaichiLLVMContext(Arch arch) : arch(arch) {
  llvm::InitializeAllTargets();
  llvm::remove_fatal_error_handler();
  llvm::install_fatal_error_handler(
      [](void *user_data, const std::string &reason, bool gen_crash_diag) {
        TC_ERROR("LLVM Fatal Error: {}", reason);
      },
      nullptr);

  if (arch == Arch::x86_64) {
    llvm::InitializeNativeTarget();
    llvm::InitializeNativeTargetAsmPrinter();
    llvm::InitializeNativeTargetAsmParser();
  } else {
    LLVMInitializeNVPTXTarget();
    LLVMInitializeNVPTXTargetMC();
    LLVMInitializeNVPTXTargetInfo();
    LLVMInitializeNVPTXAsmPrinter();
  }
  ctx = std::make_unique<llvm::LLVMContext>();
  TC_INFO("Creating llvm context for arch: {}", arch_name(arch));
  jit = exit_on_err(TaichiLLVMJIT::create(arch));
}

TaichiLLVMContext::~TaichiLLVMContext() {
}

llvm::Type *TaichiLLVMContext::get_data_type(DataType dt) {
  if (dt == DataType::i32) {
    return llvm::Type::getInt32Ty(*ctx);
  } else if (dt == DataType::f32) {
    return llvm::Type::getFloatTy(*ctx);
  } else if (dt == DataType::f64) {
    return llvm::Type::getDoubleTy(*ctx);
  } else {
    TC_INFO(data_type_name(dt));
    TC_NOT_IMPLEMENTED
  }
  return nullptr;
}

std::string find_existing_command(const std::vector<std::string> &commands) {
  for (auto &cmd : commands) {
    if (command_exist(cmd)) {
      return cmd;
    }
  }
  TC_P(commands);
  TC_ERROR("No command found.");
}

std::string get_runtime_fn(Arch arch) {
  return fmt::format("runtime_{}.bc", arch_name(arch));
}

void compile_runtime_bitcode(Arch arch) {
  TI_AUTO_PROF;
  static std::set<int> runtime_compiled;
  if (runtime_compiled.find((int)arch) == runtime_compiled.end()) {
    auto clang = find_existing_command({"clang-7", "clang"});
    TC_ASSERT(command_exist("llvm-as"));
    TC_TRACE("Compiling runtime module bitcode...");
    auto runtime_folder = get_repo_dir() + "/src/runtime/";
    std::string macro = fmt::format(" -D ARCH_{} ", arch_name(arch));
    int ret = std::system(
        fmt::format(
            "{} -S {}runtime.cpp -o {}runtime.ll -emit-llvm -std=c++17 {}",
            clang, runtime_folder, runtime_folder, macro)
            .c_str());
    if (ret) {
      TC_ERROR("Runtime compilation failed.");
    }
    std::system(fmt::format("llvm-as {}runtime.ll -o {}{}", runtime_folder,
                            runtime_folder, get_runtime_fn(arch))
                    .c_str());
    runtime_compiled.insert((int)arch);
  }
}

void compile_runtimes() {
  compile_runtime_bitcode(Arch::x86_64);
#if defined(TLANG_WITH_CUDA)
  compile_runtime_bitcode(Arch::gpu);
#endif
}

std::string libdevice_path() {
#if defined(TLANG_WITH_CUDA)
  auto folder =
      fmt::format("/usr/local/cuda-{}/nvvm/libdevice/", TLANG_CUDA_VERSION);
  if (is_release()) {
    folder = compiled_lib_dir;
  }
  auto cuda_version_major = int(std::atof(TLANG_CUDA_VERSION));
  return fmt::format("{}/libdevice.{}.bc", folder, cuda_version_major);
#else
  TC_NOT_IMPLEMENTED;
  return "";
#endif
}

std::unique_ptr<llvm::Module> TaichiLLVMContext::get_init_module() {
  return clone_runtime_module();
}

std::unique_ptr<llvm::Module> module_from_bitcode_file(std::string bitcode_path,
                                                       llvm::LLVMContext *ctx) {
  TI_AUTO_PROF
  std::ifstream ifs(bitcode_path);
  std::string bitcode(std::istreambuf_iterator<char>(ifs),
                      (std::istreambuf_iterator<char>()));
  auto runtime =
      parseBitcodeFile(MemoryBufferRef(bitcode, "runtime_bitcode"), *ctx);
  if (!runtime) {
    TC_ERROR("Runtime bitcode load failure.");
  }
  bool module_broken = llvm::verifyModule(*runtime.get(), &llvm::errs());
  TC_ERROR_IF(module_broken, "Module broken");
  return std::move(runtime.get());
}

int num_instructions(llvm::Function *func) {
  int counter = 0;
  for (BasicBlock &bb : *func)
    counter += std::distance(bb.begin(), bb.end());
  return counter;
}

std::unique_ptr<llvm::Module> TaichiLLVMContext::clone_runtime_module() {
  TI_AUTO_PROF
  if (!runtime_module) {
    if (is_release()) {
      runtime_module = module_from_bitcode_file(
          fmt::format("{}/{}", compiled_lib_dir, get_runtime_fn(arch)),
          ctx.get());
    } else {
      compile_runtime_bitcode(arch);
      runtime_module = module_from_bitcode_file(
          get_repo_dir() + fmt::format("/src/runtime/{}", get_runtime_fn(arch)),
          ctx.get());
    }
    if (arch == Arch::gpu) {
      runtime_module->setTargetTriple("nvptx64-nvidia-cuda");

      auto patch_intrinsic = [&](std::string name, Intrinsic::ID intrin,
                                 bool ret = true) {
        TC_PROFILER("patch intrinsic");
        auto func = runtime_module->getFunction(name);
        func->getEntryBlock().eraseFromParent();
        auto bb = llvm::BasicBlock::Create(*ctx, "entry", func);
        IRBuilder<> builder(*ctx);
        builder.SetInsertPoint(bb);
        if (ret) {
          builder.CreateRet(builder.CreateIntrinsic(intrin, {}, {}));
        } else {
          builder.CreateIntrinsic(intrin, {}, {});
          builder.CreateRetVoid();
        }
      };

      patch_intrinsic("thread_idx", Intrinsic::nvvm_read_ptx_sreg_tid_x);
      patch_intrinsic("block_idx", Intrinsic::nvvm_read_ptx_sreg_ctaid_x);
      patch_intrinsic("block_dim", Intrinsic::nvvm_read_ptx_sreg_ntid_x);
      patch_intrinsic("grid_dim", Intrinsic::nvvm_read_ptx_sreg_nctaid_x);
      patch_intrinsic("block_barrier", Intrinsic::nvvm_barrier0, false);

      link_module_with_libdevice(runtime_module);

      // runtime_module->print(llvm::errs(), nullptr);
    }
    /*
    for (auto &f : *runtime_module) {
      TC_INFO("Loaded runtime function: {} (inst. count= {})",
              std::string(f.getName()), num_instructions(&f));
    }
    */
  }
  return llvm::CloneModule(*runtime_module);
}

void TaichiLLVMContext::link_module_with_libdevice(
    std::unique_ptr<llvm::Module> &module) {
  TI_AUTO_PROF
  auto libdevice_module = module_from_bitcode_file(libdevice_path(), ctx.get());

  std::vector<std::string> libdevice_function_names;
  for (auto &f : *libdevice_module) {
    if (!f.isDeclaration()) {
      libdevice_function_names.push_back(f.getName());
    }
  }

  libdevice_module->setTargetTriple("nvptx64-nvidia-cuda");
  module->setDataLayout(libdevice_module->getDataLayout());

  bool failed = llvm::Linker::linkModules(*module, std::move(libdevice_module));
  if (failed) {
    TC_ERROR("CUDA libdevice linking failure.");
  }

  for (auto func_name : libdevice_function_names) {
    module->getFunction(func_name)->setLinkage(Function::InternalLinkage);
  }
}

std::unique_ptr<llvm::Module> TaichiLLVMContext::clone_struct_module() {
  TC_ASSERT(struct_module);
  return llvm::CloneModule(*struct_module);
}

void TaichiLLVMContext::set_struct_module(
    const std::unique_ptr<llvm::Module> &module) {
  TC_ASSERT(module);
  if (llvm::verifyModule(*module, &llvm::errs())) {
    module->print(llvm::errs(), nullptr);
    TC_ERROR("module broken");
  }
  struct_module = llvm::CloneModule(*module);
}

template <typename T>
llvm::Value *TaichiLLVMContext::get_constant(DataType dt, T t) {
  if (dt == DataType::f32) {
    return llvm::ConstantFP::get(*ctx, llvm::APFloat((float32)t));
  } else if (dt == DataType::f64) {
    return llvm::ConstantFP::get(*ctx, llvm::APFloat((float64)t));
  } else if (dt == DataType::i32) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(32, t, true));
  } else if (dt == DataType::u32) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(32, t, false));
  } else if (dt == DataType::i64) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(64, t, true));
  } else if (dt == DataType::u64) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(64, t, false));
  } else {
    TC_NOT_IMPLEMENTED
    return nullptr;
  }
}

template llvm::Value *TaichiLLVMContext::get_constant(DataType dt, int32 t);

template <typename T>
llvm::Value *TaichiLLVMContext::get_constant(T t) {
  using TargetType = T;
  if constexpr (std::is_same_v<TargetType, float32> ||
                std::is_same_v<TargetType, float64>) {
    return llvm::ConstantFP::get(*ctx, llvm::APFloat(t));
  } else if (std::is_same_v<TargetType, bool>) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(1, t, false));
  } else if (std::is_same_v<TargetType, int32> ||
             std::is_same_v<TargetType, uint32>) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(32, t, true));
  } else if (std::is_same_v<TargetType, int64> ||
             std::is_same_v<TargetType, uint64>) {
    return llvm::ConstantInt::get(*ctx, llvm::APInt(64, t, true));
  } else {
    TC_NOT_IMPLEMENTED
    return nullptr;
  }
}

std::string TaichiLLVMContext::type_name(llvm::Type *type) {
  std::string type_name;
  llvm::raw_string_ostream rso(type_name);
  type->print(rso);
  return rso.str();
}

std::size_t TaichiLLVMContext::get_type_size(llvm::Type *type) {
  return jit->get_type_size(type);
}

template llvm::Value *TaichiLLVMContext::get_constant(float32 t);
template llvm::Value *TaichiLLVMContext::get_constant(float64 t);

template llvm::Value *TaichiLLVMContext::get_constant(bool t);

template llvm::Value *TaichiLLVMContext::get_constant(int32 t);
template llvm::Value *TaichiLLVMContext::get_constant(uint32 t);

template llvm::Value *TaichiLLVMContext::get_constant(int64 t);
template llvm::Value *TaichiLLVMContext::get_constant(uint64 t);

TLANG_NAMESPACE_END
#endif
<|endoftext|>


<commit_before>
/**
 * 所有 #include <外部文件> 的头文件 都必须保证 本文件被首先#include
 * 以下项目可以由本文件保证包含：
 * 	<cstdint>
 * 	<cstddef>
 * 	并保证 全局名称空间包含所有stdint。
 * 	CS_DEBUG = 1 时保证 #include <iostream>
 */

#pragma once

/**
 * 平台兼容原则：以不大的代价换取一定的兼容性；以首要目标平台0-overhead为前提。
 */

#ifndef NDEBUG
#	ifndef CS_DEBUG
#   	define CS_DEBUG		2
#	endif
#	ifndef CS_LOG_ON
#		define CS_LOG_ON	1
#	endif
#else
#	undef CS_DEBUG
#	define CS_DEBUG			0
#	undef CS_LOG_ON
#	define CS_LOG_ON		0
#endif

#if ! CS_DEBUG
#ifndef NDEBUG
#	define NDEBUG
#endif
#	define BOOST_DISABLE_ASSERTS
#endif

// Use wchar_t/wstring wcout/wcerr/wcin wfstream or not.
// It's just a standard, not make used in this file.
#ifndef CS_USE_WCS
#	define CS_USE_WCS	0
#endif
// Use wcout/wcerr/wcin or not. This will make use in this file.
#ifndef CS_USE_WIO
#	define CS_USE_WIO	0
#endif

#if defined(CS_USE_WCS) && CS_USE_WCS
#	define CS_STR_LITER(str_liter)		L##str_liter
#else
#	define CS_STR_LITER(str_liter)		str_liter
#endif

#ifdef __cplusplus
// avoid from boost if possible.
#	if defined(__GNUC__)	\
		&& (!defined(__GXX_EXPERIMENTAL_CXX0X__) || !__GXX_EXPERIMENTAL_CXX0X__)
#		include <boost/cstdint.hpp>
#	else
#		include <cstdint>
#	endif
#	include <cstddef>
#else
#	include <stdint.h>
#	include <stddef.h>
#endif

#define CS_EXIT_STATUS_FAILED EXIT_FAILURE

#ifdef __cplusplus
#	define CS_PREF_STD(symbol)		::std::symbol
#else
#	define CS_PREF_STD(symbol)		symbol
#endif

#if CS_USE_WIO
#	define CS_STDIN		::std::wcin
#	define CS_STDOUT	::std::wcout
#	define CS_STDERR	::std::wcerr
#else
#	define CS_STDIN		::std::cin
#	define CS_STDOUT	::std::cout
#	define CS_STDERR	::std::cerr
#endif

#define CS_OC_BLACK(...)	"\033[32;30;5m" << __VA_ARGS__ << "\033[0m"
#define CS_OC_BLUE(...)		"\033[32;34;5m" << __VA_ARGS__ << "\033[0m"
#define CS_OC_RED(...)		"\033[32;31;5m" << __VA_ARGS__ << "\033[0m"
#define CS_OC_GREEN(...)	"\033[32;49;5m" << __VA_ARGS__ << "\033[0m"

// line-seperator, should only use for IO..
#ifdef __linux
#	define CS_LINESEP		'\n'
#	define CS_LINESEP_STR	"\n"
#else
#	ifdef(__windows)
#		define CS_LINESEP		"\r\n"
#		define CS_LINESEP_STR	CS_LINESEP
#	else
#		define CS_LINESEP		::std::endl
#		define CS_LINESEP_STR	::std::endl()
#	endif
#endif

#if CS_DEBUG
#   include <iostream>
#   if CS_DEBUG > 1
#		if CS_USE_WCS
#			define CS_OUT(ostream, ...)											\
				do {															\
				ostream << __FILE__ << ":" << __LINE__							\
					<< ":" << __FUNCTION__ << "()" << ":\t"						\
					<< __VA_ARGS__												\
					<< ::std::endl;												\
				} while(false)
#		else
#       	define CS_OUT(ostream, ...)											\
				do {															\
				ostream << CS_OC_BLACK(__FILE__ << ":" << __LINE__)				\
					<< ":" << CS_OC_BLUE(__FUNCTION__ << "()") << ":\t"			\
					<< CS_OC_GREEN(__VA_ARGS__)									\
					<< ::std::endl;												\
				} while(false)
#		endif
#   else
#		define CS_OUT(ostream, ...)	do {ostream << __VA_ARGS__ << ::std::endl;} while(false)
#   endif
#else
#	define CS_OUT(ostream, ...)
#endif

#define CS_SAY(...)		CS_OUT(CS_STDOUT, __VA_ARGS__)

// NOTE: 要自行 #include <iostream>
#define CS_ECHO(...)	CS_STDOUT << __VA_ARGS__ << ::std::endl;
#define CS_ERR(...)		CS_STDERR << __VA_ARGS__ << ::std::endl;

// #define CS_DUMP(...)	CS_OUT(CS_STDOUT, GOL_OC_BLUE(#__VA_ARGS__) << ": " << GOL_OC_GREEN(__VA_ARGS__))

#define CS_DUMP(...)																	\
	CS_OUT(CS_STDOUT, CS_OC_RED(#__VA_ARGS__) << ":[" << CS_OC_GREEN(__VA_ARGS__) << "]")
#define CS_DUMP_N(...) 																	\
	CS_OUT(CS_STDOUT, CS_OC_RED(#__VA_ARGS__) << ":" << ::std::endl << CS_OC_GREEN(__VA_ARGS__))

// #define CS_ERR(msg)
//	CS_STDERR << msg << ::std::endl
// --- 要自行 #include <cstdlib>
#define CS_DIE(msg)																		\
	CS_ERR(msg); 																		\
    ::std::exit(CS_EXIT_STATUS_FAILED)

#ifndef CS_USED
#	ifdef __GNUC__
#		define CS_USED __attribute__((__used__))
#	else
#		define CS_USED
#	endif
#endif

#ifndef CS_ATTR_CONST
#	ifdef __GNUC__
#		define CS_ATTR_CONST __attribute__((__const__))
#	else
#		define CS_ATTR_CONST
#	endif
#endif

#ifndef CS_MUST_CHECK
#	ifdef __GNUC__
#		define CS_MUST_CHECK __attribute__((warn_unused_result))
#	else
#		define CS_MUST_CHECK
#	endif
#endif

#ifndef CS_DEPRECATED
#	ifdef __GNUC__
#		define CS_DEPRECATED __attribute__((deprecated))
#	else
#		define CS_DEPRECATED
#	endif
#endif

#ifndef CS_FORCE_INLINE
#   ifdef __GNUC__
#       define CS_FORCE_INLINE __attribute__((always_inline))
#   else
#       define CS_FORCE_INLINE inline
#   endif
#endif

#if !defined(CS_LIKELY) && !defined(CS_UNLIKELY)
#	ifdef __GNUC__
#		define CS_LIKELY(...)		__builtin_expect(!!(__VA_ARGS__), 1)
#		define CS_UNLIKELY(...)		__builtin_expect(!!(__VA_ARGS__), 0)
#		define CS_BLIKELY(...)		__builtin_expect((__VA_ARGS__), true)
#		define CS_BUNLIKELY(...)	__builtin_expect((__VA_ARGS__), false)
#	else
#		define CS_LIKELY(expr)		(expr)
#		define CS_UNLIKELY(expr)	(expr)
#		define CS_BLIKELY(expr)		(expr)
#		define CS_BUNLIKELY(expr)	(expr)
#	endif
#endif

#ifndef CS_PREFETCH
#	ifdef __GNUC__
#		define CS_PREFETCH(addr,rw,clean) __builtin_prefetch(addr,rw,clean)
#	else
#		define CS_PREFETCH(addr,rw,clean)
#	endif
#endif

#define CS_RETURN_IF_NORMAL(cond, ...)		\
if (cond)									\
{											\
	return __VA_ARGS__;						\
}
#define CS_RETURN_IF(cond, ...) CS_RETURN_IF_NORMAL(CS_BUNLIKELY(cond), __VA_ARGS__)

#define CS_ABORT_IF_NORMAL(cond)			\
if (cond)									\
{											\
	::std::exit(CS_EXIT_STATUS_FAILED);		\
}
#define CS_ABORT_IF(cond) CS_ABORT_IF_NORMAL(CS_BUNLIKELY(cond))

#define _CS_DIE_IF_NORMAL(cond, info, why)	\
if (cond)									\
{											\
	CS_DIE(__FILE__ << ":" << __LINE__ << "::" << __FUNCTION__ << "() die on " << info << ", results from {" << #why << "}.");  \
}
#define CS_DIE_IF_NORMAL(cond, info) CS_DIE_IF_NORMAL(cond, info, cond)
#define CS_DIE_IF(cond, info) _CS_DIE_IF_NORMAL(CS_BUNLIKELY(cond), info, cond)
<commit_msg>revision<commit_after>
/**
 * 所有 #include <外部文件> 的头文件 都必须保证 本文件被首先#include
 * 以下项目可以由本文件保证包含：
 * 	<cstdint>
 * 	<cstddef>
 * 	并保证 全局名称空间包含所有stdint。
 * 	CS_DEBUG = 1 时保证 #include <iostream>
 */

#pragma once

/**
 * 平台兼容原则：以不大的代价换取一定的兼容性；以首要目标平台0-overhead为前提。
 */

#ifndef NDEBUG
#	ifndef CS_DEBUG
#   	define CS_DEBUG		2
#	endif
#	ifndef CS_LOG_ON
#		define CS_LOG_ON	1
#	endif
#else
#	undef CS_DEBUG
#	define CS_DEBUG			0
#	undef CS_LOG_ON
#	define CS_LOG_ON		0
#endif

#if ! CS_DEBUG
#ifndef NDEBUG
#	define NDEBUG
#endif
#	define BOOST_DISABLE_ASSERTS
#endif

// Use wchar_t/wstring wcout/wcerr/wcin wfstream or not.
// It's just a standard, not make used in this file.
#ifndef CS_USE_WCS
#	define CS_USE_WCS	0
#endif
// Use wcout/wcerr/wcin or not. This will make use in this file.
#ifndef CS_USE_WIO
#	define CS_USE_WIO	0
#endif

#if defined(CS_USE_WCS) && CS_USE_WCS
#	define CS_STR_LITER(str_liter)		L##str_liter
#else
#	define CS_STR_LITER(str_liter)		#str_liter
#endif

#ifdef __cplusplus
// avoid from boost if possible.
#	if defined(__GNUC__)	\
		&& (!defined(__GXX_EXPERIMENTAL_CXX0X__) || !__GXX_EXPERIMENTAL_CXX0X__)
#		include <boost/cstdint.hpp>
#	else
#		include <cstdint>
#	endif
#	include <cstddef>
#else
#	include <stdint.h>
#	include <stddef.h>
#endif

#define CS_EXIT_STATUS_FAILED EXIT_FAILURE

#ifdef __cplusplus
#	define CS_PREF_STD(symbol)		::std::symbol
#else
#	define CS_PREF_STD(symbol)		symbol
#endif

#if CS_USE_WIO
#	define CS_STDIN		::std::wcin
#	define CS_STDOUT	::std::wcout
#	define CS_STDERR	::std::wcerr
#else
#	define CS_STDIN		::std::cin
#	define CS_STDOUT	::std::cout
#	define CS_STDERR	::std::cerr
#endif

#define CS_OC_BLACK(...)	"\033[32;30;5m" << __VA_ARGS__ << "\033[0m"
#define CS_OC_BLUE(...)		"\033[32;34;5m" << __VA_ARGS__ << "\033[0m"
#define CS_OC_RED(...)		"\033[32;31;5m" << __VA_ARGS__ << "\033[0m"
#define CS_OC_GREEN(...)	"\033[32;49;5m" << __VA_ARGS__ << "\033[0m"

// line-seperator, should only use for IO..
#ifdef __linux
#	define CS_LINESEP		'\n'
#	define CS_LINESEP_STR	"\n"
#else
#	ifdef(__windows)
#		define CS_LINESEP		"\r\n"
#		define CS_LINESEP_STR	CS_LINESEP
#	else
#		define CS_LINESEP		::std::endl
#		define CS_LINESEP_STR	::std::endl()
#	endif
#endif

#if CS_DEBUG
#   include <iostream>
#   if CS_DEBUG > 1
#		if CS_USE_WCS
#			define CS_OUT(ostream, ...)											\
				do {															\
				ostream << __FILE__ << ":" << __LINE__							\
					<< ":" << __FUNCTION__ << "()" << ":\t"						\
					<< __VA_ARGS__												\
					<< ::std::endl;												\
				} while(false)
#		else
#       	define CS_OUT(ostream, ...)											\
				do {															\
				ostream << CS_OC_BLACK(__FILE__ << ":" << __LINE__)				\
					<< ":" << CS_OC_BLUE(__FUNCTION__ << "()") << ":\t"			\
					<< CS_OC_GREEN(__VA_ARGS__)									\
					<< ::std::endl;												\
				} while(false)
#		endif
#   else
#		define CS_OUT(ostream, ...)	do {ostream << __VA_ARGS__ << ::std::endl;} while(false)
#   endif
#else
#	define CS_OUT(ostream, ...)
#endif

#define CS_SAY(...)		CS_OUT(CS_STDOUT, __VA_ARGS__)

// NOTE: 要自行 #include <iostream>
#define CS_ECHO(...)	CS_STDOUT << __VA_ARGS__ << ::std::endl;
#define CS_ERR(...)		CS_STDERR << __VA_ARGS__ << ::std::endl;

// #define CS_DUMP(...)	CS_OUT(CS_STDOUT, GOL_OC_BLUE(#__VA_ARGS__) << ": " << GOL_OC_GREEN(__VA_ARGS__))

#define CS_DUMP(...)																	\
	CS_OUT(CS_STDOUT, CS_OC_RED(#__VA_ARGS__) << ":[" << CS_OC_GREEN(__VA_ARGS__) << "]")
#define CS_DUMP_N(...) 																	\
	CS_OUT(CS_STDOUT, CS_OC_RED(#__VA_ARGS__) << ":" << ::std::endl << CS_OC_GREEN(__VA_ARGS__))

// #define CS_ERR(msg)
//	CS_STDERR << msg << ::std::endl
// --- 要自行 #include <cstdlib>
#define CS_DIE(msg)																		\
	CS_ERR(msg); 																		\
    ::std::exit(CS_EXIT_STATUS_FAILED)

#ifndef CS_USED
#	ifdef __GNUC__
#		define CS_USED __attribute__((__used__))
#	else
#		define CS_USED
#	endif
#endif

#ifndef CS_ATTR_CONST
#	ifdef __GNUC__
#		define CS_ATTR_CONST __attribute__((__const__))
#	else
#		define CS_ATTR_CONST
#	endif
#endif

#ifndef CS_MUST_CHECK
#	ifdef __GNUC__
#		define CS_MUST_CHECK __attribute__((warn_unused_result))
#	else
#		define CS_MUST_CHECK
#	endif
#endif

#ifndef CS_DEPRECATED
#	ifdef __GNUC__
#		define CS_DEPRECATED __attribute__((deprecated))
#	else
#		define CS_DEPRECATED
#	endif
#endif

#ifndef CS_FORCE_INLINE
#   ifdef __GNUC__
#       define CS_FORCE_INLINE __attribute__((always_inline))
#   else
#       define CS_FORCE_INLINE inline
#   endif
#endif

#if !defined(CS_LIKELY) && !defined(CS_UNLIKELY)
#	ifdef __GNUC__
#		define CS_LIKELY(...)		__builtin_expect(!!(__VA_ARGS__), 1)
#		define CS_UNLIKELY(...)		__builtin_expect(!!(__VA_ARGS__), 0)
#		define CS_BLIKELY(...)		__builtin_expect((__VA_ARGS__), true)
#		define CS_BUNLIKELY(...)	__builtin_expect((__VA_ARGS__), false)
#	else
#		define CS_LIKELY(expr)		(expr)
#		define CS_UNLIKELY(expr)	(expr)
#		define CS_BLIKELY(expr)		(expr)
#		define CS_BUNLIKELY(expr)	(expr)
#	endif
#endif

#ifndef CS_PREFETCH
#	ifdef __GNUC__
#		define CS_PREFETCH(addr,rw,clean) __builtin_prefetch(addr,rw,clean)
#	else
#		define CS_PREFETCH(addr,rw,clean)
#	endif
#endif

#define CS_RETURN_IF_NORMAL(cond, ...)		\
if (cond)									\
{											\
	return __VA_ARGS__;						\
}
#define CS_RETURN_IF(cond, ...) CS_RETURN_IF_NORMAL(CS_BUNLIKELY(cond), __VA_ARGS__)

#define CS_ABORT_IF_NORMAL(cond)			\
if (cond)									\
{											\
	::std::exit(CS_EXIT_STATUS_FAILED);		\
}
#define CS_ABORT_IF(cond) CS_ABORT_IF_NORMAL(CS_BUNLIKELY(cond))

#define _CS_DIE_IF_NORMAL(cond, info, why)	\
if (cond)									\
{											\
	CS_DIE(__FILE__ << ":" << __LINE__ << "::" << __FUNCTION__ << "() die on " << info << ", results from {" << #why << "}.");  \
}
#define CS_DIE_IF_NORMAL(cond, info) CS_DIE_IF_NORMAL(cond, info, cond)
#define CS_DIE_IF(cond, info) _CS_DIE_IF_NORMAL(CS_BUNLIKELY(cond), info, cond)
<|endoftext|>


<commit_before>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <string>

using namespace std;

class Date {
	int d, m, y;
	mutable bool cache_valid;
	mutable string cache;

	void compute_cache_value() const;

	public:
	void init(int dd, int mm, int yy);
	Date(int dd, int mm, int yy);
	void add_year(int n);
	void add_month(int n);
	void add_day(int n);
	int day() const { return d; }
	int month() const { return m; }
	int year() const { return y; }
	string string_rep() const;
};

void Date::compute_cache_value() const {
	char tmp[20];
	sprintf(tmp, "%04d-%02d-%02d", y, m, d);
	string tmpstr(tmp);
	cache = tmpstr;
	cache_valid = true;
	return;
}

string Date::string_rep() const {
	if (!cache_valid) {
		compute_cache_value();
	}
	return cache;
}

void f() {
	Date today(16, 10, 1996);
	cout << today.string_rep() << endl;
	today.init(16, 10, 1996);
	Date my_birthday(29, 7, 1985);
	cout << my_birthday.string_rep() << endl;

	Date tomorrow = today;
	tomorrow.add_day(1);

	Date d1 = today;
	Date d2(today);
}

void f2(Date &d, const Date &cd) {
	int i = d.day();
	d.add_year(1);

	i = cd.day();
	// cd.add_year(1); // Error

	return;
}

int main() {
	f();
	return 0;
}

void Date::init(int dd, int mm, int yy) {
	d = dd;
	m = mm;
	y = yy;
}

Date::Date(int dd, int mm, int yy) :
	cache_valid(false),
	cache("")
{
	d = dd;
	m = mm;
	y = yy;
}

void Date::add_day(int n) {
	d += n;
}

void Date::add_year(int n) {
	y += n;
}
<commit_msg>Added code for 16.2.9.4<commit_after>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <string>

using namespace std;

struct cache {
	cache() : cache_valid(false), rep("") {}
	bool cache_valid;
	string rep;
};

class Date {
	int d, m, y;
	mutable bool cache_valid;
	mutable string rep;
	cache *c;

	void compute_cache_value() const;
	void compute_cache_value_indirection() const;

	public:
	void init(int dd, int mm, int yy);
	Date(int dd, int mm, int yy);
	Date();
	Date(const Date &);
	~Date();
	void add_year(int n);
	void add_month(int n);
	void add_day(int n);
	int day() const { return d; }
	int month() const { return m; }
	int year() const { return y; }
	string string_rep() const;
	string string_rep_indirection() const;

	const Date &operator=(const Date &);
};

Date::Date(const Date &d1) {
	d = d1.d;
	m = d1.m;
	y = d1.y;
	cache_valid = d1.cache_valid;
	rep = d1.rep;

	if ((c = new cache()) == NULL) {
		cerr << "Can not get space for cache!" << endl;
		exit(66);
	}
	c->cache_valid = d1.c->cache_valid;
	c->rep = d1.c->rep;
}

const Date & Date::operator=(const Date &d1) {
	d = d1.d;
	m = d1.m;
	y = d1.y;
	cache_valid = d1.cache_valid;
	rep = d1.rep;

	if (c == NULL) {
		if ((c = new cache()) == NULL) {
			cerr << "Can not get space for cache!" << endl;
			exit(66);
		}
	}
	c->cache_valid = d1.c->cache_valid;
	c->rep = d1.c->rep;

	return *this;
}

void Date::compute_cache_value() const {
	char tmp[20];
	sprintf(tmp, "%04d-%02d-%02d", y, m, d);
	string tmpstr(tmp);
	rep = tmpstr;
	cache_valid = true;
	return;
}

void Date::compute_cache_value_indirection() const {
	char tmp[20];
	sprintf(tmp, "%04d-%02d-%02d", y, m, d);
	string tmpstr(tmp);
	c->rep = tmpstr;
	c->cache_valid = true;
	return;
}

string Date::string_rep_indirection() const {
	if (!c->cache_valid) {
		compute_cache_value_indirection();
	}
	return c->rep;
}

string Date::string_rep() const {
	if (!cache_valid) {
		compute_cache_value();
	}
	return rep;
}

void f() {
	Date today(16, 10, 1996);
	cout << today.string_rep() << endl;
	cout << today.string_rep_indirection() << endl;
	today.init(16, 10, 1996);
	Date my_birthday(29, 7, 1985);
	cout << my_birthday.string_rep() << endl;
	cout << my_birthday.string_rep_indirection() << endl;

	Date tomorrow = today;
	tomorrow.add_day(1);
	cout << tomorrow.string_rep() << endl;
	cout << tomorrow.string_rep_indirection() << endl;

	Date d1 = today;
	Date d2(today);
}

void f2(Date &d, const Date &cd) {
	int i = d.day();
	d.add_year(1);

	i = cd.day();
	// cd.add_year(1); // Error

	return;
}

int main() {
	f();
	return 0;
}

void Date::init(int dd, int mm, int yy) {
	d = dd;
	m = mm;
	y = yy;
}

Date::Date(int dd, int mm, int yy) :
	cache_valid(false),
	rep("")
{
	d = dd;
	m = mm;
	y = yy;
	c = new cache();
}

Date::Date() :
	cache_valid(false),
	rep(""),
	c(NULL)
{
	d = 0;
	m = 0;
	y = 0;
}

void Date::add_day(int n) {
	d += n;
}

void Date::add_year(int n) {
	y += n;
}

Date::~Date() {
	delete c;
}
<|endoftext|>


<commit_before>#include "systems/inventorysystem.h"
#include "cmapserver.h"
#include "cmapclient.h"
#include "makevector.h"
#include "cli_equipitem.h"
#include "srv_equipitem.h"
#include "srv_setitem.h"
#include "itemdb.h"
#include "systems/luasystem.h"
#include "itemapi.h"

using namespace Systems;
using namespace RoseCommon;

InventorySystem::InventorySystem(SystemManager &manager) : System(manager) {
    manager.registerDispatcher(ePacketType::PAKCS_EQUIP_ITEM, &InventorySystem::processEquip);
}

void InventorySystem::update(EntityManager&, double) {}

uint8_t InventorySystem::findNextEmptySlot(Entity entity) {
    if (!entity || !entity.component<Inventory>())
        return 0;
    const auto &inventory = entity.component<Inventory>()->items_;
    for (size_t i = Inventory::MAX_EQUIP_ITEMS; i < Inventory::maxItems; ++i)
        if (!inventory[i].id_)
            return i;
    return 0;
}

bool InventorySystem::swapItems(Entity entity, uint8_t &a, uint8_t &b) {
    if (!entity || !entity.component<Inventory>())
        return false;
    if (a >= Inventory::maxItems || b >= Inventory::maxItems)
        return false;
    if (a == b || (!a && !b))
        return false;
    if (!a)
        a = findNextEmptySlot(entity);
    if (!b)
        b = findNextEmptySlot(entity);
    auto inventory = entity.component<Inventory>();
    std::swap(inventory->items_[a], inventory->items_[b]);
    return true;
}

void InventorySystem::processEquip(CMapClient& client, Entity entity, const RoseCommon::CliEquipItem &packet) {
    logger_->trace("InventorySystem::processEquip");
    if (packet.slotTo() < 1 || packet.slotTo() >= Inventory::MAX_EQUIP_ITEMS) {
        logger_->warn("When requesting to change equipped item, the destination wasn't good");
        return;
    }
    uint8_t to = packet.slotTo();
    uint8_t from = packet.slotFrom();
    if (!swapItems(entity, to, from)) {
        logger_->warn("There was an error while swapping items for client {}", getId(entity));
        return;
    }
    CMapServer::SendPacket(client, CMapServer::eSendType::EVERYONE,
            *makePacket<ePacketType::PAKWC_EQUIP_ITEM>(entity, packet.slotTo()));
    client.send(*makePacket<ePacketType::PAKWC_SET_ITEM>(entity, Core::make_vector(to, from)));
}

bool InventorySystem::addItem(Entity e, Item&& item) {
    uint8_t slot = findNextEmptySlot(e);
    if (!slot) return false;
    auto inv = e.component<Inventory>();
    inv->items_[slot] = std::move(item);

    auto client = getClient(e);

    client->send(*makePacket<ePacketType::PAKWC_SET_ITEM>(e, Core::make_vector(slot)));
    return true;
}

Item InventorySystem::removeItem(Entity entity, uint8_t slot) {
    if (!slot || slot >= Inventory::maxItems) return {};
    auto inv = entity.component<Inventory>();
    Item item{std::move(inv->items_[slot])};
    inv->items_[slot] = {};
    return item;
}

Item InventorySystem::buildItem(uint8_t type, uint16_t id, uint16_t life, bool isAppraised) {
    auto &itemDb = ItemDatabase::getInstance();
    auto def = itemDb.getItemDef(type, id);
    Item item(def);
    item.life_ = life;
    item.isAppraised_ = isAppraised;
    auto luaSystem = manager_.get<LuaSystem>();
    item.lua_ = luaSystem->loadScript<RoseCommon::ItemAPI>("");
    //const auto& env = item.lua_.getEnv();
    // TODO: get durability from lua env
    return item;
}
<commit_msg>Update inventorysystem.cpp<commit_after>#include "systems/inventorysystem.h"
#include "cmapserver.h"
#include "cmapclient.h"
#include "makevector.h"
#include "cli_equipitem.h"
#include "srv_equipitem.h"
#include "srv_setitem.h"
#include "itemdb.h"
#include "systems/luasystem.h"
#include "itemapi.h"

using namespace Systems;
using namespace RoseCommon;

InventorySystem::InventorySystem(SystemManager &manager) : System(manager) {
    manager.registerDispatcher(ePacketType::PAKCS_EQUIP_ITEM, &InventorySystem::processEquip);
}

void InventorySystem::update(EntityManager&, double) {}

uint8_t InventorySystem::findNextEmptySlot(Entity entity) {
    if (!entity || !entity.component<Inventory>())
        return 0;
    const auto &inventory = entity.component<Inventory>()->items_;
    for (size_t i = Inventory::MAX_EQUIP_ITEMS; i < Inventory::maxItems; ++i)
        if (!inventory[i].id_)
            return i;
    return 0;
}

bool InventorySystem::swapItems(Entity entity, uint8_t &a, uint8_t &b) {
    if (!entity || !entity.component<Inventory>())
        return false;
    if (a >= Inventory::maxItems || b >= Inventory::maxItems)
        return false;
    if (a == b || (!a && !b))
        return false;
    if (!a)
        a = findNextEmptySlot(entity);
    if (!b)
        b = findNextEmptySlot(entity);
    auto inventory = entity.component<Inventory>();
    std::swap(inventory->items_[a], inventory->items_[b]);
    return true;
}

void InventorySystem::processEquip(CMapClient& client, Entity entity, const RoseCommon::CliEquipItem &packet) {
    logger_->trace("InventorySystem::processEquip");
    if (packet.slotTo() < 1 || packet.slotTo() >= Inventory::MAX_EQUIP_ITEMS) {
        logger_->warn("When requesting to change equipped item, the destination wasn't good");
        return;
    }
    uint8_t to = packet.slotTo();
    uint8_t from = packet.slotFrom();
    if (!swapItems(entity, to, from)) {
        logger_->warn("There was an error while swapping items for client {}", getId(entity));
        return;
    }
    if (from < Inventory::MAX_EQUIP_ITEMS)
        entity.component<Inventory>()->items_[to].onUnequip(&entity);
    if (to < Inventory::MAX_EQUIP_ITEMS)
        entity.component<Inventory>()->items_[from].onEquip(&entity);
    CMapServer::SendPacket(client, CMapServer::eSendType::EVERYONE,
            *makePacket<ePacketType::PAKWC_EQUIP_ITEM>(entity, packet.slotTo()));
    client.send(*makePacket<ePacketType::PAKWC_SET_ITEM>(entity, Core::make_vector(to, from)));
}

bool InventorySystem::addItem(Entity e, Item&& item) {
    uint8_t slot = findNextEmptySlot(e);
    if (!slot) return false;
    item.lua_.onPickup();
    auto inv = e.component<Inventory>();
    inv->items_[slot] = std::move(item);

    auto client = getClient(e);

    client->send(*makePacket<ePacketType::PAKWC_SET_ITEM>(e, Core::make_vector(slot)));
    return true;
}

Item InventorySystem::removeItem(Entity entity, uint8_t slot) {
    if (!slot || slot >= Inventory::maxItems) return {};
    auto inv = entity.component<Inventory>();
    Item item{std::move(inv->items_[slot])};
    inv->items_[slot] = {};
    item.lua_.onDrop();
    return item;
}

Item InventorySystem::buildItem(uint8_t type, uint16_t id, uint16_t life, bool isAppraised) {
    auto &itemDb = ItemDatabase::getInstance();
    auto def = itemDb.getItemDef(type, id);
    Item item(def);
    item.life_ = life;
    item.isAppraised_ = isAppraised;
    auto luaSystem = manager_.get<LuaSystem>();
    item.lua_ = luaSystem->loadScript<RoseCommon::ItemAPI>("");
    //const auto& env = item.lua_.getEnv();
    // TODO: get durability from lua env
    item.lua_.onInit();
    return item;
}
<|endoftext|>


<commit_before>#include <assert.h>
#include <string.h>

#ifdef WIN32
	#include <winsock2.h>
	#include <ws2tcpip.h>
	#define SHUT_RDWR SD_BOTH
#else // WIN32
	#include <sys/socket.h>
	#include <netinet/in.h>
	#include <netinet/tcp.h>
	#include <netdb.h>
	#include <fcntl.h>
#endif // !WIN32

#include <unordered_map>

#include "connection.h"

#include "utils.h"

Connection::~Connection() {
	assert(disconnectFlags & DISCONNECT_COMPLETE);
	if (disconnectFlags & DISCONNECT_FROM_WRITE_THREAD)
		write_thread->join();
	else if (disconnectFlags & DISCONNECT_FROM_READ_THREAD)
		read_thread->join();
	else
		assert(!"DISCONNECT_COMPLETE set but not from either thread?");
	close(sock);
	delete read_thread;
	delete write_thread;
}

void Connection::do_send_bytes(const std::shared_ptr<std::vector<unsigned char> >& bytes, int send_mutex_token) {
	if (!send_mutex_token)
		send_mutex.lock();
	else
		ALWAYS_ASSERT(send_mutex_token == outside_send_mutex_token);

	if (initial_outbound_throttle && send_mutex_token)
		initial_outbound_bytes += bytes->size();

	if (total_waiting_size - (initial_outbound_throttle ? initial_outbound_bytes : 0) > 4000000) {
		if (!send_mutex_token)
			send_mutex.unlock();
		return disconnect_from_outside("total_waiting_size blew up :(", false);
	}

	outbound_primary_queue.push_back(bytes);
	total_waiting_size += bytes->size();
	cv.notify_all();
	if (!send_mutex_token)
		send_mutex.unlock();
}

void Connection::maybe_send_bytes(const std::shared_ptr<std::vector<unsigned char> >& bytes, int send_mutex_token) {
	if (!send_mutex_token) {
		if (!send_mutex.try_lock())
			return;
	} else
		ALWAYS_ASSERT(send_mutex_token == outside_send_mutex_token);

	if (total_waiting_size - (initial_outbound_throttle ? initial_outbound_bytes : 0) > 4000000) {
		if (!send_mutex_token)
			send_mutex.unlock();
		return disconnect_from_outside("total_waiting_size blew up :(", false);
	}

	outbound_secondary_queue.push_back(bytes);
	total_waiting_size += bytes->size();
	cv.notify_all();
	if (!send_mutex_token)
		send_mutex.unlock();
}

void Connection::disconnect_from_outside(const char* reason, bool push_send) {
	if (disconnectFlags.fetch_or(DISCONNECT_PRINT_AND_CLOSE) & DISCONNECT_PRINT_AND_CLOSE)
		return;

	printf("%s Disconnect: %s (%s)\n", host.c_str(), reason, strerror(errno));
	shutdown(sock, SHUT_RDWR);

	if (push_send)
		do_send_bytes(std::make_shared<std::vector<unsigned char> >(1));
}

void Connection::disconnect(const char* reason) {
	if (disconnectFlags.fetch_or(DISCONNECT_STARTED) & DISCONNECT_STARTED)
		return;

	if (!(disconnectFlags.fetch_or(DISCONNECT_PRINT_AND_CLOSE) & DISCONNECT_PRINT_AND_CLOSE)) {
		printf("%s Disconnect: %s (%s)\n", host.c_str(), reason, strerror(errno));
		shutdown(sock, SHUT_RDWR);
	}

	if (std::this_thread::get_id() != read_thread->get_id()) {
		disconnectFlags |= DISCONNECT_FROM_WRITE_THREAD;
		read_thread->join();
	} else {
		disconnectFlags |= DISCONNECT_FROM_READ_THREAD;
		{
			/* Wake up the write thread */
			std::lock_guard<std::mutex> lock(send_mutex);
			outbound_secondary_queue.push_back(std::make_shared<std::vector<unsigned char> >(1));
			cv.notify_all();
		}
		write_thread->join();
	}

	outbound_secondary_queue.clear();
	outbound_primary_queue.clear();
	inbound_queue.clear();

	disconnectFlags |= DISCONNECT_COMPLETE;

	if (on_disconnect)
		std::thread(on_disconnect).detach();
}

class GlobalNetProcess {
public:
	std::mutex fd_map_mutex;
	std::unordered_map<int, Connection*> fd_map;

	static void do_net_process(GlobalNetProcess* me) {
		fd_set fd_set_read, fd_set_write;
		struct timeval timeout;

		while (true) {
			timeout.tv_sec = 1;
			timeout.tv_usec = 0;

			FD_ZERO(&fd_set_read); FD_ZERO(&fd_set_write);
			int max = 0;
			{
				std::lock_guard<std::mutex> lock(me->fd_map_mutex);
				for (const auto& e : me->fd_map) {
					ALWAYS_ASSERT(e.first < FD_SETSIZE);
					FD_SET(e.first, &fd_set_read);
					//FD_SET(e.first, &fd_set_write);
					max = std::max(e.first, max);
				}
			}

			assert(select(max + 1, &fd_set_read, &fd_set_write, NULL, &timeout) >= 0);

			std::list<int> remove_list;
			std::lock_guard<std::mutex> lock(me->fd_map_mutex);
			for (const auto& e : me->fd_map) {
				if (FD_ISSET(e.first, &fd_set_read)) {
					unsigned char buf[4096];
					ssize_t count = recv(e.second->sock, (char*)buf, 4096, 0);

					std::lock_guard<std::mutex> lock(e.second->read_mutex);
					if (count <= 0) {
						e.second->inbound_queue.emplace_back((std::nullptr_t)NULL);
						remove_list.push_back(e.first);
					} else
						e.second->inbound_queue.emplace_back(new std::vector<unsigned char>(buf, buf + count));
					e.second->read_cv.notify_all();
				}
				if (FD_ISSET(e.first, &fd_set_write)) {
					//TODO: Move write here
				}
			}

			for (const int fd : remove_list)
				me->fd_map.erase(fd);
		}
	}

	GlobalNetProcess() {
		std::thread(do_net_process, this).detach();
	}
};
static GlobalNetProcess processor;

void Connection::do_setup_and_read(Connection* me) {
	#ifdef WIN32
		unsigned long nonblocking = 0;
		ioctlsocket(me->sock, FIONBIO, &nonblocking);
	#else
		fcntl(me->sock, F_SETFL, fcntl(me->sock, F_GETFL) & ~O_NONBLOCK);
	#endif

	#ifdef X86_BSD
		int nosigpipe = 1;
		setsockopt(me->sock, SOL_SOCKET, SO_NOSIGPIPE, (void *)&nosigpipe, sizeof(int));
	#endif

	int nodelay = 1;
	setsockopt(me->sock, IPPROTO_TCP, TCP_NODELAY, (char*)&nodelay, sizeof(nodelay));

	if (errno)
		return me->disconnect("error during connect");

	{
		std::lock_guard<std::mutex>(processor.fd_map_mutex);
		processor.fd_map[me->sock] = me;
	}

	me->net_process([&](const char* reason) { me->disconnect(reason); });
}

ssize_t Connection::read_all(char *buf, size_t nbyte) {
	size_t total = 0;
	while (total < nbyte) {
		std::unique_lock<std::mutex> lock(read_mutex);
		while (!inbound_queue.size())
			read_cv.wait(lock);

		if (!inbound_queue.front())
			return -1;

		size_t readamt = std::min(nbyte - total, inbound_queue.front()->size() - readpos);
		memcpy(buf + total, &(*inbound_queue.front())[readpos], readamt);
		if (readpos + readamt == inbound_queue.front()->size()) {
			readpos = 0;
			inbound_queue.pop_front();
		} else
			readpos += readamt;
		total += readamt;
	}
	assert(total == nbyte);
	return nbyte;
}

void Connection::do_write(Connection* me) {
	me->net_write();
}

void Connection::net_write() {
	//TODO: Have a single global thread doing this
	while (true) {
		std::shared_ptr<std::vector<unsigned char> > msg;
		bool sleepFirst = false;
		{
			std::unique_lock<std::mutex> write_lock(send_mutex);
			while (!outbound_secondary_queue.size() && !outbound_primary_queue.size())
				cv.wait(write_lock);

			if (disconnectFlags)
				return disconnect("disconnect started elsewhere");

			if (outbound_primary_queue.size()) {
				msg = outbound_primary_queue.front();
				outbound_primary_queue.pop_front();
			} else {
				msg = outbound_secondary_queue.front();
				outbound_secondary_queue.pop_front();
			}

			total_waiting_size -= msg->size();
			if (!total_waiting_size)
				initial_outbound_throttle = false;
			else if (initial_outbound_throttle)
				sleepFirst = true;
		}
		if (sleepFirst)
			std::this_thread::sleep_for(std::chrono::milliseconds(20)); // Limit outbound to avg 5Mbps worst-case
		if (send_all(sock, (char*)&(*msg)[0], msg->size()) != int64_t(msg->size()))
			return disconnect("failed to send msg");
	}
}

void OutboundPersistentConnection::reconnect(std::string disconnectReason) {
	OutboundConnection* old = (OutboundConnection*) connection.fetch_and(0);
	if (old)
		old->disconnect_from_outside(disconnectReason.c_str());

	on_disconnect();

	std::this_thread::sleep_for(std::chrono::seconds(1));
	while (old && !(old->getDisconnectFlags() & DISCONNECT_COMPLETE)) {
		printf("Disconnect of outbound connection still not complete (status is %d)\n", old->getDisconnectFlags());
		std::this_thread::sleep_for(std::chrono::seconds(1));
	}

	std::thread(do_connect, this).detach();
	delete old;
}

void OutboundPersistentConnection::do_connect(OutboundPersistentConnection* me) {
	int sock = socket(AF_INET6, SOCK_STREAM, 0);
	if (sock <= 0)
		return me->reconnect("unable to create socket");

	sockaddr_in6 addr;
	if (!lookup_address(me->serverHost.c_str(), &addr)) {
		close(sock);
		return me->reconnect("unable to lookup host");
	}

	int v6only = 0;
	setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY, (char *)&v6only, sizeof(v6only));

	addr.sin6_port = htons(me->serverPort);
	if (connect(sock, (struct sockaddr*)&addr, sizeof(addr))) {
		close(sock);
		return me->reconnect("failed to connect()");
	}

	OutboundConnection* new_conn = new OutboundConnection(sock, me);
#ifndef NDEBUG
	unsigned long old_val =
#endif
		me->connection.exchange((unsigned long)new_conn);
	assert(old_val == 0);
	new_conn->construction_done();
}
<commit_msg>Add pipe-wakeup to event thread<commit_after>#include <assert.h>
#include <string.h>

#ifdef WIN32
	#include <winsock2.h>
	#include <ws2tcpip.h>
	#define SHUT_RDWR SD_BOTH
#else // WIN32
	#include <sys/socket.h>
	#include <netinet/in.h>
	#include <netinet/tcp.h>
	#include <netdb.h>
	#include <fcntl.h>
#endif // !WIN32

#include <unordered_map>

#include "connection.h"

#include "utils.h"

Connection::~Connection() {
	assert(disconnectFlags & DISCONNECT_COMPLETE);
	if (disconnectFlags & DISCONNECT_FROM_WRITE_THREAD)
		write_thread->join();
	else if (disconnectFlags & DISCONNECT_FROM_READ_THREAD)
		read_thread->join();
	else
		assert(!"DISCONNECT_COMPLETE set but not from either thread?");
	close(sock);
	delete read_thread;
	delete write_thread;
}

void Connection::do_send_bytes(const std::shared_ptr<std::vector<unsigned char> >& bytes, int send_mutex_token) {
	if (!send_mutex_token)
		send_mutex.lock();
	else
		ALWAYS_ASSERT(send_mutex_token == outside_send_mutex_token);

	if (initial_outbound_throttle && send_mutex_token)
		initial_outbound_bytes += bytes->size();

	if (total_waiting_size - (initial_outbound_throttle ? initial_outbound_bytes : 0) > 4000000) {
		if (!send_mutex_token)
			send_mutex.unlock();
		return disconnect_from_outside("total_waiting_size blew up :(", false);
	}

	outbound_primary_queue.push_back(bytes);
	total_waiting_size += bytes->size();
	cv.notify_all();
	if (!send_mutex_token)
		send_mutex.unlock();
}

void Connection::maybe_send_bytes(const std::shared_ptr<std::vector<unsigned char> >& bytes, int send_mutex_token) {
	if (!send_mutex_token) {
		if (!send_mutex.try_lock())
			return;
	} else
		ALWAYS_ASSERT(send_mutex_token == outside_send_mutex_token);

	if (total_waiting_size - (initial_outbound_throttle ? initial_outbound_bytes : 0) > 4000000) {
		if (!send_mutex_token)
			send_mutex.unlock();
		return disconnect_from_outside("total_waiting_size blew up :(", false);
	}

	outbound_secondary_queue.push_back(bytes);
	total_waiting_size += bytes->size();
	cv.notify_all();
	if (!send_mutex_token)
		send_mutex.unlock();
}

void Connection::disconnect_from_outside(const char* reason, bool push_send) {
	if (disconnectFlags.fetch_or(DISCONNECT_PRINT_AND_CLOSE) & DISCONNECT_PRINT_AND_CLOSE)
		return;

	printf("%s Disconnect: %s (%s)\n", host.c_str(), reason, strerror(errno));
	shutdown(sock, SHUT_RDWR);

	if (push_send)
		do_send_bytes(std::make_shared<std::vector<unsigned char> >(1));
}

void Connection::disconnect(const char* reason) {
	if (disconnectFlags.fetch_or(DISCONNECT_STARTED) & DISCONNECT_STARTED)
		return;

	if (!(disconnectFlags.fetch_or(DISCONNECT_PRINT_AND_CLOSE) & DISCONNECT_PRINT_AND_CLOSE)) {
		printf("%s Disconnect: %s (%s)\n", host.c_str(), reason, strerror(errno));
		shutdown(sock, SHUT_RDWR);
	}

	if (std::this_thread::get_id() != read_thread->get_id()) {
		disconnectFlags |= DISCONNECT_FROM_WRITE_THREAD;
		read_thread->join();
	} else {
		disconnectFlags |= DISCONNECT_FROM_READ_THREAD;
		{
			/* Wake up the write thread */
			std::lock_guard<std::mutex> lock(send_mutex);
			outbound_secondary_queue.push_back(std::make_shared<std::vector<unsigned char> >(1));
			cv.notify_all();
		}
		write_thread->join();
	}

	outbound_secondary_queue.clear();
	outbound_primary_queue.clear();
	inbound_queue.clear();

	disconnectFlags |= DISCONNECT_COMPLETE;

	if (on_disconnect)
		std::thread(on_disconnect).detach();
}

class GlobalNetProcess {
public:
	std::mutex fd_map_mutex;
	std::unordered_map<int, Connection*> fd_map;
#ifndef WIN32
	int pipe_write;
#endif

	static void do_net_process(GlobalNetProcess* me) {
		fd_set fd_set_read, fd_set_write;
		struct timeval timeout;

#ifndef WIN32
		int pipefd[2];
		assert(!pipe(pipefd));
		fcntl(pipefd[1], F_SETFL, fcntl(pipefd[1], F_GETFL) | O_NONBLOCK);
		fcntl(pipefd[0], F_SETFL, fcntl(pipefd[0], F_GETFL) | O_NONBLOCK);
		me->pipe_write = pipefd[1];
#endif

		while (true) {
#ifndef WIN32
			timeout.tv_sec = 60;
			timeout.tv_usec = 0;
#else
			timeout.tv_sec = 0;
			timeout.tv_usec = 1000;
#endif

			FD_ZERO(&fd_set_read); FD_ZERO(&fd_set_write);
#ifndef WIN32
			int max = pipefd[0];
			FD_SET(pipefd[0], &fd_set_read);
#else
			int max = 0;
#endif
			{
				std::lock_guard<std::mutex> lock(me->fd_map_mutex);
				for (const auto& e : me->fd_map) {
					ALWAYS_ASSERT(e.first < FD_SETSIZE);
					FD_SET(e.first, &fd_set_read);
					//FD_SET(e.first, &fd_set_write);
					max = std::max(e.first, max);
				}
			}

			assert(select(max + 1, &fd_set_read, &fd_set_write, NULL, &timeout) >= 0);

			unsigned char buf[4096];
			{
				std::list<int> remove_list;
				std::lock_guard<std::mutex> lock(me->fd_map_mutex);
				for (const auto& e : me->fd_map) {
					if (FD_ISSET(e.first, &fd_set_read)) {
						ssize_t count = recv(e.second->sock, (char*)buf, 4096, 0);

						std::lock_guard<std::mutex> lock(e.second->read_mutex);
						if (count <= 0) {
							e.second->inbound_queue.emplace_back((std::nullptr_t)NULL);
							remove_list.push_back(e.first);
						} else
							e.second->inbound_queue.emplace_back(new std::vector<unsigned char>(buf, buf + count));
						e.second->read_cv.notify_all();
					}
					if (FD_ISSET(e.first, &fd_set_write)) {
						//TODO: Move write here
					}
				}

				for (const int fd : remove_list)
					me->fd_map.erase(fd);
			}
#ifndef WIN32
			if (FD_ISSET(pipefd[0], &fd_set_read))
				while (read(pipefd[0], buf, 4096) > 0);
#endif
		}
	}

	GlobalNetProcess() {
		std::thread(do_net_process, this).detach();
	}
};
static GlobalNetProcess processor;

void Connection::do_setup_and_read(Connection* me) {
	#ifdef WIN32
		unsigned long nonblocking = 0;
		ioctlsocket(me->sock, FIONBIO, &nonblocking);
	#else
		fcntl(me->sock, F_SETFL, fcntl(me->sock, F_GETFL) & ~O_NONBLOCK);
	#endif

	#ifdef X86_BSD
		int nosigpipe = 1;
		setsockopt(me->sock, SOL_SOCKET, SO_NOSIGPIPE, (void *)&nosigpipe, sizeof(int));
	#endif

	int nodelay = 1;
	setsockopt(me->sock, IPPROTO_TCP, TCP_NODELAY, (char*)&nodelay, sizeof(nodelay));

	if (errno)
		return me->disconnect("error during connect");

	{
		std::lock_guard<std::mutex>(processor.fd_map_mutex);
		processor.fd_map[me->sock] = me;
#ifndef WIN32
		write(processor.pipe_write, "1", 1);
#endif
	}

	me->net_process([&](const char* reason) { me->disconnect(reason); });
}

ssize_t Connection::read_all(char *buf, size_t nbyte) {
	size_t total = 0;
	while (total < nbyte) {
		std::unique_lock<std::mutex> lock(read_mutex);
		while (!inbound_queue.size())
			read_cv.wait(lock);

		if (!inbound_queue.front())
			return -1;

		size_t readamt = std::min(nbyte - total, inbound_queue.front()->size() - readpos);
		memcpy(buf + total, &(*inbound_queue.front())[readpos], readamt);
		if (readpos + readamt == inbound_queue.front()->size()) {
			readpos = 0;
			inbound_queue.pop_front();
		} else
			readpos += readamt;
		total += readamt;
	}
	assert(total == nbyte);
	return nbyte;
}

void Connection::do_write(Connection* me) {
	me->net_write();
}

void Connection::net_write() {
	//TODO: Have a single global thread doing this
	while (true) {
		std::shared_ptr<std::vector<unsigned char> > msg;
		bool sleepFirst = false;
		{
			std::unique_lock<std::mutex> write_lock(send_mutex);
			while (!outbound_secondary_queue.size() && !outbound_primary_queue.size())
				cv.wait(write_lock);

			if (disconnectFlags)
				return disconnect("disconnect started elsewhere");

			if (outbound_primary_queue.size()) {
				msg = outbound_primary_queue.front();
				outbound_primary_queue.pop_front();
			} else {
				msg = outbound_secondary_queue.front();
				outbound_secondary_queue.pop_front();
			}

			total_waiting_size -= msg->size();
			if (!total_waiting_size)
				initial_outbound_throttle = false;
			else if (initial_outbound_throttle)
				sleepFirst = true;
		}
		if (sleepFirst)
			std::this_thread::sleep_for(std::chrono::milliseconds(20)); // Limit outbound to avg 5Mbps worst-case
		if (send_all(sock, (char*)&(*msg)[0], msg->size()) != int64_t(msg->size()))
			return disconnect("failed to send msg");
	}
}

void OutboundPersistentConnection::reconnect(std::string disconnectReason) {
	OutboundConnection* old = (OutboundConnection*) connection.fetch_and(0);
	if (old)
		old->disconnect_from_outside(disconnectReason.c_str());

	on_disconnect();

	std::this_thread::sleep_for(std::chrono::seconds(1));
	while (old && !(old->getDisconnectFlags() & DISCONNECT_COMPLETE)) {
		printf("Disconnect of outbound connection still not complete (status is %d)\n", old->getDisconnectFlags());
		std::this_thread::sleep_for(std::chrono::seconds(1));
	}

	std::thread(do_connect, this).detach();
	delete old;
}

void OutboundPersistentConnection::do_connect(OutboundPersistentConnection* me) {
	int sock = socket(AF_INET6, SOCK_STREAM, 0);
	if (sock <= 0)
		return me->reconnect("unable to create socket");

	sockaddr_in6 addr;
	if (!lookup_address(me->serverHost.c_str(), &addr)) {
		close(sock);
		return me->reconnect("unable to lookup host");
	}

	int v6only = 0;
	setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY, (char *)&v6only, sizeof(v6only));

	addr.sin6_port = htons(me->serverPort);
	if (connect(sock, (struct sockaddr*)&addr, sizeof(addr))) {
		close(sock);
		return me->reconnect("failed to connect()");
	}

	OutboundConnection* new_conn = new OutboundConnection(sock, me);
#ifndef NDEBUG
	unsigned long old_val =
#endif
		me->connection.exchange((unsigned long)new_conn);
	assert(old_val == 0);
	new_conn->construction_done();
}
<|endoftext|>


<commit_before>#include "Mocks.hpp"

using namespace blackhole;

TEST(file_t, Class) {
    sink::file_t<> sink("test.log");
    UNUSED(sink);
}

TEST(file_t, WritesToTheFile) {
    sink::file_t<NiceMock<mock::backend_t>> sink("test.log");
    EXPECT_CALL(sink.backend(), write(std::string("formatted message"))).
            Times(1);
    sink.consume("formatted message");
}

TEST(file_t, OpensFileIfItClosedWhenWriting) {
    sink::file_t<mock::backend_t> sink("test.log");
    EXPECT_CALL(sink.backend(), opened()).
            Times(1).
            WillOnce(Return(false));
    EXPECT_CALL(sink.backend(), open()).
            Times(1).
            WillOnce(Return(true));
    EXPECT_CALL(sink.backend(), write(_)).
            Times(1);
    sink.consume("message");
}

TEST(file_t, ThrowsExceptionIfFileCannotBeOpened) {
    sink::file_t<NiceMock<mock::backend_t>> sink("test.log");
    EXPECT_CALL(sink.backend(), opened())
            .Times(1)
            .WillOnce(Return(false));
    EXPECT_CALL(sink.backend(), open())
            .Times(1)
            .WillOnce(Return(false));
    EXPECT_CALL(sink.backend(), write(_))
            .Times(0);
    EXPECT_THROW(sink.consume("message"), error_t);
}
<commit_msg>Ambiguity resolving.<commit_after>#include "Mocks.hpp"

using namespace blackhole;

TEST(file_t, Class) {
    sink::file_t<> sink("test.log");
    UNUSED(sink);
}

TEST(file_t, WritesToTheFile) {
    sink::file_t<NiceMock<mock::backend_t>> sink("test.log");
    EXPECT_CALL(sink.backend(), write(std::string("formatted message"))).
            Times(1);
    sink.consume("formatted message");
}

TEST(file_t, OpensFileIfItClosedWhenWriting) {
    sink::file_t<mock::backend_t> sink("test.log");
    EXPECT_CALL(sink.backend(), opened()).
            Times(1).
            WillOnce(Return(false));
    EXPECT_CALL(sink.backend(), open()).
            Times(1).
            WillOnce(Return(true));
    EXPECT_CALL(sink.backend(), write(_)).
            Times(1);
    sink.consume("message");
}

TEST(file_t, ThrowsExceptionIfFileCannotBeOpened) {
    sink::file_t<NiceMock<mock::backend_t>> sink("test.log");
    EXPECT_CALL(sink.backend(), opened())
            .Times(1)
            .WillOnce(Return(false));
    EXPECT_CALL(sink.backend(), open())
            .Times(1)
            .WillOnce(Return(false));
    EXPECT_CALL(sink.backend(), write(_))
            .Times(0);
    EXPECT_THROW(sink.consume("message"), blackhole::error_t);
}
<|endoftext|>


<commit_before>/* Copyright (c) 2013-2015, AOYAMA Kazuharu
 * All rights reserved.
 *
 * This software may be used and distributed according to the terms of
 * the New BSD License, which is incorporated herein by reference.
 */

#include <QList>
#include <QMutexLocker>
#include "tfileaiowriter.h"
#include "tfcore_unix.h"

const int MAX_NUM_BUFFERING_DATA = 10000;

class TFileAioWriterData
{
public:
    mutable QMutex mutex;
    QString fileName;
    int fileDescriptor;
    QList<struct aiocb *> syncBuffer;

    TFileAioWriterData() : mutex(QMutex::Recursive), fileName(), fileDescriptor(0), syncBuffer() { }
    void clearSyncBuffer();
};

void TFileAioWriterData::clearSyncBuffer()
{
    if (!syncBuffer.isEmpty()) {
        for (QListIterator<struct aiocb *> it(syncBuffer); it.hasNext(); ) {
            struct aiocb *cb = it.next();
            delete (char *)cb->aio_buf;
            delete cb;
        }
        syncBuffer.clear();
    }
}

/*!
  Constructor.
 */
TFileAioWriter::TFileAioWriter(const QString &name)
    : d(new TFileAioWriterData)
{
    d->fileName = name;
}

TFileAioWriter::~TFileAioWriter()
{
    close();
    delete d;
}

bool TFileAioWriter::open()
{
    QMutexLocker locker(&d->mutex);

    if (d->fileDescriptor <= 0) {
        if (d->fileName.isEmpty())
            return false;

        d->fileDescriptor = ::open(qPrintable(d->fileName), (O_CREAT | O_WRONLY | O_APPEND | O_CLOEXEC), 0666);
        if (d->fileDescriptor < 0) {
            //fprintf(stderr, "file open failed: %s\n", qPrintable(d->fileName));
        }
    }

    return (d->fileDescriptor > 0);
}

void TFileAioWriter::close()
{
    QMutexLocker locker(&d->mutex);

    flush();

    if (d->fileDescriptor > 0) {
        tf_close(d->fileDescriptor);
    }
    d->fileDescriptor = 0;
}

bool TFileAioWriter::isOpen() const
{
    return (d->fileDescriptor > 0);
}

int TFileAioWriter::write(const char *data, int length)
{
    QMutexLocker locker(&d->mutex);

    if (!isOpen()) {
        return -1;
    }

    if (length <= 0)
        return -1;

    // check whether last writing is finished
    if (d->syncBuffer.count() > 0) {
        struct aiocb *lastcb = d->syncBuffer.last();
        if (aio_error(lastcb) != EINPROGRESS) {
            d->clearSyncBuffer();
        } else {
            if (d->syncBuffer.count() > MAX_NUM_BUFFERING_DATA) {
                flush();
            }
        }
    }

    struct aiocb *cb = new struct aiocb;
    memset(cb, 0, sizeof(struct aiocb));

    cb->aio_fildes = d->fileDescriptor;
    cb->aio_nbytes = length;
    cb->aio_buf = new char[length];
    memcpy((void *)cb->aio_buf, data, length);

    int ret = tf_aio_write(cb);
    int err = errno;

    if (ret < 0) {
        //fprintf(stderr, "aio_write error fd:%d (pid:%d tid:%d) ret:%d errno:%d\n", d->fileDescriptor, getpid(), gettid(), ret, err);
        //fprintf(stderr, "aio_write str: %s\n", data);
        delete (char *)cb->aio_buf;
        delete cb;

        if (err != EAGAIN) {
            close();
        } else {
#ifdef Q_OS_DARWIN
            // try sync-write
            return (tf_write(d->fileDescriptor, data, length) > 0) ? 0 : -1;
#endif
        }
        return ret;
    }

    d->syncBuffer << cb;
    return 0;
}

void TFileAioWriter::flush()
{
    QMutexLocker locker(&d->mutex);

    if (d->syncBuffer.count() > 0) {
        struct aiocb *lastcb = d->syncBuffer.last();
        while (aio_error(lastcb) == EINPROGRESS) { }
        d->clearSyncBuffer();
    }
}

void TFileAioWriter::setFileName(const QString &name)
{
    QMutexLocker locker(&d->mutex);

    if (isOpen()) {
        close();
    }
    d->fileName = name;
}

QString TFileAioWriter::fileName() const
{
    QMutexLocker locker(&d->mutex);
    return d->fileName;
}
<commit_msg>Add valgrind error on mixing new[] and delete<commit_after>/* Copyright (c) 2013-2015, AOYAMA Kazuharu
 * All rights reserved.
 *
 * This software may be used and distributed according to the terms of
 * the New BSD License, which is incorporated herein by reference.
 */

#include <QList>
#include <QMutexLocker>
#include "tfileaiowriter.h"
#include "tfcore_unix.h"

const int MAX_NUM_BUFFERING_DATA = 10000;

class TFileAioWriterData
{
public:
    mutable QMutex mutex;
    QString fileName;
    int fileDescriptor;
    QList<struct aiocb *> syncBuffer;

    TFileAioWriterData() : mutex(QMutex::Recursive), fileName(), fileDescriptor(0), syncBuffer() { }
    void clearSyncBuffer();
};

void TFileAioWriterData::clearSyncBuffer()
{
    if (!syncBuffer.isEmpty()) {
        for (QListIterator<struct aiocb *> it(syncBuffer); it.hasNext(); ) {
            struct aiocb *cb = it.next();
            delete [] (char *) cb->aio_buf;
            delete cb;
        }
        syncBuffer.clear();
    }
}

/*!
  Constructor.
 */
TFileAioWriter::TFileAioWriter(const QString &name)
    : d(new TFileAioWriterData)
{
    d->fileName = name;
}

TFileAioWriter::~TFileAioWriter()
{
    close();
    delete d;
}

bool TFileAioWriter::open()
{
    QMutexLocker locker(&d->mutex);

    if (d->fileDescriptor <= 0) {
        if (d->fileName.isEmpty())
            return false;

        d->fileDescriptor = ::open(qPrintable(d->fileName), (O_CREAT | O_WRONLY | O_APPEND | O_CLOEXEC), 0666);
        if (d->fileDescriptor < 0) {
            //fprintf(stderr, "file open failed: %s\n", qPrintable(d->fileName));
        }
    }

    return (d->fileDescriptor > 0);
}

void TFileAioWriter::close()
{
    QMutexLocker locker(&d->mutex);

    flush();

    if (d->fileDescriptor > 0) {
        tf_close(d->fileDescriptor);
    }
    d->fileDescriptor = 0;
}

bool TFileAioWriter::isOpen() const
{
    return (d->fileDescriptor > 0);
}

int TFileAioWriter::write(const char *data, int length)
{
    QMutexLocker locker(&d->mutex);

    if (!isOpen()) {
        return -1;
    }

    if (length <= 0)
        return -1;

    // check whether last writing is finished
    if (d->syncBuffer.count() > 0) {
        struct aiocb *lastcb = d->syncBuffer.last();
        if (aio_error(lastcb) != EINPROGRESS) {
            d->clearSyncBuffer();
        } else {
            if (d->syncBuffer.count() > MAX_NUM_BUFFERING_DATA) {
                flush();
            }
        }
    }

    struct aiocb *cb = new struct aiocb;
    memset(cb, 0, sizeof(struct aiocb));

    cb->aio_fildes = d->fileDescriptor;
    cb->aio_nbytes = length;
    cb->aio_buf = new char[length];
    memcpy((void *)cb->aio_buf, data, length);

    int ret = tf_aio_write(cb);
    int err = errno;

    if (ret < 0) {
        //fprintf(stderr, "aio_write error fd:%d (pid:%d tid:%d) ret:%d errno:%d\n", d->fileDescriptor, getpid(), gettid(), ret, err);
        //fprintf(stderr, "aio_write str: %s\n", data);
        delete (char *)cb->aio_buf;
        delete cb;

        if (err != EAGAIN) {
            close();
        } else {
#ifdef Q_OS_DARWIN
            // try sync-write
            return (tf_write(d->fileDescriptor, data, length) > 0) ? 0 : -1;
#endif
        }
        return ret;
    }

    d->syncBuffer << cb;
    return 0;
}

void TFileAioWriter::flush()
{
    QMutexLocker locker(&d->mutex);

    if (d->syncBuffer.count() > 0) {
        struct aiocb *lastcb = d->syncBuffer.last();
        while (aio_error(lastcb) == EINPROGRESS) { }
        d->clearSyncBuffer();
    }
}

void TFileAioWriter::setFileName(const QString &name)
{
    QMutexLocker locker(&d->mutex);

    if (isOpen()) {
        close();
    }
    d->fileName = name;
}

QString TFileAioWriter::fileName() const
{
    QMutexLocker locker(&d->mutex);
    return d->fileName;
}
<|endoftext|>


<commit_before>//
//  Copyright (C) 2015  Marcus Comstedt <marcus@mc.pp.se>
//
//  Permission to use, copy, modify, and/or distribute this software for any
//  purpose with or without fee is hereby granted, provided that the above
//  copyright notice and this permission notice appear in all copies.
//
//  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
//  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
//  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
//  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
//  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
//  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
//  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
//

#include <fstream>
#include <iostream>
#include <cstdint>
#include <memory>

#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define log(...) fprintf(stderr, __VA_ARGS__);
#define error(...) do { fprintf(stderr, "%s: ", program_short_name); fprintf(stderr, __VA_ARGS__); exit(EXIT_FAILURE); } while (0)

static char *program_short_name;

static const int NUM_IMAGES = 4;
static const int HEADER_SIZE = 32;

static void align_offset(uint32_t &offset, int bits)
{
    uint32_t mask = (1 << bits) - 1;
    if (offset & mask)
        offset = (offset | mask) + 1;
}

static void write_byte(std::ostream &ofs, uint32_t &file_offset, uint8_t byte)
{
    ofs << byte;
    file_offset++;
}

static void write_bytes(std::ostream &ofs, uint32_t &file_offset,
                        const uint8_t *buf, size_t n)
{
    if (n > 0) {
        ofs.write(reinterpret_cast<const char*>(buf), n);
        file_offset += n;
    }
}

static void write_file(std::ostream &ofs, uint32_t &file_offset,
                       std::istream &ifs, const char *filename)
{
    const size_t bufsize = 8192;
    uint8_t *buffer = new uint8_t[bufsize];

    while(!ifs.eof()) {
        ifs.read(reinterpret_cast<char *>(buffer), bufsize);
        if (ifs.bad())
            error("can't read input image `%s': %s\n", filename, strerror(errno));
        write_bytes(ofs, file_offset, buffer, ifs.gcount());
    }

    delete[] buffer;
}

static void pad_to(std::ostream &ofs, uint32_t &file_offset, uint32_t target)
{
    if (target < file_offset)
        error("Trying to pad backwards!\n");
    while(file_offset < target)
        write_byte(ofs, file_offset, 0xff);
}

class Image {
    const char *filename;
    std::ifstream ifs;
    uint32_t offs;

public:
    Image(const char *filename);
    size_t size();
    void write(std::ostream &ofs, uint32_t &file_offset);
    void place(uint32_t o) { offs = o; }
    uint32_t offset() const { return offs; }
};

Image::Image(const char *filename) : filename(filename), ifs(filename, std::ifstream::binary)
{
    if (ifs.fail())
        error("can't open input image `%s': %s\n", filename, strerror(errno));
}

size_t Image::size()
{
    ifs.seekg (0, ifs.end);
    if (ifs.fail())
        error("can't seek on input image `%s': %s\n", filename, strerror(errno));
    size_t length = ifs.tellg();
    ifs.seekg (0, ifs.beg);
    if (ifs.fail())
        error("can't seek on input image `%s': %s\n", filename, strerror(errno));

    if (length == 0)
        error("input image `%s' doesn't contain any data\n", filename);
    return length;
}

void Image::write(std::ostream &ofs, uint32_t &file_offset)
{
    write_file(ofs, file_offset, ifs, filename);
}

static void write_header(std::ostream &ofs, uint32_t &file_offset,
                         Image const *image, bool coldboot)
{
    // Preamble
    write_byte(ofs, file_offset, 0x7e);
    write_byte(ofs, file_offset, 0xaa);
    write_byte(ofs, file_offset, 0x99);
    write_byte(ofs, file_offset, 0x7e);

    // Boot mode
    write_byte(ofs, file_offset, 0x92);
    write_byte(ofs, file_offset, 0x00);
    write_byte(ofs, file_offset, coldboot ? 0x10 : 0x00);

    // Boot address
    write_byte(ofs, file_offset, 0x44);
    write_byte(ofs, file_offset, 0x03);
    write_byte(ofs, file_offset, (image->offset() >> 16) & 0xff);
    write_byte(ofs, file_offset, (image->offset() >> 8) & 0xff);
    write_byte(ofs, file_offset, image->offset() & 0xff);

    // Bank offset
    write_byte(ofs, file_offset, 0x82);
    write_byte(ofs, file_offset, 0x00);
    write_byte(ofs, file_offset, 0x00);

    // Reboot
    write_byte(ofs, file_offset, 0x01);
    write_byte(ofs, file_offset, 0x08);

    // Zero out any unused bytes
    while (file_offset & (HEADER_SIZE - 1))
        write_byte(ofs, file_offset, 0x00);
}

void usage()
{
    log("\n");
    log("Usage: icemulti [options] input-files\n");
    log("\n");
    log(" -c\n");
    log(" coldboot mode, power on reset image is selected by CBSEL0/CBSEL1\n");
    log("\n");
    log(" -p0, -p1, -p2, -p3\n");
    log(" select power on reset image when not using coldboot mode\n");
    log("\n");
    log(" -a<n>, -A<n>\n");
    log(" align images at 2^<n> bytes. -A also aligns image 0.\n");
    log("\n");
    log(" -o filename\n");
    log(" write output image to file instead of stdout\n");
    log("\n");
    log(" -v\n");
    log(" verbose (repeat to increase verbosity)\n");
    log("\n");
    exit(EXIT_FAILURE);
}

int main(int argc, char **argv)
{
    int c;
    char *endptr = NULL;
    bool coldboot = false;
    int por_image = 0;
    int header_count = 0;
    int image_count = 0;
    int align_bits = 0;
    bool align_first = false;
    Image *header_images[NUM_IMAGES];
    std::unique_ptr<Image> images[NUM_IMAGES];
    const char *outfile_name = NULL;
    bool print_offsets = false;

    static struct option long_options[] = {
        {NULL, 0, NULL, 0}
    };

    program_short_name = strrchr(argv[0], '/');
    if (program_short_name == NULL)
        program_short_name = argv[0];
    else
        program_short_name++;

    while ((c = getopt_long(argc, argv, "cp:a:A:o:v",
                long_options, NULL)) != -1)
        switch (c) {
            case 'c':
                coldboot = true;
                break;
            case 'p':
                if (optarg[0] == '0' && optarg[1] == '\0')
                    por_image = 0;
                else if (optarg[0] == '1' && optarg[1] == '\0')
                    por_image = 1;
                else if (optarg[0] == '2' && optarg[1] == '\0')
                    por_image = 2;
                else if (optarg[0] == '3' && optarg[1] == '\0')
                    por_image = 3;
                else
                    error("`%s' is not a valid power-on/reset image (must be 0, 1, 2, or 3)\n", optarg);
                break;
            case 'A':
                align_first = true;
                /* fallthrough */
            case 'a':
                align_bits = strtol(optarg, &endptr, 0);
                if (*endptr != '\0')
                    error("`%s' is not a valid number\n", optarg);
                if (align_bits < 0)
                    error("argument to `-%c' must be non-negative\n", c);
                break;
            case 'o':
                outfile_name = optarg;
                break;
            case 'v':
                print_offsets = true;
                break;
            default:
                usage();
        }

    if (optind == argc) {
        fprintf(stderr, "%s: missing argument\n", program_short_name);
        usage();
    }

    while (optind != argc) {
        if (header_count >= NUM_IMAGES)
            error("Too many images supplied\n");
        images[image_count].reset(new Image(argv[optind++]));
        header_images[header_count] = &*images[image_count];
        header_count++;
        image_count++;
    }

    if (coldboot && por_image != 0)
        error("Can't select power on reset boot image in cold boot mode\n");

    if (por_image >= header_count)
        error("Specified non-existing image for power on reset\n");

    // Place images
    uint32_t offs = (NUM_IMAGES + 1) * HEADER_SIZE;
    if (align_first)
        align_offset(offs, align_bits);
    for (int i=0; i<image_count; i++) {
        images[i]->place(offs);
        offs += images[i]->size();
        align_offset(offs, align_bits);
        if (print_offsets)
            fprintf(stderr, "Place image %d at %06x .. %06x.\n", i, int(images[i]->offset()), int(offs));
    }

    // Populate headers
    for (int i=header_count; i < NUM_IMAGES; i++)
        header_images[i] = header_images[por_image];

    std::ofstream ofs;
    std::ostream *osp;

    if (outfile_name != NULL) {
        ofs.open(outfile_name, std::ofstream::binary);
        if (!ofs.is_open())
            error("can't open output file `%s': %s\n", outfile_name, strerror(errno));
        osp = &ofs;
    } else {
        osp = &std::cout;
    }

    uint32_t file_offset = 0;
    for (int i=0; i<NUM_IMAGES + 1; i++)
    {
        pad_to(*osp, file_offset, i * HEADER_SIZE);
        if (i == 0)
            write_header(*osp, file_offset, header_images[por_image], coldboot);
        else
            write_header(*osp, file_offset, header_images[i - 1], false);
    }
    for (int i=0; i<image_count; i++)
    {
        pad_to(*osp, file_offset, images[i]->offset());
        images[i]->write(*osp, file_offset);
    }

    return EXIT_SUCCESS;
}
<commit_msg>icemulti: Print image filenames along with offsets<commit_after>//
//  Copyright (C) 2015  Marcus Comstedt <marcus@mc.pp.se>
//
//  Permission to use, copy, modify, and/or distribute this software for any
//  purpose with or without fee is hereby granted, provided that the above
//  copyright notice and this permission notice appear in all copies.
//
//  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
//  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
//  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
//  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
//  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
//  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
//  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
//

#include <fstream>
#include <iostream>
#include <cstdint>
#include <memory>

#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define log(...) fprintf(stderr, __VA_ARGS__);
#define error(...) do { fprintf(stderr, "%s: ", program_short_name); fprintf(stderr, __VA_ARGS__); exit(EXIT_FAILURE); } while (0)

static char *program_short_name;

static const int NUM_IMAGES = 4;
static const int HEADER_SIZE = 32;

static void align_offset(uint32_t &offset, int bits)
{
    uint32_t mask = (1 << bits) - 1;
    if (offset & mask)
        offset = (offset | mask) + 1;
}

static void write_byte(std::ostream &ofs, uint32_t &file_offset, uint8_t byte)
{
    ofs << byte;
    file_offset++;
}

static void write_bytes(std::ostream &ofs, uint32_t &file_offset,
                        const uint8_t *buf, size_t n)
{
    if (n > 0) {
        ofs.write(reinterpret_cast<const char*>(buf), n);
        file_offset += n;
    }
}

static void write_file(std::ostream &ofs, uint32_t &file_offset,
                       std::istream &ifs, const char *filename)
{
    const size_t bufsize = 8192;
    uint8_t *buffer = new uint8_t[bufsize];

    while(!ifs.eof()) {
        ifs.read(reinterpret_cast<char *>(buffer), bufsize);
        if (ifs.bad())
            error("can't read input image `%s': %s\n", filename, strerror(errno));
        write_bytes(ofs, file_offset, buffer, ifs.gcount());
    }

    delete[] buffer;
}

static void pad_to(std::ostream &ofs, uint32_t &file_offset, uint32_t target)
{
    if (target < file_offset)
        error("Trying to pad backwards!\n");
    while(file_offset < target)
        write_byte(ofs, file_offset, 0xff);
}

class Image {
    std::ifstream ifs;
    uint32_t offs;

public:
    const char *const filename;

    Image(const char *filename);
    size_t size();
    void write(std::ostream &ofs, uint32_t &file_offset);
    void place(uint32_t o) { offs = o; }
    uint32_t offset() const { return offs; }
};

Image::Image(const char *filename) : ifs(filename, std::ifstream::binary), filename(filename)
{
    if (ifs.fail())
        error("can't open input image `%s': %s\n", filename, strerror(errno));
}

size_t Image::size()
{
    ifs.seekg (0, ifs.end);
    if (ifs.fail())
        error("can't seek on input image `%s': %s\n", filename, strerror(errno));
    size_t length = ifs.tellg();
    ifs.seekg (0, ifs.beg);
    if (ifs.fail())
        error("can't seek on input image `%s': %s\n", filename, strerror(errno));

    if (length == 0)
        error("input image `%s' doesn't contain any data\n", filename);
    return length;
}

void Image::write(std::ostream &ofs, uint32_t &file_offset)
{
    write_file(ofs, file_offset, ifs, filename);
}

static void write_header(std::ostream &ofs, uint32_t &file_offset,
                         Image const *image, bool coldboot)
{
    // Preamble
    write_byte(ofs, file_offset, 0x7e);
    write_byte(ofs, file_offset, 0xaa);
    write_byte(ofs, file_offset, 0x99);
    write_byte(ofs, file_offset, 0x7e);

    // Boot mode
    write_byte(ofs, file_offset, 0x92);
    write_byte(ofs, file_offset, 0x00);
    write_byte(ofs, file_offset, coldboot ? 0x10 : 0x00);

    // Boot address
    write_byte(ofs, file_offset, 0x44);
    write_byte(ofs, file_offset, 0x03);
    write_byte(ofs, file_offset, (image->offset() >> 16) & 0xff);
    write_byte(ofs, file_offset, (image->offset() >> 8) & 0xff);
    write_byte(ofs, file_offset, image->offset() & 0xff);

    // Bank offset
    write_byte(ofs, file_offset, 0x82);
    write_byte(ofs, file_offset, 0x00);
    write_byte(ofs, file_offset, 0x00);

    // Reboot
    write_byte(ofs, file_offset, 0x01);
    write_byte(ofs, file_offset, 0x08);

    // Zero out any unused bytes
    while (file_offset & (HEADER_SIZE - 1))
        write_byte(ofs, file_offset, 0x00);
}

void usage()
{
    log("\n");
    log("Usage: icemulti [options] input-files\n");
    log("\n");
    log(" -c\n");
    log(" coldboot mode, power on reset image is selected by CBSEL0/CBSEL1\n");
    log("\n");
    log(" -p0, -p1, -p2, -p3\n");
    log(" select power on reset image when not using coldboot mode\n");
    log("\n");
    log(" -a<n>, -A<n>\n");
    log(" align images at 2^<n> bytes. -A also aligns image 0.\n");
    log("\n");
    log(" -o filename\n");
    log(" write output image to file instead of stdout\n");
    log("\n");
    log(" -v\n");
    log(" verbose (repeat to increase verbosity)\n");
    log("\n");
    exit(EXIT_FAILURE);
}

int main(int argc, char **argv)
{
    int c;
    char *endptr = NULL;
    bool coldboot = false;
    int por_image = 0;
    int header_count = 0;
    int image_count = 0;
    int align_bits = 0;
    bool align_first = false;
    Image *header_images[NUM_IMAGES];
    std::unique_ptr<Image> images[NUM_IMAGES];
    const char *outfile_name = NULL;
    bool print_offsets = false;

    static struct option long_options[] = {
        {NULL, 0, NULL, 0}
    };

    program_short_name = strrchr(argv[0], '/');
    if (program_short_name == NULL)
        program_short_name = argv[0];
    else
        program_short_name++;

    while ((c = getopt_long(argc, argv, "cp:a:A:o:v",
                long_options, NULL)) != -1)
        switch (c) {
            case 'c':
                coldboot = true;
                break;
            case 'p':
                if (optarg[0] == '0' && optarg[1] == '\0')
                    por_image = 0;
                else if (optarg[0] == '1' && optarg[1] == '\0')
                    por_image = 1;
                else if (optarg[0] == '2' && optarg[1] == '\0')
                    por_image = 2;
                else if (optarg[0] == '3' && optarg[1] == '\0')
                    por_image = 3;
                else
                    error("`%s' is not a valid power-on/reset image (must be 0, 1, 2, or 3)\n", optarg);
                break;
            case 'A':
                align_first = true;
                /* fallthrough */
            case 'a':
                align_bits = strtol(optarg, &endptr, 0);
                if (*endptr != '\0')
                    error("`%s' is not a valid number\n", optarg);
                if (align_bits < 0)
                    error("argument to `-%c' must be non-negative\n", c);
                break;
            case 'o':
                outfile_name = optarg;
                break;
            case 'v':
                print_offsets = true;
                break;
            default:
                usage();
        }

    if (optind == argc) {
        fprintf(stderr, "%s: missing argument\n", program_short_name);
        usage();
    }

    while (optind != argc) {
        if (header_count >= NUM_IMAGES)
            error("Too many images supplied\n");
        images[image_count].reset(new Image(argv[optind++]));
        header_images[header_count] = &*images[image_count];
        header_count++;
        image_count++;
    }

    if (coldboot && por_image != 0)
        error("Can't select power on reset boot image in cold boot mode\n");

    if (por_image >= header_count)
        error("Specified non-existing image for power on reset\n");

    // Place images
    uint32_t offs = (NUM_IMAGES + 1) * HEADER_SIZE;
    if (align_first)
        align_offset(offs, align_bits);
    for (int i=0; i<image_count; i++) {
        images[i]->place(offs);
        offs += images[i]->size();
        align_offset(offs, align_bits);
        if (print_offsets)
            fprintf(stderr, "Place image %d at %06x .. %06x (`%s')\n", i, int(images[i]->offset()), int(offs), images[i]->filename);
    }

    // Populate headers
    for (int i=header_count; i < NUM_IMAGES; i++)
        header_images[i] = header_images[por_image];

    std::ofstream ofs;
    std::ostream *osp;

    if (outfile_name != NULL) {
        ofs.open(outfile_name, std::ofstream::binary);
        if (!ofs.is_open())
            error("can't open output file `%s': %s\n", outfile_name, strerror(errno));
        osp = &ofs;
    } else {
        osp = &std::cout;
    }

    uint32_t file_offset = 0;
    for (int i=0; i<NUM_IMAGES + 1; i++)
    {
        pad_to(*osp, file_offset, i * HEADER_SIZE);
        if (i == 0)
            write_header(*osp, file_offset, header_images[por_image], coldboot);
        else
            write_header(*osp, file_offset, header_images[i - 1], false);
    }
    for (int i=0; i<image_count; i++)
    {
        pad_to(*osp, file_offset, images[i]->offset());
        images[i]->write(*osp, file_offset);
    }

    return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before>#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct GrutState {
  char* start_of_string;
  // TODO: We need to handle the correct sizing of this.
  char* captures_and_registers[64];
};

extern "C" int grut(GrutState* state, const char* s);

static void dump(GrutState* state, int capturePairs) {
  *state->captures_and_registers[1] = '\0';
  printf("«%s»\n", state->captures_and_registers[0]);
  for (int i = 2; i < capturePairs; i += 2) {
    if (state->captures_and_registers[i] == nullptr) {
      printf("%d: «null»\n", i >> 1);
    } else {
      char* end = state->captures_and_registers[i + 1];
      char saved = *end;
      *end = '\0';
      printf("%d: «%s»\n", i >> 1, state->captures_and_registers[i]);
      *end = saved;
    }
  }
}

int main(int argc, const char* const* argv) {
  bool multiline = false;
  if (argc >= 1 && !strcmp(argv[1], "-m")) {
    multiline = true;
    argc--;
    argv++;
  }
  if (argc != 2) {
    fprintf(stderr, "Usage: grutdump <inputfile> > outputfile\n");
    exit(5);
  }
  GrutState state;
  const char* fn = argv[1];
  FILE* fp = fopen(fn, "r");
  if (fp == NULL) {
    perror(fn);
    exit(1);
  }
#define BUF_SIZ (4 << 20)
  static char buffer[BUF_SIZ];
  size_t bytes = fread(buffer, 1, BUF_SIZ, fp);
  if (bytes == BUF_SIZ) {
    fprintf(stderr, "Input too large: %s\n", fn);
    exit(2);
  }
  const char* input = buffer;
  if (bytes == 0) {
    if(!feof(fp)) {
      perror(fn);
      exit(3);
    }
    input = "";
  } else {
    // Remove one trailing newline from the input because many editors like
    // to add a newline to the last line,
    if (buffer[bytes - 1] == '\n') buffer[bytes - 1] = '\0';
  }
  int capturePairs = grut(nullptr, nullptr);
  if (multiline) {
    while (true) {
      char* newline = const_cast<char*>(strchr(input, '\n'));
      if (newline != nullptr) *newline = '\0';
      if (grut(&state, input)) {
	dump(&state, capturePairs);
      } else {
	printf("\n");
      }
      if (newline == nullptr) break;
      input = newline + 1;
    }
  } else {
    if (grut(&state, input)) {
      dump(&state, capturePairs);
    }
  }
  return 0;
}
<commit_msg>Use guard page in tests to reveal out of bounds reads or writes<commit_after>#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

struct GrutState {
  char* start_of_string;
  // TODO: We need to handle the correct sizing of this.
  char* captures_and_registers[64];
};

extern "C" int grut(GrutState* state, const char* s);

static const intptr_t kBufSize = 4 << 20;
static const intptr_t kMargin = 64 << 10;  // 64k.

static char tmp[kBufSize];

static void dump(GrutState* state, int capturePairs) {
  int len = state->captures_and_registers[1] - state->captures_and_registers[0];
  memcpy(tmp, state->captures_and_registers[0], len);
  tmp[len] = '\0';
  printf("«%s»\n", tmp);
  for (int i = 2; i < capturePairs; i += 2) {
    if (state->captures_and_registers[i] == nullptr) {
      printf("%d: «null»\n", i >> 1);
    } else {
      char* start = state->captures_and_registers[i];
      char* end = state->captures_and_registers[i + 1];
      memcpy(tmp, start, end - start);
      tmp[end - start] = '\0';
      printf("%d: «%s»\n", i >> 1, tmp);
    }
  }
}

int main(int argc, const char* const* argv) {
  void* arena = mmap(nullptr, kBufSize, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
  if (arena == MAP_FAILED) {
    perror("mmap");
    exit(8);
  }
  char* start = reinterpret_cast<char*>(arena) + kMargin;
  char* end = reinterpret_cast<char*>(arena) + kBufSize - kMargin;
  if (mprotect(arena, kMargin, PROT_NONE)) {
    perror("mprotect1");
    exit(8);
  }
  if (mprotect(end, kMargin, PROT_NONE)) {
    perror("mprotect2");
    exit(8);
  }

  bool multiline = false;
  if (argc >= 1 && !strcmp(argv[1], "-m")) {
    multiline = true;
    argc--;
    argv++;
  }
  if (argc != 2) {
    fprintf(stderr, "Usage: grutdump <inputfile> > outputfile\n");
    exit(5);
  }
  GrutState state;
  const char* fn = argv[1];
  FILE* fp = fopen(fn, "r");
  if (fp == NULL) {
    perror(fn);
    exit(1);
  }
  static char buffer[kBufSize];
  size_t bytes = fread(buffer, 1, kBufSize, fp);
  if (bytes == kBufSize) {
    fprintf(stderr, "Input too large: %s\n", fn);
    exit(2);
  }
  const char* input = buffer;
  if (bytes == 0) {
    if(!feof(fp)) {
      perror(fn);
      exit(3);
    }
    input = "";
  } else {
    // Remove one trailing newline from the input because many editors like
    // to add a newline to the last line,
    if (buffer[bytes - 1] == '\n') buffer[bytes - 1] = '\0';
    bytes--;
  }
  int capturePairs = grut(nullptr, nullptr);
  if (multiline) {
    while (true) {
      bool done = false;
      char* newline = const_cast<char*>(strchr(input, '\n'));
      if (newline == nullptr) {
	done = true;
	newline = const_cast<char*>(input + strlen(input));
      } else {
	if (*newline != '\0') *newline = '\0';
      }
      int len = newline + 1 - input;
      memcpy(start, input, len);
      memcpy(end - len, input, len);
      if (mprotect(start, kBufSize - 2 * kMargin, PROT_READ)) {
	perror("mprotect3");
	exit(42);
      }
      if (grut(&state, start)) {
	dump(&state, capturePairs);
      } else {
	printf("\n");
      }
      grut(&state, end - len);  // Check if we hit the guard.
      if (mprotect(start, kBufSize - 2 * kMargin, PROT_READ | PROT_WRITE)) {
	perror("mprotect4");
	exit(103);
      }
      if (done) break;
      input = newline + 1;
    }
  } else {
    memcpy(start, input, bytes + 1);
    memcpy(end - bytes - 1, input, bytes + 1);
    if (mprotect(start, kBufSize - 2 * kMargin, PROT_READ)) {
      perror("mprotect3");
      exit(42);
    }
    if (grut(&state, start)) {
      dump(&state, capturePairs);
    }
    grut(&state, end - bytes - 1);  // Check if we hit the guard.
    if (mprotect(start, kBufSize - 2 * kMargin, PROT_READ | PROT_WRITE)) {
      perror("mprotect4");
      exit(103);
    }
  }
  return 0;
}
<|endoftext|>