<commit_before>// from lib/THC/THCTensorMasked.cu:

#if defined(__APPLE__) || defined(__MACOSX)
#include <OpenCL/cl.h>
#else
#include <CL/cl.h>
#endif

#include <boost/compute/core.hpp>
#include <boost/compute/iterator/buffer_iterator.hpp>
#include <boost/compute/algorithm/reverse.hpp>
#include <boost/compute/algorithm/transform_if.hpp>
#include <boost/compute/lambda.hpp>
#include <boost/compute/iterator/zip_iterator.hpp>

#include "THClTensorMath.h"
#include "THClGeneral.h"
#include "THClBlas.h"
#include "THClTensorCopy.h"
//#include "THClTensorRandom.h"
#include "THClApply.h"
#include "THClReduce.h"

#include <iostream>
using namespace std;

// The largest consecutive integer representable in float32 (2^24)
#define FLOAT32_MAX_CONSECUTIVE_INT 16777216.0f

class TensorMaskedFillOp : public HasOperator2, public HasScalars {
public:
  int getNumScalars() const { return 1; }
  float getScalar( int index ) const { return value; }
  TensorMaskedFillOp(float v) : value(v) {}
  std::string operator2() const {
    return "if( *in1 != 0.0f ) { *out = val1; }";
  }
  float value;
};

class TensorMaskedCopyOp : public HasOperator2, public HasGlobalTensors {
public:
  TensorMaskedCopyOp(THClTensor *src, THClTensor *baseMask, THClTensor *maskPrefixSum) :
    src(src), baseMask(baseMask), maskPrefixSum(maskPrefixSum) {
  }
  int getNumGlobalTensors() const { return 3; }
  THClTensor *getTensor(int index) const {
    if(index == 0){ return src; }
    if(index == 1){ return baseMask; }
    if(index == 2){ return maskPrefixSum; }
    THError("index not recognized %i", index);
    return 0;
  }
  std::string getTensorName(int index) const {
    if(index == 0){ return "src"; }
    if(index == 1){ return "baseMask"; }
    if(index == 2){ return "maskPrefixSum"; }
    THError("index not recognized %i", index);
    return "";
  }
  std::string operator2() const {
    return "if( *in1 != 0.0f ) { *out = src[(int)maskPrefixSum[srcOffset] ]; }";
  }
  THClTensor *src;
  THClTensor *baseMask;
  THClTensor *maskPrefixSum;
};

//struct TensorMaskedCopyOp {
//  TensorMaskedCopyOp(float* s, float* bm, float* ps)
//      : src(s),
//        baseMask(bm),
//        maskPrefixSum(ps) {
//  }

//  /*__device__*/ /*__forceline__*/ void operator()(float* out, float* mask) {
//    // Really mask should be `0` or `1` but we can't propagate errors here.
//    if (*mask != 0.0f) {
//      // We've already checked that this offset is <= 2^24, so this is ok.
//      int srcOffset = (int) (mask - baseMask);
//      *out = src[(int) maskPrefixSum[srcOffset]];
//    }
//  }

//  // Where we are copying from
//  float* src;

//  // The base address of mask so we can calculate offset
//  float* baseMask;

//  // The index we are copying from
//  float* maskPrefixSum;
//};

//class TensorMaskedSelectOp : public HasOperator3, public HasScalars {
//public:
//  int getNumScalars() const { return 1; }
//  string operator3() const {
//    return "if(*out != 0.0f){out[(int)*in1] = *in2; }";
//  }
//  TensorMaskedSelectOp(float* t) : out(t) {}
//  void operator()(float* mask, float* maskPrefixSum, float* in) {
//    // Really mask should be `0` or `1` but we can't propagate errors here.
//    if (*mask != 0.0f) {
//      out[(int) *maskPrefixSum] = *in;
//    }
//  }

//  float* out;
//};

void THClTensor_maskedFill(THClState* state,
                             THClTensor *tensor, THClTensor *mask, float value)
{
  THAssert(THClTensor_checkGPU(state, 2, tensor, mask));
  THArgCheck(THClTensor_nElement(state, tensor) ==
             THClTensor_nElement(state, mask),
             2, "sizes do not match");

  TensorMaskedFillOp op(value);
  if (!THClTensor_pointwiseApply2(state, tensor, mask, &op)) {
    THArgCheck(false, 2, CLTORCH_DIM_WARNING);
  }
}

void THClTensor_maskedCopy(THClState* state,
                             THClTensor *tensor, THClTensor *mask, THClTensor *src)
{
  THAssert(THClTensor_checkGPU(state, 3, tensor, src, mask));
  long maskSize = THClTensor_nElement(state, mask);
  long tensorSize = THClTensor_nElement(state, tensor);
  long srcSize = THClTensor_nElement(state, src);

  // Since we are performing a prefix sum of mask, it cannot exceed
  // the size allowed in consecutive integers in float32
  THArgCheck(maskSize <= (long) FLOAT32_MAX_CONSECUTIVE_INT,
             3, "mask nElements exceeds single-precision float "
             "consecutive integer precision size (2^24)");

  // `mask` and `tensor` must have the same number of elements
  THArgCheck(maskSize == tensorSize, 2,
             "mask and tensor must have the same number of elements");

  THClTensor* contigMask = THClTensor_newContiguous(state, mask);
  long oneElements = (long) THClTensor_sumall(state, contigMask);

  // The number of `1` elements present in the mask must be <= the
  // number of elements available in `src`
  if (oneElements > srcSize) {
    THClTensor_free(state, contigMask);
    THArgCheck(false, 2, "source nElements must be == mask `1` elements");
  }

  // Use a prefix sum to determine the copy locations of the masked elements
  THClTensor* maskPrefixSum = THClTensor_newv2(state, src->storage->device);
  THClTensor_resizeAs(state, maskPrefixSum, contigMask);

  // We are getting elements from `src` based on an offset from
  // `maskPrefixSum`, so that should be made contiguous too
  THClTensor* contigSrc = THClTensor_newContiguous(state, src);

//  TensorAddOp cumOp;
//  THAssert(THClTensor_checkGPU(state, 2, maskData, maskPrefixSumData));
//  THClTensor_scanDim(state, maskPrefixSumData, maskData, dimension, 0.0f, &cumOp);
  // hmmmm: this is scanDim, but what we really need is somsething like: scanAll

//   thrust::device_ptr<float>
//    maskData(THClTensor_data(state, contigMask));
//   thrust::device_ptr<float>
//    maskPrefixSumData(THClTensor_data(state, maskPrefixSum));
//   thrust::exclusive_scan(maskData,
//                         maskData + THClTensor_nElement(state, contigMask),
//                         maskPrefixSumData);
  THError("Not implemented");

  // update `tensor` where `mask` == 1 but pull from `src` at
  // maskPrefixSum
  TensorMaskedCopyOp maskedCopyOp(
     contigSrc, contigMask, maskPrefixSum);
  bool status = THClTensor_pointwiseApply2(
    state, tensor, contigMask, &maskedCopyOp);
  THError("Not implemented");

  THClTensor_free(state, contigSrc);
  THClTensor_free(state, maskPrefixSum);
  THClTensor_free(state, contigMask);

  THArgCheck(status, 2, CLTORCH_DIM_WARNING);

  THError("Not implemented");
}

void THClTensor_maskedSelect(THClState* state,
                              THClTensor *tensor, THClTensor *src, THClTensor *mask)
  {
  THAssert(THClTensor_checkGPU(state, 3, tensor, src, mask));
  THArgCheck(THClTensor_nElement(state, mask) == THClTensor_nElement(state, src),
            2, "sizes do not match");

  // Since we are performing a prefix sum of mask, it cannot exceed
  // the size allowed in consecutive integers in float32
  // THArgCheck(THClTensor_nElement(state, mask) <=
  //           (long) FLOAT32_MAX_CONSECUTIVE_INT,
  //           3, "mask nElements exceeds single-precision float "
  //           "consecutive integer precision size (2^24)");

  THClTensor* contigSrc = THClTensor_newContiguous(state, src);

  // Determine our output size
  THClTensor* contigMask = THClTensor_newContiguous(state, mask);
  long totalElements = (long) THClTensor_sumall(state, contigMask);
  cout << "totalElements " << totalElements << endl;

  // This should be contiguous already, so no need to make it contig
  // for the apply kernel
  THClTensor_resize1d(state, tensor, totalElements);

  // Use a prefix sum to determine the output locations of the masked elements
  // THClTensor* maskPrefixSum = THClTensor_new(state);
  // THClTensor_resizeAs(state, maskPrefixSum, contigMask);

  // ==== boost compute bit starts ========

    EasyCL *cl = src->storage->cl;

    boost::compute::context boost_context(*cl->context);
    boost::compute::command_queue boost_queue(*cl->queue);

    boost::compute::buffer boostData(*contigSrc->storage->wrapper->getDeviceArray());
    boost::compute::buffer boostMask(*contigMask->storage->wrapper->getDeviceArray());
    boost::compute::buffer boostOut(*tensor->storage->wrapper->getDeviceArray());

    transform_if(
      make_zip_iterator(
        boost::make_tuple(
        boost::compute::make_buffer_iterator<float>(boostData, 0),
        boost::compute::make_buffer_iterator<float>(boostMask, 0)
        )
      ),
      make_zip_iterator(
        boost::make_tuple(
        boost::compute::make_buffer_iterator<float>(boostData, THClTensor_nElement(state, mask)),
        boost::compute::make_buffer_iterator<float>(boostMask, THClTensor_nElement(state, mask))
        )
      ),
      boost::compute::make_buffer_iterator<float>(boostOut, 0),
      boost::compute::get<0>(), // function that return input value
      boost::compute::lambda::get<1>(boost::compute::_1) == 1, // lambda function that checks if mask is 1
      boost_queue // command queue (boost::compute::command_queue object)
    );
    tensor->storage->wrapper->markDeviceDirty();

  // ==== boost compute bit ends ========

  // thrust::device_ptr<float>
  //    maskData(THClTensor_data(state, contigMask));
  // thrust::device_ptr<float>
  //   maskPrefixSumData(THClTensor_data(state, maskPrefixSum));
  // thrust::exclusive_scan(maskData,
  //                       maskData + THClTensor_nElement(state, contigMask),
  //                       maskPrefixSumData);

  // Then copy over the masked elements at their desired output index
  // bool status = THClTensor_pointwiseApply3(
  //  state, contigMask, maskPrefixSum,
  //  src, TensorMaskedSelectOp(THClTensor_data(state, tensor)));

  THClTensor_free(state, contigSrc);
  THClTensor_free(state, contigMask);
  // THClTensor_free(state, maskPrefixSum);

  // THArgCheck(status, 2, CLTORCH_DIM_WARNING);

  // THError("Not implemented");
}

