<commit_before>/*
 * SessionConsoleProcess.cpp
 *
 * Copyright (C) 2009-11 by RStudio, Inc.
 *
 * This program is licensed to you under the terms of version 3 of the
 * GNU Affero General Public License. This program is distributed WITHOUT
 * ANY EXPRESS OR IMPLIED WARRANTY, INCLUDING THOSE OF NON-INFRINGEMENT,
 * MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. Please refer to the
 * AGPL (http://www.gnu.org/licenses/agpl-3.0.txt) for more details.
 *
 */
#include "SessionConsoleProcess.hpp"

#include <core/json/JsonRpc.hpp>
#include <core/system/Process.hpp>
#include <core/Exec.hpp>
#include <core/SafeConvert.hpp>
#include <core/Settings.hpp>

#include <core/system/Environment.hpp>

#include <session/SessionModuleContext.hpp>

#include "config.h"

#ifdef RSTUDIO_SERVER
#include <core/system/Crypto.hpp>
#endif

using namespace core;

namespace session {
namespace modules {
namespace console_process {

namespace {
   typedef std::map<std::string, boost::shared_ptr<ConsoleProcess> > ProcTable;
   ProcTable s_procs;

   core::system::ProcessOptions procOptions()
   {
      core::system::ProcessOptions options;
#ifndef _WIN32
      options.detachSession = true;
#endif
      return options;
   }
} // anonymous namespace

#ifdef _WIN32
const int kDefaultMaxOutputLines = 140;
#else
const int kDefaultMaxOutputLines = 500;
#endif

ConsoleProcess::ConsoleProcess()
   : dialog_(false), interactionMode_(InteractionNever),
     maxOutputLines_(kDefaultMaxOutputLines), started_(true),
     interrupt_(false), outputBuffer_(maxOutputLines_)
{
   // When we retrieve from outputBuffer, we only want complete lines. Add a
   // dummy \n so we can tell the first line is a complete line.
   outputBuffer_.push_back('\n');
}

ConsoleProcess::ConsoleProcess(const std::string& command,
                               const core::system::ProcessOptions& options,
                               const std::string& caption,
                               bool dialog,
                               InteractionMode interactionMode,
                               int maxOutputLines,
                               const boost::function<void()>& onExit)
   : command_(command), options_(options), caption_(caption), dialog_(dialog),
     interactionMode_(interactionMode), maxOutputLines_(maxOutputLines),
     started_(false), interrupt_(false),
     outputBuffer_(maxOutputLines_),
     onExit_(onExit)
{
   commonInit();
}

ConsoleProcess::ConsoleProcess(const std::string& program,
                               const std::vector<std::string>& args,
                               const core::system::ProcessOptions& options,
                               const std::string& caption,
                               bool dialog,
                               InteractionMode interactionMode,
                               int maxOutputLines,
                               const boost::function<void()>& onExit)
   : program_(program), args_(args), options_(options), caption_(caption), dialog_(dialog),
     interactionMode_(interactionMode), maxOutputLines_(maxOutputLines),
     started_(false),  interrupt_(false),
     outputBuffer_(maxOutputLines_),
     onExit_(onExit)
{
   commonInit();
}

void ConsoleProcess::commonInit()
{
   handle_ = core::system::generateUuid(false);

   // request a pseudoterminal if this is an interactive console process
#ifndef _WIN32
   if (interactionMode() != InteractionNever)
   {
      options_.pseudoterminal = core::system::Pseudoterminal(80, 1);

      // define TERM to dumb (but first make sure we have an environment
      // block to modify)
      if (!options_.environment)
      {
         core::system::Options childEnv;
         core::system::environment(&childEnv);
         options_.environment = childEnv;
      }
      core::system::setenv(&(options_.environment.get()), "TERM", "dumb");

   }
#endif

   // When we retrieve from outputBuffer, we only want complete lines. Add a
   // dummy \n so we can tell the first line is a complete line.
   outputBuffer_.push_back('\n');
}

std::string ConsoleProcess::bufferedOutput() const
{
   boost::circular_buffer<char>::const_iterator pos =
         std::find(outputBuffer_.begin(), outputBuffer_.end(), '\n');

   std::string result;
   if (pos != outputBuffer_.end())
      pos++;
   std::copy(pos, outputBuffer_.end(), std::back_inserter(result));
   // Will be empty if the buffer was overflowed by a single line
   return result;
}

Error ConsoleProcess::start()
{
   if (started_)
      return Success();

   Error error;
   if (!command_.empty())
   {
      error = module_context::processSupervisor().runCommand(
                                 command_, options_, createProcessCallbacks());
   }
   else
   {
      error = module_context::processSupervisor().runProgram(
                          program_, args_, options_, createProcessCallbacks());
   }
   if (!error)
      started_ = true;
   return error;
}

void ConsoleProcess::enqueueInput(const std::string &input)
{
   inputQueue_.push(Input(input));
}

void ConsoleProcess::interrupt()
{
   interrupt_ = true;
}

void ConsoleProcess::enquePtyInterrupt()
{
   inputQueue_.push(Input());
}

bool ConsoleProcess::onContinue(core::system::ProcessOperations& ops)
{
   // full stop interrupt if requested
   if (interrupt_)
      return false;

   // process input queue
   while (!inputQueue_.empty())
   {
      // pop input
      Input input = inputQueue_.front();
      inputQueue_.pop();

      // pty interrupt
      if (input.interrupt)
      {
         Error error = ops.ptyInterrupt();
         if (error)
            LOG_ERROR(error);
      }

      // text input
      else
      {
         Error error = ops.writeToStdin(input.text, false);
         if (error)
            LOG_ERROR(error);

         // add a newline to the output buffer (we don't echo back the
         // actual input because it is very likely to be a password
         // or passphrase)
         outputBuffer_.push_back('\n');
      }
   }

   // continue
   return true;
}

void ConsoleProcess::enqueOutputEvent(const std::string &output, bool error)
{
   // copy to output buffer
   std::copy(output.begin(), output.end(),
             std::back_inserter(outputBuffer_));

   // If there's more output than the client can even show, then
   // truncate it to the amount that the client can show. Too much
   // output can overwhelm the client, making it unresponsive.
   std::string trimmedOutput = output;
   string_utils::trimLeadingLines(maxOutputLines_, &trimmedOutput);

   json::Object data;
   data["handle"] = handle_;
   data["error"] = error;
   data["output"] = trimmedOutput;
   module_context::enqueClientEvent(
         ClientEvent(client_events::kConsoleProcessOutput, data));
}

void ConsoleProcess::onStdout(core::system::ProcessOperations& ops,
                              const std::string& output)
{
   enqueOutputEvent(output, false);
}

void ConsoleProcess::onStderr(core::system::ProcessOperations& ops,
                                     const std::string& output)
{
   enqueOutputEvent(output, true);
}

void ConsoleProcess::onExit(int exitCode)
{
   exitCode_.reset(exitCode);

   json::Object data;
   data["handle"] = handle_;
   data["exitCode"] = exitCode;
   module_context::enqueClientEvent(
         ClientEvent(client_events::kConsoleProcessExit, data));

   if (onExit_)
      onExit_();
}

core::json::Object ConsoleProcess::toJson() const
{
   json::Object result;
   result["handle"] = handle_;
   result["caption"] = caption_;
   result["dialog"] = dialog_;
   result["interaction_mode"] = static_cast<int>(interactionMode_);
   result["max_output_lines"] = maxOutputLines_;
   result["buffered_output"] = bufferedOutput();
   if (exitCode_)
      result["exit_code"] = *exitCode_;
   else
      result["exit_code"] = json::Value();
   return result;
}

boost::shared_ptr<ConsoleProcess> ConsoleProcess::fromJson(
                                             core::json::Object &obj)
{
   boost::shared_ptr<ConsoleProcess> pProc(new ConsoleProcess());
   pProc->handle_ = obj["handle"].get_str();
   pProc->caption_ = obj["caption"].get_str();
   pProc->dialog_ = obj["dialog"].get_bool();

   json::Value mode = obj["interaction_mode"];
   if (!mode.is_null())
      pProc->interactionMode_ = static_cast<InteractionMode>(mode.get_int());
   else
      pProc->interactionMode_ = InteractionNever;

   json::Value maxLines = obj["max_output_lines"];
   if (!maxLines.is_null())
      pProc->maxOutputLines_ = maxLines.get_int();
   else
      pProc->maxOutputLines_ = kDefaultMaxOutputLines;

   std::string bufferedOutput = obj["buffered_output"].get_str();
   std::copy(bufferedOutput.begin(), bufferedOutput.end(),
             std::back_inserter(pProc->outputBuffer_));
   json::Value exitCode = obj["exit_code"];
   if (exitCode.is_null())
      pProc->exitCode_.reset();
   else
      pProc->exitCode_.reset(exitCode.get_int());

   return pProc;
}

core::system::ProcessCallbacks ConsoleProcess::createProcessCallbacks()
{
   core::system::ProcessCallbacks cb;
   cb.onContinue = boost::bind(&ConsoleProcess::onContinue, ConsoleProcess::shared_from_this(), _1);
   cb.onStdout = boost::bind(&ConsoleProcess::onStdout, ConsoleProcess::shared_from_this(), _1, _2);
   cb.onStderr = boost::bind(&ConsoleProcess::onStderr, ConsoleProcess::shared_from_this(), _1, _2);
   cb.onExit = boost::bind(&ConsoleProcess::onExit, ConsoleProcess::shared_from_this(), _1);
   return cb;
}

Error procStart(const json::JsonRpcRequest& request,
                json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;
   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
      return pos->second->start();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

Error procInterrupt(const json::JsonRpcRequest& request,
                    json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;
   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
      pos->second->interrupt();
      return Success();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

Error procPtyInterrupt(const json::JsonRpcRequest& request,
                       json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;
   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
      pos->second->enquePtyInterrupt();
      return Success();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

Error procReap(const json::JsonRpcRequest& request,
               json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;

   if (!s_procs.erase(handle))
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
   else
   {
      return Success();
   }
}

Error procWriteStdin(const json::JsonRpcRequest& request,
                     json::JsonRpcResponse* pResponse)
{
   std::string handle, input;
   Error error = json::readParams(request.params, &handle, &input);
   if (error)
      return error;

   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
#ifdef RSTUDIO_SERVER
      error = core::system::crypto::rsaPrivateDecrypt(input, &input);
      if (error)
         return error;
#endif

      pos->second->enqueueInput(input);

      return Success();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

boost::shared_ptr<ConsoleProcess> ConsoleProcess::create(
      const std::string& command,
      core::system::ProcessOptions options,
      const std::string& caption,
      bool dialog,
      InteractionMode interactionMode,
      int maxOutputLines,
      const boost::function<void()>& onExit)
{
   options.terminateChildren = true;
   boost::shared_ptr<ConsoleProcess> ptrProc(
         new ConsoleProcess(command,
                            options,
                            caption,
                            dialog,
                            interactionMode,
                            maxOutputLines,
                            onExit));
   s_procs[ptrProc->handle()] = ptrProc;
   return ptrProc;
}

boost::shared_ptr<ConsoleProcess> ConsoleProcess::create(
      const std::string& program,
      const std::vector<std::string>& args,
      core::system::ProcessOptions options,
      const std::string& caption,
      bool dialog,
      InteractionMode interactionMode,
      int maxOutputLines,
      const boost::function<void()>& onExit)
{
   options.terminateChildren = true;
   boost::shared_ptr<ConsoleProcess> ptrProc(
         new ConsoleProcess(program,
                            args,
                            options,
                            caption,
                            dialog,
                            interactionMode,
                            maxOutputLines,
                            onExit));
   s_procs[ptrProc->handle()] = ptrProc;
   return ptrProc;
}

core::json::Array processesAsJson()
{
   json::Array procInfos;
   for (ProcTable::const_iterator it = s_procs.begin();
        it != s_procs.end();
        it++)
   {
      procInfos.push_back(it->second->toJson());
   }
   return procInfos;
}

void onSuspend(core::Settings* pSettings)
{
   json::Array array;
   for (ProcTable::const_iterator it = s_procs.begin();
        it != s_procs.end();
        it++)
   {
      array.push_back(it->second->toJson());
   }

   std::ostringstream ostr;
   json::write(array, ostr);
   pSettings->set("console_procs", ostr.str());
}

void onResume(const core::Settings& settings)
{
   std::string strVal = settings.get("console_procs");
   if (strVal.empty())
      return;

   json::Value value;
   if (!json::parse(strVal, &value))
      return;

   json::Array procs = value.get_array();
   for (json::Array::iterator it = procs.begin();
        it != procs.end();
        it++)
   {
      boost::shared_ptr<ConsoleProcess> proc =
                                    ConsoleProcess::fromJson(it->get_obj());
      s_procs[proc->handle()] = proc;
   }
}

Error initialize()
{
   using boost::bind;
   using namespace module_context;

   // add suspend/resume handler
   addSuspendHandler(SuspendHandler(onSuspend, onResume));

   // install rpc methods
   ExecBlock initBlock ;
   initBlock.addFunctions()
      (bind(registerRpcMethod, "process_start", procStart))
      (bind(registerRpcMethod, "process_interrupt", procInterrupt))
      (bind(registerRpcMethod, "process_pty_interrupt", procInterrupt))
      (bind(registerRpcMethod, "process_reap", procReap))
      (bind(registerRpcMethod, "process_write_stdin", procWriteStdin));

   return initBlock.execute();
}

} // namespace console_process
} // namespace modules
} // namespace session
<commit_msg>restore 8k console buffer size<commit_after>/*
 * SessionConsoleProcess.cpp
 *
 * Copyright (C) 2009-11 by RStudio, Inc.
 *
 * This program is licensed to you under the terms of version 3 of the
 * GNU Affero General Public License. This program is distributed WITHOUT
 * ANY EXPRESS OR IMPLIED WARRANTY, INCLUDING THOSE OF NON-INFRINGEMENT,
 * MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. Please refer to the
 * AGPL (http://www.gnu.org/licenses/agpl-3.0.txt) for more details.
 *
 */
#include "SessionConsoleProcess.hpp"

#include <core/json/JsonRpc.hpp>
#include <core/system/Process.hpp>
#include <core/Exec.hpp>
#include <core/SafeConvert.hpp>
#include <core/Settings.hpp>

#include <core/system/Environment.hpp>

#include <session/SessionModuleContext.hpp>

#include "config.h"

#ifdef RSTUDIO_SERVER
#include <core/system/Crypto.hpp>
#endif

using namespace core;

namespace session {
namespace modules {
namespace console_process {

namespace {
   const size_t OUTPUT_BUFFER_SIZE = 8192;
   typedef std::map<std::string, boost::shared_ptr<ConsoleProcess> > ProcTable;
   ProcTable s_procs;

   core::system::ProcessOptions procOptions()
   {
      core::system::ProcessOptions options;
#ifndef _WIN32
      options.detachSession = true;
#endif
      return options;
   }
} // anonymous namespace

#ifdef _WIN32
const int kDefaultMaxOutputLines = 140;
#else
const int kDefaultMaxOutputLines = 500;
#endif

ConsoleProcess::ConsoleProcess()
   : dialog_(false), interactionMode_(InteractionNever),
     maxOutputLines_(kDefaultMaxOutputLines), started_(true),
     interrupt_(false), outputBuffer_(OUTPUT_BUFFER_SIZE)
{
   // When we retrieve from outputBuffer, we only want complete lines. Add a
   // dummy \n so we can tell the first line is a complete line.
   outputBuffer_.push_back('\n');
}

ConsoleProcess::ConsoleProcess(const std::string& command,
                               const core::system::ProcessOptions& options,
                               const std::string& caption,
                               bool dialog,
                               InteractionMode interactionMode,
                               int maxOutputLines,
                               const boost::function<void()>& onExit)
   : command_(command), options_(options), caption_(caption), dialog_(dialog),
     interactionMode_(interactionMode), maxOutputLines_(maxOutputLines),
     started_(false), interrupt_(false),
     outputBuffer_(OUTPUT_BUFFER_SIZE),
     onExit_(onExit)
{
   commonInit();
}

ConsoleProcess::ConsoleProcess(const std::string& program,
                               const std::vector<std::string>& args,
                               const core::system::ProcessOptions& options,
                               const std::string& caption,
                               bool dialog,
                               InteractionMode interactionMode,
                               int maxOutputLines,
                               const boost::function<void()>& onExit)
   : program_(program), args_(args), options_(options), caption_(caption), dialog_(dialog),
     interactionMode_(interactionMode), maxOutputLines_(maxOutputLines),
     started_(false),  interrupt_(false),
     outputBuffer_(OUTPUT_BUFFER_SIZE),
     onExit_(onExit)
{
   commonInit();
}

void ConsoleProcess::commonInit()
{
   handle_ = core::system::generateUuid(false);

   // request a pseudoterminal if this is an interactive console process
#ifndef _WIN32
   if (interactionMode() != InteractionNever)
   {
      options_.pseudoterminal = core::system::Pseudoterminal(80, 1);

      // define TERM to dumb (but first make sure we have an environment
      // block to modify)
      if (!options_.environment)
      {
         core::system::Options childEnv;
         core::system::environment(&childEnv);
         options_.environment = childEnv;
      }
      core::system::setenv(&(options_.environment.get()), "TERM", "dumb");

   }
#endif

   // When we retrieve from outputBuffer, we only want complete lines. Add a
   // dummy \n so we can tell the first line is a complete line.
   outputBuffer_.push_back('\n');
}

std::string ConsoleProcess::bufferedOutput() const
{
   boost::circular_buffer<char>::const_iterator pos =
         std::find(outputBuffer_.begin(), outputBuffer_.end(), '\n');

   std::string result;
   if (pos != outputBuffer_.end())
      pos++;
   std::copy(pos, outputBuffer_.end(), std::back_inserter(result));
   // Will be empty if the buffer was overflowed by a single line
   return result;
}

Error ConsoleProcess::start()
{
   if (started_)
      return Success();

   Error error;
   if (!command_.empty())
   {
      error = module_context::processSupervisor().runCommand(
                                 command_, options_, createProcessCallbacks());
   }
   else
   {
      error = module_context::processSupervisor().runProgram(
                          program_, args_, options_, createProcessCallbacks());
   }
   if (!error)
      started_ = true;
   return error;
}

void ConsoleProcess::enqueueInput(const std::string &input)
{
   inputQueue_.push(Input(input));
}

void ConsoleProcess::interrupt()
{
   interrupt_ = true;
}

void ConsoleProcess::enquePtyInterrupt()
{
   inputQueue_.push(Input());
}

bool ConsoleProcess::onContinue(core::system::ProcessOperations& ops)
{
   // full stop interrupt if requested
   if (interrupt_)
      return false;

   // process input queue
   while (!inputQueue_.empty())
   {
      // pop input
      Input input = inputQueue_.front();
      inputQueue_.pop();

      // pty interrupt
      if (input.interrupt)
      {
         Error error = ops.ptyInterrupt();
         if (error)
            LOG_ERROR(error);
      }

      // text input
      else
      {
         Error error = ops.writeToStdin(input.text, false);
         if (error)
            LOG_ERROR(error);

         // add a newline to the output buffer (we don't echo back the
         // actual input because it is very likely to be a password
         // or passphrase)
         outputBuffer_.push_back('\n');
      }
   }

   // continue
   return true;
}

void ConsoleProcess::enqueOutputEvent(const std::string &output, bool error)
{
   // copy to output buffer
   std::copy(output.begin(), output.end(),
             std::back_inserter(outputBuffer_));

   // If there's more output than the client can even show, then
   // truncate it to the amount that the client can show. Too much
   // output can overwhelm the client, making it unresponsive.
   std::string trimmedOutput = output;
   string_utils::trimLeadingLines(maxOutputLines_, &trimmedOutput);

   json::Object data;
   data["handle"] = handle_;
   data["error"] = error;
   data["output"] = trimmedOutput;
   module_context::enqueClientEvent(
         ClientEvent(client_events::kConsoleProcessOutput, data));
}

void ConsoleProcess::onStdout(core::system::ProcessOperations& ops,
                              const std::string& output)
{
   enqueOutputEvent(output, false);
}

void ConsoleProcess::onStderr(core::system::ProcessOperations& ops,
                                     const std::string& output)
{
   enqueOutputEvent(output, true);
}

void ConsoleProcess::onExit(int exitCode)
{
   exitCode_.reset(exitCode);

   json::Object data;
   data["handle"] = handle_;
   data["exitCode"] = exitCode;
   module_context::enqueClientEvent(
         ClientEvent(client_events::kConsoleProcessExit, data));

   if (onExit_)
      onExit_();
}

core::json::Object ConsoleProcess::toJson() const
{
   json::Object result;
   result["handle"] = handle_;
   result["caption"] = caption_;
   result["dialog"] = dialog_;
   result["interaction_mode"] = static_cast<int>(interactionMode_);
   result["max_output_lines"] = maxOutputLines_;
   result["buffered_output"] = bufferedOutput();
   if (exitCode_)
      result["exit_code"] = *exitCode_;
   else
      result["exit_code"] = json::Value();
   return result;
}

boost::shared_ptr<ConsoleProcess> ConsoleProcess::fromJson(
                                             core::json::Object &obj)
{
   boost::shared_ptr<ConsoleProcess> pProc(new ConsoleProcess());
   pProc->handle_ = obj["handle"].get_str();
   pProc->caption_ = obj["caption"].get_str();
   pProc->dialog_ = obj["dialog"].get_bool();

   json::Value mode = obj["interaction_mode"];
   if (!mode.is_null())
      pProc->interactionMode_ = static_cast<InteractionMode>(mode.get_int());
   else
      pProc->interactionMode_ = InteractionNever;

   json::Value maxLines = obj["max_output_lines"];
   if (!maxLines.is_null())
      pProc->maxOutputLines_ = maxLines.get_int();
   else
      pProc->maxOutputLines_ = kDefaultMaxOutputLines;

   std::string bufferedOutput = obj["buffered_output"].get_str();
   std::copy(bufferedOutput.begin(), bufferedOutput.end(),
             std::back_inserter(pProc->outputBuffer_));
   json::Value exitCode = obj["exit_code"];
   if (exitCode.is_null())
      pProc->exitCode_.reset();
   else
      pProc->exitCode_.reset(exitCode.get_int());

   return pProc;
}

core::system::ProcessCallbacks ConsoleProcess::createProcessCallbacks()
{
   core::system::ProcessCallbacks cb;
   cb.onContinue = boost::bind(&ConsoleProcess::onContinue, ConsoleProcess::shared_from_this(), _1);
   cb.onStdout = boost::bind(&ConsoleProcess::onStdout, ConsoleProcess::shared_from_this(), _1, _2);
   cb.onStderr = boost::bind(&ConsoleProcess::onStderr, ConsoleProcess::shared_from_this(), _1, _2);
   cb.onExit = boost::bind(&ConsoleProcess::onExit, ConsoleProcess::shared_from_this(), _1);
   return cb;
}

Error procStart(const json::JsonRpcRequest& request,
                json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;
   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
      return pos->second->start();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

Error procInterrupt(const json::JsonRpcRequest& request,
                    json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;
   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
      pos->second->interrupt();
      return Success();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

Error procPtyInterrupt(const json::JsonRpcRequest& request,
                       json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;
   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
      pos->second->enquePtyInterrupt();
      return Success();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

Error procReap(const json::JsonRpcRequest& request,
               json::JsonRpcResponse* pResponse)
{
   std::string handle;
   Error error = json::readParams(request.params, &handle);
   if (error)
      return error;

   if (!s_procs.erase(handle))
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
   else
   {
      return Success();
   }
}

Error procWriteStdin(const json::JsonRpcRequest& request,
                     json::JsonRpcResponse* pResponse)
{
   std::string handle, input;
   Error error = json::readParams(request.params, &handle, &input);
   if (error)
      return error;

   ProcTable::const_iterator pos = s_procs.find(handle);
   if (pos != s_procs.end())
   {
#ifdef RSTUDIO_SERVER
      error = core::system::crypto::rsaPrivateDecrypt(input, &input);
      if (error)
         return error;
#endif

      pos->second->enqueueInput(input);

      return Success();
   }
   else
   {
      return systemError(boost::system::errc::invalid_argument,
                         ERROR_LOCATION);
   }
}

boost::shared_ptr<ConsoleProcess> ConsoleProcess::create(
      const std::string& command,
      core::system::ProcessOptions options,
      const std::string& caption,
      bool dialog,
      InteractionMode interactionMode,
      int maxOutputLines,
      const boost::function<void()>& onExit)
{
   options.terminateChildren = true;
   boost::shared_ptr<ConsoleProcess> ptrProc(
         new ConsoleProcess(command,
                            options,
                            caption,
                            dialog,
                            interactionMode,
                            maxOutputLines,
                            onExit));
   s_procs[ptrProc->handle()] = ptrProc;
   return ptrProc;
}

boost::shared_ptr<ConsoleProcess> ConsoleProcess::create(
      const std::string& program,
      const std::vector<std::string>& args,
      core::system::ProcessOptions options,
      const std::string& caption,
      bool dialog,
      InteractionMode interactionMode,
      int maxOutputLines,
      const boost::function<void()>& onExit)
{
   options.terminateChildren = true;
   boost::shared_ptr<ConsoleProcess> ptrProc(
         new ConsoleProcess(program,
                            args,
                            options,
                            caption,
                            dialog,
                            interactionMode,
                            maxOutputLines,
                            onExit));
   s_procs[ptrProc->handle()] = ptrProc;
   return ptrProc;
}

core::json::Array processesAsJson()
{
   json::Array procInfos;
   for (ProcTable::const_iterator it = s_procs.begin();
        it != s_procs.end();
        it++)
   {
      procInfos.push_back(it->second->toJson());
   }
   return procInfos;
}

void onSuspend(core::Settings* pSettings)
{
   json::Array array;
   for (ProcTable::const_iterator it = s_procs.begin();
        it != s_procs.end();
        it++)
   {
      array.push_back(it->second->toJson());
   }

   std::ostringstream ostr;
   json::write(array, ostr);
   pSettings->set("console_procs", ostr.str());
}

void onResume(const core::Settings& settings)
{
   std::string strVal = settings.get("console_procs");
   if (strVal.empty())
      return;

   json::Value value;
   if (!json::parse(strVal, &value))
      return;

   json::Array procs = value.get_array();
   for (json::Array::iterator it = procs.begin();
        it != procs.end();
        it++)
   {
      boost::shared_ptr<ConsoleProcess> proc =
                                    ConsoleProcess::fromJson(it->get_obj());
      s_procs[proc->handle()] = proc;
   }
}

Error initialize()
{
   using boost::bind;
   using namespace module_context;

   // add suspend/resume handler
   addSuspendHandler(SuspendHandler(onSuspend, onResume));

   // install rpc methods
   ExecBlock initBlock ;
   initBlock.addFunctions()
      (bind(registerRpcMethod, "process_start", procStart))
      (bind(registerRpcMethod, "process_interrupt", procInterrupt))
      (bind(registerRpcMethod, "process_pty_interrupt", procInterrupt))
      (bind(registerRpcMethod, "process_reap", procReap))
      (bind(registerRpcMethod, "process_write_stdin", procWriteStdin));

   return initBlock.execute();
}

} // namespace console_process
} // namespace modules
} // namespace session
<|endoftext|>


<commit_before>// Copyright (c) 2016 Greenheart Games Pty. Ltd. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include <memory>

#include "nan.h"
#include "steam/steam_api.h"
#include "v8.h"

#include "greenworks_utils.h"
#include "steam_api_registry.h"
#include "steam_id.h"

namespace greenworks {
namespace api {
namespace {

void InitFriendFlags(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> friend_flags = Nan::New<v8::Object>();
  SET_TYPE(friend_flags, "None", k_EFriendFlagNone);
  SET_TYPE(friend_flags, "Blocked", k_EFriendFlagBlocked);
  SET_TYPE(friend_flags, "FriendshipRequested",
           k_EFriendFlagFriendshipRequested);
  SET_TYPE(friend_flags, "Immediate", k_EFriendFlagImmediate);
  SET_TYPE(friend_flags, "ClanMember", k_EFriendFlagClanMember);
  SET_TYPE(friend_flags, "OnGameServer", k_EFriendFlagOnGameServer);
  SET_TYPE(friend_flags, "RequestingFriendship",
           k_EFriendFlagRequestingFriendship);
  SET_TYPE(friend_flags, "RequestingInfo", k_EFriendFlagRequestingInfo);
  SET_TYPE(friend_flags, "Ignored", k_EFriendFlagIgnored);
  SET_TYPE(friend_flags, "IgnoredFriend", k_EFriendFlagIgnoredFriend);
  // k_EFriendFlagSuggested has been removed since steamworks sdk 1.39.
  // SET_TYPE(friend_flags, "Suggested", k_EFriendFlagSuggested);
  SET_TYPE(friend_flags, "ChatMember", k_EFriendFlagChatMember);
  SET_TYPE(friend_flags, "All", k_EFriendFlagAll);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(friend_flags);
  Nan::Set(exports,
           Nan::New("FriendFlags").ToLocalChecked(),
           friend_flags);
}

void InitFriendRelationship(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> relationship = Nan::New<v8::Object>();
  SET_TYPE(relationship, "None", k_EFriendRelationshipNone);
  SET_TYPE(relationship, "Blocked", k_EFriendRelationshipBlocked);
  SET_TYPE(relationship, "RequestRecipient",
           k_EFriendRelationshipRequestRecipient);
  SET_TYPE(relationship, "Friend", k_EFriendRelationshipFriend);
  SET_TYPE(relationship, "RequestInitiator",
           k_EFriendRelationshipRequestInitiator);
  SET_TYPE(relationship, "Ignored", k_EFriendRelationshipIgnored);
  SET_TYPE(relationship, "IgnoredFriend", k_EFriendRelationshipIgnoredFriend);
  SET_TYPE(relationship, "Suggested",
           k_EFriendRelationshipSuggested_DEPRECATED);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(relationship);
  Nan::Set(exports,
           Nan::New("FriendRelationship").ToLocalChecked(),
           relationship);
}

void InitFriendPersonaChange(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> persona_change = Nan::New<v8::Object>();
  SET_TYPE(persona_change, "Name", k_EPersonaChangeName);
  SET_TYPE(persona_change, "Status", k_EPersonaChangeStatus);
  SET_TYPE(persona_change, "ComeOnline", k_EPersonaChangeComeOnline);
  SET_TYPE(persona_change, "GoneOffline", k_EPersonaChangeGoneOffline);
  SET_TYPE(persona_change, "GamePlayed", k_EPersonaChangeGamePlayed);
  SET_TYPE(persona_change, "GameServer", k_EPersonaChangeGameServer);
  SET_TYPE(persona_change, "Avator", k_EPersonaChangeAvatar);
  SET_TYPE(persona_change, "JoinedSource", k_EPersonaChangeJoinedSource);
  SET_TYPE(persona_change, "LeftSource", k_EPersonaChangeLeftSource);
  SET_TYPE(persona_change, "RelationshipChanged",
           k_EPersonaChangeRelationshipChanged);
  SET_TYPE(persona_change, "NameFirstSet", k_EPersonaChangeNameFirstSet);
  SET_TYPE(persona_change, "FacebookInfo", k_EPersonaChangeFacebookInfo);
  SET_TYPE(persona_change, "NickName", k_EPersonaChangeNickname);
  SET_TYPE(persona_change, "SteamLevel", k_EPersonaChangeSteamLevel);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(persona_change);
  Nan::Set(exports,
           Nan::New("PersonaChange").ToLocalChecked(),
           persona_change);
}

void InitAccountType(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> account_type = Nan::New<v8::Object>();
  SET_TYPE(account_type, "Invalid", k_EAccountTypeInvalid);
  SET_TYPE(account_type, "Individual", k_EAccountTypeIndividual);
  SET_TYPE(account_type, "Multiseat", k_EAccountTypeMultiseat);
  SET_TYPE(account_type, "GameServer", k_EAccountTypeGameServer);
  SET_TYPE(account_type, "AnonymousGameServer", k_EAccountTypeAnonGameServer);
  SET_TYPE(account_type, "Pending", k_EAccountTypePending);
  SET_TYPE(account_type, "ContentServer", k_EAccountTypeContentServer);
  SET_TYPE(account_type, "Clan", k_EAccountTypeClan);
  SET_TYPE(account_type, "Chat", k_EAccountTypeChat);
  SET_TYPE(account_type, "ConsoleUser", k_EAccountTypeConsoleUser);
  SET_TYPE(account_type, "AnonymousUser", k_EAccountTypeAnonUser);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(account_type);
  Nan::Set(exports,
           Nan::New("AccountType").ToLocalChecked(),
           account_type);
}

void InitChatEntryType(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> chat_entry_type = Nan::New<v8::Object>();
  SET_TYPE(chat_entry_type, "Invalid", k_EChatEntryTypeInvalid);
  SET_TYPE(chat_entry_type, "ChatMsg", k_EChatEntryTypeChatMsg);
  SET_TYPE(chat_entry_type, "Typing", k_EChatEntryTypeTyping);
  SET_TYPE(chat_entry_type, "InviteGame", k_EChatEntryTypeInviteGame);
  SET_TYPE(chat_entry_type, "Emote", k_EChatEntryTypeEmote);
  SET_TYPE(chat_entry_type, "LeftConversation",
           k_EChatEntryTypeLeftConversation);
  SET_TYPE(chat_entry_type, "Entered", k_EChatEntryTypeEntered);
  SET_TYPE(chat_entry_type, "WasKicked", k_EChatEntryTypeWasKicked);
  SET_TYPE(chat_entry_type, "WasBanned", k_EChatEntryTypeWasBanned);
  SET_TYPE(chat_entry_type, "Disconnected", k_EChatEntryTypeDisconnected);
  SET_TYPE(chat_entry_type, "HistoricalChat", k_EChatEntryTypeHistoricalChat);
  SET_TYPE(chat_entry_type, "LinkBlocked", k_EChatEntryTypeLinkBlocked);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(chat_entry_type);
  Nan::Set(exports,
           Nan::New("ChatEntryType").ToLocalChecked(),
           chat_entry_type);
}

NAN_METHOD(GetFriendCount) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsInt32()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  auto friend_flag = static_cast<EFriendFlags>(info[0]->Int32Value());

  info.GetReturnValue().Set(Nan::New<v8::Integer>(
    SteamFriends()->GetFriendCount(friend_flag)));
}

NAN_METHOD(GetFriends) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsInt32()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  auto friend_flag = static_cast<EFriendFlags>(info[0]->Int32Value());
  int friends_count = SteamFriends()->GetFriendCount(friend_flag);
  v8::Local<v8::Array> friends = Nan::New<v8::Array>(friends_count);

  for (int i = 0; i < friends_count; ++i) {
    CSteamID steam_id = SteamFriends()->GetFriendByIndex(i, friend_flag);
    friends->Set(i, greenworks::SteamID::Create(steam_id));
  }
  info.GetReturnValue().Set(friends);
}

NAN_METHOD(GetSmallFriendAvatar) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->GetSmallFriendAvatar(steam_id));
}

NAN_METHOD(GetMediumFriendAvatar) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->GetMediumFriendAvatar(steam_id));
}

NAN_METHOD(GetLargeFriendAvatar) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->GetLargeFriendAvatar(steam_id));
}

NAN_METHOD(RequestUserInformation) {
  Nan::HandleScope scope;
  if (info.Length() < 2 || !info[0]->IsString() || !info[1]->IsBoolean()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  bool require_name_only = info[1]->BooleanValue();
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  SteamFriends()->RequestUserInformation(steam_id, require_name_only);
}

NAN_METHOD(SetListenForFriendsMessages) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsBoolean()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  bool intercept_enabled = info[0]->BooleanValue();
  info.GetReturnValue().Set(
      SteamFriends()->SetListenForFriendsMessages(intercept_enabled));
}

NAN_METHOD(ReplyToFriendMessage) {
  Nan::HandleScope scope;
  if (info.Length() < 2 || !info[0]->IsString() || !info[1]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  std::string message_being_sent(*(v8::String::Utf8Value(info[1])));
  info.GetReturnValue().Set(SteamFriends()->ReplyToFriendMessage(
      steam_id,
      message_being_sent.c_str()));
}

NAN_METHOD(GetFriendMessage) {
  Nan::HandleScope scope;
  if (info.Length() < 3 || !info[0]->IsString() || !info[1]->IsInt32() ||
      !info[2]->IsInt32()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  int message_id = info[1]->Int32Value();
  int maximam_size = info[2]->Int32Value();

  EChatEntryType chat_type;
  std::unique_ptr<char[]>message(new char[maximam_size]);

  int message_size = SteamFriends()->GetFriendMessage(
      steam_id, message_id, message.get(),
      maximam_size, &chat_type);

  v8::Local<v8::Object> result = Nan::New<v8::Object>();
  result->Set(Nan::New("message").ToLocalChecked(),
              Nan::New(message.get(), message_size).ToLocalChecked());
  result->Set(Nan::New("chatEntryType").ToLocalChecked(),
              Nan::New(chat_type));
  info.GetReturnValue().Set(result);
}

void RegisterAPIs(v8::Handle<v8::Object> exports) {
  InitFriendFlags(exports);
  InitFriendRelationship(exports);
  InitFriendPersonaChange(exports);
  InitAccountType(exports);
  InitChatEntryType(exports);

  Nan::Set(exports,
           Nan::New("getFriendCount").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetFriendCount)->GetFunction());
  Nan::Set(exports,
           Nan::New("getFriends").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetFriends)->GetFunction());
  Nan::Set(exports, Nan::New("getSmallFriendAvatar").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetSmallFriendAvatar)->GetFunction());
  Nan::Set(
      exports, Nan::New("getMediumFriendAvatar").ToLocalChecked(),
      Nan::New<v8::FunctionTemplate>(GetMediumFriendAvatar)->GetFunction());
  Nan::Set(
      exports, Nan::New("getLargeFriendAvatar").ToLocalChecked(),
      Nan::New<v8::FunctionTemplate>(GetLargeFriendAvatar)->GetFunction());
  Nan::Set(
      exports, Nan::New("requestUserInformation").ToLocalChecked(),
      Nan::New<v8::FunctionTemplate>(RequestUserInformation)->GetFunction());
  Nan::Set(exports, Nan::New("setListenForFriendsMessage").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(SetListenForFriendsMessages)
               ->GetFunction());
  Nan::Set(exports, Nan::New("replyToFriendMessage").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(ReplyToFriendMessage)
               ->GetFunction());
  Nan::Set(exports, Nan::New("getFriendMessage").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetFriendMessage)
               ->GetFunction());
}

