<commit_before>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "proto_rpc_adapter.h"
#include "searchapi.h"
#include "docsumapi.h"
#include "monitorapi.h"
#include <vespa/fnet/frt/rpcrequest.h>
#include <vespa/fnet/frt/supervisor.h>
#include <vespa/vespalib/util/compressor.h>
#include <vespa/searchlib/util/slime_output_raw_buf_adapter.h>
#include <vespa/vespalib/data/databuffer.h>
#include <vespa/searchlib/common/packets.h>

#include <vespa/log/log.h>
LOG_SETUP(".engine.proto_rpc_adapter");

namespace search::engine {

using vespalib::DataBuffer;
using vespalib::ConstBufferRef;
using vespalib::compression::CompressionConfig;
using ProtoSearchRequest = ProtoConverter::ProtoSearchRequest;
using ProtoSearchReply = ProtoConverter::ProtoSearchReply;
using ProtoDocsumRequest = ProtoConverter::ProtoDocsumRequest;
using ProtoDocsumReply = ProtoConverter::ProtoDocsumReply;
using ProtoMonitorRequest = ProtoConverter::ProtoMonitorRequest;
using ProtoMonitorReply = ProtoConverter::ProtoMonitorReply;
using QueryStats = SearchProtocolMetrics::QueryStats;
using DocsumStats = SearchProtocolMetrics::DocsumStats;

namespace {

CompressionConfig get_compression_config() {
    using search::fs4transport::FS4PersistentPacketStreamer;
    const FS4PersistentPacketStreamer & streamer = FS4PersistentPacketStreamer::Instance;
    return CompressionConfig(streamer.getCompressionType(), streamer.getCompressionLevel(), 80, streamer.getCompressionLimit());
}

template <typename MSG>
void encode_message(const MSG &src, FRT_Values &dst) {
    using vespalib::compression::compress;
    auto output = src.SerializeAsString();
    ConstBufferRef buf(output.data(), output.size());
    DataBuffer compressed(output.data(), output.size());
    CompressionConfig::Type type = compress(get_compression_config(), buf, compressed, true);
    dst.AddInt8(type);
    dst.AddInt32(buf.size());
    dst.AddData(compressed.getData(), compressed.getDataLen());
}

void encode_search_reply(const ProtoSearchReply &src, FRT_Values &dst) {
    using vespalib::compression::compress;
    auto output = src.SerializeAsString();
    if (src.grouping_blob().empty()) {
        dst.AddInt8(CompressionConfig::Type::NONE);
        dst.AddInt32(output.size());
        dst.AddData(output.data(), output.size());
    } else {
        ConstBufferRef buf(output.data(), output.size());
        DataBuffer compressed(output.data(), output.size());
        CompressionConfig::Type type = compress(get_compression_config(), buf, compressed, true);
        dst.AddInt8(type);
        dst.AddInt32(buf.size());
        dst.AddData(compressed.getData(), compressed.getDataLen());
    }
}

template <typename MSG>
bool decode_message(const FRT_Values &src, MSG &dst) {
    using vespalib::compression::decompress;
    uint8_t encoding = src[0]._intval8;
    uint32_t uncompressed_size = src[1]._intval32;
    DataBuffer uncompressed(src[2]._data._buf, src[2]._data._len);
    ConstBufferRef blob(src[2]._data._buf, src[2]._data._len);
    decompress(CompressionConfig::toType(encoding), uncompressed_size, blob, uncompressed, true);
    assert(uncompressed_size == uncompressed.getDataLen());
    return dst.ParseFromArray(uncompressed.getData(), uncompressed.getDataLen());
}

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

struct SearchRequestDecoder : SearchRequest::Source::Decoder {
    FRT_RPCRequest &rpc; // valid until Return is called
    QueryStats &stats;
    RelativeTime relative_time;
    SearchRequestDecoder(FRT_RPCRequest &rpc_in, QueryStats &stats_in)
        : rpc(rpc_in), stats(stats_in), relative_time(std::make_unique<SteadyClock>()) {}
    std::unique_ptr<SearchRequest> decode() override {
        ProtoSearchRequest msg;
        stats.request_size = (*rpc.GetParams())[2]._data._len;
        if (!decode_message(*rpc.GetParams(), msg)) {
            LOG(warning, "got bad protobuf search request over rpc (unable to decode)");
            return std::unique_ptr<SearchRequest>(nullptr);
        }
        auto req = std::make_unique<SearchRequest>(std::move(relative_time));
        ProtoConverter::search_request_from_proto(msg, *req);
        return req;
    }
};

std::unique_ptr<SearchRequest::Source::Decoder> search_request_decoder(FRT_RPCRequest &rpc, QueryStats &stats) {
    return std::make_unique<SearchRequestDecoder>(rpc, stats);
}

// allocated in the stash of the request it is completing; no self-delete needed
struct SearchCompletionHandler : SearchClient {
    FRT_RPCRequest &req;
    SearchProtocolMetrics &metrics;
    QueryStats stats;
    SearchCompletionHandler(FRT_RPCRequest &req_in, SearchProtocolMetrics &metrics_in)
        : req(req_in), metrics(metrics_in), stats() {}
    void searchDone(SearchReply::UP reply) override {
        ProtoSearchReply msg;
        ProtoConverter::search_reply_to_proto(*reply, msg);
        encode_search_reply(msg, *req.GetReturn());
        stats.reply_size = (*req.GetReturn())[2]._data._len;
        if (reply->request) {
            stats.latency = vespalib::to_s(reply->request->getTimeUsed());
            metrics.update_query_metrics(stats);
        }
        req.Return();
    }
};

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

struct DocsumRequestDecoder : DocsumRequest::Source::Decoder {
    FRT_RPCRequest &rpc; // valid until Return is called
    DocsumStats &stats;
    RelativeTime relative_time;
    DocsumRequestDecoder(FRT_RPCRequest &rpc_in, DocsumStats &stats_in)
        : rpc(rpc_in), stats(stats_in), relative_time(std::make_unique<SteadyClock>()) {}
    std::unique_ptr<DocsumRequest> decode() override {
        ProtoDocsumRequest msg;
        stats.request_size = (*rpc.GetParams())[2]._data._len;
        if (!decode_message(*rpc.GetParams(), msg)) {
            LOG(warning, "got bad protobuf docsum request over rpc (unable to decode)");
            return std::unique_ptr<DocsumRequest>(nullptr);
        }
        stats.requested_documents = msg.global_ids_size();
        auto req = std::make_unique<DocsumRequest>(std::move(relative_time));
        ProtoConverter::docsum_request_from_proto(msg, *req);
        return req;
    }
};

std::unique_ptr<DocsumRequest::Source::Decoder> docsum_request_decoder(FRT_RPCRequest &rpc, DocsumStats &stats) {
    return std::make_unique<DocsumRequestDecoder>(rpc, stats);
}

// allocated in the stash of the request it is completing; no self-delete needed
struct GetDocsumsCompletionHandler : DocsumClient {
    FRT_RPCRequest &req;
    SearchProtocolMetrics &metrics;
    DocsumStats stats;
    GetDocsumsCompletionHandler(FRT_RPCRequest &req_in, SearchProtocolMetrics &metrics_in)
        : req(req_in), metrics(metrics_in), stats() {}
    void getDocsumsDone(DocsumReply::UP reply) override {
        ProtoDocsumReply msg;
        ProtoConverter::docsum_reply_to_proto(*reply, msg);
        encode_message(msg, *req.GetReturn());
        stats.reply_size = (*req.GetReturn())[2]._data._len;
        if (reply->hasRequest()) {
            stats.latency = vespalib::to_s(reply->request().getTimeUsed());
            metrics.update_docsum_metrics(stats);
        }
        req.Return();
    }
};

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

// allocated in the stash of the request it is completing; no self-delete needed
struct PingCompletionHandler : MonitorClient {
    FRT_RPCRequest &req;
    PingCompletionHandler(FRT_RPCRequest &req_in) : req(req_in) {}
    void pingDone(std::unique_ptr<MonitorReply> reply) override {
        ProtoMonitorReply msg;
        ProtoConverter::monitor_reply_to_proto(*reply, msg);
        encode_message(msg, *req.GetReturn());
        req.Return();
    }
};

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

void describe_bix_param_return(FRT_ReflectionBuilder &rb) {
    rb.ParamDesc("encoding", "0=raw, 6=lz4, 7=zstd");
    rb.ParamDesc("uncompressed_size", "uncompressed size of serialized request");
    rb.ParamDesc("request", "possibly compressed serialized request");
    rb.ReturnDesc("encoding",  "0=raw, 6=lz4, 7=zstd");
    rb.ReturnDesc("uncompressed_size", "uncompressed size of serialized reply");
    rb.ReturnDesc("reply", "possibly compressed serialized reply");
}

}

ProtoRpcAdapter::ProtoRpcAdapter(SearchServer &search_server,
                                 DocsumServer &docsum_server,
                                 MonitorServer &monitor_server,
                                 FRT_Supervisor &orb)
    : _search_server(search_server),
      _docsum_server(docsum_server),
      _monitor_server(monitor_server),
      _online(false),
      _metrics()
{
    FRT_ReflectionBuilder rb(&orb);
    //-------------------------------------------------------------------------
    rb.DefineMethod("vespa.searchprotocol.search", "bix", "bix",
                    FRT_METHOD(ProtoRpcAdapter::rpc_search), this);
    rb.MethodDesc("perform a search against this back-end");
    describe_bix_param_return(rb);
    //-------------------------------------------------------------------------
    rb.DefineMethod("vespa.searchprotocol.getDocsums", "bix", "bix",
                    FRT_METHOD(ProtoRpcAdapter::rpc_getDocsums), this);
    rb.MethodDesc("fetch document summaries from this back-end");
    describe_bix_param_return(rb);
    //-------------------------------------------------------------------------
    rb.DefineMethod("vespa.searchprotocol.ping", "bix", "bix",
                    FRT_METHOD(ProtoRpcAdapter::rpc_ping), this);
    rb.MethodDesc("ping this back-end");
    describe_bix_param_return(rb);
    //-------------------------------------------------------------------------
}

void
ProtoRpcAdapter::rpc_search(FRT_RPCRequest *req)
{
    if (!is_online()) {
        return req->SetError(FRTE_RPC_METHOD_FAILED, "Server not online");
    }
    req->Detach();
    auto &client = req->getStash().create<SearchCompletionHandler>(*req, _metrics);
    auto reply = _search_server.search(search_request_decoder(*req, client.stats), client);
    if (reply) {
        client.searchDone(std::move(reply));
    }
}

void
ProtoRpcAdapter::rpc_getDocsums(FRT_RPCRequest *req)
{
    if (!is_online()) {
        return req->SetError(FRTE_RPC_METHOD_FAILED, "Server not online");
    }
    req->Detach();
    auto &client = req->getStash().create<GetDocsumsCompletionHandler>(*req, _metrics);
    auto reply = _docsum_server.getDocsums(docsum_request_decoder(*req, client.stats), client);
    if (reply) {
        client.getDocsumsDone(std::move(reply));
    }
}

void
ProtoRpcAdapter::rpc_ping(FRT_RPCRequest *rpc)
{
    if (!is_online()) {
        return rpc->SetError(FRTE_RPC_METHOD_FAILED, "Server not online");
    }
    rpc->Detach();
    ProtoMonitorRequest msg;
    if (decode_message(*rpc->GetParams(), msg)) {
        auto req = std::make_unique<MonitorRequest>();
        ProtoConverter::monitor_request_from_proto(msg, *req);
        auto &client = rpc->getStash().create<PingCompletionHandler>(*rpc);
        auto reply = _monitor_server.ping(std::move(req), client);
        if (reply) {
            client.pingDone(std::move(reply));
        }
    } else {
        LOG(warning, "got bad protobuf monitor request over rpc (unable to decode)");
        rpc->SetError(FRTE_RPC_METHOD_FAILED, "malformed monitor request");
        rpc->Return();
    }
}

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

void
ProtoRpcAdapter::encode_search_request(const ProtoSearchRequest &src, FRT_RPCRequest &dst)
{
    dst.SetMethodName("vespa.searchprotocol.search");
    encode_message(src, *dst.GetParams());
}

bool
ProtoRpcAdapter::decode_search_reply(FRT_RPCRequest &src, ProtoSearchReply &dst)
{
    return (src.CheckReturnTypes("bix") && decode_message(*src.GetReturn(), dst));
}

void
ProtoRpcAdapter::encode_docsum_request(const ProtoDocsumRequest &src, FRT_RPCRequest &dst)
{
    dst.SetMethodName("vespa.searchprotocol.getDocsums");
    encode_message(src, *dst.GetParams());
}

bool
ProtoRpcAdapter::decode_docsum_reply(FRT_RPCRequest &src, ProtoDocsumReply &dst)
{
    return (src.CheckReturnTypes("bix") && decode_message(*src.GetReturn(), dst));
}

void
ProtoRpcAdapter::encode_monitor_request(const ProtoMonitorRequest &src, FRT_RPCRequest &dst)
{
    dst.SetMethodName("vespa.searchprotocol.ping");
    encode_message(src, *dst.GetParams());
}

bool
ProtoRpcAdapter::decode_monitor_reply(FRT_RPCRequest &src, ProtoMonitorReply &dst)
{
    return (src.CheckReturnTypes("bix") && decode_message(*src.GetReturn(), dst));
}

}
<commit_msg>Add request capability filters to search API functions<commit_after>// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "proto_rpc_adapter.h"
#include "searchapi.h"
#include "docsumapi.h"
#include "monitorapi.h"
#include <vespa/fnet/frt/require_capabilities.h>
#include <vespa/fnet/frt/rpcrequest.h>
#include <vespa/fnet/frt/supervisor.h>
#include <vespa/vespalib/util/compressor.h>
#include <vespa/searchlib/util/slime_output_raw_buf_adapter.h>
#include <vespa/vespalib/data/databuffer.h>
#include <vespa/searchlib/common/packets.h>

