<commit_before>/*
 *            Copyright 2009-2019 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      Licensed under the Apache License, Version 2.0 (the "License")
 *
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *              http://www.apache.org/licenses/LICENSE-2.0
 *
 *Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include <chrono>
#include <votca/xtp/aomatrix.h>
#include <votca/xtp/logger.h>
#include <votca/xtp/threecenter.h>

using std::flush;

namespace votca {
namespace xtp {

void TCMatrix_gwbse::Initialize(int basissize, int mmin, int mmax, int nmin,
                                int nmax) {

  // here as storage indices starting from zero
  _nmin = nmin;
  _nmax = nmax;
  _ntotal = nmax - nmin + 1;
  _mmin = mmin;
  _mmax = mmax;
  _mtotal = mmax - mmin + 1;
  _basissize = basissize;

  // vector has mtotal elements
  _matrix = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, _basissize));
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using either CUDA or Openmp
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrix(const Eigen::MatrixXd& matrix) {

  // Try to run the operation in a Nvidia GPU, otherwise is the default Openmp
  // implementation
  std::chrono::time_point<std::chrono::system_clock> start, end;
  start = std::chrono::system_clock::now();
#if defined(USE_CUDA)
  XTP_LOG_SAVE(logDEBUG, _log)
      << TimeStamp()
      << " Using CUDA/OpenMP for tensor matrix multiplication: " << flush;
  // Preallocated memory for all the matrices
  CudaPipeline cuda_pip;
  const Eigen::MatrixXd& head = _matrix.front();
  CudaMatrix cuma_A{head.rows(), head.cols()};
  CudaMatrix cuma_B{matrix, cuda_pip.get_stream()};
  CudaMatrix cuma_C{head.rows(), matrix.cols()};

#pragma omp parallel for schedule(dynamic)
  for (int i_occ = 0; i_occ < _mtotal; i_occ++) {
    // All the GPU communication happens through a single thread that reuses all
    // memory allocated in the GPU and it's dynamically load-balanced by OpenMP.
    // The rest of the threads use the default CPU matrix multiplication
    if (OPENMP::getThreadId() == 0) {
      cuma_A.copy_to_gpu(_matrix[i_occ]);
      cuda_pip.gemm(cuma_A, cuma_B, cuma_C);
      _matrix[i_occ] = cuma_C;
    } else {
      Eigen::MatrixXd temp = _matrix[i_occ] * matrix;
      _matrix[i_occ] = temp;
    }
  }
#else
  XTP_LOG_SAVE(logDEBUG, _log)
      << TimeStamp()
      << " Using Default OpenMP for tensor matrix multiplication: " << flush;
  MultiplyRightWithAuxMatrixOpenMP(matrix);
#endif
  end = std::chrono::system_clock::now();
  std::chrono::duration<double> elapsed_time = end - start;
  std::cout << "MultiplyRightWithAuxMatrix: " << elapsed_time.count() << "\n";
  return;
}

/*
 * Fill the 3-center object by looping over shells of GW basis set and
 * calling FillBlock, which calculates all 3-center overlap integrals
 * associated to a particular shell, convoluted with the DFT orbital
 * coefficients
 */
void TCMatrix_gwbse::Fill(const AOBasis& gwbasis, const AOBasis& dftbasis,
                          const Eigen::MatrixXd& dft_orbitals) {
  std::chrono::time_point<std::chrono::system_clock> start, end;
  // needed for Rebuild())
  _auxbasis = &gwbasis;
  _dftbasis = &dftbasis;
  _dft_orbitals = &dft_orbitals;

  // If cuda is enabled the dft orbitals are sent first to the cuda gpu
  // and memory in the cuda gpu is allocated for the intermediate matrices
#if defined(USE_CUDA)
  CudaPipeline cuda_pip;
  std::array<CudaMatrix, 2> cuda_matrices =
      SendDFTMatricesToGPU(dft_orbitals, cuda_pip);
  std::array<CudaMatrix, 3> cuda_inter_matrices =
      CreateIntermediateCudaMatrices(dft_orbitals.rows());
#endif

  start = std::chrono::system_clock::now();
  // loop over all shells in the GW basis and get _Mmn for that shell
#pragma omp parallel for schedule(guided)  // private(_block)
  for (int is = 0; is < gwbasis.getNumofShells(); is++) {
    const AOShell& shell = gwbasis.getShell(is);

    // Fill block for this shell (3-center overlap with _dft_basis +
    // multiplication with _dft_orbitals )
    std::vector<Eigen::MatrixXd> symmstorage =
        ComputeSymmStorage(shell, dftbasis, dft_orbitals);

    // If cuda is enable all the GPU communication happens through a single
    // thread that reuses all memory allocated in the GPU and it's dynamically
    // load-balanced by OpenMP. The remaining threads will perform the
    // convolution using the default CPU method.
    std::vector<Eigen::MatrixXd> block;
#if defined(USE_CUDA)
    if (OPENMP::getThreadId() == 0) {
      block = FillBlockCUDA(symmstorage, cuda_matrices, cuda_inter_matrices,
                            cuda_pip);
    } else {
      block = FillBlock(symmstorage, dft_orbitals);
    }
#else
    // Otherwise the convolution is performed by Eigen
    block = FillBlock(symmstorage, dft_orbitals);
#endif
    // put into correct position
    for (int m_level = 0; m_level < _mtotal; m_level++) {
      _matrix[m_level].block(0, shell.getStartIndex(), _ntotal,
                             shell.getNumFunc()) = block[m_level];
    }  // m-th DFT orbital
  }    // shells of GW basis set
  end = std::chrono::system_clock::now();
  std::chrono::duration<double> elapsed_time = end - start;
  std::cout << "fillblock: " << elapsed_time.count() << "\n";
  AOOverlap auxoverlap;
  auxoverlap.Fill(gwbasis);
  AOCoulomb auxcoulomb;
  auxcoulomb.Fill(gwbasis);
  Eigen::MatrixXd inv_sqrt = auxcoulomb.Pseudo_InvSqrt_GWBSE(auxoverlap, 5e-7);
  _removedfunctions = auxcoulomb.Removedfunctions();
  MultiplyRightWithAuxMatrix(inv_sqrt);

  return;
}

/*
 * Determines the 3-center integrals for a given shell in the GW basis
 * by calculating the 3-center overlap integral of the functions in the
 * GW shell with ALL functions in the DFT basis set (FillThreeCenterOLBlock)
 */
std::vector<Eigen::MatrixXd> TCMatrix_gwbse::ComputeSymmStorage(
    const AOShell& auxshell, const AOBasis& dftbasis,
    const Eigen::MatrixXd& dft_orbitals) const {
  std::vector<Eigen::MatrixXd> symmstorage = std::vector<Eigen::MatrixXd>(
      auxshell.getNumFunc(),
      Eigen::MatrixXd::Zero(dftbasis.AOBasisSize(), dftbasis.AOBasisSize()));
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);
  // alpha-loop over the "left" DFT basis function
  for (int row = 0; row < dftbasis.getNumofShells(); row++) {

    const AOShell& shell_row = dftbasis.getShell(row);
    const int row_start = shell_row.getStartIndex();
    // ThreecMatrix is symmetric, restrict explicit calculation to triangular
    // matrix
    for (int col = 0; col <= row; col++) {
      const AOShell& shell_col = dftbasis.getShell(col);
      const int col_start = shell_col.getStartIndex();

      Eigen::Tensor<double, 3> threec_block(auxshell.getNumFunc(),
                                            shell_row.getNumFunc(),
                                            shell_col.getNumFunc());
      threec_block.setZero();

      bool nonzero =
          FillThreeCenterRepBlock(threec_block, auxshell, shell_row, shell_col);
      if (nonzero) {
        for (int aux_c = 0; aux_c < auxshell.getNumFunc(); aux_c++) {
          for (int row_c = 0; row_c < shell_row.getNumFunc(); row_c++) {
            for (int col_c = 0; col_c < shell_col.getNumFunc(); col_c++) {
              // symmetry
              if ((col_start + col_c) > (row_start + row_c)) {
                break;
              }
              symmstorage[aux_c](row_start + row_c, col_start + col_c) =
                  threec_block(aux_c, row_c, col_c);
            }  // ROW copy
          }    // COL copy
        }      // AUX copy
      }
    }  // gamma-loop
  }    // alpha-loop

  return symmstorage;
}

/*
 *  Convolution of the GW shell with ALL functions with the DFT orbital
 * coefficients
 */
std::vector<Eigen::MatrixXd> TCMatrix_gwbse::FillBlock(
    const std::vector<Eigen::MatrixXd>& symmstorage,
    const Eigen::MatrixXd& dft_orbitals) const {

  std::vector<Eigen::MatrixXd> block = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, symmstorage.size()));
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);

  int dim = static_cast<int>(symmstorage.size());
  for (auto k = 0; k < dim; ++k) {
    const Eigen::MatrixXd& matrix = symmstorage[k];
    Eigen::MatrixXd threec_inMo =
        dftn.transpose() * matrix.selfadjointView<Eigen::Lower>() * dftm;
    for (int i = 0; i < threec_inMo.cols(); ++i) {
      block[i].col(k) = threec_inMo.col(i);
    }
  }
  return block;
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using OPENMP parallelization
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrixOpenMP(
    const Eigen::MatrixXd& matrix) {
#pragma omp parallel for
  for (int i_occ = 0; i_occ < _mtotal; i_occ++) {
    Eigen::MatrixXd temp = _matrix[i_occ] * matrix;
    _matrix[i_occ] = temp;
  }
  return;
}

#if defined(USE_CUDA)
/*
 * Convolution of the GW shell with the DFT orbital coefficients using an Nvidia
 * GPU. The Cuda device behaves like a server that is receiving matrix-matrix
 * multiplications from a single stream (an Nvidia queue) and handle them
 * in an asynchronous way. It performs the following operations when recieving a
 * request:
 *  1. Check that there is enough space for the arrays
 *  2. Allocate memory for each matrix
 *  3. Copy the matrix to the allocated space
 *  4. Perform the matrix multiplication
 *  5. Return the result matrix
 * The Cuda device knows to which memory address it needs to copy back the
 * result. see: https://docs.nvidia.com/cuda/cublas/index.html#thread-safety2
 */
std::vector<Eigen::MatrixXd> TCMatrix_gwbse::FillBlockCUDA(
    const std::vector<Eigen::MatrixXd>& symmstorage,
    const std::array<CudaMatrix, 2>& cuda_matrices,
    std::array<CudaMatrix, 3>& cuda_inter_matrices,
    const CudaPipeline& cuda_pip) const {

  std::vector<Eigen::MatrixXd> block = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, symmstorage.size()));

  try {
    const CudaMatrix& cuma_A = cuda_matrices[0];
    const CudaMatrix& cuma_C = cuda_matrices[1];
    CudaMatrix& cuma_B = cuda_inter_matrices[0];
    CudaMatrix& cuma_X = cuda_inter_matrices[1];
    CudaMatrix& cuma_Y = cuda_inter_matrices[2];

    int dim = static_cast<int>(symmstorage.size());
    for (int k = 0; k < dim; ++k) {
      const Eigen::MatrixXd& matrix = symmstorage[k];
      cuma_B.copy_to_gpu(matrix.selfadjointView<Eigen::Lower>());
      cuda_pip.gemm(cuma_A, cuma_B, cuma_X);
      cuda_pip.gemm(cuma_X, cuma_C, cuma_Y);
      Eigen::MatrixXd threec_inMo = cuma_Y;
      for (int i = 0; i < threec_inMo.cols(); ++i) {
        block[i].col(k) = threec_inMo.col(i);
      }
    }
  } catch (const std::runtime_error& error) {
    XTP_LOG_SAVE(logDEBUG, _log)
        << TimeStamp() << " FillBlockCUDA failed due to: " << error.what()
        << flush;
    throw;
  }
  return block;
}

std::array<CudaMatrix, 2> TCMatrix_gwbse::SendDFTMatricesToGPU(
    const Eigen::MatrixXd& dft_orbitals, const CudaPipeline& cuda_pip) const {
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);

  // Smart Pointers to the cuda arrays
  const cudaStream_t& stream = cuda_pip.get_stream();

  return {CudaMatrix{dftn.transpose(), stream}, CudaMatrix{dftm, stream}};
}

