<commit_before>// Copyright (c) 2018 Kenton Varda and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#include "json-rpc.h"
#include <kj/compat/http.h>
#include <capnp/compat/json-rpc.capnp.h>

namespace capnp {

static constexpr uint64_t JSON_NAME_ANNOTATION_ID = 0xfa5b1fd61c2e7c3dull;
static constexpr uint64_t JSON_NOTIFICATION_ANNOTATION_ID = 0xa0a054dea32fd98cull;

class JsonRpc::CapabilityImpl final: public DynamicCapability::Server {
public:
  CapabilityImpl(JsonRpc& parent, InterfaceSchema schema)
      : DynamicCapability::Server(schema), parent(parent) {}

  kj::Promise<void> call(InterfaceSchema::Method method,
                         CallContext<DynamicStruct, DynamicStruct> context) override {
    auto proto = method.getProto();
    bool isNotification = false;
    kj::StringPtr name = proto.getName();
    for (auto annotation: proto.getAnnotations()) {
      switch (annotation.getId()) {
        case JSON_NAME_ANNOTATION_ID:
          name = annotation.getValue().getText();
          break;
        case JSON_NOTIFICATION_ANNOTATION_ID:
          isNotification = true;
          break;
      }
    }

    capnp::MallocMessageBuilder message;
    auto value = message.getRoot<json::Value>();
    auto list = value.initObject(3 + !isNotification);

    uint index = 0;

    auto jsonrpc = list[index++];
    jsonrpc.setName("jsonrpc");
    jsonrpc.initValue().setString("2.0");

    uint callId = parent.callCount++;

    if (!isNotification) {
      auto id = list[index++];
      id.setName("id");
      id.initValue().setNumber(callId);
    }

    auto methodName = list[index++];
    methodName.setName("method");
    methodName.initValue().setString(name);

    auto params = list[index++];
    params.setName("params");
    parent.codec.encode(context.getParams(), params.initValue());

    auto writePromise = parent.queueWrite(parent.codec.encode(value));

    if (isNotification) {
      auto sproto = context.getResultsType().getProto().getStruct();
      MessageSize size { sproto.getDataWordCount(), sproto.getPointerCount() };
      context.initResults(size);
      return kj::mv(writePromise);
    } else {
      auto paf = kj::newPromiseAndFulfiller<void>();
      parent.awaitedResponses.insert(callId, AwaitedResponse { context, kj::mv(paf.fulfiller) });
      auto promise = writePromise.then([p = kj::mv(paf.promise)]() mutable { return kj::mv(p); });
      auto& parentRef = parent;
      return promise.attach(kj::defer([&parentRef,callId]() {
        parentRef.awaitedResponses.erase(callId);
      }));
    }
  }

private:
  JsonRpc& parent;
};

JsonRpc::JsonRpc(Transport& transport, DynamicCapability::Client interface)
    : JsonRpc(transport, kj::mv(interface), kj::newPromiseAndFulfiller<void>()) {}
JsonRpc::JsonRpc(Transport& transport, DynamicCapability::Client interfaceParam,
                 kj::PromiseFulfillerPair<void> paf)
    : transport(transport),
      interface(kj::mv(interfaceParam)),
      errorPromise(paf.promise.fork()),
      errorFulfiller(kj::mv(paf.fulfiller)),
      readTask(readLoop().eagerlyEvaluate([this](kj::Exception&& e) {
        errorFulfiller->reject(kj::mv(e));
      })),
      tasks(*this) {
  codec.handleByAnnotation(interface.getSchema());
  codec.handleByAnnotation<json::RpcMessage>();

  for (auto method: interface.getSchema().getMethods()) {
    auto proto = method.getProto();
    kj::StringPtr name = proto.getName();
    for (auto annotation: proto.getAnnotations()) {
      switch (annotation.getId()) {
        case JSON_NAME_ANNOTATION_ID:
          name = annotation.getValue().getText();
          break;
      }
    }
    methodMap.insert(name, method);
  }
}

DynamicCapability::Client JsonRpc::getPeer(InterfaceSchema schema) {
  codec.handleByAnnotation(interface.getSchema());
  return kj::heap<CapabilityImpl>(*this, schema);
}

static kj::HttpHeaderTable& staticHeaderTable() {
  static kj::HttpHeaderTable HEADER_TABLE;
  return HEADER_TABLE;
}

kj::Promise<void> JsonRpc::queueWrite(kj::String text) {
  auto fork = writeQueue.then([this, text = kj::mv(text)]() mutable {
    auto promise = transport.send(text);
    return promise.attach(kj::mv(text));
  }).eagerlyEvaluate([this](kj::Exception&& e) {
    errorFulfiller->reject(kj::mv(e));
  }).fork();
  writeQueue = fork.addBranch();
  return fork.addBranch();
}

void JsonRpc::queueError(kj::Maybe<json::Value::Reader> id, int code, kj::StringPtr message) {
  MallocMessageBuilder capnpMessage;
  auto jsonResponse = capnpMessage.getRoot<json::RpcMessage>();
  jsonResponse.setJsonrpc("2.0");
  KJ_IF_MAYBE(i, id) {
    jsonResponse.setId(*i);
  } else {
    jsonResponse.initId().setNull();
  }
  auto error = jsonResponse.initError();
  error.setCode(code);
  error.setMessage(message);

  // OK to discard result of queueWrite() since it's just one branch of a fork.
  queueWrite(codec.encode(jsonResponse));
}

kj::Promise<void> JsonRpc::readLoop() {
  return transport.receive().then([this](kj::String message) -> kj::Promise<void> {
    MallocMessageBuilder capnpMessage;
    auto rpcMessageBuilder = capnpMessage.getRoot<json::RpcMessage>();

    KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
      codec.decode(message, rpcMessageBuilder);
    })) {
      queueError(nullptr, -32700, kj::str("Parse error: ", exception->getDescription()));
      return readLoop();
    }

    KJ_CONTEXT("decoding JSON-RPC message", message);

    auto rpcMessage = rpcMessageBuilder.asReader();

    if (!rpcMessage.hasJsonrpc()) {
      queueError(nullptr, -32700, kj::str("Missing 'jsonrpc' field."));
      return readLoop();
    } else if (rpcMessage.getJsonrpc() != "2.0") {
      queueError(nullptr, -32700,
          kj::str("Unknown JSON-RPC version. This peer implements version '2.0'."));
      return readLoop();
    }

    switch (rpcMessage.which()) {
      case json::RpcMessage::NONE:
        queueError(nullptr, -32700, kj::str("message has none of params, result, or error"));
        break;

      case json::RpcMessage::PARAMS: {
        // a call
        auto schema = interface.getSchema();
        KJ_IF_MAYBE(method, schema.findMethodByName(rpcMessage.getMethod())) {
          auto req = interface.newRequest(*method);
          KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
            codec.decode(rpcMessage.getParams(), req);
          })) {
            kj::Maybe<JsonValue::Reader> id;
            if (rpcMessage.hasId()) id = rpcMessage.getId();
            queueError(id, -32602,
                kj::str("Type error in method params: ", exception->getDescription()));
            break;
          }

          if (rpcMessage.hasId()) {
            auto id = rpcMessage.getId();
            auto idCopy = kj::heapArray<word>(id.totalSize().wordCount + 1);
            memset(idCopy.begin(), 0, idCopy.asBytes().size());
            copyToUnchecked(id, idCopy);
            auto idPtr = readMessageUnchecked<json::Value>(idCopy.begin());

            auto promise = req.send()
                .then([this,idPtr](Response<DynamicStruct> response) mutable {
              MallocMessageBuilder capnpMessage;
              auto jsonResponse = capnpMessage.getRoot<json::RpcMessage>();
              jsonResponse.setJsonrpc("2.0");
              jsonResponse.setId(idPtr);
              codec.encode(DynamicStruct::Reader(response), jsonResponse.initResult());
              return queueWrite(codec.encode(jsonResponse));
            }, [this,idPtr](kj::Exception&& e) {
              MallocMessageBuilder capnpMessage;
              auto jsonResponse = capnpMessage.getRoot<json::RpcMessage>();
              jsonResponse.setJsonrpc("2.0");
              jsonResponse.setId(idPtr);
              auto error = jsonResponse.initError();
              switch (e.getType()) {
                case kj::Exception::Type::FAILED:
                  error.setCode(-32000);
                  break;
                case kj::Exception::Type::DISCONNECTED:
                  error.setCode(-32001);
                  break;
                case kj::Exception::Type::OVERLOADED:
                  error.setCode(-32002);
                  break;
                case kj::Exception::Type::UNIMPLEMENTED:
                  error.setCode(-32601);  // method not found
                  break;
              }
              error.setMessage(e.getDescription());
              return queueWrite(codec.encode(jsonResponse));
            });
            tasks.add(promise.attach(kj::mv(idCopy)));
          } else {
            // No 'id', so this is a notification.
            tasks.add(req.send().ignoreResult().catch_([](kj::Exception&& exception) {
              if (exception.getType() != kj::Exception::Type::UNIMPLEMENTED) {
                KJ_LOG(ERROR, "JSON-RPC notification threw exception into the abyss", exception);
              }
            }));
          }
        } else {
          if (rpcMessage.hasId()) {
            queueError(rpcMessage.getId(), -32601, "Method not found");
          } else {
            // Ignore notification for unknown method.
          }
        }
        break;
      }

      case json::RpcMessage::RESULT: {
        auto id = rpcMessage.getId();
        if (!id.isNumber()) {
          // JSON-RPC doesn't define what to do if receiving a response with an invalid id.
          KJ_LOG(ERROR, "JSON-RPC response has invalid ID");
        } else KJ_IF_MAYBE(awaited, awaitedResponses.find((uint)id.getNumber())) {
          KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
            codec.decode(rpcMessage.getResult(), awaited->context.getResults());
            awaited->fulfiller->fulfill();
          })) {
            // Errors always propagate from callee to caller, so we don't want to throw this error
            // back to the server.
            awaited->fulfiller->reject(kj::mv(*exception));
          }
        } else {
          // Probably, this is the response to a call that was canceled.
        }
        break;
      }

      case json::RpcMessage::ERROR: {
        auto id = rpcMessage.getId();
        if (id.isNull()) {
          // Error message will be logged by KJ_CONTEXT, above.
          KJ_LOG(ERROR, "peer reports JSON-RPC protocol error");
        } else if (!id.isNumber()) {
          // JSON-RPC doesn't define what to do if receiving a response with an invalid id.
          KJ_LOG(ERROR, "JSON-RPC response has invalid ID");
        } else KJ_IF_MAYBE(awaited, awaitedResponses.find((uint)id.getNumber())) {
          auto error = rpcMessage.getError();
          auto code = error.getCode();
          kj::Exception::Type type =
              code == -32601 ? kj::Exception::Type::UNIMPLEMENTED
                             : kj::Exception::Type::FAILED;
          awaited->fulfiller->reject(kj::Exception(
              type, __FILE__, __LINE__, kj::str(error.getMessage())));
        } else {
          // Probably, this is the response to a call that was canceled.
        }
        break;
      }
    }

    return readLoop();
  });
}