#include <vespa/log/log.h>
LOG_SETUP(".engine.proto_rpc_adapter");

namespace search::engine {

using vespalib::DataBuffer;
using vespalib::ConstBufferRef;
using vespalib::compression::CompressionConfig;
using ProtoSearchRequest = ProtoConverter::ProtoSearchRequest;
using ProtoSearchReply = ProtoConverter::ProtoSearchReply;
using ProtoDocsumRequest = ProtoConverter::ProtoDocsumRequest;
using ProtoDocsumReply = ProtoConverter::ProtoDocsumReply;
using ProtoMonitorRequest = ProtoConverter::ProtoMonitorRequest;
using ProtoMonitorReply = ProtoConverter::ProtoMonitorReply;
using QueryStats = SearchProtocolMetrics::QueryStats;
using DocsumStats = SearchProtocolMetrics::DocsumStats;

namespace {

CompressionConfig get_compression_config() {
    using search::fs4transport::FS4PersistentPacketStreamer;
    const FS4PersistentPacketStreamer & streamer = FS4PersistentPacketStreamer::Instance;
    return CompressionConfig(streamer.getCompressionType(), streamer.getCompressionLevel(), 80, streamer.getCompressionLimit());
}

template <typename MSG>
void encode_message(const MSG &src, FRT_Values &dst) {
    using vespalib::compression::compress;
    auto output = src.SerializeAsString();
    ConstBufferRef buf(output.data(), output.size());
    DataBuffer compressed(output.data(), output.size());
    CompressionConfig::Type type = compress(get_compression_config(), buf, compressed, true);
    dst.AddInt8(type);
    dst.AddInt32(buf.size());
    dst.AddData(compressed.getData(), compressed.getDataLen());
}

void encode_search_reply(const ProtoSearchReply &src, FRT_Values &dst) {
    using vespalib::compression::compress;
    auto output = src.SerializeAsString();
    if (src.grouping_blob().empty()) {
        dst.AddInt8(CompressionConfig::Type::NONE);
        dst.AddInt32(output.size());
        dst.AddData(output.data(), output.size());
    } else {
        ConstBufferRef buf(output.data(), output.size());
        DataBuffer compressed(output.data(), output.size());
        CompressionConfig::Type type = compress(get_compression_config(), buf, compressed, true);
        dst.AddInt8(type);
        dst.AddInt32(buf.size());
        dst.AddData(compressed.getData(), compressed.getDataLen());
    }
}

template <typename MSG>
bool decode_message(const FRT_Values &src, MSG &dst) {
    using vespalib::compression::decompress;
    uint8_t encoding = src[0]._intval8;
    uint32_t uncompressed_size = src[1]._intval32;
    DataBuffer uncompressed(src[2]._data._buf, src[2]._data._len);
    ConstBufferRef blob(src[2]._data._buf, src[2]._data._len);
    decompress(CompressionConfig::toType(encoding), uncompressed_size, blob, uncompressed, true);
    assert(uncompressed_size == uncompressed.getDataLen());
    return dst.ParseFromArray(uncompressed.getData(), uncompressed.getDataLen());
}

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

struct SearchRequestDecoder : SearchRequest::Source::Decoder {
    FRT_RPCRequest &rpc; // valid until Return is called
    QueryStats &stats;
    RelativeTime relative_time;
    SearchRequestDecoder(FRT_RPCRequest &rpc_in, QueryStats &stats_in)
        : rpc(rpc_in), stats(stats_in), relative_time(std::make_unique<SteadyClock>()) {}
    std::unique_ptr<SearchRequest> decode() override {
        ProtoSearchRequest msg;
        stats.request_size = (*rpc.GetParams())[2]._data._len;
        if (!decode_message(*rpc.GetParams(), msg)) {
            LOG(warning, "got bad protobuf search request over rpc (unable to decode)");
            return std::unique_ptr<SearchRequest>(nullptr);
        }
        auto req = std::make_unique<SearchRequest>(std::move(relative_time));
        ProtoConverter::search_request_from_proto(msg, *req);
        return req;
    }
};

std::unique_ptr<SearchRequest::Source::Decoder> search_request_decoder(FRT_RPCRequest &rpc, QueryStats &stats) {
    return std::make_unique<SearchRequestDecoder>(rpc, stats);
}

// allocated in the stash of the request it is completing; no self-delete needed
struct SearchCompletionHandler : SearchClient {
    FRT_RPCRequest &req;
    SearchProtocolMetrics &metrics;
    QueryStats stats;
    SearchCompletionHandler(FRT_RPCRequest &req_in, SearchProtocolMetrics &metrics_in)
        : req(req_in), metrics(metrics_in), stats() {}
    void searchDone(SearchReply::UP reply) override {
        ProtoSearchReply msg;
        ProtoConverter::search_reply_to_proto(*reply, msg);
        encode_search_reply(msg, *req.GetReturn());
        stats.reply_size = (*req.GetReturn())[2]._data._len;
        if (reply->request) {
            stats.latency = vespalib::to_s(reply->request->getTimeUsed());
            metrics.update_query_metrics(stats);
        }
        req.Return();
    }
};

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

struct DocsumRequestDecoder : DocsumRequest::Source::Decoder {
    FRT_RPCRequest &rpc; // valid until Return is called
    DocsumStats &stats;
    RelativeTime relative_time;
    DocsumRequestDecoder(FRT_RPCRequest &rpc_in, DocsumStats &stats_in)
        : rpc(rpc_in), stats(stats_in), relative_time(std::make_unique<SteadyClock>()) {}
    std::unique_ptr<DocsumRequest> decode() override {
        ProtoDocsumRequest msg;
        stats.request_size = (*rpc.GetParams())[2]._data._len;
        if (!decode_message(*rpc.GetParams(), msg)) {
            LOG(warning, "got bad protobuf docsum request over rpc (unable to decode)");
            return std::unique_ptr<DocsumRequest>(nullptr);
        }
        stats.requested_documents = msg.global_ids_size();
        auto req = std::make_unique<DocsumRequest>(std::move(relative_time));
        ProtoConverter::docsum_request_from_proto(msg, *req);
        return req;
    }
};

std::unique_ptr<DocsumRequest::Source::Decoder> docsum_request_decoder(FRT_RPCRequest &rpc, DocsumStats &stats) {
    return std::make_unique<DocsumRequestDecoder>(rpc, stats);
}

// allocated in the stash of the request it is completing; no self-delete needed
struct GetDocsumsCompletionHandler : DocsumClient {
    FRT_RPCRequest &req;
    SearchProtocolMetrics &metrics;
    DocsumStats stats;
    GetDocsumsCompletionHandler(FRT_RPCRequest &req_in, SearchProtocolMetrics &metrics_in)
        : req(req_in), metrics(metrics_in), stats() {}
    void getDocsumsDone(DocsumReply::UP reply) override {
        ProtoDocsumReply msg;
        ProtoConverter::docsum_reply_to_proto(*reply, msg);
        encode_message(msg, *req.GetReturn());
        stats.reply_size = (*req.GetReturn())[2]._data._len;
        if (reply->hasRequest()) {
            stats.latency = vespalib::to_s(reply->request().getTimeUsed());
            metrics.update_docsum_metrics(stats);
        }
        req.Return();
    }
};

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

// allocated in the stash of the request it is completing; no self-delete needed
struct PingCompletionHandler : MonitorClient {
    FRT_RPCRequest &req;
    PingCompletionHandler(FRT_RPCRequest &req_in) : req(req_in) {}
    void pingDone(std::unique_ptr<MonitorReply> reply) override {
        ProtoMonitorReply msg;
        ProtoConverter::monitor_reply_to_proto(*reply, msg);
        encode_message(msg, *req.GetReturn());
        req.Return();
    }
};

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

void describe_bix_param_return(FRT_ReflectionBuilder &rb) {
    rb.ParamDesc("encoding", "0=raw, 6=lz4, 7=zstd");
    rb.ParamDesc("uncompressed_size", "uncompressed size of serialized request");
    rb.ParamDesc("request", "possibly compressed serialized request");
    rb.ReturnDesc("encoding",  "0=raw, 6=lz4, 7=zstd");
    rb.ReturnDesc("uncompressed_size", "uncompressed size of serialized reply");
    rb.ReturnDesc("reply", "possibly compressed serialized reply");
}

std::unique_ptr<FRT_RequireCapabilities> make_search_api_capability_filter() {
    return FRT_RequireCapabilities::of(vespalib::net::tls::Capability::content_search_api());
}

}

ProtoRpcAdapter::ProtoRpcAdapter(SearchServer &search_server,
                                 DocsumServer &docsum_server,
                                 MonitorServer &monitor_server,
                                 FRT_Supervisor &orb)
    : _search_server(search_server),
      _docsum_server(docsum_server),
      _monitor_server(monitor_server),
      _online(false),
      _metrics()
{
    FRT_ReflectionBuilder rb(&orb);
    //-------------------------------------------------------------------------
    rb.DefineMethod("vespa.searchprotocol.search", "bix", "bix",
                    FRT_METHOD(ProtoRpcAdapter::rpc_search), this);
    rb.MethodDesc("perform a search against this back-end");
    rb.RequestAccessFilter(make_search_api_capability_filter());
    describe_bix_param_return(rb);
    //-------------------------------------------------------------------------
    rb.DefineMethod("vespa.searchprotocol.getDocsums", "bix", "bix",
                    FRT_METHOD(ProtoRpcAdapter::rpc_getDocsums), this);
    rb.MethodDesc("fetch document summaries from this back-end");
    rb.RequestAccessFilter(make_search_api_capability_filter());
    describe_bix_param_return(rb);
    //-------------------------------------------------------------------------
    rb.DefineMethod("vespa.searchprotocol.ping", "bix", "bix",
                    FRT_METHOD(ProtoRpcAdapter::rpc_ping), this);
    rb.MethodDesc("ping this back-end");
    rb.RequestAccessFilter(make_search_api_capability_filter());
    describe_bix_param_return(rb);
    //-------------------------------------------------------------------------
}

void
ProtoRpcAdapter::rpc_search(FRT_RPCRequest *req)
{
    if (!is_online()) {
        return req->SetError(FRTE_RPC_METHOD_FAILED, "Server not online");
    }
    req->Detach();
    auto &client = req->getStash().create<SearchCompletionHandler>(*req, _metrics);
    auto reply = _search_server.search(search_request_decoder(*req, client.stats), client);
    if (reply) {
        client.searchDone(std::move(reply));
    }
}

void
ProtoRpcAdapter::rpc_getDocsums(FRT_RPCRequest *req)
{
    if (!is_online()) {
        return req->SetError(FRTE_RPC_METHOD_FAILED, "Server not online");
    }
    req->Detach();
    auto &client = req->getStash().create<GetDocsumsCompletionHandler>(*req, _metrics);
    auto reply = _docsum_server.getDocsums(docsum_request_decoder(*req, client.stats), client);
    if (reply) {
        client.getDocsumsDone(std::move(reply));
    }
}

void
ProtoRpcAdapter::rpc_ping(FRT_RPCRequest *rpc)
{
    if (!is_online()) {
        return rpc->SetError(FRTE_RPC_METHOD_FAILED, "Server not online");
    }
    rpc->Detach();
    ProtoMonitorRequest msg;
    if (decode_message(*rpc->GetParams(), msg)) {
        auto req = std::make_unique<MonitorRequest>();
        ProtoConverter::monitor_request_from_proto(msg, *req);
        auto &client = rpc->getStash().create<PingCompletionHandler>(*rpc);
        auto reply = _monitor_server.ping(std::move(req), client);
        if (reply) {
            client.pingDone(std::move(reply));
        }
    } else {
        LOG(warning, "got bad protobuf monitor request over rpc (unable to decode)");
        rpc->SetError(FRTE_RPC_METHOD_FAILED, "malformed monitor request");
        rpc->Return();
    }
}

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

void
ProtoRpcAdapter::encode_search_request(const ProtoSearchRequest &src, FRT_RPCRequest &dst)
{
    dst.SetMethodName("vespa.searchprotocol.search");
    encode_message(src, *dst.GetParams());
}

bool
ProtoRpcAdapter::decode_search_reply(FRT_RPCRequest &src, ProtoSearchReply &dst)
{
    return (src.CheckReturnTypes("bix") && decode_message(*src.GetReturn(), dst));
}

void
ProtoRpcAdapter::encode_docsum_request(const ProtoDocsumRequest &src, FRT_RPCRequest &dst)
{
    dst.SetMethodName("vespa.searchprotocol.getDocsums");
    encode_message(src, *dst.GetParams());
}

bool
ProtoRpcAdapter::decode_docsum_reply(FRT_RPCRequest &src, ProtoDocsumReply &dst)
{
    return (src.CheckReturnTypes("bix") && decode_message(*src.GetReturn(), dst));
}

void
ProtoRpcAdapter::encode_monitor_request(const ProtoMonitorRequest &src, FRT_RPCRequest &dst)
{
    dst.SetMethodName("vespa.searchprotocol.ping");
    encode_message(src, *dst.GetParams());
}

bool
ProtoRpcAdapter::decode_monitor_reply(FRT_RPCRequest &src, ProtoMonitorReply &dst)
{
    return (src.CheckReturnTypes("bix") && decode_message(*src.GetReturn(), dst));
}

}
<|endoftext|>