std::array<CudaMatrix, 3> TCMatrix_gwbse::CreateIntermediateCudaMatrices(
    long basissize) const {
  long mcols = _mtotal - _mmin;
  long ncols = _ntotal - _nmin;

  return {CudaMatrix{basissize, basissize}, CudaMatrix{ncols, basissize},
          CudaMatrix{ncols, mcols}};
}
#endif

}  // namespace xtp
}  // namespace votca
<commit_msg>removed chrono<commit_after>/*
 *            Copyright 2009-2019 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      Licensed under the Apache License, Version 2.0 (the "License")
 *
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *              http://www.apache.org/licenses/LICENSE-2.0
 *
 *Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include <votca/xtp/aomatrix.h>
#include <votca/xtp/logger.h>
#include <votca/xtp/threecenter.h>

using std::flush;

namespace votca {
namespace xtp {

void TCMatrix_gwbse::Initialize(int basissize, int mmin, int mmax, int nmin,
                                int nmax) {

  // here as storage indices starting from zero
  _nmin = nmin;
  _nmax = nmax;
  _ntotal = nmax - nmin + 1;
  _mmin = mmin;
  _mmax = mmax;
  _mtotal = mmax - mmin + 1;
  _basissize = basissize;

  // vector has mtotal elements
  _matrix = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, _basissize));
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using either CUDA or Openmp
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrix(const Eigen::MatrixXd& matrix) {

  // Try to run the operation in a Nvidia GPU, otherwise is the default Openmp
  // implementation
#if defined(USE_CUDA)
  XTP_LOG_SAVE(logDEBUG, _log)
      << TimeStamp()
      << " Using CUDA/OpenMP for tensor matrix multiplication: " << flush;
  // Preallocated memory for all the matrices
  CudaPipeline cuda_pip;
  const Eigen::MatrixXd& head = _matrix.front();
  CudaMatrix cuma_A{head.rows(), head.cols()};
  CudaMatrix cuma_B{matrix, cuda_pip.get_stream()};
  CudaMatrix cuma_C{head.rows(), matrix.cols()};

#pragma omp parallel for schedule(dynamic)
  for (int i_occ = 0; i_occ < _mtotal; i_occ++) {
    // All the GPU communication happens through a single thread that reuses all
    // memory allocated in the GPU and it's dynamically load-balanced by OpenMP.
    // The rest of the threads use the default CPU matrix multiplication
    if (OPENMP::getThreadId() == 0) {
      cuma_A.copy_to_gpu(_matrix[i_occ]);
      cuda_pip.gemm(cuma_A, cuma_B, cuma_C);
      _matrix[i_occ] = cuma_C;
    } else {
      Eigen::MatrixXd temp = _matrix[i_occ] * matrix;
      _matrix[i_occ] = temp;
    }
  }
#else
  XTP_LOG_SAVE(logDEBUG, _log)
      << TimeStamp()
      << " Using Default OpenMP for tensor matrix multiplication: " << flush;
  MultiplyRightWithAuxMatrixOpenMP(matrix);
#endif
  return;
}

/*
 * Fill the 3-center object by looping over shells of GW basis set and
 * calling FillBlock, which calculates all 3-center overlap integrals
 * associated to a particular shell, convoluted with the DFT orbital
 * coefficients
 */
void TCMatrix_gwbse::Fill(const AOBasis& gwbasis, const AOBasis& dftbasis,
                          const Eigen::MatrixXd& dft_orbitals) {
  // needed for Rebuild())
  _auxbasis = &gwbasis;
  _dftbasis = &dftbasis;
  _dft_orbitals = &dft_orbitals;

  // If cuda is enabled the dft orbitals are sent first to the cuda gpu
  // and memory in the cuda gpu is allocated for the intermediate matrices
#if defined(USE_CUDA)
  CudaPipeline cuda_pip;
  std::array<CudaMatrix, 2> cuda_matrices =
      SendDFTMatricesToGPU(dft_orbitals, cuda_pip);
  std::array<CudaMatrix, 3> cuda_inter_matrices =
      CreateIntermediateCudaMatrices(dft_orbitals.rows());
#endif

  // loop over all shells in the GW basis and get _Mmn for that shell
#pragma omp parallel for schedule(guided)  // private(_block)
  for (int is = 0; is < gwbasis.getNumofShells(); is++) {
    const AOShell& shell = gwbasis.getShell(is);

    // Fill block for this shell (3-center overlap with _dft_basis +
    // multiplication with _dft_orbitals )
    std::vector<Eigen::MatrixXd> symmstorage =
        ComputeSymmStorage(shell, dftbasis, dft_orbitals);

    // If cuda is enable all the GPU communication happens through a single
    // thread that reuses all memory allocated in the GPU and it's dynamically
    // load-balanced by OpenMP. The remaining threads will perform the
    // convolution using the default CPU method.
    std::vector<Eigen::MatrixXd> block;
#if defined(USE_CUDA)
    if (OPENMP::getThreadId() == 0) {
      block = FillBlockCUDA(symmstorage, cuda_matrices, cuda_inter_matrices,
                            cuda_pip);
    } else {
      block = FillBlock(symmstorage, dft_orbitals);
    }
#else
    // Otherwise the convolution is performed by Eigen
    block = FillBlock(symmstorage, dft_orbitals);
#endif
    // put into correct position
    for (int m_level = 0; m_level < _mtotal; m_level++) {
      _matrix[m_level].block(0, shell.getStartIndex(), _ntotal,
                             shell.getNumFunc()) = block[m_level];
    }  // m-th DFT orbital
  }    // shells of GW basis set
  AOOverlap auxoverlap;
  auxoverlap.Fill(gwbasis);
  AOCoulomb auxcoulomb;
  auxcoulomb.Fill(gwbasis);
  Eigen::MatrixXd inv_sqrt = auxcoulomb.Pseudo_InvSqrt_GWBSE(auxoverlap, 5e-7);
  _removedfunctions = auxcoulomb.Removedfunctions();
  MultiplyRightWithAuxMatrix(inv_sqrt);

  return;
}

/*
 * Determines the 3-center integrals for a given shell in the GW basis
 * by calculating the 3-center overlap integral of the functions in the
 * GW shell with ALL functions in the DFT basis set (FillThreeCenterOLBlock)
 */
std::vector<Eigen::MatrixXd> TCMatrix_gwbse::ComputeSymmStorage(
    const AOShell& auxshell, const AOBasis& dftbasis,
    const Eigen::MatrixXd& dft_orbitals) const {
  std::vector<Eigen::MatrixXd> symmstorage = std::vector<Eigen::MatrixXd>(
      auxshell.getNumFunc(),
      Eigen::MatrixXd::Zero(dftbasis.AOBasisSize(), dftbasis.AOBasisSize()));
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);
  // alpha-loop over the "left" DFT basis function
  for (int row = 0; row < dftbasis.getNumofShells(); row++) {

    const AOShell& shell_row = dftbasis.getShell(row);
    const int row_start = shell_row.getStartIndex();
    // ThreecMatrix is symmetric, restrict explicit calculation to triangular
    // matrix
    for (int col = 0; col <= row; col++) {
      const AOShell& shell_col = dftbasis.getShell(col);
      const int col_start = shell_col.getStartIndex();

      Eigen::Tensor<double, 3> threec_block(auxshell.getNumFunc(),
                                            shell_row.getNumFunc(),
                                            shell_col.getNumFunc());
      threec_block.setZero();

      bool nonzero =
          FillThreeCenterRepBlock(threec_block, auxshell, shell_row, shell_col);
      if (nonzero) {
        for (int aux_c = 0; aux_c < auxshell.getNumFunc(); aux_c++) {
          for (int row_c = 0; row_c < shell_row.getNumFunc(); row_c++) {
            for (int col_c = 0; col_c < shell_col.getNumFunc(); col_c++) {
              // symmetry
              if ((col_start + col_c) > (row_start + row_c)) {
                break;
              }
              symmstorage[aux_c](row_start + row_c, col_start + col_c) =
                  threec_block(aux_c, row_c, col_c);
            }  // ROW copy
          }    // COL copy
        }      // AUX copy
      }
    }  // gamma-loop
  }    // alpha-loop

  return symmstorage;
}

/*
 *  Convolution of the GW shell with ALL functions with the DFT orbital
 * coefficients
 */
std::vector<Eigen::MatrixXd> TCMatrix_gwbse::FillBlock(
    const std::vector<Eigen::MatrixXd>& symmstorage,
    const Eigen::MatrixXd& dft_orbitals) const {

  std::vector<Eigen::MatrixXd> block = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, symmstorage.size()));
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);

  int dim = static_cast<int>(symmstorage.size());
  for (auto k = 0; k < dim; ++k) {
    const Eigen::MatrixXd& matrix = symmstorage[k];
    Eigen::MatrixXd threec_inMo =
        dftn.transpose() * matrix.selfadjointView<Eigen::Lower>() * dftm;
    for (int i = 0; i < threec_inMo.cols(); ++i) {
      block[i].col(k) = threec_inMo.col(i);
    }
  }
  return block;
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using OPENMP parallelization
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrixOpenMP(
    const Eigen::MatrixXd& matrix) {
#pragma omp parallel for
  for (int i_occ = 0; i_occ < _mtotal; i_occ++) {
    Eigen::MatrixXd temp = _matrix[i_occ] * matrix;
    _matrix[i_occ] = temp;
  }
  return;
}

#if defined(USE_CUDA)
/*
 * Convolution of the GW shell with the DFT orbital coefficients using an Nvidia
 * GPU. The Cuda device behaves like a server that is receiving matrix-matrix
 * multiplications from a single stream (an Nvidia queue) and handle them
 * in an asynchronous way. It performs the following operations when recieving a
 * request:
 *  1. Check that there is enough space for the arrays
 *  2. Allocate memory for each matrix
 *  3. Copy the matrix to the allocated space
 *  4. Perform the matrix multiplication
 *  5. Return the result matrix
 * The Cuda device knows to which memory address it needs to copy back the
 * result. see: https://docs.nvidia.com/cuda/cublas/index.html#thread-safety2
 */
std::vector<Eigen::MatrixXd> TCMatrix_gwbse::FillBlockCUDA(
    const std::vector<Eigen::MatrixXd>& symmstorage,
    const std::array<CudaMatrix, 2>& cuda_matrices,
    std::array<CudaMatrix, 3>& cuda_inter_matrices,
    const CudaPipeline& cuda_pip) const {

  std::vector<Eigen::MatrixXd> block = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, symmstorage.size()));

  try {
    const CudaMatrix& cuma_A = cuda_matrices[0];
    const CudaMatrix& cuma_C = cuda_matrices[1];
    CudaMatrix& cuma_B = cuda_inter_matrices[0];
    CudaMatrix& cuma_X = cuda_inter_matrices[1];
    CudaMatrix& cuma_Y = cuda_inter_matrices[2];

    int dim = static_cast<int>(symmstorage.size());
    for (int k = 0; k < dim; ++k) {
      const Eigen::MatrixXd& matrix = symmstorage[k];
      cuma_B.copy_to_gpu(matrix.selfadjointView<Eigen::Lower>());
      cuda_pip.gemm(cuma_A, cuma_B, cuma_X);
      cuda_pip.gemm(cuma_X, cuma_C, cuma_Y);
      Eigen::MatrixXd threec_inMo = cuma_Y;
      for (int i = 0; i < threec_inMo.cols(); ++i) {
        block[i].col(k) = threec_inMo.col(i);
      }
    }
  } catch (const std::runtime_error& error) {
    XTP_LOG_SAVE(logDEBUG, _log)
        << TimeStamp() << " FillBlockCUDA failed due to: " << error.what()
        << flush;
    throw;
  }
  return block;
}

std::array<CudaMatrix, 2> TCMatrix_gwbse::SendDFTMatricesToGPU(
    const Eigen::MatrixXd& dft_orbitals, const CudaPipeline& cuda_pip) const {
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);

  // Smart Pointers to the cuda arrays
  const cudaStream_t& stream = cuda_pip.get_stream();

  return {CudaMatrix{dftn.transpose(), stream}, CudaMatrix{dftm, stream}};
}

