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

#include "net/base/host_resolver.h"

#if defined(OS_WIN)
#include <ws2tcpip.h>
#include <wspiapi.h>
#elif defined(OS_POSIX)
#include <netdb.h>
#endif

#include <string>

#include "base/compiler_specific.h"
#include "base/message_loop.h"
#include "base/ref_counted.h"
#include "net/base/address_list.h"
#include "net/base/completion_callback.h"
#include "net/base/host_resolver_unittest.h"
#include "net/base/net_errors.h"
#include "testing/gtest/include/gtest/gtest.h"

using net::RuleBasedHostMapper;
using net::ScopedHostMapper;
using net::WaitingHostMapper;

namespace {

class HostResolverTest : public testing::Test {
 public:
  HostResolverTest()
      : callback_called_(false),
        ALLOW_THIS_IN_INITIALIZER_LIST(
            callback_(this, &HostResolverTest::OnLookupFinished)) {
  }

 protected:
  bool callback_called_;
  int callback_result_;
  net::CompletionCallbackImpl<HostResolverTest> callback_;

 private:
  void OnLookupFinished(int result) {
    callback_called_ = true;
    callback_result_ = result;
    MessageLoop::current()->Quit();
  }
};

TEST_F(HostResolverTest, SynchronousLookup) {
  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 80;

  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AddRule("just.testing", "192.168.1.42");
  ScopedHostMapper scoped_mapper(mapper.get());

  int err = host_resolver.Resolve("just.testing", kPortnum, &adrlist, NULL);
  EXPECT_EQ(net::OK, err);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in* sa_in = (const struct sockaddr_in*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in->sin_port);
  EXPECT_TRUE(htonl(0xc0a8012a) == sa_in->sin_addr.s_addr);
}

TEST_F(HostResolverTest, AsynchronousLookup) {
  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 80;

  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AddRule("just.testing", "192.168.1.42");
  ScopedHostMapper scoped_mapper(mapper.get());

  int err = host_resolver.Resolve("just.testing", kPortnum, &adrlist,
                                  &callback_);
  EXPECT_EQ(net::ERR_IO_PENDING, err);

  MessageLoop::current()->Run();

  ASSERT_TRUE(callback_called_);
  ASSERT_EQ(net::OK, callback_result_);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in* sa_in = (const struct sockaddr_in*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in->sin_port);
  EXPECT_TRUE(htonl(0xc0a8012a) == sa_in->sin_addr.s_addr);
}

TEST_F(HostResolverTest, CanceledAsynchronousLookup) {
  scoped_refptr<WaitingHostMapper> mapper = new WaitingHostMapper();
  ScopedHostMapper scoped_mapper(mapper.get());

  {
    net::HostResolver host_resolver;
    net::AddressList adrlist;
    const int kPortnum = 80;

    int err = host_resolver.Resolve("just.testing", kPortnum, &adrlist,
                                    &callback_);
    EXPECT_EQ(net::ERR_IO_PENDING, err);

    // Make sure we will exit the queue even when callback is not called.
    MessageLoop::current()->PostDelayedTask(FROM_HERE,
                                            new MessageLoop::QuitTask(),
                                            1000);
    MessageLoop::current()->Run();
  }

  mapper->Signal();

  EXPECT_FALSE(callback_called_);
}

TEST_F(HostResolverTest, NumericAddresses) {
  // Stevens says dotted quads with AI_UNSPEC resolve to a single sockaddr_in.

  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 5555;
  int err = host_resolver.Resolve("127.0.0.1", kPortnum, &adrlist, NULL);
  EXPECT_EQ(net::OK, err);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in* sa_in = (const struct sockaddr_in*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in->sin_port);
  EXPECT_TRUE(htonl(0x7f000001) == sa_in->sin_addr.s_addr);
}

} // namespace
<commit_msg>Fix up some unit tests for HostResolver:<commit_after>// Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/base/host_resolver.h"

#if defined(OS_WIN)
#include <ws2tcpip.h>
#include <wspiapi.h>
#elif defined(OS_POSIX)
#include <netdb.h>
#endif

#include <string>

#include "base/compiler_specific.h"
#include "base/message_loop.h"
#include "base/ref_counted.h"
#include "net/base/address_list.h"
#include "net/base/completion_callback.h"
#include "net/base/host_resolver_unittest.h"
#include "net/base/net_errors.h"
#include "testing/gtest/include/gtest/gtest.h"

using net::RuleBasedHostMapper;
using net::ScopedHostMapper;
using net::WaitingHostMapper;

namespace {

class HostResolverTest : public testing::Test {
 public:
  HostResolverTest()
      : callback_called_(false),
        ALLOW_THIS_IN_INITIALIZER_LIST(
            callback_(this, &HostResolverTest::OnLookupFinished)) {
  }

 protected:
  bool callback_called_;
  int callback_result_;
  net::CompletionCallbackImpl<HostResolverTest> callback_;

 private:
  void OnLookupFinished(int result) {
    callback_called_ = true;
    callback_result_ = result;
    MessageLoop::current()->Quit();
  }
};

TEST_F(HostResolverTest, SynchronousLookup) {
  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 80;

  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AddRule("just.testing", "192.168.1.42");
  ScopedHostMapper scoped_mapper(mapper.get());

  int err = host_resolver.Resolve("just.testing", kPortnum, &adrlist, NULL);
  EXPECT_EQ(net::OK, err);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in* sa_in = (const struct sockaddr_in*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in->sin_port);
  EXPECT_TRUE(htonl(0xc0a8012a) == sa_in->sin_addr.s_addr);
}

TEST_F(HostResolverTest, AsynchronousLookup) {
  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 80;

  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AddRule("just.testing", "192.168.1.42");
  ScopedHostMapper scoped_mapper(mapper.get());

  int err = host_resolver.Resolve("just.testing", kPortnum, &adrlist,
                                  &callback_);
  EXPECT_EQ(net::ERR_IO_PENDING, err);

  MessageLoop::current()->Run();

  ASSERT_TRUE(callback_called_);
  ASSERT_EQ(net::OK, callback_result_);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in* sa_in = (const struct sockaddr_in*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in->sin_port);
  EXPECT_TRUE(htonl(0xc0a8012a) == sa_in->sin_addr.s_addr);
}

TEST_F(HostResolverTest, CanceledAsynchronousLookup) {
  scoped_refptr<WaitingHostMapper> mapper = new WaitingHostMapper();
  ScopedHostMapper scoped_mapper(mapper.get());

  {
    net::HostResolver host_resolver;
    net::AddressList adrlist;
    const int kPortnum = 80;

    int err = host_resolver.Resolve("just.testing", kPortnum, &adrlist,
                                    &callback_);
    EXPECT_EQ(net::ERR_IO_PENDING, err);

    // Make sure we will exit the queue even when callback is not called.
    MessageLoop::current()->PostDelayedTask(FROM_HERE,
                                            new MessageLoop::QuitTask(),
                                            1000);
    MessageLoop::current()->Run();
  }

  mapper->Signal();

  EXPECT_FALSE(callback_called_);
}

TEST_F(HostResolverTest, NumericIPv4Address) {
  // Stevens says dotted quads with AI_UNSPEC resolve to a single sockaddr_in.

  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AllowDirectLookup("*");
  ScopedHostMapper scoped_mapper(mapper.get());

  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 5555;
  int err = host_resolver.Resolve("127.1.2.3", kPortnum, &adrlist, NULL);
  EXPECT_EQ(net::OK, err);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in* sa_in = (const struct sockaddr_in*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in->sin_port);
  EXPECT_TRUE(htonl(0x7f010203) == sa_in->sin_addr.s_addr);
}

TEST_F(HostResolverTest, NumericIPv6Address) {
  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AllowDirectLookup("*");
  ScopedHostMapper scoped_mapper(mapper.get());

  // Resolve a plain IPv6 address.  Don't worry about [brackets], because
  // the caller should have removed them.
  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 5555;
  int err = host_resolver.Resolve("2001:db8::1", kPortnum, &adrlist, NULL);
  // On computers without IPv6 support, getaddrinfo cannot convert IPv6
  // address literals to addresses (getaddrinfo returns EAI_NONAME).  So this
  // test has to allow host_resolver.Resolve to fail.
  if (err == net::ERR_NAME_NOT_RESOLVED)
    return;
  EXPECT_EQ(net::OK, err);

  const struct addrinfo* ainfo = adrlist.head();
  EXPECT_EQ(static_cast<addrinfo*>(NULL), ainfo->ai_next);
  EXPECT_EQ(sizeof(struct sockaddr_in6), ainfo->ai_addrlen);

  const struct sockaddr* sa = ainfo->ai_addr;
  const struct sockaddr_in6* sa_in6 = (const struct sockaddr_in6*) sa;
  EXPECT_TRUE(htons(kPortnum) == sa_in6->sin6_port);

  const uint8 expect_addr[] = {
    0x20, 0x01, 0x0d, 0xb8, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01
  };
  for (int i = 0; i < 16; i++) {
    EXPECT_EQ(expect_addr[i], sa_in6->sin6_addr.s6_addr[i]);
  }
}

TEST_F(HostResolverTest, EmptyHost) {
  scoped_refptr<RuleBasedHostMapper> mapper = new RuleBasedHostMapper();
  mapper->AllowDirectLookup("*");
  ScopedHostMapper scoped_mapper(mapper.get());

  net::HostResolver host_resolver;
  net::AddressList adrlist;
  const int kPortnum = 5555;
  int err = host_resolver.Resolve("", kPortnum, &adrlist, NULL);
  EXPECT_EQ(net::ERR_NAME_NOT_RESOLVED, err);
}

}  // namespace
<|endoftext|>


<commit_before>#include "signverifymessagedialog.h"
#include "ui_signverifymessagedialog.h"

#include "addressbookpage.h"
#include "base58.h"
#include "guiutil.h"
#include "init.h"
#include "main.h"
#include "optionsmodel.h"
#include "walletmodel.h"
#include "wallet.h"

#include <string>
#include <vector>

#include <QClipboard>

SignVerifyMessageDialog::SignVerifyMessageDialog(QWidget *parent) :
    QDialog(parent),
    ui(new Ui::SignVerifyMessageDialog),
    model(0)
{
    ui->setupUi(this);

#if (QT_VERSION >= 0x040700)
    /* Do not move this to the XML file, Qt before 4.7 will choke on it */
    ui->addressIn_SM->setPlaceholderText(tr("Enter a YaCoin address (e.g. 4Zo1ga6xuKuQ7JV7M9rGDoxdbYwV5zgQJ5)"));
    ui->signatureOut_SM->setPlaceholderText(tr("Click \"Sign Message\" to generate signature"));

    ui->addressIn_VM->setPlaceholderText(tr("Enter a YaCoin address (e.g. 4Zo1ga6xuKuQ7JV7M9rGDoxdbYwV5zgQJ5)"));
    ui->signatureIn_VM->setPlaceholderText(tr("Enter YaCoin signature"));
#endif

    GUIUtil::setupAddressWidget(ui->addressIn_SM, this);
    GUIUtil::setupAddressWidget(ui->addressIn_VM, this);

    ui->addressIn_SM->installEventFilter(this);
    ui->messageIn_SM->installEventFilter(this);
    ui->signatureOut_SM->installEventFilter(this);
    ui->addressIn_VM->installEventFilter(this);
    ui->messageIn_VM->installEventFilter(this);
    ui->signatureIn_VM->installEventFilter(this);

    ui->signatureOut_SM->setFont(GUIUtil::bitcoinAddressFont());
    ui->signatureIn_VM->setFont(GUIUtil::bitcoinAddressFont());
}

SignVerifyMessageDialog::~SignVerifyMessageDialog()
{
    delete ui;
}

void SignVerifyMessageDialog::setModel(WalletModel *model)
{
    this->model = model;
}

void SignVerifyMessageDialog::setAddress_SM(QString address)
{
    ui->addressIn_SM->setText(address);
    ui->messageIn_SM->setFocus();
}

void SignVerifyMessageDialog::setAddress_VM(QString address)
{
    ui->addressIn_VM->setText(address);
    ui->messageIn_VM->setFocus();
}