<commit_before>#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <chrono>
#include "config.hpp"
#include "requestHandler.hpp"
#include "response.hpp"
#include "torrent.hpp"
#include "utility.hpp"
#include "mysql.hpp"

TorrentMap RequestHandler::torMap;
UserMap RequestHandler::usrMap;
Database* RequestHandler::db;

std::string RequestHandler::handle(std::string str, std::string ip)
{
	std::pair<Request, std::forward_list<std::string>> infos = Parser::parse(str); // parse the request
	Request* req = &infos.first;
	std::forward_list<std::string>* infoHashes = &infos.second;
	bool gzip = false;
	try { // check if the client accepts gzip
		if (req->at("accept-encoding").find("gzip") != std::string::npos)
			gzip = true;
	} catch (const std::exception& e) {}
	std::string check = Parser::check(*req); // check if we have all we need to process (saves time if not the case
	if (infoHashes->begin() == infoHashes->end())
		return error("missing info_hash", gzip);
	if (check != "success") // missing params
		return error(check, gzip);
	if (Config::get("type") == "private" && getUser(req->at("passkey")) == nullptr)
		return error("passkey not found", gzip);
	try {
		req->at("ip") = Utility::ip_hex_encode(req->at("ip")) + Utility::port_hex_encode(req->at("port"));
	} catch (const std::exception& e) {
		req->emplace("ip", ip + Utility::port_hex_encode(req->at("port")));
	}
	if (req->at("action") == "announce") {
		req->emplace("event", "updating");
		return announce(req, infoHashes->front(), gzip);
	}
	else if (req->at("action") == "scrape")
		return scrape(infoHashes, gzip);
	return error("invalid action", gzip); // not possible, since the request is checked, but, well, who knows :3
}

std::string RequestHandler::announce(const Request* req, const std::string& infoHash, const bool& gzip)
{
	auto duration = std::chrono::system_clock::now().time_since_epoch();
	long long now = std::chrono::duration_cast<std::chrono::seconds>(duration).count();
	if (Config::get("type") != "private")
		torMap.emplace(infoHash, Torrent(0));
	Torrent *tor = nullptr;
	try {
		tor = &torMap.at(infoHash);
	} catch (const std::exception& e) {
		return error("unregistered torrent", gzip);
	}
	Peers *peers = nullptr;
	Peer *peer = nullptr;
	if (req->at("left") != "0") {
		peer = tor->getLeechers()->getPeer(req->at("peer_id"), now);
		if (req->at("event") == "stopped") {
			if (peer != nullptr) {
				db->record(peer->record(std::stoul(req->at("left")), req->at("peer_id")));
				db->record(Peer::remove(req->at("peer_id"), tor->getID()));
				tor->getLeechers()->removePeer(*req);
			}
		} else if (req->at("event") == "started") {
			if (peer == nullptr)
				tor->getLeechers()->addPeer(*req, tor->getID(), now);
		} else if (peer != nullptr) {
			peer->updateStats(std::stoul(req->at("downloaded"))/(1-(tor->getFree()/100)), now);
		}
		peers = tor->getSeeders();
	} else {
		peer = tor->getSeeders()->getPeer(req->at("peer_id"), now);
		if (req->at("event") == "stopped" || req->at("event") == "completed") {
			if (peer != nullptr) {
				if (req->at("event") == "completed") {
					tor->downloadedpp();
					tor->getSeeders()->addPeer(*req, tor->getID(), now);
					tor->getLeechers()->removePeer(*req);
				} else {
					db->record(peer->record(std::stoul(req->at("left")), req->at("peer_id")));
					db->record(Peer::remove(req->at("peer_id"), tor->getID()));
				}
			}
		} else if (req->at("event") == "started") {
			if (peer == nullptr)
				tor->getSeeders()->addPeer(*req, tor->getID(), now);
		} else if (peer != nullptr) {
			tor->getLeechers()->getPeer(req->at("peer_id"), now)->updateStats(std::stoul(req->at("downloaded"))/(1-(tor->getFree()/100)), now);
		}
		peers = tor->getLeechers();
	}
	std::string peerlist;
	unsigned long i = 0;
	if (req->at("event") != "stopped") {
		i = 10;
		try {
			i = std::stoi(req->at("numwant"));
		} catch (const std::exception& e) {}
		i = std::min(i, peers->size());
	}
	while (i-- > 0) {
		Peer* p = peers->nextPeer(now);
		if (p != nullptr)
			peerlist.append(*p->getHexIP());
	}
	return response(
			("d8:completei"
			 + std::to_string(tor->getSeeders()->size())
			 + "e10:incompletei"
			 + std::to_string(tor->getLeechers()->size())
			 + "e10:downloadedi"
			 + std::to_string(tor->getDownloaded())
			 + "e8:intervali"
			 + std::to_string(900)
			 + "e12:min intervali"
			 + std::to_string(300)
			 + "e5:peers"
			 + std::to_string(peerlist.length())
			 + ":"
			 + peerlist
			 + "e"),
			gzip
			); // doesn't look as bad as it is stated on ocelot, needs stresstesting to check
}

std::string RequestHandler::scrape(const std::forward_list<std::string>* infoHashes, const bool& gzip)
{
	std::string resp("d5:filesd");

	for (const auto& infoHash : *infoHashes) {
		const TorrentMap::iterator it = torMap.find(infoHash);
		if (it != torMap.end()) {
			resp += std::to_string(infoHash.length())
				+ ":"
				+ infoHash
				+ "d8:completei"
				+ std::to_string(it->second.getSeeders()->size())
				+ "e10:incompletei"
				+ std::to_string(it->second.getLeechers()->size())
				+ "e10:downloadedi"
				+ std::to_string(it->second.getDownloaded())
				+ "ee";
		}
	}
	resp += "ee";

	return response(resp, gzip);
}

std::string RequestHandler::update(const Request* req)
{
	std::string resp;
	const std::string& type = req->at("type");

	if (type == "change_passkey")
		resp = changePasskey(req);
	else if (type == "add_torrent")
		resp = addTorrent(req);
	else if (type == "delete_torrent")
		resp = deleteTorrent(req);
	else if (type == "update_torrent")
		resp = updateTorrent(req);
	else if (type == "add_user")
		resp = addUser(req);
	else if (type == "remove_user")
		resp = removeUser(req);

	return resp;
}

std::string RequestHandler::changePasskey(const Request* req)
{
	try {
		const std::string& old = req->at("oldpasskey");
		User* user = usrMap.at(old);
		usrMap.erase(old);
		usrMap.emplace(req->at("newpasskey"), user);

		return "success";
	}
	catch (const std::exception& e) {
		return "failure";
	}
}

std::string RequestHandler::addTorrent(const Request* req)
{
	try {
		auto t = torMap.emplace(req->at("info_hash"), Torrent(std::stoul(req->at("id"))));
		if (!t.second)
			return "failure";

		try {
			if (req->at("freetorrent") == "1")
				t.first->second.setFree(1);
			else
				t.first->second.setFree(0);
		}
		catch (const std::exception& e) {}

		return "success";
	}
	catch (const std::exception& e) {
		return "failure";
	}
}

std::string RequestHandler::deleteTorrent(const Request* req)
{
	const auto it = torMap.find(req->at("info_hash"));
	if (it != torMap.end())
		torMap.erase(it);

	return "success";
}

std::string RequestHandler::updateTorrent(const Request* req)
{
	auto it = torMap.find(req->at("info_hash"));
	if (it != torMap.end()) {
		if (req->at("freetorrent") == "1")
			it->second.setFree(1);
		else
			it->second.setFree(0);

		return "success";
	}

	return "failure";
}

std::string RequestHandler::addUser(const Request* req)
{
	try {
		return (usrMap.emplace(req->at("passkey"), new User(std::stoul(req->at("id")))).second) ? "success" : "failure";
	}
	catch (const std::exception& e) {
		return "failure";
	}
}

std::string RequestHandler::removeUser(const Request* req)
{
	return (usrMap.erase(req->at("passkey")) == 1) ? "success" : "failure";
}

User* RequestHandler::getUser(const std::string& passkey) {
	try {
		return usrMap.at(passkey);
	} catch (const std::exception& e) {
		return nullptr;
	}
}

void RequestHandler::init() {
	db = new MySQL();
	db->connect();
	db->loadUsers(usrMap);
	db->loadTorrents(torMap);
}

void RequestHandler::stop() {
	// TODO record every changes before flushing
	db->flush();
	db->disconnect();
}

void RequestHandler::clearTorrentPeers(ev::timer& timer, int revents)
{
	auto duration = std::chrono::system_clock::now().time_since_epoch();
	long long now = std::chrono::duration_cast<std::chrono::seconds>(duration).count();
	for (auto& t : torMap) {
		t.second.getSeeders()->timedOut(now);
		t.second.getLeechers()->timedOut(now);
	}
}
<commit_msg>was getting a leecher instead of a seeder<commit_after>#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <chrono>
#include "config.hpp"
#include "requestHandler.hpp"
#include "response.hpp"
#include "torrent.hpp"
#include "utility.hpp"
#include "mysql.hpp"

TorrentMap RequestHandler::torMap;
UserMap RequestHandler::usrMap;
Database* RequestHandler::db;

