<commit_before>#include <bitcoin/storage/postgresql_storage.hpp>

#include <bitcoin/block.hpp>
#include <bitcoin/transaction.hpp>
#include <bitcoin/util/assert.hpp>
#include <bitcoin/util/logger.hpp>

#include "postgresql_blockchain.hpp"

namespace libbitcoin {

hash_digest hash_from_bytea(std::string byte_stream);
data_chunk bytes_from_bytea(std::string byte_stream);

uint32_t extract_bits_head(uint32_t bits);
uint32_t extract_bits_body(uint32_t bits);

postgresql_storage::postgresql_storage(kernel_ptr kernel,
    std::string database, std::string user, std::string password)
  : sql_(std::string("postgresql:dbname=") + database +
        ";user=" + user + ";password=" + password)
{
    blockchain_.reset(new pq_blockchain(sql_, service(), kernel));
    // Organise/validate old blocks in case of unclean shutdown
    strand()->post(std::bind(&pq_blockchain::start, blockchain_));
}

size_t postgresql_storage::insert(const message::transaction_input& input,
    size_t transaction_id, size_t index_in_parent)
{
    std::string hash = pretty_hex(input.previous_output.hash),
        pretty_script = pretty_hex(save_script(input.input_script));
    static cppdb::statement statement = sql_.prepare(
        "INSERT INTO inputs (input_id, transaction_id, index_in_parent, \
            script, previous_output_hash, previous_output_index, sequence) \
        VALUES (DEFAULT, ?, ?, decode(?, 'hex'), decode(?, 'hex'), ?, ?) \
        RETURNING input_id"
        );
    statement.reset();
    statement.bind(transaction_id);
    statement.bind(index_in_parent);
    statement.bind(pretty_script);
    statement.bind(hash);
    statement.bind(input.previous_output.index);
    statement.bind(input.sequence);
    return statement.row().get<size_t>(0);
}

size_t postgresql_storage::insert(const message::transaction_output& output,
        size_t transaction_id, size_t index_in_parent)
{
    std::string pretty_script = pretty_hex(save_script(output.output_script));
    static cppdb::statement statement = sql_.prepare(
        "INSERT INTO outputs ( \
            output_id, transaction_id, index_in_parent, script, value) \
        VALUES (DEFAULT, ?, ?, decode(?, 'hex'), internal_to_sql(?)) \
        RETURNING output_id"
        );
    statement.reset();
    statement.bind(transaction_id);
    statement.bind(index_in_parent);
    statement.bind(pretty_script);
    statement.bind(output.value);
    return statement.row().get<size_t>(0);
}

size_t postgresql_storage::insert(const message::transaction& transaction,
    std::vector<size_t>& input_ids, std::vector<size_t>& output_ids)
{
    hash_digest transaction_hash = hash_transaction(transaction);
    std::string transaction_hash_repr = pretty_hex(transaction_hash);
    // We use special function to insert txs.
    // Some blocks contain duplicates. See SQL for more details.
    static cppdb::statement statement = sql_.prepare(
        "SELECT insert_transaction(decode(?, 'hex'), ?, ?, ?)"
        );
    statement.reset();
    statement.bind(transaction_hash_repr);
    statement.bind(transaction.version);
    statement.bind(transaction.locktime);
    statement.bind(is_coinbase(transaction));
    cppdb::result result = statement.row();
    size_t transaction_id = result.get<size_t>(0);
    if (transaction_id == 0)
    {
        cppdb::result old_transaction_id = sql_ <<
            "SELECT transaction_id \
            FROM transactions \
            WHERE transaction_hash=decode(?, 'hex')"
            << transaction_hash_repr
            << cppdb::row;
        return old_transaction_id.get<size_t>(0);
    }

    for (size_t i = 0; i < transaction.inputs.size(); ++i)
    {
        size_t input_id = insert(transaction.inputs[i], transaction_id, i);
        input_ids.push_back(input_id);
    }
    for (size_t i = 0; i < transaction.outputs.size(); ++i)
    {
        size_t output_id = insert(transaction.outputs[i], transaction_id, i);
        output_ids.push_back(output_id);
    }
    return transaction_id;
}

void postgresql_storage::store(const message::block& block,
        store_handler handle_store)
{
    strand()->post(std::bind(
        &postgresql_storage::do_store_block, shared_from_this(),
            block, handle_store));
}
void postgresql_storage::do_store_block(const message::block& block,
        store_handler handle_store)
{
    hash_digest block_hash = hash_block_header(block);
    std::string block_hash_repr = pretty_hex(block_hash),
            prev_block_repr = pretty_hex(block.prev_block),
            merkle_repr = pretty_hex(block.merkle_root);

    cppdb::transaction guard(sql_);
    cppdb::result result = sql_ <<
        "SELECT block_id FROM blocks WHERE block_hash=decode(?, 'hex')"
        << block_hash_repr << cppdb::row;
    if (!result.empty())
    {
        log_warning() << "Block '" << block_hash_repr << "' already exists";
        blockchain_->organizer()->refresh_block(result.get<size_t>(0));
        blockchain_->raise_barrier();
        handle_store(error::object_already_exists);
        return;
    }

    static cppdb::statement statement = sql_.prepare(
        "INSERT INTO blocks( \
            block_id, \
            block_hash, \
            space, \
            depth, \
            span_left, \
            span_right, \
            version, \
            prev_block_hash, \
            merkle, \
            when_created, \
            bits_head, \
            bits_body, \
            nonce \
        ) VALUES ( \
            DEFAULT, \
            decode(?, 'hex'), \
            nextval('blocks_space_sequence'), \
            0, \
            0, \
            0, \
            ?, \
            decode(?, 'hex'), \
            decode(?, 'hex'), \
            TO_TIMESTAMP(?), \
            ?, \
            ?, \
            ? \
        ) \
        RETURNING block_id"
        );

    statement.reset();
    statement.bind(block_hash_repr);
    statement.bind(block.version);
    statement.bind(prev_block_repr);
    statement.bind(merkle_repr);
    statement.bind(block.timestamp);
    uint32_t bits_head = extract_bits_head(block.bits),
        bits_body = extract_bits_body(block.bits);
    statement.bind(bits_head);
    statement.bind(bits_body);
    statement.bind(block.nonce);

    result = statement.row();
    pq_block_info block_info;
    block_info.block_id = result.get<size_t>(0);
    for (size_t i = 0; i < block.transactions.size(); ++i)
    {
        message::transaction transaction = block.transactions[i];
        pq_transaction_info tx_info;
        tx_info.transaction_id =
            insert(transaction, tx_info.input_ids, tx_info.output_ids);
        // Create block <-> txn mapping
        static cppdb::statement link_txs = sql_.prepare(
            "INSERT INTO transactions_parents ( \
                transaction_id, block_id, index_in_block) \
            VALUES (?, ?, ?)"
            );
        link_txs.reset();
        link_txs.bind(tx_info.transaction_id);
        link_txs.bind(block_info.block_id);
        link_txs.bind(i);
        link_txs.exec();
        block_info.transactions.push_back(tx_info);
    }
    blockchain_->buffer_block(std::make_pair(block_info, block));
    blockchain_->organizer()->refresh_block(block_info.block_id);
    blockchain_->raise_barrier();
    guard.commit();
    handle_store(std::error_code());
}

void postgresql_storage::fetch_block_locator(
        fetch_handler_block_locator handle_fetch)
{
    strand()->post(std::bind(
        &postgresql_storage::do_fetch_block_locator, shared_from_this(),
            handle_fetch));
}
void postgresql_storage::do_fetch_block_locator(
        fetch_handler_block_locator handle_fetch)
{
    cppdb::result number_blocks_result = sql_ <<
        "SELECT depth \
        FROM chains \
        ORDER BY work DESC \
        LIMIT 1"
        << cppdb::row;
    if (number_blocks_result.empty())
    {
        handle_fetch(error::object_doesnt_exist, message::block_locator());
        return;
    }
    // Start at max_depth
    int top_depth = number_blocks_result.get<size_t>(0);
    std::vector<size_t> indices;
    // Push last 10 indices first
    size_t step = 1, start = 0;
    for (int i = top_depth; i > 0; i -= step, ++start)
    {
        if (start >= 10)
            step *= 2;
        indices.push_back(i);
    }
    indices.push_back(0);
    // Now actually fetch the hashes for these blocks
    std::stringstream hack_sql;
    hack_sql <<
        "SELECT encode(block_hash, 'hex') \
        FROM main_chain \
        WHERE depth IN (";
    for (size_t i = 0; i < indices.size(); ++i)
    {
        if (i != 0)
            hack_sql << ", ";
        hack_sql << indices[i];
    }
    hack_sql << ") ORDER BY depth DESC";
    // ----------------------------------------------
    cppdb::result block_hashes_result = sql_ << hack_sql.str();
    message::block_locator locator;
    while (block_hashes_result.next())
    {
        std::string block_hash_repr = block_hashes_result.get<std::string>(0);
        locator.push_back(hash_from_bytea(block_hash_repr));
    }
    BITCOIN_ASSERT(locator.size() == indices.size());
    handle_fetch(std::error_code(), locator);
}

void postgresql_storage::fetch_balance(const short_hash& pubkey_hash,
    fetch_handler_balance handle_fetch)
{
    strand()->post(std::bind(
        &postgresql_storage::do_fetch_balance, shared_from_this(),
            pubkey_hash, handle_fetch));
}
void postgresql_storage::do_fetch_balance(const short_hash& pubkey_hash,
    fetch_handler_balance handle_fetch)
{
    std::string total_script = "76 a9 14 " + pretty_hex(pubkey_hash) + " 88 ac";
    static cppdb::statement statement = sql_.prepare(
        "WITH outs AS ( \
            SELECT \
                transaction_hash,  \
                outputs.index_in_parent,  \
                outputs.value \
            FROM \
            outputs, \
            transactions \
            WHERE  \
                script=decode(?, 'hex') \
                AND outputs.transaction_id=transactions.transaction_id \
            ) \
        SELECT \
            sql_to_internal(SUM(outs.value)) \
        FROM outs \
        LEFT JOIN inputs \
        ON \
            inputs.previous_output_hash=transaction_hash \
            AND inputs.previous_output_index=outs.index_in_parent \
        WHERE inputs IS NULL"
        );
    statement.reset();
    statement.bind(total_script);
    cppdb::result result = statement.row();
    uint64_t value = 0;
    if (!result.is_null(0))
        value = result.get<uint64_t>(0);
    handle_fetch(std::error_code(), value);
}

} // libbitcoin

<commit_msg>if dupli tx stop block caching<commit_after>#include <bitcoin/storage/postgresql_storage.hpp>

#include <bitcoin/block.hpp>
#include <bitcoin/transaction.hpp>
#include <bitcoin/util/assert.hpp>
#include <bitcoin/util/logger.hpp>

#include "postgresql_blockchain.hpp"