SteamAPIRegistry::Add X(RegisterAPIs);

}  // namespace
}  // namespace api
}  // namespace greenworks
<commit_msg>Fixed RequestUserInformation not returning anything.<commit_after>// Copyright (c) 2016 Greenheart Games Pty. Ltd. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include <memory>

#include "nan.h"
#include "steam/steam_api.h"
#include "v8.h"

#include "greenworks_utils.h"
#include "steam_api_registry.h"
#include "steam_id.h"

namespace greenworks {
namespace api {
namespace {

void InitFriendFlags(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> friend_flags = Nan::New<v8::Object>();
  SET_TYPE(friend_flags, "None", k_EFriendFlagNone);
  SET_TYPE(friend_flags, "Blocked", k_EFriendFlagBlocked);
  SET_TYPE(friend_flags, "FriendshipRequested",
           k_EFriendFlagFriendshipRequested);
  SET_TYPE(friend_flags, "Immediate", k_EFriendFlagImmediate);
  SET_TYPE(friend_flags, "ClanMember", k_EFriendFlagClanMember);
  SET_TYPE(friend_flags, "OnGameServer", k_EFriendFlagOnGameServer);
  SET_TYPE(friend_flags, "RequestingFriendship",
           k_EFriendFlagRequestingFriendship);
  SET_TYPE(friend_flags, "RequestingInfo", k_EFriendFlagRequestingInfo);
  SET_TYPE(friend_flags, "Ignored", k_EFriendFlagIgnored);
  SET_TYPE(friend_flags, "IgnoredFriend", k_EFriendFlagIgnoredFriend);
  // k_EFriendFlagSuggested has been removed since steamworks sdk 1.39.
  // SET_TYPE(friend_flags, "Suggested", k_EFriendFlagSuggested);
  SET_TYPE(friend_flags, "ChatMember", k_EFriendFlagChatMember);
  SET_TYPE(friend_flags, "All", k_EFriendFlagAll);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(friend_flags);
  Nan::Set(exports,
           Nan::New("FriendFlags").ToLocalChecked(),
           friend_flags);
}

void InitFriendRelationship(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> relationship = Nan::New<v8::Object>();
  SET_TYPE(relationship, "None", k_EFriendRelationshipNone);
  SET_TYPE(relationship, "Blocked", k_EFriendRelationshipBlocked);
  SET_TYPE(relationship, "RequestRecipient",
           k_EFriendRelationshipRequestRecipient);
  SET_TYPE(relationship, "Friend", k_EFriendRelationshipFriend);
  SET_TYPE(relationship, "RequestInitiator",
           k_EFriendRelationshipRequestInitiator);
  SET_TYPE(relationship, "Ignored", k_EFriendRelationshipIgnored);
  SET_TYPE(relationship, "IgnoredFriend", k_EFriendRelationshipIgnoredFriend);
  SET_TYPE(relationship, "Suggested",
           k_EFriendRelationshipSuggested_DEPRECATED);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(relationship);
  Nan::Set(exports,
           Nan::New("FriendRelationship").ToLocalChecked(),
           relationship);
}

void InitFriendPersonaChange(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> persona_change = Nan::New<v8::Object>();
  SET_TYPE(persona_change, "Name", k_EPersonaChangeName);
  SET_TYPE(persona_change, "Status", k_EPersonaChangeStatus);
  SET_TYPE(persona_change, "ComeOnline", k_EPersonaChangeComeOnline);
  SET_TYPE(persona_change, "GoneOffline", k_EPersonaChangeGoneOffline);
  SET_TYPE(persona_change, "GamePlayed", k_EPersonaChangeGamePlayed);
  SET_TYPE(persona_change, "GameServer", k_EPersonaChangeGameServer);
  SET_TYPE(persona_change, "Avator", k_EPersonaChangeAvatar);
  SET_TYPE(persona_change, "JoinedSource", k_EPersonaChangeJoinedSource);
  SET_TYPE(persona_change, "LeftSource", k_EPersonaChangeLeftSource);
  SET_TYPE(persona_change, "RelationshipChanged",
           k_EPersonaChangeRelationshipChanged);
  SET_TYPE(persona_change, "NameFirstSet", k_EPersonaChangeNameFirstSet);
  SET_TYPE(persona_change, "FacebookInfo", k_EPersonaChangeFacebookInfo);
  SET_TYPE(persona_change, "NickName", k_EPersonaChangeNickname);
  SET_TYPE(persona_change, "SteamLevel", k_EPersonaChangeSteamLevel);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(persona_change);
  Nan::Set(exports,
           Nan::New("PersonaChange").ToLocalChecked(),
           persona_change);
}

void InitAccountType(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> account_type = Nan::New<v8::Object>();
  SET_TYPE(account_type, "Invalid", k_EAccountTypeInvalid);
  SET_TYPE(account_type, "Individual", k_EAccountTypeIndividual);
  SET_TYPE(account_type, "Multiseat", k_EAccountTypeMultiseat);
  SET_TYPE(account_type, "GameServer", k_EAccountTypeGameServer);
  SET_TYPE(account_type, "AnonymousGameServer", k_EAccountTypeAnonGameServer);
  SET_TYPE(account_type, "Pending", k_EAccountTypePending);
  SET_TYPE(account_type, "ContentServer", k_EAccountTypeContentServer);
  SET_TYPE(account_type, "Clan", k_EAccountTypeClan);
  SET_TYPE(account_type, "Chat", k_EAccountTypeChat);
  SET_TYPE(account_type, "ConsoleUser", k_EAccountTypeConsoleUser);
  SET_TYPE(account_type, "AnonymousUser", k_EAccountTypeAnonUser);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(account_type);
  Nan::Set(exports,
           Nan::New("AccountType").ToLocalChecked(),
           account_type);
}

void InitChatEntryType(v8::Handle<v8::Object> exports) {
  v8::Local<v8::Object> chat_entry_type = Nan::New<v8::Object>();
  SET_TYPE(chat_entry_type, "Invalid", k_EChatEntryTypeInvalid);
  SET_TYPE(chat_entry_type, "ChatMsg", k_EChatEntryTypeChatMsg);
  SET_TYPE(chat_entry_type, "Typing", k_EChatEntryTypeTyping);
  SET_TYPE(chat_entry_type, "InviteGame", k_EChatEntryTypeInviteGame);
  SET_TYPE(chat_entry_type, "Emote", k_EChatEntryTypeEmote);
  SET_TYPE(chat_entry_type, "LeftConversation",
           k_EChatEntryTypeLeftConversation);
  SET_TYPE(chat_entry_type, "Entered", k_EChatEntryTypeEntered);
  SET_TYPE(chat_entry_type, "WasKicked", k_EChatEntryTypeWasKicked);
  SET_TYPE(chat_entry_type, "WasBanned", k_EChatEntryTypeWasBanned);
  SET_TYPE(chat_entry_type, "Disconnected", k_EChatEntryTypeDisconnected);
  SET_TYPE(chat_entry_type, "HistoricalChat", k_EChatEntryTypeHistoricalChat);
  SET_TYPE(chat_entry_type, "LinkBlocked", k_EChatEntryTypeLinkBlocked);
  Nan::Persistent<v8::Object> constructor;
  constructor.Reset(chat_entry_type);
  Nan::Set(exports,
           Nan::New("ChatEntryType").ToLocalChecked(),
           chat_entry_type);
}

NAN_METHOD(GetFriendCount) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsInt32()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  auto friend_flag = static_cast<EFriendFlags>(info[0]->Int32Value());