void THClTensor_maskedFillByte(THClState* state, THClTensor *tensor, THByteTensor *mask, float value)
{
 THAssert(THClTensor_checkGPU(state, 1, tensor));
 THLongStorage* maskSize = THByteTensor_newSizeOf(mask);
 const int device = tensor->storage->device;
 THClTensor* maskCl = THClTensor_newWithSize(state, device, maskSize, NULL);
 THLongStorage_free(maskSize);
 THClTensor_copyByte(state, maskCl, mask);
 THClTensor_maskedFill(state, tensor, maskCl, value);
 THClTensor_free(state, maskCl);
}

//void THClTensor_maskedCopyByte(THClState* state, THClTensor *tensor, THByteTensor *mask, THClTensor *src)
//{
//  THAssert(THClTensor_checkGPU(state, 2, tensor, src));
//  THLongStorage* maskSize = THByteTensor_newSizeOf(mask);
//  THClTensor* maskCl = THClTensor_newWithSize(state, maskSize, NULL);
//  THLongStorage_free(maskSize);
//  THClTensor_copyByte(state, maskCl, mask);
//  THClTensor_maskedCopy(state, tensor, maskCl, src);
//  THClTensor_free(state, maskCl);
//}

void THClTensor_maskedSelectByte(THClState* state, THClTensor *tensor, THClTensor *src, THByteTensor *mask)
{
 THAssert(THClTensor_checkGPU(state, 2, tensor, src));
 THLongStorage* maskSize = THByteTensor_newSizeOf(mask);
 const int device = src->storage->device;
 THClTensor* maskCl = THClTensor_newWithSize(state, device, maskSize, NULL);
 THLongStorage_free(maskSize);
 THClTensor_copyByte(state, maskCl, mask);
 THClTensor_maskedSelect(state, tensor, src, maskCl);
 THClTensor_free(state, maskCl);
}
<commit_msg>use boost compute to calculate size of output tensor for maskedSelect<commit_after>// from lib/THC/THCTensorMasked.cu:

#if defined(__APPLE__) || defined(__MACOSX)
#include <OpenCL/cl.h>
#else
#include <CL/cl.h>
#endif

#include <boost/compute/core.hpp>
#include <boost/compute/iterator/buffer_iterator.hpp>
#include <boost/compute/algorithm/reverse.hpp>
#include <boost/compute/algorithm/transform_if.hpp>
#include <boost/compute/lambda.hpp>
#include <boost/compute/iterator/zip_iterator.hpp>

#include "THClTensorMath.h"
#include "THClGeneral.h"
#include "THClBlas.h"
#include "THClTensorCopy.h"
//#include "THClTensorRandom.h"
#include "THClApply.h"
#include "THClReduce.h"

#include <iostream>
using namespace std;

// The largest consecutive integer representable in float32 (2^24)
#define FLOAT32_MAX_CONSECUTIVE_INT 16777216.0f

class TensorMaskedFillOp : public HasOperator2, public HasScalars {
public:
  int getNumScalars() const { return 1; }
  float getScalar( int index ) const { return value; }
  TensorMaskedFillOp(float v) : value(v) {}
  std::string operator2() const {
    return "if( *in1 != 0.0f ) { *out = val1; }";
  }
  float value;
};

class TensorMaskedCopyOp : public HasOperator2, public HasGlobalTensors {
public:
  TensorMaskedCopyOp(THClTensor *src, THClTensor *baseMask, THClTensor *maskPrefixSum) :
    src(src), baseMask(baseMask), maskPrefixSum(maskPrefixSum) {
  }
  int getNumGlobalTensors() const { return 3; }
  THClTensor *getTensor(int index) const {
    if(index == 0){ return src; }
    if(index == 1){ return baseMask; }
    if(index == 2){ return maskPrefixSum; }
    THError("index not recognized %i", index);
    return 0;
  }
  std::string getTensorName(int index) const {
    if(index == 0){ return "src"; }
    if(index == 1){ return "baseMask"; }
    if(index == 2){ return "maskPrefixSum"; }
    THError("index not recognized %i", index);
    return "";
  }
  std::string operator2() const {
    return "if( *in1 != 0.0f ) { *out = src[(int)maskPrefixSum[srcOffset] ]; }";
  }
  THClTensor *src;
  THClTensor *baseMask;
  THClTensor *maskPrefixSum;
};

//struct TensorMaskedCopyOp {
//  TensorMaskedCopyOp(float* s, float* bm, float* ps)
//      : src(s),
//        baseMask(bm),
//        maskPrefixSum(ps) {
//  }

//  /*__device__*/ /*__forceline__*/ void operator()(float* out, float* mask) {
//    // Really mask should be `0` or `1` but we can't propagate errors here.
//    if (*mask != 0.0f) {
//      // We've already checked that this offset is <= 2^24, so this is ok.
//      int srcOffset = (int) (mask - baseMask);
//      *out = src[(int) maskPrefixSum[srcOffset]];
//    }
//  }

//  // Where we are copying from
//  float* src;

//  // The base address of mask so we can calculate offset
//  float* baseMask;

//  // The index we are copying from
//  float* maskPrefixSum;
//};

//class TensorMaskedSelectOp : public HasOperator3, public HasScalars {
//public:
//  int getNumScalars() const { return 1; }
//  string operator3() const {
//    return "if(*out != 0.0f){out[(int)*in1] = *in2; }";
//  }
//  TensorMaskedSelectOp(float* t) : out(t) {}
//  void operator()(float* mask, float* maskPrefixSum, float* in) {
//    // Really mask should be `0` or `1` but we can't propagate errors here.
//    if (*mask != 0.0f) {
//      out[(int) *maskPrefixSum] = *in;
//    }
//  }

//  float* out;
//};

void THClTensor_maskedFill(THClState* state,
                             THClTensor *tensor, THClTensor *mask, float value)
{
  THAssert(THClTensor_checkGPU(state, 2, tensor, mask));
  THArgCheck(THClTensor_nElement(state, tensor) ==
             THClTensor_nElement(state, mask),
             2, "sizes do not match");

  TensorMaskedFillOp op(value);
  if (!THClTensor_pointwiseApply2(state, tensor, mask, &op)) {
    THArgCheck(false, 2, CLTORCH_DIM_WARNING);
  }
}

void THClTensor_maskedCopy(THClState* state,
                             THClTensor *tensor, THClTensor *mask, THClTensor *src)
{
  THAssert(THClTensor_checkGPU(state, 3, tensor, src, mask));
  long maskSize = THClTensor_nElement(state, mask);
  long tensorSize = THClTensor_nElement(state, tensor);
  long srcSize = THClTensor_nElement(state, src);

  // Since we are performing a prefix sum of mask, it cannot exceed
  // the size allowed in consecutive integers in float32
  THArgCheck(maskSize <= (long) FLOAT32_MAX_CONSECUTIVE_INT,
             3, "mask nElements exceeds single-precision float "
             "consecutive integer precision size (2^24)");

  // `mask` and `tensor` must have the same number of elements
  THArgCheck(maskSize == tensorSize, 2,
             "mask and tensor must have the same number of elements");

  THClTensor* contigMask = THClTensor_newContiguous(state, mask);
  long oneElements = (long) THClTensor_sumall(state, contigMask);

  // The number of `1` elements present in the mask must be <= the
  // number of elements available in `src`
  if (oneElements > srcSize) {
    THClTensor_free(state, contigMask);
    THArgCheck(false, 2, "source nElements must be == mask `1` elements");
  }

  // Use a prefix sum to determine the copy locations of the masked elements
  THClTensor* maskPrefixSum = THClTensor_newv2(state, src->storage->device);
  THClTensor_resizeAs(state, maskPrefixSum, contigMask);

  // We are getting elements from `src` based on an offset from
  // `maskPrefixSum`, so that should be made contiguous too
  THClTensor* contigSrc = THClTensor_newContiguous(state, src);

//  TensorAddOp cumOp;
//  THAssert(THClTensor_checkGPU(state, 2, maskData, maskPrefixSumData));
//  THClTensor_scanDim(state, maskPrefixSumData, maskData, dimension, 0.0f, &cumOp);
  // hmmmm: this is scanDim, but what we really need is somsething like: scanAll

//   thrust::device_ptr<float>
//    maskData(THClTensor_data(state, contigMask));
//   thrust::device_ptr<float>
//    maskPrefixSumData(THClTensor_data(state, maskPrefixSum));
//   thrust::exclusive_scan(maskData,
//                         maskData + THClTensor_nElement(state, contigMask),
//                         maskPrefixSumData);
  THError("Not implemented");

  // update `tensor` where `mask` == 1 but pull from `src` at
  // maskPrefixSum
  TensorMaskedCopyOp maskedCopyOp(
     contigSrc, contigMask, maskPrefixSum);
  bool status = THClTensor_pointwiseApply2(
    state, tensor, contigMask, &maskedCopyOp);
  THError("Not implemented");

  THClTensor_free(state, contigSrc);
  THClTensor_free(state, maskPrefixSum);
  THClTensor_free(state, contigMask);

  THArgCheck(status, 2, CLTORCH_DIM_WARNING);

  THError("Not implemented");
}