namespace libbitcoin {

hash_digest hash_from_bytea(std::string byte_stream);
data_chunk bytes_from_bytea(std::string byte_stream);

uint32_t extract_bits_head(uint32_t bits);
uint32_t extract_bits_body(uint32_t bits);

postgresql_storage::postgresql_storage(kernel_ptr kernel,
    std::string database, std::string user, std::string password)
  : sql_(std::string("postgresql:dbname=") + database +
        ";user=" + user + ";password=" + password)
{
    blockchain_.reset(new pq_blockchain(sql_, service(), kernel));
    // Organise/validate old blocks in case of unclean shutdown
    strand()->post(std::bind(&pq_blockchain::start, blockchain_));
}

size_t postgresql_storage::insert(const message::transaction_input& input,
    size_t transaction_id, size_t index_in_parent)
{
    std::string hash = pretty_hex(input.previous_output.hash),
        pretty_script = pretty_hex(save_script(input.input_script));
    static cppdb::statement statement = sql_.prepare(
        "INSERT INTO inputs (input_id, transaction_id, index_in_parent, \
            script, previous_output_hash, previous_output_index, sequence) \
        VALUES (DEFAULT, ?, ?, decode(?, 'hex'), decode(?, 'hex'), ?, ?) \
        RETURNING input_id"
        );
    statement.reset();
    statement.bind(transaction_id);
    statement.bind(index_in_parent);
    statement.bind(pretty_script);
    statement.bind(hash);
    statement.bind(input.previous_output.index);
    statement.bind(input.sequence);
    return statement.row().get<size_t>(0);
}

size_t postgresql_storage::insert(const message::transaction_output& output,
        size_t transaction_id, size_t index_in_parent)
{
    std::string pretty_script = pretty_hex(save_script(output.output_script));
    static cppdb::statement statement = sql_.prepare(
        "INSERT INTO outputs ( \
            output_id, transaction_id, index_in_parent, script, value) \
        VALUES (DEFAULT, ?, ?, decode(?, 'hex'), internal_to_sql(?)) \
        RETURNING output_id"
        );
    statement.reset();
    statement.bind(transaction_id);
    statement.bind(index_in_parent);
    statement.bind(pretty_script);
    statement.bind(output.value);
    return statement.row().get<size_t>(0);
}

size_t postgresql_storage::insert(const message::transaction& transaction,
    std::vector<size_t>& input_ids, std::vector<size_t>& output_ids)
{
    hash_digest transaction_hash = hash_transaction(transaction);
    std::string transaction_hash_repr = pretty_hex(transaction_hash);
    // We use special function to insert txs.
    // Some blocks contain duplicates. See SQL for more details.
    static cppdb::statement statement = sql_.prepare(
        "SELECT insert_transaction(decode(?, 'hex'), ?, ?, ?)"
        );
    statement.reset();
    statement.bind(transaction_hash_repr);
    statement.bind(transaction.version);
    statement.bind(transaction.locktime);
    statement.bind(is_coinbase(transaction));
    cppdb::result result = statement.row();
    size_t transaction_id = result.get<size_t>(0);
    if (transaction_id == 0)
    {
        cppdb::result old_transaction_id = sql_ <<
            "SELECT transaction_id \
            FROM transactions \
            WHERE transaction_hash=decode(?, 'hex')"
            << transaction_hash_repr
            << cppdb::row;
        return old_transaction_id.get<size_t>(0);
    }

    for (size_t i = 0; i < transaction.inputs.size(); ++i)
    {
        size_t input_id = insert(transaction.inputs[i], transaction_id, i);
        input_ids.push_back(input_id);
    }
    for (size_t i = 0; i < transaction.outputs.size(); ++i)
    {
        size_t output_id = insert(transaction.outputs[i], transaction_id, i);
        output_ids.push_back(output_id);
    }
    return transaction_id;
}

void postgresql_storage::store(const message::block& block,
        store_handler handle_store)
{
    strand()->post(std::bind(
        &postgresql_storage::do_store_block, shared_from_this(),
            block, handle_store));
}
void postgresql_storage::do_store_block(const message::block& block,
        store_handler handle_store)
{
    hash_digest block_hash = hash_block_header(block);
    std::string block_hash_repr = pretty_hex(block_hash),
            prev_block_repr = pretty_hex(block.prev_block),
            merkle_repr = pretty_hex(block.merkle_root);

    cppdb::transaction guard(sql_);
    cppdb::result result = sql_ <<
        "SELECT block_id FROM blocks WHERE block_hash=decode(?, 'hex')"
        << block_hash_repr << cppdb::row;
    if (!result.empty())
    {
        log_warning() << "Block '" << block_hash_repr << "' already exists";
        blockchain_->organizer()->refresh_block(result.get<size_t>(0));
        blockchain_->raise_barrier();
        handle_store(error::object_already_exists);
        return;
    }

    static cppdb::statement statement = sql_.prepare(
        "INSERT INTO blocks( \
            block_id, \
            block_hash, \
            space, \
            depth, \
            span_left, \
            span_right, \
            version, \
            prev_block_hash, \
            merkle, \
            when_created, \
            bits_head, \
            bits_body, \
            nonce \
        ) VALUES ( \
            DEFAULT, \
            decode(?, 'hex'), \
            nextval('blocks_space_sequence'), \
            0, \
            0, \
            0, \
            ?, \
            decode(?, 'hex'), \
            decode(?, 'hex'), \
            TO_TIMESTAMP(?), \
            ?, \
            ?, \
            ? \
        ) \
        RETURNING block_id"
        );

    statement.reset();
    statement.bind(block_hash_repr);
    statement.bind(block.version);
    statement.bind(prev_block_repr);
    statement.bind(merkle_repr);
    statement.bind(block.timestamp);
    uint32_t bits_head = extract_bits_head(block.bits),
        bits_body = extract_bits_body(block.bits);
    statement.bind(bits_head);
    statement.bind(bits_body);
    statement.bind(block.nonce);

    result = statement.row();
    pq_block_info block_info;
    block_info.block_id = result.get<size_t>(0);
    bool buffer_block = true;
    for (size_t i = 0; i < block.transactions.size(); ++i)
    {
        message::transaction transaction = block.transactions[i];
        pq_transaction_info tx_info;
        tx_info.transaction_id =
            insert(transaction, tx_info.input_ids, tx_info.output_ids);
        if (tx_info.input_ids.empty())
        {
            BITCOIN_ASSERT(tx_info.output_ids.empty());
            buffer_block = false;
        }
        // Create block <-> txn mapping
        static cppdb::statement link_txs = sql_.prepare(
            "INSERT INTO transactions_parents ( \
                transaction_id, block_id, index_in_block) \
            VALUES (?, ?, ?)"
            );
        link_txs.reset();
        link_txs.bind(tx_info.transaction_id);
        link_txs.bind(block_info.block_id);
        link_txs.bind(i);
        link_txs.exec();
        block_info.transactions.push_back(tx_info);
    }
    if (buffer_block)
        blockchain_->buffer_block(std::make_pair(block_info, block));
    blockchain_->organizer()->refresh_block(block_info.block_id);
    blockchain_->raise_barrier();
    guard.commit();
    handle_store(std::error_code());
}

void postgresql_storage::fetch_block_locator(
        fetch_handler_block_locator handle_fetch)
{
    strand()->post(std::bind(
        &postgresql_storage::do_fetch_block_locator, shared_from_this(),
            handle_fetch));
}
void postgresql_storage::do_fetch_block_locator(
        fetch_handler_block_locator handle_fetch)
{
    cppdb::result number_blocks_result = sql_ <<
        "SELECT depth \
        FROM chains \
        ORDER BY work DESC \
        LIMIT 1"
        << cppdb::row;
    if (number_blocks_result.empty())
    {
        handle_fetch(error::object_doesnt_exist, message::block_locator());
        return;
    }
    // Start at max_depth
    int top_depth = number_blocks_result.get<size_t>(0);
    std::vector<size_t> indices;
    // Push last 10 indices first
    size_t step = 1, start = 0;
    for (int i = top_depth; i > 0; i -= step, ++start)
    {
        if (start >= 10)
            step *= 2;
        indices.push_back(i);
    }
    indices.push_back(0);
    // Now actually fetch the hashes for these blocks
    std::stringstream hack_sql;
    hack_sql <<
        "SELECT encode(block_hash, 'hex') \
        FROM main_chain \
        WHERE depth IN (";
    for (size_t i = 0; i < indices.size(); ++i)
    {
        if (i != 0)
            hack_sql << ", ";
        hack_sql << indices[i];
    }
    hack_sql << ") ORDER BY depth DESC";
    // ----------------------------------------------
    cppdb::result block_hashes_result = sql_ << hack_sql.str();
    message::block_locator locator;
    while (block_hashes_result.next())
    {
        std::string block_hash_repr = block_hashes_result.get<std::string>(0);
        locator.push_back(hash_from_bytea(block_hash_repr));
    }
    BITCOIN_ASSERT(locator.size() == indices.size());
    handle_fetch(std::error_code(), locator);
}

void postgresql_storage::fetch_balance(const short_hash& pubkey_hash,
    fetch_handler_balance handle_fetch)
{
    strand()->post(std::bind(
        &postgresql_storage::do_fetch_balance, shared_from_this(),
            pubkey_hash, handle_fetch));
}
void postgresql_storage::do_fetch_balance(const short_hash& pubkey_hash,
    fetch_handler_balance handle_fetch)
{
    std::string total_script = "76 a9 14 " + pretty_hex(pubkey_hash) + " 88 ac";
    static cppdb::statement statement = sql_.prepare(
        "WITH outs AS ( \
            SELECT \
                transaction_hash,  \
                outputs.index_in_parent,  \
                outputs.value \
            FROM \
            outputs, \
            transactions \
            WHERE  \
                script=decode(?, 'hex') \
                AND outputs.transaction_id=transactions.transaction_id \
            ) \
        SELECT \
            sql_to_internal(SUM(outs.value)) \
        FROM outs \
        LEFT JOIN inputs \
        ON \
            inputs.previous_output_hash=transaction_hash \
            AND inputs.previous_output_index=outs.index_in_parent \
        WHERE inputs IS NULL"
        );
    statement.reset();
    statement.bind(total_script);
    cppdb::result result = statement.row();
    uint64_t value = 0;
    if (!result.is_null(0))
        value = result.get<uint64_t>(0);
    handle_fetch(std::error_code(), value);
}

} // libbitcoin

<|endoftext|>


<commit_before>// Copyright (c) 2014 GitHub, Inc.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/ui/win/notify_icon.h"

#include "atom/browser/ui/win/notify_icon_host.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/utf_string_conversions.h"
#include "base/win/windows_version.h"
#include "third_party/skia/include/core/SkBitmap.h"
#include "ui/gfx/image/image.h"
#include "ui/gfx/geometry/point.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/gfx/screen.h"
#include "ui/views/controls/menu/menu_runner.h"

namespace atom {

NotifyIcon::NotifyIcon(NotifyIconHost* host,
                       UINT id,
                       HWND window,
                       UINT message)
    : host_(host),
      icon_id_(id),
      window_(window),
      message_id_(message),
      menu_model_(NULL) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  // This can happen if the explorer process isn't running when we try to
  // create the icon for some reason (for example, at startup).
  if (!result)
    LOG(WARNING) << "Unable to create status tray icon.";
}

NotifyIcon::~NotifyIcon() {
  // Remove our icon.
  host_->Remove(this);
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
}

void NotifyIcon::HandleClickEvent(int modifiers,
                                  bool left_mouse_click,
                                  bool double_button_click) {
  gfx::Rect bounds = GetBounds();

  if (left_mouse_click) {
    if (double_button_click)  // double left click
      NotifyDoubleClicked(bounds, modifiers);
    else  // single left click
      NotifyClicked(bounds, modifiers);
    return;
  } else if (!double_button_click) {  // single right click
    if (menu_model_)
      PopUpContextMenu(gfx::Point(), menu_model_);
    else
      NotifyRightClicked(bounds, modifiers);
  }
}