std::array<CudaMatrix, 3> TCMatrix_gwbse::CreateIntermediateCudaMatrices(
    long basissize) const {
  long mcols = _mtotal - _mmin;
  long ncols = _ntotal - _nmin;

  return {CudaMatrix{basissize, basissize}, CudaMatrix{ncols, basissize},
          CudaMatrix{ncols, mcols}};
}
#endif

}  // namespace xtp
}  // namespace votca
<|endoftext|>


<commit_before>/*
 *            Copyright 2009-2019 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      Licensed under the Apache License, Version 2.0 (the "License")
 *
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *              http://www.apache.org/licenses/LICENSE-2.0
 *
 *Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include <votca/xtp/aomatrix.h>
#include <votca/xtp/logger.h>
#include <votca/xtp/multiarray.h>
#include <votca/xtp/threecenter.h>

using std::flush;

namespace votca {
namespace xtp {

void TCMatrix_gwbse::Initialize(int basissize, int mmin, int mmax, int nmin,
                                int nmax) {

  // here as storage indices starting from zero
  _nmin = nmin;
  _nmax = nmax;
  _ntotal = nmax - nmin + 1;
  _mmin = mmin;
  _mmax = mmax;
  _mtotal = mmax - mmin + 1;
  _basissize = basissize;

  // vector has mtotal elements
  _matrix = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, _basissize));
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using either CUDA or Openmp
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrix(const Eigen::MatrixXd& matrix) {

  // Try to run the operation in a GPU, otherwise is the default Openmp
  // implementation
#if defined(USE_GPU)
  EigenCuda gpu_handle;
  try {
    _matrix = gpu_handle.right_matrix_tensor_mult(_matrix, matrix);
  } catch (const std::runtime_error& error) {
    XTP_LOG_SAVE(logDEBUG, _log)
        << TimeStamp()
        << " GPU Multiplyrightwithauxmatrix failed due to: " << error.what()
        << " Using default OpenMP implementation!" << flush;
    this->MultiplyRightWithAuxMatrixOpenMP(matrix);
  }
#else
  this->MultiplyRightWithAuxMatrixOpenMP(matrix);
#endif
  return;
}

/*
 * Fill the 3-center object by looping over shells of GW basis set and
 * calling FillBlock, which calculates all 3-center overlap integrals
 * associated to a particular shell, convoluted with the DFT orbital
 * coefficients
 */
void TCMatrix_gwbse::Fill(const AOBasis& gwbasis, const AOBasis& dftbasis,
                          const Eigen::MatrixXd& dft_orbitals) {
  // needed for Rebuild())
  _auxbasis = &gwbasis;
  _dftbasis = &dftbasis;
  _dft_orbitals = &dft_orbitals;

  // loop over all shells in the GW basis and get _Mmn for that shell
#pragma omp parallel for schedule(guided)  // private(_block)
  for (int is = 0; is < gwbasis.getNumofShells(); is++) {
    const AOShell& shell = gwbasis.getShell(is);
    std::vector<Eigen::MatrixXd> block;
    for (int i = 0; i < _mtotal; i++) {
      block.push_back(Eigen::MatrixXd::Zero(_ntotal, shell.getNumFunc()));
    }
    // Fill block for this shell (3-center overlap with _dft_basis +
    // multiplication with _dft_orbitals )
    FillBlock(block, shell, dftbasis, dft_orbitals);

    // put into correct position
    for (int m_level = 0; m_level < _mtotal; m_level++) {
      _matrix[m_level].block(0, shell.getStartIndex(), _ntotal,
                             shell.getNumFunc()) = block[m_level];
    }  // m-th DFT orbital
  }    // shells of GW basis set

  AOOverlap auxoverlap;
  auxoverlap.Fill(gwbasis);
  AOCoulomb auxcoulomb;
  auxcoulomb.Fill(gwbasis);
  Eigen::MatrixXd inv_sqrt = auxcoulomb.Pseudo_InvSqrt_GWBSE(auxoverlap, 5e-7);
  _removedfunctions = auxcoulomb.Removedfunctions();
  MultiplyRightWithAuxMatrix(inv_sqrt);
  return;
}

/*
 * Determines the 3-center integrals for a given shell in the GW basis
 * by calculating the 3-center overlap integral of the functions in the
 * GW shell with ALL functions in the DFT basis set (FillThreeCenterOLBlock),
 * followed by a convolution of those with the DFT orbital coefficients
 */

void TCMatrix_gwbse::FillBlock(std::vector<Eigen::MatrixXd>& block,
                               const AOShell& auxshell, const AOBasis& dftbasis,
                               const Eigen::MatrixXd& dft_orbitals) {
  tensor3d::extent_gen extents;
  std::vector<Eigen::MatrixXd> symmstorage;
  for (int i = 0; i < auxshell.getNumFunc(); ++i) {
    symmstorage.push_back(
        Eigen::MatrixXd::Zero(dftbasis.AOBasisSize(), dftbasis.AOBasisSize()));
  }
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);
  // alpha-loop over the "left" DFT basis function
  for (int row = 0; row < dftbasis.getNumofShells(); row++) {

    const AOShell& shell_row = dftbasis.getShell(row);
    const int row_start = shell_row.getStartIndex();
    // ThreecMatrix is symmetric, restrict explicit calculation to triangular
    // matrix
    for (int col = 0; col <= row; col++) {
      const AOShell& shell_col = dftbasis.getShell(col);
      const int col_start = shell_col.getStartIndex();

      tensor3d threec_block(extents[range(0, auxshell.getNumFunc())][range(
          0, shell_row.getNumFunc())][range(0, shell_col.getNumFunc())]);
      std::fill_n(threec_block.data(), threec_block.num_elements(), 0.0);

      bool nonzero =
          FillThreeCenterRepBlock(threec_block, auxshell, shell_row, shell_col);
      if (nonzero) {
        for (int aux_c = 0; aux_c < auxshell.getNumFunc(); aux_c++) {
          for (int row_c = 0; row_c < shell_row.getNumFunc(); row_c++) {
            for (int col_c = 0; col_c < shell_col.getNumFunc(); col_c++) {
              // symmetry
              if ((col_start + col_c) > (row_start + row_c)) {
                break;
              }
              symmstorage[aux_c](row_start + row_c, col_start + col_c) =
                  threec_block[aux_c][row_c][col_c];
            }  // ROW copy
          }    // COL copy
        }      // AUX copy
      }
    }  // gamma-loop
  }    // alpha-loop
  for (int k = 0; k < auxshell.getNumFunc(); ++k) {
    const Eigen::MatrixXd& matrix = symmstorage[k];
    Eigen::MatrixXd threec_inMo =
        dftn.transpose() * matrix.selfadjointView<Eigen::Lower>() * dftm;
    for (int i = 0; i < threec_inMo.cols(); ++i) {
      block[i].col(k) = threec_inMo.col(i);
    }
  }
  return;
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using OPENMP parallelization
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrixOpenMP(
    const Eigen::MatrixXd& matrix) {
#pragma omp parallel for
  for (int i_occ = 0; i_occ < _mtotal; i_occ++) {
    Eigen::MatrixXd temp = _matrix[i_occ] * matrix;
    _matrix[i_occ] = temp;
  }
  return;
}

}  // namespace xtp
}  // namespace votca
<commit_msg>removed mat_tensor_mat_mult method #262<commit_after>/*
 *            Copyright 2009-2019 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      Licensed under the Apache License, Version 2.0 (the "License")
 *
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *              http://www.apache.org/licenses/LICENSE-2.0
 *
 *Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include <votca/xtp/aomatrix.h>
#include <votca/xtp/logger.h>
#include <votca/xtp/multiarray.h>
#include <votca/xtp/threecenter.h>

using std::flush;