  info.GetReturnValue().Set(Nan::New<v8::Integer>(
    SteamFriends()->GetFriendCount(friend_flag)));
}

NAN_METHOD(GetFriends) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsInt32()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  auto friend_flag = static_cast<EFriendFlags>(info[0]->Int32Value());
  int friends_count = SteamFriends()->GetFriendCount(friend_flag);
  v8::Local<v8::Array> friends = Nan::New<v8::Array>(friends_count);

  for (int i = 0; i < friends_count; ++i) {
    CSteamID steam_id = SteamFriends()->GetFriendByIndex(i, friend_flag);
    friends->Set(i, greenworks::SteamID::Create(steam_id));
  }
  info.GetReturnValue().Set(friends);
}

NAN_METHOD(GetSmallFriendAvatar) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->GetSmallFriendAvatar(steam_id));
}

NAN_METHOD(GetMediumFriendAvatar) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->GetMediumFriendAvatar(steam_id));
}

NAN_METHOD(GetLargeFriendAvatar) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->GetLargeFriendAvatar(steam_id));
}

NAN_METHOD(RequestUserInformation) {
  Nan::HandleScope scope;
  if (info.Length() < 2 || !info[0]->IsString() || !info[1]->IsBoolean()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  bool require_name_only = info[1]->BooleanValue();
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  info.GetReturnValue().Set(
      SteamFriends()->RequestUserInformation(steam_id, require_name_only));
}

NAN_METHOD(SetListenForFriendsMessages) {
  Nan::HandleScope scope;
  if (info.Length() < 1 || !info[0]->IsBoolean()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  bool intercept_enabled = info[0]->BooleanValue();
  info.GetReturnValue().Set(
      SteamFriends()->SetListenForFriendsMessages(intercept_enabled));
}

NAN_METHOD(ReplyToFriendMessage) {
  Nan::HandleScope scope;
  if (info.Length() < 2 || !info[0]->IsString() || !info[1]->IsString()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  std::string message_being_sent(*(v8::String::Utf8Value(info[1])));
  info.GetReturnValue().Set(SteamFriends()->ReplyToFriendMessage(
      steam_id,
      message_being_sent.c_str()));
}

NAN_METHOD(GetFriendMessage) {
  Nan::HandleScope scope;
  if (info.Length() < 3 || !info[0]->IsString() || !info[1]->IsInt32() ||
      !info[2]->IsInt32()) {
    THROW_BAD_ARGS("Bad arguments");
  }
  std::string steam_id_str(*(v8::String::Utf8Value(info[0])));
  CSteamID steam_id(utils::strToUint64(steam_id_str));
  if (!steam_id.IsValid()) {
    THROW_BAD_ARGS("Steam ID is invalid");
  }
  int message_id = info[1]->Int32Value();
  int maximam_size = info[2]->Int32Value();

  EChatEntryType chat_type;
  std::unique_ptr<char[]>message(new char[maximam_size]);

  int message_size = SteamFriends()->GetFriendMessage(
      steam_id, message_id, message.get(),
      maximam_size, &chat_type);

  v8::Local<v8::Object> result = Nan::New<v8::Object>();
  result->Set(Nan::New("message").ToLocalChecked(),
              Nan::New(message.get(), message_size).ToLocalChecked());
  result->Set(Nan::New("chatEntryType").ToLocalChecked(),
              Nan::New(chat_type));
  info.GetReturnValue().Set(result);
}

void RegisterAPIs(v8::Handle<v8::Object> exports) {
  InitFriendFlags(exports);
  InitFriendRelationship(exports);
  InitFriendPersonaChange(exports);
  InitAccountType(exports);
  InitChatEntryType(exports);

  Nan::Set(exports,
           Nan::New("getFriendCount").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetFriendCount)->GetFunction());
  Nan::Set(exports,
           Nan::New("getFriends").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetFriends)->GetFunction());
  Nan::Set(exports, Nan::New("getSmallFriendAvatar").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetSmallFriendAvatar)->GetFunction());
  Nan::Set(
      exports, Nan::New("getMediumFriendAvatar").ToLocalChecked(),
      Nan::New<v8::FunctionTemplate>(GetMediumFriendAvatar)->GetFunction());
  Nan::Set(
      exports, Nan::New("getLargeFriendAvatar").ToLocalChecked(),
      Nan::New<v8::FunctionTemplate>(GetLargeFriendAvatar)->GetFunction());
  Nan::Set(
      exports, Nan::New("requestUserInformation").ToLocalChecked(),
      Nan::New<v8::FunctionTemplate>(RequestUserInformation)->GetFunction());
  Nan::Set(exports, Nan::New("setListenForFriendsMessage").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(SetListenForFriendsMessages)
               ->GetFunction());
  Nan::Set(exports, Nan::New("replyToFriendMessage").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(ReplyToFriendMessage)
               ->GetFunction());
  Nan::Set(exports, Nan::New("getFriendMessage").ToLocalChecked(),
           Nan::New<v8::FunctionTemplate>(GetFriendMessage)
               ->GetFunction());
}

