<commit_before>#include <cassert>
#include <fstream>
#include <iostream>
#include <memory>
#include <set>
#include <string>

#include <gemmi/cif.hpp>
#include <gemmi/mmcif.hpp>

#define GEMMI_WRITE_IMPLEMENTATION
#include <gemmi/to_mmcif.hpp>
#undef GEMMI_WRITE_IMPLEMENTATION

#include "cif.hh"
#include "freesasa.h"

struct ModelDiscriminator {
    ModelDiscriminator(const std::string &model_name,
                       const int model_col = 11)
        : _model_name(model_name), _model_col(model_col)
    {
    }

    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        return _model_name != site[_model_col];
    }

private:
    const std::string _model_name;
    int _model_col;
};

struct ModelSetDiscriminator {
    ModelSetDiscriminator(const std::set<int> models,
                          const int model_col = 11)
        : _models(models), _model_col(model_col)
    {
    }

    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        return _models.count(std::stoi(site[_model_col])) == 0;
    }

private:
    const std::set<int> _models;
    int _model_col;
};

struct ChainDiscriminator {
    ChainDiscriminator(const std::string &model_name, const std::string &chain_name,
                       const int model_col = 11, const int chain_col = 1)
        : _model_name(model_name), _chain_name(chain_name),
          _model_col(model_col), _chain_col(chain_col)
    {
    }

    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        return _model_name != site[_model_col] || _chain_name != site[_chain_col];
    }

private:
    const std::string _model_name, _chain_name;
    int _model_col, _chain_col;
};

static std::unique_ptr<std::set<int>>
get_models(const gemmi::cif::Document &doc)
{
    auto models = std::make_unique<std::set<int>>();
    for (auto block : doc.blocks) {
        for (auto site : block.find("_atom_site.", {"pdbx_PDB_model_num"})) {
            models->insert(gemmi::cif::as_int(site[0]));
        }
    }
    return models;
}

static std::unique_ptr<std::set<std::string>>
get_chains(const gemmi::cif::Document &doc)
{
    auto chains = std::make_unique<std::set<std::string>>();
    for (auto block : doc.blocks) {
        for (auto site : block.find("_atom_site.", {"auth_asym_id"})) {
            chains->insert(site[0]);
        }
    }
    return chains;
}

static std::unique_ptr<std::set<std::string>>
get_chains(const gemmi::Model &model)
{
    auto chains = std::make_unique<std::set<std::string>>();

    for (auto &chain : model.chains) {
        chains->insert(chain.name);
    }

    return chains;
}

static const auto atom_site_columns = std::vector<std::string>({
    "group_PDB",
    "auth_asym_id",
    "auth_seq_id",
    "pdbx_PDB_ins_code",
    "auth_comp_id",
    "auth_atom_id",
    "label_alt_id",
    "type_symbol",
    "Cartn_x",
    "Cartn_y",
    "Cartn_z",
    "pdbx_PDB_model_num",
});

static freesasa_cif_atom
freesasa_atom_from_site(const gemmi::cif::Table::Row &site)
{

    std::unique_ptr<std::string> auth_atom_id;
    // remove quotation marks if necessary
    if (site[5][0] == '"') {
        auth_atom_id = std::make_unique<std::string>(site[5].substr(1, site[5].size() - 2));
    }

    return {
        .group_PDB = site[0].c_str(),
        .auth_asym_id = site[1][0],
        .auth_seq_id = site[2].c_str(),
        .pdbx_PDB_ins_code = site[3].c_str(),
        .auth_comp_id = site[4].c_str(),
        .auth_atom_id = auth_atom_id ? std::move(*auth_atom_id).c_str() : site[5].c_str(),
        .label_alt_id = site[6].c_str(),
        .type_symbol = site[7].c_str(),
        .Cartn_x = atof(site[8].c_str()),
        .Cartn_y = atof(site[9].c_str()),
        .Cartn_z = atof(site[10].c_str())};
}

template <typename T>
static freesasa_structure *
freesasa_structure_from_pred(const gemmi::cif::Document &doc,
                             const T &discriminator,
                             const freesasa_classifier *classifier,
                             int structure_options)
{
    freesasa_structure *structure = freesasa_structure_new();
<<<<<<< HEAD
    std::string auth_atom_id;
    char prevAltId = '.';
=======
    // std::string auth_atom_id;  // Should be deleted due to refactoring.
>>>>>>> 59a514a... Successfully parsing tree structure and input cif file; matching tree atom to cif atom.

    for (auto block : doc.blocks) {
        for (auto site : block.find("_atom_site.", atom_site_columns)) {
            if (site[0] != "ATOM" && !(structure_options & FREESASA_INCLUDE_HETATM)) {
                continue;
            }

            if (discriminator(site)) continue;

            freesasa_cif_atom atom = freesasa_atom_from_site(site);

            if (!(structure_options & FREESASA_INCLUDE_HYDROGEN) && std::string(atom.type_symbol) == "H") {
                continue;
            }

            // Pick the first alternative conformation for an atom
            auto currentAltId = site[6][0];
            if ((currentAltId != '.' && prevAltId == '.') || currentAltId == '.') {
                prevAltId = currentAltId;
            } else if (currentAltId != '.' && currentAltId != prevAltId) {
                continue;
            }

            freesasa_structure_add_cif_atom(structure, &atom, classifier, structure_options);
        }
    }
    return structure;
}

freesasa_structure *
freesasa_structure_from_cif(std::FILE *input,
                            const freesasa_classifier *classifier,
                            int structure_options)
{
    const auto doc = gemmi::cif::read_cstream(input, 8192, "cif-input");
    const auto models = get_models(doc);

    std::unique_ptr<const ModelSetDiscriminator> discriminator;
    if (structure_options & FREESASA_JOIN_MODELS) {
        discriminator = std::make_unique<const ModelSetDiscriminator>(std::move(*models));
    } else {
        auto firstModel = models->begin();
        auto singleModel = std::set<int>{*firstModel};
        discriminator = std::make_unique<const ModelSetDiscriminator>(singleModel);
    }
    return freesasa_structure_from_pred(doc, *discriminator, classifier, structure_options);
}

freesasa_structure *
freesasa_structure_from_model(const gemmi::cif::Document &doc,
                              const std::string &model_name,
                              const freesasa_classifier *classifier,
                              int structure_options)
{
    const ModelDiscriminator discriminator(model_name);
    return freesasa_structure_from_pred(doc, discriminator, classifier, structure_options);
    freesasa_structure *structure = freesasa_structure_new();
}

freesasa_structure *
freesasa_structure_from_chain(const gemmi::cif::Document doc,
                              const std::string &model_name,
                              const std::string &chain_name,
                              const freesasa_classifier *classifier,
                              int structure_options)
{
    const ChainDiscriminator discriminator(model_name, chain_name);
    return freesasa_structure_from_pred(doc, discriminator, classifier, structure_options);
}

std::vector<freesasa_structure *>
freesasa_cif_structure_array(std::FILE *input,
                             int *n,
                             const freesasa_classifier *classifier,
                             int options)
{
    int n_models = 0, n_chains = 0;

    std::vector<freesasa_structure *> ss;

    const auto doc = gemmi::cif::read_cstream(input, 8192, "cif-input");

    gemmi::Structure gemmi_struct = gemmi::make_structure_from_block(doc.blocks[0]);

    const auto models = gemmi_struct.models;

    n_models = models.size();

    /* only keep first model if option not provided */
    if (!(options & FREESASA_SEPARATE_MODELS)) n_models = 1;

    /* for each model read chains if requested */
    if (options & FREESASA_SEPARATE_CHAINS) {
        for (int i = 0; i < n_models; ++i) {
            auto chain_names = get_chains(models[i]);
            int n_new_chains = chain_names->size();
            n_chains += n_new_chains;

            if (n_new_chains == 0) {
                freesasa_warn("in %s(): no chains found (in model %s)", __func__, models[i].name.c_str());
                continue;
            }

            ss.reserve(n_new_chains);
            for (auto &chain_name : *chain_names) {
                ss.emplace_back(
                    freesasa_structure_from_chain(doc, models[i].name, chain_name, classifier, options));
                freesasa_structure_set_model(ss.back(), i + 1);
            }
        }
        if (n_chains == 0) freesasa_fail("In %s(): No chains in any model in protein: %s.", __func__, gemmi_struct.name.c_str());
        *n = n_chains;
    } else {
        ss.reserve(n_models);
        for (int i = 0; i < n_models; ++i) {
            ss.emplace_back(
                freesasa_structure_from_model(doc, models[i].name, classifier, options));
            freesasa_structure_set_model(ss.back(), i + 1);
        }
        *n = n_models;
    }
    return ss;
}

struct freesasa_MCRA{

    freesasa_MCRA(const int model, const std::string &chain, const std::string &residue, const std::string &res_num, const std::string &atom)
    : _model(model), _chain(chain), _residue(residue), _res_num(res_num), _atom(atom)
    {}

    int find_row(gemmi::cif::Table &table, int start_idx=0) const {
        int idx = start_idx, total_rows = table.length();

        assert(idx < total_rows);

        while (idx < total_rows - 1){
            if ((*this)(table[idx])) {
                return idx;
            }
            ++idx;
        }

        std::cout << "Did not find row for: " << *this << std::endl;
        std::cout << "Looping through entire table... " << std::endl;

        idx = 0;
        for (auto site : table) {
            if ((*this)(site)) {
                return idx;
            }
            ++idx;
        }
        std::cout << "Still couldnt find: " << *this << std::endl;
        return -1;
    }

    friend std::ostream& operator << (std::ostream& os, const freesasa_MCRA& mcra){
        os << "Atom(" << mcra._model << " " << mcra._chain << " "  << mcra._residue << " [" << mcra._res_num << "] " << mcra._atom << ")";
        return os;
    }

private:
    bool operator () (const gemmi::cif::Table::Row &site) const {

        int model                  = std::stoul(site[11]);
        const std::string &chain   = site[1];
        const std::string &residue = site[4];
        const std::string &res_num = site[2];
        const std::string &atom    = site[5][0] != '"' ? site[5] : site[5].substr(1, site[5].size() - 2);

        if (_model == model){
            if (_chain == chain){
                if (_res_num == res_num && _residue == residue){
                    if (_atom == atom) {
                        // std::cout << "Found! " << *this << std::endl;
                        return true;
                    }
                }
            }
        }
        return false;

    }

    const int _model;
    const std::string &_chain, &_residue, &_res_num,&_atom;
};

int freesasa_write_cif(std::FILE *output,
                       freesasa_node *root,
                       int options)
{
    freesasa_node *result{freesasa_node_children(root)};
    freesasa_node *structure;

    assert(output);
    assert(root);
    assert(freesasa_node_type(root) == FREESASA_NODE_ROOT);

    std::string inp_file{freesasa_node_name(result)};
    std::string out_file = inp_file.substr(0, inp_file.find_last_of(".")) + ".sasa.cif";
    std::cout << "Input file: " << inp_file << std::endl << "Output File: " << out_file << std::endl;

    const auto doc = gemmi::cif::read_file(inp_file);
    auto block     = doc.sole_block();
    auto table     = block.find("_atom_site.", atom_site_columns);
    std::cout << "Number of rows in the table: " << table.length() << std::endl;

    while (result) {
        structure = freesasa_node_children(result);
        std::cout << "New Result!" << std::endl;
        while (structure) {
            int rowNum = 0, idx = 0;
            auto model = freesasa_node_structure_model(structure);
            auto chain = freesasa_node_children(structure);
            std::cout << "New Structure!" << std::endl;
            while (chain) {
                auto cName = freesasa_node_name(chain);
                auto residue = freesasa_node_children(chain);
                while (residue) {
                    auto rName = freesasa_node_name(residue);
                    auto rNum  = freesasa_node_residue_number(residue);
                    auto atom = freesasa_node_children(residue);
                    while (atom) {
                        auto aName   = freesasa_node_name(atom);
                        auto area    = freesasa_node_area(atom);
                        auto radius  = freesasa_node_atom_radius(atom);

                        idx = freesasa_MCRA{model, cName, rName, rNum, aName}.find_row(table, rowNum);
                        rowNum = idx != -1 ? idx : rowNum;
                        if (idx == -1) exit(1);

                        atom = freesasa_node_next(atom);
                    }
                    auto last_res_name = freesasa_node_name(residue);
                    auto last_res_number = freesasa_node_residue_number(residue);
                    residue = freesasa_node_next(residue);
                }
                auto last_chain = freesasa_node_name(chain);
                chain = freesasa_node_next(chain);
                std::cout << "Finished chain: " << cName << std::endl;
            }
            structure = freesasa_node_next(structure);
        }
        result = freesasa_node_next(result);
    }
    return FREESASA_SUCCESS;
}
<commit_msg>Cleaned up code; prepping it for adding the sasa and radius columns to the cif table.<commit_after>#include <algorithm>
#include <cassert>
#include <fstream>
#include <iostream>
#include <memory>
#include <set>
#include <string>