void SignVerifyMessageDialog::showTab_SM(bool fShow)
{
    ui->tabWidget->setCurrentIndex(0);

    if (fShow)
        this->show();
}

void SignVerifyMessageDialog::showTab_VM(bool fShow)
{
    ui->tabWidget->setCurrentIndex(1);
    if (fShow)
        this->show();
}

void SignVerifyMessageDialog::on_addressBookButton_SM_clicked()
{
    if (model && model->getAddressTableModel())
    {
        AddressBookPage dlg(AddressBookPage::ForSending, AddressBookPage::ReceivingTab, this);
        dlg.setModel(model->getAddressTableModel());
        if (dlg.exec())
        {
            setAddress_SM(dlg.getReturnValue());
        }
    }
}

void SignVerifyMessageDialog::on_pasteButton_SM_clicked()
{
    setAddress_SM(QApplication::clipboard()->text());
}

void SignVerifyMessageDialog::on_signMessageButton_SM_clicked()
{
    /* Clear old signature to ensure users don't get confused on error with an old signature displayed */
    ui->signatureOut_SM->clear();

    CBitcoinAddress addr(ui->addressIn_SM->text().toStdString());
    if (!addr.IsValid())
    {
        ui->addressIn_SM->setValid(false);
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("The entered address is invalid.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }
    CKeyID keyID;
    if (!addr.GetKeyID(keyID))
    {
        ui->addressIn_SM->setValid(false);
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("The entered address does not refer to a key.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }

    WalletModel::UnlockContext ctx(model->requestUnlock());
    if (!ctx.isValid())
    {
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("Wallet unlock was cancelled."));
        return;
    }

    CKey key;
    if (!pwalletMain->GetKey(keyID, key))
    {
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("Private key for the entered address is not available."));
        return;
    }

    CDataStream ss(SER_GETHASH, 0);
    ss << strMessageMagic;
    ss << ui->messageIn_SM->document()->toPlainText().toStdString();

    std::vector<unsigned char> vchSig;
    if (!key.SignCompact(Hash(ss.begin(), ss.end()), vchSig))
    {
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(QString("<nobr>") + tr("Message signing failed.") + QString("</nobr>"));
        return;
    }

    ui->statusLabel_SM->setStyleSheet("QLabel { color: green; }");
    ui->statusLabel_SM->setText(QString("<nobr>") + tr("Message signed.") + QString("</nobr>"));

    ui->signatureOut_SM->setText(QString::fromStdString(EncodeBase64(&vchSig[0], vchSig.size())));
}

void SignVerifyMessageDialog::on_copySignatureButton_SM_clicked()
{
    QApplication::clipboard()->setText(ui->signatureOut_SM->text());
}

void SignVerifyMessageDialog::on_clearButton_SM_clicked()
{
    ui->addressIn_SM->clear();
    ui->messageIn_SM->clear();
    ui->signatureOut_SM->clear();
    ui->statusLabel_SM->clear();

    ui->addressIn_SM->setFocus();
}

void SignVerifyMessageDialog::on_addressBookButton_VM_clicked()
{
    if (model && model->getAddressTableModel())
    {
        AddressBookPage dlg(AddressBookPage::ForSending, AddressBookPage::SendingTab, this);
        dlg.setModel(model->getAddressTableModel());
        if (dlg.exec())
        {
            setAddress_VM(dlg.getReturnValue());
        }
    }
}

void SignVerifyMessageDialog::on_verifyMessageButton_VM_clicked()
{
    CBitcoinAddress addr(ui->addressIn_VM->text().toStdString());
    if (!addr.IsValid())
    {
        ui->addressIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The entered address is invalid.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }
    CKeyID keyID;
    if (!addr.GetKeyID(keyID))
    {
        ui->addressIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The entered address does not refer to a key.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }

    bool fInvalid = false;
    std::vector<unsigned char> vchSig = DecodeBase64(ui->signatureIn_VM->text().toStdString().c_str(), &fInvalid);

    if (fInvalid)
    {
        ui->signatureIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The signature could not be decoded.") + QString(" ") + tr("Please check the signature and try again."));
        return;
    }

    CDataStream ss(SER_GETHASH, 0);
    ss << strMessageMagic;
    ss << ui->messageIn_VM->document()->toPlainText().toStdString();

    CKey key;
    if (!key.SetCompactSignature(Hash(ss.begin(), ss.end()), vchSig))
    {
        ui->signatureIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The signature did not match the message digest.") + QString(" ") + tr("Please check the signature and try again."));
        return;
    }

    if (!(CBitcoinAddress(key.GetPubKey().GetID()) == addr))
    {
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(QString("<nobr>") + tr("Message verification failed.") + QString("</nobr>"));
        return;
    }

    ui->statusLabel_VM->setStyleSheet("QLabel { color: green; }");
    ui->statusLabel_VM->setText(QString("<nobr>") + tr("Message verified.") + QString("</nobr>"));
}

void SignVerifyMessageDialog::on_clearButton_VM_clicked()
{
    ui->addressIn_VM->clear();
    ui->signatureIn_VM->clear();
    ui->messageIn_VM->clear();
    ui->statusLabel_VM->clear();

    ui->addressIn_VM->setFocus();
}

bool SignVerifyMessageDialog::eventFilter(QObject *object, QEvent *event)
{
    if (event->type() == QEvent::MouseButtonPress || event->type() == QEvent::FocusIn)
    {
        if (ui->tabWidget->currentIndex() == 0)
        {
            /* Clear status message on focus change */
            ui->statusLabel_SM->clear();

            /* Select generated signature */
            if (object == ui->signatureOut_SM)
            {
                ui->signatureOut_SM->selectAll();
                return true;
            }
        }
        else if (ui->tabWidget->currentIndex() == 1)
        {
            /* Clear status message on focus change */
            ui->statusLabel_VM->clear();
        }
    }
    return QDialog::eventFilter(object, event);
}
<commit_msg>Update signverifymessagedialog.cpp<commit_after>#include "signverifymessagedialog.h"
#include "ui_signverifymessagedialog.h"

#include "addressbookpage.h"
#include "base58.h"
#include "guiutil.h"
#include "init.h"
#include "main.h"
#include "optionsmodel.h"
#include "walletmodel.h"
#include "wallet.h"

#include <string>
#include <vector>

#include <QClipboard>

SignVerifyMessageDialog::SignVerifyMessageDialog(QWidget *parent) :
    QDialog(parent),
    ui(new Ui::SignVerifyMessageDialog),
    model(0)
{
    ui->setupUi(this);

#if (QT_VERSION >= 0x040700)
    /* Do not move this to the XML file, Qt before 4.7 will choke on it */
    ui->addressIn_SM->setPlaceholderText(tr("Enter a YaCoin address (e.g. Y9H3tANhMroxU4rCre4bkoeDRhHKRSGwsi)"));
    ui->signatureOut_SM->setPlaceholderText(tr("Click \"Sign Message\" to generate signature"));

    ui->addressIn_VM->setPlaceholderText(tr("Enter a YaCoin address (e.g. Y9H3tANhMroxU4rCre4bkoeDRhHKRSGwsi)"));
    ui->signatureIn_VM->setPlaceholderText(tr("Enter YaCoin signature"));
#endif

    GUIUtil::setupAddressWidget(ui->addressIn_SM, this);
    GUIUtil::setupAddressWidget(ui->addressIn_VM, this);

    ui->addressIn_SM->installEventFilter(this);
    ui->messageIn_SM->installEventFilter(this);
    ui->signatureOut_SM->installEventFilter(this);
    ui->addressIn_VM->installEventFilter(this);
    ui->messageIn_VM->installEventFilter(this);
    ui->signatureIn_VM->installEventFilter(this);

    ui->signatureOut_SM->setFont(GUIUtil::bitcoinAddressFont());
    ui->signatureIn_VM->setFont(GUIUtil::bitcoinAddressFont());
}

SignVerifyMessageDialog::~SignVerifyMessageDialog()
{
    delete ui;
}

void SignVerifyMessageDialog::setModel(WalletModel *model)
{
    this->model = model;
}

void SignVerifyMessageDialog::setAddress_SM(QString address)
{
    ui->addressIn_SM->setText(address);
    ui->messageIn_SM->setFocus();
}

void SignVerifyMessageDialog::setAddress_VM(QString address)
{
    ui->addressIn_VM->setText(address);
    ui->messageIn_VM->setFocus();
}

void SignVerifyMessageDialog::showTab_SM(bool fShow)
{
    ui->tabWidget->setCurrentIndex(0);

    if (fShow)
        this->show();
}

void SignVerifyMessageDialog::showTab_VM(bool fShow)
{
    ui->tabWidget->setCurrentIndex(1);
    if (fShow)
        this->show();
}

void SignVerifyMessageDialog::on_addressBookButton_SM_clicked()
{
    if (model && model->getAddressTableModel())
    {
        AddressBookPage dlg(AddressBookPage::ForSending, AddressBookPage::ReceivingTab, this);
        dlg.setModel(model->getAddressTableModel());
        if (dlg.exec())
        {
            setAddress_SM(dlg.getReturnValue());
        }
    }
}

void SignVerifyMessageDialog::on_pasteButton_SM_clicked()
{
    setAddress_SM(QApplication::clipboard()->text());
}

void SignVerifyMessageDialog::on_signMessageButton_SM_clicked()
{
    /* Clear old signature to ensure users don't get confused on error with an old signature displayed */
    ui->signatureOut_SM->clear();

    CBitcoinAddress addr(ui->addressIn_SM->text().toStdString());
    if (!addr.IsValid())
    {
        ui->addressIn_SM->setValid(false);
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("The entered address is invalid.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }
    CKeyID keyID;
    if (!addr.GetKeyID(keyID))
    {
        ui->addressIn_SM->setValid(false);
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("The entered address does not refer to a key.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }

    WalletModel::UnlockContext ctx(model->requestUnlock());
    if (!ctx.isValid())
    {
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("Wallet unlock was cancelled."));
        return;
    }

    CKey key;
    if (!pwalletMain->GetKey(keyID, key))
    {
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(tr("Private key for the entered address is not available."));
        return;
    }

    CDataStream ss(SER_GETHASH, 0);
    ss << strMessageMagic;
    ss << ui->messageIn_SM->document()->toPlainText().toStdString();

    std::vector<unsigned char> vchSig;
    if (!key.SignCompact(Hash(ss.begin(), ss.end()), vchSig))
    {
        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_SM->setText(QString("<nobr>") + tr("Message signing failed.") + QString("</nobr>"));
        return;
    }

    ui->statusLabel_SM->setStyleSheet("QLabel { color: green; }");
    ui->statusLabel_SM->setText(QString("<nobr>") + tr("Message signed.") + QString("</nobr>"));

    ui->signatureOut_SM->setText(QString::fromStdString(EncodeBase64(&vchSig[0], vchSig.size())));
}

void SignVerifyMessageDialog::on_copySignatureButton_SM_clicked()
{
    QApplication::clipboard()->setText(ui->signatureOut_SM->text());
}

void SignVerifyMessageDialog::on_clearButton_SM_clicked()
{
    ui->addressIn_SM->clear();
    ui->messageIn_SM->clear();
    ui->signatureOut_SM->clear();
    ui->statusLabel_SM->clear();

    ui->addressIn_SM->setFocus();
}

void SignVerifyMessageDialog::on_addressBookButton_VM_clicked()
{
    if (model && model->getAddressTableModel())
    {
        AddressBookPage dlg(AddressBookPage::ForSending, AddressBookPage::SendingTab, this);
        dlg.setModel(model->getAddressTableModel());
        if (dlg.exec())
        {
            setAddress_VM(dlg.getReturnValue());
        }
    }
}

void SignVerifyMessageDialog::on_verifyMessageButton_VM_clicked()
{
    CBitcoinAddress addr(ui->addressIn_VM->text().toStdString());
    if (!addr.IsValid())
    {
        ui->addressIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The entered address is invalid.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }
    CKeyID keyID;
    if (!addr.GetKeyID(keyID))
    {
        ui->addressIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The entered address does not refer to a key.") + QString(" ") + tr("Please check the address and try again."));
        return;
    }

    bool fInvalid = false;
    std::vector<unsigned char> vchSig = DecodeBase64(ui->signatureIn_VM->text().toStdString().c_str(), &fInvalid);

    if (fInvalid)
    {
        ui->signatureIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The signature could not be decoded.") + QString(" ") + tr("Please check the signature and try again."));
        return;
    }

    CDataStream ss(SER_GETHASH, 0);
    ss << strMessageMagic;
    ss << ui->messageIn_VM->document()->toPlainText().toStdString();

    CKey key;
    if (!key.SetCompactSignature(Hash(ss.begin(), ss.end()), vchSig))
    {
        ui->signatureIn_VM->setValid(false);
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(tr("The signature did not match the message digest.") + QString(" ") + tr("Please check the signature and try again."));
        return;
    }

    if (!(CBitcoinAddress(key.GetPubKey().GetID()) == addr))
    {
        ui->statusLabel_VM->setStyleSheet("QLabel { color: red; }");
        ui->statusLabel_VM->setText(QString("<nobr>") + tr("Message verification failed.") + QString("</nobr>"));
        return;
    }

    ui->statusLabel_VM->setStyleSheet("QLabel { color: green; }");
    ui->statusLabel_VM->setText(QString("<nobr>") + tr("Message verified.") + QString("</nobr>"));
}

void SignVerifyMessageDialog::on_clearButton_VM_clicked()
{
    ui->addressIn_VM->clear();
    ui->signatureIn_VM->clear();
    ui->messageIn_VM->clear();
    ui->statusLabel_VM->clear();

    ui->addressIn_VM->setFocus();
}

bool SignVerifyMessageDialog::eventFilter(QObject *object, QEvent *event)
{
    if (event->type() == QEvent::MouseButtonPress || event->type() == QEvent::FocusIn)
    {
        if (ui->tabWidget->currentIndex() == 0)
        {
            /* Clear status message on focus change */
            ui->statusLabel_SM->clear();

            /* Select generated signature */
            if (object == ui->signatureOut_SM)
            {
                ui->signatureOut_SM->selectAll();
                return true;
            }
        }
        else if (ui->tabWidget->currentIndex() == 1)
        {
            /* Clear status message on focus change */
            ui->statusLabel_VM->clear();
        }
    }
    return QDialog::eventFilter(object, event);
}
<|endoftext|>


<commit_before>/*
   Copyright 2017 Thomas Krause <thomaskrause@posteo.de>

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

#include "precedence.h"
#include <annis/wrapper.h>                    // for ListWrapper
#include <algorithm>                          // for min, move
#include <utility>                            // for pair
#include "annis/db.h"                         // for DB
#include "annis/graphstorage/graphstorage.h"  // for ReadableGraphStorage
#include "annis/iterators.h"                  // for EdgeIterator, AnnoIt
#include "annis/operators/operator.h"         // for Operator
#include "annis/util/helper.h"                // for TokenHelper

using namespace annis;

Precedence::Precedence(const DB &db, DB::GetGSFuncT getGraphStorageFunc, unsigned int minDistance, unsigned int maxDistance)
  : tokHelper(getGraphStorageFunc, db),
    gsOrder(getGraphStorageFunc(ComponentType::ORDERING, annis_ns, "")),
    gsLeft(getGraphStorageFunc(ComponentType::LEFT_TOKEN, annis_ns, "")),
    anyTokAnno(Init::initAnnotation(db.getTokStringID(), 0, db.getNamespaceStringID())),
    anyNodeAnno(Init::initAnnotation(db.getNodeNameStringID(), 0, db.getNamespaceStringID())),
    minDistance(minDistance), maxDistance(maxDistance)
{
}

Precedence::Precedence(const DB &db, DB::GetGSFuncT getGraphStorageFunc,
                       std::string segmentation,
                       unsigned int minDistance, unsigned int maxDistance)
  : tokHelper(getGraphStorageFunc, db),
    gsOrder(getGraphStorageFunc(ComponentType::ORDERING, annis_ns, segmentation)),
    gsLeft(getGraphStorageFunc(ComponentType::LEFT_TOKEN, annis_ns, "")),
    anyTokAnno(Init::initAnnotation(db.getTokStringID(), 0, db.getNamespaceStringID())),
    anyNodeAnno(Init::initAnnotation(db.getNodeNameStringID(), 0, db.getNamespaceStringID())),
    minDistance(minDistance), maxDistance(maxDistance),
    segmentation(segmentation)
{
}

std::unique_ptr<AnnoIt> Precedence::retrieveMatches(const Match &lhs)
{
  std::unique_ptr<ListWrapper> w = std::unique_ptr<ListWrapper>(new ListWrapper());

  std::unique_ptr<EdgeIterator> edgeIterator;
  nodeid_t startNode;
  if(segmentation)
  {
    startNode = lhs.node;
  }
  else
  {
    startNode = tokHelper.rightTokenForNode(lhs.node);
  }

  edgeIterator = gsOrder->findConnected(startNode,
                                        minDistance, maxDistance);
  // materialize a list of all matches and wrap it
  for(std::pair<bool, nodeid_t> matchedToken = edgeIterator->next();
      matchedToken.first; matchedToken = edgeIterator->next())
  {
    // get all nodes that are left-aligned to this token
    for(const auto& n : gsLeft->getOutgoingEdges(matchedToken.second))
    {
      w->addMatch(Init::initMatch(anyNodeAnno, n));
    }
    // add the actual token to the list as well
    w->addMatch(Init::initMatch(anyNodeAnno, matchedToken.second));
  }

  return std::move(w);
}

bool Precedence::filter(const Match &lhs, const Match &rhs)
{
  nodeid_t startNode;
  nodeid_t endNode;
  if(segmentation)
  {
    startNode = lhs.node;
    endNode = rhs.node;
  }
  else
  {
    startNode = tokHelper.rightTokenForNode(lhs.node);
    endNode = tokHelper.leftTokenForNode(rhs.node);
  }

  if(gsOrder->isConnected(Init::initEdge(startNode, endNode),
                           minDistance, maxDistance))
  {
    return true;
  }
  return false;

}

std::string Precedence::description()
{
  if(minDistance == 1 && maxDistance == 1)
  {
    if(segmentation)
    {
      return "." + *segmentation;
    }
    else
    {
      return ".";
    }
  }
  else if(minDistance == 0 && maxDistance == 0)
  {
    if(segmentation)
    {
      return "." + *segmentation + "*";
    }
    else
    {
      return ".*";
    }
  }
  else if(minDistance == maxDistance)
  {
    if(segmentation)
    {
      return "." + *segmentation + "," + std::to_string(minDistance);
    }
    else
    {
      return "." + std::to_string(minDistance);
    }
  }
  else
  {
    if(segmentation)
    {
      return "." + *segmentation + "," + std::to_string(minDistance) + "," + std::to_string(maxDistance);
    }
    else
    {
      return "." + std::to_string(minDistance) + "," + std::to_string(maxDistance);
    }
  }
}

double Precedence::selectivity()
{
  if(gsOrder == nullptr)
  {
    return Operator::selectivity();
  }

  GraphStatistic stats = gsOrder->getStatistics();
  unsigned int maxPossibleDist = std::min(maxDistance, stats.maxDepth);
  unsigned int numOfDescendants = maxPossibleDist - minDistance + 1;
  return (double) numOfDescendants / (double) (stats.nodes/2);
}

Precedence::~Precedence()
{

}
<commit_msg>Only search for precedence if the specific graph storage component exists.<commit_after>/*
   Copyright 2017 Thomas Krause <thomaskrause@posteo.de>

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

#include "precedence.h"
#include <annis/wrapper.h>                    // for ListWrapper
#include <algorithm>                          // for min, move
#include <utility>                            // for pair
#include "annis/db.h"                         // for DB
#include "annis/graphstorage/graphstorage.h"  // for ReadableGraphStorage
#include "annis/iterators.h"                  // for EdgeIterator, AnnoIt
#include "annis/operators/operator.h"         // for Operator
#include "annis/util/helper.h"                // for TokenHelper

using namespace annis;

Precedence::Precedence(const DB &db, DB::GetGSFuncT getGraphStorageFunc, unsigned int minDistance, unsigned int maxDistance)
  : tokHelper(getGraphStorageFunc, db),
    gsOrder(getGraphStorageFunc(ComponentType::ORDERING, annis_ns, "")),
    gsLeft(getGraphStorageFunc(ComponentType::LEFT_TOKEN, annis_ns, "")),
    anyTokAnno(Init::initAnnotation(db.getTokStringID(), 0, db.getNamespaceStringID())),
    anyNodeAnno(Init::initAnnotation(db.getNodeNameStringID(), 0, db.getNamespaceStringID())),
    minDistance(minDistance), maxDistance(maxDistance)
{
}

Precedence::Precedence(const DB &db, DB::GetGSFuncT getGraphStorageFunc,
                       std::string segmentation,
                       unsigned int minDistance, unsigned int maxDistance)
  : tokHelper(getGraphStorageFunc, db),
    gsOrder(getGraphStorageFunc(ComponentType::ORDERING, annis_ns, segmentation)),
    gsLeft(getGraphStorageFunc(ComponentType::LEFT_TOKEN, annis_ns, "")),
    anyTokAnno(Init::initAnnotation(db.getTokStringID(), 0, db.getNamespaceStringID())),
    anyNodeAnno(Init::initAnnotation(db.getNodeNameStringID(), 0, db.getNamespaceStringID())),
    minDistance(minDistance), maxDistance(maxDistance),
    segmentation(segmentation)
{
}

std::unique_ptr<AnnoIt> Precedence::retrieveMatches(const Match &lhs)
{
  std::unique_ptr<ListWrapper> w = std::unique_ptr<ListWrapper>(new ListWrapper());

  if(gsOrder)
  {
    std::unique_ptr<EdgeIterator> edgeIterator;
    nodeid_t startNode;
    if(segmentation)
    {
      startNode = lhs.node;
    }
    else
    {
      startNode = tokHelper.rightTokenForNode(lhs.node);
    }

    edgeIterator = gsOrder->findConnected(startNode,
                                          minDistance, maxDistance);
    // materialize a list of all matches and wrap it
    for(std::pair<bool, nodeid_t> matchedToken = edgeIterator->next();
        matchedToken.first; matchedToken = edgeIterator->next())
    {
      // get all nodes that are left-aligned to this token
      for(const auto& n : gsLeft->getOutgoingEdges(matchedToken.second))
      {
        w->addMatch(Init::initMatch(anyNodeAnno, n));
      }
      // add the actual token to the list as well
      w->addMatch(Init::initMatch(anyNodeAnno, matchedToken.second));
    }
  }

  return std::move(w);
}

bool Precedence::filter(const Match &lhs, const Match &rhs)
{
  nodeid_t startNode;
  nodeid_t endNode;
  if(segmentation)
  {
    startNode = lhs.node;
    endNode = rhs.node;
  }
  else
  {
    startNode = tokHelper.rightTokenForNode(lhs.node);
    endNode = tokHelper.leftTokenForNode(rhs.node);
  }

  if(gsOrder->isConnected(Init::initEdge(startNode, endNode),
                           minDistance, maxDistance))
  {
    return true;
  }
  return false;

}

std::string Precedence::description()
{
  if(minDistance == 1 && maxDistance == 1)
  {
    if(segmentation)
    {
      return "." + *segmentation;
    }
    else
    {
      return ".";
    }
  }
  else if(minDistance == 0 && maxDistance == 0)
  {
    if(segmentation)
    {
      return "." + *segmentation + "*";
    }
    else
    {
      return ".*";
    }
  }
  else if(minDistance == maxDistance)
  {
    if(segmentation)
    {
      return "." + *segmentation + "," + std::to_string(minDistance);
    }
    else
    {
      return "." + std::to_string(minDistance);
    }
  }
  else
  {
    if(segmentation)
    {
      return "." + *segmentation + "," + std::to_string(minDistance) + "," + std::to_string(maxDistance);
    }
    else
    {
      return "." + std::to_string(minDistance) + "," + std::to_string(maxDistance);
    }
  }
}

double Precedence::selectivity()
{
  if(gsOrder == nullptr)
  {
    return Operator::selectivity();
  }

  GraphStatistic stats = gsOrder->getStatistics();
  unsigned int maxPossibleDist = std::min(maxDistance, stats.maxDepth);
  unsigned int numOfDescendants = maxPossibleDist - minDistance + 1;
  return (double) numOfDescendants / (double) (stats.nodes/2);
}

Precedence::~Precedence()
{

}
<|endoftext|>


<commit_before>#include "xchainer/context.h"

#include <cstdlib>
#include <future>

#include <gtest/gtest.h>
#include <nonstd/optional.hpp>

#include "xchainer/backend.h"
#include "xchainer/device.h"
#include "xchainer/native/native_backend.h"
#include "xchainer/native/native_device.h"
#include "xchainer/testing/threading.h"

namespace xchainer {
namespace {

TEST(ContextTest, Ctor) {
    Context();  // no throw
}

TEST(ContextTest, GetBackend) {
    Context ctx;
    Backend& backend = ctx.GetBackend("native");
    EXPECT_EQ(&backend, &ctx.GetBackend("native"));
}

TEST(ContextTest, NativeBackend) {
    Context ctx;
    native::NativeBackend& backend = ctx.GetNativeBackend();
    EXPECT_EQ(&ctx.GetBackend("native"), &backend);
}

TEST(ContextTest, GetBackendThreadSafe) {
    static constexpr size_t kRepeat = 1;
    static constexpr size_t kThreadCount = 2;
    std::string backend_name{"native"};

    xchainer::testing::CheckThreadSafety(
            kRepeat,
            kThreadCount,
            [](size_t /*repeat*/) { return std::make_unique<Context>(); },
            [&backend_name](size_t /*thread_index*/, const std::unique_ptr<Context>& ctx) {
                Backend& backend = ctx->GetBackend(backend_name);
                return &backend;
            },
            [&backend_name](const std::vector<Backend*>& results) {
                for (Backend* backend : results) {
                    ASSERT_EQ(backend, results.front());
                    ASSERT_EQ(backend_name, backend->GetName());
                }
            });
}

TEST(ContextTest, BackendNotFound) {
    Context ctx;
    EXPECT_THROW(ctx.GetBackend("something_that_does_not_exist"), BackendError);
}

TEST(ContextTest, GetDevice) {
    Context ctx;
    Device& device = ctx.GetDevice({"native", 0});
    EXPECT_EQ(&device, &ctx.GetDevice({"native:0"}));
}

TEST(ContextTest, GetDeviceThreadSafe) {
    static constexpr size_t kRepeat = 1;
    static constexpr int kDeviceCount = 4;
    static constexpr size_t kThreadCountPerDevice = 2;
    static constexpr size_t kThreadCount = kDeviceCount * kThreadCountPerDevice;

    xchainer::testing::CheckThreadSafety(
            kRepeat,
            kThreadCount,
            [](size_t /*repeat*/) { return std::make_unique<Context>(); },
            [](size_t thread_index, const std::unique_ptr<Context>& ctx) {
                int device_index = thread_index / kThreadCountPerDevice;
                Device& device = ctx->GetDevice({"native", device_index});
                return &device;
            },
            [](const std::vector<Device*>& results) {
                // Check device pointers are identical within each set of threads corresponding to one device
                for (int device_index = 0; device_index < kDeviceCount; ++device_index) {
                    auto it_first = std::next(results.begin(), device_index * kThreadCountPerDevice);
                    auto it_last = std::next(results.begin(), (device_index + 1) * kThreadCountPerDevice);
                    Device* ref_device = *it_first;

                    // Check the device index
                    ASSERT_EQ(device_index, ref_device->index());

                    for (auto it = it_first; it != it_last; ++it) {
                        ASSERT_EQ(ref_device, *it);
                    }
                }
            });
}

TEST(ContextTest, DefaultContext) {
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(nullptr);
    ASSERT_THROW(GetDefaultContext(), XchainerError);

    Context ctx;
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(&ctx);
    ASSERT_EQ(&ctx, &GetDefaultContext());

    Context global_ctx;
    SetGlobalDefaultContext(&global_ctx);
    SetDefaultContext(nullptr);
    ASSERT_EQ(&global_ctx, &GetDefaultContext());

    SetGlobalDefaultContext(&global_ctx);
    SetDefaultContext(&ctx);
    ASSERT_EQ(&ctx, &GetDefaultContext());
}

TEST(ContextTest, DefaultContextThreadSafe) {
    static constexpr size_t kRepeat = 1;
    static constexpr size_t kThreadCount = 2;

    xchainer::testing::CheckThreadSafety(
            kRepeat,
            kThreadCount,
            [](size_t /*repeat*/) { return nullptr; },
            [](size_t /*thread_index*/, std::nullptr_t) {
                Context ctx{};
                SetDefaultContext(&ctx);
                Context& ctx2 = GetDefaultContext();
                EXPECT_EQ(&ctx, &ctx2);
                return nullptr;
            },
            [](const std::vector<std::nullptr_t>& /*results*/) {});
}

TEST(ContextTest, GlobalDefaultContextThreadSafe) {
    static constexpr size_t kRepeat = 1;
    static constexpr size_t kThreadCount = 2;

    // Each of SetGlobalDefaultContext() and GetGlobalDefaultContext() must be thread-safe, but a pair of these calls is not guaranteed to
    // be so. In this check, a single context is set as the global context simultaneously in many threads and it only checks that
    // the succeeding Get...() call returns the same instance.
    xchainer::testing::CheckThreadSafety(
            kRepeat,
            kThreadCount,
            [](size_t /*repeat*/) { return std::make_unique<Context>(); },
            [](size_t /*thread_index*/, const std::unique_ptr<Context>& ctx) {
                SetGlobalDefaultContext(ctx.get());
                Context& ctx2 = GetGlobalDefaultContext();
                EXPECT_EQ(ctx.get(), &ctx2);
                return nullptr;
            },
            [](const std::vector<std::nullptr_t>& /*results*/) {});
}

TEST(ContextTest, GlobalDefaultContext) {
    SetGlobalDefaultContext(nullptr);
    ASSERT_THROW(GetGlobalDefaultContext(), XchainerError);

    Context ctx;
    SetGlobalDefaultContext(&ctx);
    ASSERT_EQ(&ctx, &GetGlobalDefaultContext());
}

TEST(ContextTest, ThreadLocal) {
    Context ctx;
    SetDefaultContext(&ctx);

    Context ctx2;
    auto future = std::async(std::launch::async, [&ctx2] {
        SetDefaultContext(&ctx2);
        return &GetDefaultContext();
    });
    ASSERT_NE(&GetDefaultContext(), future.get());
}

TEST(ContextTest, ContextScopeCtor) {
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(nullptr);
    Context ctx1;
    {
        // ContextScope should work even if default context is not set
        ContextScope scope(ctx1);
        EXPECT_EQ(&ctx1, &GetDefaultContext());
    }
    ASSERT_THROW(GetDefaultContext(), XchainerError);
    SetDefaultContext(&ctx1);
    {
        Context ctx2;
        ContextScope scope(ctx2);
        EXPECT_EQ(&ctx2, &GetDefaultContext());
    }
    ASSERT_EQ(&ctx1, &GetDefaultContext());
    {
        ContextScope scope;
        EXPECT_EQ(&ctx1, &GetDefaultContext());
        Context ctx2;
        SetDefaultContext(&ctx2);
    }
    ASSERT_EQ(&ctx1, &GetDefaultContext());
    Context ctx2;
    {
        ContextScope scope(ctx2);
        scope.Exit();
        EXPECT_EQ(&ctx1, &GetDefaultContext());
        SetDefaultContext(&ctx2);
        // not recovered here because the scope has already existed
    }
    ASSERT_EQ(&ctx2, &GetDefaultContext());
}

TEST(ContextTest, ContextScopeResetDevice) {
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(nullptr);
    Context ctx1;
    Context ctx2;
    {
        ContextScope ctx_scope1{ctx1};
        Device& device1 = ctx1.GetDevice({"native", 0});
        DeviceScope dev_scope1{device1};

        {
            ContextScope ctx_scope2{ctx2};
            ASSERT_NE(&device1, &GetDefaultDevice());
            Device& device2 = ctx2.GetDevice({"native", 0});
            SetDefaultDevice(&device2);
        }

        EXPECT_EQ(&device1, &GetDefaultDevice());
    }
}

TEST(ContextTest, UserDefinedBackend) {
    ::setenv("XCHAINER_PATH", XCHAINER_TEST_DIR "/context_testdata", 1);
    Context ctx;
    Backend& backend0 = ctx.GetBackend("backend0");
    EXPECT_EQ("backend0", backend0.GetName());
    Backend& backend0_2 = ctx.GetBackend("backend0");
    EXPECT_EQ(&backend0, &backend0_2);
    Backend& backend1 = ctx.GetBackend("backend1");
    EXPECT_EQ("backend1", backend1.GetName());

    Device& device0 = ctx.GetDevice(std::string("backend0:0"));
    EXPECT_EQ(&backend0, &device0.backend());
}

TEST(ContextTest, GetBackendOnDefaultContext) {
    // xchainer::GetBackend
    Context ctx;
    SetDefaultContext(&ctx);
    Backend& backend = GetBackend("native");
    EXPECT_EQ(&ctx, &backend.context());
    EXPECT_EQ("native", backend.GetName());
}

TEST(ContextTest, GetNativeBackendOnDefaultContext) {
    // xchainer::GetNativeBackend
    Context ctx;
    SetDefaultContext(&ctx);
    native::NativeBackend& backend = GetNativeBackend();
    EXPECT_EQ(&ctx.GetNativeBackend(), &backend);
}

TEST(ContextTest, GetDeviceOnDefaultContext) {
    // xchainer::GetDevice
    Context ctx;
    SetDefaultContext(&ctx);
    Device& device = GetDevice({"native:0"});
    EXPECT_EQ(&ctx, &device.backend().context());
    EXPECT_EQ("native:0", device.name());
}

}  // namespace
}  // namespace xchainer
<commit_msg>context_test.cc: remove using CheckThreadSafety<commit_after>#include "xchainer/context.h"