SteamAPIRegistry::Add X(RegisterAPIs);

}  // namespace
}  // namespace api
}  // namespace greenworks
<|endoftext|>


<commit_before>/* Copyright 2017-present Facebook, Inc.
 * Licensed under the Apache License, Version 2.0 */

#include "watchman.h"
#include <folly/String.h>
#include <folly/memory/Malloc.h>

using namespace watchman;

#if (defined(USE_JEMALLOC) || defined(FOLLY_USE_JEMALLOC)) && !FOLLY_SANITIZE

/** This command is present to manually trigger a
 * heap profile dump when jemalloc is in use.
 * Since there is a complicated relationship with our build system,
 * it is only included in the folly enabled portions of watchman.
 */
static void cmd_debug_prof_dump(
    struct watchman_client* client,
    const json_ref&) {
  if (!folly::usingJEMalloc()) {
    throw std::runtime_error("jemalloc is not in use");
  }

  auto result = mallctl("prof.dump", nullptr, nullptr, nullptr, 0);
  auto resp = make_response();
  resp.set(
      "prof.dump",
      w_string_to_json(
          folly::to<std::string>(
              "mallctl prof.dump returned: ", folly::errnoStr(result))
              .c_str()));
  send_and_dispose_response(client, std::move(resp));
}
W_CMD_REG("debug-prof-dump", cmd_debug_prof_dump, CMD_DAEMON, NULL)
#endif
<commit_msg>replace COMPILER_FLAGS with a Watchman build_mode<commit_after>/* Copyright 2017-present Facebook, Inc.
 * Licensed under the Apache License, Version 2.0 */