namespace votca {
namespace xtp {

void TCMatrix_gwbse::Initialize(int basissize, int mmin, int mmax, int nmin,
                                int nmax) {

  // here as storage indices starting from zero
  _nmin = nmin;
  _nmax = nmax;
  _ntotal = nmax - nmin + 1;
  _mmin = mmin;
  _mmax = mmax;
  _mtotal = mmax - mmin + 1;
  _basissize = basissize;

  // vector has mtotal elements
  _matrix = std::vector<Eigen::MatrixXd>(
      _mtotal, Eigen::MatrixXd::Zero(_ntotal, _basissize));
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using either CUDA or Openmp
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrix(const Eigen::MatrixXd& matrix) {

  // Try to run the operation in a GPU, otherwise is the default Openmp
  // implementation
#if defined(USE_GPU)
  EigenCuda gpu_handle;
  try {
    _matrix = gpu_handle.right_matrix_tensor_mult(_matrix, matrix);
  } catch (const std::runtime_error& error) {
    XTP_LOG_SAVE(logDEBUG, _log)
        << TimeStamp()
        << " GPU Multiplyrightwithauxmatrix failed due to: " << error.what()
        << " Using default OpenMP implementation!" << flush;
    this->MultiplyRightWithAuxMatrixOpenMP(matrix);
  }
#else
  this->MultiplyRightWithAuxMatrixOpenMP(matrix);
#endif
  return;
}

/*
 * Fill the 3-center object by looping over shells of GW basis set and
 * calling FillBlock, which calculates all 3-center overlap integrals
 * associated to a particular shell, convoluted with the DFT orbital
 * coefficients
 */
void TCMatrix_gwbse::Fill(const AOBasis& gwbasis, const AOBasis& dftbasis,
                          const Eigen::MatrixXd& dft_orbitals) {
  // needed for Rebuild())
  _auxbasis = &gwbasis;
  _dftbasis = &dftbasis;
  _dft_orbitals = &dft_orbitals;

  // loop over all shells in the GW basis and get _Mmn for that shell
#pragma omp parallel for schedule(guided)  // private(_block)
  for (int is = 0; is < gwbasis.getNumofShells(); is++) {
    const AOShell& shell = gwbasis.getShell(is);
    std::vector<Eigen::MatrixXd> block;
    for (int i = 0; i < _mtotal; i++) {
      block.push_back(Eigen::MatrixXd::Zero(_ntotal, shell.getNumFunc()));
    }
    // Fill block for this shell (3-center overlap with _dft_basis +
    // multiplication with _dft_orbitals )
    FillBlock(block, shell, dftbasis, dft_orbitals);

    // put into correct position
    for (int m_level = 0; m_level < _mtotal; m_level++) {
      _matrix[m_level].block(0, shell.getStartIndex(), _ntotal,
                             shell.getNumFunc()) = block[m_level];
    }  // m-th DFT orbital
  }    // shells of GW basis set

  AOOverlap auxoverlap;
  auxoverlap.Fill(gwbasis);
  AOCoulomb auxcoulomb;
  auxcoulomb.Fill(gwbasis);
  Eigen::MatrixXd inv_sqrt = auxcoulomb.Pseudo_InvSqrt_GWBSE(auxoverlap, 5e-7);
  _removedfunctions = auxcoulomb.Removedfunctions();
  MultiplyRightWithAuxMatrix(inv_sqrt);
  return;
}

/*
 * Determines the 3-center integrals for a given shell in the GW basis
 * by calculating the 3-center overlap integral of the functions in the
 * GW shell with ALL functions in the DFT basis set (FillThreeCenterOLBlock),
 * followed by a convolution of those with the DFT orbital coefficients
 */

void TCMatrix_gwbse::FillBlock(std::vector<Eigen::MatrixXd>& block,
                               const AOShell& auxshell, const AOBasis& dftbasis,
                               const Eigen::MatrixXd& dft_orbitals) {
  tensor3d::extent_gen extents;
  std::vector<Eigen::MatrixXd> symmstorage;
  for (int i = 0; i < auxshell.getNumFunc(); ++i) {
    symmstorage.push_back(
        Eigen::MatrixXd::Zero(dftbasis.AOBasisSize(), dftbasis.AOBasisSize()));
  }
  const Eigen::MatrixXd dftm =
      dft_orbitals.block(0, _mmin, dft_orbitals.rows(), _mtotal);
  const Eigen::MatrixXd dftn =
      dft_orbitals.block(0, _nmin, dft_orbitals.rows(), _ntotal);
  // alpha-loop over the "left" DFT basis function
  for (int row = 0; row < dftbasis.getNumofShells(); row++) {

    const AOShell& shell_row = dftbasis.getShell(row);
    const int row_start = shell_row.getStartIndex();
    // ThreecMatrix is symmetric, restrict explicit calculation to triangular
    // matrix
    for (int col = 0; col <= row; col++) {
      const AOShell& shell_col = dftbasis.getShell(col);
      const int col_start = shell_col.getStartIndex();

      tensor3d threec_block(extents[range(0, auxshell.getNumFunc())][range(
          0, shell_row.getNumFunc())][range(0, shell_col.getNumFunc())]);
      std::fill_n(threec_block.data(), threec_block.num_elements(), 0.0);

      bool nonzero =
          FillThreeCenterRepBlock(threec_block, auxshell, shell_row, shell_col);
      if (nonzero) {
        for (int aux_c = 0; aux_c < auxshell.getNumFunc(); aux_c++) {
          for (int row_c = 0; row_c < shell_row.getNumFunc(); row_c++) {
            for (int col_c = 0; col_c < shell_col.getNumFunc(); col_c++) {
              // symmetry
              if ((col_start + col_c) > (row_start + row_c)) {
                break;
              }
              symmstorage[aux_c](row_start + row_c, col_start + col_c) =
                  threec_block[aux_c][row_c][col_c];
            }  // ROW copy
          }    // COL copy
        }      // AUX copy
      }
    }  // gamma-loop
  }    // alpha-loop
  for (int k = 0; k < auxshell.getNumFunc(); ++k) {
    const Eigen::MatrixXd& matrix = symmstorage[k];
    Eigen::MatrixXd threec_inMo =
        dftn.transpose() * matrix.selfadjointView<Eigen::Lower>() * dftm;
    for (int i = 0; i < threec_inMo.cols(); ++i) {
      block[i].col(k) = threec_inMo.col(i);
    }
  }
  return;
}

/*
 * Modify 3-center matrix elements consistent with use of symmetrized
 * Coulomb interaction using OPENMP parallelization
 */
void TCMatrix_gwbse::MultiplyRightWithAuxMatrixOpenMP(
    const Eigen::MatrixXd& matrix) {
#pragma omp parallel for
  for (int i_occ = 0; i_occ < _mtotal; i_occ++) {
    Eigen::MatrixXd temp = _matrix[i_occ] * matrix;
    _matrix[i_occ] = temp;
  }
  return;
}

}  // namespace xtp
}  // namespace votca
<|endoftext|>


<commit_before>/*
	This file is part of cpp-ethereum.

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

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

	You should have received a copy of the GNU General Public License
	along with cpp-ethereum.  If not, see <http://www.gnu.org/licenses/>.
*/
/**
 * @author Christian <c@ethdev.com>
 * @date 2014
 * Solidity commandline compiler.
 */

#include <string>
#include <iostream>
#include <fstream>

#include <boost/program_options.hpp>
#include <boost/filesystem.hpp>

#include "BuildInfo.h"
#include <libdevcore/Common.h>
#include <libdevcore/CommonData.h>
#include <libdevcore/CommonIO.h>
#include <libevmcore/Instruction.h>
#include <libsolidity/Scanner.h>
#include <libsolidity/Parser.h>
#include <libsolidity/ASTPrinter.h>
#include <libsolidity/NameAndTypeResolver.h>
#include <libsolidity/Exceptions.h>
#include <libsolidity/CompilerStack.h>
#include <libsolidity/SourceReferenceFormatter.h>

using namespace std;
using namespace dev;
using namespace solidity;
namespace po = boost::program_options;

void version()
{
	cout << "solc, the solidity complier commandline interface " << dev::Version << endl
		 << "  by Christian <c@ethdev.com> and Lefteris <lefteris@ethdev.com>, (c) 2014." << endl
		 << "Build: " << DEV_QUOTED(ETH_BUILD_PLATFORM) << "/" << DEV_QUOTED(ETH_BUILD_TYPE) << endl;
	exit(0);
}

enum class OutputType
{
	STDOUT,
	FILE,
	BOTH
};

#define outputTypeStr "Legal values:\n"\
	 "\tstdout: Print it to standard output\n"\
	 "\tfile: Print it to a file with same name\n"\
	"\tboth: Print both to a file and the stdout\n"

std::istream& operator>>(std::istream& _in, OutputType& io_output)
{
	std::string token;
	_in >> token;
	if (token == "stdout")
		io_output = OutputType::STDOUT;
	else if (token == "file")
		io_output = OutputType::FILE;
	else if (token == "both")
		io_output = OutputType::BOTH;
	else
		throw boost::program_options::invalid_option_value(token);
	return _in;
}

int main(int argc, char** argv)
{
	// Declare the supported options.
	po::options_description desc("Allowed options");
	desc.add_options()
	("help", "Show help message and exit")
	("version", "Show version and exit")
	("optimize", po::value<bool>()->default_value(false), "Optimize bytecode for size")
	("input-file", po::value<vector<string>>(), "input file")
	("ast", po::value<OutputType>(),
	 "Request to output the AST of the contract. " outputTypeStr)
	("asm", po::value<OutputType>(),
	 "Request to output the EVM assembly of the contract. "  outputTypeStr)
	("opcodes", po::value<OutputType>(),
	 "Request to output the Opcodes of the contract. "  outputTypeStr)
	("binary", po::value<OutputType>(),
	 "Request to output the contract in binary (hexadecimal). "  outputTypeStr);

	// All positional options should be interpreted as input files
	po::positional_options_description p;
	p.add("input-file", -1);

	// parse the compiler arguments
	po::variables_map vm;
	try
	{
		po::store(po::command_line_parser(argc, argv).options(desc).positional(p).allow_unregistered().run(), vm);
	}
	catch (po::error const& exception)
	{
		cout << exception.what() << endl;
		return -1;
	}
	po::notify(vm);

	if (vm.count("help")) {
		cout << desc;
		return 0;
	}

	if (vm.count("version")) {
		version();
		return 0;
	}
	map<string, string> sourceCodes;
	if (!vm.count("input-file"))
	{
		string s;
		while (!cin.eof())
		{
			getline(cin, s);
			sourceCodes["<stdin>"].append(s);
		}
	}
	else
		for (string const& infile: vm["input-file"].as<vector<string>>())
			sourceCodes[infile] = asString(dev::contents(infile));

	CompilerStack compiler;
	try
	{
		for (auto const& sourceCode: sourceCodes)
			compiler.addSource(sourceCode.first, sourceCode.second);
		compiler.compile( vm["optimize"].as<bool>());
	}
	catch (ParserError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Parser error", compiler);
		return -1;
	}
	catch (DeclarationError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Declaration error", compiler);
		return -1;
	}
	catch (TypeError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Type error", compiler);
		return -1;
	}
	catch (CompilerError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Compiler error", compiler);
		return -1;
	}
	catch (InternalCompilerError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Internal compiler error", compiler);
		return -1;
	}
	catch (Exception const& exception)
	{
		cerr << "Exception during compilation: " << boost::diagnostic_information(exception) << endl;
		return -1;
	}
	catch (...)
	{
		cerr << "Unknown exception during compilation." << endl;
		return -1;
	}

	// do we need AST output?
	if (vm.count("ast"))
	{
		auto choice = vm["ast"].as<OutputType>();
		if (choice != OutputType::FILE)
		{
			cout << "Syntax trees:" << endl << endl;
			for (auto const& sourceCode: sourceCodes)
			{
				cout << endl << "======= " << sourceCode.first << " =======" << endl;
				ASTPrinter printer(compiler.getAST(sourceCode.first), sourceCode.second);
				printer.print(cout);
			}
		}

		if (choice != OutputType::STDOUT)
		{
			for (auto const& sourceCode: sourceCodes)
			{
				boost::filesystem::path p(sourceCode.first);
				ofstream outFile(p.stem().string() + ".ast");
				ASTPrinter printer(compiler.getAST(sourceCode.first), sourceCode.second);
				printer.print(outFile);
				outFile.close();
			}
		}
	}

	vector<string> contracts = compiler.getContractNames();
	for (string const& contract: contracts)
	{
		cout << endl << "======= " << contract << " =======" << endl;

		// do we need EVM assembly?
		if (vm.count("asm"))
		{
			auto choice = vm["asm"].as<OutputType>();
			if (choice != OutputType::FILE)
			{
				cout << "EVM assembly:" << endl;
				compiler.streamAssembly(cout, contract);
			}

			if (choice != OutputType::STDOUT)
			{
				ofstream outFile(contract + ".evm");
				compiler.streamAssembly(outFile, contract);
				outFile.close();
			}
		}

		// do we need opcodes?
		if (vm.count("opcodes"))
		{
			auto choice = vm["opcodes"].as<OutputType>();
			if (choice != OutputType::FILE)
			{
				cout << "Opcodes:" << endl;
				cout << eth::disassemble(compiler.getBytecode(contract)) << endl;
			}

			if (choice != OutputType::STDOUT)
			{
				ofstream outFile(contract + ".opcodes");
				outFile << eth::disassemble(compiler.getBytecode(contract));
				outFile.close();
			}
		}

		// do we need binary?
		if (vm.count("binary"))
		{
			auto choice = vm["binary"].as<OutputType>();
			if (choice != OutputType::FILE)
			{
				cout << "Binary:" << endl;
				cout << toHex(compiler.getBytecode(contract)) << endl;
			}

			if (choice != OutputType::STDOUT)
			{   // TODO: Think, if we really want that? Could simply output to normal binary
				ofstream outFile(contract + ".bin");
				outFile << toHex(compiler.getBytecode(contract));
				outFile.close();
			}
		}

	} // end of contracts iteration

	cout << endl << "Contracts:" << endl;
	for (string const& contract: contracts)
	{
			 cout << "Interface specification: " << compiler.getJsonDocumentation(contract, DocumentationType::ABI_INTERFACE) << endl
			 << "Natspec user documentation: " << compiler.getJsonDocumentation(contract, DocumentationType::NATSPEC_USER) << endl
			 << "Natspec developer documentation: " << compiler.getJsonDocumentation(contract, DocumentationType::NATSPEC_DEV) << endl;
	}

	return 0;
}
<commit_msg>Solc gets arguments for interface and documentation related output<commit_after>/*
	This file is part of cpp-ethereum.

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

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

	You should have received a copy of the GNU General Public License
	along with cpp-ethereum.  If not, see <http://www.gnu.org/licenses/>.
*/
/**
 * @author Christian <c@ethdev.com>
 * @date 2014
 * Solidity commandline compiler.
 */

#include <string>
#include <iostream>
#include <fstream>

#include <boost/program_options.hpp>
#include <boost/filesystem.hpp>

#include "BuildInfo.h"
#include <libdevcore/Common.h>
#include <libdevcore/CommonData.h>
#include <libdevcore/CommonIO.h>
#include <libevmcore/Instruction.h>
#include <libsolidity/Scanner.h>
#include <libsolidity/Parser.h>
#include <libsolidity/ASTPrinter.h>
#include <libsolidity/NameAndTypeResolver.h>
#include <libsolidity/Exceptions.h>
#include <libsolidity/CompilerStack.h>
#include <libsolidity/SourceReferenceFormatter.h>

using namespace std;
using namespace dev;
using namespace solidity;
namespace po = boost::program_options;

void version()
{
	cout << "solc, the solidity complier commandline interface " << dev::Version << endl
		 << "  by Christian <c@ethdev.com> and Lefteris <lefteris@ethdev.com>, (c) 2014." << endl
		 << "Build: " << DEV_QUOTED(ETH_BUILD_PLATFORM) << "/" << DEV_QUOTED(ETH_BUILD_TYPE) << endl;
	exit(0);
}

enum class OutputType
{
	STDOUT,
	FILE,
	BOTH
};

#define outputTypeStr "Legal values:\n"\
	 "\tstdout: Print it to standard output\n"\
	 "\tfile: Print it to a file with same name\n"\
	"\tboth: Print both to a file and the stdout\n"

std::istream& operator>>(std::istream& _in, OutputType& io_output)
{
	std::string token;
	_in >> token;
	if (token == "stdout")
		io_output = OutputType::STDOUT;
	else if (token == "file")
		io_output = OutputType::FILE;
	else if (token == "both")
		io_output = OutputType::BOTH;
	else
		throw boost::program_options::invalid_option_value(token);
	return _in;
}

static void handleBytecode(po::variables_map const& vm,
						   char const* _argName,
						   string const& _title,
						   string const& _contract,
						   CompilerStack&_compiler,
						   string const& _suffix)
{
	if (vm.count(_argName))
	{
		auto choice = vm[_argName].as<OutputType>();
		if (choice != OutputType::FILE)
		{
			cout << _title << endl;
			if (_suffix == "opcodes")
				cout << _compiler.getBytecode(_contract) << endl;
			else
				cout << toHex(_compiler.getBytecode(_contract)) << endl;
		}

		if (choice != OutputType::STDOUT)
		{
			ofstream outFile(_contract + _suffix);
			if (_suffix == "opcodes")
				outFile << _compiler.getBytecode(_contract);
			else
				outFile << toHex(_compiler.getBytecode(_contract));
			outFile.close();
		}
	}
}

static void handleJson(po::variables_map const& _vm,
					   DocumentationType _type,
					   string const& _contract,
					   CompilerStack&_compiler)
{
	std::string argName;
	std::string suffix;
	std::string title;
	switch(_type)
	{
	case DocumentationType::ABI_INTERFACE:
		argName = "abi";
		suffix = ".abi";
		title = "Contract ABI";
		break;
	case DocumentationType::NATSPEC_USER:
		argName = "natspec-user";
		suffix = ".docuser";
		title = "User Documentation";
		break;
	case DocumentationType::NATSPEC_DEV:
		argName = "natspec-dev";
		suffix = ".docdev";
		title = "Developer Documentation";
		break;
	default:
		// should never happen
		BOOST_THROW_EXCEPTION(InternalCompilerError() << errinfo_comment("Unknown documentation _type"));
	}

	if (_vm.count(argName.c_str()))
	{
		auto choice = _vm[argName.c_str()].as<OutputType>();
		if (choice != OutputType::FILE)
		{
			cout << title << endl;
			cout << _compiler.getJsonDocumentation(_contract, _type);
		}

		if (choice != OutputType::STDOUT)
		{
			ofstream outFile(_contract + suffix);
			outFile << _compiler.getJsonDocumentation(_contract, _type);
			outFile.close();
		}
	}
}

static inline bool argToStdout(po::variables_map const& _vm, const char* _name)
{
	return _vm.count(_name) && _vm[_name].as<OutputType>() != OutputType::FILE;
}

static bool needStdout(po::variables_map const& _vm)
{
	return argToStdout(_vm, "abi") || argToStdout(_vm, "natspec-user") || argToStdout(_vm, "natspec-dev") ||
		   argToStdout(_vm, "asm") || argToStdout(_vm, "opcodes") || argToStdout(_vm, "binary");
}

int main(int argc, char** argv)
{
	// Declare the supported options.
	po::options_description desc("Allowed options");
	desc.add_options()
	("help", "Show help message and exit")
	("version", "Show version and exit")
	("optimize", po::value<bool>()->default_value(false), "Optimize bytecode for size")
	("input-file", po::value<vector<string>>(), "input file")
	("ast", po::value<OutputType>(),
	 "Request to output the AST of the contract. " outputTypeStr)
	("asm", po::value<OutputType>(),
	 "Request to output the EVM assembly of the contract. "  outputTypeStr)
	("opcodes", po::value<OutputType>(),
	 "Request to output the Opcodes of the contract. "  outputTypeStr)
	("binary", po::value<OutputType>(),
	 "Request to output the contract in binary (hexadecimal). "  outputTypeStr)
	("abi", po::value<OutputType>(),
	 "Request to output the contract's ABI interface. "  outputTypeStr)
	("natspec-user", po::value<OutputType>(),
	 "Request to output the contract's Natspec user documentation. "  outputTypeStr)
	("natspec-dev", po::value<OutputType>(),
	 "Request to output the contract's Natspec developer documentation. "  outputTypeStr);

	// All positional options should be interpreted as input files
	po::positional_options_description p;
	p.add("input-file", -1);

	// parse the compiler arguments
	po::variables_map vm;
	try
	{
		po::store(po::command_line_parser(argc, argv).options(desc).positional(p).allow_unregistered().run(), vm);
	}
	catch (po::error const& exception)
	{
		cout << exception.what() << endl;
		return -1;
	}
	po::notify(vm);

	if (vm.count("help")) {
		cout << desc;
		return 0;
	}

	if (vm.count("version")) {
		version();
		return 0;
	}

	// create a map of input files to source code strings
	map<string, string> sourceCodes;
	if (!vm.count("input-file"))
	{
		string s;
		while (!cin.eof())
		{
			getline(cin, s);
			sourceCodes["<stdin>"].append(s);
		}
	}
	else
		for (string const& infile: vm["input-file"].as<vector<string>>())
			sourceCodes[infile] = asString(dev::contents(infile));

	// parse the files
	CompilerStack compiler;
	try
	{
		for (auto const& sourceCode: sourceCodes)
			compiler.addSource(sourceCode.first, sourceCode.second);
		compiler.compile( vm["optimize"].as<bool>());
	}
	catch (ParserError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Parser error", compiler);
		return -1;
	}
	catch (DeclarationError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Declaration error", compiler);
		return -1;
	}
	catch (TypeError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Type error", compiler);
		return -1;
	}
	catch (CompilerError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Compiler error", compiler);
		return -1;
	}
	catch (InternalCompilerError const& exception)
	{
		SourceReferenceFormatter::printExceptionInformation(cerr, exception, "Internal compiler error", compiler);
		return -1;
	}
	catch (Exception const& exception)
	{
		cerr << "Exception during compilation: " << boost::diagnostic_information(exception) << endl;
		return -1;
	}
	catch (...)
	{
		cerr << "Unknown exception during compilation." << endl;
		return -1;
	}

	/* -- act depending on the provided arguments -- */

	// do we need AST output?
	if (vm.count("ast"))
	{
		auto choice = vm["ast"].as<OutputType>();
		if (choice != OutputType::FILE)
		{
			cout << "Syntax trees:" << endl << endl;
			for (auto const& sourceCode: sourceCodes)
			{
				cout << endl << "======= " << sourceCode.first << " =======" << endl;
				ASTPrinter printer(compiler.getAST(sourceCode.first), sourceCode.second);
				printer.print(cout);
			}
		}

		if (choice != OutputType::STDOUT)
		{
			for (auto const& sourceCode: sourceCodes)
			{
				boost::filesystem::path p(sourceCode.first);
				ofstream outFile(p.stem().string() + ".ast");
				ASTPrinter printer(compiler.getAST(sourceCode.first), sourceCode.second);
				printer.print(outFile);
				outFile.close();
			}
		}
	}

	vector<string> contracts = compiler.getContractNames();
	for (string const& contract: contracts)
	{
		if (needStdout(vm))
			cout << endl << "======= " << contract << " =======" << endl;

		// do we need EVM assembly?
		if (vm.count("asm"))
		{
			auto choice = vm["asm"].as<OutputType>();
			if (choice != OutputType::FILE)
			{
				cout << "EVM assembly:" << endl;
				compiler.streamAssembly(cout, contract);
			}

			if (choice != OutputType::STDOUT)
			{
				ofstream outFile(contract + ".evm");
				compiler.streamAssembly(outFile, contract);
				outFile.close();
			}
		}

		handleBytecode(vm, "opcodes", "Opcodes:", contract, compiler, ".opcodes");
		handleBytecode(vm, "binary", "Binary:", contract, compiler, ".binary");
		handleJson(vm, DocumentationType::ABI_INTERFACE, contract, compiler);
		handleJson(vm, DocumentationType::NATSPEC_DEV, contract, compiler);
		handleJson(vm, DocumentationType::NATSPEC_USER, contract, compiler);
	} // end of contracts iteration

	return 0;
}
<|endoftext|>


<commit_before>#ifdef _WIN32

#include "net_rubygrapefruit_platform_internal_jni_WindowsFileEventFunctions.h"
#include "generic.h"
#include "win.h"
#include <process.h>
#include <list>
#include <vector>
#include <string>

using namespace std;

// TODO Find the right size for this
#define EVENT_BUFFER_SIZE (16*1024)

#define CREATE_SHARE (FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE)
#define CREATE_FLAGS (FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OVERLAPPED)
#define EVENT_MASK (FILE_NOTIFY_CHANGE_FILE_NAME | FILE_NOTIFY_CHANGE_DIR_NAME | \
                    FILE_NOTIFY_CHANGE_ATTRIBUTES | FILE_NOTIFY_CHANGE_SIZE | FILE_NOTIFY_CHANGE_LAST_WRITE)