#include <cstdlib>
#include <future>
#include <vector>

#include <gtest/gtest.h>
#include <nonstd/optional.hpp>

#include "xchainer/backend.h"
#include "xchainer/device.h"
#include "xchainer/native/native_backend.h"
#include "xchainer/native/native_device.h"
#include "xchainer/testing/threading.h"

namespace xchainer {
namespace {

TEST(ContextTest, Ctor) {
    Context();  // no throw
}

TEST(ContextTest, GetBackend) {
    Context ctx;
    Backend& backend = ctx.GetBackend("native");
    EXPECT_EQ(&backend, &ctx.GetBackend("native"));
}

TEST(ContextTest, NativeBackend) {
    Context ctx;
    native::NativeBackend& backend = ctx.GetNativeBackend();
    EXPECT_EQ(&ctx.GetBackend("native"), &backend);
}

TEST(ContextTest, GetBackendThreadSafe) {
    static constexpr size_t kThreadCount = 2;
    Context ctx{};
    std::string backend_name{"native"};

    std::vector<Backend*> backends = testing::RunThreads(kThreadCount, [&ctx, &backend_name](size_t /*thread_index*/) {
        Backend& backend = ctx.GetBackend(backend_name);
        return &backend;
    });
    for (Backend* backend : backends) {
        ASSERT_EQ(backend, backends.front());
        ASSERT_EQ(backend_name, backend->GetName());
    }
}

TEST(ContextTest, BackendNotFound) {
    Context ctx;
    EXPECT_THROW(ctx.GetBackend("something_that_does_not_exist"), BackendError);
}

TEST(ContextTest, GetDevice) {
    Context ctx;
    Device& device = ctx.GetDevice({"native", 0});
    EXPECT_EQ(&device, &ctx.GetDevice({"native:0"}));
}

TEST(ContextTest, GetDeviceThreadSafe) {
    static constexpr int kDeviceCount = 4;
    static constexpr size_t kThreadCountPerDevice = 2;
    static constexpr size_t kThreadCount = kDeviceCount * kThreadCountPerDevice;
    Context ctx{};

    std::vector<Device*> devices = testing::RunThreads(kThreadCount, [&ctx](size_t thread_index) {
        int device_index = thread_index / kThreadCountPerDevice;
        Device& device = ctx.GetDevice({"native", device_index});
        return &device;
    });
    // Check device pointers are identical within each set of threads corresponding to one device
    for (int device_index = 0; device_index < kDeviceCount; ++device_index) {
        auto it_first = std::next(devices.begin(), device_index * kThreadCountPerDevice);
        auto it_last = std::next(devices.begin(), (device_index + 1) * kThreadCountPerDevice);
        Device* ref_device = *it_first;

        // Check the device index
        ASSERT_EQ(device_index, ref_device->index());

        for (auto it = it_first; it != it_last; ++it) {
            ASSERT_EQ(ref_device, *it);
        }
    }
}

TEST(ContextTest, DefaultContext) {
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(nullptr);
    ASSERT_THROW(GetDefaultContext(), XchainerError);

    Context ctx;
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(&ctx);
    ASSERT_EQ(&ctx, &GetDefaultContext());

    Context global_ctx;
    SetGlobalDefaultContext(&global_ctx);
    SetDefaultContext(nullptr);
    ASSERT_EQ(&global_ctx, &GetDefaultContext());

    SetGlobalDefaultContext(&global_ctx);
    SetDefaultContext(&ctx);
    ASSERT_EQ(&ctx, &GetDefaultContext());
}

TEST(ContextTest, DefaultContextThreadSafe) {
    static constexpr int kThreadCount = 2;

    testing::RunThreads(kThreadCount, [](size_t /*thread_index*/) {
        Context ctx{};
        SetDefaultContext(&ctx);
        Context& ctx2 = GetDefaultContext();
        EXPECT_EQ(&ctx, &ctx2);
        return nullptr;
    });
}

TEST(ContextTest, GlobalDefaultContextThreadSafe) {
    static constexpr int kThreadCount = 2;
    Context ctx{};

    // Each of SetGlobalDefaultContext() and GetGlobalDefaultContext() must be thread-safe, but a pair of these calls is not guaranteed
    // to be so. In this check, a single context is set as the global context simultaneously in many threads and it only checks that the
    // succeeding Get...() call returns the same instance.
    testing::RunThreads(kThreadCount, [&ctx](size_t /*thread_index*/) {
        SetGlobalDefaultContext(&ctx);
        Context& ctx2 = GetGlobalDefaultContext();
        EXPECT_EQ(&ctx, &ctx2);
        return nullptr;
    });
}

TEST(ContextTest, GlobalDefaultContext) {
    SetGlobalDefaultContext(nullptr);
    ASSERT_THROW(GetGlobalDefaultContext(), XchainerError);

    Context ctx;
    SetGlobalDefaultContext(&ctx);
    ASSERT_EQ(&ctx, &GetGlobalDefaultContext());
}

TEST(ContextTest, ThreadLocal) {
    Context ctx;
    SetDefaultContext(&ctx);

    Context ctx2;
    auto future = std::async(std::launch::async, [&ctx2] {
        SetDefaultContext(&ctx2);
        return &GetDefaultContext();
    });
    ASSERT_NE(&GetDefaultContext(), future.get());
}

TEST(ContextTest, ContextScopeCtor) {
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(nullptr);
    Context ctx1;
    {
        // ContextScope should work even if default context is not set
        ContextScope scope(ctx1);
        EXPECT_EQ(&ctx1, &GetDefaultContext());
    }
    ASSERT_THROW(GetDefaultContext(), XchainerError);
    SetDefaultContext(&ctx1);
    {
        Context ctx2;
        ContextScope scope(ctx2);
        EXPECT_EQ(&ctx2, &GetDefaultContext());
    }
    ASSERT_EQ(&ctx1, &GetDefaultContext());
    {
        ContextScope scope;
        EXPECT_EQ(&ctx1, &GetDefaultContext());
        Context ctx2;
        SetDefaultContext(&ctx2);
    }
    ASSERT_EQ(&ctx1, &GetDefaultContext());
    Context ctx2;
    {
        ContextScope scope(ctx2);
        scope.Exit();
        EXPECT_EQ(&ctx1, &GetDefaultContext());
        SetDefaultContext(&ctx2);
        // not recovered here because the scope has already existed
    }
    ASSERT_EQ(&ctx2, &GetDefaultContext());
}

TEST(ContextTest, ContextScopeResetDevice) {
    SetGlobalDefaultContext(nullptr);
    SetDefaultContext(nullptr);
    Context ctx1;
    Context ctx2;
    {
        ContextScope ctx_scope1{ctx1};
        Device& device1 = ctx1.GetDevice({"native", 0});
        DeviceScope dev_scope1{device1};

        {
            ContextScope ctx_scope2{ctx2};
            ASSERT_NE(&device1, &GetDefaultDevice());
            Device& device2 = ctx2.GetDevice({"native", 0});
            SetDefaultDevice(&device2);
        }

        EXPECT_EQ(&device1, &GetDefaultDevice());
    }
}

TEST(ContextTest, UserDefinedBackend) {
    ::setenv("XCHAINER_PATH", XCHAINER_TEST_DIR "/context_testdata", 1);
    Context ctx;
    Backend& backend0 = ctx.GetBackend("backend0");
    EXPECT_EQ("backend0", backend0.GetName());
    Backend& backend0_2 = ctx.GetBackend("backend0");
    EXPECT_EQ(&backend0, &backend0_2);
    Backend& backend1 = ctx.GetBackend("backend1");
    EXPECT_EQ("backend1", backend1.GetName());

    Device& device0 = ctx.GetDevice(std::string("backend0:0"));
    EXPECT_EQ(&backend0, &device0.backend());
}

TEST(ContextTest, GetBackendOnDefaultContext) {
    // xchainer::GetBackend
    Context ctx;
    SetDefaultContext(&ctx);
    Backend& backend = GetBackend("native");
    EXPECT_EQ(&ctx, &backend.context());
    EXPECT_EQ("native", backend.GetName());
}

TEST(ContextTest, GetNativeBackendOnDefaultContext) {
    // xchainer::GetNativeBackend
    Context ctx;
    SetDefaultContext(&ctx);
    native::NativeBackend& backend = GetNativeBackend();
    EXPECT_EQ(&ctx.GetNativeBackend(), &backend);
}

TEST(ContextTest, GetDeviceOnDefaultContext) {
    // xchainer::GetDevice
    Context ctx;
    SetDefaultContext(&ctx);
    Device& device = GetDevice({"native:0"});
    EXPECT_EQ(&ctx, &device.backend().context());
    EXPECT_EQ("native:0", device.name());
}

}  // namespace
}  // namespace xchainer
<|endoftext|>