void JsonRpc::taskFailed(kj::Exception&& exception) {
  errorFulfiller->reject(kj::mv(exception));
}

// =======================================================================================

JsonRpc::ContentLengthTransport::ContentLengthTransport(kj::AsyncIoStream& stream)
    : stream(stream), input(kj::newHttpInputStream(stream, staticHeaderTable())) {}
JsonRpc::ContentLengthTransport::~ContentLengthTransport() noexcept(false) {}

kj::Promise<void> JsonRpc::ContentLengthTransport::send(kj::StringPtr text) {
  auto headers = kj::str("Content-Length: ", text.size(), "\r\n\r\n");
  parts[0] = headers.asBytes();
  parts[1] = text.asBytes();
  return stream.write(parts).attach(kj::mv(headers));
}

kj::Promise<kj::String> JsonRpc::ContentLengthTransport::receive() {
  return input->readMessage()
      .then([this](kj::HttpInputStream::Message&& message) {
    auto promise = message.body->readAllText();
    return promise.attach(kj::mv(message.body));
  });
}

}  // namespace capnp
<commit_msg>Silence compiler warning<commit_after>// Copyright (c) 2018 Kenton Varda and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#include "json-rpc.h"
#include <kj/compat/http.h>
#include <capnp/compat/json-rpc.capnp.h>

namespace capnp {

static constexpr uint64_t JSON_NAME_ANNOTATION_ID = 0xfa5b1fd61c2e7c3dull;
static constexpr uint64_t JSON_NOTIFICATION_ANNOTATION_ID = 0xa0a054dea32fd98cull;

class JsonRpc::CapabilityImpl final: public DynamicCapability::Server {
public:
  CapabilityImpl(JsonRpc& parent, InterfaceSchema schema)
      : DynamicCapability::Server(schema), parent(parent) {}

  kj::Promise<void> call(InterfaceSchema::Method method,
                         CallContext<DynamicStruct, DynamicStruct> context) override {
    auto proto = method.getProto();
    bool isNotification = false;
    kj::StringPtr name = proto.getName();
    for (auto annotation: proto.getAnnotations()) {
      switch (annotation.getId()) {
        case JSON_NAME_ANNOTATION_ID:
          name = annotation.getValue().getText();
          break;
        case JSON_NOTIFICATION_ANNOTATION_ID:
          isNotification = true;
          break;
      }
    }

    capnp::MallocMessageBuilder message;
    auto value = message.getRoot<json::Value>();
    auto list = value.initObject(3 + !isNotification);

    uint index = 0;

    auto jsonrpc = list[index++];
    jsonrpc.setName("jsonrpc");
    jsonrpc.initValue().setString("2.0");

    uint callId = parent.callCount++;

    if (!isNotification) {
      auto id = list[index++];
      id.setName("id");
      id.initValue().setNumber(callId);
    }

    auto methodName = list[index++];
    methodName.setName("method");
    methodName.initValue().setString(name);

    auto params = list[index++];
    params.setName("params");
    parent.codec.encode(context.getParams(), params.initValue());

    auto writePromise = parent.queueWrite(parent.codec.encode(value));

    if (isNotification) {
      auto sproto = context.getResultsType().getProto().getStruct();
      MessageSize size { sproto.getDataWordCount(), sproto.getPointerCount() };
      context.initResults(size);
      return kj::mv(writePromise);
    } else {
      auto paf = kj::newPromiseAndFulfiller<void>();
      parent.awaitedResponses.insert(callId, AwaitedResponse { context, kj::mv(paf.fulfiller) });
      auto promise = writePromise.then([p = kj::mv(paf.promise)]() mutable { return kj::mv(p); });
      auto& parentRef = parent;
      return promise.attach(kj::defer([&parentRef,callId]() {
        parentRef.awaitedResponses.erase(callId);
      }));
    }
  }

private:
  JsonRpc& parent;
};

JsonRpc::JsonRpc(Transport& transport, DynamicCapability::Client interface)
    : JsonRpc(transport, kj::mv(interface), kj::newPromiseAndFulfiller<void>()) {}
JsonRpc::JsonRpc(Transport& transport, DynamicCapability::Client interfaceParam,
                 kj::PromiseFulfillerPair<void> paf)
    : transport(transport),
      interface(kj::mv(interfaceParam)),
      errorPromise(paf.promise.fork()),
      errorFulfiller(kj::mv(paf.fulfiller)),
      readTask(readLoop().eagerlyEvaluate([this](kj::Exception&& e) {
        errorFulfiller->reject(kj::mv(e));
      })),
      tasks(*this) {
  codec.handleByAnnotation(interface.getSchema());
  codec.handleByAnnotation<json::RpcMessage>();

  for (auto method: interface.getSchema().getMethods()) {
    auto proto = method.getProto();
    kj::StringPtr name = proto.getName();
    for (auto annotation: proto.getAnnotations()) {
      switch (annotation.getId()) {
        case JSON_NAME_ANNOTATION_ID:
          name = annotation.getValue().getText();
          break;
      }
    }
    methodMap.insert(name, method);
  }
}

DynamicCapability::Client JsonRpc::getPeer(InterfaceSchema schema) {
  codec.handleByAnnotation(interface.getSchema());
  return kj::heap<CapabilityImpl>(*this, schema);
}

static kj::HttpHeaderTable& staticHeaderTable() {
  static kj::HttpHeaderTable HEADER_TABLE;
  return HEADER_TABLE;
}

kj::Promise<void> JsonRpc::queueWrite(kj::String text) {
  auto fork = writeQueue.then([this, text = kj::mv(text)]() mutable {
    auto promise = transport.send(text);
    return promise.attach(kj::mv(text));
  }).eagerlyEvaluate([this](kj::Exception&& e) {
    errorFulfiller->reject(kj::mv(e));
  }).fork();
  writeQueue = fork.addBranch();
  return fork.addBranch();
}

void JsonRpc::queueError(kj::Maybe<json::Value::Reader> id, int code, kj::StringPtr message) {
  MallocMessageBuilder capnpMessage;
  auto jsonResponse = capnpMessage.getRoot<json::RpcMessage>();
  jsonResponse.setJsonrpc("2.0");
  KJ_IF_MAYBE(i, id) {
    jsonResponse.setId(*i);
  } else {
    jsonResponse.initId().setNull();
  }
  auto error = jsonResponse.initError();
  error.setCode(code);
  error.setMessage(message);

  // OK to discard result of queueWrite() since it's just one branch of a fork.
  queueWrite(codec.encode(jsonResponse));
}

kj::Promise<void> JsonRpc::readLoop() {
  return transport.receive().then([this](kj::String message) -> kj::Promise<void> {
    MallocMessageBuilder capnpMessage;
    auto rpcMessageBuilder = capnpMessage.getRoot<json::RpcMessage>();

    KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
      codec.decode(message, rpcMessageBuilder);
    })) {
      queueError(nullptr, -32700, kj::str("Parse error: ", exception->getDescription()));
      return readLoop();
    }

    KJ_CONTEXT("decoding JSON-RPC message", message);

    auto rpcMessage = rpcMessageBuilder.asReader();

    if (!rpcMessage.hasJsonrpc()) {
      queueError(nullptr, -32700, kj::str("Missing 'jsonrpc' field."));
      return readLoop();
    } else if (rpcMessage.getJsonrpc() != "2.0") {
      queueError(nullptr, -32700,
          kj::str("Unknown JSON-RPC version. This peer implements version '2.0'."));
      return readLoop();
    }

    switch (rpcMessage.which()) {
      case json::RpcMessage::NONE:
        queueError(nullptr, -32700, kj::str("message has none of params, result, or error"));
        break;

      case json::RpcMessage::PARAMS: {
        // a call
        auto schema = interface.getSchema();
        KJ_IF_MAYBE(method, schema.findMethodByName(rpcMessage.getMethod())) {
          auto req = interface.newRequest(*method);
          KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
            codec.decode(rpcMessage.getParams(), req);
          })) {
            kj::Maybe<JsonValue::Reader> id;
            if (rpcMessage.hasId()) id = rpcMessage.getId();
            queueError(id, -32602,
                kj::str("Type error in method params: ", exception->getDescription()));
            break;
          }

          if (rpcMessage.hasId()) {
            auto id = rpcMessage.getId();
            auto idCopy = kj::heapArray<word>(id.totalSize().wordCount + 1);
            memset(idCopy.begin(), 0, idCopy.asBytes().size());
            copyToUnchecked(id, idCopy);
            auto idPtr = readMessageUnchecked<json::Value>(idCopy.begin());

            auto promise = req.send()
                .then([this,idPtr](Response<DynamicStruct> response) mutable {
              MallocMessageBuilder capnpMessage;
              auto jsonResponse = capnpMessage.getRoot<json::RpcMessage>();
              jsonResponse.setJsonrpc("2.0");
              jsonResponse.setId(idPtr);
              codec.encode(DynamicStruct::Reader(response), jsonResponse.initResult());
              return queueWrite(codec.encode(jsonResponse));
            }, [this,idPtr](kj::Exception&& e) {
              MallocMessageBuilder capnpMessage;
              auto jsonResponse = capnpMessage.getRoot<json::RpcMessage>();
              jsonResponse.setJsonrpc("2.0");
              jsonResponse.setId(idPtr);
              auto error = jsonResponse.initError();
              switch (e.getType()) {
                case kj::Exception::Type::FAILED:
                  error.setCode(-32000);
                  break;
                case kj::Exception::Type::DISCONNECTED:
                  error.setCode(-32001);
                  break;
                case kj::Exception::Type::OVERLOADED:
                  error.setCode(-32002);
                  break;
                case kj::Exception::Type::UNIMPLEMENTED:
                  error.setCode(-32601);  // method not found
                  break;
              }
              error.setMessage(e.getDescription());
              return queueWrite(codec.encode(jsonResponse));
            });
            tasks.add(promise.attach(kj::mv(idCopy)));
          } else {
            // No 'id', so this is a notification.
            tasks.add(req.send().ignoreResult().catch_([](kj::Exception&& exception) {
              if (exception.getType() != kj::Exception::Type::UNIMPLEMENTED) {
                KJ_LOG(ERROR, "JSON-RPC notification threw exception into the abyss", exception);
              }
            }));
          }
        } else {
          if (rpcMessage.hasId()) {
            queueError(rpcMessage.getId(), -32601, "Method not found");
          } else {
            // Ignore notification for unknown method.
          }
        }
        break;
      }

      case json::RpcMessage::RESULT: {
        auto id = rpcMessage.getId();
        if (!id.isNumber()) {
          // JSON-RPC doesn't define what to do if receiving a response with an invalid id.
          KJ_LOG(ERROR, "JSON-RPC response has invalid ID");
        } else KJ_IF_MAYBE(awaited, awaitedResponses.find((uint)id.getNumber())) {
          KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
            codec.decode(rpcMessage.getResult(), awaited->context.getResults());
            awaited->fulfiller->fulfill();
          })) {
            // Errors always propagate from callee to caller, so we don't want to throw this error
            // back to the server.
            awaited->fulfiller->reject(kj::mv(*exception));
          }
        } else {
          // Probably, this is the response to a call that was canceled.
        }
        break;
      }

      case json::RpcMessage::ERROR: {
        auto id = rpcMessage.getId();
        if (id.isNull()) {
          // Error message will be logged by KJ_CONTEXT, above.
          KJ_LOG(ERROR, "peer reports JSON-RPC protocol error");
        } else if (!id.isNumber()) {
          // JSON-RPC doesn't define what to do if receiving a response with an invalid id.
          KJ_LOG(ERROR, "JSON-RPC response has invalid ID");
        } else KJ_IF_MAYBE(awaited, awaitedResponses.find((uint)id.getNumber())) {
          auto error = rpcMessage.getError();
          auto code = error.getCode();
          kj::Exception::Type type =
              code == -32601 ? kj::Exception::Type::UNIMPLEMENTED
                             : kj::Exception::Type::FAILED;
          awaited->fulfiller->reject(kj::Exception(
              type, __FILE__, __LINE__, kj::str(error.getMessage())));
        } else {
          // Probably, this is the response to a call that was canceled.
        }
        break;
      }
    }

    return readLoop();
  });
}