#include "watchman.h"
#include <folly/String.h>
#include <folly/memory/Malloc.h>

using namespace watchman;

#if defined(FOLLY_USE_JEMALLOC) && !FOLLY_SANITIZE

/** This command is present to manually trigger a
 * heap profile dump when jemalloc is in use.
 * Since there is a complicated relationship with our build system,
 * it is only included in the folly enabled portions of watchman.
 */
static void cmd_debug_prof_dump(
    struct watchman_client* client,
    const json_ref&) {
  if (!folly::usingJEMalloc()) {
    throw std::runtime_error("jemalloc is not in use");
  }

  auto result = mallctl("prof.dump", nullptr, nullptr, nullptr, 0);
  auto resp = make_response();
  resp.set(
      "prof.dump",
      w_string_to_json(
          folly::to<std::string>(
              "mallctl prof.dump returned: ", folly::errnoStr(result))
              .c_str()));
  send_and_dispose_response(client, std::move(resp));
}
W_CMD_REG("debug-prof-dump", cmd_debug_prof_dump, CMD_DAEMON, NULL)
#endif
<|endoftext|>


<commit_before>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkFRPROptimizer.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef _itkFRPROptimizer_cxx
#define _itkFRPROptimizer_cxx

#include "itkFRPROptimizer.h"

namespace itk
{

const double FRPR_TINY = 1e-20;

FRPROptimizer
::FRPROptimizer()
{
  m_UseUnitLengthGradient = false;
  m_OptimizationType = PolakRibiere;
}

FRPROptimizer
::~FRPROptimizer()
{
}

void
FRPROptimizer
::GetValueAndDerivative(ParametersType & p, double * val,
                        ParametersType * xi)
{
  this->m_CostFunction->GetValueAndDerivative(p, *val, *xi);
  if(this->GetMaximize())
    {
    (*val) = -(*val);
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      (*xi)[i] = -(*xi)[i];
      }
    }
  if(this->GetUseUnitLengthGradient())
    {
    double len = (*xi)[0]*(*xi)[0];
    for(unsigned int i=1; i<this->GetSpaceDimension(); i++)
      {
      len += (*xi)[i]*(*xi)[i];
      }
    len = sqrt(len);
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      (*xi)[i] /= len;
      }
    }
}

void
FRPROptimizer
::LineOptimize(ParametersType * p, ParametersType & xi, double * val)
{
  ParametersType tempCoord( this->GetSpaceDimension() );
  this->LineOptimize( p, xi, val, tempCoord );
}

void
FRPROptimizer
::LineOptimize(ParametersType * p, ParametersType & xi, double * val,
               ParametersType & tempCoord )
{
  this->SetLine(*p, xi);

  double ax = 0.0;
  double fa = (*val);
  double xx = this->GetStepLength();
  double fx;
  double bx;
  double fb;

  this->LineBracket(&ax, &xx, &bx, &fa, &fx, &fb, tempCoord);
  this->SetCurrentLinePoint(xx, fx);

  double extX = 0;
  double extVal = 0;

  this->BracketedLineOptimize(ax, xx, bx, fa, fx, fb, &extX, &extVal,
                              tempCoord);
  this->SetCurrentLinePoint(extX, extVal);

  (*p) = this->GetCurrentPosition();
  (*val) = extVal;
}