class Server;
class WatchPoint;

class WatchPoint {
public:
    WatchPoint(Server *server, wstring path) {
        this->server = server;
        this->path = path;
        this->buffer.resize(EVENT_BUFFER_SIZE);
        ZeroMemory(&this->overlapped, sizeof(OVERLAPPED));
        this->overlapped.hEvent = this;
        this->directoryHandle = CreateFileW(
            path.c_str(),                       // pointer to the file name
            FILE_LIST_DIRECTORY,                // access (read/write) mode
            CREATE_SHARE,                       // share mode
            NULL,                               // security descriptor
            OPEN_EXISTING,                      // how to create
            CREATE_FLAGS,                       // file attributes
            NULL                                // file with attributes to copy
        );

        // TODO Error handling
    }

    void close();
    void listen();

private:
    Server *server;
    friend static void CALLBACK startWatchCallback(_In_ ULONG_PTR arg);
    wstring path;
    HANDLE directoryHandle;
    OVERLAPPED overlapped;
    vector<BYTE> buffer;

    void handleEvent(DWORD errorCode, DWORD bytesTransfered);
    void handlePathChanged(FILE_NOTIFY_INFORMATION *info);
    friend static void CALLBACK handleEventCallback(DWORD errorCode, DWORD bytesTransfered, LPOVERLAPPED overlapped);
};

class Server {
public:
    Server(
        JavaVM *jvm,
        JNIEnv *env,
        jobject watcherCallback
    ) {
        this->jvm = jvm;
        this->watcherCallback = env->NewGlobalRef(watcherCallback);
        this->threadHandle = (HANDLE)_beginthreadex(
            NULL,                   // default security attributes
            0,                      // use default stack size
            EventProcessingThread,  // thread function name
            this,                   // argument to thread function
            0,                      // use default creation flags
            NULL                    // the thread identifier
        );
        // TODO Error handling
        SetThreadPriority(this->threadHandle, THREAD_PRIORITY_ABOVE_NORMAL);
    }

    void startWatching(wchar_t *path);
    void reportEvent(jint type, const wstring changedPath);
    void reportFinished(WatchPoint* watchPoint);

    void close(JNIEnv *env);

private:
    struct StartWatchRequest {
        Server *server;
        wstring path;
    };
    friend static void CALLBACK startWatchCallback(_In_ ULONG_PTR arg);

    JavaVM *jvm;
    list<WatchPoint*> watchPoints;
    jobject watcherCallback;

    HANDLE threadHandle;
    bool terminate = false;

    friend static void CALLBACK requestTerminationCallback(_In_ ULONG_PTR arg);
    void requestTermination();

    friend static unsigned CALLBACK EventProcessingThread(void* data);
    void run();
};

void WatchPoint::close() {
    CancelIo(directoryHandle);
    CloseHandle(directoryHandle);
}

void WatchPoint::listen() {
    BOOL success = ReadDirectoryChangesW(
        directoryHandle,                    // handle to directory
        &buffer[0],                         // read results buffer
        buffer.size(),                      // length of buffer
        TRUE,                               // include children
        EVENT_MASK,                         // filter conditions
        NULL,                               // bytes returned
        &overlapped,                        // overlapped buffer
        &handleEventCallback                // completion routine
    );
    // TODO Error handling
}

static void CALLBACK handleEventCallback(DWORD errorCode, DWORD bytesTransfered, LPOVERLAPPED overlapped) {
    WatchPoint* watchPoint = (WatchPoint*)overlapped->hEvent;
    watchPoint->handleEvent(errorCode, bytesTransfered);
}