#include <gemmi/cif.hpp>
#include <gemmi/mmcif.hpp>

#define GEMMI_WRITE_IMPLEMENTATION
#include <gemmi/to_mmcif.hpp>
#undef GEMMI_WRITE_IMPLEMENTATION

#include "cif.hh"
#include "freesasa.h"

struct ModelDiscriminator {
    ModelDiscriminator(const std::string &model_name,
                       const int model_col = 11)
        : _model_name(model_name), _model_col(model_col)
    {
    }

    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        return _model_name != site[_model_col];
    }

private:
    const std::string _model_name;
    int _model_col;
};

struct ModelSetDiscriminator {
    ModelSetDiscriminator(const std::set<int> models,
                          const int model_col = 11)
        : _models(models), _model_col(model_col)
    {
    }

    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        return _models.count(std::stoi(site[_model_col])) == 0;
    }

private:
    const std::set<int> _models;
    int _model_col;
};

struct ChainDiscriminator {
    ChainDiscriminator(const std::string &model_name, const std::string &chain_name,
                       const int model_col = 11, const int chain_col = 1)
        : _model_name(model_name), _chain_name(chain_name),
          _model_col(model_col), _chain_col(chain_col)
    {
    }

    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        return _model_name != site[_model_col] || _chain_name != site[_chain_col];
    }

private:
    const std::string _model_name, _chain_name;
    int _model_col, _chain_col;
};

static std::unique_ptr<std::set<int>>
get_models(const gemmi::cif::Document &doc)
{
    auto models = std::make_unique<std::set<int>>();
    for (auto block : doc.blocks) {
        for (auto site : block.find("_atom_site.", {"pdbx_PDB_model_num"})) {
            models->insert(gemmi::cif::as_int(site[0]));
        }
    }
    return models;
}

static std::unique_ptr<std::set<std::string>>
get_chains(const gemmi::cif::Document &doc)
{
    auto chains = std::make_unique<std::set<std::string>>();
    for (auto block : doc.blocks) {
        for (auto site : block.find("_atom_site.", {"auth_asym_id"})) {
            chains->insert(site[0]);
        }
    }
    return chains;
}

static std::unique_ptr<std::set<std::string>>
get_chains(const gemmi::Model &model)
{
    auto chains = std::make_unique<std::set<std::string>>();

    for (auto &chain : model.chains) {
        chains->insert(chain.name);
    }

    return chains;
}

static const auto atom_site_columns = std::vector<std::string>({
    "group_PDB",
    "auth_asym_id",
    "auth_seq_id",
    "pdbx_PDB_ins_code",
    "auth_comp_id",
    "auth_atom_id",
    "label_alt_id",
    "type_symbol",
    "Cartn_x",
    "Cartn_y",
    "Cartn_z",
    "pdbx_PDB_model_num",
});

static freesasa_cif_atom
freesasa_atom_from_site(const gemmi::cif::Table::Row &site)
{

    std::unique_ptr<std::string> auth_atom_id;
    // remove quotation marks if necessary
    if (site[5][0] == '"') {
        auth_atom_id = std::make_unique<std::string>(site[5].substr(1, site[5].size() - 2));
    }

    return {
        .group_PDB = site[0].c_str(),
        .auth_asym_id = site[1][0],
        .auth_seq_id = site[2].c_str(),
        .pdbx_PDB_ins_code = site[3].c_str(),
        .auth_comp_id = site[4].c_str(),
        .auth_atom_id = auth_atom_id ? std::move(*auth_atom_id).c_str() : site[5].c_str(),
        .label_alt_id = site[6].c_str(),
        .type_symbol = site[7].c_str(),
        .Cartn_x = atof(site[8].c_str()),
        .Cartn_y = atof(site[9].c_str()),
        .Cartn_z = atof(site[10].c_str())};
}

template <typename T>
static freesasa_structure *
freesasa_structure_from_pred(const gemmi::cif::Document &doc,
                             const T &discriminator,
                             const freesasa_classifier *classifier,
                             int structure_options)
{
    freesasa_structure *structure = freesasa_structure_new();
<<<<<<< HEAD
    std::string auth_atom_id;
    char prevAltId = '.';
=======
    // std::string auth_atom_id;  // Should be deleted due to refactoring.
>>>>>>> 59a514a... Successfully parsing tree structure and input cif file; matching tree atom to cif atom.

    for (auto block : doc.blocks) {
        for (auto site : block.find("_atom_site.", atom_site_columns)) {
            if (site[0] != "ATOM" && !(structure_options & FREESASA_INCLUDE_HETATM)) {
                continue;
            }

            if (discriminator(site)) continue;

            freesasa_cif_atom atom = freesasa_atom_from_site(site);

            if (!(structure_options & FREESASA_INCLUDE_HYDROGEN) && std::string(atom.type_symbol) == "H") {
                continue;
            }

            // Pick the first alternative conformation for an atom
            auto currentAltId = site[6][0];
            if ((currentAltId != '.' && prevAltId == '.') || currentAltId == '.') {
                prevAltId = currentAltId;
            } else if (currentAltId != '.' && currentAltId != prevAltId) {
                continue;
            }

            freesasa_structure_add_cif_atom(structure, &atom, classifier, structure_options);
        }
    }
    return structure;
}

freesasa_structure *
freesasa_structure_from_cif(std::FILE *input,
                            const freesasa_classifier *classifier,
                            int structure_options)
{
    const auto doc = gemmi::cif::read_cstream(input, 8192, "cif-input");
    const auto models = get_models(doc);

    std::unique_ptr<const ModelSetDiscriminator> discriminator;
    if (structure_options & FREESASA_JOIN_MODELS) {
        discriminator = std::make_unique<const ModelSetDiscriminator>(std::move(*models));
    } else {
        auto firstModel = models->begin();
        auto singleModel = std::set<int>{*firstModel};
        discriminator = std::make_unique<const ModelSetDiscriminator>(singleModel);
    }
    return freesasa_structure_from_pred(doc, *discriminator, classifier, structure_options);
}

freesasa_structure *
freesasa_structure_from_model(const gemmi::cif::Document &doc,
                              const std::string &model_name,
                              const freesasa_classifier *classifier,
                              int structure_options)
{
    const ModelDiscriminator discriminator(model_name);
    return freesasa_structure_from_pred(doc, discriminator, classifier, structure_options);
    freesasa_structure *structure = freesasa_structure_new();
}

freesasa_structure *
freesasa_structure_from_chain(const gemmi::cif::Document doc,
                              const std::string &model_name,
                              const std::string &chain_name,
                              const freesasa_classifier *classifier,
                              int structure_options)
{
    const ChainDiscriminator discriminator(model_name, chain_name);
    return freesasa_structure_from_pred(doc, discriminator, classifier, structure_options);
}

std::vector<freesasa_structure *>
freesasa_cif_structure_array(std::FILE *input,
                             int *n,
                             const freesasa_classifier *classifier,
                             int options)
{
    int n_models = 0, n_chains = 0;

    std::vector<freesasa_structure *> ss;

    const auto doc = gemmi::cif::read_cstream(input, 8192, "cif-input");

    gemmi::Structure gemmi_struct = gemmi::make_structure_from_block(doc.blocks[0]);

    const auto models = gemmi_struct.models;

    n_models = models.size();

    /* only keep first model if option not provided */
    if (!(options & FREESASA_SEPARATE_MODELS)) n_models = 1;

    /* for each model read chains if requested */
    if (options & FREESASA_SEPARATE_CHAINS) {
        for (int i = 0; i < n_models; ++i) {
            auto chain_names = get_chains(models[i]);
            int n_new_chains = chain_names->size();
            n_chains += n_new_chains;

            if (n_new_chains == 0) {
                freesasa_warn("in %s(): no chains found (in model %s)", __func__, models[i].name.c_str());
                continue;
            }

            ss.reserve(n_new_chains);
            for (auto &chain_name : *chain_names) {
                ss.emplace_back(
                    freesasa_structure_from_chain(doc, models[i].name, chain_name, classifier, options));
                freesasa_structure_set_model(ss.back(), i + 1);
            }
        }
        if (n_chains == 0) freesasa_fail("In %s(): No chains in any model in protein: %s.", __func__, gemmi_struct.name.c_str());
        *n = n_chains;
    } else {
        ss.reserve(n_models);
        for (int i = 0; i < n_models; ++i) {
            ss.emplace_back(
                freesasa_structure_from_model(doc, models[i].name, classifier, options));
            freesasa_structure_set_model(ss.back(), i + 1);
        }
        *n = n_models;
    }
    return ss;
}

struct freesasa_MCRA {

    freesasa_MCRA(const int model, const std::string &chain, const std::string &residue, const std::string &res_num, const std::string &atom)
        : _model(model), _chain(chain), _residue(residue), _res_num(res_num), _atom(atom)
    {
    }

    int find_row(gemmi::cif::Table &table, int start_idx = 0) const
    {
        int idx = start_idx, total_rows = table.length();

        assert(idx < total_rows);

        while (idx < total_rows) {
            if ((*this)(table[idx])) {
                return idx;
            }
            ++idx;
        }

        std::cout << "Did not find row for: " << *this << std::endl;
        std::cout << "Looping through entire table... " << std::endl;

        idx = 0;
        for (auto site : table) {
            if ((*this)(site)) {
                return idx;
            }
            ++idx;
        }
        std::cout << "Still couldnt find: " << *this << std::endl;
        return -1;
    }

    friend std::ostream &operator<<(std::ostream &os, const freesasa_MCRA &mcra)
    {
        os << "Atom(" << mcra._model << " " << mcra._chain << " " << mcra._residue << " [" << mcra._res_num << "] " << mcra._atom << ")";
        return os;
    }

private:
    bool operator()(const gemmi::cif::Table::Row &site) const
    {
        int model = std::stoul(site[11]);
        const std::string &chain = site[1];
        const std::string &residue = site[4];
        const std::string &res_num = site[2];
        const std::string &atom = site[5][0] != '"' ? site[5] : site[5].substr(1, site[5].size() - 2);

        if (_model == model) {
            if (_chain == chain) {
                if (_res_num == res_num && _residue == residue) {
                    if (_atom == atom) {
                        // std::cout << "Found! " << *this << std::endl;
                        return true;
                    }
                }
            }
        }
        return false;
    }

    const int _model;
    const std::string &_chain, &_residue, &_res_num, &_atom;
};

int freesasa_write_cif(std::FILE *output,
                       freesasa_node *root,
                       int options)
{
    freesasa_node *result{freesasa_node_children(root)};
    freesasa_node *structure;

    assert(output);
    assert(root);
    assert(freesasa_node_type(root) == FREESASA_NODE_ROOT);

    std::string inp_file{freesasa_node_name(result)};
    std::string out_file = inp_file.substr(0, inp_file.find_last_of(".")) + ".sasa.cif";
    std::cout << "Input file: " << inp_file << std::endl
              << "Output File: " << out_file << std::endl;

    const auto doc = gemmi::cif::read_file(inp_file);
    auto block = doc.sole_block();
    auto table = block.find("_atom_site.", atom_site_columns);
    auto &loop = *table.get_loop();
    std::cout << "Number of rows x columns in the table: " << table.length() << " x " << loop.tags.size() << std::endl;

    // TODO allocate new vector<vector<string>> with the proper number of columns

    std::vector<double> sasa_vals, sasa_radii;
    sasa_vals.reserve(table.length()), sasa_radii.reserve(table.length());
    while (result) {
        structure = freesasa_node_children(result);
        std::cout << "New Result!" << std::endl;
        while (structure) {
            int rowNum = 0, idx = 0;
            auto model = freesasa_node_structure_model(structure);
            auto chain = freesasa_node_children(structure);
            std::cout << "New Structure!" << std::endl;
            while (chain) {
                auto cName = freesasa_node_name(chain);
                auto residue = freesasa_node_children(chain);
                while (residue) {
                    auto rName = freesasa_node_name(residue);
                    auto rNum = freesasa_node_residue_number(residue);
                    auto atom = freesasa_node_children(residue);
                    while (atom) {
                        auto aName = freesasa_node_name(atom);
                        auto area = freesasa_node_area(atom);
                        auto radius = freesasa_node_atom_radius(atom);

                        idx = freesasa_MCRA{model, cName, rName, rNum, aName}.find_row(table, rowNum);
                        if (idx != -1) {
                            rowNum = idx;
                            sasa_vals[rowNum] = area->total;
                            sasa_radii[rowNum] = radius;
                            std::cout << "Atom FREESASA Value and Radius: " << area->total << " " << radius << std::endl;
                        } else {
                            exit(1);
                        }
                        atom = freesasa_node_next(atom);
                    }
                    auto last_res_name = freesasa_node_name(residue);
                    auto last_res_number = freesasa_node_residue_number(residue);
                    residue = freesasa_node_next(residue);
                }
                auto last_chain = freesasa_node_name(chain);
                chain = freesasa_node_next(chain);
                std::cout << "Finished chain: " << cName << std::endl;
            }
            structure = freesasa_node_next(structure);
        }
        result = freesasa_node_next(result);
    }
    std::cout << "Max freesasa value: " << sasa_vals[std::distance(sasa_vals.begin(), std::max_element(sasa_vals.begin(), sasa_vals.end()))] << std::endl;
    return FREESASA_SUCCESS;
}
<|endoftext|>