std::string RequestHandler::handle(std::string str, std::string ip)
{
	std::pair<Request, std::forward_list<std::string>> infos = Parser::parse(str); // parse the request
	Request* req = &infos.first;
	std::forward_list<std::string>* infoHashes = &infos.second;
	bool gzip = false;
	try { // check if the client accepts gzip
		if (req->at("accept-encoding").find("gzip") != std::string::npos)
			gzip = true;
	} catch (const std::exception& e) {}
	std::string check = Parser::check(*req); // check if we have all we need to process (saves time if not the case
	if (infoHashes->begin() == infoHashes->end())
		return error("missing info_hash", gzip);
	if (check != "success") // missing params
		return error(check, gzip);
	if (Config::get("type") == "private" && getUser(req->at("passkey")) == nullptr)
		return error("passkey not found", gzip);
	try {
		req->at("ip") = Utility::ip_hex_encode(req->at("ip")) + Utility::port_hex_encode(req->at("port"));
	} catch (const std::exception& e) {
		req->emplace("ip", ip + Utility::port_hex_encode(req->at("port")));
	}
	if (req->at("action") == "announce") {
		req->emplace("event", "updating");
		return announce(req, infoHashes->front(), gzip);
	}
	else if (req->at("action") == "scrape")
		return scrape(infoHashes, gzip);
	return error("invalid action", gzip); // not possible, since the request is checked, but, well, who knows :3
}

std::string RequestHandler::announce(const Request* req, const std::string& infoHash, const bool& gzip)
{
	auto duration = std::chrono::system_clock::now().time_since_epoch();
	long long now = std::chrono::duration_cast<std::chrono::seconds>(duration).count();
	if (Config::get("type") != "private")
		torMap.emplace(infoHash, Torrent(0));
	Torrent *tor = nullptr;
	try {
		tor = &torMap.at(infoHash);
	} catch (const std::exception& e) {
		return error("unregistered torrent", gzip);
	}
	Peers *peers = nullptr;
	Peer *peer = nullptr;
	if (req->at("left") != "0") {
		peer = tor->getLeechers()->getPeer(req->at("peer_id"), now);
		if (req->at("event") == "stopped") {
			if (peer != nullptr) {
				db->record(peer->record(std::stoul(req->at("left")), req->at("peer_id")));
				db->record(Peer::remove(req->at("peer_id"), tor->getID()));
				tor->getLeechers()->removePeer(*req);
			}
		} else if (req->at("event") == "started") {
			if (peer == nullptr)
				tor->getLeechers()->addPeer(*req, tor->getID(), now);
		} else {
			peer->updateStats(std::stoul(req->at("downloaded"))/(1-(tor->getFree()/100)), now);
		}
		peers = tor->getSeeders();
	} else {
		peer = tor->getSeeders()->getPeer(req->at("peer_id"), now);
		if (req->at("event") == "stopped" || req->at("event") == "completed") {
			if (peer != nullptr) {
				if (req->at("event") == "completed") {
					tor->downloadedpp();
					tor->getSeeders()->addPeer(*req, tor->getID(), now);
					tor->getLeechers()->removePeer(*req);
				} else {
					db->record(peer->record(std::stoul(req->at("left")), req->at("peer_id")));
					db->record(Peer::remove(req->at("peer_id"), tor->getID()));
				}
			}
		} else if (req->at("event") == "started") {
			if (peer == nullptr)
				tor->getSeeders()->addPeer(*req, tor->getID(), now);
		} else {
			tor->getSeeders()->getPeer(req->at("peer_id"), now)->updateStats(std::stoul(req->at("downloaded"))/(1-(tor->getFree()/100)), now);
		}
		peers = tor->getLeechers();
	}
	std::string peerlist;
	unsigned long i = 0;
	if (req->at("event") != "stopped") {
		i = 10;
		try {
			i = std::stoi(req->at("numwant"));
		} catch (const std::exception& e) {}
		i = std::min(i, peers->size());
	}
	while (i-- > 0) {
		Peer* p = peers->nextPeer(now);
		if (p != nullptr)
			peerlist.append(*p->getHexIP());
	}
	return response(
			("d8:completei"
			 + std::to_string(tor->getSeeders()->size())
			 + "e10:incompletei"
			 + std::to_string(tor->getLeechers()->size())
			 + "e10:downloadedi"
			 + std::to_string(tor->getDownloaded())
			 + "e8:intervali"
			 + std::to_string(900)
			 + "e12:min intervali"
			 + std::to_string(300)
			 + "e5:peers"
			 + std::to_string(peerlist.length())
			 + ":"
			 + peerlist
			 + "e"),
			gzip
			); // doesn't look as bad as it is stated on ocelot, needs stresstesting to check
}

std::string RequestHandler::scrape(const std::forward_list<std::string>* infoHashes, const bool& gzip)
{
	std::string resp("d5:filesd");

	for (const auto& infoHash : *infoHashes) {
		const TorrentMap::iterator it = torMap.find(infoHash);
		if (it != torMap.end()) {
			resp += std::to_string(infoHash.length())
				+ ":"
				+ infoHash
				+ "d8:completei"
				+ std::to_string(it->second.getSeeders()->size())
				+ "e10:incompletei"
				+ std::to_string(it->second.getLeechers()->size())
				+ "e10:downloadedi"
				+ std::to_string(it->second.getDownloaded())
				+ "ee";
		}
	}
	resp += "ee";

	return response(resp, gzip);
}

std::string RequestHandler::update(const Request* req)
{
	std::string resp;
	const std::string& type = req->at("type");

	if (type == "change_passkey")
		resp = changePasskey(req);
	else if (type == "add_torrent")
		resp = addTorrent(req);
	else if (type == "delete_torrent")
		resp = deleteTorrent(req);
	else if (type == "update_torrent")
		resp = updateTorrent(req);
	else if (type == "add_user")
		resp = addUser(req);
	else if (type == "remove_user")
		resp = removeUser(req);

	return resp;
}

std::string RequestHandler::changePasskey(const Request* req)
{
	try {
		const std::string& old = req->at("oldpasskey");
		User* user = usrMap.at(old);
		usrMap.erase(old);
		usrMap.emplace(req->at("newpasskey"), user);

		return "success";
	}
	catch (const std::exception& e) {
		return "failure";
	}
}

std::string RequestHandler::addTorrent(const Request* req)
{
	try {
		auto t = torMap.emplace(req->at("info_hash"), Torrent(std::stoul(req->at("id"))));
		if (!t.second)
			return "failure";

		try {
			if (req->at("freetorrent") == "1")
				t.first->second.setFree(1);
			else
				t.first->second.setFree(0);
		}
		catch (const std::exception& e) {}

		return "success";
	}
	catch (const std::exception& e) {
		return "failure";
	}
}

std::string RequestHandler::deleteTorrent(const Request* req)
{
	const auto it = torMap.find(req->at("info_hash"));
	if (it != torMap.end())
		torMap.erase(it);

	return "success";
}

std::string RequestHandler::updateTorrent(const Request* req)
{
	auto it = torMap.find(req->at("info_hash"));
	if (it != torMap.end()) {
		if (req->at("freetorrent") == "1")
			it->second.setFree(1);
		else
			it->second.setFree(0);

		return "success";
	}

	return "failure";
}

std::string RequestHandler::addUser(const Request* req)
{
	try {
		return (usrMap.emplace(req->at("passkey"), new User(std::stoul(req->at("id")))).second) ? "success" : "failure";
	}
	catch (const std::exception& e) {
		return "failure";
	}
}

std::string RequestHandler::removeUser(const Request* req)
{
	return (usrMap.erase(req->at("passkey")) == 1) ? "success" : "failure";
}

User* RequestHandler::getUser(const std::string& passkey) {
	try {
		return usrMap.at(passkey);
	} catch (const std::exception& e) {
		return nullptr;
	}
}

void RequestHandler::init() {
	db = new MySQL();
	db->connect();
	db->loadUsers(usrMap);
	db->loadTorrents(torMap);
}

void RequestHandler::stop() {
	// TODO record every changes before flushing
	db->flush();
	db->disconnect();
}

void RequestHandler::clearTorrentPeers(ev::timer& timer, int revents)
{
	auto duration = std::chrono::system_clock::now().time_since_epoch();
	long long now = std::chrono::duration_cast<std::chrono::seconds>(duration).count();
	for (auto& t : torMap) {
		t.second.getSeeders()->timedOut(now);
		t.second.getLeechers()->timedOut(now);
	}
}
<|endoftext|>


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

// Perform decompression from *.jp2 to *.pnm format

#include <libgen.h>
#include <syscall.h>

#include <cstdlib>
#include <iostream>
#include <vector>
#include <iterator>

#include "absl/algorithm/container.h"
#include "convert_helper.h"
#include "openjp2_sapi.sapi.h"

class Openjp2SapiSandbox : public Openjp2Sandbox {
 public:
  Openjp2SapiSandbox(const std::string in_file)
      : in_file_(std::move(in_file)) {}

  std::unique_ptr<sandbox2::Policy> ModifyPolicy(
      sandbox2::PolicyBuilder*) override {
    return sandbox2::PolicyBuilder()
        .AllowStaticStartup()
        .AllowOpen()
        .AllowRead()
        .AllowWrite()
        .AllowStat()
        .AllowSystemMalloc()
        .AllowExit()
        .AllowSyscalls({
            __NR_futex,
            __NR_close,
            __NR_lseek,
        })
        .AddFile(in_file_)
        .BuildOrDie();
  }

 private:
  std::string in_file_;
};

int main(int argc, char* argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  google::InitGoogleLogging(argv[0]);

  if (argc != 3) {
    std::cerr << "usage: " << basename(argv[0]) << " absolute/path/to/INPUT.jp2"
              << " absolute/path/to/OUTPUT.pnm\n";
    return EXIT_FAILURE;
  }

  std::string in_file(argv[1]);

  // initialize sandbox
  Openjp2SapiSandbox sandbox(in_file);
  absl::Status status = sandbox.Init();
  if (!status.ok()) {
    LOG(FATAL) << "sandbox initialization status: " << status;
    return EXIT_FAILURE;
  }

  Openjp2Api api(&sandbox);
  sapi::v::ConstCStr in_file_v(in_file.c_str());

  // initialize library's main data-holders
  sapi::StatusOr<opj_stream_t*> stream =
      api.opj_stream_create_default_file_stream(in_file_v.PtrBefore(), 1);
  if (!stream.ok()) {
    LOG(FATAL) << "opj_stream initialization failed: " << stream.status();
    return EXIT_FAILURE;
  }
  sapi::v::RemotePtr stream_pointer(stream.value());

  sapi::StatusOr<opj_codec_t*> codec = api.opj_create_decompress(OPJ_CODEC_JP2);
  if (!codec.ok()) {
    LOG(FATAL) << "opj_codec initialization failed: " << codec.status();
    return EXIT_FAILURE;
  }
  sapi::v::RemotePtr codec_pointer(codec.value());

  sapi::v::Struct<opj_dparameters_t> parameters;
  status = api.opj_set_default_decoder_parameters(parameters.PtrBoth());
  if (!status.ok()) {
    LOG(FATAL) << "parameters initialization failed: " << status;
    return EXIT_FAILURE;
  }

  sapi::StatusOr<OPJ_BOOL> bool_status =
      api.opj_setup_decoder(&codec_pointer, parameters.PtrBefore());
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "decoder setup failed";
    return EXIT_FAILURE;
  }

  // start reading image from the input file
  sapi::v::GenericPtr image_pointer;
  bool_status = api.opj_read_header(&stream_pointer, &codec_pointer,
                                    image_pointer.PtrAfter());
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "reading image header failed";
    return EXIT_FAILURE;
  }

  sapi::v::Struct<opj_image_t> image;
  image.SetRemote((void*)image_pointer.GetValue());
  if (!sandbox.TransferFromSandboxee(&image).ok()) {
    LOG(FATAL) << "transfer from sandboxee failed";
    return EXIT_FAILURE;
  }

  sapi::v::RemotePtr image_remote_pointer(image.GetRemote());
  bool_status = api.opj_decode(&codec_pointer, &stream_pointer,
                              &image_remote_pointer);
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "decoding failed";
    return EXIT_FAILURE;
  }

  bool_status = api.opj_end_decompress(&codec_pointer, &stream_pointer);
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "ending decompress failed";
    return EXIT_FAILURE;
  }

  int components = image.data().numcomps;

  // transfer the read data to the main process
  sapi::v::Array<opj_image_comp_t> image_components(components);
  image_components.SetRemote(image.data().comps);
  if (!sandbox.TransferFromSandboxee(&image_components).ok()) {
    LOG(FATAL) << "transfer from sandboxee failed";
    return EXIT_FAILURE;
  }

  image.mutable_data()->comps = (opj_image_comp_t*)image_components.GetLocal();

  int width = (int)image.data().comps[0].w;
  int height = (int)image.data().comps[0].h;

  OPJ_INT32 data[components][width * height];
  sapi::v::Array<OPJ_INT32> image_components_data(width * height);

  for (int i = 0; i < components; ++i) {
    image_components_data.SetRemote(image.data().comps[i].data);
    if (!sandbox.TransferFromSandboxee(&image_components_data).ok()) {
      LOG(FATAL) << "transfer from sandboxee failed";
      return EXIT_FAILURE;
    }
    for (int j = 0; j < width * height; ++j) {
      data[i][j] = image_components_data[j];
    }
    image_components[i].data = data[i];
  }

  // convert the image to the desired format and save it to the file
  int error = imagetopnm((opj_image_t*)image.GetLocal(), argv[2], 0);
  if (error) LOG(FATAL) << "image convert failed";

  // cleanup
  status = api.opj_image_destroy(image.PtrNone());
  if (!status.ok()) LOG(FATAL) << "image destroy failed: " << status;

  status = api.opj_stream_destroy(&stream_pointer);
  if (!status.ok()) LOG(FATAL) << "stream destroy failed: " << status;

  status = api.opj_destroy_codec(&codec_pointer);
  if (!status.ok()) LOG(FATAL) << "codec destroy failed: " << status;

  return EXIT_SUCCESS;
}
<commit_msg>changed VLA to vectors<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Perform decompression from *.jp2 to *.pnm format