<commit_before>#include "xchainer/python/array.h"

#include <algorithm>
#include <cstdint>

#include <pybind11/numpy.h>
#include <pybind11/operators.h>

#include "xchainer/array.h"
#include "xchainer/dtype.h"
#include "xchainer/error.h"

#include "xchainer/python/common.h"

namespace xchainer {

namespace py = pybind11;

namespace {

using ArrayBodyPtr = std::shared_ptr<internal::ArrayBody>;
using ConstArrayBodyPtr = std::shared_ptr<const internal::ArrayBody>;

Dtype NumpyDtypeToDtype(const py::dtype& npdtype) {
    switch (npdtype.kind()) {
        case 'b':
            return Dtype::kBool;
        case 'i':
            switch (npdtype.itemsize()) {
                case 1:
                    return Dtype::kInt8;
                case 2:
                    return Dtype::kInt16;
                case 4:
                    return Dtype::kInt32;
                case 8:
                    return Dtype::kInt64;
                default:
                    break;
            }
            break;
        case 'u':
            switch (npdtype.itemsize()) {
                case 1:
                    return Dtype::kUInt8;
                default:
                    break;
            }
            break;
        case 'f':
            switch (npdtype.itemsize()) {
                case 4:
                    return Dtype::kFloat32;
                case 8:
                    return Dtype::kFloat64;
                default:
                    break;
            }
            break;
        default:
            break;
    }
    throw DtypeError("unsupported NumPy dtype");
}

ArrayBodyPtr MakeArray(const Shape& shape, Dtype dtype, py::list list) {
    auto total_size = shape.total_size();
    auto bytes = GetElementSize(dtype) * total_size;
    if (static_cast<size_t>(total_size) != list.size()) {
        throw DimensionError("Invalid data length");
    }

    // Allocate a buffer and copy data
    std::shared_ptr<void> ptr = std::make_unique<uint8_t[]>(bytes);
    VisitDtype(dtype, [&](auto pt) {
        using T = typename decltype(pt)::type;
        std::transform(list.begin(), list.end(), static_cast<T*>(ptr.get()), [](auto& item) { return py::cast<T>(item); });
    });
    return Array::FromBuffer(shape, dtype, ptr).move_body();
}

ArrayBodyPtr MakeArray(py::array array) {
    if ((array.flags() & py::array::c_style) == 0) {
        throw DimensionError("cannot convert non-contiguous NumPy array to Array");
    }

    Dtype dtype = NumpyDtypeToDtype(array.dtype());
    py::buffer_info info = array.request();
    Shape shape(info.shape);

    // data holds the copy of py::array which in turn references the NumPy array and the buffer is therefore not released
    std::shared_ptr<void> data(std::make_shared<py::array>(std::move(array)), array.mutable_data());

    return Array::FromBuffer(shape, dtype, data).move_body();
}

py::buffer_info MakeNumpyArrayFromArray(internal::ArrayBody& self) {
    // Used as a temporary accessor
    Array array{std::move(ArrayBodyPtr(&self, [](internal::ArrayBody* ptr) {
        (void)ptr;  // unused
    }))};

    if (!array.is_contiguous()) {
        throw DimensionError("cannot convert non-contiguous Array to NumPy array");
    }

    int64_t itemsize{GetElementSize(array.dtype())};
    const Shape& shape = array.shape();

    // compute C-contiguous strides
    size_t ndim = array.ndim();
    std::vector<size_t> strides(ndim);
    if (ndim > 0) {
        std::partial_sum(shape.crbegin(), shape.crend() - 1, strides.rbegin() + 1, std::multiplies<size_t>());
        strides.back() = 1;
        std::transform(strides.crbegin(), strides.crend(), strides.rbegin(), [&itemsize](size_t item) { return item * itemsize; });
    }

    return py::buffer_info(array.data().get(), itemsize, std::string(1, GetCharCode(array.dtype())), ndim, shape, strides);
}

}  // namespace

void InitXchainerArray(pybind11::module& m) {
    py::class_<internal::ArrayBody, ArrayBodyPtr>{m, "Array", py::buffer_protocol()}
        .def(py::init(py::overload_cast<const Shape&, Dtype, py::list>(&MakeArray)))
        .def(py::init(py::overload_cast<py::array>(&MakeArray)))
        .def_buffer(&MakeNumpyArrayFromArray)
        .def("view",
             [](const ArrayBodyPtr& self) {
                 // Duplicate the array body
                 return std::make_shared<internal::ArrayBody>(*self);
             })
        .def("__iadd__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} += Array{rhs}).move_body(); })
        .def("__imul__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} *= Array{rhs}).move_body(); })
        .def("__add__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} + Array{rhs}).move_body(); })
        .def("__mul__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} * Array{rhs}).move_body(); })
        .def("__repr__", [](const ArrayBodyPtr& self) { return Array{self}.ToString(); })
        .def("copy", [](const ArrayBodyPtr& self) { return Array{self}.Copy().move_body(); })
        .def_property("requires_grad", [](const ArrayBodyPtr& self) { return Array{self}.IsGradRequired(); },
                      [](const ArrayBodyPtr& self, bool value) {
                          // Cannot unset required gradients
                          if (value && !self->HasNode()) {
                              Array{self}.RequireGrad();
                          }
                      })
        .def_property("grad",
                      [](const ArrayBodyPtr& self) -> ConstArrayBodyPtr {
                          if (self->HasNode()) {
                              return Array{self}.GetGrad()->body();
                          } else {
                              return nullptr;
                          }
                      },
                      [](const ArrayBodyPtr& self, const ArrayBodyPtr& grad) {
                          if (grad) {
                              if (!self->HasNode()) {
                                  Array{self}.RequireGrad().SetGrad(Array{grad});
                              } else {
                                  Array{self}.SetGrad(Array{grad});
                              }
                          } else {
                              Array{self}.ClearGrad();
                          }
                      })
        .def_property_readonly("dtype", [](const ArrayBodyPtr& self) { return Array{self}.dtype(); })
        .def_property_readonly("element_bytes", [](const ArrayBodyPtr& self) { return Array{self}.element_bytes(); })
        .def_property_readonly("is_contiguous", [](const ArrayBodyPtr& self) { return Array{self}.is_contiguous(); })
        .def_property_readonly("ndim", [](const ArrayBodyPtr& self) { return Array{self}.ndim(); })
        .def_property_readonly("offset", [](const ArrayBodyPtr& self) { return Array{self}.offset(); })
        .def_property_readonly("shape", [](const ArrayBodyPtr& self) { return Array{self}.shape(); })
        .def_property_readonly("total_bytes", [](const ArrayBodyPtr& self) { return Array{self}.total_bytes(); })
        .def_property_readonly("total_size", [](const ArrayBodyPtr& self) { return Array{self}.total_size(); })
        .def_property_readonly("_debug_data_memory_address",  // These methods starting with `_debug_` are stubs for testing
                               [](const ArrayBodyPtr& self) { return reinterpret_cast<std::uintptr_t>(Array{self}.data().get()); })
        .def_property_readonly("_debug_flat_data", [](const ArrayBodyPtr& self) {
            py::list list;
            Array array{self};

            // Copy data into the list
            VisitDtype(array.dtype(), [&array, &list](auto pt) {
                using T = typename decltype(pt)::type;
                auto size = array.total_size();
                const T& data = *std::static_pointer_cast<const T>(array.data());
                for (int64_t i = 0; i < size; ++i) {
                    list.append((&data)[i]);
                }
            });

            return list;
        });
}

}  // namespace xchainer
<commit_msg>Fix unclear comment<commit_after>#include "xchainer/python/array.h"

#include <algorithm>
#include <cstdint>

#include <pybind11/numpy.h>
#include <pybind11/operators.h>

#include "xchainer/array.h"
#include "xchainer/dtype.h"
#include "xchainer/error.h"

#include "xchainer/python/common.h"

namespace xchainer {

namespace py = pybind11;

namespace {

using ArrayBodyPtr = std::shared_ptr<internal::ArrayBody>;
using ConstArrayBodyPtr = std::shared_ptr<const internal::ArrayBody>;

Dtype NumpyDtypeToDtype(const py::dtype& npdtype) {
    switch (npdtype.kind()) {
        case 'b':
            return Dtype::kBool;
        case 'i':
            switch (npdtype.itemsize()) {
                case 1:
                    return Dtype::kInt8;
                case 2:
                    return Dtype::kInt16;
                case 4:
                    return Dtype::kInt32;
                case 8:
                    return Dtype::kInt64;
                default:
                    break;
            }
            break;
        case 'u':
            switch (npdtype.itemsize()) {
                case 1:
                    return Dtype::kUInt8;
                default:
                    break;
            }
            break;
        case 'f':
            switch (npdtype.itemsize()) {
                case 4:
                    return Dtype::kFloat32;
                case 8:
                    return Dtype::kFloat64;
                default:
                    break;
            }
            break;
        default:
            break;
    }
    throw DtypeError("unsupported NumPy dtype");
}

ArrayBodyPtr MakeArray(const Shape& shape, Dtype dtype, py::list list) {
    auto total_size = shape.total_size();
    auto bytes = GetElementSize(dtype) * total_size;
    if (static_cast<size_t>(total_size) != list.size()) {
        throw DimensionError("Invalid data length");
    }

    // Allocate a buffer and copy data
    std::shared_ptr<void> ptr = std::make_unique<uint8_t[]>(bytes);
    VisitDtype(dtype, [&](auto pt) {
        using T = typename decltype(pt)::type;
        std::transform(list.begin(), list.end(), static_cast<T*>(ptr.get()), [](auto& item) { return py::cast<T>(item); });
    });
    return Array::FromBuffer(shape, dtype, ptr).move_body();
}

ArrayBodyPtr MakeArray(py::array array) {
    if ((array.flags() & py::array::c_style) == 0) {
        throw DimensionError("cannot convert non-contiguous NumPy array to Array");
    }

    Dtype dtype = NumpyDtypeToDtype(array.dtype());
    py::buffer_info info = array.request();
    Shape shape(info.shape);

    // data holds the copy of py::array which in turn references the NumPy array and the buffer is therefore not released
    std::shared_ptr<void> data(std::make_shared<py::array>(std::move(array)), array.mutable_data());

    return Array::FromBuffer(shape, dtype, data).move_body();
}

py::buffer_info MakeNumpyArrayFromArray(internal::ArrayBody& self) {
    // Used as a temporary accessor
    Array array{std::move(ArrayBodyPtr(&self, [](internal::ArrayBody* ptr) {
        (void)ptr;  // unused
    }))};

    if (!array.is_contiguous()) {
        throw DimensionError("cannot convert non-contiguous Array to NumPy array");
    }

    int64_t itemsize{GetElementSize(array.dtype())};
    const Shape& shape = array.shape();

    // compute C-contiguous strides
    size_t ndim = array.ndim();
    std::vector<size_t> strides(ndim);
    if (ndim > 0) {
        std::partial_sum(shape.crbegin(), shape.crend() - 1, strides.rbegin() + 1, std::multiplies<size_t>());
        strides.back() = 1;
        std::transform(strides.crbegin(), strides.crend(), strides.rbegin(), [&itemsize](size_t item) { return item * itemsize; });
    }

    return py::buffer_info(array.data().get(), itemsize, std::string(1, GetCharCode(array.dtype())), ndim, shape, strides);
}

}  // namespace

void InitXchainerArray(pybind11::module& m) {
    py::class_<internal::ArrayBody, ArrayBodyPtr>{m, "Array", py::buffer_protocol()}
        .def(py::init(py::overload_cast<const Shape&, Dtype, py::list>(&MakeArray)))
        .def(py::init(py::overload_cast<py::array>(&MakeArray)))
        .def_buffer(&MakeNumpyArrayFromArray)
        .def("view",
             [](const ArrayBodyPtr& self) {
                 // Duplicate the array body
                 return std::make_shared<internal::ArrayBody>(*self);
             })
        .def("__iadd__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} += Array{rhs}).move_body(); })
        .def("__imul__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} *= Array{rhs}).move_body(); })
        .def("__add__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} + Array{rhs}).move_body(); })
        .def("__mul__", [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return (Array{self} * Array{rhs}).move_body(); })
        .def("__repr__", [](const ArrayBodyPtr& self) { return Array{self}.ToString(); })
        .def("copy", [](const ArrayBodyPtr& self) { return Array{self}.Copy().move_body(); })
        .def_property("requires_grad", [](const ArrayBodyPtr& self) { return Array{self}.IsGradRequired(); },
                      [](const ArrayBodyPtr& self, bool value) {
                          // TODO(hvy): requires_grad should not be a boolean property but a method that takes a graph id argument, aligning
                          // to the c++ interface. Currently, this property is broken in the sense that once the required_grad flag is set
                          // to true (and an ArrayNode is created internally) it cannot be uset.
                          if (value && !self->HasNode()) {
                              Array{self}.RequireGrad();
                          }
                      })
        .def_property("grad",
                      [](const ArrayBodyPtr& self) -> ConstArrayBodyPtr {
                          if (self->HasNode()) {
                              return Array{self}.GetGrad()->body();
                          } else {
                              return nullptr;
                          }
                      },
                      [](const ArrayBodyPtr& self, const ArrayBodyPtr& grad) {
                          if (grad) {
                              if (!self->HasNode()) {
                                  Array{self}.RequireGrad().SetGrad(Array{grad});
                              } else {
                                  Array{self}.SetGrad(Array{grad});
                              }
                          } else {
                              Array{self}.ClearGrad();
                          }
                      })
        .def_property_readonly("dtype", [](const ArrayBodyPtr& self) { return Array{self}.dtype(); })
        .def_property_readonly("element_bytes", [](const ArrayBodyPtr& self) { return Array{self}.element_bytes(); })
        .def_property_readonly("is_contiguous", [](const ArrayBodyPtr& self) { return Array{self}.is_contiguous(); })
        .def_property_readonly("ndim", [](const ArrayBodyPtr& self) { return Array{self}.ndim(); })
        .def_property_readonly("offset", [](const ArrayBodyPtr& self) { return Array{self}.offset(); })
        .def_property_readonly("shape", [](const ArrayBodyPtr& self) { return Array{self}.shape(); })
        .def_property_readonly("total_bytes", [](const ArrayBodyPtr& self) { return Array{self}.total_bytes(); })
        .def_property_readonly("total_size", [](const ArrayBodyPtr& self) { return Array{self}.total_size(); })
        .def_property_readonly("_debug_data_memory_address",  // These methods starting with `_debug_` are stubs for testing
                               [](const ArrayBodyPtr& self) { return reinterpret_cast<std::uintptr_t>(Array{self}.data().get()); })
        .def_property_readonly("_debug_flat_data", [](const ArrayBodyPtr& self) {
            py::list list;
            Array array{self};

            // Copy data into the list
            VisitDtype(array.dtype(), [&array, &list](auto pt) {
                using T = typename decltype(pt)::type;
                auto size = array.total_size();
                const T& data = *std::static_pointer_cast<const T>(array.data());
                for (int64_t i = 0; i < size; ++i) {
                    list.append((&data)[i]);
                }
            });

            return list;
        });
}

}  // namespace xchainer
<|endoftext|>


<commit_before>#include <common/buffer.h>
#include <common/endian.h>

#include <xcodec/xcodec.h>
#include <xcodec/xcodec_cache.h>
#include <xcodec/xcodec_decoder.h>
#include <xcodec/xcodec_hash.h>

XCodecDecoder::XCodecDecoder(XCodec *codec)
: log_("/xcodec/decoder"),
  cache_(codec->cache_),
  window_()
{ }

XCodecDecoder::~XCodecDecoder()
{ }

/*
 * Decode an XCodec-encoded stream.  Returns false if there was an
 * inconsistency in the stream or something malformed.  Returns true
 * if we were able to process the stream or expect to be able to
 * once more data arrives.  The input buffer is cleared of anything
 * we can process right now.
 */
bool
XCodecDecoder::decode(Buffer *output, Buffer *input)
{
	BufferSegment *seg, *oseg;
	uint64_t hash;
	size_t inlen;
	uint8_t ch;

	inlen = input->length();
	while (inlen != 0) {
		ch = input->peek();

		while (!XCODEC_CHAR_SPECIAL(ch)) {
			inlen--;
			input->skip(1);
			output->append(ch);

			if (inlen == 0)
				break;

			ch = input->peek();
		}

		ASSERT(inlen == input->length());

		if (inlen == 0)
			break;

		switch (ch) {
			/*
			 * A reference to a hash we are expected to have
			 * in our database.
			 */
		case XCODEC_HASHREF_CHAR:
			if (inlen < 9)
				return (true);

			input->moveout((uint8_t *)&hash, 1, sizeof hash);
			hash = LittleEndian::decode(hash);

			seg = cache_->lookup(hash);
			if (seg == NULL) {
				ERROR(log_) << "Stream referenced unseen hash: " << hash;
				return (false);
			}
			window_.declare(hash, seg);

			output->append(seg);
			inlen -= 9;
			break;

			/*
			 * A literal character will follow that would
			 * otherwise have seemed magic.
			 */
		case XCODEC_ESCAPE_CHAR:
			if (inlen < 2)
				return (true);
			input->skip(1);
			output->append(input->peek());
			input->skip(1);
			inlen -= 2;
			break;

			/*
			 * A learning opportunity -- a hash we should put
			 * into our database for future use and the data
			 * it will be used in reference to next time.
			 */
		case XCODEC_DECLARE_CHAR:
			if (inlen < 9 + XCODEC_SEGMENT_LENGTH)
				return (true);

			input->moveout((uint8_t *)&hash, 1, sizeof hash);
			hash = LittleEndian::decode(hash);

			input->copyout(&seg, XCODEC_SEGMENT_LENGTH);
			input->skip(XCODEC_SEGMENT_LENGTH);

			if (XCodecHash<XCODEC_SEGMENT_LENGTH>::hash(seg->data()) != hash) {
				seg->unref();
				ERROR(log_) << "Data in stream does not have expected hash.";
				return (false);
			}

			/*
			 * Chech whether we already know a hash by this name.
			 */
			oseg = cache_->lookup(hash);
			if (oseg != NULL) {
				/*
				 * We already know a hash by this name.  Check
				 * whether it is for the same data.  That would
				 * be nice.  In that case all we need to do is
				 * update the backreference window.
				 */
				if (seg->match(oseg)) {
					seg->unref();
					/*
					 * NB: Use the existing segment since
					 * it's nice to share.
					 */
					window_.declare(hash, oseg);
					oseg->unref();
				} else {
					seg->unref();
					oseg->unref();
					ERROR(log_) << "Stream includes hash with different local data.";
					return (false);
				}
			} else {
				/*
				 * This is a brand new hash.  Enter it into the
				 * database and the backref window.
				 */
				cache_->enter(hash, seg);
				window_.declare(hash, seg);
				seg->unref();
			}

			inlen -= 9 + XCODEC_SEGMENT_LENGTH;
			break;

			/*
			 * A reference to a recently-used hash.  We expect
			 * to see a lot of these whenever we see declarations
			 * since the declarations are desorted in the stream
			 * and are not literally inserted.
			 */
		case XCODEC_BACKREF_CHAR:
			if (inlen < 2)
				return (true);
			input->skip(1);
			ch = input->peek();
			input->skip(1);

			seg = window_.dereference(ch);
			if (seg == NULL) {
				ERROR(log_) << "Stream included invalid backref.";
				return (false);
			}

			output->append(seg);
			inlen -= 2;
			break;

		default:
			NOTREACHED();
		}
	}

	return (true);
}
<commit_msg>Drop references to BufferSegments appropriately.  Thanks, Valgrind!<commit_after>#include <common/buffer.h>
#include <common/endian.h>

#include <xcodec/xcodec.h>
#include <xcodec/xcodec_cache.h>
#include <xcodec/xcodec_decoder.h>
#include <xcodec/xcodec_hash.h>

XCodecDecoder::XCodecDecoder(XCodec *codec)
: log_("/xcodec/decoder"),
  cache_(codec->cache_),
  window_()
{ }

XCodecDecoder::~XCodecDecoder()
{ }

/*
 * Decode an XCodec-encoded stream.  Returns false if there was an
 * inconsistency in the stream or something malformed.  Returns true
 * if we were able to process the stream or expect to be able to
 * once more data arrives.  The input buffer is cleared of anything
 * we can process right now.
 */
bool
XCodecDecoder::decode(Buffer *output, Buffer *input)
{
	BufferSegment *seg, *oseg;
	uint64_t hash;
	size_t inlen;
	uint8_t ch;

	inlen = input->length();
	while (inlen != 0) {
		ch = input->peek();

		while (!XCODEC_CHAR_SPECIAL(ch)) {
			inlen--;
			input->skip(1);
			output->append(ch);

			if (inlen == 0)
				break;

			ch = input->peek();
		}

		ASSERT(inlen == input->length());

		if (inlen == 0)
			break;

		switch (ch) {
			/*
			 * A reference to a hash we are expected to have
			 * in our database.
			 */
		case XCODEC_HASHREF_CHAR:
			if (inlen < 9)
				return (true);

			input->moveout((uint8_t *)&hash, 1, sizeof hash);
			hash = LittleEndian::decode(hash);

			seg = cache_->lookup(hash);
			if (seg == NULL) {
				ERROR(log_) << "Stream referenced unseen hash: " << hash;
				return (false);
			}
			window_.declare(hash, seg);

			output->append(seg);
			seg->unref();
			inlen -= 9;
			break;

			/*
			 * A literal character will follow that would
			 * otherwise have seemed magic.
			 */
		case XCODEC_ESCAPE_CHAR:
			if (inlen < 2)
				return (true);
			input->skip(1);
			output->append(input->peek());
			input->skip(1);
			inlen -= 2;
			break;

			/*
			 * A learning opportunity -- a hash we should put
			 * into our database for future use and the data
			 * it will be used in reference to next time.
			 */
		case XCODEC_DECLARE_CHAR:
			if (inlen < 9 + XCODEC_SEGMENT_LENGTH)
				return (true);

			input->moveout((uint8_t *)&hash, 1, sizeof hash);
			hash = LittleEndian::decode(hash);

			input->copyout(&seg, XCODEC_SEGMENT_LENGTH);
			input->skip(XCODEC_SEGMENT_LENGTH);

			if (XCodecHash<XCODEC_SEGMENT_LENGTH>::hash(seg->data()) != hash) {
				seg->unref();
				ERROR(log_) << "Data in stream does not have expected hash.";
				return (false);
			}

			/*
			 * Chech whether we already know a hash by this name.
			 */
			oseg = cache_->lookup(hash);
			if (oseg != NULL) {
				/*
				 * We already know a hash by this name.  Check
				 * whether it is for the same data.  That would
				 * be nice.  In that case all we need to do is
				 * update the backreference window.
				 */
				if (seg->match(oseg)) {
					seg->unref();
					/*
					 * NB: Use the existing segment since
					 * it's nice to share.
					 */
					window_.declare(hash, oseg);
					oseg->unref();
				} else {
					seg->unref();
					oseg->unref();
					ERROR(log_) << "Stream includes hash with different local data.";
					return (false);
				}
			} else {
				/*
				 * This is a brand new hash.  Enter it into the
				 * database and the backref window.
				 */
				cache_->enter(hash, seg);
				window_.declare(hash, seg);
				seg->unref();
			}

			inlen -= 9 + XCODEC_SEGMENT_LENGTH;
			break;

			/*
			 * A reference to a recently-used hash.  We expect
			 * to see a lot of these whenever we see declarations
			 * since the declarations are desorted in the stream
			 * and are not literally inserted.
			 */
		case XCODEC_BACKREF_CHAR:
			if (inlen < 2)
				return (true);
			input->skip(1);
			ch = input->peek();
			input->skip(1);

			seg = window_.dereference(ch);
			if (seg == NULL) {
				ERROR(log_) << "Stream included invalid backref.";
				return (false);
			}

			output->append(seg);
			seg->unref();
			inlen -= 2;
			break;

		default:
			NOTREACHED();
		}
	}

	return (true);
}
<|endoftext|>


<commit_before>/*
* Prime Generation
* (C) 1999-2007,2018,2019 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/numthry.h>
#include <botan/rng.h>
#include <botan/internal/bit_ops.h>
#include <botan/loadstor.h>
#include <botan/reducer.h>
#include <botan/internal/primality.h>
#include <algorithm>

namespace Botan {

namespace {

class Prime_Sieve final
   {
   public:
      Prime_Sieve(const BigInt& init_value, size_t sieve_size) :
         m_sieve(std::min(sieve_size, PRIME_TABLE_SIZE))
         {
         for(size_t i = 0; i != m_sieve.size(); ++i)
            m_sieve[i] = static_cast<uint16_t>(init_value % PRIMES[i]);
         }

      void step(word increment)
         {
         for(size_t i = 0; i != m_sieve.size(); ++i)
            {
            m_sieve[i] = (m_sieve[i] + increment) % PRIMES[i];
            }
         }

      bool passes(bool check_2p1 = false) const
         {
         for(size_t i = 0; i != m_sieve.size(); ++i)
            {
            /*
            In this case, p is a multiple of PRIMES[i]
            */
            if(m_sieve[i] == 0)
               return false;

            if(check_2p1)
               {
               /*
               In this case, 2*p+1 will be a multiple of PRIMES[i]

               So if potentially generating a safe prime, we want to
               avoid this value because 2*p+1 will certainly not be prime.

               See "Safe Prime Generation with a Combined Sieve" M. Wiener
               https://eprint.iacr.org/2003/186.pdf
               */
               if(m_sieve[i] == (PRIMES[i] - 1) / 2)
                  return false;
               }
            }

         return true;
         }

   private:
      std::vector<uint16_t> m_sieve;
   };

}

/*
* Generate a random prime
*/
BigInt random_prime(RandomNumberGenerator& rng,
                    size_t bits, const BigInt& coprime,
                    size_t equiv, size_t modulo,
                    size_t prob)
   {
   if(coprime.is_negative() || (!coprime.is_zero() && coprime.is_even()) || coprime.bits() >= bits)
      {
      throw Invalid_Argument("random_prime: invalid coprime");
      }
   if(modulo == 0)
      {
      throw Invalid_Argument("random_prime: Invalid modulo value");
      }

   equiv %= modulo;

   if(equiv == 0)
      throw Invalid_Argument("random_prime Invalid value for equiv/modulo");

   // Handle small values:

   if(bits <= 16)
      {
      if(equiv != 1 || modulo != 2 || coprime != 0)
         throw Not_Implemented("random_prime equiv/modulo/coprime options not usable for small primes");

      if(bits <= 1)
         {
         throw Invalid_Argument("random_prime: Can't make a prime of " +
                                std::to_string(bits) + " bits");
         }
      else if(bits == 2)
         {
         return ((rng.next_byte() % 2) ? 2 : 3);
         }
      else if(bits == 3)
         {
         return ((rng.next_byte() % 2) ? 5 : 7);
         }
      else if(bits == 4)
         {
         return ((rng.next_byte() % 2) ? 11 : 13);
         }
      else
         {
         for(;;)
            {
            // This is slightly biased, but for small primes it does not seem to matter
            const uint8_t b0 = rng.next_byte();
            const uint8_t b1 = rng.next_byte();
            const size_t idx = make_uint16(b0, b1) % PRIME_TABLE_SIZE;
            const uint16_t small_prime = PRIMES[idx];

            if(high_bit(small_prime) == bits)
               return small_prime;
            }
         }
      }

   const size_t MAX_ATTEMPTS = 32*1024;

   while(true)
      {
      BigInt p(rng, bits);

      // Force lowest and two top bits on
      p.set_bit(bits - 1);
      p.set_bit(bits - 2);
      p.set_bit(0);

      // Force p to be equal to equiv mod modulo
      p += (modulo - (p % modulo)) + equiv;

      Prime_Sieve sieve(p, bits);

      for(size_t attempt = 0; attempt <= MAX_ATTEMPTS; ++attempt)
         {
         p += modulo;

         sieve.step(modulo);

         // p can be even if modulo is odd, continue on in that case
         if(p.is_even() || sieve.passes(true) == false)
            continue;

         Modular_Reducer mod_p(p);

         /*
         First do a single M-R iteration to quickly elimate most non-primes,
         before doing the coprimality check which is expensive
         */
         if(is_miller_rabin_probable_prime(p, mod_p, rng, 1) == false)
            continue;

         if(coprime > 1)
            {
            /*
            * Check if gcd(p - 1, coprime) != 1 by attempting to compute a
            * modular inverse. We are assured coprime is odd (since if it was
            * even, it would always have a common factor with p - 1), so
            * take off the factors of 2 from (p-1) then compute the inverse
            * of coprime modulo that integer.
            */
            BigInt p1 = p - 1;
            p1 >>= low_zero_bits(p1);
            if(ct_inverse_mod_odd_modulus(coprime, p1).is_zero())
               {
               BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) > 1);
               continue;
               }

            BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) == 1);
            }

         if(p.bits() > bits)
            break;

         const size_t t = miller_rabin_test_iterations(bits, prob, true);

         if(is_miller_rabin_probable_prime(p, mod_p, rng, t) == false)
            continue;

         if(prob > 32 && !is_lucas_probable_prime(p, mod_p))
            continue;

         return p;
         }
      }
   }