void
FRPROptimizer
::StartOptimization()
{
  unsigned int i;

  if( m_CostFunction.IsNull() )
    {
    return;
    }

  this->InvokeEvent( StartEvent() );
  this->SetStop(false);

  this->SetSpaceDimension(m_CostFunction->GetNumberOfParameters());

  FRPROptimizer::ParametersType tempCoord( this->GetSpaceDimension() );

  double gg, gam, dgg;
  FRPROptimizer::ParametersType g( this->GetSpaceDimension() );
  FRPROptimizer::ParametersType h( this->GetSpaceDimension() );
  FRPROptimizer::ParametersType xi( this->GetSpaceDimension() );

  FRPROptimizer::ParametersType p( this->GetSpaceDimension() );
  p = this->GetInitialPosition();
  this->SetCurrentPosition(p);

  double fp;
  this->GetValueAndDerivative(p, &fp, &xi);

  for( i = 0; i < this->GetSpaceDimension(); i++ )
    {
    g[i] = -xi[i];
    xi[i] = g[i];
    h[i] = g[i];
    }

  unsigned int limitCount = 0;

  for (unsigned int currentIteration = 0;
       currentIteration <= this->GetMaximumIteration();
       currentIteration++)
    {
    this->SetCurrentIteration(currentIteration);

    double fret;
    fret = fp;
    this->LineOptimize(&p, xi, &fret, tempCoord);

    if ( 2.0 * vcl_abs(fret - fp) <=
      this->GetValueTolerance() * (vcl_abs(fret)+ vcl_abs(fp) + FRPR_TINY) )
      {
      if( limitCount <  this->GetSpaceDimension() )
        {
        this->GetValueAndDerivative(p, &fp, &xi);
        xi[limitCount] = 1;
        limitCount++;
        }
      else
        {
        this->SetCurrentPosition(p);
        this->InvokeEvent( EndEvent() );
        return;
        }
      }
    else
      {
      limitCount = 0;
      this->GetValueAndDerivative(p, &fp, &xi);
      }

    gg = 0.0;
    dgg = 0.0;

    if( m_OptimizationType == PolakRibiere )
      {
      for( i=0; i< this->GetSpaceDimension(); i++ )
        {
        gg += g[i] * g[i];
        dgg += (xi[i] + g[i]) * xi[i];
        }
      }
    if( m_OptimizationType == FletchReeves )
      {
      for( i=0; i< this->GetSpaceDimension(); i++ )
        {
        gg += g[i] * g[i];
        dgg += xi[i] * xi[i];
        }
      }

    if(gg == 0)
      {
      this->SetCurrentPosition(p);
      this->InvokeEvent( EndEvent() );
      return;
      }

    gam = dgg/gg;
    for( i = 0; i < this->GetSpaceDimension(); i++)
      {
      g[i] = -xi[i];
      xi[i] = g[i] + gam * h[i];
      h[i] = xi[i];
      }

    this->SetCurrentPosition(p);
    this->InvokeEvent( IterationEvent() );
    }

  this->InvokeEvent( EndEvent() );
}

/**
 *
 */
void
FRPROptimizer
::SetToPolakRibiere()
{
  m_OptimizationType = PolakRibiere;
}

/**
 *
 */
void
FRPROptimizer
::SetToFletchReeves()
{
  m_OptimizationType = FletchReeves;
}

/**
 *
 */
void
FRPROptimizer
::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os,indent);
  os << indent << "Optimization Type = " << m_OptimizationType << std::endl;
  os << indent << "0=FletchReeves, 1=PolakRibiere" << std::endl;
  os << indent << "Use unit length gradient = " << m_UseUnitLengthGradient << std::endl;
}

} // end of namespace itk
#endif
<commit_msg>ENH: Gradient can be a fixed-step gradient (mean magnitude of each component is 1)<commit_after>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkFRPROptimizer.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef _itkFRPROptimizer_cxx
#define _itkFRPROptimizer_cxx

#include "itkFRPROptimizer.h"

namespace itk
{

const double FRPR_TINY = 1e-20;

FRPROptimizer
::FRPROptimizer()
{
  m_UseUnitLengthGradient = false;
  m_OptimizationType = PolakRibiere;
}

FRPROptimizer
::~FRPROptimizer()
{
}

void
FRPROptimizer
::GetValueAndDerivative(ParametersType & p, double * val,
                        ParametersType * xi)
{
  this->m_CostFunction->GetValueAndDerivative(p, *val, *xi);
  if(this->GetMaximize())
    {
    (*val) = -(*val);
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      (*xi)[i] = -(*xi)[i];
      }
    }
  if(this->GetUseUnitLengthGradient())
    {
    double len = (*xi)[0]*(*xi)[0];
    for(unsigned int i=1; i<this->GetSpaceDimension(); i++)
      {
      len += (*xi)[i]*(*xi)[i];
      }
    len = sqrt(len/this->GetSpaceDimension());
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      (*xi)[i] /= len;
      }
    }
}

void
FRPROptimizer
::LineOptimize(ParametersType * p, ParametersType & xi, double * val)
{
  ParametersType tempCoord( this->GetSpaceDimension() );
  this->LineOptimize( p, xi, val, tempCoord );
}

void
FRPROptimizer
::LineOptimize(ParametersType * p, ParametersType & xi, double * val,
               ParametersType & tempCoord )
{
  this->SetLine(*p, xi);

  double ax = 0.0;
  double fa = (*val);
  double xx = this->GetStepLength();
  double fx;
  double bx;
  double fb;

  this->LineBracket(&ax, &xx, &bx, &fa, &fx, &fb, tempCoord);
  this->SetCurrentLinePoint(xx, fx);

  double extX = 0;
  double extVal = 0;

  this->BracketedLineOptimize(ax, xx, bx, fa, fx, fb, &extX, &extVal,
                              tempCoord);
  this->SetCurrentLinePoint(extX, extVal);

  (*p) = this->GetCurrentPosition();
  (*val) = extVal;
}

void
FRPROptimizer
::StartOptimization()
{
  unsigned int i;

  if( m_CostFunction.IsNull() )
    {
    return;
    }

  this->InvokeEvent( StartEvent() );
  this->SetStop(false);

  this->SetSpaceDimension(m_CostFunction->GetNumberOfParameters());

  FRPROptimizer::ParametersType tempCoord( this->GetSpaceDimension() );

  double gg, gam, dgg;
  FRPROptimizer::ParametersType g( this->GetSpaceDimension() );
  FRPROptimizer::ParametersType h( this->GetSpaceDimension() );
  FRPROptimizer::ParametersType xi( this->GetSpaceDimension() );

  FRPROptimizer::ParametersType p( this->GetSpaceDimension() );
  p = this->GetInitialPosition();
  this->SetCurrentPosition(p);

  double fp;
  this->GetValueAndDerivative(p, &fp, &xi);

  for( i = 0; i < this->GetSpaceDimension(); i++ )
    {
    g[i] = -xi[i];
    xi[i] = g[i];
    h[i] = g[i];
    }

  unsigned int limitCount = 0;

  for (unsigned int currentIteration = 0;
       currentIteration <= this->GetMaximumIteration();
       currentIteration++)
    {
    this->SetCurrentIteration(currentIteration);

    double fret;
    fret = fp;
    this->LineOptimize(&p, xi, &fret, tempCoord);

    if ( 2.0 * vcl_abs(fret - fp) <=
      this->GetValueTolerance() * (vcl_abs(fret)+ vcl_abs(fp) + FRPR_TINY) )
      {
      if( limitCount <  this->GetSpaceDimension() )
        {
        this->GetValueAndDerivative(p, &fp, &xi);
        xi[limitCount] = 1;
        limitCount++;
        }
      else
        {
        this->SetCurrentPosition(p);
        this->InvokeEvent( EndEvent() );
        return;
        }
      }
    else
      {
      limitCount = 0;
      this->GetValueAndDerivative(p, &fp, &xi);
      }

    gg = 0.0;
    dgg = 0.0;

    if( m_OptimizationType == PolakRibiere )
      {
      for( i=0; i< this->GetSpaceDimension(); i++ )
        {
        gg += g[i] * g[i];
        dgg += (xi[i] + g[i]) * xi[i];
        }
      }
    if( m_OptimizationType == FletchReeves )
      {
      for( i=0; i< this->GetSpaceDimension(); i++ )
        {
        gg += g[i] * g[i];
        dgg += xi[i] * xi[i];
        }
      }

    if(gg == 0)
      {
      this->SetCurrentPosition(p);
      this->InvokeEvent( EndEvent() );
      return;
      }

    gam = dgg/gg;
    for( i = 0; i < this->GetSpaceDimension(); i++)
      {
      g[i] = -xi[i];
      xi[i] = g[i] + gam * h[i];
      h[i] = xi[i];
      }

    this->SetCurrentPosition(p);
    this->InvokeEvent( IterationEvent() );
    }

  this->InvokeEvent( EndEvent() );
}

/**
 *
 */
void
FRPROptimizer
::SetToPolakRibiere()
{
  m_OptimizationType = PolakRibiere;
}

/**
 *
 */
void
FRPROptimizer
::SetToFletchReeves()
{
  m_OptimizationType = FletchReeves;
}

/**
 *
 */
void
FRPROptimizer
::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os,indent);
  os << indent << "Optimization Type = " << m_OptimizationType << std::endl;
  os << indent << "0=FletchReeves, 1=PolakRibiere" << std::endl;
  os << indent << "Use unit length gradient = " << m_UseUnitLengthGradient << std::endl;
}

} // end of namespace itk
#endif
<|endoftext|>


<commit_before>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkFRPROptimizer.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef _itkFRPROptimizer_cxx
#define _itkFRPROptimizer_cxx

#include "itkFRPROptimizer.h"

#define FRPR_TINY 1e-20

namespace itk
{

FRPROptimizer
::FRPROptimizer()
{
}

FRPROptimizer
::~FRPROptimizer()
{
}

void
FRPROptimizer
::GetValueAndDerivative(ParametersType p, double * val,
                        ParametersType * xi)
{
  this->m_CostFunction->GetValueAndDerivative(p, *val, *xi);
  if(this->GetMaximize())
    {
    (*val) = -(*val);
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      (*xi)[i] = -(*xi)[i];
      }
    }
}

void
FRPROptimizer
::LineOptimize(ParametersType * p, ParametersType xi, double * val)
{
  this->SetLine(*p, xi);

  double ax = 0.0;
  double fa = (*val);
  double xx = this->GetStepLength();
  double fx;
  double bx;
  double fb;

  ParametersType pp = (*p);

  this->LineBracket(&ax, &xx, &bx, &fa, &fx, &fb);
  this->SetCurrentLinePoint(xx, fx);

  double extX = 0;
  double extVal = 0;

  this->BracketedLineOptimize(ax, xx, bx, fa, fx, fb, &extX, &extVal);
  this->SetCurrentLinePoint(extX, extVal);

  (*p) = this->GetCurrentPosition();
  (*val) = extVal;
}