#include <libgen.h>
#include <syscall.h>

#include <cstdlib>
#include <iostream>
#include <vector>

#include "absl/algorithm/container.h"
#include "convert_helper.h"
#include "openjp2_sapi.sapi.h"

class Openjp2SapiSandbox : public Openjp2Sandbox {
 public:
  Openjp2SapiSandbox(const std::string in_file)
      : in_file_(std::move(in_file)) {}

  std::unique_ptr<sandbox2::Policy> ModifyPolicy(
      sandbox2::PolicyBuilder*) override {
    return sandbox2::PolicyBuilder()
        .AllowStaticStartup()
        .AllowOpen()
        .AllowRead()
        .AllowWrite()
        .AllowStat()
        .AllowSystemMalloc()
        .AllowExit()
        .AllowSyscalls({
            __NR_futex,
            __NR_close,
            __NR_lseek,
        })
        .AddFile(in_file_)
        .BuildOrDie();
  }

 private:
  std::string in_file_;
};

int main(int argc, char* argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  google::InitGoogleLogging(argv[0]);

  if (argc != 3) {
    std::cerr << "usage: " << basename(argv[0]) << " absolute/path/to/INPUT.jp2"
              << " absolute/path/to/OUTPUT.pnm\n";
    return EXIT_FAILURE;
  }

  std::string in_file(argv[1]);

  // initialize sandbox
  Openjp2SapiSandbox sandbox(in_file);
  absl::Status status = sandbox.Init();
  if (!status.ok()) {
    LOG(FATAL) << "sandbox initialization status: " << status;
    return EXIT_FAILURE;
  }

  Openjp2Api api(&sandbox);
  sapi::v::ConstCStr in_file_v(in_file.c_str());

  // initialize library's main data-holders
  sapi::StatusOr<opj_stream_t*> stream =
      api.opj_stream_create_default_file_stream(in_file_v.PtrBefore(), 1);
  if (!stream.ok()) {
    LOG(FATAL) << "opj_stream initialization failed: " << stream.status();
    return EXIT_FAILURE;
  }
  sapi::v::RemotePtr stream_pointer(stream.value());

  sapi::StatusOr<opj_codec_t*> codec = api.opj_create_decompress(OPJ_CODEC_JP2);
  if (!codec.ok()) {
    LOG(FATAL) << "opj_codec initialization failed: " << codec.status();
    return EXIT_FAILURE;
  }
  sapi::v::RemotePtr codec_pointer(codec.value());

  sapi::v::Struct<opj_dparameters_t> parameters;
  status = api.opj_set_default_decoder_parameters(parameters.PtrBoth());
  if (!status.ok()) {
    LOG(FATAL) << "parameters initialization failed: " << status;
    return EXIT_FAILURE;
  }

  sapi::StatusOr<OPJ_BOOL> bool_status =
      api.opj_setup_decoder(&codec_pointer, parameters.PtrBefore());
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "decoder setup failed";
    return EXIT_FAILURE;
  }

  // start reading image from the input file
  sapi::v::GenericPtr image_pointer;
  bool_status = api.opj_read_header(&stream_pointer, &codec_pointer,
                                    image_pointer.PtrAfter());
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "reading image header failed";
    return EXIT_FAILURE;
  }

  sapi::v::Struct<opj_image_t> image;
  image.SetRemote((void*)image_pointer.GetValue());
  if (!sandbox.TransferFromSandboxee(&image).ok()) {
    LOG(FATAL) << "transfer from sandboxee failed";
    return EXIT_FAILURE;
  }

  sapi::v::RemotePtr image_remote_pointer(image.GetRemote());
  bool_status =
      api.opj_decode(&codec_pointer, &stream_pointer, &image_remote_pointer);
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "decoding failed";
    return EXIT_FAILURE;
  }

  bool_status = api.opj_end_decompress(&codec_pointer, &stream_pointer);
  if (!bool_status.ok() || !bool_status.value()) {
    LOG(FATAL) << "ending decompress failed";
    return EXIT_FAILURE;
  }

  int components = image.data().numcomps;

  // transfer the read data to the main process
  sapi::v::Array<opj_image_comp_t> image_components(components);
  image_components.SetRemote(image.data().comps);
  if (!sandbox.TransferFromSandboxee(&image_components).ok()) {
    LOG(FATAL) << "transfer from sandboxee failed";
    return EXIT_FAILURE;
  }

  image.mutable_data()->comps = (opj_image_comp_t*)image_components.GetLocal();

  int width = (int)image.data().comps[0].w;
  int height = (int)image.data().comps[0].h;

  std::vector<std::vector<OPJ_INT32>> data;
  sapi::v::Array<OPJ_INT32> image_components_data(width * height);

  for (int i = 0; i < components; ++i) {
    image_components_data.SetRemote(image.data().comps[i].data);
    if (!sandbox.TransferFromSandboxee(&image_components_data).ok()) {
      LOG(FATAL) << "transfer from sandboxee failed";
      return EXIT_FAILURE;
    }
    std::vector<OPJ_INT32> component_data(
        image_components_data.GetData(),
        image_components_data.GetData() + (width * height));
    data.push_back(component_data);
  }

  for (int i = 0; i < components; ++i) {
    image_components[i].data = &data[i][0];
  }

  // convert the image to the desired format and save it to the file
  int error = imagetopnm((opj_image_t*)image.GetLocal(), argv[2], 0);
  if (error) LOG(FATAL) << "image convert failed";

  // cleanup
  status = api.opj_image_destroy(image.PtrNone());
  if (!status.ok()) LOG(FATAL) << "image destroy failed: " << status;

  status = api.opj_stream_destroy(&stream_pointer);
  if (!status.ok()) LOG(FATAL) << "stream destroy failed: " << status;

  status = api.opj_destroy_codec(&codec_pointer);
  if (!status.ok()) LOG(FATAL) << "codec destroy failed: " << status;

  return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before><commit_msg>`Species::render`: render only if `this->should_be_rendered`.<commit_after><|endoftext|>


<commit_before>#include <iostream>

#include "Block.hpp"
#include "Matrix4x4.hpp"
#include "Matrix3x3.hpp"
#include "vec.hpp"
#include "PhongMaterial.hpp"

using namespace rffalcon;

// _element_indices are in the order of 3 sides: xmin, ymin, ymax
GLuint Block::_element_indices[12] = { 1, 0, 7, 6, 0, 2, 6, 4, 1, 7, 3, 5 };

Block::Block(float x, float y, float z, float lengthX, float lengthY, float lengthZ)
{
	_xmin = x;
	_xmax = x + lengthX;
	_ymin = y;
	_ymax = y + lengthY;
	_zmin = z;
	_zmax = z + lengthZ;
	_phongMaterial = std::make_shared<PhongMaterial>(PhongMaterial::PolishedCopper);
}

Block::~Block()
{
	glDeleteBuffers(2, _vbo);
	glDeleteVertexArrays(1, _vao);
}

void Block::getMCBoundingBox(double* xyzBounds) const
{
	double buffer = 0.25;
	xyzBounds[0] = _xmin - buffer;
	xyzBounds[1] = _xmax + buffer;
	xyzBounds[2] = _ymin - buffer;
	xyzBounds[3] = _ymax + buffer;
	xyzBounds[4] = _zmin - buffer;
	xyzBounds[5] = _zmax + buffer;
}

void Block::render()
{
	if (!_initialized)
	{
		_initBlock();
		_initialized = true;
	}

	glBindVertexArray(_vao[0]);

	GLint programId;
	glGetIntegerv(GL_CURRENT_PROGRAM, &programId);
	GLint pvaLoc_mcNormal = glGetAttribLocation(programId, "mcNormal");
	if (pvaLoc_mcNormal  < 0)
	{
		std::cerr << "Unable to find per-vertex attribute: 'mcNormal'" << std::endl;
	}
	else
	{
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, 0.0, 1.0);
		glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

		//float color2[] = { 0.8f, 0.8f, 0.0f };
		//glUniform3fv(_ppuLoc_kd, 1, color2);
		glVertexAttrib3f(pvaLoc_mcNormal, 1.0, 0.0, 0.0);
		glDrawArrays(GL_TRIANGLE_STRIP, 2, 4);

		//float color3[] = { 0.0f, 0.8f, 0.0f };
		//glUniform3fv(_ppuLoc_kd, 1, color3);
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, 0.0, -1.0);
		glDrawArrays(GL_TRIANGLE_STRIP, 4, 4);

		glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, _vbo[1]);

		//float color4[] = { 0.8f, 0.0f, 0.0f };
		//glUniform3fv(_ppuLoc_kd, 1, color4);
		glVertexAttrib3f(pvaLoc_mcNormal, -1.0, 0.0, 0.0);
		glDrawElements(GL_TRIANGLE_STRIP, 4, GL_UNSIGNED_INT, (void*)0);

		//float color5[] = { 0.8f, 0.0f, 0.8f };
		//glUniform3fv(_ppuLoc_kd, 1, color5);
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, -1.0, 0.0);
		glDrawElements(GL_TRIANGLE_STRIP, 4, GL_UNSIGNED_INT, (void*)(4 * sizeof(int)));

		//float color6[] = { 0.0f, 0.8f, 0.8f };
		//glUniform3fv(_ppuLoc_kd, 1, color6);
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, 1.0, 0.0);
		glDrawElements(GL_TRIANGLE_STRIP, 4, GL_UNSIGNED_INT, (void*)(8 * sizeof(int)));
	}
}

void Block::_initBlock()
{
	vec3 vertices[] = {
		{ _xmin, _ymin, _zmax }, { _xmin, _ymax, _zmax },
		{ _xmax, _ymin, _zmax }, { _xmax, _ymax, _zmax },
		{ _xmax, _ymin, _zmin }, { _xmax, _ymax, _zmin },
		{ _xmin, _ymin, _zmin }, { _xmin, _ymax, _zmin }
	};
	glGenVertexArrays(1, _vao);
	glBindVertexArray(_vao[0]);
	glGenBuffers(2, _vbo);
	glBindBuffer(GL_ARRAY_BUFFER, _vbo[0]);
	glBufferData(GL_ARRAY_BUFFER, 8 * sizeof(vec3), vertices, GL_STATIC_DRAW);
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, _vbo[1]);
	glBufferData(GL_ELEMENT_ARRAY_BUFFER, 12 * sizeof(int), _element_indices, GL_STATIC_DRAW);

	GLint programId;
	glGetIntegerv(GL_CURRENT_PROGRAM, &programId);
	GLint pvaLoc_mcNormal = glGetAttribLocation(programId, "mcNormal");
	GLint pvaLoc_mcPosition = glGetAttribLocation(programId, "mcPosition");

	if (pvaLoc_mcPosition < 0)
	{
		std::cerr << "Unable to find per-vertex attribute: 'mcPosition'" << std::endl;
	}
	else
	{
		glVertexAttribPointer(pvaLoc_mcPosition, 3, GL_FLOAT, GL_FALSE, 0, 0);
		glEnableVertexAttribArray(pvaLoc_mcPosition);
	}

	if (pvaLoc_mcNormal  < 0)
	{
		std::cerr << "Unable to find per-vertex attribute: 'mcNormal'" << std::endl;
	}
	else
	{
		glDisableVertexAttribArray(pvaLoc_mcNormal);
	}
}
<commit_msg>Removing unnecessary include<commit_after>#include <iostream>