void JsonRpc::taskFailed(kj::Exception&& exception) {
  errorFulfiller->reject(kj::mv(exception));
}

// =======================================================================================

JsonRpc::ContentLengthTransport::ContentLengthTransport(kj::AsyncIoStream& stream)
    : stream(stream), input(kj::newHttpInputStream(stream, staticHeaderTable())) {}
JsonRpc::ContentLengthTransport::~ContentLengthTransport() noexcept(false) {}

kj::Promise<void> JsonRpc::ContentLengthTransport::send(kj::StringPtr text) {
  auto headers = kj::str("Content-Length: ", text.size(), "\r\n\r\n");
  parts[0] = headers.asBytes();
  parts[1] = text.asBytes();
  return stream.write(parts).attach(kj::mv(headers));
}

kj::Promise<kj::String> JsonRpc::ContentLengthTransport::receive() {
  return input->readMessage()
      .then([](kj::HttpInputStream::Message&& message) {
    auto promise = message.body->readAllText();
    return promise.attach(kj::mv(message.body));
  });
}

}  // namespace capnp
<|endoftext|>


<commit_before>// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <windows.h>

#include "chrome_frame/chrome_launcher.h"

// We want to keep this EXE tiny, so we avoid all dependencies and link to no
// libraries, and we do not use the C runtime.
//
// To catch errors in debug builds, we define an extremely simple assert macro.
#ifndef NDEBUG
#define CLM_ASSERT(x) do { if (!(x)) { ::DebugBreak(); } } while (false)
#else
#define CLM_ASSERT(x)
#endif  // NDEBUG