void THClTensor_maskedSelect(THClState* state,
                              THClTensor *tensor, THClTensor *src, THClTensor *mask)
  {
  THAssert(THClTensor_checkGPU(state, 3, tensor, src, mask));
  THArgCheck(THClTensor_nElement(state, mask) == THClTensor_nElement(state, src),
            2, "sizes do not match");

  // Since we are performing a prefix sum of mask, it cannot exceed
  // the size allowed in consecutive integers in float32
  // THArgCheck(THClTensor_nElement(state, mask) <=
  //           (long) FLOAT32_MAX_CONSECUTIVE_INT,
  //           3, "mask nElements exceeds single-precision float "
  //           "consecutive integer precision size (2^24)");

  THClTensor* contigSrc = THClTensor_newContiguous(state, src);

  // Determine our output size
  THClTensor* contigMask = THClTensor_newContiguous(state, mask);

  int sourceSize = THClTensor_nElement(state, contigMask);

  EasyCL *cl = src->storage->cl;

  boost::compute::context boost_context(*cl->context);
  boost::compute::command_queue boost_queue(*cl->queue);

  boost::compute::buffer boostData(*contigSrc->storage->wrapper->getDeviceArray());
  boost::compute::buffer boostMask(*contigMask->storage->wrapper->getDeviceArray());

  // long totalElements = (long) THClTensor_sumall(state, contigMask);
  // since sumall doesnt work if we have more than something like ~2^24 elements, because exceeds float precision
  // we need to calculate the sumall in blocks, of say 2^20 elements
  // int totalElements = 0;
  // int device = mask->device;
  // const int blockSize = 1 << 21;
  // cout << "blockSize " << blockSize << endl;
  // int numBlocks = (sourceSize + blockSize - 1) / blockSize;
  // THClTensor *flatContigMask = THClTensor_newWithStorage1d(state, device, contigMask->storage, 0,
  //                              totalElements, 1);
  // for(int block = 0; block < numBlocks; block++) {
  //   int blockStart = block * blockSize;
  //   int blockEndExclusive = blockStart + blockSize;
  //   if(blockEndExclusive > totalElements) {
  //     blockEndExclusive = totalElements;
  //   }

  //   THClTensor *blockTensor = THClTensor_newNarrow(state, flatContigMask, 0, 0, blockEndExclusive - blockStart);
  //   int blockCount = THClTensor_sumall(state, blockTensor)
  //   THClTensor_free(state, blockTensor);
  //   totalElements += blockCount;
  // }
  // cout << "totalElements " << totalElements << endl;
  // THClTensor_free(state, flatContigMask);

  size_t copied_element_count = ::boost::compute::count(
    boost::compute::make_buffer_iterator<float>(boostMask, 0),
    boost::compute::make_buffer_iterator<float>(boostMask, sourceSize),
    1,
    boost_queue
  );
  int totalElements = (int)copied_element_count;
  cout << "totalElements " << totalElements << endl;

  THClTensor_resize1d(state, tensor, totalElements);
  boost::compute::buffer boostOut(*tensor->storage->wrapper->getDeviceArray());

  transform_if(
    make_zip_iterator(
      boost::make_tuple(
      boost::compute::make_buffer_iterator<float>(boostData, 0),
      boost::compute::make_buffer_iterator<float>(boostMask, 0)
      )
    ),
    make_zip_iterator(
      boost::make_tuple(
      boost::compute::make_buffer_iterator<float>(boostData, sourceSize),
      boost::compute::make_buffer_iterator<float>(boostMask, sourceSize)
      )
    ),
    boost::compute::make_buffer_iterator<float>(boostOut, 0),
    boost::compute::get<0>(), // function that return input value
    boost::compute::lambda::get<1>(boost::compute::_1) == 1, // lambda function that checks if mask is 1
    boost_queue // command queue (boost::compute::command_queue object)
  );
  tensor->storage->wrapper->markDeviceDirty();

  // ==== boost compute bit ends ========

  // thrust::device_ptr<float>
  //    maskData(THClTensor_data(state, contigMask));
  // thrust::device_ptr<float>
  //   maskPrefixSumData(THClTensor_data(state, maskPrefixSum));
  // thrust::exclusive_scan(maskData,
  //                       maskData + THClTensor_nElement(state, contigMask),
  //                       maskPrefixSumData);

  // Then copy over the masked elements at their desired output index
  // bool status = THClTensor_pointwiseApply3(
  //  state, contigMask, maskPrefixSum,
  //  src, TensorMaskedSelectOp(THClTensor_data(state, tensor)));

  THClTensor_free(state, contigSrc);
  THClTensor_free(state, contigMask);
  // THClTensor_free(state, maskPrefixSum);

  // THArgCheck(status, 2, CLTORCH_DIM_WARNING);

  // THError("Not implemented");
}

void THClTensor_maskedFillByte(THClState* state, THClTensor *tensor, THByteTensor *mask, float value)
{
 THAssert(THClTensor_checkGPU(state, 1, tensor));
 THLongStorage* maskSize = THByteTensor_newSizeOf(mask);
 const int device = tensor->storage->device;
 THClTensor* maskCl = THClTensor_newWithSize(state, device, maskSize, NULL);
 THLongStorage_free(maskSize);
 THClTensor_copyByte(state, maskCl, mask);
 THClTensor_maskedFill(state, tensor, maskCl, value);
 THClTensor_free(state, maskCl);
}

//void THClTensor_maskedCopyByte(THClState* state, THClTensor *tensor, THByteTensor *mask, THClTensor *src)
//{
//  THAssert(THClTensor_checkGPU(state, 2, tensor, src));
//  THLongStorage* maskSize = THByteTensor_newSizeOf(mask);
//  THClTensor* maskCl = THClTensor_newWithSize(state, maskSize, NULL);
//  THLongStorage_free(maskSize);
//  THClTensor_copyByte(state, maskCl, mask);
//  THClTensor_maskedCopy(state, tensor, maskCl, src);
//  THClTensor_free(state, maskCl);
//}

void THClTensor_maskedSelectByte(THClState* state, THClTensor *tensor, THClTensor *src, THByteTensor *mask)
{
 THAssert(THClTensor_checkGPU(state, 2, tensor, src));
 THLongStorage* maskSize = THByteTensor_newSizeOf(mask);
 const int device = src->storage->device;
 THClTensor* maskCl = THClTensor_newWithSize(state, device, maskSize, NULL);
 THLongStorage_free(maskSize);
 THClTensor_copyByte(state, maskCl, mask);
 THClTensor_maskedSelect(state, tensor, src, maskCl);
 THClTensor_free(state, maskCl);
}
<|endoftext|>


<commit_before>/**
 * \file
 * \brief Mutex class implementation
 *
 * \author Copyright (C) 2014 Kamil Szczygiel http://www.distortec.com http://www.freddiechopin.info
 *
 * \par License
 * This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. If a copy of the MPL was not
 * distributed with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * \date 2014-09-26
 */

#include "distortos/scheduler/Mutex.hpp"

#include "distortos/scheduler/schedulerInstance.hpp"
#include "distortos/scheduler/Scheduler.hpp"

#include "distortos/architecture/InterruptMaskingLock.hpp"

#include <cerrno>

namespace distortos
{

namespace scheduler
{

/*---------------------------------------------------------------------------------------------------------------------+
| public functions
+---------------------------------------------------------------------------------------------------------------------*/

Mutex::Mutex(const Type type) :
		blockedList_{schedulerInstance.getThreadControlBlockListAllocator(), ThreadControlBlock::State::BlockedOnMutex},
		owner_{},
		type_{type}
{

}

int Mutex::lock()
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (owner_ != nullptr)
	{
		if (type_ == Type::ErrorChecking && owner_ == &schedulerInstance.getCurrentThreadControlBlock())
			return EDEADLK;

		schedulerInstance.block(blockedList_);
	}
	else
		owner_ = &schedulerInstance.getCurrentThreadControlBlock();

	return 0;
}

int Mutex::tryLock()
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (owner_ != nullptr)
		return EBUSY;

	owner_ = &schedulerInstance.getCurrentThreadControlBlock();
	return 0;
}

int Mutex::tryLockFor(const TickClock::duration duration)
{
	return tryLockUntil(TickClock::now() + duration + TickClock::duration{1});
}

int Mutex::tryLockUntil(const TickClock::time_point timePoint)
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (owner_ != nullptr)
	{
		return schedulerInstance.blockUntil(blockedList_, timePoint);
	}

	owner_ = &schedulerInstance.getCurrentThreadControlBlock();
	return 0;
}

int Mutex::unlock()
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (blockedList_.empty() == false)
	{
		owner_ = &blockedList_.begin()->get();	// pass ownership to the unblocked thread
		schedulerInstance.unblock(blockedList_.begin());
	}
	else
		owner_ = nullptr;

	return 0;
}

}	// namespace scheduler

}	// namespace distortos
<commit_msg>Mutex: implement POSIX error handling in Mutex::unlock()<commit_after>/**
 * \file
 * \brief Mutex class implementation
 *
 * \author Copyright (C) 2014 Kamil Szczygiel http://www.distortec.com http://www.freddiechopin.info
 *
 * \par License
 * This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. If a copy of the MPL was not
 * distributed with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * \date 2014-09-27
 */

#include "distortos/scheduler/Mutex.hpp"

#include "distortos/scheduler/schedulerInstance.hpp"
#include "distortos/scheduler/Scheduler.hpp"

#include "distortos/architecture/InterruptMaskingLock.hpp"

#include <cerrno>

namespace distortos
{

namespace scheduler
{

/*---------------------------------------------------------------------------------------------------------------------+
| public functions
+---------------------------------------------------------------------------------------------------------------------*/

Mutex::Mutex(const Type type) :
		blockedList_{schedulerInstance.getThreadControlBlockListAllocator(), ThreadControlBlock::State::BlockedOnMutex},
		owner_{},
		type_{type}
{

}

int Mutex::lock()
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (owner_ != nullptr)
	{
		if (type_ == Type::ErrorChecking && owner_ == &schedulerInstance.getCurrentThreadControlBlock())
			return EDEADLK;

		schedulerInstance.block(blockedList_);
	}
	else
		owner_ = &schedulerInstance.getCurrentThreadControlBlock();

	return 0;
}

int Mutex::tryLock()
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (owner_ != nullptr)
		return EBUSY;

	owner_ = &schedulerInstance.getCurrentThreadControlBlock();
	return 0;
}

int Mutex::tryLockFor(const TickClock::duration duration)
{
	return tryLockUntil(TickClock::now() + duration + TickClock::duration{1});
}

int Mutex::tryLockUntil(const TickClock::time_point timePoint)
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (owner_ != nullptr)
	{
		return schedulerInstance.blockUntil(blockedList_, timePoint);
	}

	owner_ = &schedulerInstance.getCurrentThreadControlBlock();
	return 0;
}

int Mutex::unlock()
{
	architecture::InterruptMaskingLock interruptMaskingLock;

	if (type_ != Type::Normal && owner_ != &schedulerInstance.getCurrentThreadControlBlock())
		return EPERM;

	if (blockedList_.empty() == false)
	{
		owner_ = &blockedList_.begin()->get();	// pass ownership to the unblocked thread
		schedulerInstance.unblock(blockedList_.begin());
	}
	else
		owner_ = nullptr;

	return 0;
}

}	// namespace scheduler

}	// namespace distortos
<|endoftext|>


<commit_before>#ifndef BANAL_GPU_DOT_CFG_H
#define BANAL_GPU_DOT_CFG_H

#include <algorithm>
#include <iostream>
#include <regex>
#include <sstream>
#include <string>
#include <unordered_set>
#include <vector>
// dyninst
#include <CFG.h>

#include <include/hpctoolkit-config.h>

namespace GPUParse {

typedef Dyninst::Architecture Arch;

struct Inst {
  int offset;
  int size;
  bool dual_first;
  bool dual_second;
  bool is_call;
  bool is_jump;
  bool is_sync;
  std::string predicate;
  std::string opcode;
  std::string port;
  std::string target;
  std::vector<std::string> operands;
  Arch arch;

  // Constructor for dummy inst
  Inst(int offset, int size, Arch arch) : offset(offset), size(size), dual_first(false), dual_second(false),
    is_call(false), is_jump(false), is_sync(false), arch(arch) {}

  Inst(int offset, int size) : Inst(offset, size, Dyninst::Arch_none) {}

  explicit Inst(int offset) : Inst(offset, 0) {}
};

#ifdef DYNINST_SUPPORTS_INTEL_GPU

struct IntelInst : public Inst {
  // Constructor for dummy inst
  IntelInst(int offset, int size) : Inst(offset, size, Dyninst::Arch_intelGen9) {}

  explicit IntelInst(int offset) : Inst(offset, 0, Dyninst::Arch_intelGen9) {}
};

#endif

struct CudaInst : public Inst {
  // Constructor for dummy inst
  CudaInst(int offset, int size) : Inst(offset, size, Dyninst::Arch_cuda) {}

  explicit CudaInst(int offset) : Inst(offset, 0, Dyninst::Arch_cuda) {}