#include "Block.hpp"
#include "Matrix3x3.hpp"
#include "vec.hpp"
#include "PhongMaterial.hpp"

using namespace rffalcon;

// _element_indices are in the order of 3 sides: xmin, ymin, ymax
GLuint Block::_element_indices[12] = { 1, 0, 7, 6, 0, 2, 6, 4, 1, 7, 3, 5 };

Block::Block(float x, float y, float z, float lengthX, float lengthY, float lengthZ)
{
	_xmin = x;
	_xmax = x + lengthX;
	_ymin = y;
	_ymax = y + lengthY;
	_zmin = z;
	_zmax = z + lengthZ;
	_phongMaterial = std::make_shared<PhongMaterial>(PhongMaterial::PolishedCopper);
}

Block::~Block()
{
	glDeleteBuffers(2, _vbo);
	glDeleteVertexArrays(1, _vao);
}

void Block::getMCBoundingBox(double* xyzBounds) const
{
	double buffer = 0.25;
	xyzBounds[0] = _xmin - buffer;
	xyzBounds[1] = _xmax + buffer;
	xyzBounds[2] = _ymin - buffer;
	xyzBounds[3] = _ymax + buffer;
	xyzBounds[4] = _zmin - buffer;
	xyzBounds[5] = _zmax + buffer;
}

void Block::render()
{
	if (!_initialized)
	{
		_initBlock();
		_initialized = true;
	}

	glBindVertexArray(_vao[0]);

	GLint programId;
	glGetIntegerv(GL_CURRENT_PROGRAM, &programId);
	GLint pvaLoc_mcNormal = glGetAttribLocation(programId, "mcNormal");
	if (pvaLoc_mcNormal  < 0)
	{
		std::cerr << "Unable to find per-vertex attribute: 'mcNormal'" << std::endl;
	}
	else
	{
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, 0.0, 1.0);
		glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

		//float color2[] = { 0.8f, 0.8f, 0.0f };
		//glUniform3fv(_ppuLoc_kd, 1, color2);
		glVertexAttrib3f(pvaLoc_mcNormal, 1.0, 0.0, 0.0);
		glDrawArrays(GL_TRIANGLE_STRIP, 2, 4);

		//float color3[] = { 0.0f, 0.8f, 0.0f };
		//glUniform3fv(_ppuLoc_kd, 1, color3);
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, 0.0, -1.0);
		glDrawArrays(GL_TRIANGLE_STRIP, 4, 4);

		glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, _vbo[1]);

		//float color4[] = { 0.8f, 0.0f, 0.0f };
		//glUniform3fv(_ppuLoc_kd, 1, color4);
		glVertexAttrib3f(pvaLoc_mcNormal, -1.0, 0.0, 0.0);
		glDrawElements(GL_TRIANGLE_STRIP, 4, GL_UNSIGNED_INT, (void*)0);

		//float color5[] = { 0.8f, 0.0f, 0.8f };
		//glUniform3fv(_ppuLoc_kd, 1, color5);
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, -1.0, 0.0);
		glDrawElements(GL_TRIANGLE_STRIP, 4, GL_UNSIGNED_INT, (void*)(4 * sizeof(int)));

		//float color6[] = { 0.0f, 0.8f, 0.8f };
		//glUniform3fv(_ppuLoc_kd, 1, color6);
		glVertexAttrib3f(pvaLoc_mcNormal, 0.0, 1.0, 0.0);
		glDrawElements(GL_TRIANGLE_STRIP, 4, GL_UNSIGNED_INT, (void*)(8 * sizeof(int)));
	}
}

void Block::_initBlock()
{
	vec3 vertices[] = {
		{ _xmin, _ymin, _zmax }, { _xmin, _ymax, _zmax },
		{ _xmax, _ymin, _zmax }, { _xmax, _ymax, _zmax },
		{ _xmax, _ymin, _zmin }, { _xmax, _ymax, _zmin },
		{ _xmin, _ymin, _zmin }, { _xmin, _ymax, _zmin }
	};
	glGenVertexArrays(1, _vao);
	glBindVertexArray(_vao[0]);
	glGenBuffers(2, _vbo);
	glBindBuffer(GL_ARRAY_BUFFER, _vbo[0]);
	glBufferData(GL_ARRAY_BUFFER, 8 * sizeof(vec3), vertices, GL_STATIC_DRAW);
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, _vbo[1]);
	glBufferData(GL_ELEMENT_ARRAY_BUFFER, 12 * sizeof(int), _element_indices, GL_STATIC_DRAW);

	GLint programId;
	glGetIntegerv(GL_CURRENT_PROGRAM, &programId);
	GLint pvaLoc_mcNormal = glGetAttribLocation(programId, "mcNormal");
	GLint pvaLoc_mcPosition = glGetAttribLocation(programId, "mcPosition");

	if (pvaLoc_mcPosition < 0)
	{
		std::cerr << "Unable to find per-vertex attribute: 'mcPosition'" << std::endl;
	}
	else
	{
		glVertexAttribPointer(pvaLoc_mcPosition, 3, GL_FLOAT, GL_FALSE, 0, 0);
		glEnableVertexAttribArray(pvaLoc_mcPosition);
	}

	if (pvaLoc_mcNormal  < 0)
	{
		std::cerr << "Unable to find per-vertex attribute: 'mcNormal'" << std::endl;
	}
	else
	{
		glDisableVertexAttribArray(pvaLoc_mcNormal);
	}
}
<|endoftext|>