BigInt generate_rsa_prime(RandomNumberGenerator& keygen_rng,
                          RandomNumberGenerator& prime_test_rng,
                          size_t bits,
                          const BigInt& coprime,
                          size_t prob)
   {
   if(bits < 512)
      throw Invalid_Argument("generate_rsa_prime bits too small");

   /*
   * The restriction on coprime <= 64 bits is arbitrary but generally speaking
   * very large RSA public exponents are a bad idea both for performance and due
   * to attacks on small d.
   */
   if(coprime <= 1 || coprime.is_even() || coprime.bits() > 64)
      throw Invalid_Argument("generate_rsa_prime coprime must be small odd positive integer");

   const size_t MAX_ATTEMPTS = 32*1024;

   while(true)
      {
      BigInt p(keygen_rng, bits);

      // Force lowest and two top bits on
      p.set_bit(bits - 1);
      p.set_bit(bits - 2);
      p.set_bit(0);

      Prime_Sieve sieve(p, bits);

      const word step = 2;

      for(size_t attempt = 0; attempt <= MAX_ATTEMPTS; ++attempt)
         {
         p += step;

         sieve.step(step);

         if(sieve.passes() == false)
            continue;

         Modular_Reducer mod_p(p);

         /*
         * Do a single primality test first before checking coprimality, since
         * currently a single Miller-Rabin test is faster than computing modular
         * inverse, and this eliminates almost all wasted modular inverses.
         */
         if(is_miller_rabin_probable_prime(p, mod_p, prime_test_rng, 1) == false)
            continue;

         /*
         * Check if p - 1 and coprime are relatively prime by computing the inverse.
         *
         * We avoid gcd here because that algorithm is not constant time.
         * Modular inverse is (for odd modulus anyway).
         *
         * We earlier verified that coprime argument is odd. Thus the factors of 2
         * in (p - 1) cannot possibly be factors in coprime, so remove them from p - 1.
         * Using an odd modulus allows the const time algorithm to be used.
         *
         * This assumes coprime < p - 1 which is always true for RSA.
         */
         BigInt p1 = p - 1;
         p1 >>= low_zero_bits(p1);
         if(ct_inverse_mod_odd_modulus(coprime, p1).is_zero())
            {
            BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) > 1);
            continue;
            }

         BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) == 1);

         if(p.bits() > bits)
            break;

         const size_t t = miller_rabin_test_iterations(bits, prob, true);

         if(is_miller_rabin_probable_prime(p, mod_p, prime_test_rng, t) == true)
            return p;
         }
      }
   }

