<commit_before>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2014-2016 ArangoDB GmbH, Cologne, Germany
/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
///
/// 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
///
///     vpp://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.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Jan Christoph Uhde
////////////////////////////////////////////////////////////////////////////////

#include "VppCommTask.h"

#include "Basics/StringBuffer.h"
#include "Basics/HybridLogicalClock.h"
#include "Basics/VelocyPackHelper.h"
#include "GeneralServer/GeneralServer.h"
#include "GeneralServer/GeneralServerFeature.h"
#include "GeneralServer/RestHandler.h"
#include "GeneralServer/RestHandlerFactory.h"
#include "Scheduler/Scheduler.h"
#include "Scheduler/SchedulerFeature.h"
#include "Logger/LoggerFeature.h"
#include "VocBase/ticks.h"

#include <velocypack/Validator.h>
#include <velocypack/velocypack-aliases.h>

#include <iostream>
#include <stdexcept>

using namespace arangodb;
using namespace arangodb::basics;
using namespace arangodb::rest;

namespace {
std::size_t findAndValidateVPacks(char const* vpHeaderStart,
                                  char const* chunkEnd) {
  VPackValidator validator;
  // check for slice start to the end of Chunk
  // isSubPart allows the slice to be shorter than the checked buffer.
  validator.validate(vpHeaderStart, std::distance(vpHeaderStart, chunkEnd),
                     /*isSubPart =*/true);

  // check if there is payload and locate the start
  VPackSlice vpHeader(vpHeaderStart);
  auto vpHeaderLen = vpHeader.byteSize();
  auto vpPayloadStart = vpHeaderStart + vpHeaderLen;
  if (vpPayloadStart == chunkEnd) {
    return 0;  // no payload available
  }

  // validate Payload VelocyPack
  validator.validate(vpPayloadStart, std::distance(vpPayloadStart, chunkEnd),
                     /*isSubPart =*/false);
  return std::distance(vpHeaderStart, vpPayloadStart);
}

std::unique_ptr<basics::StringBuffer> createChunkForNetworkDetail(
    std::vector<VPackSlice> const& slices, bool isFirstChunk, uint32_t chunk,
    uint64_t id, uint32_t totalMessageLength = 0) {
  using basics::StringBuffer;
  bool firstOfMany = false;

  // if we have more than one chunk and the chunk is the first
  // then we are sending the first in a series. if this condition
  // is true we also send extra 8 bytes for the messageLength
  // (length of all VPackData)
  if (isFirstChunk && chunk > 1) {
    firstOfMany = true;
  }

  // build chunkX -- see VelocyStream Documentaion
  chunk <<= 1;
  chunk |= isFirstChunk ? 0x1 : 0x0;

  // get the lenght of VPack data
  uint32_t dataLength = 0;
  for (auto& slice : slices) {
    dataLength += slice.byteSize();
  }

  // calculate length of current chunk
  uint32_t chunkLength = dataLength;
  chunkLength += (sizeof(chunkLength) + sizeof(chunk) + sizeof(id));
  if (firstOfMany) {
    chunkLength += sizeof(totalMessageLength);
  }

  auto buffer =
      std::make_unique<StringBuffer>(TRI_UNKNOWN_MEM_ZONE, chunkLength, false);

  char cChunkLength[sizeof(chunkLength)];
  char const* cChunkLengthPtr = cChunkLength;
  std::memcpy(&cChunkLength, &chunkLength, sizeof(chunkLength));
  buffer->appendText(cChunkLengthPtr, sizeof(chunkLength));

  char cChunk[sizeof(chunk)];
  char const* cChunkPtr = cChunk;
  std::memcpy(&cChunk, &chunk, sizeof(chunk));
  buffer->appendText(cChunkPtr, sizeof(chunk));

  char cId[sizeof(id)];
  char const* cIdPtr = cId;
  std::memcpy(&cId, &id, sizeof(id));
  buffer->appendText(cIdPtr, sizeof(id));

  if (firstOfMany) {
    char cTotalMessageLength[sizeof(totalMessageLength)];
    char const* cTotalMessageLengthPtr = cTotalMessageLength;
    std::memcpy(&cTotalMessageLength, &totalMessageLength,
                sizeof(totalMessageLength));
    buffer->appendText(cTotalMessageLengthPtr, sizeof(totalMessageLength));
  }

  // append data in slices
  for (auto const& slice : slices) {
    buffer->appendText(std::string(slice.startAs<char>(), slice.byteSize()));
  }

  return buffer;
}

std::unique_ptr<basics::StringBuffer> createChunkForNetworkSingle(
    std::vector<VPackSlice> const& slices, uint64_t id) {
  return createChunkForNetworkDetail(slices, true, 1, id, 0 /*unused*/);
}

// TODO FIXME make use of these functions
// std::unique_ptr<basics::StringBuffer> createChunkForNetworkMultiFirst(
//     std::vector<VPackSlice> const& slices, uint64_t id, uint32_t
//     numberOfChunks,
//     uint32_t totalMessageLength) {
//   return createChunkForNetworkDetail(slices, true, numberOfChunks, id,
//                                      totalMessageLength);
// }
//
// std::unique_ptr<basics::StringBuffer> createChunkForNetworkMultiFollow(
//     std::vector<VPackSlice> const& slices, uint64_t id, uint32_t chunkNumber,
//     uint32_t totalMessageLength) {
//   return createChunkForNetworkDetail(slices, false, chunkNumber, id, 0);
// }
}

VppCommTask::VppCommTask(GeneralServer* server, TRI_socket_t sock,
                         ConnectionInfo&& info, double timeout)
    : Task("VppCommTask"),
      GeneralCommTask(server, sock, std::move(info), timeout) {
  _protocol = "vpp";
  _readBuffer->reserve(
      _bufferLength);  // ATTENTION <- this is required so we do not
                       // loose information during a resize
                       // connectionStatisticsAgentSetVpp();
}

void VppCommTask::addResponse(VppResponse* response, bool isError) {
  if (isError) {
    // FIXME (obi)
    // what do we need to do?
    // clean read buffer? reset process read cursor
  }

  VPackMessageNoOwnBuffer response_message = response->prepareForNetwork();
  uint64_t& id = response_message._id;

  std::vector<VPackSlice> slices;
  slices.push_back(response_message._header);

  VPackBuilder builder;
  if (response_message._generateBody) {
    builder = basics::VelocyPackHelper::stanitizeExternalsChecked(
        response_message._payload);
    slices.push_back(builder.slice());
  }

  // FIXME (obi)
  // If the message is big we will create many small chunks in a loop.
  // For the first tests we just send single Messages

  LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "got response:";
  for (auto const& slice : slices) {
    try {
      LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << slice.toJson();
    } catch (arangodb::velocypack::Exception const& e) {
      std::cout << "obi exception" << e.what();
    }
  }

  // adds chunk header infromation and creates SingBuffer* that can be
  // used with _writeBuffers
  auto buffer = createChunkForNetworkSingle(slices, id);
  _writeBuffers.push_back(buffer.get());
  buffer.release();

  fillWriteBuffer();  // move data from _writebuffers to _writebuffer
                      // implemented in base
}

VppCommTask::ChunkHeader VppCommTask::readChunkHeader() {
  VppCommTask::ChunkHeader header;

  auto cursor = _processReadVariables._readBufferCursor;

  std::memcpy(&header._chunkLength, cursor, sizeof(header._chunkLength));
  cursor += sizeof(header._chunkLength);

  uint32_t chunkX;
  std::memcpy(&chunkX, cursor, sizeof(chunkX));
  cursor += sizeof(chunkX);

  header._isFirst = chunkX & 0x1;
  header._chunk = chunkX >> 1;

  std::memcpy(&header._messageID, cursor, sizeof(header._messageID));
  cursor += sizeof(header._messageID);

  // extract total len of message
  if (header._isFirst && header._chunk > 1) {
    std::memcpy(&header._messageLength, cursor, sizeof(header._messageLength));
    cursor += sizeof(header._messageLength);
  } else {
    header._messageLength = 0;  // not needed
  }

  header._headerLength =
      std::distance(_processReadVariables._readBufferCursor, cursor);

  return header;
}

bool VppCommTask::isChunkComplete(char* start) {
  std::size_t length = std::distance(start, _readBuffer->end());
  auto& prv = _processReadVariables;

  if (!prv._currentChunkLength && length < sizeof(uint32_t)) {
    return false;
  }
  if (!prv._currentChunkLength) {
    // read chunk length
    std::memcpy(&prv._currentChunkLength, start, sizeof(uint32_t));
  }
  if (length < prv._currentChunkLength) {
    // chunk not complete
    return false;
  }

  return true;
}