void
FRPROptimizer
::StartOptimization()
{
  unsigned int i;

  if( m_CostFunction.IsNull() )
    {
    return ;
    }

  this->InvokeEvent( StartEvent() );
  this->SetStop(false);

  this->SetSpaceDimension(m_CostFunction->GetNumberOfParameters());

  double gg, gam, dgg;
  FRPROptimizer::ParametersType g(this->GetSpaceDimension());
  FRPROptimizer::ParametersType h(this->GetSpaceDimension());
  FRPROptimizer::ParametersType xi(this->GetSpaceDimension());

  FRPROptimizer::ParametersType p(this->GetSpaceDimension());
  p = this->GetInitialPosition();

  double fp;
  this->GetValueAndDerivative(p, &fp, &xi);

  for(i=0; i<this->GetSpaceDimension(); i++)
    {
    g[i] = -xi[i];
    xi[i] = g[i];
    h[i] = g[i];
    }

  unsigned int limitCount = 0;

  for (unsigned int currentIteration = 0;
       currentIteration <= this->GetMaximumIteration();
       currentIteration++)
    {
    std::cout << currentIteration << " : " << p << " : " << fp << std::endl;
    this->SetCurrentIteration(currentIteration);

    double fret;
    fret = fp;
    this->LineOptimize(&p, xi, &fret);

    if(2.0 * vcl_abs(fret - fp) <=
       this->GetStepTolerance() * (vcl_abs(fret)*vcl_abs(fp) + FRPR_TINY))
      {
      if(limitCount < this->GetSpaceDimension())
        {
        this->GetValueAndDerivative(p, &fp, &xi);
        xi[limitCount] = 1;
        limitCount++;
        }
      else
        {
        this->SetCurrentPosition(p);
        this->InvokeEvent( EndEvent() );
        return;
        }
      }
    else
      {
      limitCount = 0;
      this->GetValueAndDerivative(p, &fp, &xi);
      }

    gg = 0.0;
    dgg = 0.0;
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      gg += g[i] * g[i];
      //dgg += xi[i] * xi[i];  // Uncomment for Fletcher-Reeves
      dgg += (xi[i] + g[i]) * xi[i];    // Uncoment for Polak-Ribiere
      }

    if(gg == 0)
      {
      this->SetCurrentPosition(p);
      this->InvokeEvent( EndEvent() );
      return;
      }

    gam = dgg/gg;
    for(i=0; i<this->GetSpaceDimension(); i++)
      {
      g[i] = -xi[i];
      xi[i] = g[i] + gam * h[i];
      h[i] = xi[i];
      }

    this->SetCurrentPosition(p);
    this->InvokeEvent( IterationEvent() );
    }

  this->SetCurrentPosition(p);
  this->InvokeEvent( EndEvent() );
}

/**
 *
 */
void
FRPROptimizer
::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os,indent);
}

} // end of namespace itk
#endif
<commit_msg>COMP: warnings.<commit_after>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkFRPROptimizer.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef _itkFRPROptimizer_cxx
#define _itkFRPROptimizer_cxx

#include "itkFRPROptimizer.h"

#define FRPR_TINY 1e-20

namespace itk
{

FRPROptimizer
::FRPROptimizer()
{
}

FRPROptimizer
::~FRPROptimizer()
{
}

void
FRPROptimizer
::GetValueAndDerivative(ParametersType p, double * val,
                        ParametersType * xi)
{
  this->m_CostFunction->GetValueAndDerivative(p, *val, *xi);
  if(this->GetMaximize())
    {
    (*val) = -(*val);
    for(unsigned int i=0; i<this->GetSpaceDimension(); i++)
      {
      (*xi)[i] = -(*xi)[i];
      }
    }
}

void
FRPROptimizer
::LineOptimize(ParametersType * p, ParametersType xi, double * val)
{
  this->SetLine(*p, xi);

  double ax = 0.0;
  double fa = (*val);
  double xx = this->GetStepLength();
  double fx;
  double bx;
  double fb;

  ParametersType pp = (*p);

  this->LineBracket(&ax, &xx, &bx, &fa, &fx, &fb);
  this->SetCurrentLinePoint(xx, fx);

  double extX = 0;
  double extVal = 0;

  this->BracketedLineOptimize(ax, xx, bx, fa, fx, fb, &extX, &extVal);
  this->SetCurrentLinePoint(extX, extVal);

  (*p) = this->GetCurrentPosition();
  (*val) = extVal;
}

void
FRPROptimizer
::StartOptimization()
{
  unsigned int i;

  if( m_CostFunction.IsNull() )
    {
    return ;
    }

  this->InvokeEvent( StartEvent() );
  this->SetStop(false);

  this->SetSpaceDimension(m_CostFunction->GetNumberOfParameters());

  double gg, gam, dgg;
  FRPROptimizer::ParametersType g(this->GetSpaceDimension());
  FRPROptimizer::ParametersType h(this->GetSpaceDimension());
  FRPROptimizer::ParametersType xi(this->GetSpaceDimension());

  FRPROptimizer::ParametersType p(this->GetSpaceDimension());
  p = this->GetInitialPosition();

  double fp;
  this->GetValueAndDerivative(p, &fp, &xi);

  for(i=0; i<this->GetSpaceDimension(); i++)
    {
    g[i] = -xi[i];
    xi[i] = g[i];
    h[i] = g[i];
    }

  unsigned int limitCount = 0;

  for (unsigned int currentIteration = 0;
       currentIteration <= this->GetMaximumIteration();
       currentIteration++)
    {
    std::cout << currentIteration << " : " << p << " : " << fp << std::endl;
    this->SetCurrentIteration(currentIteration);

    double fret;
    fret = fp;
    this->LineOptimize(&p, xi, &fret);

    if(2.0 * vcl_abs(fret - fp) <=
       this->GetStepTolerance() * (vcl_abs(fret)*vcl_abs(fp) + FRPR_TINY))
      {
      if(limitCount < this->GetSpaceDimension())
        {
        this->GetValueAndDerivative(p, &fp, &xi);
        xi[limitCount] = 1;
        limitCount++;
        }
      else
        {
        this->SetCurrentPosition(p);
        this->InvokeEvent( EndEvent() );
        return;
        }
      }
    else
      {
      limitCount = 0;
      this->GetValueAndDerivative(p, &fp, &xi);
      }

    gg = 0.0;
    dgg = 0.0;
    for(i=0; i<this->GetSpaceDimension(); i++)
      {
      gg += g[i] * g[i];
      //dgg += xi[i] * xi[i];  // Uncomment for Fletcher-Reeves
      dgg += (xi[i] + g[i]) * xi[i];    // Uncoment for Polak-Ribiere
      }

    if(gg == 0)
      {
      this->SetCurrentPosition(p);
      this->InvokeEvent( EndEvent() );
      return;
      }

    gam = dgg/gg;
    for(i=0; i<this->GetSpaceDimension(); i++)
      {
      g[i] = -xi[i];
      xi[i] = g[i] + gam * h[i];
      h[i] = xi[i];
      }

    this->SetCurrentPosition(p);
    this->InvokeEvent( IterationEvent() );
    }

  this->SetCurrentPosition(p);
  this->InvokeEvent( EndEvent() );
}

/**
 *
 */
void
FRPROptimizer
::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os,indent);
}

} // end of namespace itk
#endif
<|endoftext|>


<commit_before>Algoritmo de búsqueda en profundidad para grafos. Parte de un nodo inicial s visita a todos sus vecinos. DFS puede ser usado para contar la cantidad de componentes conexas en un grafo y puede ser modificado para que retorne información de los nodos dependiendo del problema. Permite hallar ciclos en un grafo.
SE DEBEN LIMPIAR LAS ESTRUCTURAS DE DATOS ANTES DE UTILIZARSE

int v, e; //vertices, arcos
const int MAX=100005; //Cantidad máxima de nodos del grafo
vector<int> ady[MAX]; //lista de Adyacencia del grafo
bool marked[MAX]; //Estructura auxiliar para marcar los nodos ya visitados

/*Este método se llama con el indice del nodo desde el que se desea comenzar
el recorrido.*/
static void dfs(int s){
	marked[s] = 1;
	int i, next;

	for( i = 0; i < ady[s].size(); i++ ){
		next = ady[s][i];
		if( !marked[next] )	dfs(next);
	}
}
<commit_msg>Update DFS.cpp<commit_after>Algoritmo de búsqueda en profundidad para grafos. Parte de un nodo inicial s visita a todos sus vecinos. DFS puede ser usado para contar la cantidad de componentes conexas en un grafo y puede ser modificado para que retorne información de los nodos dependiendo del problema. Permite hallar ciclos en un grafo.
SE DEBEN LIMPIAR LAS ESTRUCTURAS DE DATOS ANTES DE UTILIZARSE

#include <bits/stdc++.h>
using namespace std;

int v, e; //vertices, arcos
const int MAX=100005; //Cantidad máxima de nodos del grafo
vector<int> ady[MAX]; //lista de Adyacencia del grafo
bool marked[MAX]; //Estructura auxiliar para marcar los nodos ya visitados

void init(){
	for( int j = 0; j <= v; j++ ) {
        marked[j] = 0;
	}
}

/*Este método se llama con el indice del nodo desde el que se desea comenzar
el recorrido.*/
static void dfs(int s){
	marked[s] = 1;
	int i, next;

	for( i = 0; i < ady[s].size(); i++ ){
		next = ady[s][i];
		if( !marked[next] )	dfs(next);
	}
}

int main(){
	init();
	dfs(0);
	return 0;
}
<|endoftext|>


<commit_before>//
//  readsToPWP.cpp
//  cPWP
//
//  Created by Evan McCartney-Melstad on 12/31/14.
//  Copyright (c) 2014 Evan McCartney-Melstad. All rights reserved.
//

#include "bamsToPWP.h"
#include <iostream>
/*
#include <stdlib.h>
#include <stdio.h>

#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <algorithm>
#include <thread>
 */
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/algorithm/string.hpp>

#include <cstdlib>
#include <cmath>

int runANGSDforReadCounts (std::string bamlist, std::string angsdPrefix, std::string nThreads, std::string angsdOutputLog) {
    std::cout << "**********\nChecking if processor is available to run ANGSD...**********\n";
    if (system(NULL)) puts ("Ok");
    else exit (EXIT_FAILURE);

    //std::string angsdCommand = "/home/evan/bin/angsd0.613/angsd -bam " + bamlist + " -out " + angsdPrefix + " -uniqueOnly 1 -only_proper_pairs 1 -remove_bads 1 -doCounts 1 -nThreads " + nThreads + " -dumpCounts 4 -doMaf 1 -doMajorMinor 2 -GL 2 -minMapQ 20 -minQ 30 -SNP_pval 1e-6 > " + angsdOutputLog + " 2>&1";

    std::string angsdCommand = "/home/evan/bin/angsd0.613/angsd -bam " + bamlist + " -out " + angsdPrefix + " -uniqueOnly 1 -only_proper_pairs 1 -remove_bads 1 -doCounts 1 -nThreads " + nThreads + " -dumpCounts 4 -doMaf 1 -doMajorMinor 2 -GL 2 -minMapQ 20 -minQ 30 > " + angsdOutputLog + " 2>&1";

    if (system((angsdCommand).c_str()) != 0) {
        std::cout << "**********\nFailure running the following command: " << angsdCommand << "\n**********\n";
        exit(EXIT_FAILURE);
    } else {
        std::cout << "**********\nExecuted the following command: " << angsdCommand << "\n**********\n";
    }

    return 0;
}

unsigned char readCountToUnsignedChar(int individual, std::string allele, std::vector<std::string> countsVector) {
    /*    std::cout << "countsVector elements: ";
     for (int i=0; i < countsVector.size(); i++) {
     std::cout <<  countsVector.at(i) << " ";
     }
     std::cout << std::endl;
     */

    //    std::cout << "FunctionPassedAllele: " << allele << std::endl;

    unsigned char countChar;
    if (allele == "A") {
        //        std::cout << "Allele is an A" << std::endl;
        int countField = individual*4;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    }
    else if (allele == "C") {
        //	std::cout << "Allele is an C" << std::endl;
        int countField = (individual*4)+1;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    } else if (allele == "G") {
        //	std::cout << "Allele is a G" << std::endl;
        int countField = (individual*4) + 2;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    } else if (allele == "T") {
        //	std::cout << "Allele is a T" << std::endl;
        int countField = (individual*4) + 3;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    } else {
        std::cout << "Allele doesn't look right. The allele fed into readCountToUnsignedChar is: " << allele << ". Exiting now" << std::endl;
        exit(EXIT_FAILURE);
    }
    //    std::cout << std::to_string(countChar) << std::endl;
    return countChar;
}