void NotifyIcon::ResetIcon() {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  // Delete any previously existing icon.
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  icon_data.hIcon = icon_.get();
  // If we have an image, then set the NIF_ICON flag, which tells
  // Shell_NotifyIcon() to set the image for the status icon it creates.
  if (icon_data.hIcon)
    icon_data.uFlags |= NIF_ICON;
  // Re-add our icon.
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to re-create status tray icon.";
}

void NotifyIcon::SetImage(HICON image) {
  icon_ = base::win::ScopedHICON(CopyIcon(image));

  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_ICON;
  icon_data.hIcon = image;
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Error setting status tray icon image";
}

void NotifyIcon::SetPressedImage(HICON image) {
  // Ignore pressed images, since the standard on Windows is to not highlight
  // pressed status icons.
}

void NotifyIcon::SetToolTip(const std::string& tool_tip) {
  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_TIP;
  wcsncpy_s(icon_data.szTip, base::UTF8ToUTF16(tool_tip).c_str(), _TRUNCATE);
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to set tooltip for status tray icon";
}

void NotifyIcon::DisplayBalloon(HICON icon,
                                const base::string16& title,
                                const base::string16& contents) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_INFO;
  icon_data.dwInfoFlags = NIIF_INFO;
  wcsncpy_s(icon_data.szInfoTitle, title.c_str(), _TRUNCATE);
  wcsncpy_s(icon_data.szInfo, contents.c_str(), _TRUNCATE);
  icon_data.uTimeout = 0;
  icon_data.hBalloonIcon = icon;
  icon_data.dwInfoFlags = NIIF_USER | NIIF_LARGE_ICON;

  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to create status tray balloon.";
}

void NotifyIcon::PopUpContextMenu(const gfx::Point& pos,
                                  ui::SimpleMenuModel* menu_model) {
  // Returns if context menu isn't set.
  if (!menu_model)
    return;
  // Set our window as the foreground window, so the context menu closes when
  // we click away from it.
  if (!SetForegroundWindow(window_))
    return;

  // Show menu at mouse's position by default.
  gfx::Rect rect(pos, gfx::Size());
  if (pos.IsOrigin())
    rect.set_origin(gfx::Screen::GetScreen()->GetCursorScreenPoint());

  views::MenuRunner menu_runner(
      menu_model,
      views::MenuRunner::CONTEXT_MENU | views::MenuRunner::HAS_MNEMONICS);
  ignore_result(menu_runner.RunMenuAt(
      NULL, NULL, rect, views::MENU_ANCHOR_TOPLEFT, ui::MENU_SOURCE_MOUSE));
}

void NotifyIcon::SetContextMenu(ui::SimpleMenuModel* menu_model) {
  menu_model_ = menu_model;
}

gfx::Rect NotifyIcon::GetBounds() {
  NOTIFYICONIDENTIFIER icon_id;
  memset(&icon_id, 0, sizeof(NOTIFYICONIDENTIFIER));
  icon_id.uID = icon_id_;
  icon_id.hWnd = window_;
  icon_id.cbSize = sizeof(NOTIFYICONIDENTIFIER);

  RECT rect = { 0 };
  Shell_NotifyIconGetRect(&icon_id, &rect);
  return gfx::Rect(rect);
}

void NotifyIcon::InitIconData(NOTIFYICONDATA* icon_data) {
  memset(icon_data, 0, sizeof(NOTIFYICONDATA));
  icon_data->cbSize = sizeof(NOTIFYICONDATA);
  icon_data->hWnd = window_;
  icon_data->uID = icon_id_;
}

}  // namespace atom
<commit_msg>win: Use DIP rect for tray icon's bounds<commit_after>// Copyright (c) 2014 GitHub, Inc.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/ui/win/notify_icon.h"

#include "atom/browser/ui/win/notify_icon_host.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/utf_string_conversions.h"
#include "base/win/windows_version.h"
#include "third_party/skia/include/core/SkBitmap.h"
#include "ui/gfx/image/image.h"
#include "ui/gfx/geometry/point.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/gfx/screen.h"
#include "ui/gfx/win/dpi.h"
#include "ui/views/controls/menu/menu_runner.h"

namespace atom {

NotifyIcon::NotifyIcon(NotifyIconHost* host,
                       UINT id,
                       HWND window,
                       UINT message)
    : host_(host),
      icon_id_(id),
      window_(window),
      message_id_(message),
      menu_model_(NULL) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  // This can happen if the explorer process isn't running when we try to
  // create the icon for some reason (for example, at startup).
  if (!result)
    LOG(WARNING) << "Unable to create status tray icon.";
}

NotifyIcon::~NotifyIcon() {
  // Remove our icon.
  host_->Remove(this);
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
}

void NotifyIcon::HandleClickEvent(int modifiers,
                                  bool left_mouse_click,
                                  bool double_button_click) {
  gfx::Rect bounds = GetBounds();

  if (left_mouse_click) {
    if (double_button_click)  // double left click
      NotifyDoubleClicked(bounds, modifiers);
    else  // single left click
      NotifyClicked(bounds, modifiers);
    return;
  } else if (!double_button_click) {  // single right click
    if (menu_model_)
      PopUpContextMenu(gfx::Point(), menu_model_);
    else
      NotifyRightClicked(bounds, modifiers);
  }
}

void NotifyIcon::ResetIcon() {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  // Delete any previously existing icon.
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  icon_data.hIcon = icon_.get();
  // If we have an image, then set the NIF_ICON flag, which tells
  // Shell_NotifyIcon() to set the image for the status icon it creates.
  if (icon_data.hIcon)
    icon_data.uFlags |= NIF_ICON;
  // Re-add our icon.
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to re-create status tray icon.";
}

void NotifyIcon::SetImage(HICON image) {
  icon_ = base::win::ScopedHICON(CopyIcon(image));

  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_ICON;
  icon_data.hIcon = image;
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Error setting status tray icon image";
}

void NotifyIcon::SetPressedImage(HICON image) {
  // Ignore pressed images, since the standard on Windows is to not highlight
  // pressed status icons.
}

void NotifyIcon::SetToolTip(const std::string& tool_tip) {
  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_TIP;
  wcsncpy_s(icon_data.szTip, base::UTF8ToUTF16(tool_tip).c_str(), _TRUNCATE);
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to set tooltip for status tray icon";
}

void NotifyIcon::DisplayBalloon(HICON icon,
                                const base::string16& title,
                                const base::string16& contents) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags |= NIF_INFO;
  icon_data.dwInfoFlags = NIIF_INFO;
  wcsncpy_s(icon_data.szInfoTitle, title.c_str(), _TRUNCATE);
  wcsncpy_s(icon_data.szInfo, contents.c_str(), _TRUNCATE);
  icon_data.uTimeout = 0;
  icon_data.hBalloonIcon = icon;
  icon_data.dwInfoFlags = NIIF_USER | NIIF_LARGE_ICON;

  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to create status tray balloon.";
}

void NotifyIcon::PopUpContextMenu(const gfx::Point& pos,
                                  ui::SimpleMenuModel* menu_model) {
  // Returns if context menu isn't set.
  if (!menu_model)
    return;
  // Set our window as the foreground window, so the context menu closes when
  // we click away from it.
  if (!SetForegroundWindow(window_))
    return;

  // Show menu at mouse's position by default.
  gfx::Rect rect(pos, gfx::Size());
  if (pos.IsOrigin())
    rect.set_origin(gfx::Screen::GetScreen()->GetCursorScreenPoint());

  views::MenuRunner menu_runner(
      menu_model,
      views::MenuRunner::CONTEXT_MENU | views::MenuRunner::HAS_MNEMONICS);
  ignore_result(menu_runner.RunMenuAt(
      NULL, NULL, rect, views::MENU_ANCHOR_TOPLEFT, ui::MENU_SOURCE_MOUSE));
}

void NotifyIcon::SetContextMenu(ui::SimpleMenuModel* menu_model) {
  menu_model_ = menu_model;
}

gfx::Rect NotifyIcon::GetBounds() {
  NOTIFYICONIDENTIFIER icon_id;
  memset(&icon_id, 0, sizeof(NOTIFYICONIDENTIFIER));
  icon_id.uID = icon_id_;
  icon_id.hWnd = window_;
  icon_id.cbSize = sizeof(NOTIFYICONIDENTIFIER);

  RECT rect = { 0 };
  Shell_NotifyIconGetRect(&icon_id, &rect);
  return gfx::win::ScreenToDIPRect(gfx::Rect(rect));
}

void NotifyIcon::InitIconData(NOTIFYICONDATA* icon_data) {
  memset(icon_data, 0, sizeof(NOTIFYICONDATA));
  icon_data->cbSize = sizeof(NOTIFYICONDATA);
  icon_data->hWnd = window_;
  icon_data->uID = icon_id_;
}

}  // namespace atom
<|endoftext|>


<commit_before>// Copyright (c) 2014 GitHub, Inc.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/ui/win/notify_icon.h"

#include "atom/browser/ui/win/notify_icon_host.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/utf_string_conversions.h"
#include "base/win/windows_version.h"
#include "third_party/skia/include/core/SkBitmap.h"
#include "ui/gfx/icon_util.h"
#include "ui/gfx/image/image.h"
#include "ui/gfx/geometry/point.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/views/controls/menu/menu_runner.h"

namespace atom {

NotifyIcon::NotifyIcon(NotifyIconHost* host,
                       UINT id,
                       HWND window,
                       UINT message)
    : host_(host),
      icon_id_(id),
      window_(window),
      message_id_(message),
      menu_model_(NULL) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  // This can happen if the explorer process isn't running when we try to
  // create the icon for some reason (for example, at startup).
  if (!result)
    LOG(WARNING) << "Unable to create status tray icon.";
}

NotifyIcon::~NotifyIcon() {
  // Remove our icon.
  host_->Remove(this);
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
}

void NotifyIcon::HandleClickEvent(const gfx::Point& cursor_pos,
                                  bool left_mouse_click) {
  // Pass to the observer if appropriate.
  if (left_mouse_click) {
    NOTIFYICONIDENTIFIER icon_id;
    memset(&icon_id, 0, sizeof(NOTIFYICONIDENTIFIER));
    icon_id.uID = icon_id_;
    icon_id.hWnd = window_;
    icon_id.cbSize = sizeof(NOTIFYICONIDENTIFIER);

    RECT rect;
    Shell_NotifyIconGetRect(&icon_id, &rect);

    int width = rect.right - rect.left;
    int height =  rect.bottom - rect.top;
    NotifyClicked(gfx::Rect(rect.left, rect.top, width, height));
    return;
  }

  if (!menu_model_)
    return;

  // Set our window as the foreground window, so the context menu closes when
  // we click away from it.
  if (!SetForegroundWindow(window_))
    return;

  views::MenuRunner menu_runner(
      menu_model_,
      views::MenuRunner::CONTEXT_MENU | views::MenuRunner::HAS_MNEMONICS);
  ignore_result(menu_runner.RunMenuAt(
      NULL,
      NULL,
      gfx::Rect(cursor_pos, gfx::Size()),
      views::MENU_ANCHOR_TOPLEFT,
      ui::MENU_SOURCE_MOUSE));
}

void NotifyIcon::ResetIcon() {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  // Delete any previously existing icon.
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  icon_data.hIcon = icon_.Get();
  // If we have an image, then set the NIF_ICON flag, which tells
  // Shell_NotifyIcon() to set the image for the status icon it creates.
  if (icon_data.hIcon)
    icon_data.uFlags |= NIF_ICON;
  // Re-add our icon.
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to re-create status tray icon.";
}