/*
* Generate a random safe prime
*/
BigInt random_safe_prime(RandomNumberGenerator& rng, size_t bits)
   {
   if(bits <= 64)
      throw Invalid_Argument("random_safe_prime: Can't make a prime of " +
                             std::to_string(bits) + " bits");

   BigInt q, p;
   for(;;)
      {
      /*
      Generate q == 2 (mod 3)

      Otherwise [q == 1 (mod 3) case], 2*q+1 == 3 (mod 3) and not prime.

      Initially allow a very high error prob (1/2**8) to allow for fast checks,
      then if 2*q+1 turns out to be a prime go back and strongly check q.
      */
      q = random_prime(rng, bits - 1, 0, 2, 3, 8);
      p = (q << 1) + 1;

      if(is_prime(p, rng, 128, true))
         {
         // We did only a weak check before, go back and verify q before returning
         if(is_prime(q, rng, 128, true))
            return p;
         }
      }
   }

}
<commit_msg>Fix error message<commit_after>/*
* Prime Generation
* (C) 1999-2007,2018,2019 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/numthry.h>
#include <botan/rng.h>
#include <botan/internal/bit_ops.h>
#include <botan/loadstor.h>
#include <botan/reducer.h>
#include <botan/internal/primality.h>
#include <algorithm>

namespace Botan {

namespace {

class Prime_Sieve final
   {
   public:
      Prime_Sieve(const BigInt& init_value, size_t sieve_size) :
         m_sieve(std::min(sieve_size, PRIME_TABLE_SIZE))
         {
         for(size_t i = 0; i != m_sieve.size(); ++i)
            m_sieve[i] = static_cast<uint16_t>(init_value % PRIMES[i]);
         }

      void step(word increment)
         {
         for(size_t i = 0; i != m_sieve.size(); ++i)
            {
            m_sieve[i] = (m_sieve[i] + increment) % PRIMES[i];
            }
         }

      bool passes(bool check_2p1 = false) const
         {
         for(size_t i = 0; i != m_sieve.size(); ++i)
            {
            /*
            In this case, p is a multiple of PRIMES[i]
            */
            if(m_sieve[i] == 0)
               return false;

            if(check_2p1)
               {
               /*
               In this case, 2*p+1 will be a multiple of PRIMES[i]

               So if potentially generating a safe prime, we want to
               avoid this value because 2*p+1 will certainly not be prime.

               See "Safe Prime Generation with a Combined Sieve" M. Wiener
               https://eprint.iacr.org/2003/186.pdf
               */
               if(m_sieve[i] == (PRIMES[i] - 1) / 2)
                  return false;
               }
            }

         return true;
         }

   private:
      std::vector<uint16_t> m_sieve;
   };

}