// reads data from the socket
bool VppCommTask::processRead() {
  // TODO FIXME
  // - in case of error send an operation failed to all incomplete messages /
  //   operation and close connection (implement resetState/resetCommtask)

  auto& prv = _processReadVariables;
  if (!prv._readBufferCursor) {
    prv._readBufferCursor = _readBuffer->begin();
  }

  auto chunkBegin = prv._readBufferCursor;
  if (chunkBegin == nullptr || !isChunkComplete(chunkBegin)) {
    return false;  // no data or incomplete
  }

  ChunkHeader chunkHeader = readChunkHeader();
  auto chunkEnd = chunkBegin + chunkHeader._chunkLength;
  auto vpackBegin = chunkBegin + chunkHeader._headerLength;
  bool doExecute = false;
  bool read_maybe_only_part_of_buffer = false;
  VPackMessage message;  // filled in CASE 1 or CASE 2b

  // CASE 1: message is in one chunk
  if (chunkHeader._isFirst && chunkHeader._chunk == 1) {
    std::size_t payloadOffset = findAndValidateVPacks(vpackBegin, chunkEnd);
    message._id = chunkHeader._messageID;
    message._buffer.append(vpackBegin, std::distance(vpackBegin, chunkEnd));
    message._header = VPackSlice(message._buffer.data());
    if (payloadOffset) {
      message._payload = VPackSlice(message._buffer.data() + payloadOffset);
    }
    VPackValidator val;
    val.validate(message._header.begin(), message._header.byteSize());

    doExecute = true;
    LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 1";
  }
  // CASE 2:  message is in multiple chunks
  auto incompleteMessageItr = _incompleteMessages.find(chunkHeader._messageID);

  // CASE 2a: chunk starts new message
  if (chunkHeader._isFirst) {  // first chunk of multi chunk message
    if (incompleteMessageItr != _incompleteMessages.end()) {
      throw std::logic_error(
          "Message should be first but is already in the Map of incomplete "
          "messages");
    }

    IncompleteVPackMessage message(chunkHeader._messageLength,
                                   chunkHeader._chunk /*number of chunks*/);
    message._buffer.append(vpackBegin, std::distance(vpackBegin, chunkEnd));
    auto insertPair = _incompleteMessages.emplace(
        std::make_pair(chunkHeader._messageID, std::move(message)));
    if (!insertPair.second) {
      throw std::logic_error("insert failed");
    }

    LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 2a";

    // CASE 2b: chunk continues a message
  } else {  // followup chunk of some mesage
    if (incompleteMessageItr == _incompleteMessages.end()) {
      throw std::logic_error("found message without previous part");
    }
    auto& im = incompleteMessageItr->second;  // incomplete Message
    im._currentChunk++;
    assert(im._currentChunk == chunkHeader._chunk);
    im._buffer.append(vpackBegin, std::distance(vpackBegin, chunkEnd));
    // check buffer longer than length

    // MESSAGE COMPLETE
    if (im._currentChunk == im._numberOfChunks - 1 /* zero based counting */) {
      std::size_t payloadOffset = findAndValidateVPacks(
          reinterpret_cast<char const*>(im._buffer.data()),
          reinterpret_cast<char const*>(im._buffer.data() +
                                        im._buffer.byteSize()));
      message._id = chunkHeader._messageID;
      message._buffer = std::move(im._buffer);
      message._header = VPackSlice(message._buffer.data());
      if (payloadOffset) {
        message._payload = VPackSlice(message._buffer.data() + payloadOffset);
      }
      _incompleteMessages.erase(incompleteMessageItr);
      // check length

      doExecute = true;
      LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 2b - complete";
    }
    LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 2b - still incomplete";
  }

  // clean buffer up to length of chunk
  read_maybe_only_part_of_buffer = true;
  prv._currentChunkLength = 0;
  prv._readBufferCursor = chunkEnd;
  std::size_t processedDataLen =
      std::distance(_readBuffer->begin(), prv._readBufferCursor);
  if (processedDataLen > prv._cleanupLength) {
    _readBuffer->move_front(prv._cleanupLength);
    prv._readBufferCursor = nullptr;  // the positon will be set at the
                                      // begin of this function
  }

  if (doExecute) {
    //    return false;  // we have no complete request, so we return early
    // for now we can handle only one request at a time
    // lock _request???? REVIEW (fc)
    LOG_TOPIC(DEBUG, Logger::COMMUNICATION)
        << "got request:" << message._header.toJson();
    _request = new VppRequest(_connectionInfo, std::move(message));
    GeneralServerFeature::HANDLER_FACTORY->setRequestContext(_request);
    if (_request->requestContext() == nullptr) {
      handleSimpleError(GeneralResponse::ResponseCode::NOT_FOUND,
                        TRI_ERROR_ARANGO_DATABASE_NOT_FOUND,
                        TRI_errno_string(TRI_ERROR_ARANGO_DATABASE_NOT_FOUND));
    } else {
      _request->setClientTaskId(_taskId);
      _protocolVersion = _request->protocolVersion();
      executeRequest(
          _request, new VppResponse(GeneralResponse::ResponseCode::SERVER_ERROR,
                                    chunkHeader._messageID));
    }
  }

  if (read_maybe_only_part_of_buffer) {
    if (prv._readBufferCursor == _readBuffer->end()) {
      return false;
    }
    return true;
  }
  return doExecute;
}

void VppCommTask::completedWriteBuffer() {
  // REVIEW (fc)
}

void VppCommTask::resetState(bool close) {
  // REVIEW (fc)
  // is there a case where we do not want to close the connection
  replyToIncompleteMessages();
}

// GeneralResponse::ResponseCode VppCommTask::authenticateRequest() {
//   auto context = (_request == nullptr) ? nullptr :
//   _request->requestContext();
//
//   if (context == nullptr && _request != nullptr) {
//     bool res =
//         GeneralServerFeature::HANDLER_FACTORY->setRequestContext(_request);
//
//     if (!res) {
//       return GeneralResponse::ResponseCode::NOT_FOUND;
//     }
//
//     context = _request->requestContext();
//   }
//
//   if (context == nullptr) {
//     return GeneralResponse::ResponseCode::SERVER_ERROR;
//   }
//
//   return context->authenticate();
// }

// convert internal GeneralRequest to VppRequest
VppRequest* VppCommTask::requestAsVpp() {
  VppRequest* request = dynamic_cast<VppRequest*>(_request);
  if (request == nullptr) {
    THROW_ARANGO_EXCEPTION(TRI_ERROR_INTERNAL);
  }
  return request;
};
<commit_msg>fix buffer clean up<commit_after>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2014-2016 ArangoDB GmbH, Cologne, Germany
/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
///
/// 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
///
///     vpp://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.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Jan Christoph Uhde
////////////////////////////////////////////////////////////////////////////////

#include "VppCommTask.h"

#include "Basics/StringBuffer.h"
#include "Basics/HybridLogicalClock.h"
#include "Basics/VelocyPackHelper.h"
#include "GeneralServer/GeneralServer.h"
#include "GeneralServer/GeneralServerFeature.h"
#include "GeneralServer/RestHandler.h"
#include "GeneralServer/RestHandlerFactory.h"
#include "Scheduler/Scheduler.h"
#include "Scheduler/SchedulerFeature.h"
#include "Logger/LoggerFeature.h"
#include "VocBase/ticks.h"

#include <velocypack/Validator.h>
#include <velocypack/velocypack-aliases.h>

#include <iostream>
#include <stdexcept>

using namespace arangodb;
using namespace arangodb::basics;
using namespace arangodb::rest;

namespace {
std::size_t findAndValidateVPacks(char const* vpHeaderStart,
                                  char const* chunkEnd) {
  VPackValidator validator;
  // check for slice start to the end of Chunk
  // isSubPart allows the slice to be shorter than the checked buffer.
  validator.validate(vpHeaderStart, std::distance(vpHeaderStart, chunkEnd),
                     /*isSubPart =*/true);

  // check if there is payload and locate the start
  VPackSlice vpHeader(vpHeaderStart);
  auto vpHeaderLen = vpHeader.byteSize();
  auto vpPayloadStart = vpHeaderStart + vpHeaderLen;
  if (vpPayloadStart == chunkEnd) {
    return 0;  // no payload available
  }

  // validate Payload VelocyPack
  validator.validate(vpPayloadStart, std::distance(vpPayloadStart, chunkEnd),
                     /*isSubPart =*/false);
  return std::distance(vpHeaderStart, vpPayloadStart);
}

std::unique_ptr<basics::StringBuffer> createChunkForNetworkDetail(
    std::vector<VPackSlice> const& slices, bool isFirstChunk, uint32_t chunk,
    uint64_t id, uint32_t totalMessageLength = 0) {
  using basics::StringBuffer;
  bool firstOfMany = false;

  // if we have more than one chunk and the chunk is the first
  // then we are sending the first in a series. if this condition
  // is true we also send extra 8 bytes for the messageLength
  // (length of all VPackData)
  if (isFirstChunk && chunk > 1) {
    firstOfMany = true;
  }

  // build chunkX -- see VelocyStream Documentaion
  chunk <<= 1;
  chunk |= isFirstChunk ? 0x1 : 0x0;

  // get the lenght of VPack data
  uint32_t dataLength = 0;
  for (auto& slice : slices) {
    dataLength += slice.byteSize();
  }

  // calculate length of current chunk
  uint32_t chunkLength = dataLength;
  chunkLength += (sizeof(chunkLength) + sizeof(chunk) + sizeof(id));
  if (firstOfMany) {
    chunkLength += sizeof(totalMessageLength);
  }

  auto buffer =
      std::make_unique<StringBuffer>(TRI_UNKNOWN_MEM_ZONE, chunkLength, false);

  char cChunkLength[sizeof(chunkLength)];
  char const* cChunkLengthPtr = cChunkLength;
  std::memcpy(&cChunkLength, &chunkLength, sizeof(chunkLength));
  buffer->appendText(cChunkLengthPtr, sizeof(chunkLength));

  char cChunk[sizeof(chunk)];
  char const* cChunkPtr = cChunk;
  std::memcpy(&cChunk, &chunk, sizeof(chunk));
  buffer->appendText(cChunkPtr, sizeof(chunk));

  char cId[sizeof(id)];
  char const* cIdPtr = cId;
  std::memcpy(&cId, &id, sizeof(id));
  buffer->appendText(cIdPtr, sizeof(id));

  if (firstOfMany) {
    char cTotalMessageLength[sizeof(totalMessageLength)];
    char const* cTotalMessageLengthPtr = cTotalMessageLength;
    std::memcpy(&cTotalMessageLength, &totalMessageLength,
                sizeof(totalMessageLength));
    buffer->appendText(cTotalMessageLengthPtr, sizeof(totalMessageLength));
  }

  // append data in slices
  for (auto const& slice : slices) {
    buffer->appendText(std::string(slice.startAs<char>(), slice.byteSize()));
  }

  return buffer;
}

std::unique_ptr<basics::StringBuffer> createChunkForNetworkSingle(
    std::vector<VPackSlice> const& slices, uint64_t id) {
  return createChunkForNetworkDetail(slices, true, 1, id, 0 /*unused*/);
}

// TODO FIXME make use of these functions
// std::unique_ptr<basics::StringBuffer> createChunkForNetworkMultiFirst(
//     std::vector<VPackSlice> const& slices, uint64_t id, uint32_t
//     numberOfChunks,
//     uint32_t totalMessageLength) {
//   return createChunkForNetworkDetail(slices, true, numberOfChunks, id,
//                                      totalMessageLength);
// }
//
// std::unique_ptr<basics::StringBuffer> createChunkForNetworkMultiFollow(
//     std::vector<VPackSlice> const& slices, uint64_t id, uint32_t chunkNumber,
//     uint32_t totalMessageLength) {
//   return createChunkForNetworkDetail(slices, false, chunkNumber, id, 0);
// }
}