void NotifyIcon::SetImage(const gfx::Image& image) {
  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_ICON;
  icon_.Set(IconUtil::CreateHICONFromSkBitmap(image.AsBitmap()));
  icon_data.hIcon = icon_.Get();
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Error setting status tray icon image";
}

void NotifyIcon::SetPressedImage(const gfx::Image& image) {
  // Ignore pressed images, since the standard on Windows is to not highlight
  // pressed status icons.
}

void NotifyIcon::SetToolTip(const std::string& tool_tip) {
  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_TIP;
  wcscpy_s(icon_data.szTip, base::UTF8ToUTF16(tool_tip).c_str());
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to set tooltip for status tray icon";
}

void NotifyIcon::DisplayBalloon(const gfx::Image& icon,
                                const base::string16& title,
                                const base::string16& contents) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_INFO;
  icon_data.dwInfoFlags = NIIF_INFO;
  wcscpy_s(icon_data.szInfoTitle, title.c_str());
  wcscpy_s(icon_data.szInfo, contents.c_str());
  icon_data.uTimeout = 0;

  base::win::Version win_version = base::win::GetVersion();
  if (!icon.IsEmpty() && win_version != base::win::VERSION_PRE_XP) {
    balloon_icon_.Set(IconUtil::CreateHICONFromSkBitmap(icon.AsBitmap()));
    icon_data.hBalloonIcon = balloon_icon_.Get();
    icon_data.dwInfoFlags = NIIF_USER | NIIF_LARGE_ICON;
  }

  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to create status tray balloon.";
}

void NotifyIcon::SetContextMenu(ui::SimpleMenuModel* menu_model) {
  menu_model_ = menu_model;
}

void NotifyIcon::InitIconData(NOTIFYICONDATA* icon_data) {
  memset(icon_data, 0, sizeof(NOTIFYICONDATA));
  icon_data->cbSize = sizeof(NOTIFYICONDATA);
  icon_data->hWnd = window_;
  icon_data->uID = icon_id_;
}

}  // namespace atom
<commit_msg>Ow :poop:, where did that extra space come from?<commit_after>// Copyright (c) 2014 GitHub, Inc.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/ui/win/notify_icon.h"

#include "atom/browser/ui/win/notify_icon_host.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/utf_string_conversions.h"
#include "base/win/windows_version.h"
#include "third_party/skia/include/core/SkBitmap.h"
#include "ui/gfx/icon_util.h"
#include "ui/gfx/image/image.h"
#include "ui/gfx/geometry/point.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/views/controls/menu/menu_runner.h"

namespace atom {

NotifyIcon::NotifyIcon(NotifyIconHost* host,
                       UINT id,
                       HWND window,
                       UINT message)
    : host_(host),
      icon_id_(id),
      window_(window),
      message_id_(message),
      menu_model_(NULL) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  // This can happen if the explorer process isn't running when we try to
  // create the icon for some reason (for example, at startup).
  if (!result)
    LOG(WARNING) << "Unable to create status tray icon.";
}

NotifyIcon::~NotifyIcon() {
  // Remove our icon.
  host_->Remove(this);
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
}

void NotifyIcon::HandleClickEvent(const gfx::Point& cursor_pos,
                                  bool left_mouse_click) {
  // Pass to the observer if appropriate.
  if (left_mouse_click) {
    NOTIFYICONIDENTIFIER icon_id;
    memset(&icon_id, 0, sizeof(NOTIFYICONIDENTIFIER));
    icon_id.uID = icon_id_;
    icon_id.hWnd = window_;
    icon_id.cbSize = sizeof(NOTIFYICONIDENTIFIER);

    RECT rect;
    Shell_NotifyIconGetRect(&icon_id, &rect);

    int width = rect.right - rect.left;
    int height = rect.bottom - rect.top;
    NotifyClicked(gfx::Rect(rect.left, rect.top, width, height));
    return;
  }

  if (!menu_model_)
    return;

  // Set our window as the foreground window, so the context menu closes when
  // we click away from it.
  if (!SetForegroundWindow(window_))
    return;

  views::MenuRunner menu_runner(
      menu_model_,
      views::MenuRunner::CONTEXT_MENU | views::MenuRunner::HAS_MNEMONICS);
  ignore_result(menu_runner.RunMenuAt(
      NULL,
      NULL,
      gfx::Rect(cursor_pos, gfx::Size()),
      views::MENU_ANCHOR_TOPLEFT,
      ui::MENU_SOURCE_MOUSE));
}

void NotifyIcon::ResetIcon() {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  // Delete any previously existing icon.
  Shell_NotifyIcon(NIM_DELETE, &icon_data);
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_MESSAGE;
  icon_data.uCallbackMessage = message_id_;
  icon_data.hIcon = icon_.Get();
  // If we have an image, then set the NIF_ICON flag, which tells
  // Shell_NotifyIcon() to set the image for the status icon it creates.
  if (icon_data.hIcon)
    icon_data.uFlags |= NIF_ICON;
  // Re-add our icon.
  BOOL result = Shell_NotifyIcon(NIM_ADD, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to re-create status tray icon.";
}

void NotifyIcon::SetImage(const gfx::Image& image) {
  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_ICON;
  icon_.Set(IconUtil::CreateHICONFromSkBitmap(image.AsBitmap()));
  icon_data.hIcon = icon_.Get();
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Error setting status tray icon image";
}

void NotifyIcon::SetPressedImage(const gfx::Image& image) {
  // Ignore pressed images, since the standard on Windows is to not highlight
  // pressed status icons.
}

void NotifyIcon::SetToolTip(const std::string& tool_tip) {
  // Create the icon.
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_TIP;
  wcscpy_s(icon_data.szTip, base::UTF8ToUTF16(tool_tip).c_str());
  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to set tooltip for status tray icon";
}

void NotifyIcon::DisplayBalloon(const gfx::Image& icon,
                                const base::string16& title,
                                const base::string16& contents) {
  NOTIFYICONDATA icon_data;
  InitIconData(&icon_data);
  icon_data.uFlags = NIF_INFO;
  icon_data.dwInfoFlags = NIIF_INFO;
  wcscpy_s(icon_data.szInfoTitle, title.c_str());
  wcscpy_s(icon_data.szInfo, contents.c_str());
  icon_data.uTimeout = 0;

  base::win::Version win_version = base::win::GetVersion();
  if (!icon.IsEmpty() && win_version != base::win::VERSION_PRE_XP) {
    balloon_icon_.Set(IconUtil::CreateHICONFromSkBitmap(icon.AsBitmap()));
    icon_data.hBalloonIcon = balloon_icon_.Get();
    icon_data.dwInfoFlags = NIIF_USER | NIIF_LARGE_ICON;
  }

  BOOL result = Shell_NotifyIcon(NIM_MODIFY, &icon_data);
  if (!result)
    LOG(WARNING) << "Unable to create status tray balloon.";
}

void NotifyIcon::SetContextMenu(ui::SimpleMenuModel* menu_model) {
  menu_model_ = menu_model;
}

void NotifyIcon::InitIconData(NOTIFYICONDATA* icon_data) {
  memset(icon_data, 0, sizeof(NOTIFYICONDATA));
  icon_data->cbSize = sizeof(NOTIFYICONDATA);
  icon_data->hWnd = window_;
  icon_data->uID = icon_id_;
}

}  // namespace atom
<|endoftext|>


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

/**
 * @file ecl_yaw_controller.cpp
 * Implementation of a simple orthogonal coordinated turn yaw PID controller.
 *
 * Authors and acknowledgements in header.
 */

#include "ecl_yaw_controller.h"
#include <stdint.h>
#include <float.h>
#include <geo/geo.h>
#include <ecl/ecl.h>
#include <mathlib/mathlib.h>
#include <systemlib/err.h>
#include <ecl/ecl.h>

ECL_YawController::ECL_YawController() :
	ECL_Controller("yaw"),
	_coordinated_min_speed(1.0f),
	_max_rate(0.0f), /* disable by default */
	_coordinated_method(0)
{
}

float ECL_YawController::control_attitude(const struct ECL_ControlData &ctl_data)
{
	switch (_coordinated_method) {
	case COORD_METHOD_OPEN:
		return control_attitude_impl_openloop(ctl_data);

	case COORD_METHOD_CLOSEACC:
		return control_attitude_impl_accclosedloop(ctl_data);

	default:
		static hrt_abstime last_print = 0;

		if (ecl_elapsed_time(&last_print) > 5e6) {
			warnx("invalid param setting FW_YCO_METHOD");
			last_print = ecl_absolute_time();
		}
	}

	return _rate_setpoint;
}

float ECL_YawController::control_attitude_impl_openloop(const struct ECL_ControlData &ctl_data)
{
	/* Do not calculate control signal with bad inputs */
	if (!(PX4_ISFINITE(ctl_data.roll) &&
	      PX4_ISFINITE(ctl_data.pitch) &&
	      PX4_ISFINITE(ctl_data.speed_body_u) &&
	      PX4_ISFINITE(ctl_data.speed_body_v) &&
	      PX4_ISFINITE(ctl_data.speed_body_w) &&
	      PX4_ISFINITE(ctl_data.roll_rate_setpoint) &&
	      PX4_ISFINITE(ctl_data.pitch_rate_setpoint))) {
		return _rate_setpoint;
	}

	float constrained_roll;
	/* roll is used as feedforward term and inverted flight needs to be considered */
	if (fabsf(ctl_data.roll) < math::radians(90.0f)) {
		/* not inverted, but numerically still potentially close to infinity */
		constrained_roll = math::constrain(ctl_data.roll, math::radians(-80.0f), math::radians(80.0f));

	} else {
		 // inverted flight, constrain on the two extremes of -pi..+pi to avoid infinity
		 //note: the ranges are extended by 10 deg here to avoid numeric resolution effects
		if (ctl_data.roll > 0.0f) {
			/* right hemisphere */
			constrained_roll = math::constrain(ctl_data.roll, math::radians(100.0f), math::radians(180.0f));

		} else {
			/* left hemisphere */
			constrained_roll = math::constrain(ctl_data.roll, math::radians(-180.0f), math::radians(-100.0f));
		}
	}

	/* Calculate desired yaw rate from coordinated turn constraint / (no side forces) */
	_rate_setpoint = tanf(constrained_roll) * cosf(ctl_data.pitch) * 9.81f / (ctl_data.airspeed < ctl_data.airspeed_min ? ctl_data.airspeed_min : ctl_data.airspeed);

	/* limit the rate */ //XXX: move to body angluar rates
	if (_max_rate > 0.01f) {
		_rate_setpoint = (_rate_setpoint > _max_rate) ? _max_rate : _rate_setpoint;
		_rate_setpoint = (_rate_setpoint < -_max_rate) ? -_max_rate : _rate_setpoint;
	}

	if (!PX4_ISFINITE(_rate_setpoint)) {
		warnx("yaw rate sepoint not finite");
		_rate_setpoint = 0.0f;
	}

	return _rate_setpoint;
}