<commit_before>/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 *     Copyright 2012 Couchbase, Inc
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

#include "config.h"

#include "common.h"
#include "couch-kvstore/couch-fs-stats.h"
#include "histo.h"

extern "C" {
static couch_file_handle cfs_construct(couchstore_error_info_t*, void* cookie);
static couchstore_error_t cfs_open(couchstore_error_info_t*,
                                   couch_file_handle*, const char*, int);
static void cfs_close(couchstore_error_info_t*, couch_file_handle);
static ssize_t cfs_pread(couchstore_error_info_t*, couch_file_handle,
                         void *, size_t, cs_off_t);
static ssize_t cfs_pwrite(couchstore_error_info_t*, couch_file_handle,
                          const void *, size_t, cs_off_t);
static cs_off_t cfs_goto_eof(couchstore_error_info_t*, couch_file_handle);
static couchstore_error_t cfs_sync(couchstore_error_info_t*, couch_file_handle);
static couchstore_error_t cfs_advise(couchstore_error_info_t*,
                                     couch_file_handle,
                                     cs_off_t, cs_off_t,
                                     couchstore_file_advice_t);
static void cfs_destroy(couchstore_error_info_t*,couch_file_handle);
}

couch_file_ops getCouchstoreStatsOps(CouchstoreStats* stats) {
    couch_file_ops ops = {
        5,
        cfs_construct,
        cfs_open,
        cfs_close,
        cfs_pread,
        cfs_pwrite,
        cfs_goto_eof,
        cfs_sync,
        cfs_advise,
        cfs_destroy,
        stats
    };
    return ops;
}

struct StatFile {
    const couch_file_ops* orig_ops;
    couch_file_handle orig_handle;
    CouchstoreStats* stats;
    cs_off_t last_offs;
};

extern "C" {
    static couch_file_handle cfs_construct(couchstore_error_info_t *errinfo,
                                           void* cookie) {
        StatFile* sf = new StatFile;
        sf->stats = static_cast<CouchstoreStats*>(cookie);
        sf->orig_ops = couchstore_get_default_file_ops();
        sf->orig_handle = sf->orig_ops->constructor(errinfo,
                                                    sf->orig_ops->cookie);
        sf->last_offs = 0;
        return reinterpret_cast<couch_file_handle>(sf);
    }

    static couchstore_error_t cfs_open(couchstore_error_info_t *errinfo,
                                       couch_file_handle* h,
                                       const char* path,
                                       int flags) {
        StatFile* sf = reinterpret_cast<StatFile*>(*h);
        return sf->orig_ops->open(errinfo, &sf->orig_handle, path, flags);
    }

    static void cfs_close(couchstore_error_info_t *errinfo,
                          couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->orig_ops->close(errinfo, sf->orig_handle);
    }

    static ssize_t cfs_pread(couchstore_error_info_t *errinfo,
                             couch_file_handle h,
                             void* buf,
                             size_t sz,
                             cs_off_t off) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->stats->readSizeHisto.add(sz);
        if(sf->last_offs) {
            sf->stats->readSeekHisto.add(abs(off - sf->last_offs));
        }
        sf->last_offs = off;
        BlockTimer bt(&sf->stats->readTimeHisto);
        return sf->orig_ops->pread(errinfo, sf->orig_handle, buf, sz, off);
    }

    static ssize_t cfs_pwrite(couchstore_error_info_t *errinfo,
                              couch_file_handle h,
                              const void* buf,
                              size_t sz,
                              cs_off_t off) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->stats->writeSizeHisto.add(sz);
        BlockTimer bt(&sf->stats->writeTimeHisto);
        return sf->orig_ops->pwrite(errinfo, sf->orig_handle, buf, sz, off);
    }

    static cs_off_t cfs_goto_eof(couchstore_error_info_t *errinfo,
                                 couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        return sf->orig_ops->goto_eof(errinfo, sf->orig_handle);
    }

    static couchstore_error_t cfs_sync(couchstore_error_info_t *errinfo,
                                       couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        BlockTimer bt(&sf->stats->syncTimeHisto);
        return sf->orig_ops->sync(errinfo, sf->orig_handle);
    }

    static couchstore_error_t cfs_advise(couchstore_error_info_t *errinfo,
                                         couch_file_handle h,
                                         cs_off_t offs,
                                         cs_off_t len,
                                         couchstore_file_advice_t adv) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        return sf->orig_ops->advise(errinfo, sf->orig_handle, offs, len, adv);
    }

    static void cfs_destroy(couchstore_error_info_t *errinfo,
                            couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->orig_ops->destructor(errinfo, sf->orig_handle);
        delete sf;
    }

}
<commit_msg>Address compiler warning on mac<commit_after>/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 *     Copyright 2012 Couchbase, Inc
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */

#include "config.h"

#include "common.h"
#include "couch-kvstore/couch-fs-stats.h"
#include "histo.h"

extern "C" {
static couch_file_handle cfs_construct(couchstore_error_info_t*, void* cookie);
static couchstore_error_t cfs_open(couchstore_error_info_t*,
                                   couch_file_handle*, const char*, int);
static void cfs_close(couchstore_error_info_t*, couch_file_handle);
static ssize_t cfs_pread(couchstore_error_info_t*, couch_file_handle,
                         void *, size_t, cs_off_t);
static ssize_t cfs_pwrite(couchstore_error_info_t*, couch_file_handle,
                          const void *, size_t, cs_off_t);
static cs_off_t cfs_goto_eof(couchstore_error_info_t*, couch_file_handle);
static couchstore_error_t cfs_sync(couchstore_error_info_t*, couch_file_handle);
static couchstore_error_t cfs_advise(couchstore_error_info_t*,
                                     couch_file_handle,
                                     cs_off_t, cs_off_t,
                                     couchstore_file_advice_t);
static void cfs_destroy(couchstore_error_info_t*,couch_file_handle);
}

couch_file_ops getCouchstoreStatsOps(CouchstoreStats* stats) {
    couch_file_ops ops = {
        5,
        cfs_construct,
        cfs_open,
        cfs_close,
        cfs_pread,
        cfs_pwrite,
        cfs_goto_eof,
        cfs_sync,
        cfs_advise,
        cfs_destroy,
        stats
    };
    return ops;
}

struct StatFile {
    const couch_file_ops* orig_ops;
    couch_file_handle orig_handle;
    CouchstoreStats* stats;
    cs_off_t last_offs;
};

extern "C" {
    static couch_file_handle cfs_construct(couchstore_error_info_t *errinfo,
                                           void* cookie) {
        StatFile* sf = new StatFile;
        sf->stats = static_cast<CouchstoreStats*>(cookie);
        sf->orig_ops = couchstore_get_default_file_ops();
        sf->orig_handle = sf->orig_ops->constructor(errinfo,
                                                    sf->orig_ops->cookie);
        sf->last_offs = 0;
        return reinterpret_cast<couch_file_handle>(sf);
    }

    static couchstore_error_t cfs_open(couchstore_error_info_t *errinfo,
                                       couch_file_handle* h,
                                       const char* path,
                                       int flags) {
        StatFile* sf = reinterpret_cast<StatFile*>(*h);
        return sf->orig_ops->open(errinfo, &sf->orig_handle, path, flags);
    }

    static void cfs_close(couchstore_error_info_t *errinfo,
                          couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->orig_ops->close(errinfo, sf->orig_handle);
    }

    static ssize_t cfs_pread(couchstore_error_info_t *errinfo,
                             couch_file_handle h,
                             void* buf,
                             size_t sz,
                             cs_off_t off) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->stats->readSizeHisto.add(sz);
        if(sf->last_offs) {
            sf->stats->readSeekHisto.add(std::abs(off - sf->last_offs));
        }
        sf->last_offs = off;
        BlockTimer bt(&sf->stats->readTimeHisto);
        return sf->orig_ops->pread(errinfo, sf->orig_handle, buf, sz, off);
    }

    static ssize_t cfs_pwrite(couchstore_error_info_t *errinfo,
                              couch_file_handle h,
                              const void* buf,
                              size_t sz,
                              cs_off_t off) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->stats->writeSizeHisto.add(sz);
        BlockTimer bt(&sf->stats->writeTimeHisto);
        return sf->orig_ops->pwrite(errinfo, sf->orig_handle, buf, sz, off);
    }

    static cs_off_t cfs_goto_eof(couchstore_error_info_t *errinfo,
                                 couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        return sf->orig_ops->goto_eof(errinfo, sf->orig_handle);
    }

    static couchstore_error_t cfs_sync(couchstore_error_info_t *errinfo,
                                       couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        BlockTimer bt(&sf->stats->syncTimeHisto);
        return sf->orig_ops->sync(errinfo, sf->orig_handle);
    }

    static couchstore_error_t cfs_advise(couchstore_error_info_t *errinfo,
                                         couch_file_handle h,
                                         cs_off_t offs,
                                         cs_off_t len,
                                         couchstore_file_advice_t adv) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        return sf->orig_ops->advise(errinfo, sf->orig_handle, offs, len, adv);
    }

    static void cfs_destroy(couchstore_error_info_t *errinfo,
                            couch_file_handle h) {
        StatFile* sf = reinterpret_cast<StatFile*>(h);
        sf->orig_ops->destructor(errinfo, sf->orig_handle);
        delete sf;
    }

}
<|endoftext|>


<commit_before>/*
 * PosixCrypto.cpp
 *
 * Copyright (C) 2009-12 by RStudio, Inc.
 *
 * Unless you have received this program directly from RStudio pursuant
 * to the terms of a commercial license agreement with RStudio, then
 * this program is licensed to you under the terms of version 3 of the
 * GNU Affero General Public License. This program is distributed WITHOUT
 * ANY EXPRESS OR IMPLIED WARRANTY, INCLUDING THOSE OF NON-INFRINGEMENT,
 * MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. Please refer to the
 * AGPL (http://www.gnu.org/licenses/agpl-3.0.txt) for more details.
 *
 */

#include <core/system/Crypto.hpp>

#include <fcntl.h>

#include <openssl/err.h>
#include <openssl/hmac.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/rsa.h>

#include <algorithm>

#include <boost/utility.hpp>

#include <core/Log.hpp>
#include <core/Error.hpp>

// openssl calls on lion are are all marked as deprecated
#ifdef __clang__
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
#endif

using namespace rstudio::core;

namespace rstudio {
namespace core {
namespace system {
namespace crypto {

namespace {

// NOTE: we've never see the error codepath in spite of trying to get it to
// return an error by tweaking params -- NULL params caused a crash rather
// than returning an error). this is presumably because we are calling
// such low level functions. we may want to look at the src code to see
// if there is a way to test/engineer an error (or remove the error
// checking if there is no way to get one)

Error lastCryptoError(const ErrorLocation& location)
{
   // get the error code
   unsigned long ec = ::ERR_get_error();
   if (ec == 0)
   {
      LOG_WARNING_MESSAGE("lastCrytpoError called with no pending error");
      return systemError(boost::system::errc::not_supported,
                         "lastCrytpoError called with no pending error",
                         location);
   }

   // get the error message (docs say max len is 120)
   const int ERR_BUFF_SIZE = 250;
   char errorBuffer[ERR_BUFF_SIZE];
   ::ERR_error_string_n(ec, errorBuffer, ERR_BUFF_SIZE);

   // return the error
   return systemError(boost::system::errc::bad_message,
                      errorBuffer,
                      location);
}

class BIOFreeAllScope : boost::noncopyable
{
public:
   BIOFreeAllScope(BIO* pMem)
      : pMem_(pMem)
   {
   }