/*
* Generate a random prime
*/
BigInt random_prime(RandomNumberGenerator& rng,
                    size_t bits, const BigInt& coprime,
                    size_t equiv, size_t modulo,
                    size_t prob)
   {
   if(bits <= 1)
      {
      throw Invalid_Argument("random_prime: Can't make a prime of " +
                             std::to_string(bits) + " bits");
      }
   if(coprime.is_negative() || (!coprime.is_zero() && coprime.is_even()) || coprime.bits() >= bits)
      {
      throw Invalid_Argument("random_prime: invalid coprime");
      }
   if(modulo == 0)
      {
      throw Invalid_Argument("random_prime: Invalid modulo value");
      }

   equiv %= modulo;

   if(equiv == 0)
      throw Invalid_Argument("random_prime Invalid value for equiv/modulo");

   // Handle small values:

   if(bits <= 16)
      {
      if(equiv != 1 || modulo != 2 || coprime != 0)
         throw Not_Implemented("random_prime equiv/modulo/coprime options not usable for small primes");

      if(bits == 2)
         {
         return ((rng.next_byte() % 2) ? 2 : 3);
         }
      else if(bits == 3)
         {
         return ((rng.next_byte() % 2) ? 5 : 7);
         }
      else if(bits == 4)
         {
         return ((rng.next_byte() % 2) ? 11 : 13);
         }
      else
         {
         for(;;)
            {
            // This is slightly biased, but for small primes it does not seem to matter
            const uint8_t b0 = rng.next_byte();
            const uint8_t b1 = rng.next_byte();
            const size_t idx = make_uint16(b0, b1) % PRIME_TABLE_SIZE;
            const uint16_t small_prime = PRIMES[idx];

            if(high_bit(small_prime) == bits)
               return small_prime;
            }
         }
      }

   const size_t MAX_ATTEMPTS = 32*1024;

   while(true)
      {
      BigInt p(rng, bits);

      // Force lowest and two top bits on
      p.set_bit(bits - 1);
      p.set_bit(bits - 2);
      p.set_bit(0);

      // Force p to be equal to equiv mod modulo
      p += (modulo - (p % modulo)) + equiv;

      Prime_Sieve sieve(p, bits);

      for(size_t attempt = 0; attempt <= MAX_ATTEMPTS; ++attempt)
         {
         p += modulo;

         sieve.step(modulo);

         // p can be even if modulo is odd, continue on in that case
         if(p.is_even() || sieve.passes(true) == false)
            continue;

         Modular_Reducer mod_p(p);

         /*
         First do a single M-R iteration to quickly elimate most non-primes,
         before doing the coprimality check which is expensive
         */
         if(is_miller_rabin_probable_prime(p, mod_p, rng, 1) == false)
            continue;

         if(coprime > 1)
            {
            /*
            * Check if gcd(p - 1, coprime) != 1 by attempting to compute a
            * modular inverse. We are assured coprime is odd (since if it was
            * even, it would always have a common factor with p - 1), so
            * take off the factors of 2 from (p-1) then compute the inverse
            * of coprime modulo that integer.
            */
            BigInt p1 = p - 1;
            p1 >>= low_zero_bits(p1);
            if(ct_inverse_mod_odd_modulus(coprime, p1).is_zero())
               {
               BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) > 1);
               continue;
               }

            BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) == 1);
            }

         if(p.bits() > bits)
            break;

         const size_t t = miller_rabin_test_iterations(bits, prob, true);

         if(is_miller_rabin_probable_prime(p, mod_p, rng, t) == false)
            continue;

         if(prob > 32 && !is_lucas_probable_prime(p, mod_p))
            continue;

         return p;
         }
      }
   }