float ECL_YawController::control_bodyrate(const struct ECL_ControlData &ctl_data)
{
	/* Do not calculate control signal with bad inputs */
	if (!(PX4_ISFINITE(ctl_data.roll) && PX4_ISFINITE(ctl_data.pitch) && PX4_ISFINITE(ctl_data.body_y_rate) &&
	      PX4_ISFINITE(ctl_data.body_z_rate) && PX4_ISFINITE(ctl_data.pitch_rate_setpoint) &&
	      PX4_ISFINITE(ctl_data.airspeed_min) && PX4_ISFINITE(ctl_data.airspeed_max) &&
	      PX4_ISFINITE(ctl_data.scaler))) {
		return math::constrain(_last_output, -1.0f, 1.0f);
	}

	/* get the usual dt estimate */
	uint64_t dt_micros = ecl_elapsed_time(&_last_run);
	_last_run = ecl_absolute_time();
	float dt = (float)dt_micros * 1e-6f;

	/* lock integral for long intervals */
	bool lock_integrator = ctl_data.lock_integrator;

	if (dt_micros > 500000) {
		lock_integrator = true;
	}

	/* input conditioning */
	float airspeed = ctl_data.airspeed;

	if (!PX4_ISFINITE(airspeed)) {
		/* airspeed is NaN, +- INF or not available, pick center of band */
		airspeed = 0.5f * (ctl_data.airspeed_min + ctl_data.airspeed_max);

	} else if (airspeed < ctl_data.airspeed_min) {
		airspeed = ctl_data.airspeed_min;
	}

	/* Close the acceleration loop if _coordinated_method wants this: change body_rate setpoint */
	if (_coordinated_method == COORD_METHOD_CLOSEACC) {
		// XXX lateral acceleration needs to go into integrator with a gain
		//_bodyrate_setpoint -= (ctl_data.acc_body_y / (airspeed * cosf(ctl_data.pitch)));
	}

	/* Calculate body angular rate error */
	_rate_error = _bodyrate_setpoint - ctl_data.body_z_rate; //body angular rate error

	if (!lock_integrator && _k_i > 0.0f && airspeed > 0.5f * ctl_data.airspeed_min) {

		float id = _rate_error * dt;

		/*
		 * anti-windup: do not allow integrator to increase if actuator is at limit
		 */
		if (_last_output < -1.0f) {
			/* only allow motion to center: increase value */
			id = math::max(id, 0.0f);

		} else if (_last_output > 1.0f) {
			/* only allow motion to center: decrease value */
			id = math::min(id, 0.0f);
		}

		_integrator += id * _k_i;
	}

	/* integrator limit */
	//xxx: until start detection is available: integral part in control signal is limited here
	float integrator_constrained = math::constrain(_integrator, -_integrator_max, _integrator_max);

	/* Apply PI rate controller and store non-limited output */
	_last_output = (_bodyrate_setpoint * _k_ff + _rate_error * _k_p + integrator_constrained) * ctl_data.scaler *
		       ctl_data.scaler;  //scaler is proportional to 1/airspeed

	return math::constrain(_last_output, -1.0f, 1.0f);
}

float ECL_YawController::control_attitude_impl_accclosedloop(const struct ECL_ControlData &ctl_data)
{
	/* dont set a rate setpoint */
	return 0.0f;
}

float ECL_YawController::control_euler_rate(const struct ECL_ControlData &ctl_data)
{
	/* Transform setpoint to body angular rates (jacobian) */
	_bodyrate_setpoint = -sinf(ctl_data.roll) * ctl_data.pitch_rate_setpoint +
			     cosf(ctl_data.roll) * cosf(ctl_data.pitch) * _rate_setpoint;

	return control_bodyrate(ctl_data);

}
<commit_msg>yaw controller: for now do not do turn compensation when inverted and use the roll setpoint as limit for turn compensation<commit_after>/****************************************************************************
 *
 *   Copyright (c) 2013-2016 Estimation and Control Library (ECL). All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 * 3. Neither the name ECL nor the names of its contributors may be
 *    used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
 * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 ****************************************************************************/

/**
 * @file ecl_yaw_controller.cpp
 * Implementation of a simple orthogonal coordinated turn yaw PID controller.
 *
 * Authors and acknowledgements in header.
 */

#include "ecl_yaw_controller.h"
#include <stdint.h>
#include <float.h>
#include <geo/geo.h>
#include <ecl/ecl.h>
#include <mathlib/mathlib.h>
#include <systemlib/err.h>
#include <ecl/ecl.h>

ECL_YawController::ECL_YawController() :
	ECL_Controller("yaw"),
	_coordinated_min_speed(1.0f),
	_max_rate(0.0f), /* disable by default */
	_coordinated_method(0)
{
}

float ECL_YawController::control_attitude(const struct ECL_ControlData &ctl_data)
{
	switch (_coordinated_method) {
	case COORD_METHOD_OPEN:
		return control_attitude_impl_openloop(ctl_data);

	case COORD_METHOD_CLOSEACC:
		return control_attitude_impl_accclosedloop(ctl_data);

	default:
		static hrt_abstime last_print = 0;

		if (ecl_elapsed_time(&last_print) > 5e6) {
			warnx("invalid param setting FW_YCO_METHOD");
			last_print = ecl_absolute_time();
		}
	}

	return _rate_setpoint;
}

float ECL_YawController::control_attitude_impl_openloop(const struct ECL_ControlData &ctl_data)
{
	/* Do not calculate control signal with bad inputs */
	if (!(PX4_ISFINITE(ctl_data.roll) &&
	      PX4_ISFINITE(ctl_data.pitch) &&
	      PX4_ISFINITE(ctl_data.speed_body_u) &&
	      PX4_ISFINITE(ctl_data.speed_body_v) &&
	      PX4_ISFINITE(ctl_data.speed_body_w) &&
	      PX4_ISFINITE(ctl_data.roll_rate_setpoint) &&
	      PX4_ISFINITE(ctl_data.pitch_rate_setpoint))) {
		return _rate_setpoint;
	}

	float constrained_roll;
	bool inverted = false;
	/* roll is used as feedforward term and inverted flight needs to be considered */
	if (fabsf(ctl_data.roll) < math::radians(90.0f)) {
		/* not inverted, but numerically still potentially close to infinity */
		constrained_roll = math::constrain(ctl_data.roll, math::radians(-80.0f), math::radians(80.0f));

	} else {
		inverted = true;
		 // inverted flight, constrain on the two extremes of -pi..+pi to avoid infinity
		 //note: the ranges are extended by 10 deg here to avoid numeric resolution effects
		if (ctl_data.roll > 0.0f) {
			/* right hemisphere */
			constrained_roll = math::constrain(ctl_data.roll, math::radians(100.0f), math::radians(180.0f));

		} else {
			/* left hemisphere */
			constrained_roll = math::constrain(ctl_data.roll, math::radians(-180.0f), math::radians(-100.0f));
		}
	}

	constrained_roll = math::constrain(constrained_roll, -ctl_data.roll_setpoint, ctl_data.roll_setpoint);

	if (!inverted) {
		/* Calculate desired yaw rate from coordinated turn constraint / (no side forces) */
		_rate_setpoint = tanf(constrained_roll) * cosf(ctl_data.pitch) * 9.81f / (ctl_data.airspeed < ctl_data.airspeed_min ? ctl_data.airspeed_min : ctl_data.airspeed);
	}

	/* limit the rate */ //XXX: move to body angluar rates
	if (_max_rate > 0.01f) {
		_rate_setpoint = (_rate_setpoint > _max_rate) ? _max_rate : _rate_setpoint;
		_rate_setpoint = (_rate_setpoint < -_max_rate) ? -_max_rate : _rate_setpoint;
	}

	if (!PX4_ISFINITE(_rate_setpoint)) {
		warnx("yaw rate sepoint not finite");
		_rate_setpoint = 0.0f;
	}

	return _rate_setpoint;
}

float ECL_YawController::control_bodyrate(const struct ECL_ControlData &ctl_data)
{
	/* Do not calculate control signal with bad inputs */
	if (!(PX4_ISFINITE(ctl_data.roll) && PX4_ISFINITE(ctl_data.pitch) && PX4_ISFINITE(ctl_data.body_y_rate) &&
	      PX4_ISFINITE(ctl_data.body_z_rate) && PX4_ISFINITE(ctl_data.pitch_rate_setpoint) &&
	      PX4_ISFINITE(ctl_data.airspeed_min) && PX4_ISFINITE(ctl_data.airspeed_max) &&
	      PX4_ISFINITE(ctl_data.scaler))) {
		return math::constrain(_last_output, -1.0f, 1.0f);
	}

	/* get the usual dt estimate */
	uint64_t dt_micros = ecl_elapsed_time(&_last_run);
	_last_run = ecl_absolute_time();
	float dt = (float)dt_micros * 1e-6f;

	/* lock integral for long intervals */
	bool lock_integrator = ctl_data.lock_integrator;

	if (dt_micros > 500000) {
		lock_integrator = true;
	}

	/* input conditioning */
	float airspeed = ctl_data.airspeed;

	if (!PX4_ISFINITE(airspeed)) {
		/* airspeed is NaN, +- INF or not available, pick center of band */
		airspeed = 0.5f * (ctl_data.airspeed_min + ctl_data.airspeed_max);

	} else if (airspeed < ctl_data.airspeed_min) {
		airspeed = ctl_data.airspeed_min;
	}

	/* Close the acceleration loop if _coordinated_method wants this: change body_rate setpoint */
	if (_coordinated_method == COORD_METHOD_CLOSEACC) {
		// XXX lateral acceleration needs to go into integrator with a gain
		//_bodyrate_setpoint -= (ctl_data.acc_body_y / (airspeed * cosf(ctl_data.pitch)));
	}

	/* Calculate body angular rate error */
	_rate_error = _bodyrate_setpoint - ctl_data.body_z_rate; // body angular rate error

	if (!lock_integrator && _k_i > 0.0f && airspeed > 0.5f * ctl_data.airspeed_min) {

		float id = _rate_error * dt;

		/*
		 * anti-windup: do not allow integrator to increase if actuator is at limit
		 */
		if (_last_output < -1.0f) {
			/* only allow motion to center: increase value */
			id = math::max(id, 0.0f);

		} else if (_last_output > 1.0f) {
			/* only allow motion to center: decrease value */
			id = math::min(id, 0.0f);
		}

		_integrator += id * _k_i;
	}

	/* integrator limit */
	//xxx: until start detection is available: integral part in control signal is limited here
	float integrator_constrained = math::constrain(_integrator, -_integrator_max, _integrator_max);

	/* Apply PI rate controller and store non-limited output */
	_last_output = (_bodyrate_setpoint * _k_ff + _rate_error * _k_p + integrator_constrained) * ctl_data.scaler *
		       ctl_data.scaler;  //scaler is proportional to 1/airspeed

	return math::constrain(_last_output, -1.0f, 1.0f);
}

float ECL_YawController::control_attitude_impl_accclosedloop(const struct ECL_ControlData &ctl_data)
{
	/* dont set a rate setpoint */
	return 0.0f;
}

float ECL_YawController::control_euler_rate(const struct ECL_ControlData &ctl_data)
{
	/* Transform setpoint to body angular rates (jacobian) */
	_bodyrate_setpoint = -sinf(ctl_data.roll) * ctl_data.pitch_rate_setpoint +
			     cosf(ctl_data.roll) * cosf(ctl_data.pitch) * _rate_setpoint;

	return control_bodyrate(ctl_data);

}
<|endoftext|>