VppCommTask::VppCommTask(GeneralServer* server, TRI_socket_t sock,
                         ConnectionInfo&& info, double timeout)
    : Task("VppCommTask"),
      GeneralCommTask(server, sock, std::move(info), timeout) {
  _protocol = "vpp";
  _readBuffer->reserve(
      _bufferLength);  // ATTENTION <- this is required so we do not
                       // loose information during a resize
                       // connectionStatisticsAgentSetVpp();
}

void VppCommTask::addResponse(VppResponse* response, bool isError) {
  if (isError) {
    // FIXME (obi)
    // what do we need to do?
    // clean read buffer? reset process read cursor
  }

  VPackMessageNoOwnBuffer response_message = response->prepareForNetwork();
  uint64_t& id = response_message._id;

  std::vector<VPackSlice> slices;
  slices.push_back(response_message._header);

  VPackBuilder builder;
  if (response_message._generateBody) {
    builder = basics::VelocyPackHelper::stanitizeExternalsChecked(
        response_message._payload);
    slices.push_back(builder.slice());
  }

  // FIXME (obi)
  // If the message is big we will create many small chunks in a loop.
  // For the first tests we just send single Messages

  LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "got response:";
  for (auto const& slice : slices) {
    try {
      LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << slice.toJson();
    } catch (arangodb::velocypack::Exception const& e) {
      std::cout << "obi exception" << e.what();
    }
  }

  // adds chunk header infromation and creates SingBuffer* that can be
  // used with _writeBuffers
  auto buffer = createChunkForNetworkSingle(slices, id);
  _writeBuffers.push_back(buffer.get());
  buffer.release();

  fillWriteBuffer();  // move data from _writebuffers to _writebuffer
                      // implemented in base
}

VppCommTask::ChunkHeader VppCommTask::readChunkHeader() {
  VppCommTask::ChunkHeader header;

  auto cursor = _processReadVariables._readBufferCursor;

  std::memcpy(&header._chunkLength, cursor, sizeof(header._chunkLength));
  cursor += sizeof(header._chunkLength);

  uint32_t chunkX;
  std::memcpy(&chunkX, cursor, sizeof(chunkX));
  cursor += sizeof(chunkX);

  header._isFirst = chunkX & 0x1;
  header._chunk = chunkX >> 1;

  std::memcpy(&header._messageID, cursor, sizeof(header._messageID));
  cursor += sizeof(header._messageID);

  // extract total len of message
  if (header._isFirst && header._chunk > 1) {
    std::memcpy(&header._messageLength, cursor, sizeof(header._messageLength));
    cursor += sizeof(header._messageLength);
  } else {
    header._messageLength = 0;  // not needed
  }

  header._headerLength =
      std::distance(_processReadVariables._readBufferCursor, cursor);

  return header;
}

bool VppCommTask::isChunkComplete(char* start) {
  std::size_t length = std::distance(start, _readBuffer->end());
  auto& prv = _processReadVariables;

  if (!prv._currentChunkLength && length < sizeof(uint32_t)) {
    return false;
  }
  if (!prv._currentChunkLength) {
    // read chunk length
    std::memcpy(&prv._currentChunkLength, start, sizeof(uint32_t));
  }
  if (length < prv._currentChunkLength) {
    // chunk not complete
    return false;
  }

  return true;
}

// reads data from the socket
bool VppCommTask::processRead() {
  // TODO FIXME
  // - in case of error send an operation failed to all incomplete messages /
  //   operation and close connection (implement resetState/resetCommtask)

  auto& prv = _processReadVariables;
  if (!prv._readBufferCursor) {
    prv._readBufferCursor = _readBuffer->begin();
  }

  auto chunkBegin = prv._readBufferCursor;
  if (chunkBegin == nullptr || !isChunkComplete(chunkBegin)) {
    return false;  // no data or incomplete
  }

  ChunkHeader chunkHeader = readChunkHeader();
  auto chunkEnd = chunkBegin + chunkHeader._chunkLength;
  auto vpackBegin = chunkBegin + chunkHeader._headerLength;
  bool doExecute = false;
  bool read_maybe_only_part_of_buffer = false;
  VPackMessage message;  // filled in CASE 1 or CASE 2b

  // CASE 1: message is in one chunk
  if (chunkHeader._isFirst && chunkHeader._chunk == 1) {
    std::size_t payloadOffset = findAndValidateVPacks(vpackBegin, chunkEnd);
    message._id = chunkHeader._messageID;
    message._buffer.append(vpackBegin, std::distance(vpackBegin, chunkEnd));
    message._header = VPackSlice(message._buffer.data());
    if (payloadOffset) {
      message._payload = VPackSlice(message._buffer.data() + payloadOffset);
    }
    VPackValidator val;
    val.validate(message._header.begin(), message._header.byteSize());

    doExecute = true;
    LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 1";
  }
  // CASE 2:  message is in multiple chunks
  auto incompleteMessageItr = _incompleteMessages.find(chunkHeader._messageID);

  // CASE 2a: chunk starts new message
  if (chunkHeader._isFirst) {  // first chunk of multi chunk message
    if (incompleteMessageItr != _incompleteMessages.end()) {
      throw std::logic_error(
          "Message should be first but is already in the Map of incomplete "
          "messages");
    }

    IncompleteVPackMessage message(chunkHeader._messageLength,
                                   chunkHeader._chunk /*number of chunks*/);
    message._buffer.append(vpackBegin, std::distance(vpackBegin, chunkEnd));
    auto insertPair = _incompleteMessages.emplace(
        std::make_pair(chunkHeader._messageID, std::move(message)));
    if (!insertPair.second) {
      throw std::logic_error("insert failed");
    }

    LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 2a";

    // CASE 2b: chunk continues a message
  } else {  // followup chunk of some mesage
    if (incompleteMessageItr == _incompleteMessages.end()) {
      throw std::logic_error("found message without previous part");
    }
    auto& im = incompleteMessageItr->second;  // incomplete Message
    im._currentChunk++;
    assert(im._currentChunk == chunkHeader._chunk);
    im._buffer.append(vpackBegin, std::distance(vpackBegin, chunkEnd));
    // check buffer longer than length

    // MESSAGE COMPLETE
    if (im._currentChunk == im._numberOfChunks - 1 /* zero based counting */) {
      std::size_t payloadOffset = findAndValidateVPacks(
          reinterpret_cast<char const*>(im._buffer.data()),
          reinterpret_cast<char const*>(im._buffer.data() +
                                        im._buffer.byteSize()));
      message._id = chunkHeader._messageID;
      message._buffer = std::move(im._buffer);
      message._header = VPackSlice(message._buffer.data());
      if (payloadOffset) {
        message._payload = VPackSlice(message._buffer.data() + payloadOffset);
      }
      _incompleteMessages.erase(incompleteMessageItr);
      // check length

      doExecute = true;
      LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 2b - complete";
    }
    LOG_TOPIC(DEBUG, Logger::COMMUNICATION) << "CASE 2b - still incomplete";
  }

  // clean buffer up to length of chunk
  read_maybe_only_part_of_buffer = true;
  prv._currentChunkLength = 0;
  prv._readBufferCursor = chunkEnd;
  std::size_t processedDataLen =
      std::distance(_readBuffer->begin(), prv._readBufferCursor);
  if (processedDataLen > prv._cleanupLength) {
    _readBuffer->move_front(processedDataLen);
    prv._readBufferCursor = nullptr;  // the positon will be set at the
                                      // begin of this function
  }

  if (doExecute) {
    //    return false;  // we have no complete request, so we return early
    // for now we can handle only one request at a time
    // lock _request???? REVIEW (fc)
    LOG_TOPIC(DEBUG, Logger::COMMUNICATION)
        << "got request:" << message._header.toJson();
    _request = new VppRequest(_connectionInfo, std::move(message));
    GeneralServerFeature::HANDLER_FACTORY->setRequestContext(_request);
    if (_request->requestContext() == nullptr) {
      handleSimpleError(GeneralResponse::ResponseCode::NOT_FOUND,
                        TRI_ERROR_ARANGO_DATABASE_NOT_FOUND,
                        TRI_errno_string(TRI_ERROR_ARANGO_DATABASE_NOT_FOUND));
    } else {
      _request->setClientTaskId(_taskId);
      _protocolVersion = _request->protocolVersion();
      executeRequest(
          _request, new VppResponse(GeneralResponse::ResponseCode::SERVER_ERROR,
                                    chunkHeader._messageID));
    }
  }

  if (read_maybe_only_part_of_buffer) {
    if (prv._readBufferCursor == _readBuffer->end()) {
      return false;
    }
    return true;
  }
  return doExecute;
}

void VppCommTask::completedWriteBuffer() {
  // REVIEW (fc)
}

void VppCommTask::resetState(bool close) {
  // REVIEW (fc)
  // is there a case where we do not want to close the connection
  replyToIncompleteMessages();
}

// GeneralResponse::ResponseCode VppCommTask::authenticateRequest() {
//   auto context = (_request == nullptr) ? nullptr :
//   _request->requestContext();
//
//   if (context == nullptr && _request != nullptr) {
//     bool res =
//         GeneralServerFeature::HANDLER_FACTORY->setRequestContext(_request);
//
//     if (!res) {
//       return GeneralResponse::ResponseCode::NOT_FOUND;
//     }
//
//     context = _request->requestContext();
//   }
//
//   if (context == nullptr) {
//     return GeneralResponse::ResponseCode::SERVER_ERROR;
//   }
//
//   return context->authenticate();
// }