   ~BIOFreeAllScope()
   {
      try
      {
         ::BIO_free_all(pMem_);
      }
      catch(...)
      {
      }
   }
private:
   BIO* pMem_;
};

} // anonymous namespace

void initialize()
{
   // load global error string table
   ::ERR_load_crypto_strings();
}

Error HMAC_SHA2(const std::string& data,
                const std::string& key,
                std::vector<unsigned char>* pHMAC)
{
   // copy data into vector
   std::vector<unsigned char> keyVector;
   std::copy(key.begin(), key.end(), std::back_inserter(keyVector));

   // call core
   return HMAC_SHA2(data, keyVector, pHMAC);
}

Error HMAC_SHA2(const std::string& data,
                const std::vector<unsigned char>& key,
                std::vector<unsigned char>* pHMAC)
{
   // copy data into data vector
   std::vector<unsigned char> dataVector;
   std::copy(data.begin(), data.end(), std::back_inserter(dataVector));

   // perform the hash
   unsigned int md_len = 0;
   pHMAC->resize(EVP_MAX_MD_SIZE);
   unsigned char* pResult = ::HMAC(EVP_sha256(),
                                   &(key[0]),
                                   key.size(),
                                   &(dataVector[0]),
                                   dataVector.size(),
                                   &(pHMAC->operator[](0)),
                                   &md_len);
   if (pResult != NULL)
   {
      pHMAC->resize(md_len);
      return Success();
   }
   else
   {
      return lastCryptoError(ERROR_LOCATION);
   }
}

Error base64Encode(const std::vector<unsigned char>& data,
                   std::string* pEncoded)
{
   return base64Encode(&(data[0]), data.size(), pEncoded);
}

Error base64Encode(const unsigned char* pData,
                   int len,
                   std::string* pEncoded)
{
   // allocate BIO
   BIO* pB64 = ::BIO_new(BIO_f_base64());
   if (pB64 == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // no newlines
   BIO_set_flags(pB64, BIO_FLAGS_BASE64_NO_NL);

   // make sure it is freed prior to exit from the function
   BIOFreeAllScope freeB64Scope(pB64);

   // allocate memory stream
   BIO* pMem = ::BIO_new(BIO_s_mem());
   if (pMem == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // tie the stream to the b64 stream
   pB64 = ::BIO_push(pB64, pMem);

   // perform the encoding
   int written = ::BIO_write(pB64, pData, len);
   if (written != len)
      return lastCryptoError(ERROR_LOCATION);

   // flush all writes
   int result = BIO_flush(pB64);
   if (result <= 0)
      return lastCryptoError(ERROR_LOCATION);

   // seek to beginning of memory stream
   result = BIO_seek(pMem, 0);
   if (result == -1)
      return lastCryptoError(ERROR_LOCATION);

   // read the memory stream
   std::vector<char> buffer(len *2); // plenty more than len * 1.37 + padding
   int bytesRead = ::BIO_read(pMem, &(buffer[0]), buffer.capacity());
   if (bytesRead < 0 && ::ERR_get_error() != 0)
      return lastCryptoError(ERROR_LOCATION);

   // copy to out param
   buffer.resize(bytesRead);
   pEncoded->assign(buffer.begin(), buffer.end());

   // return success
   return Success();
}

Error base64Decode(const std::string& data,
                   std::vector<unsigned char>* pDecoded)
{
   // allocate b64 BIO
   BIO* pB64 = ::BIO_new(BIO_f_base64());
   if (pB64 == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // no newlines
   BIO_set_flags(pB64, BIO_FLAGS_BASE64_NO_NL);

   // make sure it is freed prior to exit from the function
   BIOFreeAllScope freeB64Scope(pB64);

   // allocate buffer
   BIO* pMem = BIO_new_mem_buf((void*)data.data(), data.length());
   if (pMem == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // tie the stream to the b64 stream
   pB64 = ::BIO_push(pB64, pMem);

   // reserve adequate memory in the decoded buffer and read into it
   pDecoded->clear();
   pDecoded->resize(data.length());
   int bytesRead = ::BIO_read(pB64,
                              &(pDecoded->operator[](0)),
                              pDecoded->size());
   if (bytesRead < 0)
      return lastCryptoError(ERROR_LOCATION);

   // resize the out buffer to the number of bytes actually read
   pDecoded->resize(bytesRead);

   // return success
   return Success();

}

namespace {
RSA* s_pRSA;
std::string s_modulo;
std::string s_exponent;
}

core::Error rsaInit()
{
   const int KEY_SIZE = 1024;
   const int ENTROPY_BYTES = 4096;

   BIGNUM *bn = BN_new();
   BN_set_word(bn, RSA_F4);
   const BIGNUM *bn_n;
   const BIGNUM *bn_e;

   int rnd = ::open("/dev/urandom", O_RDONLY);
   if (rnd == -1)
      return systemError(errno, ERROR_LOCATION);

   char entropy[ENTROPY_BYTES];
   if (-1 == ::read(rnd, entropy, ENTROPY_BYTES))
   {
      ::close(rnd);
      return systemError(errno, ERROR_LOCATION);
   }
   ::close(rnd);

   RAND_seed(entropy, ENTROPY_BYTES);

   #if OPENSSL_VERSION_NUMBER < 0x10100000L
      s_pRSA = ::RSA_generate_key(KEY_SIZE, 0x10001, NULL, NULL);
      if (!s_pRSA)
         return lastCryptoError(ERROR_LOCATION);

      bn_n = s_pRSA->n;
      bn_e = s_pRSA->e;
   #else
      s_pRSA = RSA_new();
      int rc = ::RSA_generate_key_ex(s_pRSA, KEY_SIZE, bn, NULL);
      if (rc != 1) {
        return lastCryptoError(ERROR_LOCATION);
        RSA_free(s_pRSA);
      }
      BN_clear_free(bn);

      RSA_get0_key(s_pRSA, &bn_n, &bn_e, NULL);
   #endif

   char* n = BN_bn2hex(bn_n);
   s_modulo = n;
   OPENSSL_free(n);
   char* e = BN_bn2hex(bn_e);
   s_exponent = e;
   OPENSSL_free(e);

   return Success();
}

void rsaPublicKey(std::string* pExponent, std::string* pModulo)
{
   pModulo->assign(s_modulo.begin(), s_modulo.end());
   pExponent->assign(s_exponent.begin(), s_exponent.end());
}

core::Error rsaPrivateDecrypt(const std::string& cipherText, std::string* pPlainText)
{
   std::vector<unsigned char> cipherTextBytes;
   Error error = base64Decode(cipherText, &cipherTextBytes);
   if (error)
      return error;

   int size = RSA_size(s_pRSA);
   std::vector<unsigned char> plainTextBytes(size);
   int bytesRead = RSA_private_decrypt(cipherTextBytes.size(),
                                       &cipherTextBytes[0],
                                       &plainTextBytes[0],
                                       s_pRSA,
                                       RSA_PKCS1_PADDING);
   if (bytesRead == -1)
      return lastCryptoError(ERROR_LOCATION);

   plainTextBytes.resize(bytesRead);
   pPlainText->assign(plainTextBytes.begin(), plainTextBytes.end());

   return Success();
}

} // namespace crypto
} // namespace system
} // namespace core
} // namespace rstudio

<commit_msg>address comments<commit_after>/*
 * PosixCrypto.cpp
 *
 * Copyright (C) 2009-12 by RStudio, Inc.
 *
 * Unless you have received this program directly from RStudio pursuant
 * to the terms of a commercial license agreement with RStudio, then
 * this program is licensed to you under the terms of version 3 of the
 * GNU Affero General Public License. This program is distributed WITHOUT
 * ANY EXPRESS OR IMPLIED WARRANTY, INCLUDING THOSE OF NON-INFRINGEMENT,
 * MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. Please refer to the
 * AGPL (http://www.gnu.org/licenses/agpl-3.0.txt) for more details.
 *
 */

#include <core/system/Crypto.hpp>

#include <fcntl.h>

#include <openssl/err.h>
#include <openssl/hmac.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/rsa.h>

#include <algorithm>

#include <boost/utility.hpp>

#include <core/Log.hpp>
#include <core/Error.hpp>

// openssl calls on lion are are all marked as deprecated
#ifdef __clang__
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
#endif

using namespace rstudio::core;

namespace rstudio {
namespace core {
namespace system {
namespace crypto {

namespace {

// NOTE: we've never see the error codepath in spite of trying to get it to
// return an error by tweaking params -- NULL params caused a crash rather
// than returning an error). this is presumably because we are calling
// such low level functions. we may want to look at the src code to see
// if there is a way to test/engineer an error (or remove the error
// checking if there is no way to get one)

Error lastCryptoError(const ErrorLocation& location)
{
   // get the error code
   unsigned long ec = ::ERR_get_error();
   if (ec == 0)
   {
      LOG_WARNING_MESSAGE("lastCrytpoError called with no pending error");
      return systemError(boost::system::errc::not_supported,
                         "lastCrytpoError called with no pending error",
                         location);
   }

   // get the error message (docs say max len is 120)
   const int ERR_BUFF_SIZE = 250;
   char errorBuffer[ERR_BUFF_SIZE];
   ::ERR_error_string_n(ec, errorBuffer, ERR_BUFF_SIZE);

   // return the error
   return systemError(boost::system::errc::bad_message,
                      errorBuffer,
                      location);
}

class BIOFreeAllScope : boost::noncopyable
{
public:
   BIOFreeAllScope(BIO* pMem)
      : pMem_(pMem)
   {
   }