<commit_before>//===-- StraightLineStrengthReduce.cpp - ------------------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements straight-line strength reduction (SLSR). Unlike loop
// strength reduction, this algorithm is designed to reduce arithmetic
// redundancy in straight-line code instead of loops. It has proven to be
// effective in simplifying arithmetic statements derived from an unrolled loop.
// It can also simplify the logic of SeparateConstOffsetFromGEP.
//
// There are many optimizations we can perform in the domain of SLSR. This file
// for now contains only an initial step. Specifically, we look for strength
// reduction candidate in the form of
//
// (B + i) * S
//
// where B and S are integer constants or variables, and i is a constant
// integer. If we found two such candidates
//
// S1: X = (B + i) * S S2: Y = (B + i') * S
//
// and S1 dominates S2, we call S1 a basis of S2, and can replace S2 with
//
// Y = X + (i' - i) * S
//
// where (i' - i) * S is folded to the extent possible. When S2 has multiple
// bases, we pick the one that is closest to S2, or S2's "immediate" basis.
//
// TODO:
//
// - Handle candidates in the form of B + i * S
//
// - Handle candidates in the form of pointer arithmetics. e.g., B[i * S]
//
// - Floating point arithmetics when fast math is enabled.
//
// - SLSR may decrease ILP at the architecture level. Targets that are very
//   sensitive to ILP may want to disable it. Having SLSR to consider ILP is
//   left as future work.
#include <vector>

#include "llvm/ADT/DenseSet.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"

using namespace llvm;
using namespace PatternMatch;

namespace {

class StraightLineStrengthReduce : public FunctionPass {
 public:
  // SLSR candidate. Such a candidate must be in the form of
  //   (Base + Index) * Stride
  struct Candidate : public ilist_node<Candidate> {
    Candidate(Value *B = nullptr, ConstantInt *Idx = nullptr,
              Value *S = nullptr, Instruction *I = nullptr)
        : Base(B), Index(Idx), Stride(S), Ins(I), Basis(nullptr) {}
    Value *Base;
    ConstantInt *Index;
    Value *Stride;
    // The instruction this candidate corresponds to. It helps us to rewrite a
    // candidate with respect to its immediate basis. Note that one instruction
    // can corresponds to multiple candidates depending on how you associate the
    // expression. For instance,
    //
    // (a + 1) * (b + 2)
    //
    // can be treated as
    //
    // <Base: a, Index: 1, Stride: b + 2>
    //
    // or
    //
    // <Base: b, Index: 2, Stride: a + 1>
    Instruction *Ins;
    // Points to the immediate basis of this candidate, or nullptr if we cannot
    // find any basis for this candidate.
    Candidate *Basis;
  };

  static char ID;

  StraightLineStrengthReduce() : FunctionPass(ID), DT(nullptr) {
    initializeStraightLineStrengthReducePass(*PassRegistry::getPassRegistry());
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<DominatorTreeWrapperPass>();
    // We do not modify the shape of the CFG.
    AU.setPreservesCFG();
  }

  bool runOnFunction(Function &F) override;

 private:
  // Returns true if Basis is a basis for C, i.e., Basis dominates C and they
  // share the same base and stride.
  bool isBasisFor(const Candidate &Basis, const Candidate &C);
  // Checks whether I is in a candidate form. If so, adds all the matching forms
  // to Candidates, and tries to find the immediate basis for each of them.
  void allocateCandidateAndFindBasis(Instruction *I);
  // Given that I is in the form of "(B + Idx) * S", adds this form to
  // Candidates, and finds its immediate basis.
  void allocateCandidateAndFindBasis(Value *B, ConstantInt *Idx, Value *S,
                                     Instruction *I);
  // Rewrites candidate C with respect to Basis.
  void rewriteCandidateWithBasis(const Candidate &C, const Candidate &Basis);

  DominatorTree *DT;
  ilist<Candidate> Candidates;
  // Temporarily holds all instructions that are unlinked (but not deleted) by
  // rewriteCandidateWithBasis. These instructions will be actually removed
  // after all rewriting finishes.
  DenseSet<Instruction *> UnlinkedInstructions;
};
}  // anonymous namespace

char StraightLineStrengthReduce::ID = 0;
INITIALIZE_PASS_BEGIN(StraightLineStrengthReduce, "slsr",
                      "Straight line strength reduction", false, false)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
INITIALIZE_PASS_END(StraightLineStrengthReduce, "slsr",
                    "Straight line strength reduction", false, false)

FunctionPass *llvm::createStraightLineStrengthReducePass() {
  return new StraightLineStrengthReduce();
}

bool StraightLineStrengthReduce::isBasisFor(const Candidate &Basis,
                                            const Candidate &C) {
  return (Basis.Ins != C.Ins && // skip the same instruction
          // Basis must dominate C in order to rewrite C with respect to Basis.
          DT->dominates(Basis.Ins->getParent(), C.Ins->getParent()) &&
          // They share the same base and stride.
          Basis.Base == C.Base &&
          Basis.Stride == C.Stride);
}

// TODO: We currently implement an algorithm whose time complexity is linear to
// the number of existing candidates. However, a better algorithm exists. We
// could depth-first search the dominator tree, and maintain a hash table that
// contains all candidates that dominate the node being traversed.  This hash
// table is indexed by the base and the stride of a candidate.  Therefore,
// finding the immediate basis of a candidate boils down to one hash-table look
// up.
void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Value *B,
                                                               ConstantInt *Idx,
                                                               Value *S,
                                                               Instruction *I) {
  Candidate C(B, Idx, S, I);
  // Try to compute the immediate basis of C.
  unsigned NumIterations = 0;
  // Limit the scan radius to avoid running forever.
  static const int MaxNumIterations = 50;
  for (auto Basis = Candidates.rbegin();
       Basis != Candidates.rend() && NumIterations < MaxNumIterations;
       ++Basis, ++NumIterations) {
    if (isBasisFor(*Basis, C)) {
      C.Basis = &(*Basis);
      break;
    }
  }
  // Regardless of whether we find a basis for C, we need to push C to the
  // candidate list.
  Candidates.push_back(C);
}

void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Instruction *I) {
  Value *B = nullptr;
  ConstantInt *Idx = nullptr;
  // "(Base + Index) * Stride" must be a Mul instruction at the first hand.
  if (I->getOpcode() == Instruction::Mul) {
    if (IntegerType *ITy = dyn_cast<IntegerType>(I->getType())) {
      Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
      for (unsigned Swapped = 0; Swapped < 2; ++Swapped) {
        // Only handle the canonical operand ordering.
        if (match(LHS, m_Add(m_Value(B), m_ConstantInt(Idx)))) {
          // If LHS is in the form of "Base + Index", then I is in the form of
          // "(Base + Index) * RHS".
          allocateCandidateAndFindBasis(B, Idx, RHS, I);
        } else {
          // Otherwise, at least try the form (LHS + 0) * RHS.
          allocateCandidateAndFindBasis(LHS, ConstantInt::get(ITy, 0), RHS, I);
        }
        // Swap LHS and RHS so that we also cover the cases where LHS is the
        // stride.
        if (LHS == RHS)
          break;
        std::swap(LHS, RHS);
      }
    }
  }
}

void StraightLineStrengthReduce::rewriteCandidateWithBasis(
    const Candidate &C, const Candidate &Basis) {
  // An instruction can correspond to multiple candidates. Therefore, instead of
  // simply deleting an instruction when we rewrite it, we mark its parent as
  // nullptr (i.e. unlink it) so that we can skip the candidates whose
  // instruction is already rewritten.
  if (!C.Ins->getParent())
    return;
  assert(C.Base == Basis.Base && C.Stride == Basis.Stride);
  // Basis = (B + i) * S
  // C     = (B + i') * S
  //   ==>
  // C     = Basis + (i' - i) * S
  IRBuilder<> Builder(C.Ins);
  ConstantInt *IndexOffset = ConstantInt::get(
      C.Ins->getContext(), C.Index->getValue() - Basis.Index->getValue());
  Value *Reduced;
  // TODO: preserve nsw/nuw in some cases.
  if (IndexOffset->isOne()) {
    // If (i' - i) is 1, fold C into Basis + S.
    Reduced = Builder.CreateAdd(Basis.Ins, C.Stride);
  } else if (IndexOffset->isMinusOne()) {
    // If (i' - i) is -1, fold C into Basis - S.
    Reduced = Builder.CreateSub(Basis.Ins, C.Stride);
  } else {
    Value *Bump = Builder.CreateMul(C.Stride, IndexOffset);
    Reduced = Builder.CreateAdd(Basis.Ins, Bump);
  }
  Reduced->takeName(C.Ins);
  C.Ins->replaceAllUsesWith(Reduced);
  C.Ins->dropAllReferences();
  // Unlink C.Ins so that we can skip other candidates also corresponding to
  // C.Ins. The actual deletion is postponed to the end of runOnFunction.
  C.Ins->removeFromParent();
  UnlinkedInstructions.insert(C.Ins);
}

bool StraightLineStrengthReduce::runOnFunction(Function &F) {
  if (skipOptnoneFunction(F))
    return false;

  DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
  // Traverse the dominator tree in the depth-first order. This order makes sure
  // all bases of a candidate are in Candidates when we process it.
  for (auto node = GraphTraits<DominatorTree *>::nodes_begin(DT);
       node != GraphTraits<DominatorTree *>::nodes_end(DT); ++node) {
    BasicBlock *B = node->getBlock();
    for (auto I = B->begin(); I != B->end(); ++I) {
      allocateCandidateAndFindBasis(I);
    }
  }

  // Rewrite candidates in the reverse depth-first order. This order makes sure
  // a candidate being rewritten is not a basis for any other candidate.
  while (!Candidates.empty()) {
    const Candidate &C = Candidates.back();
    if (C.Basis != nullptr) {
      rewriteCandidateWithBasis(C, *C.Basis);
    }
    Candidates.pop_back();
  }

  // Delete all unlink instructions.
  for (auto I : UnlinkedInstructions) {
    delete I;
  }
  bool Ret = !UnlinkedInstructions.empty();
  UnlinkedInstructions.clear();
  return Ret;
}
<commit_msg>Fixing a -Wsign-compare warning; NFC<commit_after>//===-- StraightLineStrengthReduce.cpp - ------------------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements straight-line strength reduction (SLSR). Unlike loop
// strength reduction, this algorithm is designed to reduce arithmetic
// redundancy in straight-line code instead of loops. It has proven to be
// effective in simplifying arithmetic statements derived from an unrolled loop.
// It can also simplify the logic of SeparateConstOffsetFromGEP.
//
// There are many optimizations we can perform in the domain of SLSR. This file
// for now contains only an initial step. Specifically, we look for strength
// reduction candidate in the form of
//
// (B + i) * S
//
// where B and S are integer constants or variables, and i is a constant
// integer. If we found two such candidates
//
// S1: X = (B + i) * S S2: Y = (B + i') * S
//
// and S1 dominates S2, we call S1 a basis of S2, and can replace S2 with
//
// Y = X + (i' - i) * S
//
// where (i' - i) * S is folded to the extent possible. When S2 has multiple
// bases, we pick the one that is closest to S2, or S2's "immediate" basis.
//
// TODO:
//
// - Handle candidates in the form of B + i * S
//
// - Handle candidates in the form of pointer arithmetics. e.g., B[i * S]
//
// - Floating point arithmetics when fast math is enabled.
//
// - SLSR may decrease ILP at the architecture level. Targets that are very
//   sensitive to ILP may want to disable it. Having SLSR to consider ILP is
//   left as future work.
#include <vector>

#include "llvm/ADT/DenseSet.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"

using namespace llvm;
using namespace PatternMatch;