int convertANGSDcountsToBinary(std::string angsdPrefix, std::string binaryOutputFileName, int numIndividuals, int readDepthMax)
{
    //std::string countsFileName = angsdPrefix + ".counts.gz";
    std::string countsFileName = angsdPrefix + ".counts";
    //std::string mafsFileName = angsdPrefix + ".mafs.gz";
    std::string mafsFileName = angsdPrefix + ".mafs";
    //std::ifstream countsFile(countsFileName, std::ios_base::in | std::ios_base::binary);
    std::ifstream countsFile(countsFileName, std::ios_base::in);
    //std::ifstream mafsFile(mafsFileName, std::ios_base::in | std::ios_base::binary);
    std::ifstream mafsFile(mafsFileName, std::ios_base::in);
    std::ofstream outCountsFile(binaryOutputFileName, std::ios_base::out | std::ios_base::binary);

    try {
        /*
        boost::iostreams::filtering_istream in;
        boost::iostreams::filtering_istream inMafs;
        in.push(boost::iostreams::gzip_decompressor());
        in.push(countsFile);
        inMafs.push(boost::iostreams::gzip_decompressor());
        inMafs.push(mafsFile);
        */
        int counter = 0;
        //int numIndividuals = 3;
        for(std::string str; std::getline(countsFile, str); )
        {
            /* Make sure to get the maf string before we do anything else to keep it in
             sync with the counts string */
            std::string mafStr;
            std::getline(mafsFile, mafStr);

            // Skip the first line of the counts and mafs file, since it is a header
            if (counter == 0) {
                counter++;
                continue;
            }

            // Just for monitoring progress
            if (counter % 1000000 == 0) {
                std::cout << "Processed " << counter << " loci" << std::endl;
            }

            /* The first thing we'll do after getting the lines from the gzipped file is to split
             the tab-separated lines into separate vectors for the counts and mafs. The vector will
             be recreated for every line of the gzipped counts and mafs files output by angsd (each
             line represents a locus) */
            std::vector<std::string> countsFields;
            boost::algorithm::split(countsFields, str, boost::is_any_of("\t"), boost::token_compress_on);

            std::vector<std::string> mafFields;
            boost::algorithm::split(mafFields, mafStr, boost::is_any_of("\t"), boost::token_compress_on);

            /* Since we're using a weighted PWP calculation, and since we don't want to overweight
             repetitive regions, we'll throw out any locus that has more than readDepthMax reads
             for any individual at that locus */
            std::vector<int> countsFieldsInts;
            for (size_t i=0; i<countsFields.size(); i++) {
                int countInt = atoi((countsFields[i]).c_str());
                countsFieldsInts.push_back(countInt);
            }

            auto countMax = std::minmax_element(countsFieldsInts.begin(),countsFieldsInts.end());
            if (*countMax.second > readDepthMax) {
                continue; // Skip to the next locus
            }

            // Now that the lines are parsed, pull the major and minor alleles from angsd's mafs output
            std::string majorAllele = mafFields[2];
            std::string minorAllele = mafFields[3];

            /* Next, we'll count find how many reads there are supporting the major and minor
             alleles for each individual by looking through the corresponding counts file. The
             function that does this, "readCountToUnsignedChar" is implemented above. It returns
             the readcount number as an unsigned char for binary printing. The max value of an
             unsigned char is 255, but we are skipping all loci with individual read depths
             above readDepthMax (which will usually be much less than 255)*/
            for (int individual = 0; individual < numIndividuals; individual++) {
                unsigned char majorAlleleUnsignedChar = readCountToUnsignedChar(individual, majorAllele, countsFields);
                unsigned char minorAlleleUnsignedChar = readCountToUnsignedChar(individual, minorAllele, countsFields);
                outCountsFile << majorAlleleUnsignedChar << minorAlleleUnsignedChar;
            }
            counter++; // Keep track of how many loci have been processed
        }
    }
    catch(const boost::iostreams::gzip_error& e) {
        std::cout << e.what() << '\n';
    }
    outCountsFile.close();

    return 0;
}

<commit_msg>Minor changes<commit_after>//
//  readsToPWP.cpp
//  cPWP
//
//  Created by Evan McCartney-Melstad on 12/31/14.
//  Copyright (c) 2014 Evan McCartney-Melstad. All rights reserved.
//

#include "bamsToPWP.h"
#include <iostream>
#include <fstream>
#include <sstream>
/*
#include <stdlib.h>
#include <stdio.h>

#include <string>
#include <vector>
#include <algorithm>
#include <thread>
 */
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/algorithm/string.hpp>

#include <cstdlib>
#include <cmath>

int runANGSDforReadCounts (std::string bamlist, std::string angsdPrefix, std::string nThreads, std::string angsdOutputLog) {
    std::cout << "**********\nChecking if processor is available to run ANGSD...**********\n";
    if (system(NULL)) puts ("Ok");
    else exit (EXIT_FAILURE);

    //std::string angsdCommand = "/home/evan/bin/angsd0.613/angsd -bam " + bamlist + " -out " + angsdPrefix + " -uniqueOnly 1 -only_proper_pairs 1 -remove_bads 1 -doCounts 1 -nThreads " + nThreads + " -dumpCounts 4 -doMaf 1 -doMajorMinor 2 -GL 2 -minMapQ 20 -minQ 30 -SNP_pval 1e-6 > " + angsdOutputLog + " 2>&1";

    std::string angsdCommand = "/home/evan/bin/angsd0.613/angsd -bam " + bamlist + " -out " + angsdPrefix + " -uniqueOnly 1 -only_proper_pairs 1 -remove_bads 1 -doCounts 1 -nThreads " + nThreads + " -dumpCounts 4 -doMaf 1 -doMajorMinor 2 -GL 2 -minMapQ 20 -minQ 30 > " + angsdOutputLog + " 2>&1";

    if (system((angsdCommand).c_str()) != 0) {
        std::cout << "**********\nFailure running the following command: " << angsdCommand << "\n**********\n";
        exit(EXIT_FAILURE);
    } else {
        std::cout << "**********\nExecuted the following command: " << angsdCommand << "\n**********\n";
    }

    return 0;
}

unsigned char readCountToUnsignedChar(int individual, std::string allele, std::vector<std::string> countsVector) {
    /*    std::cout << "countsVector elements: ";
     for (int i=0; i < countsVector.size(); i++) {
     std::cout <<  countsVector.at(i) << " ";
     }
     std::cout << std::endl;
     */

    //    std::cout << "FunctionPassedAllele: " << allele << std::endl;

    unsigned char countChar;
    if (allele == "A") {
        //        std::cout << "Allele is an A" << std::endl;
        int countField = individual*4;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    }
    else if (allele == "C") {
        //	std::cout << "Allele is an C" << std::endl;
        int countField = (individual*4)+1;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    } else if (allele == "G") {
        //	std::cout << "Allele is a G" << std::endl;
        int countField = (individual*4) + 2;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    } else if (allele == "T") {
        //	std::cout << "Allele is a T" << std::endl;
        int countField = (individual*4) + 3;
        int countInt = atoi(countsVector[countField].c_str());
        countChar = (unsigned char) countInt;
    } else {
        std::cout << "Allele doesn't look right. The allele fed into readCountToUnsignedChar is: " << allele << ". Exiting now" << std::endl;
        exit(EXIT_FAILURE);
    }
    //    std::cout << std::to_string(countChar) << std::endl;
    return countChar;
}

int convertANGSDcountsToBinary(std::string angsdPrefix, std::string binaryOutputFileName, int numIndividuals, int readDepthMax)
{
    //std::string countsFileName = angsdPrefix + ".counts.gz";
    std::string countsFileName = angsdPrefix + ".counts";
    //std::string mafsFileName = angsdPrefix + ".mafs.gz";
    std::string mafsFileName = angsdPrefix + ".mafs";
    //std::ifstream countsFile(countsFileName, std::ios_base::in | std::ios_base::binary);
    std::ifstream countsFile(countsFileName, std::ios_base::in);
    //std::ifstream mafsFile(mafsFileName, std::ios_base::in | std::ios_base::binary);
    std::ifstream mafsFile(mafsFileName, std::ios_base::in);
    std::ofstream outCountsFile(binaryOutputFileName, std::ios_base::out | std::ios_base::binary);

    try {
        /*
        boost::iostreams::filtering_istream in;
        boost::iostreams::filtering_istream inMafs;
        in.push(boost::iostreams::gzip_decompressor());
        in.push(countsFile);
        inMafs.push(boost::iostreams::gzip_decompressor());
        inMafs.push(mafsFile);
        */
        int counter = 0;
        //int numIndividuals = 3;
        for(std::string str; std::getline(countsFile, str); )
        {
            /* Make sure to get the maf string before we do anything else to keep it in
             sync with the counts string */
            std::string mafStr;
            std::getline(mafsFile, mafStr);

            // Skip the first line of the counts and mafs file, since it is a header
            if (counter == 0) {
                counter++;
                continue;
            }

            // Just for monitoring progress
            if (counter % 1000000 == 0) {
                std::cout << "Processed " << counter << " loci" << std::endl;
            }

            /* The first thing we'll do after getting the lines from the gzipped file is to split
             the tab-separated lines into separate vectors for the counts and mafs. The vector will
             be recreated for every line of the gzipped counts and mafs files output by angsd (each
             line represents a locus) */
            std::vector<std::string> countsFields;
            boost::algorithm::split(countsFields, str, boost::is_any_of("\t"), boost::token_compress_on);

            std::vector<std::string> mafFields;
            boost::algorithm::split(mafFields, mafStr, boost::is_any_of("\t"), boost::token_compress_on);

            /* Since we're using a weighted PWP calculation, and since we don't want to overweight
             repetitive regions, we'll throw out any locus that has more than readDepthMax reads
             for any individual at that locus */
            std::vector<int> countsFieldsInts;
            for (size_t i=0; i<countsFields.size(); i++) {
                int countInt = atoi((countsFields[i]).c_str());
                countsFieldsInts.push_back(countInt);
            }

            auto countMax = std::minmax_element(countsFieldsInts.begin(),countsFieldsInts.end());
            if (*countMax.second > readDepthMax) {
                continue; // Skip to the next locus
            }

            // Now that the lines are parsed, pull the major and minor alleles from angsd's mafs output
            std::string majorAllele = mafFields[2];
            std::string minorAllele = mafFields[3];

            /* Next, we'll count find how many reads there are supporting the major and minor
             alleles for each individual by looking through the corresponding counts file. The
             function that does this, "readCountToUnsignedChar" is implemented above. It returns
             the readcount number as an unsigned char for binary printing. The max value of an
             unsigned char is 255, but we are skipping all loci with individual read depths
             above readDepthMax (which will usually be much less than 255)*/
            for (int individual = 0; individual < numIndividuals; individual++) {
                unsigned char majorAlleleUnsignedChar = readCountToUnsignedChar(individual, majorAllele, countsFields);
                unsigned char minorAlleleUnsignedChar = readCountToUnsignedChar(individual, minorAllele, countsFields);
                outCountsFile << majorAlleleUnsignedChar << minorAlleleUnsignedChar;
            }
            counter++; // Keep track of how many loci have been processed
        }
    }
    catch(const boost::iostreams::gzip_error& e) {
        std::cout << e.what() << '\n';
    }
    outCountsFile.close();

    return 0;
}

<|endoftext|>