  // Cuda instruction constructor
  CudaInst(std::string &inst_str) : CudaInst(0, 0) {
    if (inst_str.find("{") != std::string::npos) {  // Dual first
      auto pos = inst_str.find("{");
      inst_str.replace(pos, 1, " ");
      dual_first = true;
    }
    if (inst_str.find("}") != std::string::npos) {  // Dual second
      inst_str = inst_str.substr(2);
      auto pos = inst_str.find("*/");
      if (pos != std::string::npos) {
        inst_str.replace(pos, 2, ":");
        dual_second = true;
      }
    }
    std::istringstream *iss = new std::istringstream(inst_str);
    std::string s;
    if (std::getline(*iss, s, ':')) {
      if (s.find("<") != std::string::npos) {
        // Port notation in dot graph to link basic blocks
        auto pos = s.find(">");
        this->port = s.substr(1, pos - 1);
        s = s.substr(pos + 1);
      }
      std::stringstream ss;
      ss << std::hex << s;
      ss >> offset;
      if (std::getline(*iss, s, ':')) {
        s.erase(std::remove(s.begin(), s.end(), '{'), s.end());
        s.erase(std::remove(s.begin(), s.end(), '}'), s.end());
        s.erase(std::remove(s.begin(), s.end(), ';'), s.end());
        s.erase(std::remove(s.begin(), s.end(), ','), s.end());
        s.erase(std::remove(s.begin(), s.end(), '('), s.end());
        s.erase(std::remove(s.begin(), s.end(), ')'), s.end());
        s.erase(std::remove(s.begin(), s.end(), '`'), s.end());
        std::regex e("\\\\ ");
        std::string newline("\n");
        delete iss;
        iss = new std::istringstream(std::regex_replace(s, e, newline));
        while (std::getline(*iss, s)) {
          if (s != "") {
            if (opcode == "") {
              if (s.find("@") != std::string::npos) {
                predicate = s;
              } else {
                opcode = s;
                if (opcode.find("CALL") != std::string::npos || // sm_70
                  opcode.find("CAL") != std::string::npos) { // sm_60
                  this->is_call = true;
                } else if (opcode.find("BRA") != std::string::npos ||
                  opcode.find("BRX") != std::string::npos ||
                  opcode.find("JMP") != std::string::npos ||
                  opcode.find("JMX") != std::string::npos ||
                  opcode.find("BREAK") != std::string::npos ||
                  opcode.find("JMXU") != std::string::npos) {
                  this->is_jump = true;
                } else if (opcode.find("SYNC") != std::string::npos) {
                  // avoid Barrier Set Convergence Synchronization Point
                  //opcode.find("SSY") != std::string::npos ||
                  //opcode.find("BSSY") != std::string::npos)
                  // TODO(Keren): add more sync instructions
                  this->is_sync = true;
                }
              }
            } else {
              // Target
              operands.push_back(s);
              if (is_jump || is_sync) {
                auto pos = s.find(".L_");
                if (pos != std::string::npos) {
                  auto end_pos = pos + 3;
                  size_t len = 0;
                  while (end_pos != std::string::npos) {
                    if (!std::isdigit(s[end_pos])) {
                      break;
                    }
                    ++len;
                    ++end_pos;
                  }
                  this->target = s.substr(pos, len + 3);
                }
              }
            }
          }
        }
      }
    }
    delete iss;
  }
};

struct Block;

typedef Dyninst::ParseAPI::EdgeTypeEnum TargetType;

struct Target {
  Inst *inst;
  Block *block;
  TargetType type;

  Target(Inst *inst, Block *block, TargetType type) : inst(inst), block(block), type(type) {}

  bool operator<(const Target &other) const {
    return this->inst->offset < other.inst->offset;
  }
};

struct Block {
  size_t id;
  int address;
  int begin_offset;
  std::vector<Inst *> insts;
  std::vector<Target *> targets;
  std::string name;

  Block(size_t id, int address, const std::string &name) : id(id), address(address), name(name) {}

  Block(size_t id, const std::string &name) : Block(id, 0, name) {}

  bool operator<(const Block &other) const {
    if (this->insts.size() == 0) {
      return true;
    } else if (other.insts.size() == 0) {
      return false;
    } else {
      return this->insts[0]->offset < other.insts[0]->offset;
    }
  }

  ~Block() {
    for (auto *inst : insts) {
      delete inst;
    }
    for (auto *target: targets) {
      delete target;
    }
  }
};

struct Function {
  std::vector<Block *> blocks;
  size_t id;
  std::string name;
  int address;

  Function(size_t id, const std::string &name) : id(id), name(name),
    address(0) {}

  ~Function() {
    for (auto *block : blocks) {
      delete block;
    }
  }
};

struct LoopEntry {
  Block *entry_block;
  Block *back_edge_block;
  Inst *back_edge_inst;

  LoopEntry(Block *entry_block) : entry_block(entry_block) {}

  LoopEntry(Block *entry_block, Block *back_edge_block, Inst *back_edge_inst) :
    entry_block(entry_block), back_edge_block(back_edge_block), back_edge_inst(back_edge_inst) {}
};

struct Loop {
  std::vector<LoopEntry *> entries;
  std::unordered_set<Loop *> child_loops;
  std::unordered_set<Block *> blocks;
  std::unordered_set<Block *> child_blocks;
  Function *function;

  Loop(Function *function) : function(function) {}
};

struct Call {
  Inst *inst;
  Block *block;
  Function *caller_function;
  Function *callee_function;

  Call(Inst *inst, Block *block, Function *caller_function, Function *callee_function) :
    inst(inst), block(block), caller_function(caller_function), callee_function(callee_function) {}
};

}