namespace {

class StraightLineStrengthReduce : public FunctionPass {
 public:
  // SLSR candidate. Such a candidate must be in the form of
  //   (Base + Index) * Stride
  struct Candidate : public ilist_node<Candidate> {
    Candidate(Value *B = nullptr, ConstantInt *Idx = nullptr,
              Value *S = nullptr, Instruction *I = nullptr)
        : Base(B), Index(Idx), Stride(S), Ins(I), Basis(nullptr) {}
    Value *Base;
    ConstantInt *Index;
    Value *Stride;
    // The instruction this candidate corresponds to. It helps us to rewrite a
    // candidate with respect to its immediate basis. Note that one instruction
    // can corresponds to multiple candidates depending on how you associate the
    // expression. For instance,
    //
    // (a + 1) * (b + 2)
    //
    // can be treated as
    //
    // <Base: a, Index: 1, Stride: b + 2>
    //
    // or
    //
    // <Base: b, Index: 2, Stride: a + 1>
    Instruction *Ins;
    // Points to the immediate basis of this candidate, or nullptr if we cannot
    // find any basis for this candidate.
    Candidate *Basis;
  };

  static char ID;

  StraightLineStrengthReduce() : FunctionPass(ID), DT(nullptr) {
    initializeStraightLineStrengthReducePass(*PassRegistry::getPassRegistry());
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<DominatorTreeWrapperPass>();
    // We do not modify the shape of the CFG.
    AU.setPreservesCFG();
  }

  bool runOnFunction(Function &F) override;

 private:
  // Returns true if Basis is a basis for C, i.e., Basis dominates C and they
  // share the same base and stride.
  bool isBasisFor(const Candidate &Basis, const Candidate &C);
  // Checks whether I is in a candidate form. If so, adds all the matching forms
  // to Candidates, and tries to find the immediate basis for each of them.
  void allocateCandidateAndFindBasis(Instruction *I);
  // Given that I is in the form of "(B + Idx) * S", adds this form to
  // Candidates, and finds its immediate basis.
  void allocateCandidateAndFindBasis(Value *B, ConstantInt *Idx, Value *S,
                                     Instruction *I);
  // Rewrites candidate C with respect to Basis.
  void rewriteCandidateWithBasis(const Candidate &C, const Candidate &Basis);

  DominatorTree *DT;
  ilist<Candidate> Candidates;
  // Temporarily holds all instructions that are unlinked (but not deleted) by
  // rewriteCandidateWithBasis. These instructions will be actually removed
  // after all rewriting finishes.
  DenseSet<Instruction *> UnlinkedInstructions;
};
}  // anonymous namespace

char StraightLineStrengthReduce::ID = 0;
INITIALIZE_PASS_BEGIN(StraightLineStrengthReduce, "slsr",
                      "Straight line strength reduction", false, false)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
INITIALIZE_PASS_END(StraightLineStrengthReduce, "slsr",
                    "Straight line strength reduction", false, false)

FunctionPass *llvm::createStraightLineStrengthReducePass() {
  return new StraightLineStrengthReduce();
}

bool StraightLineStrengthReduce::isBasisFor(const Candidate &Basis,
                                            const Candidate &C) {
  return (Basis.Ins != C.Ins && // skip the same instruction
          // Basis must dominate C in order to rewrite C with respect to Basis.
          DT->dominates(Basis.Ins->getParent(), C.Ins->getParent()) &&
          // They share the same base and stride.
          Basis.Base == C.Base &&
          Basis.Stride == C.Stride);
}

// TODO: We currently implement an algorithm whose time complexity is linear to
// the number of existing candidates. However, a better algorithm exists. We
// could depth-first search the dominator tree, and maintain a hash table that
// contains all candidates that dominate the node being traversed.  This hash
// table is indexed by the base and the stride of a candidate.  Therefore,
// finding the immediate basis of a candidate boils down to one hash-table look
// up.
void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Value *B,
                                                               ConstantInt *Idx,
                                                               Value *S,
                                                               Instruction *I) {
  Candidate C(B, Idx, S, I);
  // Try to compute the immediate basis of C.
  unsigned NumIterations = 0;
  // Limit the scan radius to avoid running forever.
  static const unsigned MaxNumIterations = 50;
  for (auto Basis = Candidates.rbegin();
       Basis != Candidates.rend() && NumIterations < MaxNumIterations;
       ++Basis, ++NumIterations) {
    if (isBasisFor(*Basis, C)) {
      C.Basis = &(*Basis);
      break;
    }
  }
  // Regardless of whether we find a basis for C, we need to push C to the
  // candidate list.
  Candidates.push_back(C);
}

void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Instruction *I) {
  Value *B = nullptr;
  ConstantInt *Idx = nullptr;
  // "(Base + Index) * Stride" must be a Mul instruction at the first hand.
  if (I->getOpcode() == Instruction::Mul) {
    if (IntegerType *ITy = dyn_cast<IntegerType>(I->getType())) {
      Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
      for (unsigned Swapped = 0; Swapped < 2; ++Swapped) {
        // Only handle the canonical operand ordering.
        if (match(LHS, m_Add(m_Value(B), m_ConstantInt(Idx)))) {
          // If LHS is in the form of "Base + Index", then I is in the form of
          // "(Base + Index) * RHS".
          allocateCandidateAndFindBasis(B, Idx, RHS, I);
        } else {
          // Otherwise, at least try the form (LHS + 0) * RHS.
          allocateCandidateAndFindBasis(LHS, ConstantInt::get(ITy, 0), RHS, I);
        }
        // Swap LHS and RHS so that we also cover the cases where LHS is the
        // stride.
        if (LHS == RHS)
          break;
        std::swap(LHS, RHS);
      }
    }
  }
}

void StraightLineStrengthReduce::rewriteCandidateWithBasis(
    const Candidate &C, const Candidate &Basis) {
  // An instruction can correspond to multiple candidates. Therefore, instead of
  // simply deleting an instruction when we rewrite it, we mark its parent as
  // nullptr (i.e. unlink it) so that we can skip the candidates whose
  // instruction is already rewritten.
  if (!C.Ins->getParent())
    return;
  assert(C.Base == Basis.Base && C.Stride == Basis.Stride);
  // Basis = (B + i) * S
  // C     = (B + i') * S
  //   ==>
  // C     = Basis + (i' - i) * S
  IRBuilder<> Builder(C.Ins);
  ConstantInt *IndexOffset = ConstantInt::get(
      C.Ins->getContext(), C.Index->getValue() - Basis.Index->getValue());
  Value *Reduced;
  // TODO: preserve nsw/nuw in some cases.
  if (IndexOffset->isOne()) {
    // If (i' - i) is 1, fold C into Basis + S.
    Reduced = Builder.CreateAdd(Basis.Ins, C.Stride);
  } else if (IndexOffset->isMinusOne()) {
    // If (i' - i) is -1, fold C into Basis - S.
    Reduced = Builder.CreateSub(Basis.Ins, C.Stride);
  } else {
    Value *Bump = Builder.CreateMul(C.Stride, IndexOffset);
    Reduced = Builder.CreateAdd(Basis.Ins, Bump);
  }
  Reduced->takeName(C.Ins);
  C.Ins->replaceAllUsesWith(Reduced);
  C.Ins->dropAllReferences();
  // Unlink C.Ins so that we can skip other candidates also corresponding to
  // C.Ins. The actual deletion is postponed to the end of runOnFunction.
  C.Ins->removeFromParent();
  UnlinkedInstructions.insert(C.Ins);
}

bool StraightLineStrengthReduce::runOnFunction(Function &F) {
  if (skipOptnoneFunction(F))
    return false;

  DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
  // Traverse the dominator tree in the depth-first order. This order makes sure
  // all bases of a candidate are in Candidates when we process it.
  for (auto node = GraphTraits<DominatorTree *>::nodes_begin(DT);
       node != GraphTraits<DominatorTree *>::nodes_end(DT); ++node) {
    BasicBlock *B = node->getBlock();
    for (auto I = B->begin(); I != B->end(); ++I) {
      allocateCandidateAndFindBasis(I);
    }
  }

  // Rewrite candidates in the reverse depth-first order. This order makes sure
  // a candidate being rewritten is not a basis for any other candidate.
  while (!Candidates.empty()) {
    const Candidate &C = Candidates.back();
    if (C.Basis != nullptr) {
      rewriteCandidateWithBasis(C, *C.Basis);
    }
    Candidates.pop_back();
  }

  // Delete all unlink instructions.
  for (auto I : UnlinkedInstructions) {
    delete I;
  }
  bool Ret = !UnlinkedInstructions.empty();
  UnlinkedInstructions.clear();
  return Ret;
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// ----------------------------------------------------------------------------
// Copyright 2020-2021 Arm Limited
//
// 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.
// ----------------------------------------------------------------------------

/**
 * @brief Platform-specific function implementations.
 *
 * This module contains functions for querying the host extended ISA support.
 */

// Include before the defines below to pick up any auto-setup based on compiler
// built-in config, if not being set explicitly by the build system
#include "astcenc_internal.h"

#if (ASTCENC_SSE > 0)    || (ASTCENC_AVX > 0) || \
    (ASTCENC_POPCNT > 0) || (ASTCENC_F16C > 0)

static bool g_init { false };

/** Does this CPU support SSE 4.1? Set to -1 if not yet initialized. */
static bool g_cpu_has_sse41 { false };

/** Does this CPU support AVX2? Set to -1 if not yet initialized. */
static bool g_cpu_has_avx2 { false };

/** Does this CPU support POPCNT? Set to -1 if not yet initialized. */
static bool g_cpu_has_popcnt { false };

/** Does this CPU support F16C? Set to -1 if not yet initialized. */
static bool g_cpu_has_f16c { false };

/* ============================================================================
   Platform code for Visual Studio
============================================================================ */
#if !defined(__clang__) && defined(_MSC_VER)
#include <intrin.h>

/**
 * @brief Detect platform CPU ISA support and update global trackers.
 */
static void detect_cpu_isa()
{
	int data[4];

	__cpuid(data, 0);
	int num_id = data[0];

	if (num_id >= 1)
	{
		__cpuidex(data, 1, 0);
		// SSE41 = Bank 1, ECX, bit 19
		g_cpu_has_sse41 = data[2] & (1 << 19) ? true : false;
		// POPCNT = Bank 1, ECX, bit 23
		g_cpu_has_popcnt = data[2] & (1 << 23) ? true : false;
		// F16C = Bank 1, ECX, bit 29
		g_cpu_has_f16c = data[2] & (1 << 29) ? true : false;
	}

	if (num_id >= 7)
	{
		__cpuidex(data, 7, 0);
		// AVX2 = Bank 7, EBX, bit 5
		g_cpu_has_avx2 = data[1] & (1 << 5) ? true : false;
	}

	g_init = true;
}

/* ============================================================================
   Platform code for GCC and Clang
============================================================================ */
#else
#include <cpuid.h>

/**
 * @brief Detect platform CPU ISA support and update global trackers.
 */
static void detect_cpu_isa()
{
	unsigned int data[4];

	if (__get_cpuid_count(1, 0, &data[0], &data[1], &data[2], &data[3]))
	{
		// SSE41 = Bank 1, ECX, bit 19
		g_cpu_has_sse41 = data[2] & (1 << 19) ? true : false;
		// POPCNT = Bank 1, ECX, bit 23
		g_cpu_has_popcnt = data[2] & (1 << 23) ? true : false;
		// F16C = Bank 1, ECX, bit 29
		g_cpu_has_f16c = data[2] & (1 << 29) ? true : false;
	}

	g_cpu_has_avx2 = 0;
	if (__get_cpuid_count(7, 0, &data[0], &data[1], &data[2], &data[3]))
	{
		// AVX2 = Bank 7, EBX, bit 5
		g_cpu_has_avx2 = data[1] & (1 << 5) ? true : false;
	}

	g_init = true;
}
#endif

/* See header for documentation. */
bool cpu_supports_popcnt()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_popcnt;
}