   ~BIOFreeAllScope()
   {
      try
      {
         ::BIO_free_all(pMem_);
      }
      catch(...)
      {
      }
   }
private:
   BIO* pMem_;
};

} // anonymous namespace

void initialize()
{
   // load global error string table
   ::ERR_load_crypto_strings();
}

Error HMAC_SHA2(const std::string& data,
                const std::string& key,
                std::vector<unsigned char>* pHMAC)
{
   // copy data into vector
   std::vector<unsigned char> keyVector;
   std::copy(key.begin(), key.end(), std::back_inserter(keyVector));

   // call core
   return HMAC_SHA2(data, keyVector, pHMAC);
}

Error HMAC_SHA2(const std::string& data,
                const std::vector<unsigned char>& key,
                std::vector<unsigned char>* pHMAC)
{
   // copy data into data vector
   std::vector<unsigned char> dataVector;
   std::copy(data.begin(), data.end(), std::back_inserter(dataVector));

   // perform the hash
   unsigned int md_len = 0;
   pHMAC->resize(EVP_MAX_MD_SIZE);
   unsigned char* pResult = ::HMAC(EVP_sha256(),
                                   &(key[0]),
                                   key.size(),
                                   &(dataVector[0]),
                                   dataVector.size(),
                                   &(pHMAC->operator[](0)),
                                   &md_len);
   if (pResult != NULL)
   {
      pHMAC->resize(md_len);
      return Success();
   }
   else
   {
      return lastCryptoError(ERROR_LOCATION);
   }
}

Error base64Encode(const std::vector<unsigned char>& data,
                   std::string* pEncoded)
{
   return base64Encode(&(data[0]), data.size(), pEncoded);
}

Error base64Encode(const unsigned char* pData,
                   int len,
                   std::string* pEncoded)
{
   // allocate BIO
   BIO* pB64 = ::BIO_new(BIO_f_base64());
   if (pB64 == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // no newlines
   BIO_set_flags(pB64, BIO_FLAGS_BASE64_NO_NL);

   // make sure it is freed prior to exit from the function
   BIOFreeAllScope freeB64Scope(pB64);

   // allocate memory stream
   BIO* pMem = ::BIO_new(BIO_s_mem());
   if (pMem == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // tie the stream to the b64 stream
   pB64 = ::BIO_push(pB64, pMem);

   // perform the encoding
   int written = ::BIO_write(pB64, pData, len);
   if (written != len)
      return lastCryptoError(ERROR_LOCATION);

   // flush all writes
   int result = BIO_flush(pB64);
   if (result <= 0)
      return lastCryptoError(ERROR_LOCATION);

   // seek to beginning of memory stream
   result = BIO_seek(pMem, 0);
   if (result == -1)
      return lastCryptoError(ERROR_LOCATION);

   // read the memory stream
   std::vector<char> buffer(len *2); // plenty more than len * 1.37 + padding
   int bytesRead = ::BIO_read(pMem, &(buffer[0]), buffer.capacity());
   if (bytesRead < 0 && ::ERR_get_error() != 0)
      return lastCryptoError(ERROR_LOCATION);

   // copy to out param
   buffer.resize(bytesRead);
   pEncoded->assign(buffer.begin(), buffer.end());

   // return success
   return Success();
}

Error base64Decode(const std::string& data,
                   std::vector<unsigned char>* pDecoded)
{
   // allocate b64 BIO
   BIO* pB64 = ::BIO_new(BIO_f_base64());
   if (pB64 == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // no newlines
   BIO_set_flags(pB64, BIO_FLAGS_BASE64_NO_NL);

   // make sure it is freed prior to exit from the function
   BIOFreeAllScope freeB64Scope(pB64);

   // allocate buffer
   BIO* pMem = BIO_new_mem_buf((void*)data.data(), data.length());
   if (pMem == NULL)
      return lastCryptoError(ERROR_LOCATION);

   // tie the stream to the b64 stream
   pB64 = ::BIO_push(pB64, pMem);

   // reserve adequate memory in the decoded buffer and read into it
   pDecoded->clear();
   pDecoded->resize(data.length());
   int bytesRead = ::BIO_read(pB64,
                              &(pDecoded->operator[](0)),
                              pDecoded->size());
   if (bytesRead < 0)
      return lastCryptoError(ERROR_LOCATION);

   // resize the out buffer to the number of bytes actually read
   pDecoded->resize(bytesRead);

   // return success
   return Success();

}

namespace {
RSA* s_pRSA;
std::string s_modulo;
std::string s_exponent;
}

core::Error rsaInit()
{
   const int KEY_SIZE = 1024;
   const int ENTROPY_BYTES = 4096;

   BIGNUM *bn = BN_new();
   BN_set_word(bn, RSA_F4);
   const BIGNUM *bn_n;
   const BIGNUM *bn_e;

   int rnd = ::open("/dev/urandom", O_RDONLY);
   if (rnd == -1)
      return systemError(errno, ERROR_LOCATION);

   char entropy[ENTROPY_BYTES];
   if (-1 == ::read(rnd, entropy, ENTROPY_BYTES))
   {
      ::close(rnd);
      return systemError(errno, ERROR_LOCATION);
   }
   ::close(rnd);

   RAND_seed(entropy, ENTROPY_BYTES);

   #if OPENSSL_VERSION_NUMBER < 0x10100000L
      s_pRSA = ::RSA_generate_key(KEY_SIZE, 0x10001, NULL, NULL);
      if (!s_pRSA)
         return lastCryptoError(ERROR_LOCATION);

      bn_n = s_pRSA->n;
      bn_e = s_pRSA->e;

      BN_clear_free(bn);
   #else
      s_pRSA = RSA_new();
      int rc = ::RSA_generate_key_ex(s_pRSA, KEY_SIZE, bn, NULL);
      BN_clear_free(bn);
      if (rc != 1) {
        return lastCryptoError(ERROR_LOCATION);
        RSA_free(s_pRSA);
      }

      RSA_get0_key(s_pRSA, &bn_n, &bn_e, NULL);
   #endif

   char* n = BN_bn2hex(bn_n);
   s_modulo = n;
   OPENSSL_free(n);
   char* e = BN_bn2hex(bn_e);
   s_exponent = e;
   OPENSSL_free(e);

   return Success();
}

void rsaPublicKey(std::string* pExponent, std::string* pModulo)
{
   pModulo->assign(s_modulo.begin(), s_modulo.end());
   pExponent->assign(s_exponent.begin(), s_exponent.end());
}

core::Error rsaPrivateDecrypt(const std::string& cipherText, std::string* pPlainText)
{
   std::vector<unsigned char> cipherTextBytes;
   Error error = base64Decode(cipherText, &cipherTextBytes);
   if (error)
      return error;

   int size = RSA_size(s_pRSA);
   std::vector<unsigned char> plainTextBytes(size);
   int bytesRead = RSA_private_decrypt(cipherTextBytes.size(),
                                       &cipherTextBytes[0],
                                       &plainTextBytes[0],
                                       s_pRSA,
                                       RSA_PKCS1_PADDING);
   if (bytesRead == -1)
      return lastCryptoError(ERROR_LOCATION);

   plainTextBytes.resize(bytesRead);
   pPlainText->assign(plainTextBytes.begin(), plainTextBytes.end());

   return Success();
}

} // namespace crypto
} // namespace system
} // namespace core
} // namespace rstudio

<|endoftext|>


<commit_before>// Copyright 2016 Proyectos y Sistemas de Mantenimiento SL (eProsima).
//
// 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.

/*
 * Publisher.cpp
 *
 */

#include "PublisherImpl.h"
#include "../participant/ParticipantImpl.h"
#include <fastrtps/publisher/Publisher.h>
#include <fastrtps/TopicDataType.h>
#include <fastrtps/publisher/PublisherListener.h>

#include <fastrtps/rtps/writer/RTPSWriter.h>
#include <fastrtps/rtps/writer/StatefulWriter.h>

#include <fastrtps/rtps/participant/RTPSParticipant.h>
#include <fastrtps/rtps/RTPSDomain.h>

#include <fastrtps/utils/RTPSLog.h>

using namespace eprosima::fastrtps;
using namespace ::rtps;

static const char* const CLASS_NAME = "PublisherImpl";

::rtps::WriteParams WRITE_PARAM_DEFAULT;

PublisherImpl::PublisherImpl(ParticipantImpl* p,TopicDataType*pdatatype,
		PublisherAttributes& att,PublisherListener* listen ):
										mp_participant(p),
										mp_writer(nullptr),
										mp_type(pdatatype),
										m_att(att),
#pragma warning (disable : 4355 )
										m_history(this, pdatatype->m_typeSize, att.topic.historyQos, att.topic.resourceLimitsQos,att.historyMemoryPolicy),
										mp_listener(listen),
#pragma warning (disable : 4355 )
										m_writerListener(this),
										mp_userPublisher(nullptr),
										mp_rtpsParticipant(nullptr)
{

}

PublisherImpl::~PublisherImpl()
{
	const char* const METHOD_NAME = "~PublisherImpl";
	logInfo(PUBLISHER,this->getGuid().entityId << " in topic: "<<this->m_att.topic.topicName);
	RTPSDomain::removeRTPSWriter(mp_writer);
	delete(this->mp_userPublisher);
}

bool PublisherImpl::create_new_change(ChangeKind_t changeKind, void* data)
{
    return create_new_change_with_params(changeKind, data, WRITE_PARAM_DEFAULT);
}

bool PublisherImpl::create_new_change_with_params(ChangeKind_t changeKind, void* data, WriteParams &wparams)
{
	const char* const METHOD_NAME = "create_new_change";

    /// Preconditions
	if (data == nullptr)
	{
		logError(PUBLISHER, "Data pointer not valid");
		return false;
	}

	if(changeKind == NOT_ALIVE_UNREGISTERED || changeKind == NOT_ALIVE_DISPOSED ||
			changeKind == NOT_ALIVE_DISPOSED_UNREGISTERED)
	{
		if(m_att.topic.topicKind == NO_KEY)
		{
			logError(PUBLISHER,"Topic is NO_KEY, operation not permitted");
			return false;
		}
	}

	InstanceHandle_t handle;
	if(m_att.topic.topicKind == WITH_KEY)
	{
		mp_type->getKey(data,&handle);
	}

	CacheChange_t* ch = mp_writer->new_change(mp_type->getSerializedSizeProvider(data), changeKind,handle);
	if(ch != nullptr)
	{
		if(changeKind == ALIVE)
		{
			//First check that we can actually write to the buffer, then write

			//Static mode implies making sure the buffer size is enough for the maximum posible piece of data
			if( (m_att.historyMemoryPolicy == PREALLOCATED_MEMORY_MODE) && ch->serializedPayload.max_size > mp_type->m_typeSize)
			{
				logWarning(RTPS_WRITER,
					"Serialized Payload length larger than maximum type size (" <<
					ch->serializedPayload.length << "/" << mp_type->m_typeSize << ")");
				m_history.release_Cache(ch);
				return false;
			}
			//In the other modes, the cachechange is expected to at least be initialized at this point
			if(ch->serializedPayload.max_size == 0)
			{
				logWarning(RTPS_WRITER,"Serialized Payload length must be set to >0 ";);
				m_history.release_Cache(ch);
				return false;
			}
			//If these two checks are correct, we asume the cachechange is valid and thwn we can write to it.
			if(!mp_type->serialize(data,&ch->serializedPayload))
			{
				logWarning(RTPS_WRITER,"RTPSWriter:Serialization returns false";);
				m_history.release_Cache(ch);
				return false;
			}
		}

        // If it is big data, frament it.
        uint32_t high_mark = (mp_participant->getAttributes().rtps.sendSocketBufferSize - RTPSMESSAGE_COMMON_RTPS_PAYLOAD_SIZE) > PAYLOAD_MAX_SIZE ? PAYLOAD_MAX_SIZE : (mp_participant->getAttributes().rtps.sendSocketBufferSize - RTPSMESSAGE_COMMON_RTPS_PAYLOAD_SIZE);

        if(ch->serializedPayload.length > high_mark)
        {
            /// Fragment the data.
            // Set the fragment size to the cachechange.
            // Note: high_mark will always be a value that can be casted to uint16_t)
            ch->setFragmentSize((uint16_t)high_mark);
        }

        if(&wparams != &WRITE_PARAM_DEFAULT)
        {
            ch->write_params = wparams;
        }

		if(!this->m_history.add_pub_change(ch, wparams))
		{
			m_history.release_Cache(ch);
			return false;
		}

		return true;
	}
	return false;
}

bool PublisherImpl::removeMinSeqChange()
{
	return m_history.removeMinChange();
}

bool PublisherImpl::removeAllChange(size_t* removed)
{
	return m_history.removeAllChange(removed);
}

const GUID_t& PublisherImpl::getGuid()
{
	return mp_writer->getGuid();
}
//
bool PublisherImpl::updateAttributes(PublisherAttributes& att)
{
	const char* const METHOD_NAME = "updateAttributes";
	bool updated = true;
	bool missing = false;
	if(this->m_att.qos.m_reliability.kind == RELIABLE_RELIABILITY_QOS)
	{
		if(att.unicastLocatorList.size() != this->m_att.unicastLocatorList.size() ||
				att.multicastLocatorList.size() != this->m_att.multicastLocatorList.size())
		{
			logWarning(PUBLISHER,"Locator Lists cannot be changed or updated in this version");
			updated &= false;
		}
		else
		{
			for(LocatorListIterator lit1 = this->m_att.unicastLocatorList.begin();
					lit1!=this->m_att.unicastLocatorList.end();++lit1)
			{
				missing = true;
				for(LocatorListIterator lit2 = att.unicastLocatorList.begin();
						lit2!= att.unicastLocatorList.end();++lit2)
				{
					if(*lit1 == *lit2)
					{
						missing = false;
						break;
					}
				}
				if(missing)
				{
					logWarning(PUBLISHER,"Locator: "<< *lit1 << " not present in new list");
					logWarning(PUBLISHER,"Locator Lists cannot be changed or updated in this version");
				}
			}
			for(LocatorListIterator lit1 = this->m_att.multicastLocatorList.begin();
					lit1!=this->m_att.multicastLocatorList.end();++lit1)
			{
				missing = true;
				for(LocatorListIterator lit2 = att.multicastLocatorList.begin();
						lit2!= att.multicastLocatorList.end();++lit2)
				{
					if(*lit1 == *lit2)
					{
						missing = false;
						break;
					}
				}
				if(missing)
				{
					logWarning(PUBLISHER,"Locator: "<< *lit1<< " not present in new list");
					logWarning(PUBLISHER,"Locator Lists cannot be changed or updated in this version");
				}
			}
		}
	}

	//TOPIC ATTRIBUTES
	if(this->m_att.topic != att.topic)
	{
		logWarning(PUBLISHER,"Topic Attributes cannot be updated");
		updated &= false;
	}
	//QOS:
	//CHECK IF THE QOS CAN BE SET
	if(!this->m_att.qos.canQosBeUpdated(att.qos))
	{
		updated &=false;
	}
	if(updated)
	{
		if(this->m_att.qos.m_reliability.kind == RELIABLE_RELIABILITY_QOS)
		{
			//UPDATE TIMES:
			StatefulWriter* sfw = (StatefulWriter*)mp_writer;
			sfw->updateTimes(att.times);
		}
		this->m_att.qos.setQos(att.qos,false);
		this->m_att = att;
		//Notify the participant that a Writer has changed its QOS
		mp_rtpsParticipant->updateWriter(this->mp_writer,m_att.qos);
	}

	return updated;
}

void PublisherImpl::PublisherWriterListener::onWriterMatched(RTPSWriter* /*writer*/,MatchingInfo& info)
{
	if(mp_publisherImpl->mp_listener!=nullptr)
		mp_publisherImpl->mp_listener->onPublicationMatched(mp_publisherImpl->mp_userPublisher,info);
}

bool PublisherImpl::clean_history(unsigned int max)
{
    if(m_att.topic.historyQos.kind == HistoryQosPolicyKind::KEEP_ALL_HISTORY_QOS)
        return mp_writer->clean_history(max);
    else
        return mp_writer->remove_older_changes(max);
}

bool PublisherImpl::wait_for_all_acked(const Time_t& max_wait)
{
    return mp_writer->wait_for_all_acked(max_wait);
}

<commit_msg>Removing duplicate safechecks on cachechanges (PublisherImpl)<commit_after>// Copyright 2016 Proyectos y Sistemas de Mantenimiento SL (eProsima).
//
// 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.

/*
 * Publisher.cpp
 *
 */

#include "PublisherImpl.h"
#include "../participant/ParticipantImpl.h"
#include <fastrtps/publisher/Publisher.h>
#include <fastrtps/TopicDataType.h>
#include <fastrtps/publisher/PublisherListener.h>

#include <fastrtps/rtps/writer/RTPSWriter.h>
#include <fastrtps/rtps/writer/StatefulWriter.h>

#include <fastrtps/rtps/participant/RTPSParticipant.h>
#include <fastrtps/rtps/RTPSDomain.h>

#include <fastrtps/utils/RTPSLog.h>

using namespace eprosima::fastrtps;
using namespace ::rtps;

static const char* const CLASS_NAME = "PublisherImpl";

::rtps::WriteParams WRITE_PARAM_DEFAULT;

PublisherImpl::PublisherImpl(ParticipantImpl* p,TopicDataType*pdatatype,
		PublisherAttributes& att,PublisherListener* listen ):
										mp_participant(p),
										mp_writer(nullptr),
										mp_type(pdatatype),
										m_att(att),
#pragma warning (disable : 4355 )
										m_history(this, pdatatype->m_typeSize, att.topic.historyQos, att.topic.resourceLimitsQos,att.historyMemoryPolicy),
										mp_listener(listen),
#pragma warning (disable : 4355 )
										m_writerListener(this),
										mp_userPublisher(nullptr),
										mp_rtpsParticipant(nullptr)
{

}

PublisherImpl::~PublisherImpl()
{
	const char* const METHOD_NAME = "~PublisherImpl";
	logInfo(PUBLISHER,this->getGuid().entityId << " in topic: "<<this->m_att.topic.topicName);
	RTPSDomain::removeRTPSWriter(mp_writer);
	delete(this->mp_userPublisher);
}

bool PublisherImpl::create_new_change(ChangeKind_t changeKind, void* data)
{
    return create_new_change_with_params(changeKind, data, WRITE_PARAM_DEFAULT);
}

bool PublisherImpl::create_new_change_with_params(ChangeKind_t changeKind, void* data, WriteParams &wparams)
{
	const char* const METHOD_NAME = "create_new_change";

    /// Preconditions
	if (data == nullptr)
	{
		logError(PUBLISHER, "Data pointer not valid");
		return false;
	}

	if(changeKind == NOT_ALIVE_UNREGISTERED || changeKind == NOT_ALIVE_DISPOSED ||
			changeKind == NOT_ALIVE_DISPOSED_UNREGISTERED)
	{
		if(m_att.topic.topicKind == NO_KEY)
		{
			logError(PUBLISHER,"Topic is NO_KEY, operation not permitted");
			return false;
		}
	}

	InstanceHandle_t handle;
	if(m_att.topic.topicKind == WITH_KEY)
	{
		mp_type->getKey(data,&handle);
	}

	CacheChange_t* ch = mp_writer->new_change(mp_type->getSerializedSizeProvider(data), changeKind,handle);
	if(ch != nullptr)
	{
		if(changeKind == ALIVE)
		{
			//If these two checks are correct, we asume the cachechange is valid and thwn we can write to it.
			if(!mp_type->serialize(data,&ch->serializedPayload))
			{
				logWarning(RTPS_WRITER,"RTPSWriter:Serialization returns false";);
				m_history.release_Cache(ch);
				return false;
			}
		}

        // If it is big data, frament it.
        uint32_t high_mark = (mp_participant->getAttributes().rtps.sendSocketBufferSize - RTPSMESSAGE_COMMON_RTPS_PAYLOAD_SIZE) > PAYLOAD_MAX_SIZE ? PAYLOAD_MAX_SIZE : (mp_participant->getAttributes().rtps.sendSocketBufferSize - RTPSMESSAGE_COMMON_RTPS_PAYLOAD_SIZE);

        if(ch->serializedPayload.length > high_mark)
        {
            /// Fragment the data.
            // Set the fragment size to the cachechange.
            // Note: high_mark will always be a value that can be casted to uint16_t)
            ch->setFragmentSize((uint16_t)high_mark);
        }

        if(&wparams != &WRITE_PARAM_DEFAULT)
        {
            ch->write_params = wparams;
        }

		if(!this->m_history.add_pub_change(ch, wparams))
		{
			m_history.release_Cache(ch);
			return false;
		}

		return true;
	}
	return false;
}

bool PublisherImpl::removeMinSeqChange()
{
	return m_history.removeMinChange();
}

bool PublisherImpl::removeAllChange(size_t* removed)
{
	return m_history.removeAllChange(removed);
}

const GUID_t& PublisherImpl::getGuid()
{
	return mp_writer->getGuid();
}
//
bool PublisherImpl::updateAttributes(PublisherAttributes& att)
{
	const char* const METHOD_NAME = "updateAttributes";
	bool updated = true;
	bool missing = false;
	if(this->m_att.qos.m_reliability.kind == RELIABLE_RELIABILITY_QOS)
	{
		if(att.unicastLocatorList.size() != this->m_att.unicastLocatorList.size() ||
				att.multicastLocatorList.size() != this->m_att.multicastLocatorList.size())
		{
			logWarning(PUBLISHER,"Locator Lists cannot be changed or updated in this version");
			updated &= false;
		}
		else
		{
			for(LocatorListIterator lit1 = this->m_att.unicastLocatorList.begin();
					lit1!=this->m_att.unicastLocatorList.end();++lit1)
			{
				missing = true;
				for(LocatorListIterator lit2 = att.unicastLocatorList.begin();
						lit2!= att.unicastLocatorList.end();++lit2)
				{
					if(*lit1 == *lit2)
					{
						missing = false;
						break;
					}
				}
				if(missing)
				{
					logWarning(PUBLISHER,"Locator: "<< *lit1 << " not present in new list");
					logWarning(PUBLISHER,"Locator Lists cannot be changed or updated in this version");
				}
			}
			for(LocatorListIterator lit1 = this->m_att.multicastLocatorList.begin();
					lit1!=this->m_att.multicastLocatorList.end();++lit1)
			{
				missing = true;
				for(LocatorListIterator lit2 = att.multicastLocatorList.begin();
						lit2!= att.multicastLocatorList.end();++lit2)
				{
					if(*lit1 == *lit2)
					{
						missing = false;
						break;
					}
				}
				if(missing)
				{
					logWarning(PUBLISHER,"Locator: "<< *lit1<< " not present in new list");
					logWarning(PUBLISHER,"Locator Lists cannot be changed or updated in this version");
				}
			}
		}
	}

	//TOPIC ATTRIBUTES
	if(this->m_att.topic != att.topic)
	{
		logWarning(PUBLISHER,"Topic Attributes cannot be updated");
		updated &= false;
	}
	//QOS:
	//CHECK IF THE QOS CAN BE SET
	if(!this->m_att.qos.canQosBeUpdated(att.qos))
	{
		updated &=false;
	}
	if(updated)
	{
		if(this->m_att.qos.m_reliability.kind == RELIABLE_RELIABILITY_QOS)
		{
			//UPDATE TIMES:
			StatefulWriter* sfw = (StatefulWriter*)mp_writer;
			sfw->updateTimes(att.times);
		}
		this->m_att.qos.setQos(att.qos,false);
		this->m_att = att;
		//Notify the participant that a Writer has changed its QOS
		mp_rtpsParticipant->updateWriter(this->mp_writer,m_att.qos);
	}

	return updated;
}

void PublisherImpl::PublisherWriterListener::onWriterMatched(RTPSWriter* /*writer*/,MatchingInfo& info)
{
	if(mp_publisherImpl->mp_listener!=nullptr)
		mp_publisherImpl->mp_listener->onPublicationMatched(mp_publisherImpl->mp_userPublisher,info);
}

bool PublisherImpl::clean_history(unsigned int max)
{
    if(m_att.topic.historyQos.kind == HistoryQosPolicyKind::KEEP_ALL_HISTORY_QOS)
        return mp_writer->clean_history(max);
    else
        return mp_writer->remove_older_changes(max);
}

bool PublisherImpl::wait_for_all_acked(const Time_t& max_wait)
{
    return mp_writer->wait_for_all_acked(max_wait);
}

<|endoftext|>


<commit_before>#include "debug_log.hpp"
#include "connection_pool.hpp"

#include <boost/asio/io_service.hpp>
#include <boost/asio/signal_set.hpp>
#include <boost/program_options.hpp>

namespace {
void wait_on_signal() {
  boost::asio::io_service signal_io_service;
  boost::asio::io_service::work work{signal_io_service};
  boost::asio::signal_set signals(signal_io_service, SIGINT, SIGTERM);
  signals.async_wait([&](...) { signal_io_service.stop(); });
  signal_io_service.run();
}

template<class T, class U>
inline const T max(T a, U b) { return a < b ? b : a; }
}  // namespace

int main(int argc, char *argv[]) {
  using namespace riak;
  using namespace std::chrono;
  namespace po = boost::program_options;

  // Parse arguments
  std::string hostname;
  uint16_t port, num_threads, num_sockets, highwatermark;
  uint32_t nmsgs;
  po::options_description description{
    "Sends a lot of get_object requests to a Riak node using a connection pool."
  };
  description.add_options()
      ("help,h", "prints this help message")
      ("hostname,n",
       po::value<std::string>(&hostname)->default_value("localhost"),
       "hostname of Riak node")
      ("port,p",
       po::value<uint16_t>(&port)->default_value(10017),
       "port to connect on Riak node")
      ("num-threads,t",
       po::value<uint16_t>(&num_threads)->default_value(2),
       "number of I/O threads")
      ("num-sockets,s",
       po::value<uint16_t>(&num_sockets)->default_value(256),
       "number of sockets in pool")
      ("highwatermark,k",
       po::value<uint16_t>(&highwatermark)->default_value(1024),
       "max buffered requests")
      ("nmsgs,m",
       po::value<uint32_t>(&nmsgs)->default_value(1000),
       "number of messages to send to the node");
  po::variables_map variables;
  try {
    po::store(po::parse_command_line(argc, argv, description), variables);
    po::notify(variables);
  } catch (const std::exception& e) {
    DLOG << e.what();
    return -1;
  }
  if (variables.count("help")) {
    std::cerr << description << std::endl;
    return -1;
  }

  // Simple connection_pool usage:
  //   riak::connection_pool conn(hostname, port, num_threads, num_sockets,
  //                              highwatermark);
  //   conn.send(string_message1, deadline_ms, handler);
  //   conn.send(string_message2, deadline_ms, handler);
  //   etc.
  //
  // What follows is a mess because this is throwaway code.

  std::mutex num_sent_mutex;
  uint32_t num_sent = 0;
  auto start_clock = high_resolution_clock::now();
  auto first_response_clock = start_clock;

  std::string message{"\x09\x0A\01\x62\x12\x01\x6B", 7};
  DLOG << "Creating connection pool...";
  iasdf
  riak::connection_pool conn(hostname, port, num_threads, num_sockets,
                             highwatermark);

  DLOG << "Buffering messages...";
  auto log_every = max(1, nmsgs / 20);
  for (int i = 0 ; i < nmsgs ; ++ i) {
    conn.send(
        message, 1000,
        [&](std::string response, std::error_code error) {
          std::lock_guard<std::mutex> lock{num_sent_mutex};
          ++num_sent;
          if (num_sent == 1) {
            first_response_clock = high_resolution_clock::now();
            double secs = duration_cast<milliseconds>(
                first_response_clock - start_clock).count() / 1000.0;
            DLOG << error.message() << " [first message " << secs << " secs].";
          } else if (num_sent % log_every == 0 || num_sent == nmsgs) {
            if (num_sent <= num_sockets) return;
            auto total = duration_cast<milliseconds>(
                high_resolution_clock::now() - first_response_clock);
            auto msgs_per_sec =
                (num_sent - num_sockets) / (double(total.count()) / 1000.0);
            DLOG << error.message() << " [sent " << num_sent << " at "
                 << msgs_per_sec << " messages/sec]";
          }
    });

    if (i % log_every == 0) DLOG << "Buffered " << i << " messages.";
  }
  DLOG << "Buffered all the messages. Waiting on signal...";

  wait_on_signal();
  DLOG << "Signal caught.";

  return 0;
}
<commit_msg>fixed a stray char sequence<commit_after>#include "debug_log.hpp"
#include "connection_pool.hpp"

#include <boost/asio/io_service.hpp>
#include <boost/asio/signal_set.hpp>
#include <boost/program_options.hpp>

namespace {
void wait_on_signal() {
  boost::asio::io_service signal_io_service;
  boost::asio::io_service::work work{signal_io_service};
  boost::asio::signal_set signals(signal_io_service, SIGINT, SIGTERM);
  signals.async_wait([&](...) { signal_io_service.stop(); });
  signal_io_service.run();
}

template<class T, class U>
inline const T max(T a, U b) { return a < b ? b : a; }
}  // namespace

int main(int argc, char *argv[]) {
  using namespace riak;
  using namespace std::chrono;
  namespace po = boost::program_options;

  // Parse arguments
  std::string hostname;
  uint16_t port, num_threads, num_sockets, highwatermark;
  uint32_t nmsgs;
  po::options_description description{
    "Sends a lot of get_object requests to a Riak node using a connection pool."
  };
  description.add_options()
      ("help,h", "prints this help message")
      ("hostname,n",
       po::value<std::string>(&hostname)->default_value("localhost"),
       "hostname of Riak node")
      ("port,p",
       po::value<uint16_t>(&port)->default_value(10017),
       "port to connect on Riak node")
      ("num-threads,t",
       po::value<uint16_t>(&num_threads)->default_value(2),
       "number of I/O threads")
      ("num-sockets,s",
       po::value<uint16_t>(&num_sockets)->default_value(256),
       "number of sockets in pool")
      ("highwatermark,k",
       po::value<uint16_t>(&highwatermark)->default_value(1024),
       "max buffered requests")
      ("nmsgs,m",
       po::value<uint32_t>(&nmsgs)->default_value(1000),
       "number of messages to send to the node");
  po::variables_map variables;
  try {
    po::store(po::parse_command_line(argc, argv, description), variables);
    po::notify(variables);
  } catch (const std::exception& e) {
    DLOG << e.what();
    return -1;
  }
  if (variables.count("help")) {
    std::cerr << description << std::endl;
    return -1;
  }

  // Simple connection_pool usage:
  //   riak::connection_pool conn(hostname, port, num_threads, num_sockets,
  //                              highwatermark);
  //   conn.send(string_message1, deadline_ms, handler);
  //   conn.send(string_message2, deadline_ms, handler);
  //   etc.
  //
  // What follows is a mess because this is throwaway code.

  std::mutex num_sent_mutex;
  uint32_t num_sent = 0;
  auto start_clock = high_resolution_clock::now();
  auto first_response_clock = start_clock;

  std::string message{"\x09\x0A\01\x62\x12\x01\x6B", 7};
  DLOG << "Creating connection pool...";
  riak::connection_pool conn(hostname, port, num_threads, num_sockets,
                             highwatermark);

  DLOG << "Buffering messages...";
  auto log_every = max(1, nmsgs / 20);
  for (int i = 0 ; i < nmsgs ; ++ i) {
    conn.send(
        message, 1000,
        [&](std::string response, std::error_code error) {
          std::lock_guard<std::mutex> lock{num_sent_mutex};
          ++num_sent;
          if (num_sent == 1) {
            first_response_clock = high_resolution_clock::now();
            double secs = duration_cast<milliseconds>(
                first_response_clock - start_clock).count() / 1000.0;
            DLOG << error.message() << " [first message " << secs << " secs].";
          } else if (num_sent % log_every == 0 || num_sent == nmsgs) {
            if (num_sent <= num_sockets) return;
            auto total = duration_cast<milliseconds>(
                high_resolution_clock::now() - first_response_clock);
            auto msgs_per_sec =
                (num_sent - num_sockets) / (double(total.count()) / 1000.0);
            DLOG << error.message() << " [sent " << num_sent << " at "
                 << msgs_per_sec << " messages/sec]";
          }
    });

    if (i % log_every == 0) DLOG << "Buffered " << i << " messages.";
  }
  DLOG << "Buffered all the messages. Waiting on signal...";

  wait_on_signal();
  DLOG << "Signal caught.";

  return 0;
}
<|endoftext|>


<commit_before>#pragma once

#include "opencv2/opencv.hpp"

extern "C" {
#include "vl/generic.h"
#include "vl/slic.h"
}

#include "util.hpp"

using namespace cv;

typedef unsigned int uint;

void _getSLICSegments(const Mat& img, std::vector<vl_uint32>& segmentation)
{
	uint H = img.rows,
	     W = img.cols,
	     HW = H * W;

	// Convert format from LABLAB to LLAABB (for vlfeat)
	auto img_vl = new float[HW * 3];
	auto img_   = img.ptr<Vec3b>(0);
	for (uint j = 0; j < H; j++) {
		for (uint i = 0; i < W; i++) {
			img_vl[j * W + i]          = img_[j * W + i][0];
			img_vl[j * W + i + HW]     = img_[j * W + i][1];
			img_vl[j * W + i + HW * 2] = img_[j * W + i][2];
		}
	}

	vl_size regionSize    = HW / 5e3,
		minRegionSize = HW / 5e4;
	printf("\nSLIC parameters:\n- regionSize: %llu\n- minRegionSize: %llu\n",
	       regionSize, minRegionSize);

	vl_slic_segment(segmentation.data(), img_vl, W, H, img.channels(),
			regionSize, 200, minRegionSize);

	// Visualise segmentation
	Mat mat = img;
	int** labels = new int*[H];
	for (uint j = 0; j < H; j++) {
		labels[j] = new int[W];
		for (uint i = 0; i < W; i++)
			labels[j][i] = (int) segmentation[j*W + i];
	}

	int label, labelTop, labelBottom, labelLeft, labelRight;
	for (uint j = 1; j < H - 1; j++) {
		for (uint i = 1; i < W - 1; i++) {
			label       = labels[j][i];
			labelTop    = labels[j - 1][i];
			labelBottom = labels[j + 1][i];
			labelLeft   = labels[j][i - 1];
			labelRight  = labels[j][i + 1];
			if (label != labelTop  || label != labelBottom ||
			    label != labelLeft || label != labelRight) {
				mat.at<Vec3b>(j, i)[0] = 0;
				mat.at<Vec3b>(j, i)[1] = 0;
				mat.at<Vec3b>(j, i)[2] = 255;
			}
		}
	}
	showImage("SLIC", mat);
}

float _getSLICVariances(Mat& grey, std::vector<vl_uint32>& segmentation,
                        std::vector<float>& vars)
{
	uint n  = vars.size(),
	     HW = grey.cols * grey.rows;

	// 1. Aggregate pixels by super pixel
	auto spxl_vals = std::vector<std::vector<float>>(n);
	for (uint i = 0; i < n; i++) {
		spxl_vals[i] = std::vector<float>(0);
		spxl_vals[i].reserve(20);
	}

	vl_uint32 spxl_id;
	float     spxl_val;
	for (uint i = 0; i < HW; i++) {
		spxl_id = segmentation[i];
		spxl_val = (float) grey.ptr<uchar>(0)[i];
		spxl_vals[spxl_id].push_back(spxl_val);
	}

	// 2. Calculate variance of group of pixels
	for (uint i = 0; i < n; i++)
		vars[i] = var(spxl_vals[i]);

	// 3. Calculate variance threshold (25% with highest variance)
	auto vars_sorted = vars;
	std::sort(vars_sorted.begin(), vars_sorted.end());
	return vars_sorted[n - n / 4];
}

/**
 * Generates a pattern distinctiveness map
 *
 * 1) [Divide image into 9x9 patches]
 * 2) Perform PCA
 * 3) Project each patch into PCA space
 * 4) Take L1-norm and store to map
 */
Mat _getPatternDistinct(const Mat& img, std::vector<vl_uint32>& segmentation,
                        std::vector<float>& spxl_vars, float var_thresh)
{
	uint H  = img.rows,
	     W  = img.cols,
	     Y  = H - 1,    // Limit of y indexing
	     X  = W - 1,
	     IH = H - 2,    // Inner width (sans 1-pixel border)
	     IW = W - 2;    // Inner height

	const uchar* row_p = img.ptr<uchar>(0); // Pixel values of i-1th row
	const uchar* row   = img.ptr<uchar>(1); // Pixel values of ith row
	const uchar* row_n;                     // Pixel values of i+1th row

	/******************************/
	/* Create list of 9x9 patches */
	/******************************/
	auto _patches = std::vector<uchar>(0);
	_patches.reserve(X * Y * 9);

	// Patches in superpixels with var above var_thresh
	auto _distpatches = std::vector<uchar>(0);
	_distpatches.reserve(X * Y * 9);

	// Iterate over all inner pixels (patches) with variance above threshold
	uint i = 0, spxl_i;
	for (uint y = 1; y < Y; y++) {
		row_n = img.ptr<uchar>(y + 1);
		for (uint x = 1; x < X; x++) {
			spxl_i = segmentation[y*X + x];
			if (spxl_vars[spxl_i] > var_thresh) {
				_distpatches.push_back(row_p[x - 1]);
				_distpatches.push_back(row_p[x]    );
				_distpatches.push_back(row_p[x + 1]);
				_distpatches.push_back(row  [x - 1]);
				_distpatches.push_back(row  [x]    );
				_distpatches.push_back(row  [x + 1]);
				_distpatches.push_back(row_n[x - 1]);
				_distpatches.push_back(row_n[x]    );
				_distpatches.push_back(row_n[x + 1]);
				i += 9;
			}

			_patches.push_back(row_p[x - 1]);
			_patches.push_back(row_p[x]    );
			_patches.push_back(row_p[x + 1]);
			_patches.push_back(row  [x - 1]);
			_patches.push_back(row  [x]    );
			_patches.push_back(row  [x + 1]);
			_patches.push_back(row_n[x - 1]);
			_patches.push_back(row_n[x]    );
			_patches.push_back(row_n[x + 1]);

		}
		row_p = row;
		row = row_n;
	}
	_distpatches.shrink_to_fit();
	auto distpatches = Mat(i / 9, 9, CV_8U, _distpatches.data());
	auto patches     = Mat(X*Y, 9, CV_8U, _patches.data());

	/*******/
	/* PCA */
	/*******/
	auto pca = PCA(distpatches, Mat(), CV_PCA_DATA_AS_ROW);

	Mat pca_pos  = Mat::zeros(IH * IW, 9, CV_32F); // Coordinates in PCA space
	Mat pca_norm = Mat::zeros(IH * IW, 1, CV_32F); // L1 norm of pca_pos

	pca.project(patches, pca_pos); 	                  // Project patches into PCA space
	reduce(abs(pca_pos), pca_norm, 1, CV_REDUCE_SUM); // Calc L1 norm

	// Pad with 1-pixel thick black border
	Mat out_inner = Mat(IH, IW, CV_32F, pca_norm.ptr<float>(0), 0);
	Mat out;
	copyMakeBorder(out_inner, out, 1, 1, 1, 1, BORDER_CONSTANT, 0);

	return out;
}

/**
 * Generates a colour distinctiveness map
 *
 * 1)
 */
Mat _getColourDistinct(const Mat& img, std::vector<vl_uint32>& segmentation,
                       uint spxl_n)
{
	// 1. Aggregate pixels by super pixel
	auto spxl_vals = std::vector<std::vector<float>>(spxl_n);
	for (uint i = 0; i < spxl_n; i++) {
		spxl_vals[i] = std::vector<float>(0);
		spxl_vals[i].reserve(20);
	}

	auto out = Mat(img.size(), CV_32F);
	return out;
}

/**
 * Generates a saliency map using a method from Margolin et al. (2013)
 *
 * 1) Acquire pattern distinctiveness map
 * 2) Acquire colour distinctiveness map
 * 3) Calculate pixelwise multiplication of the two maps
 */
const float maxSize = 1000.f;

Mat getSaliency(const Mat& img)
{
	Mat img_BGR;

	uint H  = img.rows,
	     W  = img.cols,
	     HW = H * W;

	float scale = (float) max(H, W) / maxSize;
	if (scale > 1.f) {
		resize(img, img_BGR, Size(W / scale, H / scale));

		W = W / scale;
		H = H / scale;
	} else {
		img_BGR = img;
	}

	// Get grayscale image to work on
	auto img_grey = Mat(img_BGR.size(), img_BGR.type());
	cvtColor(img_BGR, img_grey, CV_BGR2GRAY);

	auto img_lab = Mat(img_BGR.size(), img_BGR.type());
	cvtColor(img_BGR, img_lab, CV_BGR2Lab);

	// Get SLIC superpixels
	auto segmentation = std::vector<vl_uint32>(H*W);
	_getSLICSegments(img_lab, segmentation);

	// Calculate variance of super pixels
	auto spxl_n = std::accumulate(segmentation.begin(),
		segmentation.end(), 0, [&](vl_uint32 b, vl_uint32 n) {
			return n > b ? n : b;
		}) + 1;
	printf("Calculated %d superpixels.\n", spxl_n);
	auto spxl_vars  = std::vector<float>(spxl_n);
	auto var_thresh = _getSLICVariances(img_grey, segmentation, spxl_vars);

	// Compute distinctiveness maps
	Mat patternD = _getPatternDistinct(img_grey, segmentation, spxl_vars,
					   var_thresh);
	Mat colourD = _getColourDistinct(img_lab, segmentation, spxl_n);
	//out = patternD.mul(colourD);
	Mat out = patternD;

	if (scale > 1.f) {
		Mat out_scaled = Mat();
		resize(out, out_scaled, img.size());
		std::swap(out, out_scaled);
	}
	return out;
}
<commit_msg>Change parameters to speed up getSaliency<commit_after>#pragma once

#include "opencv2/opencv.hpp"

extern "C" {
#include "vl/generic.h"
#include "vl/slic.h"
}

#include "util.hpp"

using namespace cv;

typedef unsigned int uint;

void _getSLICSegments(const Mat& img, std::vector<vl_uint32>& segmentation)
{
	uint H = img.rows,
	     W = img.cols,
	     HW = H * W;

	// Convert format from LABLAB to LLAABB (for vlfeat)
	auto img_vl = new float[HW * 3];
	auto img_   = img.ptr<Vec3b>(0);
	for (uint j = 0; j < H; j++) {
		for (uint i = 0; i < W; i++) {
			img_vl[j * W + i]          = img_[j * W + i][0];
			img_vl[j * W + i + HW]     = img_[j * W + i][1];
			img_vl[j * W + i + HW * 2] = img_[j * W + i][2];
		}
	}

	vl_size regionSize    = HW / 2e3,
		minRegionSize = HW / 5e4;
	printf("\nSLIC parameters:\n- regionSize: %llu\n- minRegionSize: %llu\n",
	       regionSize, minRegionSize);

	vl_slic_segment(segmentation.data(), img_vl, W, H, img.channels(),
			regionSize, 100, minRegionSize);

	// Visualise segmentation
	Mat mat = img;
	int** labels = new int*[H];
	for (uint j = 0; j < H; j++) {
		labels[j] = new int[W];
		for (uint i = 0; i < W; i++)
			labels[j][i] = (int) segmentation[j*W + i];
	}

	int label, labelTop, labelBottom, labelLeft, labelRight;
	for (uint j = 1; j < H - 1; j++) {
		for (uint i = 1; i < W - 1; i++) {
			label       = labels[j][i];
			labelTop    = labels[j - 1][i];
			labelBottom = labels[j + 1][i];
			labelLeft   = labels[j][i - 1];
			labelRight  = labels[j][i + 1];
			if (label != labelTop  || label != labelBottom ||
			    label != labelLeft || label != labelRight) {
				mat.at<Vec3b>(j, i)[0] = 0;
				mat.at<Vec3b>(j, i)[1] = 0;
				mat.at<Vec3b>(j, i)[2] = 255;
			}
		}
	}
	showImage("SLIC", mat);
}

float _getSLICVariances(Mat& grey, std::vector<vl_uint32>& segmentation,
                        std::vector<float>& vars)
{
	uint n  = vars.size(),
	     HW = grey.cols * grey.rows;

	// 1. Aggregate pixels by super pixel
	auto spxl_vals = std::vector<std::vector<float>>(n);
	for (uint i = 0; i < n; i++) {
		spxl_vals[i] = std::vector<float>(0);
		spxl_vals[i].reserve(20);
	}

	vl_uint32 spxl_id;
	float     spxl_val;
	for (uint i = 0; i < HW; i++) {
		spxl_id = segmentation[i];
		spxl_val = (float) grey.ptr<uchar>(0)[i];
		spxl_vals[spxl_id].push_back(spxl_val);
	}

	// 2. Calculate variance of group of pixels
	for (uint i = 0; i < n; i++)
		vars[i] = var(spxl_vals[i]);

	// 3. Calculate variance threshold (25% with highest variance)
	auto vars_sorted = vars;
	std::sort(vars_sorted.begin(), vars_sorted.end());
	return vars_sorted[n - n / 4];
}

/**
 * Generates a pattern distinctiveness map
 *
 * 1) [Divide image into 9x9 patches]
 * 2) Perform PCA
 * 3) Project each patch into PCA space
 * 4) Take L1-norm and store to map
 */
Mat _getPatternDistinct(const Mat& img, std::vector<vl_uint32>& segmentation,
                        std::vector<float>& spxl_vars, float var_thresh)
{
	uint H  = img.rows,
	     W  = img.cols,
	     Y  = H - 1,    // Limit of y indexing
	     X  = W - 1,
	     IH = H - 2,    // Inner width (sans 1-pixel border)
	     IW = W - 2;    // Inner height

	const uchar* row_p = img.ptr<uchar>(0); // Pixel values of i-1th row
	const uchar* row   = img.ptr<uchar>(1); // Pixel values of ith row
	const uchar* row_n;                     // Pixel values of i+1th row

	/******************************/
	/* Create list of 9x9 patches */
	/******************************/
	auto _patches = std::vector<uchar>(0);
	_patches.reserve(X * Y * 9);

	// Patches in superpixels with var above var_thresh
	auto _distpatches = std::vector<uchar>(0);
	_distpatches.reserve(X * Y * 9);

	// Iterate over all inner pixels (patches) with variance above threshold
	uint i = 0, spxl_i;
	for (uint y = 1; y < Y; y++) {
		row_n = img.ptr<uchar>(y + 1);
		for (uint x = 1; x < X; x++) {
			spxl_i = segmentation[y*X + x];
			if (spxl_vars[spxl_i] > var_thresh) {
				_distpatches.push_back(row_p[x - 1]);
				_distpatches.push_back(row_p[x]    );
				_distpatches.push_back(row_p[x + 1]);
				_distpatches.push_back(row  [x - 1]);
				_distpatches.push_back(row  [x]    );
				_distpatches.push_back(row  [x + 1]);
				_distpatches.push_back(row_n[x - 1]);
				_distpatches.push_back(row_n[x]    );
				_distpatches.push_back(row_n[x + 1]);
				i += 9;
			}

			_patches.push_back(row_p[x - 1]);
			_patches.push_back(row_p[x]    );
			_patches.push_back(row_p[x + 1]);
			_patches.push_back(row  [x - 1]);
			_patches.push_back(row  [x]    );
			_patches.push_back(row  [x + 1]);
			_patches.push_back(row_n[x - 1]);
			_patches.push_back(row_n[x]    );
			_patches.push_back(row_n[x + 1]);

		}
		row_p = row;
		row = row_n;
	}
	_distpatches.shrink_to_fit();
	auto distpatches = Mat(i / 9, 9, CV_8U, _distpatches.data());
	auto patches     = Mat(X*Y, 9, CV_8U, _patches.data());

	/*******/
	/* PCA */
	/*******/
	auto pca = PCA(distpatches, Mat(), CV_PCA_DATA_AS_ROW);

	Mat pca_pos  = Mat::zeros(IH * IW, 9, CV_32F); // Coordinates in PCA space
	Mat pca_norm = Mat::zeros(IH * IW, 1, CV_32F); // L1 norm of pca_pos

	pca.project(patches, pca_pos); 	                  // Project patches into PCA space
	reduce(abs(pca_pos), pca_norm, 1, CV_REDUCE_SUM); // Calc L1 norm

	// Pad with 1-pixel thick black border
	Mat out_inner = Mat(IH, IW, CV_32F, pca_norm.ptr<float>(0), 0);
	Mat out;
	copyMakeBorder(out_inner, out, 1, 1, 1, 1, BORDER_CONSTANT, 0);

	return out;
}

/**
 * Generates a colour distinctiveness map
 *
 * 1)
 */
Mat _getColourDistinct(const Mat& img, std::vector<vl_uint32>& segmentation,
                       uint spxl_n)
{
	// 1. Aggregate pixels by super pixel
	auto spxl_vals = std::vector<std::vector<float>>(spxl_n);
	for (uint i = 0; i < spxl_n; i++) {
		spxl_vals[i] = std::vector<float>(0);
		spxl_vals[i].reserve(20);
	}

	auto out = Mat(img.size(), CV_32F);
	return out;
}

/**
 * Generates a saliency map using a method from Margolin et al. (2013)
 *
 * 1) Acquire pattern distinctiveness map
 * 2) Acquire colour distinctiveness map
 * 3) Calculate pixelwise multiplication of the two maps
 */
const float maxSize = 600.f;

Mat getSaliency(const Mat& img)
{
	Mat img_BGR;

	uint H  = img.rows,
	     W  = img.cols,
	     HW = H * W;

	// Scale image to have not more than maxSize pixels on its larger
	// dimension
	float scale = (float) max(H, W) / maxSize;
	if (scale > 1.f) {
		resize(img, img_BGR, Size(W / scale, H / scale));

		W = W / scale;
		H = H / scale;
	} else {
		img_BGR = img;
	}

	// Get grayscale image to work on
	auto img_grey = Mat(img_BGR.size(), img_BGR.type());
	cvtColor(img_BGR, img_grey, CV_BGR2GRAY);

	auto img_lab = Mat(img_BGR.size(), img_BGR.type());
	cvtColor(img_BGR, img_lab, CV_BGR2Lab);

	// Get SLIC superpixels
	auto segmentation = std::vector<vl_uint32>(H*W);
	_getSLICSegments(img_lab, segmentation);

	// Calculate variance of super pixels
	auto spxl_n = std::accumulate(segmentation.begin(),
		segmentation.end(), 0, [&](vl_uint32 b, vl_uint32 n) {
			return n > b ? n : b;
		}) + 1;
	printf("Calculated %d superpixels.\n", spxl_n);
	auto spxl_vars  = std::vector<float>(spxl_n);
	auto var_thresh = _getSLICVariances(img_grey, segmentation, spxl_vars);

	// Compute distinctiveness maps
	Mat patternD = _getPatternDistinct(img_grey, segmentation, spxl_vars,
					   var_thresh);
	Mat colourD = _getColourDistinct(img_lab, segmentation, spxl_n);
	//out = patternD.mul(colourD);
	Mat out = patternD;

	// Scale back to original size for further processing
	if (scale > 1.f) {
		Mat out_scaled = Mat();
		resize(out, out_scaled, img.size());
		std::swap(out, out_scaled);
	}
	return out;
}
<|endoftext|>