void WatchPoint::handleEvent(DWORD errorCode, DWORD bytesTransfered) {
    if (errorCode == ERROR_OPERATION_ABORTED){
        server->reportFinished(this);
        delete this;
        return;
    }

    if (bytesTransfered == 0) {
        // don't send dirty too much, everything is changed anyway
        // TODO Understand what this does
        // if (WaitForSingleObject(stopEventHandle, 500) == WAIT_OBJECT_0)
        //    break;

        // Got a buffer overflow => current changes lost => send INVALIDATE on root
        server->reportEvent(FILE_EVENT_INVALIDATE, path);
    } else {
        FILE_NOTIFY_INFORMATION *info = (FILE_NOTIFY_INFORMATION *)buffer.data();
        do {
            handlePathChanged(info);
            info = (FILE_NOTIFY_INFORMATION *)((char *)info + info->NextEntryOffset);
        } while (info->NextEntryOffset != 0);
    }

    // Get the new read issued as fast as possible. The documentation
    // says that the original OVERLAPPED structure will not be used
    // again once the completion routine is called.
    listen();
}

void WatchPoint::handlePathChanged(FILE_NOTIFY_INFORMATION *info) {
    wstring changedPath = wstring(info->FileName, 0, info->FileNameLength / sizeof(wchar_t));
    if (!changedPath.empty()) {
        changedPath.insert(0, 1, L'\\');
        changedPath.insert(0, path);
    }

    log_info(details->env, L"~~~~ Changed: 0x%x %ls\n", info->Action, changedPath);

    jint type;
    if (info->Action == FILE_ACTION_ADDED || info->Action == FILE_ACTION_RENAMED_NEW_NAME) {
        type = FILE_EVENT_CREATED;
    } else if (info->Action == FILE_ACTION_REMOVED || info->Action == FILE_ACTION_RENAMED_OLD_NAME) {
        type = FILE_EVENT_REMOVED;
    } else if (info->Action == FILE_ACTION_MODIFIED) {
        type = FILE_EVENT_MODIFIED;
    } else {
        wprintf(L"~~~~ Unknown event 0x%x for %ls\n", info->Action, changedPath.c_str());
        type = FILE_EVENT_UNKNOWN;
    }

    server->reportEvent(type, changedPath);
}

void Server::startWatching(wchar_t *path) {
    WatchPoint* watchPoint = new WatchPoint(this, wstring(path));
    QueueUserAPC(&startWatchCallback, threadHandle, (ULONG_PTR) watchPoint);
}

// Called by QueueUserAPC to add another directory.
static void CALLBACK startWatchCallback(_In_ ULONG_PTR arg) {
    WatchPoint* watchPoint = (WatchPoint*)arg;
    watchPoint->server->watchPoints.push_back(watchPoint);
    watchPoint->listen();
}

void Server::reportFinished(WatchPoint* watchPoint) {
    watchPoints.remove(watchPoint);
}

static JNIEnv* getThreadEnv(JavaVM *jvm) {
    JNIEnv* env;
    // TODO Verify that JNI 1.6 is the right version
    jint ret = jvm->GetEnv((void **) &(env), JNI_VERSION_1_6);
    if (ret != JNI_OK) {
        printf("Failed to attach JNI to current thread: %d\n", ret);
        return NULL;
    }
    return env;
}

void Server::reportEvent(jint type, const wstring changedPath) {
    JNIEnv* env = getThreadEnv(jvm);
    jstring changedPathJava = wchar_to_java_path(env, changedPath.c_str());
    jclass callback_class = env->GetObjectClass(watcherCallback);
    jmethodID methodCallback = env->GetMethodID(callback_class, "pathChanged", "(ILjava/lang/String;)V");
    env->CallVoidMethod(watcherCallback, methodCallback, type, changedPathJava);
}

static unsigned CALLBACK EventProcessingThread(void* data) {
    Server *server = (Server*) data;
    server->run();
    return 0;
}

void Server::run() {
    printf("~~~~ Thread %d running\n", GetCurrentThreadId());

    // TODO Extract this logic to some shared function
    JNIEnv* env;
    jint statAttach = jvm->AttachCurrentThreadAsDaemon((void **) &(env), NULL);
    if (statAttach != JNI_OK) {
        printf("Failed to attach JNI to current thread: %d\n", statAttach);
        return;
    }

    while (!terminate || watchPoints.size() > 0) {
        SleepEx(INFINITE, true);
    }

    printf("~~~~ Stopping thread\n");

    // TODO Extract this logic to some shared function
    jint statDetach = jvm->DetachCurrentThread();
    if (statDetach != JNI_OK) {
        printf("Failed to detach JNI from current thread: %d\n", statAttach);
        return;
    }
}

static void CALLBACK requestTerminationCallback(_In_ ULONG_PTR arg) {
    Server* server = (Server*)arg;
    server->requestTermination();
}

void Server::requestTermination() {
    terminate = true;
    for (auto &watchPoint : watchPoints) {
        watchPoint->close();
    }
}

void Server::close(JNIEnv *env) {
    QueueUserAPC(requestTerminationCallback, this->threadHandle, (ULONG_PTR)this);
    WaitForSingleObject(this->threadHandle, INFINITE);
    CloseHandle(this->threadHandle);
    env->DeleteGlobalRef(this->watcherCallback);
}

JNIEXPORT jobject JNICALL
Java_net_rubygrapefruit_platform_internal_jni_WindowsFileEventFunctions_startWatching(JNIEnv *env, jclass target, jobjectArray paths, jobject javaCallback, jobject result) {

    printf("\n~~~~ Configuring...\n");

    JavaVM* jvm;
    int jvmStatus = env->GetJavaVM(&jvm);
    if (jvmStatus != JNI_OK) {
        mark_failed_with_errno(env, "Could not store jvm instance.", result);
        return NULL;
    }

    int watchPointCount = env->GetArrayLength(paths);
    if (watchPointCount == 0) {
        mark_failed_with_errno(env, "No paths given to watch.", result);
        return NULL;
    }

    Server* server = new Server(jvm, env, javaCallback);

    for (int i = 0; i < watchPointCount; i++) {
        jstring path = (jstring) env->GetObjectArrayElement(paths, i);
        wchar_t* watchPoint = java_to_wchar_path(env, path);
        int watchPointLen = wcslen(watchPoint);
        wprintf(L"~~~~ Watching %ls\n", watchPoint);
        server->startWatching(watchPoint);
        free(watchPoint);
    }

    jclass clsWatch = env->FindClass("net/rubygrapefruit/platform/internal/jni/WindowsFileEventFunctions$WatcherImpl");
    jmethodID constructor = env->GetMethodID(clsWatch, "<init>", "(Ljava/lang/Object;)V");
    return env->NewObject(clsWatch, constructor, env->NewDirectByteBuffer(server, sizeof(server)));
}

JNIEXPORT void JNICALL
Java_net_rubygrapefruit_platform_internal_jni_WindowsFileEventFunctions_stopWatching(JNIEnv *env, jclass target, jobject detailsObj, jobject result) {
    Server* server = (Server*)env->GetDirectBufferAddress(detailsObj);
    server->close(env);
    delete server;
}

#endif
<commit_msg>Make copy of watcher list when closing server<commit_after>#ifdef _WIN32

#include "net_rubygrapefruit_platform_internal_jni_WindowsFileEventFunctions.h"
#include "generic.h"
#include "win.h"
#include <process.h>
#include <list>
#include <vector>
#include <string>

using namespace std;

// TODO Find the right size for this
#define EVENT_BUFFER_SIZE (16*1024)

#define CREATE_SHARE (FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE)
#define CREATE_FLAGS (FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OVERLAPPED)
#define EVENT_MASK (FILE_NOTIFY_CHANGE_FILE_NAME | FILE_NOTIFY_CHANGE_DIR_NAME | \
                    FILE_NOTIFY_CHANGE_ATTRIBUTES | FILE_NOTIFY_CHANGE_SIZE | FILE_NOTIFY_CHANGE_LAST_WRITE)

class Server;
class WatchPoint;

class WatchPoint {
public:
    WatchPoint(Server *server, wstring path) {
        this->server = server;
        this->path = path;
        this->buffer.resize(EVENT_BUFFER_SIZE);
        ZeroMemory(&this->overlapped, sizeof(OVERLAPPED));
        this->overlapped.hEvent = this;
        this->directoryHandle = CreateFileW(
            path.c_str(),                       // pointer to the file name
            FILE_LIST_DIRECTORY,                // access (read/write) mode
            CREATE_SHARE,                       // share mode
            NULL,                               // security descriptor
            OPEN_EXISTING,                      // how to create
            CREATE_FLAGS,                       // file attributes
            NULL                                // file with attributes to copy
        );

        // TODO Error handling
    }

    void close();
    void listen();

private:
    Server *server;
    friend static void CALLBACK startWatchCallback(_In_ ULONG_PTR arg);
    wstring path;
    HANDLE directoryHandle;
    OVERLAPPED overlapped;
    vector<BYTE> buffer;

    void handleEvent(DWORD errorCode, DWORD bytesTransfered);
    void handlePathChanged(FILE_NOTIFY_INFORMATION *info);
    friend static void CALLBACK handleEventCallback(DWORD errorCode, DWORD bytesTransfered, LPOVERLAPPED overlapped);
};

class Server {
public:
    Server(
        JavaVM *jvm,
        JNIEnv *env,
        jobject watcherCallback
    ) {
        this->jvm = jvm;
        this->watcherCallback = env->NewGlobalRef(watcherCallback);
        this->threadHandle = (HANDLE)_beginthreadex(
            NULL,                   // default security attributes
            0,                      // use default stack size
            EventProcessingThread,  // thread function name
            this,                   // argument to thread function
            0,                      // use default creation flags
            NULL                    // the thread identifier
        );
        // TODO Error handling
        SetThreadPriority(this->threadHandle, THREAD_PRIORITY_ABOVE_NORMAL);
    }

    void startWatching(wchar_t *path);
    void reportEvent(jint type, const wstring changedPath);
    void reportFinished(WatchPoint* watchPoint);

    void close(JNIEnv *env);

private:
    struct StartWatchRequest {
        Server *server;
        wstring path;
    };
    friend static void CALLBACK startWatchCallback(_In_ ULONG_PTR arg);

    JavaVM *jvm;
    list<WatchPoint*> watchPoints;
    jobject watcherCallback;

    HANDLE threadHandle;
    bool terminate = false;

    friend static void CALLBACK requestTerminationCallback(_In_ ULONG_PTR arg);
    void requestTermination();

    friend static unsigned CALLBACK EventProcessingThread(void* data);
    void run();
};

void WatchPoint::close() {
    CancelIo(directoryHandle);
    CloseHandle(directoryHandle);
}

void WatchPoint::listen() {
    BOOL success = ReadDirectoryChangesW(
        directoryHandle,                    // handle to directory
        &buffer[0],                         // read results buffer
        buffer.size(),                      // length of buffer
        TRUE,                               // include children
        EVENT_MASK,                         // filter conditions
        NULL,                               // bytes returned
        &overlapped,                        // overlapped buffer
        &handleEventCallback                // completion routine
    );
    // TODO Error handling
}

static void CALLBACK handleEventCallback(DWORD errorCode, DWORD bytesTransfered, LPOVERLAPPED overlapped) {
    WatchPoint* watchPoint = (WatchPoint*)overlapped->hEvent;
    watchPoint->handleEvent(errorCode, bytesTransfered);
}

void WatchPoint::handleEvent(DWORD errorCode, DWORD bytesTransfered) {
    if (errorCode == ERROR_OPERATION_ABORTED){
        server->reportFinished(this);
        delete this;
        return;
    }

    if (bytesTransfered == 0) {
        // don't send dirty too much, everything is changed anyway
        // TODO Understand what this does
        // if (WaitForSingleObject(stopEventHandle, 500) == WAIT_OBJECT_0)
        //    break;

        // Got a buffer overflow => current changes lost => send INVALIDATE on root
        server->reportEvent(FILE_EVENT_INVALIDATE, path);
    } else {
        FILE_NOTIFY_INFORMATION *info = (FILE_NOTIFY_INFORMATION *)buffer.data();
        do {
            handlePathChanged(info);
            info = (FILE_NOTIFY_INFORMATION *)((char *)info + info->NextEntryOffset);
        } while (info->NextEntryOffset != 0);
    }

    // Get the new read issued as fast as possible. The documentation
    // says that the original OVERLAPPED structure will not be used
    // again once the completion routine is called.
    listen();
}