// convert internal GeneralRequest to VppRequest
VppRequest* VppCommTask::requestAsVpp() {
  VppRequest* request = dynamic_cast<VppRequest*>(_request);
  if (request == nullptr) {
    THROW_ARANGO_EXCEPTION(TRI_ERROR_INTERNAL);
  }
  return request;
};
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2014-2016 ArangoDB GmbH, Cologne, Germany
/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
///
/// 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.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Dr. Frank Celler
/// @author Martin Schoenert
////////////////////////////////////////////////////////////////////////////////

#include "SchedulerThread.h"

#include "Logger/Logger.h"
#include "Basics/MutexLocker.h"

#ifdef _WIN32
#include "Basics/win-utils.h"
#endif

#include "Scheduler/Scheduler.h"
#include "Scheduler/Task.h"

#include <velocypack/Value.h>
#include <velocypack/Builder.h>
#include <velocypack/velocypack-aliases.h>

using namespace arangodb::basics;
using namespace arangodb::rest;

SchedulerThread::SchedulerThread(Scheduler* scheduler, EventLoop loop,
                                 bool defaultLoop)
    : Thread("Scheduler"),
      _scheduler(scheduler),
      _defaultLoop(defaultLoop),
      _loop(loop),
      _numberTasks(0),
      _taskData(100) {}

////////////////////////////////////////////////////////////////////////////////
/// @brief begin shutdown sequence
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::beginShutdown() {
  Thread::beginShutdown();

  LOG(TRACE) << "beginning shutdown sequence of scheduler thread ("
             << threadId() << ")";

  _scheduler->wakeupLoop(_loop);
}

////////////////////////////////////////////////////////////////////////////////
/// @brief registers a task
////////////////////////////////////////////////////////////////////////////////

bool SchedulerThread::registerTask(Scheduler* scheduler, Task* task) {
  // thread has already been stopped
  if (isStopping()) {
    // do nothing
    deleteTask(task);
    return false;
  }

  TRI_ASSERT(scheduler != nullptr);
  TRI_ASSERT(task != nullptr);

  // same thread, in this case it does not matter if we are inside the loop
  if (threadId() == currentThreadId()) {
    bool ok = setupTask(task, scheduler, _loop);

    if (ok) {
      ++_numberTasks;
    } else {
      LOG(WARN) << "In SchedulerThread::registerTask setupTask has failed";
      cleanupTask(task);
      deleteTask(task);
    }

    return ok;
  }

  Work w(SETUP, scheduler, task);

  // different thread, be careful - we have to stop the event loop
  // put the register request onto the queue
  MUTEX_LOCKER(mutexLocker, _queueLock);

  _queue.push_back(w);

  scheduler->wakeupLoop(_loop);

  return true;
}

////////////////////////////////////////////////////////////////////////////////
/// @brief unregisters a task
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::unregisterTask(Task* task) {
  // thread has already been stopped
  if (isStopping()) {
    return;
  }

  // same thread, in this case it does not matter if we are inside the loop
  if (threadId() == currentThreadId()) {
    cleanupTask(task);
    --_numberTasks;
  }

  // different thread, be careful - we have to stop the event loop
  else {
    Work w(CLEANUP, nullptr, task);

    // put the unregister request into the queue
    MUTEX_LOCKER(mutexLocker, _queueLock);

    _queue.push_back(w);

    _scheduler->wakeupLoop(_loop);
  }
}

////////////////////////////////////////////////////////////////////////////////
/// @brief unregisters a task
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::destroyTask(Task* task) {
  // thread has already been stopped
  if (isStopping()) {
    deleteTask(task);
    return;
  }

  // same thread, in this case it does not matter if we are inside the loop
  if (threadId() == currentThreadId()) {
    cleanupTask(task);
    deleteTask(task);
    --_numberTasks;
  }

  // different thread, be careful - we have to stop the event loop
  else {
    // put the unregister request into the queue
    Work w(DESTROY, nullptr, task);

    MUTEX_LOCKER(mutexLocker, _queueLock);

    _queue.push_back(w);

    _scheduler->wakeupLoop(_loop);
  }
}

////////////////////////////////////////////////////////////////////////////////
/// @brief sends data to a task
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::signalTask(std::unique_ptr<TaskData>& data) {
  _taskData.push(data.release());
  _scheduler->wakeupLoop(_loop);
}

void SchedulerThread::run() {
  LOG(TRACE) << "scheduler thread started (" << threadId() << ")";

  if (_defaultLoop) {
#ifdef TRI_HAVE_POSIX_THREADS
    sigset_t all;
    sigemptyset(&all);
    pthread_sigmask(SIG_SETMASK, &all, 0);
#endif
  }

  while (!isStopping()) {
    // handle the returned data
    TaskData* data;

    while (_taskData.pop(data)) {
      Task* task = _scheduler->lookupTaskById(data->_taskId);

      if (task != nullptr) {
        task->signalTask(data);
      }

      delete data;
    }

    // handle the events
    try {
      _scheduler->eventLoop(_loop);
    } catch (std::exception const& e) {
#ifdef TRI_HAVE_POSIX_THREADS
      if (isStopping()) {
        LOG(WARN) << "caught cancelation exception during work";
        throw;
      }
#endif
      LOG(WARN) << "caught exception from ev_loop: " << e.what();
    } catch (...) {
#ifdef TRI_HAVE_POSIX_THREADS
      if (isStopping()) {
        LOG(WARN) << "caught cancelation exception during work";
        throw;
      }
#endif
      LOG(WARN) << "caught exception from ev_loop";
    }

#if defined(DEBUG_SCHEDULER_THREAD)
    LOG(TRACE) << "left scheduler loop " << threadId();
#endif

    while (true) {
      Work w;

      {
        MUTEX_LOCKER(mutexLocker,
                     _queueLock);  // TODO(fc) XXX goto boost lockfree

        if (_queue.empty()) {
          break;
        }

        w = _queue.front();
        _queue.pop_front();
      }

      // will only get here if there is something to do
      switch (w.work) {
        case CLEANUP: {
          cleanupTask(w.task);
          --_numberTasks;
          break;
        }

        case SETUP: {
          bool ok = setupTask(w.task, w.scheduler, _loop);

          if (ok) {
            ++_numberTasks;
          } else {
            cleanupTask(w.task);
            deleteTask(w.task);
          }

          break;
        }

        case DESTROY: {
          cleanupTask(w.task);
          deleteTask(w.task);
          --_numberTasks;
          break;
        }

        case INVALID: {
          LOG(ERR) << "logic error. got invalid Work item";
          break;
        }
      }
    }
  }

  LOG(TRACE) << "scheduler thread stopped (" << threadId() << ")";

  // pop all undeliviered task data
  {
    TaskData* data;

    while (_taskData.pop(data)) {
      delete data;
    }
  }

  // pop all elements from the queue and delete them
  while (true) {
    Work w;

    {
      MUTEX_LOCKER(mutexLocker, _queueLock);

      if (_queue.empty()) {
        break;
      }

      w = _queue.front();
      _queue.pop_front();
    }

    // will only get here if there is something to do
    switch (w.work) {
      case CLEANUP:
        break;

      case SETUP:
        break;

      case DESTROY:
        deleteTask(w.task);
        break;

      case INVALID:
        LOG(ERR) << "logic error. got invalid Work item";
        break;
    }
  }
}

void SchedulerThread::addStatus(VPackBuilder* b) {
  Thread::addStatus(b);
  b->add("numberTasks", VPackValue(_numberTasks.load()));
}
<commit_msg>potentially fix leak<commit_after>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2014-2016 ArangoDB GmbH, Cologne, Germany
/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
///
/// 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.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Dr. Frank Celler
/// @author Martin Schoenert
////////////////////////////////////////////////////////////////////////////////

#include "SchedulerThread.h"

#include "Logger/Logger.h"
#include "Basics/MutexLocker.h"

#ifdef _WIN32
#include "Basics/win-utils.h"
#endif

#include "Scheduler/Scheduler.h"
#include "Scheduler/Task.h"

#include <velocypack/Value.h>
#include <velocypack/Builder.h>
#include <velocypack/velocypack-aliases.h>

using namespace arangodb::basics;
using namespace arangodb::rest;

SchedulerThread::SchedulerThread(Scheduler* scheduler, EventLoop loop,
                                 bool defaultLoop)
    : Thread("Scheduler"),
      _scheduler(scheduler),
      _defaultLoop(defaultLoop),
      _loop(loop),
      _numberTasks(0),
      _taskData(100) {}

////////////////////////////////////////////////////////////////////////////////
/// @brief begin shutdown sequence
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::beginShutdown() {
  Thread::beginShutdown();

  LOG(TRACE) << "beginning shutdown sequence of scheduler thread ("
             << threadId() << ")";

  _scheduler->wakeupLoop(_loop);
}

////////////////////////////////////////////////////////////////////////////////
/// @brief registers a task
////////////////////////////////////////////////////////////////////////////////

bool SchedulerThread::registerTask(Scheduler* scheduler, Task* task) {
  // thread has already been stopped
  if (isStopping()) {
    // do nothing
    deleteTask(task);
    return false;
  }

  TRI_ASSERT(scheduler != nullptr);
  TRI_ASSERT(task != nullptr);

  // same thread, in this case it does not matter if we are inside the loop
  if (threadId() == currentThreadId()) {
    bool ok = setupTask(task, scheduler, _loop);

    if (ok) {
      ++_numberTasks;
    } else {
      LOG(WARN) << "In SchedulerThread::registerTask setupTask has failed";
      cleanupTask(task);
      deleteTask(task);
    }

    return ok;
  }

  Work w(SETUP, scheduler, task);

  // different thread, be careful - we have to stop the event loop
  // put the register request onto the queue
  MUTEX_LOCKER(mutexLocker, _queueLock);

  _queue.push_back(w);

  scheduler->wakeupLoop(_loop);

  return true;
}