<commit_before>/*
 * This file is part of the statismo library.
 *
 * Author: Marcel Luethi (marcel.luethi@unibas.ch)
 *
 * Copyright (c) 2011 University of Basel
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * Neither the name of the project's author nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#ifndef __PCAModelBuilder_TXX
#define __PCAModelBuilder_TXX

#include "PCAModelBuilder.h"

#include <iostream>

#include <Eigen/SVD>
#include <Eigen/Eigenvalues>

#include "CommonTypes.h"
#include "Exceptions.h"

namespace statismo {

template <typename T>
PCAModelBuilder<T>::PCAModelBuilder()
    : Superclass() {
}

template <typename T>
typename PCAModelBuilder<T>::StatisticalModelType*
PCAModelBuilder<T>::BuildNewModel(const DataItemListType& sampleDataList, double noiseVariance, bool computeScores, EigenValueMethod method) const {

    unsigned n = sampleDataList.size();
    if (n <= 0) {
        throw StatisticalModelException("Provided empty sample set. Cannot build the sample matrix");
    }

    unsigned p = sampleDataList.front()->GetSampleVector().rows();
    const Representer<T>* representer = sampleDataList.front()->GetRepresenter();

    // Compute the mean vector mu
    VectorType mu = VectorType::Zero(p);

    for (typename DataItemListType::const_iterator it = sampleDataList.begin();
      it != sampleDataList.end();  ++it) {
      assert((*it)->GetSampleVector().rows() == p); // all samples must have same number of rows
      assert((*it)->GetRepresenter() == representer); // all samples have the same representer
      mu += (*it)->GetSampleVector();
    }

    // Build the mean free sample matrix X0
    MatrixType X0(n, p);
    unsigned i = 0;
    for (typename DataItemListType::const_iterator it = sampleDataList.begin();
      it != sampleDataList.end(); ++it) {
      X0.row(i++) = (*it)->GetSampleVector() - mu;
    }

    // build the model
    StatisticalModelType* model = BuildNewModelInternal(representer, X0, mu, noiseVariance, method);

    // compute the scores if requested
    MatrixType scores;
    if (computeScores) {
      scores = this->ComputeScores(sampleDataList, model);
    }

    typename BuilderInfo::ParameterInfoList bi;
    bi.push_back(BuilderInfo::KeyValuePair("NoiseVariance ", Utils::toString(noiseVariance)));

    typename BuilderInfo::DataInfoList dataInfo;
    i = 0;
    for (typename DataItemListType::const_iterator it = sampleDataList.begin();
            it != sampleDataList.end();
            ++it, i++) {
        std::ostringstream os;
        os << "URI_" << i;
        dataInfo.push_back(BuilderInfo::KeyValuePair(os.str().c_str(),(*it)->GetDatasetURI()));
    }

    // finally add meta data to the model info
    BuilderInfo builderInfo("PCAModelBuilder", dataInfo, bi);

    ModelInfo::BuilderInfoList biList;
    biList.push_back(builderInfo);

    ModelInfo info(scores, biList);
    model->SetModelInfo(info);

    return model;
}

template <typename T>
typename PCAModelBuilder<T>::StatisticalModelType*
PCAModelBuilder<T>::BuildNewModelInternal(const Representer<T>* representer, const MatrixType& X0, const VectorType& mu,
    double noiseVariance, EigenValueMethod method) const
{

    unsigned n = X0.rows();
    unsigned p = X0.cols();

    switch(method) {
    case JacobiSVD:

        typedef Eigen::JacobiSVD<MatrixType> SVDType;
        typedef Eigen::JacobiSVD<MatrixTypeDoublePrecision> SVDDoublePrecisionType;

        // We destinguish the case where we have more variables than samples and
        // the case where we have more samples than variable.
        // In the first case we compute the (smaller) inner product matrix instead of the full covariance matrix.
        // It is known that this has the same non-zero singular values as the covariance matrix.
        // Furthermore, it is possible to compute the corresponding eigenvectors of the covariance matrix from the
        // decomposition.

        if (n < p) {
            // we compute the eigenvectors of the covariance matrix by computing an SVD of the
            // n x n inner product matrix 1/(n-1) X0X0^T
            MatrixType Cov = X0 * X0.transpose() * 1.0/(n-1);
            SVDDoublePrecisionType SVD(Cov.cast<double>(), Eigen::ComputeThinV);
            VectorType singularValues = SVD.singularValues().cast<ScalarType>();
            MatrixType V = SVD.matrixV().cast<ScalarType>();

            unsigned numComponentsAboveTolerance = ((singularValues.array() - noiseVariance - Superclass::TOLERANCE) > 0).count();

            // there can be at most n-1 nonzero singular values in this case. Everything else must be due to numerical inaccuracies
            unsigned numComponentsToKeep = std::min(numComponentsAboveTolerance, n - 1);
            // compute the pseudo inverse of the square root of the singular values
            // which is then needed to recompute the PCA basis
            VectorType singSqrt = singularValues.array().sqrt();
            VectorType singSqrtInv = VectorType::Zero(singSqrt.rows());
            for (unsigned i = 0; i < numComponentsToKeep; i++) {
                assert(singSqrt(i) > Superclass::TOLERANCE);
                singSqrtInv(i) = 1.0 / singSqrt(i);
            }

            if (numComponentsToKeep == 0) {
              throw StatisticalModelException("All the eigenvalues are below the given tolerance. Model cannot be built.");
            }

            // we recover the eigenvectors U of the full covariance matrix from the eigenvectors V of the inner product matrix.
            // We use the fact that if we decompose X as X=UDV^T, then we get X^TX = UD^2U^T and XX^T = VD^2V^T (exploiting the orthogonormality
            // of the matrix U and V from the SVD). The additional factor sqrt(n-1) is to compensate for the 1/sqrt(n-1) in the formula
            // for the covariance matrix.

            MatrixType pcaBasis = X0.transpose() * V * singSqrtInv.asDiagonal();
            pcaBasis /= sqrt(n - 1.0);
            pcaBasis.conservativeResize(Eigen::NoChange, numComponentsToKeep);

            VectorType sampleVarianceVector = singularValues.topRows(numComponentsToKeep);
            VectorType pcaVariance = (sampleVarianceVector - VectorType::Ones(numComponentsToKeep) * noiseVariance);

            StatisticalModelType* model = StatisticalModelType::Create(representer, mu, pcaBasis, pcaVariance, noiseVariance);

            return model;
        } else {
            // we compute an SVD of the full p x p  covariance matrix 1/(n-1) X0^TX0 directly
            SVDType SVD(X0.transpose() * X0, Eigen::ComputeThinU);
            VectorType singularValues = SVD.singularValues();
            unsigned numComponentsToKeep = ((singularValues.array() - noiseVariance - Superclass::TOLERANCE) > 0).count();
            MatrixType pcaBasis = SVD.matrixU();
            pcaBasis /= sqrt(n - 1.0);
            pcaBasis.conservativeResize(Eigen::NoChange, numComponentsToKeep);

            if (numComponentsToKeep == 0) {
                throw StatisticalModelException("All the eigenvalues are below the given tolerance. Model cannot be built.");
            }

            VectorType sampleVarianceVector = singularValues.topRows(numComponentsToKeep);
            VectorType pcaVariance = (sampleVarianceVector - VectorType::Ones(numComponentsToKeep) * noiseVariance);
            StatisticalModelType* model = StatisticalModelType::Create(representer, mu, pcaBasis, pcaVariance, noiseVariance);
            return model;
        }
        break;

    case SelfAdjointEigenSolver: {
        // we compute the eigenvalues/eigenvectors of the full p x p  covariance matrix 1/(n-1) X0^TX0 directly

        typedef Eigen::SelfAdjointEigenSolver<MatrixType> SelfAdjointEigenSolver;
        SelfAdjointEigenSolver es;
        es.compute(X0.transpose() * X0);
        VectorType eigenValues = es.eigenvalues().reverse(); // SelfAdjointEigenSolver orders the eigenvalues in increasing order

        unsigned numComponentsToKeep = ((eigenValues.array() - noiseVariance - Superclass::TOLERANCE) > 0).count();
        MatrixType pcaBasis = es.eigenvectors().rowwise().reverse();
        pcaBasis /= sqrt(n - 1.0);
        pcaBasis.conservativeResize(Eigen::NoChange, numComponentsToKeep);

        if (numComponentsToKeep == 0) {
            throw StatisticalModelException("All the eigenvalues are below the given tolerance. Model cannot be built.");
        }

        VectorType sampleVarianceVector = eigenValues.topRows(numComponentsToKeep);
        VectorType pcaVariance = (sampleVarianceVector - VectorType::Ones(numComponentsToKeep) * noiseVariance);
        StatisticalModelType* model = StatisticalModelType::Create(representer, mu, pcaBasis, pcaVariance, noiseVariance);
        return model;
    }
    break;

    default:
        throw StatisticalModelException("Unrecognized decomposition/eigenvalue solver method.");
        return 0;
        break;
    }
}

} // namespace statismo

#endif
<commit_msg>fixed bug in mean computation<commit_after>/*
 * This file is part of the statismo library.
 *
 * Author: Marcel Luethi (marcel.luethi@unibas.ch)
 *
 * Copyright (c) 2011 University of Basel
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * Neither the name of the project's author nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#ifndef __PCAModelBuilder_TXX
#define __PCAModelBuilder_TXX

#include "PCAModelBuilder.h"

#include <iostream>

#include <Eigen/SVD>
#include <Eigen/Eigenvalues>

#include "CommonTypes.h"
#include "Exceptions.h"

namespace statismo {

template <typename T>
PCAModelBuilder<T>::PCAModelBuilder()
    : Superclass() {
}

template <typename T>
typename PCAModelBuilder<T>::StatisticalModelType*
PCAModelBuilder<T>::BuildNewModel(const DataItemListType& sampleDataList, double noiseVariance, bool computeScores, EigenValueMethod method) const {

    unsigned n = sampleDataList.size();
    if (n <= 0) {
        throw StatisticalModelException("Provided empty sample set. Cannot build the sample matrix");
    }

    unsigned p = sampleDataList.front()->GetSampleVector().rows();
    const Representer<T>* representer = sampleDataList.front()->GetRepresenter();

    // Compute the mean vector mu
    VectorType mu = VectorType::Zero(p);

    for (typename DataItemListType::const_iterator it = sampleDataList.begin();
      it != sampleDataList.end();  ++it) {
      assert((*it)->GetSampleVector().rows() == p); // all samples must have same number of rows
      assert((*it)->GetRepresenter() == representer); // all samples have the same representer
      mu += (*it)->GetSampleVector();
    }
    mu /= n;

    // Build the mean free sample matrix X0
    MatrixType X0(n, p);
    unsigned i = 0;
    for (typename DataItemListType::const_iterator it = sampleDataList.begin();
      it != sampleDataList.end(); ++it) {
      X0.row(i++) = (*it)->GetSampleVector() - mu;
    }

    // build the model
    StatisticalModelType* model = BuildNewModelInternal(representer, X0, mu, noiseVariance, method);

    // compute the scores if requested
    MatrixType scores;
    computeScores = false;
    if (computeScores) {
      scores = this->ComputeScores(sampleDataList, model);
    }

    typename BuilderInfo::ParameterInfoList bi;
    bi.push_back(BuilderInfo::KeyValuePair("NoiseVariance ", Utils::toString(noiseVariance)));

    typename BuilderInfo::DataInfoList dataInfo;
    i = 0;
    for (typename DataItemListType::const_iterator it = sampleDataList.begin();
            it != sampleDataList.end();
            ++it, i++) {
        std::ostringstream os;
        os << "URI_" << i;
        dataInfo.push_back(BuilderInfo::KeyValuePair(os.str().c_str(),(*it)->GetDatasetURI()));
    }

    // finally add meta data to the model info
    BuilderInfo builderInfo("PCAModelBuilder", dataInfo, bi);

    ModelInfo::BuilderInfoList biList;
    biList.push_back(builderInfo);

    ModelInfo info(scores, biList);
    model->SetModelInfo(info);

    return model;
}

template <typename T>
typename PCAModelBuilder<T>::StatisticalModelType*
PCAModelBuilder<T>::BuildNewModelInternal(const Representer<T>* representer, const MatrixType& X0, const VectorType& mu,
    double noiseVariance, EigenValueMethod method) const
{

    unsigned n = X0.rows();
    unsigned p = X0.cols();

    switch(method) {
    case JacobiSVD:

        typedef Eigen::JacobiSVD<MatrixType> SVDType;
        typedef Eigen::JacobiSVD<MatrixTypeDoublePrecision> SVDDoublePrecisionType;

        // We destinguish the case where we have more variables than samples and
        // the case where we have more samples than variable.
        // In the first case we compute the (smaller) inner product matrix instead of the full covariance matrix.
        // It is known that this has the same non-zero singular values as the covariance matrix.
        // Furthermore, it is possible to compute the corresponding eigenvectors of the covariance matrix from the
        // decomposition.

        if (n < p) {
            // we compute the eigenvectors of the covariance matrix by computing an SVD of the
            // n x n inner product matrix 1/(n-1) X0X0^T
            MatrixType Cov = X0 * X0.transpose() * 1.0/(n-1);
            SVDDoublePrecisionType SVD(Cov.cast<double>(), Eigen::ComputeThinV);
            VectorType singularValues = SVD.singularValues().cast<ScalarType>();
            MatrixType V = SVD.matrixV().cast<ScalarType>();

            unsigned numComponentsAboveTolerance = ((singularValues.array() - noiseVariance - Superclass::TOLERANCE) > 0).count();

            // there can be at most n-1 nonzero singular values in this case. Everything else must be due to numerical inaccuracies
            unsigned numComponentsToKeep = std::min(numComponentsAboveTolerance, n - 1);
            // compute the pseudo inverse of the square root of the singular values
            // which is then needed to recompute the PCA basis
            VectorType singSqrt = singularValues.array().sqrt();
            VectorType singSqrtInv = VectorType::Zero(singSqrt.rows());
            for (unsigned i = 0; i < numComponentsToKeep; i++) {
                assert(singSqrt(i) > Superclass::TOLERANCE);
                singSqrtInv(i) = 1.0 / singSqrt(i);
            }

            if (numComponentsToKeep == 0) {
              throw StatisticalModelException("All the eigenvalues are below the given tolerance. Model cannot be built.");
            }

            // we recover the eigenvectors U of the full covariance matrix from the eigenvectors V of the inner product matrix.
            // We use the fact that if we decompose X as X=UDV^T, then we get X^TX = UD^2U^T and XX^T = VD^2V^T (exploiting the orthogonormality
            // of the matrix U and V from the SVD). The additional factor sqrt(n-1) is to compensate for the 1/sqrt(n-1) in the formula
            // for the covariance matrix.

            MatrixType pcaBasis = X0.transpose() * V * singSqrtInv.asDiagonal();
            pcaBasis /= sqrt(n - 1.0);
            pcaBasis.conservativeResize(Eigen::NoChange, numComponentsToKeep);

            VectorType sampleVarianceVector = singularValues.topRows(numComponentsToKeep);
            VectorType pcaVariance = (sampleVarianceVector - VectorType::Ones(numComponentsToKeep) * noiseVariance);

            StatisticalModelType* model = StatisticalModelType::Create(representer, mu, pcaBasis, pcaVariance, noiseVariance);

            return model;
        } else {
            // we compute an SVD of the full p x p  covariance matrix 1/(n-1) X0^TX0 directly
            SVDType SVD(X0.transpose() * X0, Eigen::ComputeThinU);
            VectorType singularValues = SVD.singularValues();
            unsigned numComponentsToKeep = ((singularValues.array() - noiseVariance - Superclass::TOLERANCE) > 0).count();
            MatrixType pcaBasis = SVD.matrixU();
            pcaBasis /= sqrt(n - 1.0);
            pcaBasis.conservativeResize(Eigen::NoChange, numComponentsToKeep);

            if (numComponentsToKeep == 0) {
                throw StatisticalModelException("All the eigenvalues are below the given tolerance. Model cannot be built.");
            }

            VectorType sampleVarianceVector = singularValues.topRows(numComponentsToKeep);
            VectorType pcaVariance = (sampleVarianceVector - VectorType::Ones(numComponentsToKeep) * noiseVariance);
            StatisticalModelType* model = StatisticalModelType::Create(representer, mu, pcaBasis, pcaVariance, noiseVariance);
            return model;
        }
        break;

    case SelfAdjointEigenSolver: {
        // we compute the eigenvalues/eigenvectors of the full p x p  covariance matrix 1/(n-1) X0^TX0 directly

        typedef Eigen::SelfAdjointEigenSolver<MatrixType> SelfAdjointEigenSolver;
        SelfAdjointEigenSolver es;
        es.compute(X0.transpose() * X0);
        VectorType eigenValues = es.eigenvalues().reverse(); // SelfAdjointEigenSolver orders the eigenvalues in increasing order

        unsigned numComponentsToKeep = ((eigenValues.array() - noiseVariance - Superclass::TOLERANCE) > 0).count();
        MatrixType pcaBasis = es.eigenvectors().rowwise().reverse();
        pcaBasis /= sqrt(n - 1.0);
        pcaBasis.conservativeResize(Eigen::NoChange, numComponentsToKeep);

        if (numComponentsToKeep == 0) {
            throw StatisticalModelException("All the eigenvalues are below the given tolerance. Model cannot be built.");
        }

        VectorType sampleVarianceVector = eigenValues.topRows(numComponentsToKeep);
        VectorType pcaVariance = (sampleVarianceVector - VectorType::Ones(numComponentsToKeep) * noiseVariance);
        StatisticalModelType* model = StatisticalModelType::Create(representer, mu, pcaBasis, pcaVariance, noiseVariance);
        return model;
    }
    break;

    default:
        throw StatisticalModelException("Unrecognized decomposition/eigenvalue solver method.");
        return 0;
        break;
    }
}

} // namespace statismo

#endif
<|endoftext|>


<commit_before>#include <limits>
#include <stdexcept>
//test
#include <iostream>
using namespace std;
#include <fstream>
#include <stdlib.h>

#include "visual.h"
#include "bvh_node.h"

BVHNode::BVHNode()
{
	this->left = NULL;
	this->right = NULL;
}

BVHNode::~BVHNode()
{
    // TODO
}

void BVHNode::insert(Mesh const& mesh, std::vector<unsigned int>* faceIDs)
{
	this->triangles.push_back(Triangle(&mesh, 0));
	Vec3f min = this->triangles.back().getAABBMin();
	Vec3f max = this->triangles.back().getAABBMax();

	for (unsigned int var = 1; var < faceIDs->size(); ++var)
	{
		this->triangles.push_back(Triangle(&mesh, var));

		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	std::vector<Triangle> left;
	std::vector<Triangle> right;

	Vec3f middle = max;
	int axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;

	AABB half = AABB(min, middle);

	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

    this->aabb = AABB(min, max);

	if (left.size() > MAX_LEAF_TRIANGLES)
	{
		this->left = new BVHNode();
		this->left->insert(&left);
	}

	if (right.size() > MAX_LEAF_TRIANGLES)
	{
		this->right = new BVHNode();
		this->right->insert(&right);
	}
}

void BVHNode::insert(std::vector<Triangle>* triangles)
{
	if (triangles->size() == 0)
		return;

	this->triangles.push_back(triangles->back());
	Vec3f min = this->triangles.back().getAABBMin();
	Vec3f max = this->triangles.back().getAABBMax();

	triangles->pop_back();

	while(triangles->size() > 0)
	{
		this->triangles.push_back(triangles->back());
		triangles->pop_back();
		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	this->aabb = AABB(min, max);

	if (this->triangles.size() <= MAX_LEAF_TRIANGLES)
		return;

	std::vector<Triangle> left;
	std::vector<Triangle> right;

	Vec3f middle = max;
	int axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;

	AABB half = AABB(min, middle);

	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

	if (left.size() == 0)
	{
		left.push_back(right.back());
		right.pop_back();
	}

	if (right.size() == 0)
	{
		right.push_back(left.back());
		left.pop_back();
	}

	this->left = new BVHNode();
	this->left->insert(&left);

	this->right = new BVHNode();
	this->right->insert(&right);
}

bool BVHNode::intersect(Ray const& ray, Intersection* intersection) const
{
	bool hitl = false;
	bool hitr = false;

	if (this->aabb.intersect(ray))
	{
		if (left == NULL && right == NULL)
		{
			for (unsigned int var = 0; var < triangles.size(); ++var)
			{
				if (triangles.at(var).intersect(ray, intersection))
					return true;
			}

		}
		else
		{
			hitl=left->intersect(ray, intersection);
			hitr=right->intersect(ray, intersection);
			if (hitl)
				return true;

			if (hitr)
				return true;
		}

	}
    return false;
}
<commit_msg>refactored intersect<commit_after>#include <limits>
#include <stdexcept>

#include <iostream>
#include <fstream>
#include <stdlib.h>

#include "visual.h"
#include "bvh_node.h"

using namespace std;

BVHNode::BVHNode()
{
	this->left = NULL;
	this->right = NULL;
}

BVHNode::~BVHNode()
{
    // TODO
}

void BVHNode::insert(Mesh const& mesh, std::vector<unsigned int>* faceIDs)
{
	this->triangles.push_back(Triangle(&mesh, 0));
	Vec3f min = this->triangles.back().getAABBMin();
	Vec3f max = this->triangles.back().getAABBMax();

	for (unsigned int var = 1; var < faceIDs->size(); ++var)
	{
		this->triangles.push_back(Triangle(&mesh, var));

		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	std::vector<Triangle> left;
	std::vector<Triangle> right;

	Vec3f middle = max;
	int axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;

	AABB half = AABB(min, middle);

	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

    this->aabb = AABB(min, max);

	if (left.size() > MAX_LEAF_TRIANGLES)
	{
		this->left = new BVHNode();
		this->left->insert(&left);
	}

	if (right.size() > MAX_LEAF_TRIANGLES)
	{
		this->right = new BVHNode();
		this->right->insert(&right);
	}
}

void BVHNode::insert(std::vector<Triangle>* triangles)
{
	this->triangles.push_back(triangles->back());
	Vec3f min = this->triangles.back().getAABBMin();
	Vec3f max = this->triangles.back().getAABBMax();

	triangles->pop_back();

	while(triangles->size() > 0)
	{
		this->triangles.push_back(triangles->back());
		triangles->pop_back();
		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	this->aabb = AABB(min, max);

	if (this->triangles.size() <= MAX_LEAF_TRIANGLES)
		return;

	std::vector<Triangle> left;
	std::vector<Triangle> right;

	Vec3f middle = max;
	int axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;

	AABB half = AABB(min, middle);

	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

	if (left.size() == 0)
	{
		left.push_back(right.back());
		right.pop_back();
	}

	if (right.size() == 0)
	{
		right.push_back(left.back());
		left.pop_back();
	}

	this->left = new BVHNode();
	this->left->insert(&left);

	this->right = new BVHNode();
	this->right->insert(&right);
}

bool BVHNode::intersect(Ray const& ray, Intersection* intersection) const
{
	bool hitl = false;
	bool hitr = false;

	if (this->aabb.intersect(ray))
	{
		if (left == NULL && right == NULL)
		{
			for (unsigned int var = 0; var < triangles.size(); ++var)
				if (triangles.at(var).intersect(ray, intersection))
					return true;
		}
		else
		{
			hitl = left->intersect(ray, intersection);
			hitr = right->intersect(ray, intersection);

			return hitl || hitr;
		}

	}
    return false;
}
<|endoftext|>


<commit_before>#include <limits>
#include <stdexcept>

#include <iostream>
#include <fstream>
#include <stdlib.h>

#include "visual.h"
#include "bvh_node.h"

using namespace std;

BVHNode::BVHNode()
{
	this->left = NULL;
	this->right = NULL;
}

BVHNode::~BVHNode()
{
    // TODO
}

void BVHNode::insert(Mesh const& mesh, std::vector<unsigned int>* faceIDs)
{
	this->triangles.push_back(Triangle(&mesh, 0));
	Vec3f min = this->triangles.back().getAABBMin();
	Vec3f max = this->triangles.back().getAABBMax();

	for (unsigned int var = 1; var < faceIDs->size(); ++var)
	{
		this->triangles.push_back(Triangle(&mesh, var));

		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	std::vector<Triangle> left;
	std::vector<Triangle> right;

	Vec3f middle = max;
	int axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;

	AABB half = AABB(min, middle);

	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

    this->aabb = AABB(min, max);

	if (left.size() > MAX_LEAF_TRIANGLES)
	{
		this->left = new BVHNode();
		this->left->insert(&left);
	}

	if (right.size() > MAX_LEAF_TRIANGLES)
	{
		this->right = new BVHNode();
		this->right->insert(&right);
	}
}

void BVHNode::insert(std::vector<Triangle>* triangles)
{
	this->triangles.push_back(triangles->back());
	Vec3f min = this->triangles.back().getAABBMin();
	Vec3f max = this->triangles.back().getAABBMax();

	triangles->pop_back();

	while(triangles->size() > 0)
	{
		this->triangles.push_back(triangles->back());
		triangles->pop_back();
		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	this->aabb = AABB(min, max);

	if (this->triangles.size() <= MAX_LEAF_TRIANGLES)
		return;

	std::vector<Triangle> left;
	std::vector<Triangle> right;

	Vec3f middle = max;
	int axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;

	AABB half = AABB(min, middle);

	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

	if (left.size() == 0)
	{
		left.push_back(right.back());
		right.pop_back();
	}

	if (right.size() == 0)
	{
		right.push_back(left.back());
		left.pop_back();
	}

	this->left = new BVHNode();
	this->left->insert(&left);

	this->right = new BVHNode();
	this->right->insert(&right);
}

bool BVHNode::intersect(Ray const& ray, Intersection* intersection) const
{
	bool hitl = false;
	bool hitr = false;

	if (this->aabb.intersect(ray))
	{
		if (left == NULL && right == NULL)
		{
			for (unsigned int var = 0; var < triangles.size(); ++var)
				if (triangles.at(var).intersect(ray, intersection))
					return true;
		}
		else
		{
			hitl = left->intersect(ray, intersection);
			hitr = right->intersect(ray, intersection);

			return hitl || hitr;
		}

	}
    return false;
}
<commit_msg>comments added<commit_after>#include <limits>
#include <stdexcept>

#include <iostream>
#include <fstream>
#include <stdlib.h>

#include "visual.h"
#include "bvh_node.h"

using namespace std;

BVHNode::BVHNode()
{
	this->left = NULL;
	this->right = NULL;
}

BVHNode::~BVHNode()
{
    // TODO
}

void BVHNode::insert(Mesh const& mesh, std::vector<unsigned int>* faceIDs)
{
	std::vector<Triangle> left;
	std::vector<Triangle> right;
	Vec3f min;
	Vec3f max;
	Vec3f middle;
	int axis;
	AABB half;

	//erstelle erste Dreieck und setzte Baseline für min/max des AABB
	this->triangles.push_back(Triangle(&mesh, 0));
	min = this->triangles.back().getAABBMin();
	max = this->triangles.back().getAABBMax();

	//füge alle Dreiecke ein und prüfe ob neues min/max gefunden wurde
	for (unsigned int var = 1; var < faceIDs->size(); ++var)
	{
		this->triangles.push_back(Triangle(&mesh, var));

		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

    this->aabb = AABB(min, max);

	//prüfe welche Achse die längste ist und teile AABB in der Mitte an dieser
	middle = max;
	axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;
	half = AABB(min, middle);

	//teile Dreiecke in Dreiecke die im neuen AABB sind und Dreicke die im anderen AABB sind
	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

	//erstelle Kinder
	if (left.size() > MAX_LEAF_TRIANGLES)
	{
		this->left = new BVHNode();
		this->left->insert(&left);}

	if (right.size() > MAX_LEAF_TRIANGLES)
	{
		this->right = new BVHNode();
		this->right->insert(&right);}
}

void BVHNode::insert(std::vector<Triangle>* triangles)
{
	std::vector<Triangle> left,right;
	Vec3f min,max,middle;
	AABB half;
	int axis;

	//prüfe Triangle größe
	if (triangles->size() == 0)
		return;

	//erstelle erste Dreieck und setzte Baseline für min/max des AABB
	this->triangles.push_back(triangles->back());
	min = this->triangles.back().getAABBMin();
	max = this->triangles.back().getAABBMax();
	triangles->pop_back();

	//füge alle übergebenen Dreiecke ein und prüfe ob neues min/max gefunden wurde
	while(triangles->size() > 0)
	{
		this->triangles.push_back(triangles->back());
		triangles->pop_back();
		for (int i = 0; i < 3; i++)
		{
			if (min[i] > this->triangles.back().getAABBMin()[i])
				min[i] = this->triangles.back().getAABBMin()[i];
			if (max[i] < this->triangles.back().getAABBMax()[i])
				max[i] = this->triangles.back().getAABBMax()[i];
		}
	}

	//erstelle AAB und prüfe ob der aktuelle Knoten
	this->aabb = AABB(min, max);
	if (this->triangles.size() <= MAX_LEAF_TRIANGLES)
		return;

	//prüfe welche Achse die längste ist und teile AABB in der Mitte an dieser
	middle = max;
	axis = aabb.getLongestAxis();
	middle[axis] = (min[axis] + max[axis]) / 2;
	half = AABB(min, middle);

	//teile Dreiecke in Dreiecke die im neuen AABB sind und Dreicke die im anderen AABB sind
	while (this->triangles.size() > 0)
	{
		if (half.inside(this->triangles.back().getCentroid()))
			left.push_back(this->triangles.back());
		else
			right.push_back(this->triangles.back());

		this->triangles.pop_back();
	}

	//prüfe ob trotz des Versuchs aufteilen kein Dreieck in
	if (left.size() == 0)
	{
		left.push_back(right.back());
		right.pop_back();}
	if (right.size() == 0)
	{
		right.push_back(left.back());
		left.pop_back();}

	//erstelle Kinder
	this->left = new BVHNode();
	this->left->insert(&left);

	this->right = new BVHNode();
	this->right->insert(&right);
}

bool BVHNode::intersect(Ray const& ray, Intersection* intersection) const
{
	bool hitl = false;
	bool hitr = false;

	//prüft ob der Strahl den aktuellen Knoten/dessen Kinder trifft
	if (this->aabb.intersect(ray))
	{
		//Blatt, prüfe für jedes Dreieck ob der Strahl es trifft
		if (left == NULL && right == NULL)
		{
			for (unsigned int var = 0; var < triangles.size(); ++var)
				if (triangles.at(var).intersect(ray, intersection))
					return true;
		}
		//kein Blatt, rek. Aufruf
		else
		{
			hitl = left->intersect(ray, intersection);
			hitr = right->intersect(ray, intersection);

			return hitl || hitr;
		}

	}
    return false;
}
<|endoftext|>