void WatchPoint::handlePathChanged(FILE_NOTIFY_INFORMATION *info) {
    wstring changedPath = wstring(info->FileName, 0, info->FileNameLength / sizeof(wchar_t));
    if (!changedPath.empty()) {
        changedPath.insert(0, 1, L'\\');
        changedPath.insert(0, path);
    }

    log_info(details->env, L"~~~~ Changed: 0x%x %ls\n", info->Action, changedPath);

    jint type;
    if (info->Action == FILE_ACTION_ADDED || info->Action == FILE_ACTION_RENAMED_NEW_NAME) {
        type = FILE_EVENT_CREATED;
    } else if (info->Action == FILE_ACTION_REMOVED || info->Action == FILE_ACTION_RENAMED_OLD_NAME) {
        type = FILE_EVENT_REMOVED;
    } else if (info->Action == FILE_ACTION_MODIFIED) {
        type = FILE_EVENT_MODIFIED;
    } else {
        wprintf(L"~~~~ Unknown event 0x%x for %ls\n", info->Action, changedPath.c_str());
        type = FILE_EVENT_UNKNOWN;
    }

    server->reportEvent(type, changedPath);
}

void Server::startWatching(wchar_t *path) {
    WatchPoint* watchPoint = new WatchPoint(this, wstring(path));
    QueueUserAPC(&startWatchCallback, threadHandle, (ULONG_PTR) watchPoint);
}

// Called by QueueUserAPC to add another directory.
static void CALLBACK startWatchCallback(_In_ ULONG_PTR arg) {
    WatchPoint* watchPoint = (WatchPoint*)arg;
    watchPoint->server->watchPoints.push_back(watchPoint);
    watchPoint->listen();
}

void Server::reportFinished(WatchPoint* watchPoint) {
    watchPoints.remove(watchPoint);
}

static JNIEnv* getThreadEnv(JavaVM *jvm) {
    JNIEnv* env;
    // TODO Verify that JNI 1.6 is the right version
    jint ret = jvm->GetEnv((void **) &(env), JNI_VERSION_1_6);
    if (ret != JNI_OK) {
        printf("Failed to attach JNI to current thread: %d\n", ret);
        return NULL;
    }
    return env;
}

void Server::reportEvent(jint type, const wstring changedPath) {
    JNIEnv* env = getThreadEnv(jvm);
    jstring changedPathJava = wchar_to_java_path(env, changedPath.c_str());
    jclass callback_class = env->GetObjectClass(watcherCallback);
    jmethodID methodCallback = env->GetMethodID(callback_class, "pathChanged", "(ILjava/lang/String;)V");
    env->CallVoidMethod(watcherCallback, methodCallback, type, changedPathJava);
}

static unsigned CALLBACK EventProcessingThread(void* data) {
    Server *server = (Server*) data;
    server->run();
    return 0;
}

void Server::run() {
    printf("~~~~ Thread %d running\n", GetCurrentThreadId());

    // TODO Extract this logic to some shared function
    JNIEnv* env;
    jint statAttach = jvm->AttachCurrentThreadAsDaemon((void **) &(env), NULL);
    if (statAttach != JNI_OK) {
        printf("Failed to attach JNI to current thread: %d\n", statAttach);
        return;
    }

    while (!terminate || watchPoints.size() > 0) {
        SleepEx(INFINITE, true);
    }

    printf("~~~~ Stopping thread\n");

    // TODO Extract this logic to some shared function
    jint statDetach = jvm->DetachCurrentThread();
    if (statDetach != JNI_OK) {
        printf("Failed to detach JNI from current thread: %d\n", statAttach);
        return;
    }
}

static void CALLBACK requestTerminationCallback(_In_ ULONG_PTR arg) {
    Server* server = (Server*)arg;
    server->requestTermination();
}

void Server::requestTermination() {
    terminate = true;
    // Make copy so terminated entries can be removed
    list<WatchPoint*> copyWatchPoints(watchPoints);
    for (auto &watchPoint : copyWatchPoints) {
        watchPoint->close();
    }
}

void Server::close(JNIEnv *env) {
    QueueUserAPC(requestTerminationCallback, this->threadHandle, (ULONG_PTR)this);
    WaitForSingleObject(this->threadHandle, INFINITE);
    CloseHandle(this->threadHandle);
    env->DeleteGlobalRef(this->watcherCallback);
}

JNIEXPORT jobject JNICALL
Java_net_rubygrapefruit_platform_internal_jni_WindowsFileEventFunctions_startWatching(JNIEnv *env, jclass target, jobjectArray paths, jobject javaCallback, jobject result) {

    printf("\n~~~~ Configuring...\n");

    JavaVM* jvm;
    int jvmStatus = env->GetJavaVM(&jvm);
    if (jvmStatus != JNI_OK) {
        mark_failed_with_errno(env, "Could not store jvm instance.", result);
        return NULL;
    }

    int watchPointCount = env->GetArrayLength(paths);
    if (watchPointCount == 0) {
        mark_failed_with_errno(env, "No paths given to watch.", result);
        return NULL;
    }

    Server* server = new Server(jvm, env, javaCallback);

    for (int i = 0; i < watchPointCount; i++) {
        jstring path = (jstring) env->GetObjectArrayElement(paths, i);
        wchar_t* watchPoint = java_to_wchar_path(env, path);
        int watchPointLen = wcslen(watchPoint);
        wprintf(L"~~~~ Watching %ls\n", watchPoint);
        server->startWatching(watchPoint);
        free(watchPoint);
    }

    jclass clsWatch = env->FindClass("net/rubygrapefruit/platform/internal/jni/WindowsFileEventFunctions$WatcherImpl");
    jmethodID constructor = env->GetMethodID(clsWatch, "<init>", "(Ljava/lang/Object;)V");
    return env->NewObject(clsWatch, constructor, env->NewDirectByteBuffer(server, sizeof(server)));
}

JNIEXPORT void JNICALL
Java_net_rubygrapefruit_platform_internal_jni_WindowsFileEventFunctions_stopWatching(JNIEnv *env, jclass target, jobject detailsObj, jobject result) {
    Server* server = (Server*)env->GetDirectBufferAddress(detailsObj);
    server->close(env);
    delete server;
}

#endif
<|endoftext|>


<commit_before>/// \file Util.cpp
//-----------------------------------------------------------------------------

#include "sha1.h"
#include "BlowFish.h"

#include <fcntl.h>
#include <errno.h>
#include <io.h>
#include <sys\stat.h>

#include "Util.h"

// for now
#define String CString
#define SecString CMyString

// used by CBC routines...
static void
xormem(unsigned char* mem1, unsigned char* mem2, int length)
{
   for (int x=0;x<length;x++)
      mem1[x] ^= mem2[x];
}

//-----------------------------------------------------------------------------
/*
  Note: A bunch of the encryption-related routines may not be Unicode
  compliant.  Which really isn't a huge problem, since they are actually
  using MFC routines anyway
*/

//-----------------------------------------------------------------------------
//Overwrite the memory

void
trashMemory(SHA1_CTX& context)
{
   trashMemory((unsigned char*)context.state, sizeof context.state);
   trashMemory((unsigned char*)context.count, sizeof context.count);
   trashMemory((unsigned char*)context.buffer, sizeof context.buffer);
}

void
trashMemory(unsigned char* buffer,
            long length,
            int numiter) // default 30
{
   for (int x=0; x<numiter; x++)
   {
      memset(buffer, 0x00, length);
      memset(buffer, 0xFF, length);
      memset(buffer, 0x00, length);
   }
}

//Generates a passkey-based hash from stuff - used to validate the passkey
void
GenRandhash(const CMyString &a_passkey,
            const unsigned char* a_randstuff,
            unsigned char* a_randhash)
{
  const LPCSTR pkey = (const LPCSTR)a_passkey;
   /*
     I'm not quite sure what this is doing, so as I figure out each piece,
     I'll add more comments {jpr}
   */

   /*
     tempSalt <- H(a_randstuff + a_passkey)
   */
   SHA1_CTX keyHash;
   SHA1Init(&keyHash);
   SHA1Update(&keyHash, a_randstuff, StuffSize);
   SHA1Update(&keyHash,
              (const unsigned char*)pkey,
              a_passkey.GetLength());

   unsigned char tempSalt[20]; // HashSize
   SHA1Final(tempSalt, &keyHash);

   /*
     tempbuf <- a_randstuff encrypted 1000 times using tempSalt as key?
   */

   BlowFish Cipher(tempSalt, sizeof(tempSalt));

   unsigned char tempbuf[StuffSize];
   memcpy((char*)tempbuf, (char*)a_randstuff, StuffSize);

   for (int x=0; x<1000; x++)
      Cipher.Encrypt(tempbuf, tempbuf);

   /*
     hmm - seems we're not done with this context
     we throw the tempbuf into the hasher, and extract a_randhash
   */
   SHA1Update(&keyHash, tempbuf, StuffSize);
   SHA1Final(a_randhash, &keyHash);
   trashMemory(keyHash);
}

unsigned char
newrand()
{
   int	r;
   while ((r = rand()) % 257 == 256)
      ; // 257?!?
   return (unsigned char)r;
}

/*
 *  RangeRand(len)
 *
 *  Returns a random number in the range 0 to (len-1).
 *  For example, RangeRand(256) returns a value from 0 to 255.
 */
unsigned int
RangeRand(size_t len)
{
   unsigned int      r;
   unsigned int      ceil = UINT_MAX - (UINT_MAX % len) - 1;

   while ((r = rand()) > ceil)
      ;
   return(r%len);
}

/*
 * Returns a BlowFish object set up for encryption or decrytion.
 *
 * The main issue here is that the BlowFish key is SHA1(passphrase|salt)
 * Aside from saving duplicate code, we win here by minimizing the exposure
 * of the actual key.
 * The lose is that the BlowFish object is now dynamically allocated.
 * This could be fixed by having a ctor of BlowFish that works without a key,
 * which would be set by another member function, but I doubt that it's worth the bother.
 *
 * Note that it's the caller's responsibility to delete the BlowFish object allocated here
 */

BlowFish *MakeBlowFish(const unsigned char *pass, int passlen,
		       const unsigned char *salt, int saltlen)
{
   unsigned char passkey[20]; // SHA1 digest is 20 bytes - why isn't there a constant for this?
   VirtualLock(passkey, sizeof(passkey));

   SHA1_CTX context;
   SHA1Init(&context);
   SHA1Update(&context, pass, passlen);
   SHA1Update(&context, salt, saltlen);
   SHA1Final(passkey, &context);
   BlowFish *retval = new BlowFish(passkey, sizeof(passkey));
   trashMemory(passkey, sizeof(passkey));
   trashMemory(context);
   VirtualUnlock(passkey, sizeof(passkey));

   return retval;
}

int
_writecbc(int fp,
          const unsigned char* buffer,
          int length, unsigned char type,
	  const unsigned char *pass, int passlen,
          const unsigned char* salt, int saltlen,
          unsigned char* cbcbuffer)
{
   int numWritten = 0;

   int BlockLength = ((length+7)/8)*8;
   if (BlockLength == 0)
      BlockLength = 8;

   BlowFish *Algorithm = MakeBlowFish(pass, passlen, salt, saltlen);

   // First encrypt and write the length of the buffer
   unsigned char lengthblock[8];
   memset(lengthblock, 0, 8);
   // XXX next line is a portability issue - what if file is read by a program
   // compiled with a different sizeof int or different endian-ness?
   memcpy(lengthblock, (unsigned char*)&length, sizeof length);

   // following new for format 2.0 - lengthblock bytes 4-7 were unused before.
   lengthblock[sizeof(length)] = type;

   xormem(lengthblock, cbcbuffer, 8); // do the CBC thing
   Algorithm->Encrypt(lengthblock, lengthblock);
   memcpy(cbcbuffer, lengthblock, 8); // update CBC for next round
   numWritten = _write(fp, lengthblock, 8);
   trashMemory(lengthblock, 8);

   // Now, encrypt and write the buffer
   unsigned char curblock[8];
   for (int x=0;x<BlockLength;x+=8)
   {
      if ((length == 0) || ((length%8 != 0) && (length-x<8)))
      {
         //This is for an uneven last block
         memset(curblock, 0, 8);
         memcpy(curblock, buffer+x, length % 8);
      }
      else
         memcpy(curblock, buffer+x, 8);
      xormem(curblock, cbcbuffer, 8);
      Algorithm->Encrypt(curblock, curblock);
      memcpy(cbcbuffer, curblock, 8);
      numWritten += _write(fp, curblock, 8);
   }
   trashMemory(curblock, 8);
   delete Algorithm;
   return numWritten;
}