////////////////////////////////////////////////////////////////////////////////
/// @brief unregisters a task
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::unregisterTask(Task* task) {
  // thread has already been stopped
  if (isStopping()) {
    return;
  }

  // same thread, in this case it does not matter if we are inside the loop
  if (threadId() == currentThreadId()) {
    cleanupTask(task);
    --_numberTasks;
  }

  // different thread, be careful - we have to stop the event loop
  else {
    Work w(CLEANUP, nullptr, task);

    // put the unregister request into the queue
    MUTEX_LOCKER(mutexLocker, _queueLock);

    _queue.push_back(w);

    _scheduler->wakeupLoop(_loop);
  }
}

////////////////////////////////////////////////////////////////////////////////
/// @brief unregisters a task
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::destroyTask(Task* task) {
  // thread has already been stopped
  if (isStopping()) {
    deleteTask(task);
    return;
  }

  // same thread, in this case it does not matter if we are inside the loop
  if (threadId() == currentThreadId()) {
    cleanupTask(task);
    deleteTask(task);
    --_numberTasks;
  }

  // different thread, be careful - we have to stop the event loop
  else {
    // put the unregister request into the queue
    Work w(DESTROY, nullptr, task);

    MUTEX_LOCKER(mutexLocker, _queueLock);

    _queue.push_back(w);

    _scheduler->wakeupLoop(_loop);
  }
}

////////////////////////////////////////////////////////////////////////////////
/// @brief sends data to a task
////////////////////////////////////////////////////////////////////////////////

void SchedulerThread::signalTask(std::unique_ptr<TaskData>& data) {
  bool result = _taskData.push(data.get());
  if (result) {
    data.release();
    _scheduler->wakeupLoop(_loop);
  }
}

void SchedulerThread::run() {
  LOG(TRACE) << "scheduler thread started (" << threadId() << ")";

  if (_defaultLoop) {
#ifdef TRI_HAVE_POSIX_THREADS
    sigset_t all;
    sigemptyset(&all);
    pthread_sigmask(SIG_SETMASK, &all, 0);
#endif
  }

  while (!isStopping()) {
    // handle the returned data
    TaskData* data;

    while (_taskData.pop(data)) {
      Task* task = _scheduler->lookupTaskById(data->_taskId);

      if (task != nullptr) {
        task->signalTask(data);
      }

      delete data;
    }

    // handle the events
    try {
      _scheduler->eventLoop(_loop);
    } catch (std::exception const& e) {
#ifdef TRI_HAVE_POSIX_THREADS
      if (isStopping()) {
        LOG(WARN) << "caught cancelation exception during work";
        throw;
      }
#endif
      LOG(WARN) << "caught exception from ev_loop: " << e.what();
    } catch (...) {
#ifdef TRI_HAVE_POSIX_THREADS
      if (isStopping()) {
        LOG(WARN) << "caught cancelation exception during work";
        throw;
      }
#endif
      LOG(WARN) << "caught exception from ev_loop";
    }

#if defined(DEBUG_SCHEDULER_THREAD)
    LOG(TRACE) << "left scheduler loop " << threadId();
#endif

    while (true) {
      Work w;

      {
        MUTEX_LOCKER(mutexLocker,
                     _queueLock);  // TODO(fc) XXX goto boost lockfree

        if (_queue.empty()) {
          break;
        }

        w = _queue.front();
        _queue.pop_front();
      }

      // will only get here if there is something to do
      switch (w.work) {
        case CLEANUP: {
          cleanupTask(w.task);
          --_numberTasks;
          break;
        }

        case SETUP: {
          bool ok = setupTask(w.task, w.scheduler, _loop);

          if (ok) {
            ++_numberTasks;
          } else {
            cleanupTask(w.task);
            deleteTask(w.task);
          }

          break;
        }

        case DESTROY: {
          cleanupTask(w.task);
          deleteTask(w.task);
          --_numberTasks;
          break;
        }

        case INVALID: {
          LOG(ERR) << "logic error. got invalid Work item";
          break;
        }
      }
    }
  }

  LOG(TRACE) << "scheduler thread stopped (" << threadId() << ")";

  // pop all undeliviered task data
  {
    TaskData* data;

    while (_taskData.pop(data)) {
      delete data;
    }
  }

  // pop all elements from the queue and delete them
  while (true) {
    Work w;

    {
      MUTEX_LOCKER(mutexLocker, _queueLock);

      if (_queue.empty()) {
        break;
      }

      w = _queue.front();
      _queue.pop_front();
    }

    // will only get here if there is something to do
    switch (w.work) {
      case CLEANUP:
        break;

      case SETUP:
        break;

      case DESTROY:
        deleteTask(w.task);
        break;

      case INVALID:
        LOG(ERR) << "logic error. got invalid Work item";
        break;
    }
  }
}

void SchedulerThread::addStatus(VPackBuilder* b) {
  Thread::addStatus(b);
  b->add("numberTasks", VPackValue(_numberTasks.load()));
}
<|endoftext|>


<commit_before><commit_msg>Cosmetic changes<commit_after><|endoftext|>


<commit_before>#include <phypp.hpp>

void print_help();