BigInt generate_rsa_prime(RandomNumberGenerator& keygen_rng,
                          RandomNumberGenerator& prime_test_rng,
                          size_t bits,
                          const BigInt& coprime,
                          size_t prob)
   {
   if(bits < 512)
      throw Invalid_Argument("generate_rsa_prime bits too small");

   /*
   * The restriction on coprime <= 64 bits is arbitrary but generally speaking
   * very large RSA public exponents are a bad idea both for performance and due
   * to attacks on small d.
   */
   if(coprime <= 1 || coprime.is_even() || coprime.bits() > 64)
      throw Invalid_Argument("generate_rsa_prime coprime must be small odd positive integer");

   const size_t MAX_ATTEMPTS = 32*1024;

   while(true)
      {
      BigInt p(keygen_rng, bits);

      // Force lowest and two top bits on
      p.set_bit(bits - 1);
      p.set_bit(bits - 2);
      p.set_bit(0);

      Prime_Sieve sieve(p, bits);

      const word step = 2;

      for(size_t attempt = 0; attempt <= MAX_ATTEMPTS; ++attempt)
         {
         p += step;

         sieve.step(step);

         if(sieve.passes() == false)
            continue;

         Modular_Reducer mod_p(p);

         /*
         * Do a single primality test first before checking coprimality, since
         * currently a single Miller-Rabin test is faster than computing modular
         * inverse, and this eliminates almost all wasted modular inverses.
         */
         if(is_miller_rabin_probable_prime(p, mod_p, prime_test_rng, 1) == false)
            continue;

         /*
         * Check if p - 1 and coprime are relatively prime by computing the inverse.
         *
         * We avoid gcd here because that algorithm is not constant time.
         * Modular inverse is (for odd modulus anyway).
         *
         * We earlier verified that coprime argument is odd. Thus the factors of 2
         * in (p - 1) cannot possibly be factors in coprime, so remove them from p - 1.
         * Using an odd modulus allows the const time algorithm to be used.
         *
         * This assumes coprime < p - 1 which is always true for RSA.
         */
         BigInt p1 = p - 1;
         p1 >>= low_zero_bits(p1);
         if(ct_inverse_mod_odd_modulus(coprime, p1).is_zero())
            {
            BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) > 1);
            continue;
            }

         BOTAN_DEBUG_ASSERT(gcd(p - 1, coprime) == 1);

         if(p.bits() > bits)
            break;

         const size_t t = miller_rabin_test_iterations(bits, prob, true);

         if(is_miller_rabin_probable_prime(p, mod_p, prime_test_rng, t) == true)
            return p;
         }
      }
   }

/*
* Generate a random safe prime
*/
BigInt random_safe_prime(RandomNumberGenerator& rng, size_t bits)
   {
   if(bits <= 64)
      throw Invalid_Argument("random_safe_prime: Can't make a prime of " +
                             std::to_string(bits) + " bits");

   BigInt q, p;
   for(;;)
      {
      /*
      Generate q == 2 (mod 3)

      Otherwise [q == 1 (mod 3) case], 2*q+1 == 3 (mod 3) and not prime.

      Initially allow a very high error prob (1/2**8) to allow for fast checks,
      then if 2*q+1 turns out to be a prime go back and strongly check q.
      */
      q = random_prime(rng, bits - 1, 0, 2, 3, 8);
      p = (q << 1) + 1;

      if(is_prime(p, rng, 128, true))
         {
         // We did only a weak check before, go back and verify q before returning
         if(is_prime(q, rng, 128, true))
            return p;
         }
      }
   }

}
<|endoftext|>