/*
 * Reads an encrypted record into buffer.
 * The first block of the record contains the encrypted record length
 * We have the usual ugly problem of fixed buffer lengths in C/C++.
 * Don't want to use CStrings, for future portability.
 * Best solution would be STL strings, but not enough experience with them (yet)
 * So for the meantime, we're going to allocate the buffer here, to ensure that it's long
 * enough.
 * *** THE CALLER MUST delete[] IT AFTER USE *** UGH++
 *
 * (unless buffer_len is zero)
 *
 * An alternative to STL strings would be to accept a buffer, and allocate a replacement
 * iff the buffer is too small. This is serious ugliness, but should be considered if
 * the new/delete performance hit is too big.
 */
int
_readcbc(int fp,
         unsigned char* &buffer, unsigned int &buffer_len, unsigned char &type,
	 const unsigned char *pass, int passlen,
         const unsigned char* salt, int saltlen,
         unsigned char* cbcbuffer)
{
   int numRead = 0;

   unsigned char lengthblock[8];
   unsigned char lcpy[8];

   buffer_len = 0;
   numRead = _read(fp, lengthblock, 8);
   if (numRead != 8)
      return 0;
   memcpy(lcpy, lengthblock, 8);

   BlowFish *Algorithm = MakeBlowFish(pass, passlen, salt, saltlen);

   Algorithm->Decrypt(lengthblock, lengthblock);
   xormem(lengthblock, cbcbuffer, 8);
   memcpy(cbcbuffer, lcpy, 8);

   // portability issue - see comment in _writecbc
   int length = *((int*)lengthblock);

   // new for 2.0 -- lengthblock[4..7] previously set to zero
   type = lengthblock[sizeof(int)]; // type is first byte after the length

   trashMemory(lengthblock, 8);
   trashMemory(lcpy, 8);

   if (length < 0) {
     delete Algorithm;
     buffer = NULL;
     buffer_len = 0;
     return 0;
   }

   int BlockLength = ((length+7)/8)*8;
   // Following is meant for lengths < 8,
   // but results in a block being read even
   // if length is zero. This is wasteful,
   // but fixing it would break all existing databases
   if (BlockLength == 0)
      BlockLength = 8;

   buffer_len = length;
   buffer = new unsigned char[BlockLength]; // so we lie a little...

   unsigned char tempcbc[8];
   numRead += _read(fp, buffer, BlockLength);
   for (int x=0;x<BlockLength;x+=8)
   {
      memcpy(tempcbc, buffer+x, 8);
      Algorithm->Decrypt(buffer+x, buffer+x);
      xormem(buffer+x, cbcbuffer, 8);
      memcpy(cbcbuffer, tempcbc, 8);
   }

   trashMemory(tempcbc, 8);
   delete Algorithm;
   if (length == 0) {
     // delete[] buffer here since caller will see zero length
     delete[] buffer;
   }

   return numRead;
}
<commit_msg>removed unneeded #includes<commit_after>/// \file Util.cpp
//-----------------------------------------------------------------------------

#include "sha1.h"
#include "BlowFish.h"

#include <io.h>

#include "Util.h"

// for now
#define String CString
#define SecString CMyString

// used by CBC routines...
static void
xormem(unsigned char* mem1, unsigned char* mem2, int length)
{
   for (int x=0;x<length;x++)
      mem1[x] ^= mem2[x];
}

//-----------------------------------------------------------------------------
/*
  Note: A bunch of the encryption-related routines may not be Unicode
  compliant.  Which really isn't a huge problem, since they are actually
  using MFC routines anyway
*/

//-----------------------------------------------------------------------------
//Overwrite the memory

void
trashMemory(SHA1_CTX& context)
{
   trashMemory((unsigned char*)context.state, sizeof context.state);
   trashMemory((unsigned char*)context.count, sizeof context.count);
   trashMemory((unsigned char*)context.buffer, sizeof context.buffer);
}

void
trashMemory(unsigned char* buffer,
            long length,
            int numiter) // default 30
{
   for (int x=0; x<numiter; x++)
   {
      memset(buffer, 0x00, length);
      memset(buffer, 0xFF, length);
      memset(buffer, 0x00, length);
   }
}

//Generates a passkey-based hash from stuff - used to validate the passkey
void
GenRandhash(const CMyString &a_passkey,
            const unsigned char* a_randstuff,
            unsigned char* a_randhash)
{
  const LPCSTR pkey = (const LPCSTR)a_passkey;
   /*
     I'm not quite sure what this is doing, so as I figure out each piece,
     I'll add more comments {jpr}
   */

   /*
     tempSalt <- H(a_randstuff + a_passkey)
   */
   SHA1_CTX keyHash;
   SHA1Init(&keyHash);
   SHA1Update(&keyHash, a_randstuff, StuffSize);
   SHA1Update(&keyHash,
              (const unsigned char*)pkey,
              a_passkey.GetLength());

   unsigned char tempSalt[20]; // HashSize
   SHA1Final(tempSalt, &keyHash);

   /*
     tempbuf <- a_randstuff encrypted 1000 times using tempSalt as key?
   */

   BlowFish Cipher(tempSalt, sizeof(tempSalt));

   unsigned char tempbuf[StuffSize];
   memcpy((char*)tempbuf, (char*)a_randstuff, StuffSize);

   for (int x=0; x<1000; x++)
      Cipher.Encrypt(tempbuf, tempbuf);

   /*
     hmm - seems we're not done with this context
     we throw the tempbuf into the hasher, and extract a_randhash
   */
   SHA1Update(&keyHash, tempbuf, StuffSize);
   SHA1Final(a_randhash, &keyHash);
   trashMemory(keyHash);
}

unsigned char
newrand()
{
   int	r;
   while ((r = rand()) % 257 == 256)
      ; // 257?!?
   return (unsigned char)r;
}

/*
 *  RangeRand(len)
 *
 *  Returns a random number in the range 0 to (len-1).
 *  For example, RangeRand(256) returns a value from 0 to 255.
 */
unsigned int
RangeRand(size_t len)
{
   unsigned int      r;
   unsigned int      ceil = UINT_MAX - (UINT_MAX % len) - 1;

   while ((r = rand()) > ceil)
      ;
   return(r%len);
}

/*
 * Returns a BlowFish object set up for encryption or decrytion.
 *
 * The main issue here is that the BlowFish key is SHA1(passphrase|salt)
 * Aside from saving duplicate code, we win here by minimizing the exposure
 * of the actual key.
 * The lose is that the BlowFish object is now dynamically allocated.
 * This could be fixed by having a ctor of BlowFish that works without a key,
 * which would be set by another member function, but I doubt that it's worth the bother.
 *
 * Note that it's the caller's responsibility to delete the BlowFish object allocated here
 */

BlowFish *MakeBlowFish(const unsigned char *pass, int passlen,
		       const unsigned char *salt, int saltlen)
{
   unsigned char passkey[20]; // SHA1 digest is 20 bytes - why isn't there a constant for this?
   VirtualLock(passkey, sizeof(passkey));

   SHA1_CTX context;
   SHA1Init(&context);
   SHA1Update(&context, pass, passlen);
   SHA1Update(&context, salt, saltlen);
   SHA1Final(passkey, &context);
   BlowFish *retval = new BlowFish(passkey, sizeof(passkey));
   trashMemory(passkey, sizeof(passkey));
   trashMemory(context);
   VirtualUnlock(passkey, sizeof(passkey));

   return retval;
}

int
_writecbc(int fp,
          const unsigned char* buffer,
          int length, unsigned char type,
	  const unsigned char *pass, int passlen,
          const unsigned char* salt, int saltlen,
          unsigned char* cbcbuffer)
{
   int numWritten = 0;

   int BlockLength = ((length+7)/8)*8;
   if (BlockLength == 0)
      BlockLength = 8;

   BlowFish *Algorithm = MakeBlowFish(pass, passlen, salt, saltlen);

   // First encrypt and write the length of the buffer
   unsigned char lengthblock[8];
   memset(lengthblock, 0, 8);
   // XXX next line is a portability issue - what if file is read by a program
   // compiled with a different sizeof int or different endian-ness?
   memcpy(lengthblock, (unsigned char*)&length, sizeof length);

   // following new for format 2.0 - lengthblock bytes 4-7 were unused before.
   lengthblock[sizeof(length)] = type;

   xormem(lengthblock, cbcbuffer, 8); // do the CBC thing
   Algorithm->Encrypt(lengthblock, lengthblock);
   memcpy(cbcbuffer, lengthblock, 8); // update CBC for next round
   numWritten = _write(fp, lengthblock, 8);
   trashMemory(lengthblock, 8);

   // Now, encrypt and write the buffer
   unsigned char curblock[8];
   for (int x=0;x<BlockLength;x+=8)
   {
      if ((length == 0) || ((length%8 != 0) && (length-x<8)))
      {
         //This is for an uneven last block
         memset(curblock, 0, 8);
         memcpy(curblock, buffer+x, length % 8);
      }
      else
         memcpy(curblock, buffer+x, 8);
      xormem(curblock, cbcbuffer, 8);
      Algorithm->Encrypt(curblock, curblock);
      memcpy(cbcbuffer, curblock, 8);
      numWritten += _write(fp, curblock, 8);
   }
   trashMemory(curblock, 8);
   delete Algorithm;
   return numWritten;
}

/*
 * Reads an encrypted record into buffer.
 * The first block of the record contains the encrypted record length
 * We have the usual ugly problem of fixed buffer lengths in C/C++.
 * Don't want to use CStrings, for future portability.
 * Best solution would be STL strings, but not enough experience with them (yet)
 * So for the meantime, we're going to allocate the buffer here, to ensure that it's long
 * enough.
 * *** THE CALLER MUST delete[] IT AFTER USE *** UGH++
 *
 * (unless buffer_len is zero)
 *
 * An alternative to STL strings would be to accept a buffer, and allocate a replacement
 * iff the buffer is too small. This is serious ugliness, but should be considered if
 * the new/delete performance hit is too big.
 */
int
_readcbc(int fp,
         unsigned char* &buffer, unsigned int &buffer_len, unsigned char &type,
	 const unsigned char *pass, int passlen,
         const unsigned char* salt, int saltlen,
         unsigned char* cbcbuffer)
{
   int numRead = 0;

   unsigned char lengthblock[8];
   unsigned char lcpy[8];

   buffer_len = 0;
   numRead = _read(fp, lengthblock, 8);
   if (numRead != 8)
      return 0;
   memcpy(lcpy, lengthblock, 8);

   BlowFish *Algorithm = MakeBlowFish(pass, passlen, salt, saltlen);

   Algorithm->Decrypt(lengthblock, lengthblock);
   xormem(lengthblock, cbcbuffer, 8);
   memcpy(cbcbuffer, lcpy, 8);

   // portability issue - see comment in _writecbc
   int length = *((int*)lengthblock);

   // new for 2.0 -- lengthblock[4..7] previously set to zero
   type = lengthblock[sizeof(int)]; // type is first byte after the length

   trashMemory(lengthblock, 8);
   trashMemory(lcpy, 8);

   if (length < 0) {
     delete Algorithm;
     buffer = NULL;
     buffer_len = 0;
     return 0;
   }

   int BlockLength = ((length+7)/8)*8;
   // Following is meant for lengths < 8,
   // but results in a block being read even
   // if length is zero. This is wasteful,
   // but fixing it would break all existing databases
   if (BlockLength == 0)
      BlockLength = 8;

   buffer_len = length;
   buffer = new unsigned char[BlockLength]; // so we lie a little...

   unsigned char tempcbc[8];
   numRead += _read(fp, buffer, BlockLength);
   for (int x=0;x<BlockLength;x+=8)
   {
      memcpy(tempcbc, buffer+x, 8);
      Algorithm->Decrypt(buffer+x, buffer+x);
      xormem(buffer+x, cbcbuffer, 8);
      memcpy(cbcbuffer, tempcbc, 8);
   }

   trashMemory(tempcbc, 8);
   delete Algorithm;
   if (length == 0) {
     // delete[] buffer here since caller will see zero length
     delete[] buffer;
   }

   return numRead;
}
<|endoftext|>