int phypp_main(int argc, char* argv[]) {
    if (argc < 3) {
        print_help();
        return 0;
    }

    vec1s suffix;
    std::string fitmask;
    std::string outdir;
    bool residuals = false;
    read_args(argc-2, argv+2, arg_list(suffix, fitmask, residuals, outdir));

    if (!outdir.empty()) {
        outdir = file::directorize(outdir);
        file::mkdir(outdir);
    }

    // Read cube
    fits::input_image fimg(argv[1]);
    vec3d flx;
    fimg.reach_hdu(1);
    fimg.read(flx);
    // Read uncertainty
    vec3d err;
    fimg.reach_hdu(2);
    fimg.read(err);

    // Get wavelength WCS
    double crpix, crval, cdelt;
    fimg.read_keyword("CRPIX3", crpix);
    fimg.read_keyword("CRVAL3", crval);
    fimg.read_keyword("CDELT3", cdelt);

    // Read fit mask (optional)
    vec2d mask; {
        if (!fitmask.empty()) {
            fits::read(fitmask, mask);
        } else {
            mask = replicate(1, flx.dims[1], flx.dims[2]);
        }
    }

    // Read models
    vec2d models;
    uint_t nmodel; {
        fits::input_image fmodel(argv[2]);
        if (fmodel.axis_count() == 2) {
            nmodel = 1;
            fits::read(argv[2], models);
            models = reform(models, 1, models.size());
        } else {
            vec3d tmp;
            fits::read(argv[2], tmp);

            nmodel = tmp.dims[0];
            models = reform(tmp, nmodel, tmp.dims[1]*tmp.dims[2]);
        }

        // Make sure each model has unit integral
        for (uint_t i : range(nmodel)) {
            vec1u idg = where(is_finite(models(i,_)));
            models(i,_) /= total(models(i,_)[idg]);
        }

        if (suffix.size() != nmodel && nmodel != 1) {
            error("please provide as many suffixes as there are models (", nmodel,
                ") in suffix=[...]");
            return 1;
        }
    }

    // Adjust suffixes to include "_"
    if (nmodel == 1 && suffix.empty()) {
        suffix = {""};
    }

    for (std::string& s : suffix) {
        if (!s.empty() && s[0] != '_') {
            s = "_"+s;
        }
    }

    // Do the fit wavelength by wavelength
    vec2d flx1d(flx.dims[0], nmodel);
    vec2d err1d(flx.dims[0], nmodel);
    for (uint_t l : range(flx.dims[0])) {
        vec1u idg = where(is_finite(flx(l,_,_)) && is_finite(err(l,_,_)) && mask > 0.0);
        auto res = linfit_pack(flx(l,_,_)[idg], err(l,_,_)[idg], models(_,idg));
        flx1d(l,_) = res.params;
        err1d(l,_) = res.errors;
    }

    // Write spectra to disk
    std::string filebase = outdir+file::get_basename(erase_end(argv[1], ".fits"));
    for (uint_t i : range(nmodel)) {
        fits::output_image fspec(filebase+suffix[i]+"_spec.fits");
        fspec.write(vec1d(0)); // empty primary extension, KMOS convention
        fspec.reach_hdu(1);
        fspec.write(flx1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
        fspec.reach_hdu(2);
        fspec.write(err1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
    }

    // Compute residuals if asked
    if (residuals) {
        fimg.reach_hdu(1);
        fits::header hdr = fimg.read_header();

        for (uint_t i : range(nmodel)) {
            vec3d res = flx;
            for (uint_t l : range(res.dims[0])) {
                flatten(res(l,_,_)) -= flx1d(l,i)*models(i,_);
            }

            fits::output_image fspec(filebase+suffix[i]+"_residual.fits");
            fspec.write(vec3d(0,0,0)); // empty primary HDU
            fspec.reach_hdu(1);
            fspec.write(res);
            fspec.write_header(hdr);
            fspec.reach_hdu(2);
            fspec.write(err);
            fspec.write_header(hdr);
        }
    }

    return 0;
}

void print_help() {
    using namespace format;

    print("multispecfit v1.0");
    print("usage: multispecfit <kmos_cube.fits> <models.fits> suffix=[...] [options]");
    print("");
    print("Main parameters:");
    paragraph("'kmos_cube.fits' must be a cube created by the KMOS pipeline, with 3 "
        "extensions: the first is empty (KMOS convention), the second contains the flux, "
        "and the third contains the uncertainty. 'models.fits' must be a cube created by "
        "yourself (e.g., with IDL or Python) containing the models to fit. Each slice of "
        "this cube corresponds to a different model. Note that you have to include a "
        "uniform model (all pixels equal to 1) if you want to fit a constant background. "
        "Finally, 'suffix' must contain the names of each model, with the following "
        "format: [name1,name1,name3,...]. There must be as many suffixes as there are "
        "models in the 'models.fits' file.");
    print("Available options:");
    bullet("fitmask=...", "Must be a 2D FITS file containing the mask to define the fitting "
        "region. Only the pixels with a non-zero value in the mask will be used "
        "(default: all valid pixels are used).");
    bullet("residuals", "Set this flag if you want the program to generate residual cubes, "
        "removing the contribution of each component separately (default: do not "
        "generate residuals)");
    bullet("outdir", "Name of the directory into which the output files should be created. "
        "Default is the current directory.");
}
<commit_msg>Make sure we only fit valid pixels in multispecfit<commit_after>#include <phypp.hpp>

void print_help();

int phypp_main(int argc, char* argv[]) {
    if (argc < 3) {
        print_help();
        return 0;
    }

    vec1s suffix;
    std::string fitmask;
    std::string outdir;
    bool residuals = false;
    read_args(argc-2, argv+2, arg_list(suffix, fitmask, residuals, outdir));

    if (!outdir.empty()) {
        outdir = file::directorize(outdir);
        file::mkdir(outdir);
    }

    // Read cube
    fits::input_image fimg(argv[1]);
    vec3d flx;
    fimg.reach_hdu(1);
    fimg.read(flx);
    // Read uncertainty
    vec3d err;
    fimg.reach_hdu(2);
    fimg.read(err);

    // Get wavelength WCS
    double crpix, crval, cdelt;
    fimg.read_keyword("CRPIX3", crpix);
    fimg.read_keyword("CRVAL3", crval);
    fimg.read_keyword("CDELT3", cdelt);

    // Read fit mask (optional)
    vec2d mask; {
        if (!fitmask.empty()) {
            fits::read(fitmask, mask);
        } else {
            mask = vec2d{partial_count(0, is_finite(flx)) > 1};
        }
    }

    // Read models
    vec2d models;
    uint_t nmodel; {
        fits::input_image fmodel(argv[2]);
        if (fmodel.axis_count() == 2) {
            nmodel = 1;
            fits::read(argv[2], models);
            models = reform(models, 1, models.size());
        } else {
            vec3d tmp;
            fits::read(argv[2], tmp);

            nmodel = tmp.dims[0];
            models = reform(tmp, nmodel, tmp.dims[1]*tmp.dims[2]);
        }

        // Make sure each model has unit integral
        for (uint_t i : range(nmodel)) {
            vec1u idg = where(is_finite(models(i,_)));
            models(i,_) /= total(models(i,_)[idg]);
        }

        if (suffix.size() != nmodel && nmodel != 1) {
            error("please provide as many suffixes as there are models (", nmodel,
                ") in suffix=[...]");
            return 1;
        }
    }

    // Adjust suffixes to include "_"
    if (nmodel == 1 && suffix.empty()) {
        suffix = {""};
    }

    for (std::string& s : suffix) {
        if (!s.empty() && s[0] != '_') {
            s = "_"+s;
        }
    }

    // Do the fit wavelength by wavelength
    vec2d flx1d(flx.dims[0], nmodel);
    vec2d err1d(flx.dims[0], nmodel);
    for (uint_t l : range(flx.dims[0])) {
        vec1u idg = where(is_finite(flx(l,_,_)) && is_finite(err(l,_,_)) && err(l,_,_) > 0 && mask > 0.0);
        auto res = linfit_pack(flx(l,_,_)[idg], err(l,_,_)[idg], models(_,idg));
        flx1d(l,_) = res.params;
        err1d(l,_) = res.errors;
    }

    // Write spectra to disk
    std::string filebase = outdir+file::get_basename(erase_end(argv[1], ".fits"));
    for (uint_t i : range(nmodel)) {
        fits::output_image fspec(filebase+suffix[i]+"_spec.fits");
        fspec.write(vec1d(0)); // empty primary extension, KMOS convention
        fspec.reach_hdu(1);
        fspec.write(flx1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
        fspec.reach_hdu(2);
        fspec.write(err1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
    }

    // Compute residuals if asked
    if (residuals) {
        fimg.reach_hdu(1);
        fits::header hdr = fimg.read_header();

        for (uint_t i : range(nmodel)) {
            vec3d res = flx;
            for (uint_t l : range(res.dims[0])) {
                flatten(res(l,_,_)) -= flx1d(l,i)*models(i,_);
            }

            fits::output_image fspec(filebase+suffix[i]+"_residual.fits");
            fspec.write(vec3d(0,0,0)); // empty primary HDU
            fspec.reach_hdu(1);
            fspec.write(res);
            fspec.write_header(hdr);
            fspec.reach_hdu(2);
            fspec.write(err);
            fspec.write_header(hdr);
        }
    }

    return 0;
}

void print_help() {
    using namespace format;

    print("multispecfit v1.0");
    print("usage: multispecfit <kmos_cube.fits> <models.fits> suffix=[...] [options]");
    print("");
    print("Main parameters:");
    paragraph("'kmos_cube.fits' must be a cube created by the KMOS pipeline, with 3 "
        "extensions: the first is empty (KMOS convention), the second contains the flux, "
        "and the third contains the uncertainty. 'models.fits' must be a cube created by "
        "yourself (e.g., with IDL or Python) containing the models to fit. Each slice of "
        "this cube corresponds to a different model. Note that you have to include a "
        "uniform model (all pixels equal to 1) if you want to fit a constant background. "
        "Finally, 'suffix' must contain the names of each model, with the following "
        "format: [name1,name1,name3,...]. There must be as many suffixes as there are "
        "models in the 'models.fits' file.");
    print("Available options:");
    bullet("fitmask=...", "Must be a 2D FITS file containing the mask to define the fitting "
        "region. Only the pixels with a non-zero value in the mask will be used "
        "(default: all valid pixels are used).");
    bullet("residuals", "Set this flag if you want the program to generate residual cubes, "
        "removing the contribution of each component separately (default: do not "
        "generate residuals)");
    bullet("outdir", "Name of the directory into which the output files should be created. "
        "Default is the current directory.");
}
<|endoftext|>


<commit_before>/*=========================================================================
*
* Copyright Marius Staring, Stefan Klein, David Doria. 2011.
*
* 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.txt
*
* 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.
*
*=========================================================================*/
/** \file
 \brief Stack images into one big vector image.

 \verbinclude imagestovectorimage.help
 */
#include "itkCommandLineArgumentParser.h"
#include "ITKToolsHelpers.h"
#include "imagestovectorimage.h"

/**
 * ******************* GetHelpString *******************
 */

std::string GetHelpString( void )
{
  std::stringstream ss;
  ss << "ITKTools v" << itktools::GetITKToolsVersion() << "\n"
    << "Usage:\n"
    << "pximagetovectorimage\n"
    << "  -in      inputFilenames, at least 2\n"
    << "  [-out]   outputFilename, default VECTOR.mhd\n"
    << "  [-s]     number of streams, default 1.\n"
    << "Supported: 2D, 3D, (unsigned) char, (unsigned) short,\n"
    << "(unsigned) int, (unsigned) long, float, double.\n"
    << "Note: make sure that the input images are of the same type, size, etc.";

  return ss.str();

} // end GetHelpString()

//-------------------------------------------------------------------------------------

int main( int argc, char ** argv )
{
  /** Create a command line argument parser. */
  itk::CommandLineArgumentParser::Pointer parser = itk::CommandLineArgumentParser::New();
  parser->SetCommandLineArguments( argc, argv );
  parser->SetProgramHelpText( GetHelpString() );

  parser->MarkArgumentAsRequired( "-in", "The input filename." );

  itk::CommandLineArgumentParser::ReturnValue validateArguments = parser->CheckForRequiredArguments();

  if( validateArguments == itk::CommandLineArgumentParser::FAILED )
  {
    return EXIT_FAILURE;
  }
  else if( validateArguments == itk::CommandLineArgumentParser::HELPREQUESTED )
  {
    return EXIT_SUCCESS;
  }

  /** Get arguments. */
  std::vector<std::string>  inputFileNames( 0, "" );
  parser->GetCommandLineArgument( "-in", inputFileNames );

  std::string outputFileName = "VECTOR.mhd";
  parser->GetCommandLineArgument( "-out", outputFileName );

  unsigned int numberOfStreams = 1;
  parser->GetCommandLineArgument( "-s", numberOfStreams );

  /** Check if the required arguments are given. */
  if( inputFileNames.size() < 2 )
  {
    std::cerr << "ERROR: You should specify at least two (2) input files." << std::endl;
    return EXIT_FAILURE;
  }

  /** Determine image properties. */
  itk::ImageIOBase::IOPixelType pixelType = itk::ImageIOBase::UNKNOWNPIXELTYPE;
  itk::ImageIOBase::IOComponentType componentType = itk::ImageIOBase::UNKNOWNCOMPONENTTYPE;
  unsigned int dim = 0;
  unsigned int numberOfComponents = 0;
  bool retgip = itktools::GetImageProperties(
    inputFileNames[ 0 ], pixelType, componentType, dim, numberOfComponents );
  if( !retgip ) return EXIT_FAILURE;

  /** Check for vector images. */
  bool retNOCCheck = itktools::NumberOfComponentsCheck( numberOfComponents );
  if( !retNOCCheck ) return EXIT_FAILURE;

  /** Class that does the work. */
  ITKToolsImagesToVectorImageBase * filter = NULL;

  try
  {
    // now call all possible template combinations.
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, float >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, double >::New( dim, componentType );

#ifdef ITKTOOLS_3D_SUPPORT
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, float >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, double >::New( dim, componentType );
#endif
    /** Check if filter was instantiated. */
    bool supported = itktools::IsFilterSupportedCheck( filter, dim, componentType );
    if( !supported ) return EXIT_FAILURE;

    /** Set the filter arguments. */
    filter->m_InputFileNames = inputFileNames;
    filter->m_OutputFileName = outputFileName;
    filter->m_NumberOfStreams = numberOfStreams;

    filter->Run();

    delete filter;
  }
  catch( itk::ExceptionObject & excp )
  {
    std::cerr << "ERROR: Caught ITK exception: " << excp << std::endl;
    delete filter;
    return EXIT_FAILURE;
  }

  /** End program. */
  return EXIT_SUCCESS;

} // end main
<commit_msg>BUG: vector input images are supported by pximagestovectorimages<commit_after>/*=========================================================================
*
* Copyright Marius Staring, Stefan Klein, David Doria. 2011.
*
* 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.txt
*
* 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.
*
*=========================================================================*/
/** \file
 \brief Stack images into one big vector image.

 \verbinclude imagestovectorimage.help
 */
#include "itkCommandLineArgumentParser.h"
#include "ITKToolsHelpers.h"
#include "imagestovectorimage.h"

/**
 * ******************* GetHelpString *******************
 */

std::string GetHelpString( void )
{
  std::stringstream ss;
  ss << "ITKTools v" << itktools::GetITKToolsVersion() << "\n"
    << "Usage:\n"
    << "pximagetovectorimage\n"
    << "  -in      inputFilenames, at least 2\n"
    << "  [-out]   outputFilename, default VECTOR.mhd\n"
    << "  [-s]     number of streams, default 1.\n"
    << "Supported: 2D, 3D, (unsigned) char, (unsigned) short,\n"
    << "(unsigned) int, (unsigned) long, float, double.\n"
    << "Note: make sure that the input images are of the same type, size, etc.";

  return ss.str();

} // end GetHelpString()

//-------------------------------------------------------------------------------------

int main( int argc, char ** argv )
{
  /** Create a command line argument parser. */
  itk::CommandLineArgumentParser::Pointer parser = itk::CommandLineArgumentParser::New();
  parser->SetCommandLineArguments( argc, argv );
  parser->SetProgramHelpText( GetHelpString() );

  parser->MarkArgumentAsRequired( "-in", "The input filename." );

  itk::CommandLineArgumentParser::ReturnValue validateArguments = parser->CheckForRequiredArguments();

  if( validateArguments == itk::CommandLineArgumentParser::FAILED )
  {
    return EXIT_FAILURE;
  }
  else if( validateArguments == itk::CommandLineArgumentParser::HELPREQUESTED )
  {
    return EXIT_SUCCESS;
  }

  /** Get arguments. */
  std::vector<std::string>  inputFileNames( 0, "" );
  parser->GetCommandLineArgument( "-in", inputFileNames );

  std::string outputFileName = "VECTOR.mhd";
  parser->GetCommandLineArgument( "-out", outputFileName );

  unsigned int numberOfStreams = 1;
  parser->GetCommandLineArgument( "-s", numberOfStreams );

  /** Check if the required arguments are given. */
  if( inputFileNames.size() < 2 )
  {
    std::cerr << "ERROR: You should specify at least two (2) input files." << std::endl;
    return EXIT_FAILURE;
  }

  /** Determine image properties. */
  itk::ImageIOBase::IOPixelType pixelType = itk::ImageIOBase::UNKNOWNPIXELTYPE;
  itk::ImageIOBase::IOComponentType componentType = itk::ImageIOBase::UNKNOWNCOMPONENTTYPE;
  unsigned int dim = 0;
  unsigned int numberOfComponents = 0;
  bool retgip = itktools::GetImageProperties(
    inputFileNames[ 0 ], pixelType, componentType, dim, numberOfComponents );
  if( !retgip ) return EXIT_FAILURE;

  /** Class that does the work. */
  ITKToolsImagesToVectorImageBase * filter = NULL;

  try
  {
    // now call all possible template combinations.
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, unsigned long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, float >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 2, double >::New( dim, componentType );

#ifdef ITKTOOLS_3D_SUPPORT
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned char >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned short >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned int >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, unsigned long >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, float >::New( dim, componentType );
    if( !filter ) filter = ITKToolsImagesToVectorImage< 3, double >::New( dim, componentType );
#endif
    /** Check if filter was instantiated. */
    bool supported = itktools::IsFilterSupportedCheck( filter, dim, componentType );
    if( !supported ) return EXIT_FAILURE;

    /** Set the filter arguments. */
    filter->m_InputFileNames = inputFileNames;
    filter->m_OutputFileName = outputFileName;
    filter->m_NumberOfStreams = numberOfStreams;

    filter->Run();

    delete filter;
  }
  catch( itk::ExceptionObject & excp )
  {
    std::cerr << "ERROR: Caught ITK exception: " << excp << std::endl;
    delete filter;
    return EXIT_FAILURE;
  }

  /** End program. */
  return EXIT_SUCCESS;

} // end main
<|endoftext|>


<commit_before>#include "itkCommandLineArgumentParser.h"
#include "CommandLineArgumentHelper.h"

#include "itkImageFileReader.h"
#include "itkImageToVectorImageFilter.h"
#include "itkImageFileWriter.h"
#include "itkVectorIndexSelectionCastImageFilter.h"

//-------------------------------------------------------------------------------------

/** run: A macro to call a function. */
#define run( function, type, dim ) \
if ( ComponentTypeIn == #type && Dimension == dim ) \
{ \
  typedef itk::VectorImage< type, dim >       InputImageType; \
  typedef itk::VectorImage< type, dim > OutputImageType; \
  function< InputImageType, OutputImageType >( inputFileNames, outputFileName, numberOfStreams ); \
  supported = true; \
}

//-------------------------------------------------------------------------------------

/** Declare ComposeVectorImage. */
template< class InputImageType, class OutputImageType >
void ComposeVectorImage(
  const std::vector<std::string> & inputFileNames,
  const std::string & outputFileName,
  const unsigned int & numberOfStreams );

/** Declare PrintHelp. */
void PrintHelp( void );

//-------------------------------------------------------------------------------------

int main( int argc, char ** argv )
{
  /** Check arguments for help. */
  if ( argc < 4 )
  {
    PrintHelp();
    return 1;
  }

  /** Create a command line argument parser. */
  itk::CommandLineArgumentParser::Pointer parser = itk::CommandLineArgumentParser::New();
  parser->SetCommandLineArguments( argc, argv );

  /** Get arguments. */
  std::vector<std::string>  inputFileNames( 0, "" );
  bool retin = parser->GetCommandLineArgument( "-in", inputFileNames );

  std::string outputFileName = "VECTOR.mhd";
  parser->GetCommandLineArgument( "-out", outputFileName );

  /** Support for streaming. */
  unsigned int numberOfStreams = 1;
  parser->GetCommandLineArgument( "-s", numberOfStreams );

  /** Check if the required arguments are given. */
  if ( !retin )
  {
    std::cerr << "ERROR: You should specify \"-in\"." << std::endl;
    return 1;
  }
  if ( inputFileNames.size() < 2 )
  {
    std::cerr << "ERROR: You should specify at least two (2) input files." << std::endl;
    return 1;
  }

  /** Determine image properties. */
  std::string ComponentTypeIn = "short";
  std::string PixelType; //we don't use this
  unsigned int Dimension = 3;
  unsigned int NumberOfComponents = 1;
  std::vector<unsigned int> imagesize( Dimension, 0 );
  int retgip = GetImageProperties(
    inputFileNames[ 0 ],
    PixelType,
    ComponentTypeIn,
    Dimension,
    NumberOfComponents,
    imagesize );
  if ( retgip != 0 )
  {
    return 1;
  }

  /** Get rid of the possible "_" in ComponentType. */
  ReplaceUnderscoreWithSpace( ComponentTypeIn );

  /** Run the program. */
  bool supported = false;
  try
  {
    run( ComposeVectorImage, char, 2 );
    run( ComposeVectorImage, unsigned char, 2 );
    run( ComposeVectorImage, short, 2 );
    run( ComposeVectorImage, unsigned short, 2 );
    run( ComposeVectorImage, int, 2 );
    run( ComposeVectorImage, unsigned int, 2 );
    run( ComposeVectorImage, long, 2 );
    run( ComposeVectorImage, unsigned long, 2 );
    run( ComposeVectorImage, float, 2 );
    run( ComposeVectorImage, double, 2 );

    run( ComposeVectorImage, char, 3 );
    run( ComposeVectorImage, unsigned char, 3 );
    run( ComposeVectorImage, short, 3 );
    run( ComposeVectorImage, unsigned short, 3 );
    run( ComposeVectorImage, int, 3 );
    run( ComposeVectorImage, unsigned int, 3 );
    run( ComposeVectorImage, long, 3 );
    run( ComposeVectorImage, unsigned long, 3 );
    run( ComposeVectorImage, float, 3 );
    run( ComposeVectorImage, double, 3 );
  }
  catch( itk::ExceptionObject &e )
  {
    std::cerr << "Caught ITK exception: " << e << std::endl;
    return 1;
  }
  if ( !supported )
  {
    std::cerr << "ERROR: this combination of pixeltype and dimension is not supported!" << std::endl;
    std::cerr
      << "pixel (component) type = " << ComponentTypeIn
      << " ; dimension = " << Dimension
      << std::endl;
    return 1;
  }

  /** End program. */
  return 0;

} // end main

  /**
   * ******************* ComposeVectorImage *******************
   */

template< class InputImageType, class OutputImageType >
void ComposeVectorImage(
  const std::vector<std::string> & inputFileNames,
  const std::string & outputFileName,
  const unsigned int & numberOfStreams )
{
  /** Typedef's. */
  typedef itk::ImageFileReader< InputImageType >          ReaderType;
  typedef itk::Image<typename InputImageType::InternalPixelType, InputImageType::ImageDimension> ScalarImageType;
  typedef itk::ImageToVectorImageFilter< ScalarImageType > ImageToVectorImageFilterType;
  typedef itk::ImageFileWriter< OutputImageType >         WriterType;

  /** Read in the input images. */
  std::vector<typename ReaderType::Pointer> readers( inputFileNames.size() );
  for ( unsigned int i = 0; i < inputFileNames.size(); ++i )
  {
    readers[ i ] = ReaderType::New();
    readers[ i ]->SetFileName( inputFileNames[ i ] );
    readers[ i ]->Update();
  }

  /** Create index extractor and writer. */
  typename ImageToVectorImageFilterType::Pointer imageToVectorImageFilter = ImageToVectorImageFilterType::New();
  // For each input image
  std::cout << "There are " << inputFileNames.size() << " input images." << std::endl;
  unsigned int currentOutputIndex = 0;
  for ( unsigned int inputImageIndex = 0; inputImageIndex < inputFileNames.size(); ++inputImageIndex )
  {
    typedef itk::VectorIndexSelectionCastImageFilter<InputImageType, ScalarImageType> ComponentExtractionType;

    // For each component of the current image
    std::cout << "There are " << readers[inputImageIndex]->GetOutput()->GetNumberOfComponentsPerPixel() << " components in image "
              << inputImageIndex << std::endl;
    for ( unsigned int component = 0; component < readers[inputImageIndex]->GetOutput()->GetNumberOfComponentsPerPixel(); ++component )
    {
      typename ComponentExtractionType::Pointer componentExtractionFilter = ComponentExtractionType::New();
      componentExtractionFilter->SetIndex(component);
      componentExtractionFilter->SetInput(readers[inputImageIndex]->GetOutput());
      componentExtractionFilter->Update();

      imageToVectorImageFilter->SetNthInput( currentOutputIndex, componentExtractionFilter->GetOutput());
      currentOutputIndex++;
    }
  }

  imageToVectorImageFilter->Update();

  std::cout << "Output image has " << imageToVectorImageFilter->GetOutput()->GetNumberOfComponentsPerPixel() << " components." << std::endl;

  /** Write vector image. */
  typename WriterType::Pointer writer = WriterType::New();
  writer->SetFileName( outputFileName );
  writer->SetInput( imageToVectorImageFilter->GetOutput() );
  writer->SetNumberOfStreamDivisions( numberOfStreams );
  writer->Update();

} // end ComposeVectorImage()

/**
 * ******************* PrintHelp *******************
 */

void PrintHelp( void )
{
  std::cout << "Usage:" << std::endl << "pximagetovectorimage\n";
  std::cout << "  -in      inputFilenames, at least 2\n";
  std::cout << "  [-out]   outputFilename, default VECTOR.mhd\n";
  std::cout << "  [-s]     number of streams, default 1.\n";
  std::cout << "Supported: 2D, 3D, (unsigned) char, (unsigned) short, "
    << "(unsigned) int, (unsigned) long, float, double.\n";
  std::cout << "Note: make sure that the input images are of the same type, size, etc." << std::endl;

} // end PrintHelp()

<commit_msg>New style arguments for imagestovectorimage<commit_after>#include "itkCommandLineArgumentParser.h"
#include "CommandLineArgumentHelper.h"

#include "itkImageFileReader.h"
#include "itkImageToVectorImageFilter.h"
#include "itkImageFileWriter.h"
#include "itkVectorIndexSelectionCastImageFilter.h"

//-------------------------------------------------------------------------------------

/** run: A macro to call a function. */
#define run( function, type, dim ) \
if ( ComponentTypeIn == #type && Dimension == dim ) \
{ \
  typedef itk::VectorImage< type, dim >       InputImageType; \
  typedef itk::VectorImage< type, dim > OutputImageType; \
  function< InputImageType, OutputImageType >( inputFileNames, outputFileName, numberOfStreams ); \
  supported = true; \
}

//-------------------------------------------------------------------------------------

/** Declare ComposeVectorImage. */
template< class InputImageType, class OutputImageType >
void ComposeVectorImage(
  const std::vector<std::string> & inputFileNames,
  const std::string & outputFileName,
  const unsigned int & numberOfStreams );

/** Declare PrintHelp. */
std::string PrintHelp( void );

//-------------------------------------------------------------------------------------

int main( int argc, char ** argv )
{
  /** Create a command line argument parser. */
  itk::CommandLineArgumentParser::Pointer parser = itk::CommandLineArgumentParser::New();
  parser->SetCommandLineArguments( argc, argv );
  parser->SetProgramHelpText(PrintHelp());

  parser->MarkArgumentAsRequired( "-in", "The input filename." );

  bool validateArguments = parser->CheckForRequiredArguments();

  if(!validateArguments)
  {
    return EXIT_FAILURE;
  }

  /** Get arguments. */
  std::vector<std::string>  inputFileNames( 0, "" );
  parser->GetCommandLineArgument( "-in", inputFileNames );

  std::string outputFileName = "VECTOR.mhd";
  parser->GetCommandLineArgument( "-out", outputFileName );

  /** Support for streaming. */
  unsigned int numberOfStreams = 1;
  parser->GetCommandLineArgument( "-s", numberOfStreams );

  /** Check if the required arguments are given. */
  if ( inputFileNames.size() < 2 )
  {
    std::cerr << "ERROR: You should specify at least two (2) input files." << std::endl;
    return 1;
  }

  /** Determine image properties. */
  std::string ComponentTypeIn = "short";
  std::string PixelType; //we don't use this
  unsigned int Dimension = 3;
  unsigned int NumberOfComponents = 1;
  std::vector<unsigned int> imagesize( Dimension, 0 );
  int retgip = GetImageProperties(
    inputFileNames[ 0 ],
    PixelType,
    ComponentTypeIn,
    Dimension,
    NumberOfComponents,
    imagesize );
  if ( retgip != 0 )
  {
    return 1;
  }

  /** Get rid of the possible "_" in ComponentType. */
  ReplaceUnderscoreWithSpace( ComponentTypeIn );

  /** Run the program. */
  bool supported = false;
  try
  {
    run( ComposeVectorImage, char, 2 );
    run( ComposeVectorImage, unsigned char, 2 );
    run( ComposeVectorImage, short, 2 );
    run( ComposeVectorImage, unsigned short, 2 );
    run( ComposeVectorImage, int, 2 );
    run( ComposeVectorImage, unsigned int, 2 );
    run( ComposeVectorImage, long, 2 );
    run( ComposeVectorImage, unsigned long, 2 );
    run( ComposeVectorImage, float, 2 );
    run( ComposeVectorImage, double, 2 );

    run( ComposeVectorImage, char, 3 );
    run( ComposeVectorImage, unsigned char, 3 );
    run( ComposeVectorImage, short, 3 );
    run( ComposeVectorImage, unsigned short, 3 );
    run( ComposeVectorImage, int, 3 );
    run( ComposeVectorImage, unsigned int, 3 );
    run( ComposeVectorImage, long, 3 );
    run( ComposeVectorImage, unsigned long, 3 );
    run( ComposeVectorImage, float, 3 );
    run( ComposeVectorImage, double, 3 );
  }
  catch( itk::ExceptionObject &e )
  {
    std::cerr << "Caught ITK exception: " << e << std::endl;
    return 1;
  }
  if ( !supported )
  {
    std::cerr << "ERROR: this combination of pixeltype and dimension is not supported!" << std::endl;
    std::cerr
      << "pixel (component) type = " << ComponentTypeIn
      << " ; dimension = " << Dimension
      << std::endl;
    return 1;
  }

  /** End program. */
  return 0;

} // end main

  /**
   * ******************* ComposeVectorImage *******************
   */

template< class InputImageType, class OutputImageType >
void ComposeVectorImage(
  const std::vector<std::string> & inputFileNames,
  const std::string & outputFileName,
  const unsigned int & numberOfStreams )
{
  /** Typedef's. */
  typedef itk::ImageFileReader< InputImageType >          ReaderType;
  typedef itk::Image<typename InputImageType::InternalPixelType, InputImageType::ImageDimension> ScalarImageType;
  typedef itk::ImageToVectorImageFilter< ScalarImageType > ImageToVectorImageFilterType;
  typedef itk::ImageFileWriter< OutputImageType >         WriterType;

  /** Read in the input images. */
  std::vector<typename ReaderType::Pointer> readers( inputFileNames.size() );
  for ( unsigned int i = 0; i < inputFileNames.size(); ++i )
  {
    readers[ i ] = ReaderType::New();
    readers[ i ]->SetFileName( inputFileNames[ i ] );
    readers[ i ]->Update();
  }

  /** Create index extractor and writer. */
  typename ImageToVectorImageFilterType::Pointer imageToVectorImageFilter = ImageToVectorImageFilterType::New();
  // For each input image
  std::cout << "There are " << inputFileNames.size() << " input images." << std::endl;
  unsigned int currentOutputIndex = 0;
  for ( unsigned int inputImageIndex = 0; inputImageIndex < inputFileNames.size(); ++inputImageIndex )
  {
    typedef itk::VectorIndexSelectionCastImageFilter<InputImageType, ScalarImageType> ComponentExtractionType;

    // For each component of the current image
    std::cout << "There are " << readers[inputImageIndex]->GetOutput()->GetNumberOfComponentsPerPixel() << " components in image "
              << inputImageIndex << std::endl;
    for ( unsigned int component = 0; component < readers[inputImageIndex]->GetOutput()->GetNumberOfComponentsPerPixel(); ++component )
    {
      typename ComponentExtractionType::Pointer componentExtractionFilter = ComponentExtractionType::New();
      componentExtractionFilter->SetIndex(component);
      componentExtractionFilter->SetInput(readers[inputImageIndex]->GetOutput());
      componentExtractionFilter->Update();

      imageToVectorImageFilter->SetNthInput( currentOutputIndex, componentExtractionFilter->GetOutput());
      currentOutputIndex++;
    }
  }

  imageToVectorImageFilter->Update();

  std::cout << "Output image has " << imageToVectorImageFilter->GetOutput()->GetNumberOfComponentsPerPixel() << " components." << std::endl;

  /** Write vector image. */
  typename WriterType::Pointer writer = WriterType::New();
  writer->SetFileName( outputFileName );
  writer->SetInput( imageToVectorImageFilter->GetOutput() );
  writer->SetNumberOfStreamDivisions( numberOfStreams );
  writer->Update();

} // end ComposeVectorImage()

/**
 * ******************* PrintHelp *******************
 */

std::string PrintHelp( void )
{
  std::string helpText = "Usage: \
  pximagetovectorimage\n \
    -in      inputFilenames, at least 2\n \
    [-out]   outputFilename, default VECTOR.mhd\n \
    [-s]     number of streams, default 1.\n \
  Supported: 2D, 3D, (unsigned) char, (unsigned) short, \
  (unsigned) int, (unsigned) long, float, double.\n \
  Note: make sure that the input images are of the same type, size, etc.";

  return helpText;

} // end PrintHelp()

<|endoftext|>