#endif
<commit_msg>CUDA-11.6-fix (#509)<commit_after>#ifndef BANAL_GPU_DOT_CFG_H
#define BANAL_GPU_DOT_CFG_H

#include <algorithm>
#include <iostream>
#include <regex>
#include <sstream>
#include <string>
#include <unordered_set>
#include <vector>
// dyninst
#include <CFG.h>

#include <include/hpctoolkit-config.h>

namespace GPUParse {

typedef Dyninst::Architecture Arch;

struct Inst {
  int offset;
  int size;
  bool dual_first;
  bool dual_second;
  bool is_call;
  bool is_jump;
  bool is_sync;
  std::string predicate;
  std::string opcode;
  std::string port;
  std::string target;
  std::vector<std::string> operands;
  Arch arch;

  // Constructor for dummy inst
  Inst(int offset, int size, Arch arch) : offset(offset), size(size), dual_first(false), dual_second(false),
    is_call(false), is_jump(false), is_sync(false), arch(arch) {}

  Inst(int offset, int size) : Inst(offset, size, Dyninst::Arch_none) {}

  explicit Inst(int offset) : Inst(offset, 0) {}
};

#ifdef DYNINST_SUPPORTS_INTEL_GPU

struct IntelInst : public Inst {
  // Constructor for dummy inst
  IntelInst(int offset, int size) : Inst(offset, size, Dyninst::Arch_intelGen9) {}

  explicit IntelInst(int offset) : Inst(offset, 0, Dyninst::Arch_intelGen9) {}
};

#endif

struct CudaInst : public Inst {
  // Constructor for dummy inst
  CudaInst(int offset, int size) : Inst(offset, size, Dyninst::Arch_cuda) {}

  explicit CudaInst(int offset) : Inst(offset, 0, Dyninst::Arch_cuda) {}

  // Cuda instruction constructor
  CudaInst(std::string &inst_str) : CudaInst(0, 0) {
    if (inst_str.find("{") != std::string::npos) {  // Dual first
      auto pos = inst_str.find("{");
      inst_str.replace(pos, 1, " ");
      dual_first = true;
    }
    if (inst_str.find("}") != std::string::npos) {  // Dual second
      inst_str = inst_str.substr(2);
      auto pos = inst_str.find("*/");
      if (pos != std::string::npos) {
        inst_str.replace(pos, 2, ":");
        dual_second = true;
      }
    }
    std::istringstream *iss = new std::istringstream(inst_str);
    std::string s;
    if (std::getline(*iss, s, ':')) {
      if (s.find("<") != std::string::npos) {
        // Port notation in dot graph to link basic blocks
        auto pos = s.find(">");
        this->port = s.substr(1, pos - 1);
        s = s.substr(pos + 1);
      }
      std::stringstream ss;
      ss << std::hex << s;
      ss >> offset;
      if (std::getline(*iss, s, ':')) {
        s.erase(std::remove(s.begin(), s.end(), '{'), s.end());
        s.erase(std::remove(s.begin(), s.end(), '}'), s.end());
        s.erase(std::remove(s.begin(), s.end(), ';'), s.end());
        s.erase(std::remove(s.begin(), s.end(), ','), s.end());
        s.erase(std::remove(s.begin(), s.end(), '('), s.end());
        s.erase(std::remove(s.begin(), s.end(), ')'), s.end());
        s.erase(std::remove(s.begin(), s.end(), '`'), s.end());
        std::regex e("\\\\ ");
        std::string newline("\n");
        delete iss;
        iss = new std::istringstream(std::regex_replace(s, e, newline));
        while (std::getline(*iss, s)) {
          if (s != "") {
            if (opcode == "") {
              if (s.find("@") != std::string::npos) {
                predicate = s;
              } else {
                opcode = s;
                if (opcode.find("CALL") != std::string::npos || // sm_70
                  opcode.find("CAL") != std::string::npos) { // sm_60
                  this->is_call = true;
                } else if (opcode.find("BRA") != std::string::npos ||
                  opcode.find("BRX") != std::string::npos ||
                  opcode.find("JMP") != std::string::npos ||
                  opcode.find("JMX") != std::string::npos ||
                  opcode.find("BREAK") != std::string::npos ||
                  opcode.find("JMXU") != std::string::npos) {
                  this->is_jump = true;
                } else if (opcode.find("SYNC") != std::string::npos) {
                  // avoid Barrier Set Convergence Synchronization Point
                  //opcode.find("SSY") != std::string::npos ||
                  //opcode.find("BSSY") != std::string::npos)
                  // TODO(Keren): add more sync instructions
                  this->is_sync = true;
                }
              }
            } else {
              // Extract target name of the jump or the sync instruction
              // The target name can be in two forms:
              // "**.L_<number>**" for CUDA < 11.5
              // "**.L_x_<number>**" for CUDA >= 11.5
              // We assuming the target label will keep the form "**.L_**<number>**" in the further release
              static const std::string TARGET_LABEL = ".L_";
              operands.push_back(s);
              if (is_jump || is_sync) {
                auto pos = s.find(TARGET_LABEL);
                if (pos != std::string::npos) {
                  auto digit_pos = pos + TARGET_LABEL.size();
                  // Find the start digit of the number
                  for (;digit_pos != std::string::npos && !std::isdigit(s[digit_pos]); ++digit_pos);
                  // Find the end digit of the number
                  for (;digit_pos != std::string::npos && std::isdigit(s[digit_pos]); ++digit_pos);
                  if (digit_pos != std::string::npos) {
                    this->target = s.substr(pos, digit_pos - pos);
                  }
                }
              }
            }
          }
        }
      }
    }
    delete iss;
  }
};

struct Block;

typedef Dyninst::ParseAPI::EdgeTypeEnum TargetType;

struct Target {
  Inst *inst;
  Block *block;
  TargetType type;

  Target(Inst *inst, Block *block, TargetType type) : inst(inst), block(block), type(type) {}

  bool operator<(const Target &other) const {
    return this->inst->offset < other.inst->offset;
  }
};

struct Block {
  size_t id;
  int address;
  int begin_offset;
  std::vector<Inst *> insts;
  std::vector<Target *> targets;
  std::string name;

  Block(size_t id, int address, const std::string &name) : id(id), address(address), name(name) {}

  Block(size_t id, const std::string &name) : Block(id, 0, name) {}

  bool operator<(const Block &other) const {
    if (this->insts.size() == 0) {
      return true;
    } else if (other.insts.size() == 0) {
      return false;
    } else {
      return this->insts[0]->offset < other.insts[0]->offset;
    }
  }

  ~Block() {
    for (auto *inst : insts) {
      delete inst;
    }
    for (auto *target: targets) {
      delete target;
    }
  }
};

struct Function {
  std::vector<Block *> blocks;
  size_t id;
  std::string name;
  int address;

  Function(size_t id, const std::string &name) : id(id), name(name),
    address(0) {}

  ~Function() {
    for (auto *block : blocks) {
      delete block;
    }
  }
};

struct LoopEntry {
  Block *entry_block;
  Block *back_edge_block;
  Inst *back_edge_inst;

  LoopEntry(Block *entry_block) : entry_block(entry_block) {}

  LoopEntry(Block *entry_block, Block *back_edge_block, Inst *back_edge_inst) :
    entry_block(entry_block), back_edge_block(back_edge_block), back_edge_inst(back_edge_inst) {}
};

struct Loop {
  std::vector<LoopEntry *> entries;
  std::unordered_set<Loop *> child_loops;
  std::unordered_set<Block *> blocks;
  std::unordered_set<Block *> child_blocks;
  Function *function;

  Loop(Function *function) : function(function) {}
};

struct Call {
  Inst *inst;
  Block *block;
  Function *caller_function;
  Function *callee_function;

  Call(Inst *inst, Block *block, Function *caller_function, Function *callee_function) :
    inst(inst), block(block), caller_function(caller_function), callee_function(callee_function) {}
};

}

#endif
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2017 Trail of Bits, Inc.
 *
 * 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 <glog/logging.h>
#include <gflags/gflags.h>

#include <iomanip>
#include <iostream>
#include <memory>
#include <string>
#include <sstream>

#include <llvm/IR/Function.h>
#include <llvm/IR/GlobalVariable.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>

#include <remill/Arch/Arch.h>
#include <remill/BC/Util.h>
#include <remill/BC/Version.h>

#include "mcsema/Arch/Arch.h"
#include "mcsema/BC/Lift.h"
#include "mcsema/BC/Util.h"

#ifndef LLVM_VERSION_STRING
# define LLVM_VERSION_STRING LLVM_VERSION_MAJOR << "." << LLVM_VERSION_MINOR
#endif

#ifndef MCSEMA_VERSION_STRING
# define MCSEMA_VERSION_STRING "unknown"
#endif  // MCSEMA_VERSION_STRING

#ifndef MCSEMA_BRANCH_NAME
# define MCSEMA_BRANCH_NAME "unknown"
#endif  // MCSEMA_BRANCH_NAME

DECLARE_string(arch);
DECLARE_string(os);

DEFINE_string(cfg, "", "Path to the CFG file containing code to lift.");

DEFINE_string(output, "", "Output bitcode file name.");

DEFINE_string(library, "", "Path to an LLVM bitcode or IR file that contains "
                           "external library definitions.");

DECLARE_bool(version);

DECLARE_bool(disable_optimizer);
DECLARE_bool(keep_memops);
DECLARE_bool(explicit_args);
DECLARE_string(pc_annotation);

DEFINE_bool(legacy_mode, false,
            "Try to make the output bitcode resemble the original McSema.");

namespace {

static void PrintVersion(void) {
  std::cout
      << "This is mcsema-lift version: " << MCSEMA_VERSION_STRING << std::endl
      << "Built from branch: " << MCSEMA_BRANCH_NAME << std::endl
      << "Using LLVM " << LLVM_VERSION_STRING << std::endl;
}

// Load in a separate bitcode or IR library, and copy function and variable
// declarations from that library into our module. We can use this feature
// to provide better type information to McSema.
static void LoadLibraryIntoModule(void) {
  std::unique_ptr<llvm::Module> lib(
      remill::LoadModuleFromFile(mcsema::gContext, FLAGS_library));

  // Declare the functions from the library in McSema's target module.
  for (auto &func : *lib) {
    auto func_name = func.getName();
    if (func_name.startswith("__mcsema") || func_name.startswith("__remill")) {
      continue;
    }
    if (mcsema::gModule->getFunction(func_name)) {
      continue;
    }

    auto dest_func = llvm::Function::Create(
        func.getFunctionType(), func.getLinkage(),
        func_name, mcsema::gModule);

    dest_func->copyAttributesFrom(&func);
    dest_func->setVisibility(func.getVisibility());
  }

  // Declare the global variables from the library in McSema's target module.
  for (auto &var : lib->globals()) {
    auto var_name = var.getName();
    if (var_name.startswith("__mcsema") || var_name.startswith("__remill")) {
      continue;
    }

    if (mcsema::gModule->getGlobalVariable(var_name)) {
      continue;
    }

    auto dest_var = new llvm::GlobalVariable(
        *mcsema::gModule, var.getType()->getElementType(),
        var.isConstant(), var.getLinkage(), nullptr,
        var_name, nullptr, var.getThreadLocalMode(),
        var.getType()->getAddressSpace());

    dest_var->copyAttributesFrom(&var);
  }
}

}  // namespace

int main(int argc, char *argv[]) {
  std::stringstream ss;
  ss << std::endl << std::endl
     << "  " << argv[0] << " \\" << std::endl
     << "    --output OUTPUT_BC_FILE \\" << std::endl
     << "    --arch ARCH_NAME \\" << std::endl
     << "    --os OS_NAME \\" << std::endl
     << "    --cfg CFG_FILE \\" << std::endl

     // This option is very useful for debugging McSema-lifted bitcode. It
     // injects so-called breakpoint functions before every lifted instruction.
     // For example, the the instruction at PC `0xf00` is lifted, then this
     // option will inject a call to `breakpoint_f00`. With this feature, we
     // can add breakpoints in a debugger on these breakpoint functions, and
     // know that they correspond to locations in the original program.
     << "    [--add_breakpoints] \\" << std::endl

     // This option injects a function call before every lifted instruction.
     // This function is implemented in the McSema runtime and it prints the
     // values of the general purpose registers to `stderr`.
     << "    [--add_reg_tracer] \\" << std::endl

     // This option tells McSema not to optimize the bitcode. This is useful
     // for debugging, especially in conjunction with `--add_breakpoints`.
     << "    [--disable_optimizer] \\" << std::endl

     // This option tells McSema not to lower Remill's memory access intrinsic
     // functions into LLVM `load` and `store` instructions.
     << "    [--keep_memops] \\" << std::endl

     // There are roughly two ways of using McSema-lifted bitcode. The default
     // use case is to compile the bitcode into an executable that behaves like
     // the original program. The other use case is to do some kind of static
     // analysis, e.g. with KLEE. In this use case, calls to external functions
     // are emulated so that we also try to explicitly lift parameter passing.
     // This mode of passing arguments explicitly is enabled by
     // `--explicit_args`, and in situations where we have no knowledge of the
     // argument counts expected by an external function, we fall back on
     // passing `--explicit_args_count` number of arguments to that function.
     << "    [--explicit_args] \\" << std::endl
     << "    [--explicit_args_count NUM_ARGS_FOR_EXTERNALS] \\" << std::endl

     // This option is most useful when using `--explicit_args` (or
     // `--legacy_mode`, which enables `--explicit_args`). In general, McSema
     // doesn't have type information about externals, and so it assumes all
     // externals operate on integer-typed arguments, and return integer values.
     // This is wrong in many ways, but tends to work out about 80% of the time.
     // To get McSema better information about externals, one should create a
     // C or C++ file with the declarations of the externals (perhaps by
     // `#include`ing standard headers). Then, should add to this file something
     // like:
     //         __attribute__((used))
     //         void *__mcsema_used_funcs[] = {
     //           (void *) external_func_name_1,
     //           (void *) external_func_name_2,
     //           ...
     //         };
     // And compile this file to bitcode using `remill-clang-M.m` (Major.minor).
     // This bitcode file will then be the source of type information for
     // McSema.
     << "    [--library BITCODE_FILE] \\" << std::endl

     // Annotate each LLVM IR instruction with some metadata that includes the
     // original program counter. The name of the LLVM metadats is
     // `PC_METADATA_ID`. This is enabled by default with `--legacy_mode`,
     // which sets `--pc_annotation` to be `mcsema_real_eip`.
     << "    [--pc_annotation PC_METADATA_ID] \\" << std::endl

     // Try to produce bitcode that looks like McSema version 1. This enables
     // `--explicit_args` and `--pc_annotation`.
     << "    [--legacy_mode] \\" << std::endl
     << "    [--version]" << std::endl
     << std::endl;

  google::InitGoogleLogging(argv[0]);
  google::SetUsageMessage(ss.str());
  google::ParseCommandLineFlags(&argc, &argv, true);

  if (FLAGS_version) {
    PrintVersion();
    return EXIT_SUCCESS;
  }

  CHECK(!FLAGS_os.empty())
      << "Must specify an operating system name to --os.";

  CHECK(!FLAGS_arch.empty())
      << "Must specify a machine code architecture name to --arch.";

  CHECK(!FLAGS_cfg.empty())
      << "Must specify the path to a CFG file to --cfg.";

  mcsema::gContext = new llvm::LLVMContext;

  CHECK(mcsema::InitArch(FLAGS_os, FLAGS_arch))
      << "Cannot initialize for arch " << FLAGS_arch
      << " and OS " << FLAGS_os << std::endl;

  if (FLAGS_legacy_mode) {
    LOG_IF(WARNING, FLAGS_keep_memops)
        << "Disabling --keep_memops in legacy mode.";
    FLAGS_keep_memops = false;

    LOG_IF(WARNING, !FLAGS_explicit_args)
        << "Enabling --explicit_args in legacy mode.";
    FLAGS_explicit_args = true;

    LOG_IF(WARNING, !FLAGS_pc_annotation.empty())
        << "Changing --pc_annotation to mcsema_real_eip in legacy mode.";
    FLAGS_pc_annotation = "mcsema_real_eip";

    LOG_IF(WARNING, FLAGS_disable_optimizer)
        << "Re-enabling the optimizer in legacy mode.";
    FLAGS_disable_optimizer = false;
  }

  auto cfg_module = mcsema::ReadProtoBuf(
      FLAGS_cfg, (mcsema::gArch->address_size / 8));
  mcsema::gModule = remill::LoadTargetSemantics(mcsema::gContext);
  mcsema::gArch->PrepareModule(mcsema::gModule);

  if (!FLAGS_library.empty()) {
    LoadLibraryIntoModule();
  }

  CHECK(mcsema::LiftCodeIntoModule(cfg_module))
      << "Unable to lift CFG from " << FLAGS_cfg << " into module "
      << FLAGS_output;

  remill::StoreModuleToFile(mcsema::gModule, FLAGS_output);
  google::ShutDownCommandLineFlags();
  google::ShutdownGoogleLogging();

  return EXIT_SUCCESS;
}
<commit_msg>Bring back the mcsema-lift option for --list-supported (#365)<commit_after>/*
 * Copyright (c) 2017 Trail of Bits, Inc.
 *
 * 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 <glog/logging.h>
#include <gflags/gflags.h>

#include <iomanip>
#include <iostream>
#include <memory>
#include <string>
#include <sstream>

#include <llvm/IR/Function.h>
#include <llvm/IR/GlobalVariable.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>

#include <remill/Arch/Arch.h>
#include <remill/BC/Util.h>
#include <remill/BC/Version.h>

#include "mcsema/Arch/Arch.h"
#include "mcsema/BC/Lift.h"
#include "mcsema/BC/Util.h"

#ifndef LLVM_VERSION_STRING
# define LLVM_VERSION_STRING LLVM_VERSION_MAJOR << "." << LLVM_VERSION_MINOR
#endif

#ifndef MCSEMA_VERSION_STRING
# define MCSEMA_VERSION_STRING "unknown"
#endif  // MCSEMA_VERSION_STRING

#ifndef MCSEMA_BRANCH_NAME
# define MCSEMA_BRANCH_NAME "unknown"
#endif  // MCSEMA_BRANCH_NAME

DECLARE_string(arch);
DECLARE_string(os);

DEFINE_string(cfg, "", "Path to the CFG file containing code to lift.");

DEFINE_string(output, "", "Output bitcode file name.");

DEFINE_string(library, "", "Path to an LLVM bitcode or IR file that contains "
                           "external library definitions.");

DECLARE_bool(version);

DECLARE_bool(disable_optimizer);
DECLARE_bool(keep_memops);
DECLARE_bool(explicit_args);
DECLARE_string(pc_annotation);

DEFINE_bool(list_supported, false,
            "List instructions that can be lifted.");
DEFINE_bool(legacy_mode, false,
            "Try to make the output bitcode resemble the original McSema.");

namespace {

static void PrintVersion(void) {
  std::cout
      << "This is mcsema-lift version: " << MCSEMA_VERSION_STRING << std::endl
      << "Built from branch: " << MCSEMA_BRANCH_NAME << std::endl
      << "Using LLVM " << LLVM_VERSION_STRING << std::endl;
}

// Print a list of instructions that Remill can lift.
static void PrintSupportedInstructions(void) {
  remill::ForEachISel(mcsema::gModule,
                      [=](llvm::GlobalVariable *isel, llvm::Function *) {
                        std::cout << isel->getName().str() << std::endl;
                      });
}

// Load in a separate bitcode or IR library, and copy function and variable
// declarations from that library into our module. We can use this feature
// to provide better type information to McSema.
static void LoadLibraryIntoModule(void) {
  std::unique_ptr<llvm::Module> lib(
      remill::LoadModuleFromFile(mcsema::gContext, FLAGS_library));

  // Declare the functions from the library in McSema's target module.
  for (auto &func : *lib) {
    auto func_name = func.getName();
    if (func_name.startswith("__mcsema") || func_name.startswith("__remill")) {
      continue;
    }
    if (mcsema::gModule->getFunction(func_name)) {
      continue;
    }

    auto dest_func = llvm::Function::Create(
        func.getFunctionType(), func.getLinkage(),
        func_name, mcsema::gModule);

    dest_func->copyAttributesFrom(&func);
    dest_func->setVisibility(func.getVisibility());
  }

  // Declare the global variables from the library in McSema's target module.
  for (auto &var : lib->globals()) {
    auto var_name = var.getName();
    if (var_name.startswith("__mcsema") || var_name.startswith("__remill")) {
      continue;
    }

    if (mcsema::gModule->getGlobalVariable(var_name)) {
      continue;
    }

    auto dest_var = new llvm::GlobalVariable(
        *mcsema::gModule, var.getType()->getElementType(),
        var.isConstant(), var.getLinkage(), nullptr,
        var_name, nullptr, var.getThreadLocalMode(),
        var.getType()->getAddressSpace());

    dest_var->copyAttributesFrom(&var);
  }
}

}  // namespace

int main(int argc, char *argv[]) {
  std::stringstream ss;
  ss << std::endl << std::endl
     << "  " << argv[0] << " \\" << std::endl
     << "    --output OUTPUT_BC_FILE \\" << std::endl
     << "    --arch ARCH_NAME \\" << std::endl
     << "    --os OS_NAME \\" << std::endl
     << "    --cfg CFG_FILE \\" << std::endl

     // This option is very useful for debugging McSema-lifted bitcode. It
     // injects so-called breakpoint functions before every lifted instruction.
     // For example, the the instruction at PC `0xf00` is lifted, then this
     // option will inject a call to `breakpoint_f00`. With this feature, we
     // can add breakpoints in a debugger on these breakpoint functions, and
     // know that they correspond to locations in the original program.
     << "    [--add_breakpoints] \\" << std::endl

     // This option injects a function call before every lifted instruction.
     // This function is implemented in the McSema runtime and it prints the
     // values of the general purpose registers to `stderr`.
     << "    [--add_reg_tracer] \\" << std::endl

     // This option tells McSema not to optimize the bitcode. This is useful
     // for debugging, especially in conjunction with `--add_breakpoints`.
     << "    [--disable_optimizer] \\" << std::endl

     // This option tells McSema not to lower Remill's memory access intrinsic
     // functions into LLVM `load` and `store` instructions.
     << "    [--keep_memops] \\" << std::endl

     // There are roughly two ways of using McSema-lifted bitcode. The default
     // use case is to compile the bitcode into an executable that behaves like
     // the original program. The other use case is to do some kind of static
     // analysis, e.g. with KLEE. In this use case, calls to external functions
     // are emulated so that we also try to explicitly lift parameter passing.
     // This mode of passing arguments explicitly is enabled by
     // `--explicit_args`, and in situations where we have no knowledge of the
     // argument counts expected by an external function, we fall back on
     // passing `--explicit_args_count` number of arguments to that function.
     << "    [--explicit_args] \\" << std::endl
     << "    [--explicit_args_count NUM_ARGS_FOR_EXTERNALS] \\" << std::endl

     // This option is most useful when using `--explicit_args` (or
     // `--legacy_mode`, which enables `--explicit_args`). In general, McSema
     // doesn't have type information about externals, and so it assumes all
     // externals operate on integer-typed arguments, and return integer values.
     // This is wrong in many ways, but tends to work out about 80% of the time.
     // To get McSema better information about externals, one should create a
     // C or C++ file with the declarations of the externals (perhaps by
     // `#include`ing standard headers). Then, should add to this file something
     // like:
     //         __attribute__((used))
     //         void *__mcsema_used_funcs[] = {
     //           (void *) external_func_name_1,
     //           (void *) external_func_name_2,
     //           ...
     //         };
     // And compile this file to bitcode using `remill-clang-M.m` (Major.minor).
     // This bitcode file will then be the source of type information for
     // McSema.
     << "    [--library BITCODE_FILE] \\" << std::endl

     // Annotate each LLVM IR instruction with some metadata that includes the
     // original program counter. The name of the LLVM metadats is
     // `PC_METADATA_ID`. This is enabled by default with `--legacy_mode`,
     // which sets `--pc_annotation` to be `mcsema_real_eip`.
     << "    [--pc_annotation PC_METADATA_ID] \\" << std::endl

     // Try to produce bitcode that looks like McSema version 1. This enables
     // `--explicit_args` and `--pc_annotation`.
     << "    [--legacy_mode] \\" << std::endl

     // Print a list of the instructions that can be lifted.
     << "    [--list-supported]" << std::endl

     // Print the version and exit.
     << "    [--version]" << std::endl
     << std::endl;

  google::InitGoogleLogging(argv[0]);
  google::SetUsageMessage(ss.str());
  google::ParseCommandLineFlags(&argc, &argv, true);

  if (FLAGS_version) {
    PrintVersion();
    return EXIT_SUCCESS;
  }

  CHECK(!FLAGS_os.empty())
      << "Must specify an operating system name to --os.";

  CHECK(!FLAGS_arch.empty())
      << "Must specify a machine code architecture name to --arch.";

  CHECK(!FLAGS_cfg.empty())
      << "Must specify the path to a CFG file to --cfg.";

  mcsema::gContext = new llvm::LLVMContext;

  CHECK(mcsema::InitArch(FLAGS_os, FLAGS_arch))
      << "Cannot initialize for arch " << FLAGS_arch
      << " and OS " << FLAGS_os << std::endl;

  if (FLAGS_legacy_mode) {
    LOG_IF(WARNING, FLAGS_keep_memops)
        << "Disabling --keep_memops in legacy mode.";
    FLAGS_keep_memops = false;

    LOG_IF(WARNING, !FLAGS_explicit_args)
        << "Enabling --explicit_args in legacy mode.";
    FLAGS_explicit_args = true;

    LOG_IF(WARNING, !FLAGS_pc_annotation.empty())
        << "Changing --pc_annotation to mcsema_real_eip in legacy mode.";
    FLAGS_pc_annotation = "mcsema_real_eip";

    LOG_IF(WARNING, FLAGS_disable_optimizer)
        << "Re-enabling the optimizer in legacy mode.";
    FLAGS_disable_optimizer = false;
  }

  auto cfg_module = mcsema::ReadProtoBuf(
      FLAGS_cfg, (mcsema::gArch->address_size / 8));
  mcsema::gModule = remill::LoadTargetSemantics(mcsema::gContext);
  mcsema::gArch->PrepareModule(mcsema::gModule);

  if (FLAGS_list_supported) {
    PrintSupportedInstructions();
  }

  if (!FLAGS_library.empty()) {
    LoadLibraryIntoModule();
  }

  CHECK(mcsema::LiftCodeIntoModule(cfg_module))
      << "Unable to lift CFG from " << FLAGS_cfg << " into module "
      << FLAGS_output;

  remill::StoreModuleToFile(mcsema::gModule, FLAGS_output);
  google::ShutDownCommandLineFlags();
  google::ShutdownGoogleLogging();

  return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before>//
// This file is part of the Marble Virtual Globe.
//
// This program is free software licensed under the GNU LGPL. You can
// find a copy of this license in LICENSE.txt in the top directory of
// the source code.
//
// Copyright 2011      Dennis Nienhüser <earthwings@gentoo.org>
//

#include "job.h"
#include "logger.h"
#include "upload.h"

#include <QtCore/QDebug>
#include <QtCore/QDateTime>
#include <QtCore/QProcess>

Job::Job(const Region &region, const JobParameters &parameters, QObject *parent) :
    QObject(parent), m_status(Waiting), m_region(region), m_parameters(parameters)
{
    // nothing to do
}

Job::Status Job::status() const
{
    return m_status;
}

QString Job::statusMessage() const
{
    return m_statusMessage;
}

Region Job::region() const
{
    return m_region;
}

void Job::setTransport(const QString &transport)
{
    m_transport = transport;
}

void Job::setProfile(const QString &profile)
{
    m_profile = profile;
}

void Job::setMonavSettings(const QString &filename)
{
    m_monavSettings = filename;
}

bool Job::operator ==(const Job &other) const
{
    return m_region == other.m_region;
}

void Job::run()
{
    if (download() && monav() && search() && package() && upload()) {
        // Nothing to do.
    }

    cleanup();
    emit finished(this);
}

void Job::changeStatus(Job::Status status, const QString &message)
{
    QString statusType;
    switch (status) {
    case Waiting: statusType = "waiting"; break;
    case Downloading: statusType = "downloading"; break;
    case Routing: statusType = "routing"; break;
    case Search: statusType = "search"; break;
    case Packaging: statusType = "packaging"; break;
    case Uploading: statusType = "uploading"; break;
    case Finished: statusType = "finished"; break;
    case Error: statusType = "error"; break;
    }

    Logger::instance().setStatus(m_region.id(), m_region.name(), statusType, message);
    m_statusMessage = message;
    m_status = status;
}

bool Job::download()
{
    changeStatus(Downloading, "Downloading data.");
    qDebug() << "Saving file to " << osmFile().absoluteFilePath();
    if (osmFile().exists()) {
        QDateTime now = QDateTime::currentDateTime();
        if (osmFile().lastModified().daysTo(now) > 7) {
            qDebug() << "Old file is outdated, re-downloading " << osmFile().absoluteFilePath();
            QFile::remove(osmFile().absoluteFilePath());
        } else {
            qDebug() << "Old file is still ok, reusing" << osmFile().absoluteFilePath();
            return true;
        }
    }

    QProcess wget;
    QStringList arguments;
    QString url = "http://download.geofabrik.de/osm/";
    url += m_region.pbfFile();
    arguments << "-O" << osmFile().absoluteFilePath() << url;
    qDebug() << "Downloading " << url;
    wget.start("wget", arguments);
    wget.waitForFinished(1000 * 60 * 60 * 12); // wait up to 12 hours for download to complete
    if (wget.exitStatus() == QProcess::NormalExit && wget.exitCode() == 0) {
        return true;
    } else {
        qDebug() << "Failed to download " << url;
        QFile::remove(osmFile().absoluteFilePath());
        changeStatus(Error, "Error downloading .osm.pbf file: " + wget.readAllStandardError());
        return false;
    }
}

//bool Job::marble()
//{
//    changeStatus(Routing, "Extracting bounding box.");
//    QStringList arguments;
//    arguments << "--name" << m_region.name();
//    arguments << "--version" << "0.2";
//    arguments << "--date" << QDateTime::currentDateTime().toString("yyyy/dd/MM");
//    arguments << "--transport" << m_transport;
//    arguments << "--payload" << targetFile().fileName();
//    arguments << m_parameters.base().absoluteFilePath("poly/" + m_region.polyFile());
//    arguments << monavDir().absoluteFilePath() + "/marble.kml";
//    QProcess poly2kml;
//    poly2kml.start("poly2kml", arguments);
//    poly2kml.waitForFinished(1000 * 60 * 30); // wait up to half an hour for poly2kml to convert the data
//    if (poly2kml.exitStatus() == QProcess::NormalExit && poly2kml.exitCode() == 0) {
//        qDebug() << "Processed kml file for marble";
//        return true;
//    } else {
//        qDebug() << "poly2kml exiting with status " << poly2kml.exitCode();
//        changeStatus(Error, "Error creating marble.kml: " + poly2kml.readAllStandardError());
//        return false;
//    }
//}

bool Job::monav()
{
    QString const status = QString("Generating offline routing map from %1 (%2).").arg(osmFile().fileName()).arg(Region::fileSize(osmFile()));
    changeStatus(Routing, status);
    QStringList arguments;
    arguments << "-s=" + m_monavSettings;
    arguments << "-i=" + osmFile().absoluteFilePath();
    arguments << "-o=" + monavDir().absoluteFilePath();
    arguments << "-pi=OpenStreetMap Importer" << "-pro=Contraction Hierarchies";
    arguments << "-pg=GPS Grid" << "-di";
    arguments << "-dro=" + m_transport;
    arguments << "--profile=" + m_profile;
    arguments << "-dd" /*<< "-dc"*/;
    QProcess monav;
    monav.start("monav-preprocessor", arguments);
    monav.waitForFinished(1000 * 60 * 60 * 2); // wait up to 2 hours for monav to convert the data
    if (monav.exitStatus() == QProcess::NormalExit && monav.exitCode() == 0) {
        qDebug() << "Processed osm file for monav";
    } else {
        qDebug() << "monav exiting with status " << monav.exitCode();
        changeStatus(Error, "Routing map conversion failed: " + monav.readAllStandardError());
        return false;
    }

    QFile pluginsFile(monavDir().absoluteFilePath() + "/plugins.ini");
    pluginsFile.open(QFile::WriteOnly | QFile::Truncate);
    QTextStream pluginsStream(&pluginsFile);
    pluginsStream << "[General]\nrouter=Contraction Hierarchies\nrenderer=Mapnik Renderer\ngpsLookup=GPS Grid\naddressLookup=Unicode Tournament Trie\n";
    pluginsFile.close();

    QFileInfo subdir = QFileInfo(monavDir().absoluteFilePath() + "/routing_" + m_transport.toLower());
    if (subdir.exists() && subdir.isDir()) {
        QFileInfoList files = QDir(subdir.absoluteFilePath()).entryInfoList(QDir::Files);
        foreach(const QFileInfo &file, files) {
            if (!QFile::rename(file.absoluteFilePath(), monavDir().absoluteFilePath() + "/" + file.fileName())) {
                changeStatus(Error, "Unable to move monav files to target directory.");
                return false;
            }
        }
        QDir("/").rmdir(subdir.absoluteFilePath());
    } else {
        changeStatus(Error, "Unable to find files created by monav");
        return false;
    }

    return true;
}

bool Job::search()
{
    QString const status = QString("Generating offline search database from %1 (%2).").arg(osmFile().fileName()).arg(Region::fileSize(osmFile()));
    changeStatus(Search, status);
    QStringList arguments;
    arguments << "--name" << m_region.name();
    arguments << "--version" << "0.3";
    arguments << "--date" << QDateTime::currentDateTime().toString("MM/dd/yy");
    arguments << "--transport" << m_transport;
    arguments << "--payload" << targetFile().fileName();
    arguments << osmFile().absoluteFilePath();
    arguments << searchFile().absoluteFilePath();
    QFileInfo kmlFile(monavDir().absoluteFilePath() + "/marble.kml");
    arguments << kmlFile.absoluteFilePath();
    QProcess osmAddresses;
    osmAddresses.start("osm-addresses", arguments);
    osmAddresses.waitForFinished(1000 * 60 * 60 * 12); // wait up to 12 hours for osm-addresses to convert the data
    if (osmAddresses.exitStatus() == QProcess::NormalExit && osmAddresses.exitCode() == 0) {
        searchFile().refresh();
        if (!searchFile().exists()) {
            qDebug() << "osm-addresses did not create the .sqlite file";
            changeStatus(Error, "Unknown error when creating the search database");
            return false;
        } else if (searchFile().size() < 8000) {
            qDebug() << "The .sqlite database has a suspiciously small size.";
            changeStatus(Error, "Search database is too small. Too little memory?");
            return false;
        }

        kmlFile.refresh();
        if (!kmlFile.exists()) {
            qDebug() << "File marble.kml has not been generated.";
            changeStatus(Error, "Failed to generate marble.kml. Too little memory?");
            return false;
        }

        return true;
    } else {
        qDebug() << "osm-addresses exiting with status " << osmAddresses.exitCode();
        changeStatus(Error, "Error creating search database: " + osmAddresses.readAllStandardError());
        return false;
    }
}

bool Job::package()
{
    changeStatus(Packaging, "Creating archive.");
    QStringList arguments;
    arguments << "czf" << targetFile().absoluteFilePath() << "earth/monav/" << "earth/placemarks";
    QProcess tar;
    tar.setWorkingDirectory(m_parameters.base().absolutePath() + "/data/" + m_region.id());
    tar.start("tar", arguments);
    tar.waitForFinished(1000 * 60 * 60); // wait up to 1 hour for tar to package things
    if (tar.exitStatus() == QProcess::NormalExit && tar.exitCode() == 0) {
        qDebug() << "Packaged tar file";
        return true;
    } else {
        changeStatus(Error, "Packaging failed: " + tar.readAllStandardError());
        return false;
    }
}

bool Job::upload()
{
    changeStatus(Uploading, "Uploading file");
    if (targetFile().exists()) {
        Upload::instance().uploadAndDelete(m_region, targetFile(), m_transport);
        return true;
    }

    changeStatus(Error, "Target file does not exist.");
    return false;
}

bool Job::cleanup()
{
    if (!m_parameters.cacheData()) {
        QFile::remove(osmFile().absoluteFilePath());
    }

    QFileInfo subdir = QFileInfo(monavDir().absoluteFilePath());
    if (subdir.exists() && subdir.isDir()) {
        QFileInfoList files = QDir(subdir.absoluteFilePath()).entryInfoList(QDir::Files);
        foreach(const QFileInfo &file, files) {
            QFile::remove(file.absoluteFilePath());
        }
    }

    QFile::remove(searchFile().absoluteFilePath());
    return true;
}

QFileInfo Job::osmFile()
{
    m_parameters.base().mkdir("download");
    QFileInfo result(m_parameters.base(), QString("download/") + m_region.id() + ".osm.pbf");
    return result;
}

QFileInfo Job::monavDir()
{
    QString const subdir = "data/" + m_region.id() + "/earth/monav/" + m_transport.toLower() + "/" + m_region.path();
    m_parameters.base().mkpath(subdir);
    QFileInfo result(m_parameters.base(), subdir);
    return result;
}

QFileInfo Job::targetFile()
{
    m_parameters.base().mkdir("finished");
    QFileInfo result(m_parameters.base(), "finished/" + m_region.id() + "_" + m_transport.toLower() + ".tar.gz");
    return result;
}

QFileInfo Job::searchFile()
{
    QString const subdir = "data/" + m_region.id() + "/earth/placemarks/" + QFileInfo(m_region.path()).path();
    m_parameters.base().mkpath(subdir);
    QFileInfo result(m_parameters.base(), subdir + "/" + m_region.id() + ".sqlite");
    return result;
}
<commit_msg>Expect full URL (to support more than one server)<commit_after>//
// This file is part of the Marble Virtual Globe.
//
// This program is free software licensed under the GNU LGPL. You can
// find a copy of this license in LICENSE.txt in the top directory of
// the source code.
//
// Copyright 2011      Dennis Nienhüser <earthwings@gentoo.org>
//

#include "job.h"
#include "logger.h"
#include "upload.h"

#include <QtCore/QDebug>
#include <QtCore/QDateTime>
#include <QtCore/QProcess>

Job::Job(const Region &region, const JobParameters &parameters, QObject *parent) :
    QObject(parent), m_status(Waiting), m_region(region), m_parameters(parameters)
{
    // nothing to do
}

Job::Status Job::status() const
{
    return m_status;
}

QString Job::statusMessage() const
{
    return m_statusMessage;
}

Region Job::region() const
{
    return m_region;
}

void Job::setTransport(const QString &transport)
{
    m_transport = transport;
}

void Job::setProfile(const QString &profile)
{
    m_profile = profile;
}

void Job::setMonavSettings(const QString &filename)
{
    m_monavSettings = filename;
}

bool Job::operator ==(const Job &other) const
{
    return m_region == other.m_region;
}

void Job::run()
{
    if (download() && monav() && search() && package() && upload()) {
        // Nothing to do.
    }

    cleanup();
    emit finished(this);
}

void Job::changeStatus(Job::Status status, const QString &message)
{
    QString statusType;
    switch (status) {
    case Waiting: statusType = "waiting"; break;
    case Downloading: statusType = "downloading"; break;
    case Routing: statusType = "routing"; break;
    case Search: statusType = "search"; break;
    case Packaging: statusType = "packaging"; break;
    case Uploading: statusType = "uploading"; break;
    case Finished: statusType = "finished"; break;
    case Error: statusType = "error"; break;
    }

    Logger::instance().setStatus(m_region.id(), m_region.name(), statusType, message);
    m_statusMessage = message;
    m_status = status;
}

bool Job::download()
{
    changeStatus(Downloading, "Downloading data.");
    qDebug() << "Saving file to " << osmFile().absoluteFilePath();
    if (osmFile().exists()) {
        QDateTime now = QDateTime::currentDateTime();
        if (osmFile().lastModified().daysTo(now) > 7) {
            qDebug() << "Old file is outdated, re-downloading " << osmFile().absoluteFilePath();
            QFile::remove(osmFile().absoluteFilePath());
        } else {
            qDebug() << "Old file is still ok, reusing" << osmFile().absoluteFilePath();
            return true;
        }
    }

    QProcess wget;
    QStringList arguments;
    QString url = m_region.pbfFile();
    arguments << "-O" << osmFile().absoluteFilePath() << url;
    qDebug() << "Downloading " << url;
    wget.start("wget", arguments);
    wget.waitForFinished(1000 * 60 * 60 * 12); // wait up to 12 hours for download to complete
    if (wget.exitStatus() == QProcess::NormalExit && wget.exitCode() == 0) {
        return true;
    } else {
        qDebug() << "Failed to download " << url;
        QFile::remove(osmFile().absoluteFilePath());
        changeStatus(Error, "Error downloading .osm.pbf file: " + wget.readAllStandardError());
        return false;
    }
}

//bool Job::marble()
//{
//    changeStatus(Routing, "Extracting bounding box.");
//    QStringList arguments;
//    arguments << "--name" << m_region.name();
//    arguments << "--version" << "0.2";
//    arguments << "--date" << QDateTime::currentDateTime().toString("yyyy/dd/MM");
//    arguments << "--transport" << m_transport;
//    arguments << "--payload" << targetFile().fileName();
//    arguments << m_parameters.base().absoluteFilePath("poly/" + m_region.polyFile());
//    arguments << monavDir().absoluteFilePath() + "/marble.kml";
//    QProcess poly2kml;
//    poly2kml.start("poly2kml", arguments);
//    poly2kml.waitForFinished(1000 * 60 * 30); // wait up to half an hour for poly2kml to convert the data
//    if (poly2kml.exitStatus() == QProcess::NormalExit && poly2kml.exitCode() == 0) {
//        qDebug() << "Processed kml file for marble";
//        return true;
//    } else {
//        qDebug() << "poly2kml exiting with status " << poly2kml.exitCode();
//        changeStatus(Error, "Error creating marble.kml: " + poly2kml.readAllStandardError());
//        return false;
//    }
//}

bool Job::monav()
{
    QString const status = QString("Generating offline routing map from %1 (%2).").arg(osmFile().fileName()).arg(Region::fileSize(osmFile()));
    changeStatus(Routing, status);
    QStringList arguments;
    arguments << "-s=" + m_monavSettings;
    arguments << "-i=" + osmFile().absoluteFilePath();
    arguments << "-o=" + monavDir().absoluteFilePath();
    arguments << "-pi=OpenStreetMap Importer" << "-pro=Contraction Hierarchies";
    arguments << "-pg=GPS Grid" << "-di";
    arguments << "-dro=" + m_transport;
    arguments << "--profile=" + m_profile;
    arguments << "-dd" /*<< "-dc"*/;
    QProcess monav;
    monav.start("monav-preprocessor", arguments);
    monav.waitForFinished(1000 * 60 * 60 * 2); // wait up to 2 hours for monav to convert the data
    if (monav.exitStatus() == QProcess::NormalExit && monav.exitCode() == 0) {
        qDebug() << "Processed osm file for monav";
    } else {
        qDebug() << "monav exiting with status " << monav.exitCode();
        changeStatus(Error, "Routing map conversion failed: " + monav.readAllStandardError());
        return false;
    }

    QFile pluginsFile(monavDir().absoluteFilePath() + "/plugins.ini");
    pluginsFile.open(QFile::WriteOnly | QFile::Truncate);
    QTextStream pluginsStream(&pluginsFile);
    pluginsStream << "[General]\nrouter=Contraction Hierarchies\nrenderer=Mapnik Renderer\ngpsLookup=GPS Grid\naddressLookup=Unicode Tournament Trie\n";
    pluginsFile.close();

    QFileInfo subdir = QFileInfo(monavDir().absoluteFilePath() + "/routing_" + m_transport.toLower());
    if (subdir.exists() && subdir.isDir()) {
        QFileInfoList files = QDir(subdir.absoluteFilePath()).entryInfoList(QDir::Files);
        foreach(const QFileInfo &file, files) {
            if (!QFile::rename(file.absoluteFilePath(), monavDir().absoluteFilePath() + "/" + file.fileName())) {
                changeStatus(Error, "Unable to move monav files to target directory.");
                return false;
            }
        }
        QDir("/").rmdir(subdir.absoluteFilePath());
    } else {
        changeStatus(Error, "Unable to find files created by monav");
        return false;
    }

    return true;
}

bool Job::search()
{
    QString const status = QString("Generating offline search database from %1 (%2).").arg(osmFile().fileName()).arg(Region::fileSize(osmFile()));
    changeStatus(Search, status);
    QStringList arguments;
    arguments << "--name" << m_region.name();
    arguments << "--version" << "0.3";
    arguments << "--date" << QDateTime::currentDateTime().toString("MM/dd/yy");
    arguments << "--transport" << m_transport;
    arguments << "--payload" << targetFile().fileName();
    arguments << osmFile().absoluteFilePath();
    arguments << searchFile().absoluteFilePath();
    QFileInfo kmlFile(monavDir().absoluteFilePath() + "/marble.kml");
    arguments << kmlFile.absoluteFilePath();
    QProcess osmAddresses;
    osmAddresses.start("osm-addresses", arguments);
    osmAddresses.waitForFinished(1000 * 60 * 60 * 12); // wait up to 12 hours for osm-addresses to convert the data
    if (osmAddresses.exitStatus() == QProcess::NormalExit && osmAddresses.exitCode() == 0) {
        searchFile().refresh();
        if (!searchFile().exists()) {
            qDebug() << "osm-addresses did not create the .sqlite file";
            changeStatus(Error, "Unknown error when creating the search database");
            return false;
        } else if (searchFile().size() < 8000) {
            qDebug() << "The .sqlite database has a suspiciously small size.";
            changeStatus(Error, "Search database is too small. Too little memory?");
            return false;
        }

        kmlFile.refresh();
        if (!kmlFile.exists()) {
            qDebug() << "File marble.kml has not been generated.";
            changeStatus(Error, "Failed to generate marble.kml. Too little memory?");
            return false;
        }

        return true;
    } else {
        qDebug() << "osm-addresses exiting with status " << osmAddresses.exitCode();
        changeStatus(Error, "Error creating search database: " + osmAddresses.readAllStandardError());
        return false;
    }
}

bool Job::package()
{
    changeStatus(Packaging, "Creating archive.");
    QStringList arguments;
    arguments << "czf" << targetFile().absoluteFilePath() << "earth/monav/" << "earth/placemarks";
    QProcess tar;
    tar.setWorkingDirectory(m_parameters.base().absolutePath() + "/data/" + m_region.id());
    tar.start("tar", arguments);
    tar.waitForFinished(1000 * 60 * 60); // wait up to 1 hour for tar to package things
    if (tar.exitStatus() == QProcess::NormalExit && tar.exitCode() == 0) {
        qDebug() << "Packaged tar file";
        return true;
    } else {
        changeStatus(Error, "Packaging failed: " + tar.readAllStandardError());
        return false;
    }
}

bool Job::upload()
{
    changeStatus(Uploading, "Uploading file");
    if (targetFile().exists()) {
        Upload::instance().uploadAndDelete(m_region, targetFile(), m_transport);
        return true;
    }

    changeStatus(Error, "Target file does not exist.");
    return false;
}

bool Job::cleanup()
{
    if (!m_parameters.cacheData()) {
        QFile::remove(osmFile().absoluteFilePath());
    }

    QFileInfo subdir = QFileInfo(monavDir().absoluteFilePath());
    if (subdir.exists() && subdir.isDir()) {
        QFileInfoList files = QDir(subdir.absoluteFilePath()).entryInfoList(QDir::Files);
        foreach(const QFileInfo &file, files) {
            QFile::remove(file.absoluteFilePath());
        }
    }

    QFile::remove(searchFile().absoluteFilePath());
    return true;
}

QFileInfo Job::osmFile()
{
    m_parameters.base().mkdir("download");
    QFileInfo result(m_parameters.base(), QString("download/") + m_region.id() + ".osm.pbf");
    return result;
}

QFileInfo Job::monavDir()
{
    QString const subdir = "data/" + m_region.id() + "/earth/monav/" + m_transport.toLower() + "/" + m_region.path();
    m_parameters.base().mkpath(subdir);
    QFileInfo result(m_parameters.base(), subdir);
    return result;
}

QFileInfo Job::targetFile()
{
    m_parameters.base().mkdir("finished");
    QFileInfo result(m_parameters.base(), "finished/" + m_region.id() + "_" + m_transport.toLower() + ".tar.gz");
    return result;
}

QFileInfo Job::searchFile()
{
    QString const subdir = "data/" + m_region.id() + "/earth/placemarks/" + QFileInfo(m_region.path()).path();
    m_parameters.base().mkpath(subdir);
    QFileInfo result(m_parameters.base(), subdir + "/" + m_region.id() + ".sqlite");
    return result;
}
<|endoftext|>