/* See header for documentation. */
bool cpu_supports_f16c()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_f16c;
}

/* See header for documentation. */
bool cpu_supports_sse41()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_sse41;
}

/* See header for documentation. */
bool cpu_supports_avx2()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_avx2;
}

#endif
<commit_msg>Add thread syncs on CPU state tracker updates<commit_after>// SPDX-License-Identifier: Apache-2.0
// ----------------------------------------------------------------------------
// Copyright 2020-2021 Arm Limited
//
// 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.
// ----------------------------------------------------------------------------

/**
 * @brief Platform-specific function implementations.
 *
 * This module contains functions for querying the host extended ISA support.
 */

// Include before the defines below to pick up any auto-setup based on compiler
// built-in config, if not being set explicitly by the build system
#include "astcenc_internal.h"

#if (ASTCENC_SSE > 0)    || (ASTCENC_AVX > 0) || \
    (ASTCENC_POPCNT > 0) || (ASTCENC_F16C > 0)

static bool g_init { false };

/** Does this CPU support SSE 4.1? Set to -1 if not yet initialized. */
static bool g_cpu_has_sse41 { false };

/** Does this CPU support AVX2? Set to -1 if not yet initialized. */
static bool g_cpu_has_avx2 { false };

/** Does this CPU support POPCNT? Set to -1 if not yet initialized. */
static bool g_cpu_has_popcnt { false };

/** Does this CPU support F16C? Set to -1 if not yet initialized. */
static bool g_cpu_has_f16c { false };

/* ============================================================================
   Platform code for Visual Studio
============================================================================ */
#if !defined(__clang__) && defined(_MSC_VER)
#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <intrin.h>

/**
 * @brief Detect platform CPU ISA support and update global trackers.
 */
static void detect_cpu_isa()
{
	int data[4];

	__cpuid(data, 0);
	int num_id = data[0];

	if (num_id >= 1)
	{
		__cpuidex(data, 1, 0);
		// SSE41 = Bank 1, ECX, bit 19
		g_cpu_has_sse41 = data[2] & (1 << 19) ? true : false;
		// POPCNT = Bank 1, ECX, bit 23
		g_cpu_has_popcnt = data[2] & (1 << 23) ? true : false;
		// F16C = Bank 1, ECX, bit 29
		g_cpu_has_f16c = data[2] & (1 << 29) ? true : false;
	}

	if (num_id >= 7)
	{
		__cpuidex(data, 7, 0);
		// AVX2 = Bank 7, EBX, bit 5
		g_cpu_has_avx2 = data[1] & (1 << 5) ? true : false;
	}

	// Ensure state bits are updated before init flag is updated
	MemoryBarrier();
	g_init = true;
}

/* ============================================================================
   Platform code for GCC and Clang
============================================================================ */
#else
#include <cpuid.h>

/**
 * @brief Detect platform CPU ISA support and update global trackers.
 */
static void detect_cpu_isa()
{
	unsigned int data[4];

	if (__get_cpuid_count(1, 0, &data[0], &data[1], &data[2], &data[3]))
	{
		// SSE41 = Bank 1, ECX, bit 19
		g_cpu_has_sse41 = data[2] & (1 << 19) ? true : false;
		// POPCNT = Bank 1, ECX, bit 23
		g_cpu_has_popcnt = data[2] & (1 << 23) ? true : false;
		// F16C = Bank 1, ECX, bit 29
		g_cpu_has_f16c = data[2] & (1 << 29) ? true : false;
	}

	g_cpu_has_avx2 = 0;
	if (__get_cpuid_count(7, 0, &data[0], &data[1], &data[2], &data[3]))
	{
		// AVX2 = Bank 7, EBX, bit 5
		g_cpu_has_avx2 = data[1] & (1 << 5) ? true : false;
	}

	// Ensure state bits are updated before init flag is updated
	__sync_synchronize();
	g_init = true;
}
#endif

/* See header for documentation. */
bool cpu_supports_popcnt()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_popcnt;
}

/* See header for documentation. */
bool cpu_supports_f16c()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_f16c;
}

/* See header for documentation. */
bool cpu_supports_sse41()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_sse41;
}

/* See header for documentation. */
bool cpu_supports_avx2()
{
	if (!g_init)
	{
		detect_cpu_isa();
	}

	return g_cpu_has_avx2;
}

#endif
<|endoftext|>


<commit_before>#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <functional>
#include <random>
#include <vector>
#include <string>

#include <gauss_multi.h>
#include <gp.h>

using namespace std;
typedef unsigned int uint;

// Random engine to use throughout
default_random_engine gen;

// How many functions to sample at each step
int n_samp = 5;

void to_tikz(vector<vector<double>> X, vector<vector<double>> Y, string fname)
{
	assert(X.size() == Y.size());
	FILE *f = fopen(fname.c_str(), "w");

	// Define colours (TODO-someday: make this parametrised...)
	fprintf(f, "\\definecolor{mycolor0}{rgb}{0.00000,0.44700,0.74100}%%\n");
	fprintf(f, "\\definecolor{mycolor1}{rgb}{0.85000,0.32500,0.09800}%%\n");
	fprintf(f, "\\definecolor{mycolor2}{rgb}{0.92900,0.69400,0.12500}%%\n");
	fprintf(f, "\\definecolor{mycolor3}{rgb}{0.49400,0.18400,0.55600}%%\n");
	fprintf(f, "\\definecolor{mycolor4}{rgb}{0.46600,0.67400,0.18800}%%\n");

	// Begin picture
	fprintf(f, "\\begin{tikzpicture}[very thick]\n");
	// Begin axis (with all parameters)
	fprintf(f, "\\begin{axis}[width=6.028in, height=4.754in,\n");
	fprintf(f, "scale only axis, xmin=-1, xmax=1, ymin=-5.0, ymax=5.0,\n");
	fprintf(f, "axis background/.style={fill=white}]\n");

	// Add plots
	for (uint i=0;i<X.size();i++)
	{
		assert(X[i].size() == Y[i].size());
		fprintf(f, "\\addplot[color=mycolor%d, solid] table[row sep=crcr]{%%\n", i);
		for (uint j=0;j<X[i].size();j++)
		{
			fprintf(f, "%.10lf %.10lf\\\\\n", X[i][j], Y[i][j]);
		}
		fprintf(f, "};\n");
	}

	fprintf(f, "\\end{axis}\n");
	fprintf(f, "\\end{tikzpicture}\n");

	fclose(f);
}

int main()
{
	// Create a new Gaussian process...
	// use zero-mean, squared-exponential covariance (hyperparam l = 1.0), no noise.
	auto m = [](double) { return 0; };
	auto k = [](double x, double y) { return exp(-(x - y) * (x - y) / (2.0 * 1.0)); };
	GP gp(m, k, 0.0);

	// points to be used to plot lines
	vector<double> xs;
	for (double x=-1.0;x<=1.0;x+=0.001) xs.push_back(x);

	// Sample the prior
	vector<double> mu = gp.get_means(xs);
	vector<vector<double>> Sigma = gp.get_covar(xs);

	/*for (uint i=0;i<Sigma.size();i++)
	{
		for (uint j=0;j<Sigma.size();j++)
		{
			printf("%.2lf ", Sigma[i][j]);
		}
		printf("\n");
	}*/

	MultiGaussian N(gen, mu, Sigma);

	vector<vector<double>> Xs(n_samp, xs);
	vector<vector<double>> Ys(n_samp);

	for (int i=0;i<n_samp;i++)
	{
		Ys[i] = N.sample();
	}

	to_tikz(Xs, Ys, "test.tex");

	return 0;
}
<commit_msg>OK parameters, need to fix posterior<commit_after>#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <functional>
#include <random>
#include <vector>
#include <string>

#include <gauss_multi.h>
#include <gp.h>

using namespace std;
typedef unsigned int uint;

// Random engine to use throughout
default_random_engine gen;

// How many functions to sample at each step
int n_samp = 5;

void to_tikz(vector<vector<double>> X, vector<vector<double>> Y, string fname)
{
	assert(X.size() == Y.size());
	FILE *f = fopen(fname.c_str(), "w");

	// Define colours (TODO-someday: make this parametrised...)
	fprintf(f, "\\definecolor{mycolor0}{rgb}{0.00000,0.44700,0.74100}%%\n");
	fprintf(f, "\\definecolor{mycolor1}{rgb}{0.85000,0.32500,0.09800}%%\n");
	fprintf(f, "\\definecolor{mycolor2}{rgb}{0.92900,0.69400,0.12500}%%\n");
	fprintf(f, "\\definecolor{mycolor3}{rgb}{0.49400,0.18400,0.55600}%%\n");
	fprintf(f, "\\definecolor{mycolor4}{rgb}{0.46600,0.67400,0.18800}%%\n");

	// Begin picture
	fprintf(f, "\\begin{tikzpicture}[very thick]\n");
	// Begin axis (with all parameters)
	fprintf(f, "\\begin{axis}[very thick, width=6.028in, height=4.754in,\n");
	fprintf(f, "scale only axis, xmin=-5, xmax=5, ymin=-2, ymax=2,\n");
	fprintf(f, "axis background/.style={fill=white}]\n");

	// Add plots
	for (uint i=0;i<X.size();i++)
	{
		assert(X[i].size() == Y[i].size());
		fprintf(f, "\\addplot[color=mycolor%d, solid, smooth] table[row sep=crcr]{%%\n", i);
		for (uint j=0;j<X[i].size();j++)
		{
			fprintf(f, "%.10lf %.10lf\\\\\n", X[i][j], Y[i][j]);
		}
		fprintf(f, "};\n");
	}

	fprintf(f, "\\end{axis}\n");
	fprintf(f, "\\end{tikzpicture}\n");

	fclose(f);
}

int main()
{
	// Create a new Gaussian process...
	// use zero-mean, squared-exponential covariance (hyperparam l = 0.5), no noise.
	auto m = [](double) { return 0; };
	auto k = [](double x, double y) { return exp(-(x - y) * (x - y) / (2.0 * 0.5 * 0.5)); };
	GP gp(m, k, 0.0);

	// points to be used to plot lines
	vector<double> xs;
	for (double x=-5.0;x<=5.0;x+=0.02) xs.push_back(x);

	// Sample the prior
	vector<double> mu = gp.get_means(xs);
	vector<vector<double>> Sigma = gp.get_covar(xs);

	/*for (uint i=0;i<Sigma.size();i++)
	{
		for (uint j=0;j<Sigma.size();j++)
		{
			printf("%.2lf ", Sigma[i][j]);
		}
		printf("\n");
	}*/

	MultiGaussian N(gen, mu, Sigma);

	vector<vector<double>> Xs(n_samp, xs);
	vector<vector<double>> Ys(n_samp);

	for (int i=0;i<n_samp;i++)
	{
		Ys[i] = N.sample();
	}

	to_tikz(Xs, Ys, "test.tex");

	return 0;
}
<|endoftext|>