// In release builds, we skip the standard library completely to minimize
// size.  This is more work in debug builds, and unnecessary, hence the
// different signatures.
#ifndef NDEBUG
int APIENTRY wWinMain(HINSTANCE, HINSTANCE, wchar_t*, int) {
#else
extern "C" void __cdecl WinMainCRTStartup() {
#endif  // NDEBUG
  // This relies on the chrome_launcher.exe residing in the same directory
  // as our DLL.  We build a full path to avoid loading it from any other
  // directory in the DLL search path.
  //
  // The code is a bit verbose because we can't use the standard library.
  const wchar_t kBaseName[] = L"npchrome_tab.dll";
  wchar_t file_path[MAX_PATH + (sizeof(kBaseName) / sizeof(kBaseName[0])) + 1];
  file_path[0] = L'\0';
  ::GetModuleFileName(::GetModuleHandle(NULL), file_path, MAX_PATH);

  // Find index of last slash, and null-terminate the string after it.
  //
  // Proof for security purposes, since we can't use the safe string
  // manipulation functions from the runtime:
  // - File_path is always null-terminated, by us initially and by
  //   ::GetModuleFileName if it puts anything into the buffer.
  // - If there is no slash in the path then it's a relative path, not an
  //   absolute one, and the code ends up creating a relative path to
  //   npchrome_tab.dll.
  // - It's safe to use lstrcatW since we know the maximum length of both
  //   parts we are concatenating, and we know the buffer will fit them in
  //   the worst case.
  int slash_index = lstrlenW(file_path);
  // Invariant: 0 <= slash_index < MAX_PATH
  CLM_ASSERT(slash_index > 0);
  while (slash_index > 0 && file_path[slash_index] != L'\\')
    --slash_index;
  // Invariant: 0 <= slash_index < MAX_PATH and it is either the index of
  // the last \ in the path, or 0.
  if (slash_index != 0)
    ++slash_index;  // don't remove the last '\'
  file_path[slash_index] = L'\0';

  lstrcatW(file_path, kBaseName);

  UINT exit_code = ERROR_FILE_NOT_FOUND;
  HMODULE chrome_tab = ::LoadLibrary(file_path);
  CLM_ASSERT(chrome_tab);
  if (chrome_tab) {
    chrome_launcher::CfLaunchChromeProc proc =
        reinterpret_cast<chrome_launcher::CfLaunchChromeProc>(
            ::GetProcAddress(chrome_tab, "CfLaunchChrome"));
    CLM_ASSERT(proc);
    if (proc) {
      exit_code = proc();
    } else {
      exit_code = ERROR_INVALID_FUNCTION;
    }

    ::FreeLibrary(chrome_tab);
  }
  ::ExitProcess(exit_code);
}
<commit_msg>Setting the SVN eol style to LF on chrome frame sources<commit_after>// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <windows.h>

#include "chrome_frame/chrome_launcher.h"

// We want to keep this EXE tiny, so we avoid all dependencies and link to no
// libraries, and we do not use the C runtime.
//
// To catch errors in debug builds, we define an extremely simple assert macro.
#ifndef NDEBUG
#define CLM_ASSERT(x) do { if (!(x)) { ::DebugBreak(); } } while (false)
#else
#define CLM_ASSERT(x)
#endif  // NDEBUG

// In release builds, we skip the standard library completely to minimize
// size.  This is more work in debug builds, and unnecessary, hence the
// different signatures.
#ifndef NDEBUG
int APIENTRY wWinMain(HINSTANCE, HINSTANCE, wchar_t*, int) {
#else
extern "C" void __cdecl WinMainCRTStartup() {
#endif  // NDEBUG
  // This relies on the chrome_launcher.exe residing in the same directory
  // as our DLL.  We build a full path to avoid loading it from any other
  // directory in the DLL search path.
  //
  // The code is a bit verbose because we can't use the standard library.
  const wchar_t kBaseName[] = L"npchrome_tab.dll";
  wchar_t file_path[MAX_PATH + (sizeof(kBaseName) / sizeof(kBaseName[0])) + 1];
  file_path[0] = L'\0';
  ::GetModuleFileName(::GetModuleHandle(NULL), file_path, MAX_PATH);

  // Find index of last slash, and null-terminate the string after it.
  //
  // Proof for security purposes, since we can't use the safe string
  // manipulation functions from the runtime:
  // - File_path is always null-terminated, by us initially and by
  //   ::GetModuleFileName if it puts anything into the buffer.
  // - If there is no slash in the path then it's a relative path, not an
  //   absolute one, and the code ends up creating a relative path to
  //   npchrome_tab.dll.
  // - It's safe to use lstrcatW since we know the maximum length of both
  //   parts we are concatenating, and we know the buffer will fit them in
  //   the worst case.
  int slash_index = lstrlenW(file_path);
  // Invariant: 0 <= slash_index < MAX_PATH
  CLM_ASSERT(slash_index > 0);
  while (slash_index > 0 && file_path[slash_index] != L'\\')
    --slash_index;
  // Invariant: 0 <= slash_index < MAX_PATH and it is either the index of
  // the last \ in the path, or 0.
  if (slash_index != 0)
    ++slash_index;  // don't remove the last '\'
  file_path[slash_index] = L'\0';

  lstrcatW(file_path, kBaseName);

  UINT exit_code = ERROR_FILE_NOT_FOUND;
  HMODULE chrome_tab = ::LoadLibrary(file_path);
  CLM_ASSERT(chrome_tab);
  if (chrome_tab) {
    chrome_launcher::CfLaunchChromeProc proc =
        reinterpret_cast<chrome_launcher::CfLaunchChromeProc>(
            ::GetProcAddress(chrome_tab, "CfLaunchChrome"));
    CLM_ASSERT(proc);
    if (proc) {
      exit_code = proc();
    } else {
      exit_code = ERROR_INVALID_FUNCTION;
    }

    ::FreeLibrary(chrome_tab);
  }
  ::ExitProcess(exit_code);
}
<|endoftext|>


<commit_before>// Copyright 2018 Google Inc. All Rights Reserved.
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA

#include "state_space_regression_model_manager.h"
#include "state_space_gaussian_model_manager.h"
#include "utils.h"

#include "r_interface/list_io.hpp"
#include "r_interface/prior_specification.hpp"
#include "Models/Glm/PosteriorSamplers/BregVsSampler.hpp"
#include "Models/Glm/PosteriorSamplers/SpikeSlabDaRegressionSampler.hpp"
#include "Models/StateSpace/PosteriorSamplers/StateSpacePosteriorSampler.hpp"
#include "Models/StateSpace/StateSpaceModel.hpp"
#include "Models/StateSpace/StateSpaceRegressionModel.hpp"
#include "cpputil/report_error.hpp"

namespace BOOM {
namespace bsts {

namespace {
typedef StateSpaceRegressionModelManager SSRMF;
}  // namespace

StateSpaceRegressionModelManager::StateSpaceRegressionModelManager()
    : predictor_dimension_(-1) {}

StateSpaceRegressionModel * SSRMF::CreateBareModel(
    SEXP r_data_list,
    SEXP r_prior,
    SEXP r_options,
    RListIoManager *io_manager) {
  Matrix predictors;
  Vector response;
  std::vector<bool> response_is_observed;
  if (!Rf_isNull(r_data_list)) {
    if (Rf_inherits(r_data_list, "bsts")) {
      predictors = ToBoomMatrix(getListElement(r_data_list, "predictors"));
      SEXP r_response = getListElement(r_data_list, "original.series");
      response = ToBoomVector(r_response);
      response_is_observed = IsObserved(r_response);
    } else {
      // If we were passed data from R then use it to build the model.
      predictors = ToBoomMatrix(getListElement(r_data_list, "predictors"));
      response = ToBoomVector(getListElement(r_data_list, "response"));
      response_is_observed = ToVectorBool(getListElement(
          r_data_list, "response.is.observed"));
    }
    UnpackTimestampInfo(r_data_list);
    if (TimestampsAreTrivial()) {
      model_.reset(new StateSpaceRegressionModel(
          response,
          predictors,
          response_is_observed));
    } else {
      // timestamps are non-trivial.
      model_.reset(new StateSpaceRegressionModel(ncol(predictors)));
      std::vector<Ptr<StateSpace::MultiplexedRegressionData>> data;
      data.reserve(NumberOfTimePoints());
      for (int i = 0; i < NumberOfTimePoints(); ++i) {
        data.push_back(new StateSpace::MultiplexedRegressionData);
      }
      for (int i = 0; i < response.size(); ++i) {
        NEW(RegressionData, observation)(response[i], predictors.row(i));
        if (!response_is_observed[i]) {
          observation->set_missing_status(Data::completely_missing);
        }
        data[TimestampMapping(i)]->add_data(observation);
      }
      for (int i = 0; i < NumberOfTimePoints(); ++i) {
        if (data[i]->observed_sample_size() == 0) {
          data[i]->set_missing_status(Data::completely_missing);
        }
        model_->add_multiplexed_data(data[i]);
      }
    }
  } else {
    // No data was passed from R, so build the model from its default
    // constructor.  We need to know the dimension of the predictors.
    if (predictor_dimension_ < 0) {
      report_error("If r_data_list is not passed, you must call "
                   "SetPredictorDimension before calling "
                   "CreateBareModel.");
    }
    model_.reset(new StateSpaceRegressionModel(predictor_dimension_));
  }

  // A NULL r_prior signals that no posterior sampler is needed.
  if (!Rf_isNull(r_prior)) {
    SetRegressionSampler(r_prior, r_options);
    Ptr<StateSpacePosteriorSampler> sampler(
        new StateSpacePosteriorSampler(model_.get()));
    if (!Rf_isNull(r_options)
        && !Rf_asLogical(getListElement(r_options, "enable.threads"))) {
      sampler->disable_threads();
    }
    model_->set_method(sampler);
  }

  // Make the io_manager aware of the model parameters.
  Ptr<RegressionModel> regression(model_->regression_model());
  io_manager->add_list_element(
      new GlmCoefsListElement(regression->coef_prm(), "coefficients"));
  io_manager->add_list_element(
      new StandardDeviationListElement(regression->Sigsq_prm(),
                                       "sigma.obs"));
  return model_.get();
}

void SSRMF::AddDataFromBstsObject(SEXP r_bsts_object) {
  AddData(ToBoomVector(getListElement(r_bsts_object, "original.series", true)),
          ToBoomMatrix(getListElement(r_bsts_object, "predictors", true)),
          IsObserved(getListElement(r_bsts_object, "original.series", true)));
}

void SSRMF::AddDataFromList(SEXP r_data_list) {
  AddData(ToBoomVector(getListElement(r_data_list, "response", true)),
          ToBoomMatrix(getListElement(r_data_list, "predictors", true)),
          ToVectorBool(getListElement(r_data_list,
                                      "response.is.observed", true)));
}

int SSRMF::UnpackForecastData(SEXP r_prediction_data) {
  forecast_predictors_ = ToBoomMatrix(getListElement(
      r_prediction_data, "predictors"));
  UnpackForecastTimestamps(r_prediction_data);
  return forecast_predictors_.nrow();
}

Vector SSRMF::SimulateForecast(const Vector &final_state) {
  if (ForecastTimestamps().empty()) {
    return model_->simulate_forecast(rng(), forecast_predictors_, final_state);
  } else {
    return model_->simulate_multiplex_forecast(rng(),
                                               forecast_predictors_,
                                               final_state,
                                               ForecastTimestamps());
  }
}

void SSRMF::SetRegressionSampler(SEXP r_regression_prior,
                                 SEXP r_options) {
  // If either the prior object or the bma method is NULL then take
  // that as a signal the model is not being specified for the
  // purposes of MCMC, and bail out.
  if (Rf_isNull(r_regression_prior)
      || Rf_isNull(r_options)
      || Rf_isNull(getListElement(r_options, "bma.method"))) {
    return;
  }
  std::string bma_method = BOOM::ToString(getListElement(
      r_options, "bma.method"));
  if (bma_method == "SSVS") {
    SetSsvsRegressionSampler(r_regression_prior);
  } else if (bma_method == "ODA") {
    SetOdaRegressionSampler(r_regression_prior, r_options);
  } else {
    std::ostringstream err;
    err << "Unrecognized value of bma_method: " << bma_method;
    BOOM::report_error(err.str());
  }
}

void SSRMF::SetSsvsRegressionSampler(SEXP r_regression_prior) {
  BOOM::RInterface::RegressionConjugateSpikeSlabPrior prior(
      r_regression_prior, model_->regression_model()->Sigsq_prm());
  DropUnforcedCoefficients(model_->regression_model(),
                           prior.prior_inclusion_probabilities());
  Ptr<BregVsSampler> sampler(new BregVsSampler(
      model_->regression_model().get(),
      prior.slab(),
      prior.siginv_prior(),
      prior.spike()));
  sampler->set_sigma_upper_limit(prior.sigma_upper_limit());
  int max_flips = prior.max_flips();
  if (max_flips > 0) {
    sampler->limit_model_selection(max_flips);
  }
  model_->regression_model()->set_method(sampler);
}

void SSRMF::SetOdaRegressionSampler(SEXP r_regression_prior,
                                    SEXP r_options) {
  SEXP r_oda_options = getListElement(r_options, "oda.options");
  BOOM::RInterface::IndependentRegressionSpikeSlabPrior prior(
      r_regression_prior, model_->regression_model()->Sigsq_prm());
  double eigenvalue_fudge_factor = 0.001;
  double fallback_probability = 0.0;
  if (!Rf_isNull(r_oda_options)) {
    eigenvalue_fudge_factor = Rf_asReal(
        getListElement(r_oda_options, "eigenvalue.fudge.factor"));
    fallback_probability = Rf_asReal(
        getListElement(r_oda_options, "fallback.probability"));
  }
  Ptr<SpikeSlabDaRegressionSampler> sampler(
      new SpikeSlabDaRegressionSampler(
          model_->regression_model().get(),
          prior.slab(),
          prior.siginv_prior(),
          prior.prior_inclusion_probabilities(),
          eigenvalue_fudge_factor,
          fallback_probability));
  sampler->set_sigma_upper_limit(prior.sigma_upper_limit());
  DropUnforcedCoefficients(model_->regression_model(),
                           prior.prior_inclusion_probabilities());
  model_->regression_model()->set_method(sampler);
}

void StateSpaceRegressionModelManager::SetPredictorDimension(int xdim) {
  predictor_dimension_ = xdim;
}

void StateSpaceRegressionModelManager::AddData(
    const Vector &response,
    const Matrix &predictors,
    const std::vector<bool> &response_is_observed) {
  if (nrow(predictors) != response.size()
      || response_is_observed.size() != response.size()) {
    std::ostringstream err;
    err << "Argument sizes do not match in "
        << "StateSpaceRegressionModelManager::AddData" << endl
        << "nrow(predictors) = " << nrow(predictors) << endl
        << "response.size()  = " << response.size() << endl
        << "observed.size()  = " << response_is_observed.size();
    report_error(err.str());
  }

  for (int i = 0; i < response.size(); ++i) {
    Ptr<RegressionData> dp(new RegressionData(response[i], predictors.row(i)));
    if (!response_is_observed[i]) {
      dp->set_missing_status(Data::partly_missing);
    }
    model_->add_regression_data(dp);
  }
}

namespace {
typedef StateSpaceRegressionHoldoutErrorSampler ErrorSampler;
}  // namespace

void ErrorSampler::sample_holdout_prediction_errors() {
  model_->sample_posterior();
  errors_->resize(niter_, model_->time_dimension() + holdout_responses_.size());
  for (int i = 0; i < niter_; ++i) {
    model_->sample_posterior();
    Vector all_errors = model_->one_step_prediction_errors(standardize_);
    all_errors.concat(model_->one_step_holdout_prediction_errors(
        holdout_predictors_, holdout_responses_, model_->final_state(), standardize_));
    errors_->row(i) = all_errors;
  }
}

HoldoutErrorSampler StateSpaceRegressionModelManager::CreateHoldoutSampler(
    SEXP r_bsts_object,
    int cutpoint,
    bool standardize,
    Matrix *prediction_error_output) {
  RListIoManager io_manager;
  Ptr<StateSpaceRegressionModel> model =
      static_cast<StateSpaceRegressionModel *>(CreateModel(
          R_NilValue,
          getListElement(r_bsts_object, "state.specification"),
          getListElement(r_bsts_object, "prior"),
          getListElement(r_bsts_object, "model.options"),
          &io_manager));
  AddDataFromBstsObject(r_bsts_object);

  std::vector<Ptr<StateSpace::MultiplexedRegressionData>> data = model->dat();
  model->clear_data();
  for (int i = 0; i <= cutpoint; ++i) {
    model->add_multiplexed_data(data[i]);
  }
  int holdout_sample_size = 0;
  for (int i = cutpoint + 1; i < data.size(); ++i) {
    holdout_sample_size += data[i]->total_sample_size();
  }
  Matrix holdout_predictors(holdout_sample_size,
                            model->observation_model()->xdim());
  Vector holdout_response(holdout_sample_size);
  int index = 0;
  for (int i = cutpoint + 1; i < data.size(); ++i) {
    for (int j = 0; j < data[i]->total_sample_size(); ++j) {
      holdout_predictors.row(index) = data[i]->regression_data(j).x();
      holdout_response[index] = data[i]->regression_data(j).y();
      ++index;
    }
  }
  return HoldoutErrorSampler(new ErrorSampler(
      model, holdout_response, holdout_predictors,
      Rf_asInteger(getListElement(r_bsts_object, "niter")),
      standardize,
      prediction_error_output));
}

}  // namespace bsts
}  // namespace BOOM
<commit_msg>Formatting and whitespace improvements.<commit_after>// Copyright 2018 Google Inc. All Rights Reserved.
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA

#include "state_space_regression_model_manager.h"
#include "state_space_gaussian_model_manager.h"
#include "utils.h"

#include "r_interface/list_io.hpp"
#include "r_interface/prior_specification.hpp"
#include "Models/Glm/PosteriorSamplers/BregVsSampler.hpp"
#include "Models/Glm/PosteriorSamplers/SpikeSlabDaRegressionSampler.hpp"
#include "Models/StateSpace/PosteriorSamplers/StateSpacePosteriorSampler.hpp"
#include "Models/StateSpace/StateSpaceModel.hpp"
#include "Models/StateSpace/StateSpaceRegressionModel.hpp"
#include "cpputil/report_error.hpp"

namespace BOOM {
  namespace bsts {

    namespace {
      typedef StateSpaceRegressionModelManager SSRMF;
    }  // namespace

    StateSpaceRegressionModelManager::StateSpaceRegressionModelManager()
        : predictor_dimension_(-1) {}

    StateSpaceRegressionModel * SSRMF::CreateBareModel(
        SEXP r_data_list,
        SEXP r_prior,
        SEXP r_options,
        RListIoManager *io_manager) {
      Matrix predictors;
      Vector response;
      std::vector<bool> response_is_observed;
      if (!Rf_isNull(r_data_list)) {
        if (Rf_inherits(r_data_list, "bsts")) {
          predictors = ToBoomMatrix(getListElement(r_data_list, "predictors"));
          SEXP r_response = getListElement(r_data_list, "original.series");
          response = ToBoomVector(r_response);
          response_is_observed = IsObserved(r_response);
        } else {
          // If we were passed data from R then use it to build the model.
          predictors = ToBoomMatrix(getListElement(r_data_list, "predictors"));
          response = ToBoomVector(getListElement(r_data_list, "response"));
          response_is_observed = ToVectorBool(getListElement(
              r_data_list, "response.is.observed"));
        }
        UnpackTimestampInfo(r_data_list);
        if (TimestampsAreTrivial()) {
          model_.reset(new StateSpaceRegressionModel(
              response,
              predictors,
              response_is_observed));
        } else {
          // timestamps are non-trivial.
          model_.reset(new StateSpaceRegressionModel(ncol(predictors)));
          std::vector<Ptr<StateSpace::MultiplexedRegressionData>> data;
          data.reserve(NumberOfTimePoints());
          for (int i = 0; i < NumberOfTimePoints(); ++i) {
            data.push_back(new StateSpace::MultiplexedRegressionData);
          }
          for (int i = 0; i < response.size(); ++i) {
            NEW(RegressionData, observation)(response[i], predictors.row(i));
            if (!response_is_observed[i]) {
              observation->set_missing_status(Data::completely_missing);
            }
            data[TimestampMapping(i)]->add_data(observation);
          }
          for (int i = 0; i < NumberOfTimePoints(); ++i) {
            if (data[i]->observed_sample_size() == 0) {
              data[i]->set_missing_status(Data::completely_missing);
            }
            model_->add_multiplexed_data(data[i]);
          }
        }
      } else {
        // No data was passed from R, so build the model from its default
        // constructor.  We need to know the dimension of the predictors.
        if (predictor_dimension_ < 0) {
          report_error("If r_data_list is not passed, you must call "
                       "SetPredictorDimension before calling "
                       "CreateBareModel.");
        }
        model_.reset(new StateSpaceRegressionModel(predictor_dimension_));
      }

      // A NULL r_prior signals that no posterior sampler is needed.
      if (!Rf_isNull(r_prior)) {
        SetRegressionSampler(r_prior, r_options);
        Ptr<StateSpacePosteriorSampler> sampler(
            new StateSpacePosteriorSampler(model_.get()));
        if (!Rf_isNull(r_options)
            && !Rf_asLogical(getListElement(r_options, "enable.threads"))) {
          sampler->disable_threads();
        }
        model_->set_method(sampler);
      }

      // Make the io_manager aware of the model parameters.
      Ptr<RegressionModel> regression(model_->regression_model());
      io_manager->add_list_element(
          new GlmCoefsListElement(regression->coef_prm(), "coefficients"));
      io_manager->add_list_element(
          new StandardDeviationListElement(regression->Sigsq_prm(),
                                           "sigma.obs"));
      return model_.get();
    }

    void SSRMF::AddDataFromBstsObject(SEXP r_bsts_object) {
      AddData(ToBoomVector(getListElement(
                  r_bsts_object, "original.series", true)),
              ToBoomMatrix(getListElement(
                  r_bsts_object, "predictors", true)),
              IsObserved(getListElement(
                  r_bsts_object, "original.series", true)));
    }

    void SSRMF::AddDataFromList(SEXP r_data_list) {
      AddData(ToBoomVector(getListElement(r_data_list, "response", true)),
              ToBoomMatrix(getListElement(r_data_list, "predictors", true)),
              ToVectorBool(getListElement(r_data_list,
                                          "response.is.observed", true)));
    }

    int SSRMF::UnpackForecastData(SEXP r_prediction_data) {
      forecast_predictors_ = ToBoomMatrix(getListElement(
          r_prediction_data, "predictors"));
      UnpackForecastTimestamps(r_prediction_data);
      return forecast_predictors_.nrow();
    }

    Vector SSRMF::SimulateForecast(const Vector &final_state) {
      if (ForecastTimestamps().empty()) {
        return model_->simulate_forecast(
            rng(), forecast_predictors_, final_state);
      } else {
        return model_->simulate_multiplex_forecast(rng(),
                                                   forecast_predictors_,
                                                   final_state,
                                                   ForecastTimestamps());
      }
    }

    void SSRMF::SetRegressionSampler(SEXP r_regression_prior,
                                     SEXP r_options) {
      // If either the prior object or the bma method is NULL then take
      // that as a signal the model is not being specified for the
      // purposes of MCMC, and bail out.
      if (Rf_isNull(r_regression_prior)
          || Rf_isNull(r_options)
          || Rf_isNull(getListElement(r_options, "bma.method"))) {
        return;
      }
      std::string bma_method = BOOM::ToString(getListElement(
          r_options, "bma.method"));
      if (bma_method == "SSVS") {
        SetSsvsRegressionSampler(r_regression_prior);
      } else if (bma_method == "ODA") {
        SetOdaRegressionSampler(r_regression_prior, r_options);
      } else {
        std::ostringstream err;
        err << "Unrecognized value of bma_method: " << bma_method;
        BOOM::report_error(err.str());
      }
    }

    void SSRMF::SetSsvsRegressionSampler(SEXP r_regression_prior) {
      BOOM::RInterface::RegressionConjugateSpikeSlabPrior prior(
          r_regression_prior, model_->regression_model()->Sigsq_prm());
      DropUnforcedCoefficients(model_->regression_model(),
                               prior.prior_inclusion_probabilities());
      Ptr<BregVsSampler> sampler(new BregVsSampler(
          model_->regression_model().get(),
          prior.slab(),
          prior.siginv_prior(),
          prior.spike()));
      sampler->set_sigma_upper_limit(prior.sigma_upper_limit());
      int max_flips = prior.max_flips();
      if (max_flips > 0) {
        sampler->limit_model_selection(max_flips);
      }
      model_->regression_model()->set_method(sampler);
    }

    void SSRMF::SetOdaRegressionSampler(SEXP r_regression_prior,
                                        SEXP r_options) {
      SEXP r_oda_options = getListElement(r_options, "oda.options");
      BOOM::RInterface::IndependentRegressionSpikeSlabPrior prior(
          r_regression_prior, model_->regression_model()->Sigsq_prm());
      double eigenvalue_fudge_factor = 0.001;
      double fallback_probability = 0.0;
      if (!Rf_isNull(r_oda_options)) {
        eigenvalue_fudge_factor = Rf_asReal(
            getListElement(r_oda_options, "eigenvalue.fudge.factor"));
        fallback_probability = Rf_asReal(
            getListElement(r_oda_options, "fallback.probability"));
      }
      Ptr<SpikeSlabDaRegressionSampler> sampler(
          new SpikeSlabDaRegressionSampler(
              model_->regression_model().get(),
              prior.slab(),
              prior.siginv_prior(),
              prior.prior_inclusion_probabilities(),
              eigenvalue_fudge_factor,
              fallback_probability));
      sampler->set_sigma_upper_limit(prior.sigma_upper_limit());
      DropUnforcedCoefficients(model_->regression_model(),
                               prior.prior_inclusion_probabilities());
      model_->regression_model()->set_method(sampler);
    }

    void StateSpaceRegressionModelManager::SetPredictorDimension(int xdim) {
      predictor_dimension_ = xdim;
    }

    void StateSpaceRegressionModelManager::AddData(
        const Vector &response,
        const Matrix &predictors,
        const std::vector<bool> &response_is_observed) {
      if (nrow(predictors) != response.size()
          || response_is_observed.size() != response.size()) {
        std::ostringstream err;
        err << "Argument sizes do not match in "
            << "StateSpaceRegressionModelManager::AddData" << endl
            << "nrow(predictors) = " << nrow(predictors) << endl
            << "response.size()  = " << response.size() << endl
            << "observed.size()  = " << response_is_observed.size();
        report_error(err.str());
      }

      for (int i = 0; i < response.size(); ++i) {
        Ptr<RegressionData> dp(new RegressionData(
            response[i], predictors.row(i)));
        if (!response_is_observed[i]) {
          dp->set_missing_status(Data::partly_missing);
        }
        model_->add_regression_data(dp);
      }
    }

    namespace {
      typedef StateSpaceRegressionHoldoutErrorSampler ErrorSampler;
    }  // namespace

    void ErrorSampler::sample_holdout_prediction_errors() {
      model_->sample_posterior();
      errors_->resize(
          niter_, model_->time_dimension() + holdout_responses_.size());
      for (int i = 0; i < niter_; ++i) {
        model_->sample_posterior();
        Vector all_errors = model_->one_step_prediction_errors(standardize_);
        all_errors.concat(model_->one_step_holdout_prediction_errors(
            holdout_predictors_,
            holdout_responses_,
            model_->final_state(),
            standardize_));
        errors_->row(i) = all_errors;
      }
    }

    HoldoutErrorSampler StateSpaceRegressionModelManager::CreateHoldoutSampler(
        SEXP r_bsts_object,
        int cutpoint,
        bool standardize,
        Matrix *prediction_error_output) {
      RListIoManager io_manager;
      Ptr<StateSpaceRegressionModel> model =
          static_cast<StateSpaceRegressionModel *>(CreateModel(
              R_NilValue,
              getListElement(r_bsts_object, "state.specification"),
              getListElement(r_bsts_object, "prior"),
              getListElement(r_bsts_object, "model.options"),
              &io_manager));
      AddDataFromBstsObject(r_bsts_object);

      std::vector<Ptr<StateSpace::MultiplexedRegressionData>> data(
          model->dat());
      model->clear_data();
      for (int i = 0; i <= cutpoint; ++i) {
        model->add_multiplexed_data(data[i]);
      }
      int holdout_sample_size = 0;
      for (int i = cutpoint + 1; i < data.size(); ++i) {
        holdout_sample_size += data[i]->total_sample_size();
      }
      Matrix holdout_predictors(holdout_sample_size,
                                model->observation_model()->xdim());
      Vector holdout_response(holdout_sample_size);
      int index = 0;
      for (int i = cutpoint + 1; i < data.size(); ++i) {
        for (int j = 0; j < data[i]->total_sample_size(); ++j) {
          holdout_predictors.row(index) = data[i]->regression_data(j).x();
          holdout_response[index] = data[i]->regression_data(j).y();
          ++index;
        }
      }
      return HoldoutErrorSampler(new ErrorSampler(
          model, holdout_response, holdout_predictors,
          Rf_asInteger(getListElement(r_bsts_object, "niter")),
          standardize,
          prediction_error_output));
    }

  }  // namespace bsts
}  // namespace BOOM
<|endoftext|>


<commit_before>#ifndef CPPEVOLVE_TREE_H_
#define CPPEVOLVE_TREE_H_

#include "cppEvolve/utils.hpp"

namespace evolve
{
    namespace Tree
    {

        template<typename rtype>
        class BaseNode
        {
        public:
            virtual ~BaseNode(){};
            virtual rtype eval()=0;

            std::vector<BaseNode<rtype>*>& getChildren() {
                return children;
            }

            virtual unsigned int getNumChildren() const =0;

            const std::string& getName() const {return name;}

            template<typename T>
            friend std::ostream& operator<< (std::ostream &out, const BaseNode<T>& node);

        protected:
            BaseNode(const std::string& _name) :
                name(_name){}
            std::vector<BaseNode<rtype>*> children;
            const std::string name;
        };

        template<typename genome>
        class Node : public BaseNode<typename genome::result_type>
        {
        public:
            Node(genome g, const std::string& _name) :
                BaseNode<typename genome::result_type>(_name),
                val(g)
            {
            }

            virtual ~Node(){}

            virtual typename genome::result_type eval() override
            {
                std::vector<typename genome::result_type> results;

                std::for_each(this->children.begin(), this->children.end(),
                        [&](BaseNode<typename genome::result_type>* child){results.push_back(child->eval());});

                return unpack_caller<typename genome::result_type>::eval(val, results);
            }

            virtual unsigned int getNumChildren() const override
            {
                return count_args<genome>::value;
            }

        private:
            genome val;
        };

        template<typename genome>
        class Terminator: public BaseNode<typename genome::result_type>
        {
        public:
            static_assert(count_args<genome>::value == 0, "The number of arguments for a terminator must be 0");

            Terminator(genome g, const std::string& _name) :
                BaseNode<typename genome::result_type>(_name),
                val(g){}

            virtual ~Terminator(){}

            virtual typename genome::result_type eval() override {
                return val();
            }

            virtual unsigned int getNumChildren() const override  {
                return 0;
            }

        private:
            genome val;
        };

        template<typename rType>
        class Tree
        {
        public:
            Tree(BaseNode<rType>* _root) :
                root(_root){}

            rType eval() const {
                return root->eval();
            }

            template<typename T>
            friend std::ostream& operator<< (std::ostream &out, const Tree<T>& tree);

            BaseNode<rType>* root;
        protected:

        };

        template<typename T>
        std::ostream& operator<< (std::ostream &out, const Tree<T>& tree)
        {
            out << *tree.root;
            return out;
        }

        template<typename T>
        std::ostream& operator<< (std::ostream &out, const BaseNode<T>& node)
        {
            out << node.name;

            if (!node.children.empty()) {
                out << "(";
                for (unsigned int i=0; i < node.children.size()-1; ++i)
                    out << *(node.children[i]) << ", ";

                out << *node.children[node.children.size()-1];
                out << ")";
            }
            return out;
        }

        template<typename rType>
        class TreeFactory
        {
        public:
            TreeFactory(){}

            template<typename... T>
            void addNode(rType(*f)(T...), const std::string& name)
            {
                static_assert(sizeof...(T) > 0, "Node function with 0 arguments should be terminator");

                std::function<BaseNode<rType>*()> func = [=]() {
                    return new Node<std::function<rType(T...)>>(f, name);
                };
                nodes.push_back(func);
            }

            void addTerminator(rType(*f)(), const std::string& name)
            {
                std::function<BaseNode<rType>*()> func = [=]() {
                    return new Terminator<std::function<rType()>>(f, name);
                };
                terminators.push_back(func);
            }

            Tree<rType>* make()
            {
                assert(!terminators.empty());
                auto tree = new Tree<rType>(createRandomNode());

                for (unsigned int i=0; i < tree->root->getNumChildren(); ++i)
                {
                    BaseNode<rType>* node = createRandomNode();
                    tree->root->getChildren().push_back(node);
                    for (unsigned int i=0; i < node->getNumChildren(); ++i)
                    {
                        node->getChildren().push_back(createRandomTerminator());
                    }
                }

                return tree;
            }

        protected:
            std::vector<std::function<BaseNode<rType>*()>> nodes;
            std::vector<std::function<BaseNode<rType>*()>> terminators;

            BaseNode<rType>* createRandomNode() const {
                return nodes[rand() % nodes.size()]();
            }

            BaseNode<rType>* createRandomTerminator() const {
                return terminators[rand() % terminators.size()]();
            }

        };
    }
}

#endif
<commit_msg>Converted node storage to map of ID to function<commit_after>#ifndef CPPEVOLVE_TREE_H_
#define CPPEVOLVE_TREE_H_

#include "cppEvolve/utils.hpp"
#include <map>

namespace evolve
{
    namespace Tree
    {

        template<typename rtype>
        class BaseNode
        {
        public:
            virtual ~BaseNode(){};
            virtual rtype eval()=0;

            std::vector<BaseNode<rtype>*>& getChildren() {
                return children;
            }

            virtual unsigned int getNumChildren() const =0;

            const std::string& getName() const {return name;}

            unsigned int getID() const {return ID;}

            template<typename T>
            friend std::ostream& operator<< (std::ostream &out, const BaseNode<T>& node);

        protected:
            BaseNode(const std::string& _name, unsigned int _ID) :
                name(_name),
                ID(_ID){}
            std::vector<BaseNode<rtype>*> children;
            const std::string name;
            const unsigned int ID;
        };

        template<typename genome>
        class Node : public BaseNode<typename genome::result_type>
        {
        public:
            Node(genome g, const std::string& _name, unsigned int _ID) :
                BaseNode<typename genome::result_type>(_name, _ID),
                val(g)
            {
            }

            virtual ~Node(){}

            virtual typename genome::result_type eval() override
            {
                std::vector<typename genome::result_type> results;

                std::for_each(this->children.begin(), this->children.end(),
                        [&](BaseNode<typename genome::result_type>* child){results.push_back(child->eval());});

                return unpack_caller<typename genome::result_type>::eval(val, results);
            }

            virtual unsigned int getNumChildren() const override
            {
                return count_args<genome>::value;
            }

        private:
            genome val;
        };

        template<typename genome>
        class Terminator: public BaseNode<typename genome::result_type>
        {
        public:
            static_assert(count_args<genome>::value == 0, "The number of arguments for a terminator must be 0");

            Terminator(genome g, const std::string& _name, unsigned int _ID) :
                BaseNode<typename genome::result_type>(_name, _ID),
                val(g){}

            virtual ~Terminator(){}

            virtual typename genome::result_type eval() override {
                return val();
            }

            virtual unsigned int getNumChildren() const override  {
                return 0;
            }

        private:
            genome val;
        };

        template<typename rType>
        class Tree
        {
        public:
            Tree(BaseNode<rType>* _root) :
                root(_root){}

            rType eval() const {
                return root->eval();
            }

            template<typename T>
            friend std::ostream& operator<< (std::ostream &out, const Tree<T>& tree);

            BaseNode<rType>* root;
        protected:

        };

        template<typename T>
        std::ostream& operator<< (std::ostream &out, const Tree<T>& tree)
        {
            out << *tree.root;
            return out;
        }

        template<typename T>
        std::ostream& operator<< (std::ostream &out, const BaseNode<T>& node)
        {
            out << node.name;

            if (!node.children.empty()) {
                out << "(";
                for (unsigned int i=0; i < node.children.size()-1; ++i)
                    out << *(node.children[i]) << ", ";

                out << *node.children[node.children.size()-1] << ")";
            }
            return out;
        }

        template<typename rType>
        class TreeFactory
        {
        public:
            TreeFactory() :
                currentID(0){}

            template<typename... T>
            void addNode(rType(*f)(T...), const std::string& name)
            {
                static_assert(sizeof...(T) > 0, "Node function with 0 arguments should be terminator");

                std::function<BaseNode<rType>*()> func = [=]() {
                    return new Node<std::function<rType(T...)>>(f, name, currentID);
                };
                nodes[currentID++] = func;
            }

            void addTerminator(rType(*f)(), const std::string& name)
            {
                std::function<BaseNode<rType>*()> func = [=]() {
                    return new Terminator<std::function<rType()>>(f, name, currentID);
                };
                terminators[currentID++] = func;
            }

            Tree<rType>* make()
            {
                assert(!terminators.empty());
                auto tree = new Tree<rType>(createRandomNode());

                for (unsigned int i=0; i < tree->root->getNumChildren(); ++i)
                {
                    BaseNode<rType>* node = createRandomNode();
                    tree->root->getChildren().push_back(node);
                    for (unsigned int i=0; i < node->getNumChildren(); ++i)
                    {
                        node->getChildren().push_back(createRandomTerminator());
                    }
                }

                return tree;
            }

            BaseNode<rType>* copyNode(const BaseNode<rType>* node) {
                if (nodes.find(node->getID()) != nodes.end())
                {
                    return nodes[node->getID()]();
                }
                else
                {
                    return terminators[node->getID()]();
                }
            }

        protected:
            unsigned int currentID;
            std::map<unsigned int, std::function<BaseNode<rType>*()>> nodes;
            std::map<unsigned int, std::function<BaseNode<rType>*()>> terminators;

            BaseNode<rType>* createRandomNode() const {
                auto loc = nodes.begin();
                std::advance(loc, rand() % nodes.size());
                return ((*loc).second)();
            }

            BaseNode<rType>* createRandomTerminator() const {
                auto loc = terminators.begin();
                std::advance(loc, rand() % terminators.size());
                return ((*loc).second)();
            }

        };
    }
}

#endif
<|endoftext|>


<commit_before>//! Libs: -lLLVM-3.1svn -lclangFrontend -lclangParse -lclangDriver -lclangSerialization -lclangSema -lclangEdit -lclangAnalysis -lclangAST -lclangLex -lclangBasic

#include <iostream>
using namespace std;

#include <llvm/Config/config.h>
#include <llvm/Support/Host.h>

#include <clang/Basic/FileManager.h>
#include <clang/Driver/ArgList.h>
#include <clang/Driver/Arg.h>
#include <clang/Driver/Driver.h>
#include <clang/Driver/Compilation.h>
#include <clang/Driver/Action.h>
#include <clang/Driver/Job.h>
#include <clang/Driver/Tool.h>
#include <clang/Lex/Preprocessor.h>
#include <clang/Frontend/FrontendActions.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Frontend/CompilerInvocation.h>
#include <clang/Frontend/Utils.h>
using namespace clang;
using llvm::dyn_cast;
using clang::driver::Arg;
using clang::driver::InputArgList;
using clang::driver::DerivedArgList;
using clang::CompilerInstance;
using llvm::sys::getDefaultTargetTriple;
using llvm::sys::Path;
using llvm::ArrayRef;
using clang::driver::Driver;
using clang::driver::Compilation;
using clang::driver::Command;
using clang::driver::ActionList;
using clang::driver::Action;
using clang::driver::InputAction;
using clang::driver::Job;
using clang::driver::JobList;
using clang::driver::Tool;

static llvm::sys::Path GetExecutablePath(const char *Argv0, bool CanonicalPrefixes);
int ExecuteCompilation(const Driver* that, const Compilation &C,
                               const Command *&FailingCommand);
int ExecuteCommand(const Compilation* that, const Command &C,
                                const Command *&FailingCommand);
int ExecuteJob(const Compilation* that, const Job &J,
                            const Command *&FailingCommand);

int main(int argc, char* argv[])
{
  if (argc < 2) {
    cerr << "Usage: " << argv[0] << " [clang-options] source.cpp" << endl;
    cerr << "Any clang option is usable, source files may be listed among in no particular order." << endl;
    return EXIT_FAILURE;
  }

  CompilerInstance ci;
  ci.createDiagnostics(argc, argv + 1);

  Path path = GetExecutablePath(argv[0], true);
  Driver driver (path.str(), getDefaultTargetTriple(), "a.out", true, ci.getDiagnostics());
  driver.CCCIsCPP = true; // only preprocess
  OwningPtr<Compilation> compil (driver.BuildCompilation(llvm::ArrayRef<const char*>(argv, argc)));
  if (!compil) {
    cerr << "Could not build a Compilation!" << endl;
    return EXIT_FAILURE;
  }

  driver.PrintActions(*compil);

  {
    const ActionList& actions = compil->getActions();
    cout << "Actions: (" << actions.size() << ")" << endl;
    for (ActionList::const_iterator it = actions.begin(), end = actions.end() ; it != end ; ++it) {
      const Action& a = *(*it);
      cout << " - " << a.getClassName() << " inputs(" << a.getInputs().size() << ") -> " << clang::driver::types::getTypeName(a.getType()) << endl;
      for (ActionList::const_iterator it2 = a.begin(), end2 = a.end() ; it2 != end2 ; ++it2) {
        const Action& a2 = *(*it2);
        const InputAction* ia = dyn_cast<const InputAction>(&a2);
        if (ia) {
          cout << "    - " << a2.getClassName() << " {";
          const Arg& inputs = ia->getInputArg();
          for (unsigned i = 0, n = inputs.getNumValues() ; i < n ; ++i) {
            cout << "\"" << inputs.getValue(compil->getArgs(), i) << "\"";
            if (i+1 < n) cout << ", ";
          }
          cout << "} -> " << clang::driver::types::getTypeName(a2.getType()) << endl;
        } else
          cout << "    - " << a2.getClassName() << " inputs(" << a2.getInputs().size() << ") -> " << clang::driver::types::getTypeName(a2.getType()) << endl;
      }
    }
  }

  cout << endl;

  //!\\ The following doesn't quite work
  //!\\ Tip: Use -### to print commands instead of executing them.
  //!\\ ExecuteCompilation exec()s the this same program, with extra arguments.
  //!\\ We should take advantage of the parsing done, and do the preprocess ourselves.

  int Res = 0;
  const Command *FailingCommand = 0;
  Res = ExecuteCompilation(&driver, *compil, FailingCommand);

  // If result status is < 0, then the driver command signalled an error.
  // In this case, generate additional diagnostic information if possible.
  if (Res < 0)
    driver.generateCompilationDiagnostics(*compil, FailingCommand);

  return EXIT_SUCCESS;
}

/// Copied from tools/driver/driver.cpp, where it was a static (local) function.
/// Turns argv[0] into the executable path.
llvm::sys::Path GetExecutablePath(const char *Argv0, bool CanonicalPrefixes)
{
  if (!CanonicalPrefixes)
    return llvm::sys::Path(Argv0);

  // symbol just needs to be some symbol in the binary
  void *symbol = (void*) (intptr_t) GetExecutablePath; // can't take ::main() address, use the current function instead
  return llvm::sys::Path::GetMainExecutable(Argv0, symbol);
}

// Some necessary includes for the following 3 functions
#include <clang/Driver/Options.h>
#include <clang/Driver/DriverDiagnostic.h>
#include <llvm/Support/Program.h>

/// Copied from Driver::ExecuteCompilation().
/// Modified to call the modified ExecuteJob().
int ExecuteCompilation(const Driver* that, const Compilation &C,
                               const Command *&FailingCommand) {
  // Just print if -### was present.
  if (C.getArgs().hasArg(clang::driver::options::OPT__HASH_HASH_HASH)) {
    C.PrintJob(llvm::errs(), C.getJobs(), "\n", true);
    return 0;
  }

  // If there were errors building the compilation, quit now.
  if (that->getDiags().hasErrorOccurred())
    return 1;

  int Res = ExecuteJob(&C, C.getJobs(), FailingCommand);

  // Remove temp files.
  C.CleanupFileList(C.getTempFiles());

  // If the command succeeded, we are done.
  if (Res == 0)
    return Res;

  // Otherwise, remove result files as well.
  if (!C.getArgs().hasArg(clang::driver::options::OPT_save_temps)) {
    C.CleanupFileList(C.getResultFiles(), true);

    // Failure result files are valid unless we crashed.
    if (Res < 0) {
      C.CleanupFileList(C.getFailureResultFiles(), true);
#ifdef _WIN32
      // Exit status should not be negative on Win32,
      // unless abnormal termination.
      Res = 1;
#endif
    }
  }

  // Print extra information about abnormal failures, if possible.
  //
  // This is ad-hoc, but we don't want to be excessively noisy. If the result
  // status was 1, assume the command failed normally. In particular, if it was
  // the compiler then assume it gave a reasonable error code. Failures in other
  // tools are less common, and they generally have worse diagnostics, so always
  // print the diagnostic there.
  const Tool &FailingTool = FailingCommand->getCreator();

  if (!FailingCommand->getCreator().hasGoodDiagnostics() || Res != 1) {
    // FIXME: See FIXME above regarding result code interpretation.
    if (Res < 0)
      that->Diag(clang::diag::err_drv_command_signalled)
        << FailingTool.getShortName();
    else
      that->Diag(clang::diag::err_drv_command_failed)
        << FailingTool.getShortName() << Res;
  }

  return Res;
}

/// Copied from Compilation::ExecuteCommand().
/// Changed to print command options, and not execute anything.
int ExecuteCommand(const Compilation* that, const Command &C,
                                const Command *&FailingCommand) {
  llvm::sys::Path Prog(C.getExecutable());
  const char **Argv = new const char*[C.getArguments().size() + 2];
  Argv[0] = C.getExecutable();
  std::copy(C.getArguments().begin(), C.getArguments().end(), Argv+1);
  Argv[C.getArguments().size() + 1] = 0;

  if ((that->getDriver().CCCEcho || that->getDriver().CCPrintOptions ||
       that->getArgs().hasArg(clang::driver::options::OPT_v)) && !that->getDriver().CCGenDiagnostics) {
    raw_ostream *OS = &llvm::errs();

    // Follow gcc implementation of CC_PRINT_OPTIONS; we could also cache the
    // output stream.
    if (that->getDriver().CCPrintOptions && that->getDriver().CCPrintOptionsFilename) {
      std::string Error;
      OS = new llvm::raw_fd_ostream(that->getDriver().CCPrintOptionsFilename,
                                    Error,
                                    llvm::raw_fd_ostream::F_Append);
      if (!Error.empty()) {
        that->getDriver().Diag(clang::diag::err_drv_cc_print_options_failure)
          << Error;
        FailingCommand = &C;
        delete OS;
        return 1;
      }
    }

    if (that->getDriver().CCPrintOptions)
      *OS << "[Logging clang options]";

    that->PrintJob(*OS, C, "\n", /*Quote=*/that->getDriver().CCPrintOptions);

    if (OS != &llvm::errs())
      delete OS;
  }

  for (const char** i = Argv ; *i ; ++i)
    cout << "\"" << *i << "\" ";
  cout << endl;

  std::string Error;
  int Res = 0;
  // Do not execute
  // Res = llvm::sys::Program::ExecuteAndWait(Prog, Argv,
  //                                    /*env*/0, /*that->Redirects (unfortunately private and inacessible)*/0,
  //                                    /*secondsToWait*/0, /*memoryLimit*/0,
  //                                    &Error);
  if (!Error.empty()) {
    assert(Res && "Error string set with 0 result code!");
    that->getDriver().Diag(clang::diag::err_drv_command_failure) << Error;
  }

  if (Res)
    FailingCommand = &C;

  delete[] Argv;
  return Res;
}

/// Copied from Compilation::ExecuteJob().
/// changed in order to call the modified ExecuteCommand().
int ExecuteJob(const Compilation* that, const Job &J,
                            const Command *&FailingCommand) {
  if (const Command *C = dyn_cast<Command>(&J)) {
    return ExecuteCommand(that, *C, FailingCommand);
  } else {
    const JobList *Jobs = cast<JobList>(&J);
    for (JobList::const_iterator
           it = Jobs->begin(), ie = Jobs->end(); it != ie; ++it) {
      if (int Res = ExecuteJob(that, **it, FailingCommand))
        return Res;
    }
    return 0;
  }
}
<commit_msg>Use the same process to run Commands<commit_after>//! Libs: -lLLVM-3.1svn -lclangFrontend -lclangParse -lclangDriver -lclangSerialization -lclangSema -lclangEdit -lclangAnalysis -lclangAST -lclangLex -lclangBasic

#include <iostream>
using namespace std;

#include <llvm/Config/config.h>
#include <llvm/Support/Host.h>

#include <clang/Basic/FileManager.h>
#include <clang/Driver/ArgList.h>
#include <clang/Driver/Arg.h>
#include <clang/Driver/Driver.h>
#include <clang/Driver/Compilation.h>
#include <clang/Driver/Action.h>
#include <clang/Driver/Job.h>
#include <clang/Driver/Tool.h>
#include <clang/Driver/ToolChain.h>
#include <clang/Lex/Preprocessor.h>
#include <clang/Frontend/FrontendActions.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Frontend/CompilerInvocation.h>
#include <clang/Frontend/Utils.h>
using namespace clang;
using llvm::dyn_cast;
using clang::driver::Arg;
using clang::driver::InputArgList;
using clang::driver::DerivedArgList;
using clang::CompilerInstance;
using llvm::sys::getDefaultTargetTriple;
using llvm::sys::Path;
using llvm::ArrayRef;
using clang::driver::Driver;
using clang::driver::Compilation;
using clang::driver::Command;
using clang::driver::ArgStringList;
using clang::driver::ActionList;
using clang::driver::Action;
using clang::driver::InputAction;
using clang::driver::PreprocessJobAction;
using clang::driver::Job;
using clang::driver::JobAction;
using clang::driver::JobList;
using clang::driver::Tool;
using clang::driver::ToolChain;

int parseArgsAndProceed(int argc, const char* argv[]);

static llvm::sys::Path GetExecutablePath(const char *Argv0, bool CanonicalPrefixes);
int ExecuteCompilation(const Driver* that, const Compilation &C,
                               const Command *&FailingCommand);
int ExecuteCommand(const Compilation* that, const Command &C,
                                const Command *&FailingCommand);
int ExecuteJob(const Compilation* that, const Job &J,
                            const Command *&FailingCommand);

int main(int argc, const char* argv[])
{
  if (argc < 2) {
    cerr << "Usage: " << argv[0] << " [clang-options] source.cpp" << endl;
    cerr << "Any clang option is usable, source files may be listed among in no particular order." << endl;
    return EXIT_FAILURE;
  }

  return parseArgsAndProceed(argc, argv);
}

int parseArgsAndProceed(int argc, const char* argv[]) {
  CompilerInstance ci;
  ci.createDiagnostics(argc, argv);

  Path path = GetExecutablePath(argv[0], true);
  Driver driver (path.str(), getDefaultTargetTriple(), "a.out", true, ci.getDiagnostics());
  driver.CCCIsCPP = true; // only preprocess
  OwningPtr<Compilation> compil (driver.BuildCompilation(llvm::ArrayRef<const char*>(argv, argc)));
  if (!compil) {
    cerr << "Cannot build the compilation." << endl;
    return EXIT_FAILURE;
  }
  if (driver.getDiags().hasErrorOccurred()) {
    cerr << "Error occurred building the compilation." << endl;
    // We must provide a Command that failed to use Driver::generateCompilationDiagnostics()
    // That Command will only be used as follows: ``if (FailingCommand->getCreator().isLinkJob()) return;''
    // Hence we'll have to create a fake command with a tool that will answer no to isLinkJob(),
    // As we only have PreprocessorJobActions, the first generated action will do.
    JobAction& fakeAction = *dyn_cast<JobAction>(*compil->getActions().begin());
    Tool& fakeCreator = compil->getDefaultToolChain().SelectTool(*compil, fakeAction, fakeAction.getInputs());
    Command fakeCommand (fakeAction, fakeCreator, /*Executable*/NULL, ArgStringList());
    cerr << "DIAGNOSTIC:" << endl;
    driver.generateCompilationDiagnostics(*compil, &fakeCommand);
    cerr << "END OF DIAGNOSTIC" << endl;
    return EXIT_FAILURE;
  }

  driver.PrintActions(*compil);

  {
    const ActionList& actions = compil->getActions();
    cout << "Actions: (" << actions.size() << ")" << endl;
    for (ActionList::const_iterator it = actions.begin(), end = actions.end() ; it != end ; ++it) {
      const Action& a = *(*it);
      cout << " - " << a.getClassName() << " inputs(" << a.getInputs().size() << ") -> " << clang::driver::types::getTypeName(a.getType()) << endl;
      for (ActionList::const_iterator it2 = a.begin(), end2 = a.end() ; it2 != end2 ; ++it2) {
        const Action& a2 = *(*it2);
        const InputAction* ia = dyn_cast<const InputAction>(&a2);
        if (ia) {
          cout << "    - " << a2.getClassName() << " {";
          const Arg& inputs = ia->getInputArg();
          for (unsigned i = 0, n = inputs.getNumValues() ; i < n ; ++i) {
            cout << "\"" << inputs.getValue(compil->getArgs(), i) << "\"";
            if (i+1 < n) cout << ", ";
          }
          cout << "} -> " << clang::driver::types::getTypeName(a2.getType()) << endl;
        } else
          cout << "    - " << a2.getClassName() << " inputs(" << a2.getInputs().size() << ") -> " << clang::driver::types::getTypeName(a2.getType()) << endl;
      }
    }
  }

  cout << endl;

  //!\\ The following doesn't quite work
  //!\\ Tip: Use -### to print commands instead of executing them.
  //!\\ ExecuteCompilation exec()s the this same program, with extra arguments.
  //!\\ We should take advantage of the parsing done, and do the preprocess ourselves.

  int Res = 0;
  const Command *FailingCommand = 0;
  Res = ExecuteCompilation(&driver, *compil, FailingCommand);

  // If result status is < 0, then the driver command signalled an error.
  // In this case, generate additional diagnostic information if possible.
  if (Res < 0)
    driver.generateCompilationDiagnostics(*compil, FailingCommand);

  return EXIT_SUCCESS;
}

/// Copied from tools/driver/driver.cpp, where it was a static (local) function.
/// Turns argv[0] into the executable path.
llvm::sys::Path GetExecutablePath(const char *Argv0, bool CanonicalPrefixes)
{
  if (!CanonicalPrefixes)
    return llvm::sys::Path(Argv0);

  // symbol just needs to be some symbol in the binary
  void *symbol = (void*) (intptr_t) GetExecutablePath; // can't take ::main() address, use the current function instead
  return llvm::sys::Path::GetMainExecutable(Argv0, symbol);
}

// Some necessary includes for the following 3 functions
#include <clang/Driver/Options.h>
#include <clang/Driver/DriverDiagnostic.h>
#include <llvm/Support/Program.h>

/// Copied from Driver::ExecuteCompilation().
/// Modified to call the modified ExecuteJob().
int ExecuteCompilation(const Driver* that, const Compilation &C,
                               const Command *&FailingCommand) {
  // Just print if -### was present.
  if (C.getArgs().hasArg(clang::driver::options::OPT__HASH_HASH_HASH)) {
    C.PrintJob(llvm::errs(), C.getJobs(), "\n", true);
    return 0;
  }

  // If there were errors building the compilation, quit now.
  if (that->getDiags().hasErrorOccurred())
    return 1;

  int Res = ExecuteJob(&C, C.getJobs(), FailingCommand);

  // If the command succeeded, we are done.
  if (Res == 0)
    return Res;

  // Print extra information about abnormal failures, if possible.
  //
  // This is ad-hoc, but we don't want to be excessively noisy. If the result
  // status was 1, assume the command failed normally. In particular, if it was
  // the compiler then assume it gave a reasonable error code. Failures in other
  // tools are less common, and they generally have worse diagnostics, so always
  // print the diagnostic there.
  const Tool &FailingTool = FailingCommand->getCreator();

  if (!FailingCommand->getCreator().hasGoodDiagnostics() || Res != 1) {
    // FIXME: See FIXME above regarding result code interpretation.
    if (Res < 0)
      that->Diag(clang::diag::err_drv_command_signalled)
        << FailingTool.getShortName();
    else
      that->Diag(clang::diag::err_drv_command_failed)
        << FailingTool.getShortName() << Res;
  }

  return Res;
}

/// Copied from Compilation::ExecuteCommand().
/// Changed to print command options, and not execute anything.
int ExecuteCommand(const Compilation* that, const Command &C,
                                const Command *&FailingCommand) {
  if ((that->getDriver().CCCEcho || that->getDriver().CCPrintOptions ||
       that->getArgs().hasArg(clang::driver::options::OPT_v)) && !that->getDriver().CCGenDiagnostics) {
    raw_ostream *OS = &llvm::errs();

    // Follow gcc implementation of CC_PRINT_OPTIONS; we could also cache the
    // output stream.
    if (that->getDriver().CCPrintOptions && that->getDriver().CCPrintOptionsFilename) {
      std::string Error;
      OS = new llvm::raw_fd_ostream(that->getDriver().CCPrintOptionsFilename,
                                    Error,
                                    llvm::raw_fd_ostream::F_Append);
      if (!Error.empty()) {
        that->getDriver().Diag(clang::diag::err_drv_cc_print_options_failure)
          << Error;
        FailingCommand = &C;
        delete OS;
        return 1;
      }
    }

    if (that->getDriver().CCPrintOptions)
      *OS << "[Logging clang options]";

    that->PrintJob(*OS, C, "\n", /*Quote=*/that->getDriver().CCPrintOptions);

    if (OS != &llvm::errs())
      delete OS;
  }

  int argc = C.getArguments().size();
  const char* *argv = new const char* [argc+1];
  argv[0] = C.getExecutable();
  std::copy(C.getArguments().begin(), C.getArguments().end(), argv+1);
  int Res = 0;
  Res = parseArgsAndProceed(argc, argv);

  if (Res)
    FailingCommand = &C;

  return Res;
}

/// Copied from Compilation::ExecuteJob().
/// changed in order to call the modified ExecuteCommand().
int ExecuteJob(const Compilation* that, const Job &J,
                            const Command *&FailingCommand) {
  if (const Command *C = dyn_cast<Command>(&J)) {
    return ExecuteCommand(that, *C, FailingCommand);
  } else {
    const JobList *Jobs = cast<JobList>(&J);
    for (JobList::const_iterator
           it = Jobs->begin(), ie = Jobs->end(); it != ie; ++it) {
      if (int Res = ExecuteJob(that, **it, FailingCommand))
        return Res;
    }
    return 0;
  }
}
<|endoftext|>