<commit_before>/* This file is part of Zanshin Todo.

   Copyright 2008 Kevin Ottens <ervin@kde.org>

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301,
   USA.
*/

#include "modeltestbase.h"
#include <qtest_kde.h>

#include <akonadi/collection.h>
#include <akonadi/item.h>
#include <akonadi/itemdeletejob.h>

#include <QtGui/QSortFilterProxyModel>

#include "todoflatmodel.h"
#include "todocategoriesmodel.h"

class TodoCategoriesModelTest : public ModelTestBase
{
    Q_OBJECT

private slots:
    void initTestCase();
    void testInitialState();
    void testSingleModification();
    void testReparentModification();
    void testSeveralReparentModification();
    void testSingleRemoved();
    void testMultipleRemoved();

private:
    TodoFlatModel m_flatModel;
    QSortFilterProxyModel m_flatSortedModel;
    TodoCategoriesModel m_model;
    QSortFilterProxyModel m_sortedModel;
};

QTEST_KDEMAIN(TodoCategoriesModelTest, GUI)

void TodoCategoriesModelTest::initTestCase()
{
    ModelTestBase::initTestCase();
    m_model.setSourceModel(&m_flatModel);
    m_model.setCollection(m_collection);
    flushNotifications();

    m_flatSortedModel.setSourceModel(&m_flatModel);
    m_flatSortedModel.sort(TodoFlatModel::RemoteId);

    m_sortedModel.setSourceModel(&m_model);
    m_sortedModel.sort(TodoFlatModel::Summary);
}

class TreeNode
{
public:
    TreeNode(const QString &i, const QString &s)
        : id(i), summary(s) { }

    TreeNode &operator<<(const TreeNode &child)
    {
        const_cast<TreeNode&>(child).parent = summary;
        children << child;
        return *this;
    }

    QString id;
    QString parent;
    QString summary;
    QList<TreeNode> children;
};

#if 0
static void dumpTree(const QList<TreeNode> &tree, int indent = 0)
{
    QString prefix;
    for (int i=0; i<indent; ++i) {
        prefix+="    ";
    }

    foreach (const TreeNode &node, tree) {
        qDebug() << prefix << "[" << node.parent << node.id << node.summary;
        dumpTree(node.children, indent+1);
        qDebug() << prefix << "]";
    }
}
#endif

static void compareTrees(QAbstractItemModel *model, const QList<TreeNode> &tree,
                         const QModelIndex &root = QModelIndex())
{
    int row = 0;
    foreach (const TreeNode &node, tree) {
        QCOMPARE(model->data(model->index(row, TodoFlatModel::RemoteId, root)).toString(), node.id);
        QCOMPARE(model->data(model->index(row, TodoFlatModel::Summary, root)).toString(), node.summary);
        if (node.parent.isEmpty()) {
            QVERIFY(model->data(model->index(row, TodoFlatModel::Categories, root)).toStringList().isEmpty());
        } else {
            QVERIFY(model->data(model->index(row, TodoFlatModel::Categories, root)).toStringList().contains(node.parent));
        }

        QCOMPARE(model->rowCount(model->index(row, 0, root)), node.children.size());
        compareTrees(model, node.children, model->index(row, 0, root));
        row++;
    }
}

void TodoCategoriesModelTest::testInitialState()
{
    QList<TreeNode> tree;

    tree << (TreeNode(QString(), "Computer")
             << TreeNode(QString(), "Online")
             << TreeNode("fake-06", "Read magazine")
            )

         << (TreeNode(QString(), "Errands")
             << TreeNode("fake-03", "Walk around with the dog")
            )

         << (TreeNode(QString(), "Home")
             << TreeNode("fake-05", "Feed the dog")
            )

         << (TreeNode(QString(), "Office")
             << TreeNode("fake-06", "Read magazine")
            )

         << TreeNode(QString(), "Phone");

    compareTrees(&m_sortedModel, tree);
}

void TodoCategoriesModelTest::testSingleModification()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(5, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Summary);

    QSignalSpy spy(&m_model, SIGNAL(dataChanged(QModelIndex, QModelIndex)));

    QVERIFY(m_model.setData(indexes.first(), "Feed the cat"));

    flushNotifications();

    foreach (const QModelIndex &index, indexes) {
        QCOMPARE(m_model.data(index).toString(), QString("Feed the cat"));
    }

    indexes = m_model.indexesForItem(item, 0);
    QCOMPARE(spy.count(), indexes.size());

    while (!spy.isEmpty()) {
        QVariantList signal = spy.takeFirst();

        QCOMPARE(signal.count(), 2);
        QModelIndex begin = signal.at(0).value<QModelIndex>();
        QModelIndex end = signal.at(1).value<QModelIndex>();

        QVERIFY(begin.row()==end.row());
        QVERIFY(begin.parent()==end.parent());
        QCOMPARE(begin.column(), 0);
        QCOMPARE(end.column(), (int)TodoFlatModel::LastColumn);

        QVERIFY(indexes.contains(begin));
    }
}

void TodoCategoriesModelTest::testReparentModification()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(5, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);

    QModelIndexList oldParents;
    QList<int> oldParentsRow;
    foreach (const QModelIndex &index, indexes) {
        oldParents << index.parent();
        oldParentsRow << index.row();
    }

    QModelIndexList newParents;
    QList<int> newParentsRow;
    newParents << m_model.indexForCategory("Errands")
               << m_model.indexForCategory("Home");
    foreach (const QModelIndex &newParent, newParents) {
        newParentsRow << m_model.rowCount(newParent);
    }

    QSignalSpy rowsInserted(&m_model, SIGNAL(rowsInserted(QModelIndex, int, int)));
    QSignalSpy rowsRemoved(&m_model, SIGNAL(rowsRemoved(QModelIndex, int, int)));

    // Note: It includes a category "Foo" which wasn't in the model previously,
    // we should get an insert event for it too!
    QVERIFY(m_model.setData(indexes.first(), "Home, Errands, Foo"));

    flushNotifications();

    QCOMPARE(rowsRemoved.count(), 2);
    while (!rowsRemoved.isEmpty()) {
        QVariantList signal = rowsRemoved.takeFirst();
        QCOMPARE(signal.count(), 3);

        QModelIndex parent = signal.at(0).value<QModelIndex>();
        QVERIFY(oldParents.contains(parent));

        int expectedRow = oldParentsRow[oldParents.indexOf(parent)];
        QCOMPARE(signal.at(1).toInt(), expectedRow);
        QCOMPARE(signal.at(2).toInt(), expectedRow);
    }

    QCOMPARE(rowsInserted.count(), 4);

    // first insert event should be for "Foo":
    {
        QVariantList signal = rowsInserted.takeFirst();
        QCOMPARE(signal.count(), 3);

        QModelIndex parent = signal.at(0).value<QModelIndex>();
        QVERIFY(!parent.isValid());

        int expectedRow = m_model.rowCount()-1;
        QCOMPARE(signal.at(1).toInt(), expectedRow);
        QCOMPARE(signal.at(2).toInt(), expectedRow);
    }

    // Adjust newParents and newParentsRow to include the new "Foo" category
    newParents << m_model.index(m_model.rowCount()-1, 0);
    newParentsRow << 0;

    // now check the other events for the item itself
    while (!rowsInserted.isEmpty()) {
        QVariantList signal = rowsInserted.takeFirst();
        QCOMPARE(signal.count(), 3);

        QModelIndex parent = signal.at(0).value<QModelIndex>();
        QVERIFY(newParents.contains(parent));

        int expectedRow = newParentsRow[newParents.indexOf(parent)];
        QCOMPARE(signal.at(1).toInt(), expectedRow);
        QCOMPARE(signal.at(2).toInt(), expectedRow);
    }
}

void TodoCategoriesModelTest::testSeveralReparentModification()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(4, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);
    QCOMPARE(indexes.size(), 1);
    QModelIndex index = indexes.takeFirst();

    QModelIndex oldParent = index.parent();
    QCOMPARE(oldParent, m_model.indexForCategory("Home"));

    QModelIndex newParent = m_model.indexForCategory("Errands");

    QSignalSpy rowsInserted(&m_model, SIGNAL(rowsInserted(QModelIndex, int, int)));
    QSignalSpy rowsRemoved(&m_model, SIGNAL(rowsRemoved(QModelIndex, int, int)));

    QVERIFY(m_model.setData(index, "Errands"));

    flushNotifications();

    QCOMPARE(rowsRemoved.count(), 1);
    QVariantList signal = rowsRemoved.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), oldParent);
    QCOMPARE(signal.at(1).toInt(), 0);
    QCOMPARE(signal.at(2).toInt(), 0);

    QCOMPARE(rowsInserted.count(), 1);
    signal = rowsInserted.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), newParent);
    QCOMPARE(signal.at(1).toInt(), 2);
    QCOMPARE(signal.at(2).toInt(), 2);

    indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);
    QCOMPARE(indexes.size(), 1);
    index = indexes.takeFirst();

    QVERIFY(m_model.setData(index, "Home"));

    flushNotifications();

    QCOMPARE(rowsRemoved.count(), 1);
    signal = rowsRemoved.takeFirst();
    QCOMPARE(signal.count(), 3);
    qDebug() << signal.at(0).value<QModelIndex>() << newParent;
    QCOMPARE(signal.at(0).value<QModelIndex>(), newParent);
    QCOMPARE(signal.at(1).toInt(), 2);
    QCOMPARE(signal.at(2).toInt(), 2);

    QCOMPARE(rowsInserted.count(), 1);
    signal = rowsInserted.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), oldParent);
    QCOMPARE(signal.at(1).toInt(), 1);
    QCOMPARE(signal.at(2).toInt(), 1);
}

void TodoCategoriesModelTest::testSingleRemoved()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(4, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);

    QCOMPARE(indexes.size(), 1);
    QModelIndex index = indexes.takeFirst();
    QModelIndex parent = index.parent();
    int count = m_model.rowCount(parent);

    QSignalSpy spy(&m_model, SIGNAL(rowsAboutToBeRemoved(QModelIndex, int, int)));

    Akonadi::ItemDeleteJob *job = new Akonadi::ItemDeleteJob(item);
    QVERIFY(job->exec());

    flushNotifications();

    QCOMPARE(m_model.rowCount(parent), count - 1);

    QCOMPARE(spy.count(), 1);
    QVariantList signal = spy.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), parent);
    QCOMPARE(signal.at(1).toInt(), 1);
    QCOMPARE(signal.at(1).toInt(), 1);
}

void TodoCategoriesModelTest::testMultipleRemoved()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(5, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);

    QCOMPARE(indexes.size(), 1);
    QModelIndex index = indexes.takeFirst();
    QModelIndex parent = index.parent();
    int count = m_model.rowCount(parent);

    QSignalSpy spy(&m_model, SIGNAL(rowsAboutToBeRemoved(QModelIndex, int, int)));

    Akonadi::ItemDeleteJob *job = new Akonadi::ItemDeleteJob(item);
    QVERIFY(job->exec());

    flushNotifications();

    QCOMPARE(m_model.rowCount(parent), count - 1);

    QCOMPARE(spy.count(), 1);
    QVariantList signal = spy.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), parent);
    QCOMPARE(signal.at(1).toInt(), 1);
    QCOMPARE(signal.at(1).toInt(), 1);
}

#include "todocategoriesmodeltest.moc"
<commit_msg>add testMultipleRemoved for todoCategoriesModelTest<commit_after>/* This file is part of Zanshin Todo.

   Copyright 2008 Kevin Ottens <ervin@kde.org>

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301,
   USA.
*/

#include "modeltestbase.h"
#include <qtest_kde.h>

#include <akonadi/collection.h>
#include <akonadi/item.h>
#include <akonadi/itemdeletejob.h>

#include <QtGui/QSortFilterProxyModel>

#include "todoflatmodel.h"
#include "todocategoriesmodel.h"

class TodoCategoriesModelTest : public ModelTestBase
{
    Q_OBJECT

private slots:
    void initTestCase();
    void testInitialState();
    void testSingleModification();
    void testReparentModification();
    void testSeveralReparentModification();
    void testSingleRemoved();
    void testMultipleRemoved();

private:
    TodoFlatModel m_flatModel;
    QSortFilterProxyModel m_flatSortedModel;
    TodoCategoriesModel m_model;
    QSortFilterProxyModel m_sortedModel;
};

QTEST_KDEMAIN(TodoCategoriesModelTest, GUI)

void TodoCategoriesModelTest::initTestCase()
{
    ModelTestBase::initTestCase();
    m_model.setSourceModel(&m_flatModel);
    m_model.setCollection(m_collection);
    flushNotifications();

    m_flatSortedModel.setSourceModel(&m_flatModel);
    m_flatSortedModel.sort(TodoFlatModel::RemoteId);

    m_sortedModel.setSourceModel(&m_model);
    m_sortedModel.sort(TodoFlatModel::Summary);
}

class TreeNode
{
public:
    TreeNode(const QString &i, const QString &s)
        : id(i), summary(s) { }

    TreeNode &operator<<(const TreeNode &child)
    {
        const_cast<TreeNode&>(child).parent = summary;
        children << child;
        return *this;
    }

    QString id;
    QString parent;
    QString summary;
    QList<TreeNode> children;
};

#if 0
static void dumpTree(const QList<TreeNode> &tree, int indent = 0)
{
    QString prefix;
    for (int i=0; i<indent; ++i) {
        prefix+="    ";
    }

    foreach (const TreeNode &node, tree) {
        qDebug() << prefix << "[" << node.parent << node.id << node.summary;
        dumpTree(node.children, indent+1);
        qDebug() << prefix << "]";
    }
}
#endif

static void compareTrees(QAbstractItemModel *model, const QList<TreeNode> &tree,
                         const QModelIndex &root = QModelIndex())
{
    int row = 0;
    foreach (const TreeNode &node, tree) {
        QCOMPARE(model->data(model->index(row, TodoFlatModel::RemoteId, root)).toString(), node.id);
        QCOMPARE(model->data(model->index(row, TodoFlatModel::Summary, root)).toString(), node.summary);
        if (node.parent.isEmpty()) {
            QVERIFY(model->data(model->index(row, TodoFlatModel::Categories, root)).toStringList().isEmpty());
        } else {
            QVERIFY(model->data(model->index(row, TodoFlatModel::Categories, root)).toStringList().contains(node.parent));
        }

        QCOMPARE(model->rowCount(model->index(row, 0, root)), node.children.size());
        compareTrees(model, node.children, model->index(row, 0, root));
        row++;
    }
}

void TodoCategoriesModelTest::testInitialState()
{
    QList<TreeNode> tree;

    tree << (TreeNode(QString(), "Computer")
             << TreeNode(QString(), "Online")
             << TreeNode("fake-06", "Read magazine")
            )

         << (TreeNode(QString(), "Errands")
             << TreeNode("fake-03", "Walk around with the dog")
            )

         << (TreeNode(QString(), "Home")
             << TreeNode("fake-05", "Feed the dog")
            )

         << (TreeNode(QString(), "Office")
             << TreeNode("fake-06", "Read magazine")
            )

         << TreeNode(QString(), "Phone");

    compareTrees(&m_sortedModel, tree);
}

void TodoCategoriesModelTest::testSingleModification()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(5, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Summary);

    QSignalSpy spy(&m_model, SIGNAL(dataChanged(QModelIndex, QModelIndex)));

    QVERIFY(m_model.setData(indexes.first(), "Feed the cat"));

    flushNotifications();

    foreach (const QModelIndex &index, indexes) {
        QCOMPARE(m_model.data(index).toString(), QString("Feed the cat"));
    }

    indexes = m_model.indexesForItem(item, 0);
    QCOMPARE(spy.count(), indexes.size());

    while (!spy.isEmpty()) {
        QVariantList signal = spy.takeFirst();

        QCOMPARE(signal.count(), 2);
        QModelIndex begin = signal.at(0).value<QModelIndex>();
        QModelIndex end = signal.at(1).value<QModelIndex>();

        QVERIFY(begin.row()==end.row());
        QVERIFY(begin.parent()==end.parent());
        QCOMPARE(begin.column(), 0);
        QCOMPARE(end.column(), (int)TodoFlatModel::LastColumn);

        QVERIFY(indexes.contains(begin));
    }
}

void TodoCategoriesModelTest::testReparentModification()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(5, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);

    QModelIndexList oldParents;
    QList<int> oldParentsRow;
    foreach (const QModelIndex &index, indexes) {
        oldParents << index.parent();
        oldParentsRow << index.row();
    }

    QModelIndexList newParents;
    QList<int> newParentsRow;
    newParents << m_model.indexForCategory("Errands")
               << m_model.indexForCategory("Home");
    foreach (const QModelIndex &newParent, newParents) {
        newParentsRow << m_model.rowCount(newParent);
    }

    QSignalSpy rowsInserted(&m_model, SIGNAL(rowsInserted(QModelIndex, int, int)));
    QSignalSpy rowsRemoved(&m_model, SIGNAL(rowsRemoved(QModelIndex, int, int)));

    // Note: It includes a category "Foo" which wasn't in the model previously,
    // we should get an insert event for it too!
    QVERIFY(m_model.setData(indexes.first(), "Home, Errands, Foo"));

    flushNotifications();

    QCOMPARE(rowsRemoved.count(), 2);
    while (!rowsRemoved.isEmpty()) {
        QVariantList signal = rowsRemoved.takeFirst();
        QCOMPARE(signal.count(), 3);

        QModelIndex parent = signal.at(0).value<QModelIndex>();
        QVERIFY(oldParents.contains(parent));

        int expectedRow = oldParentsRow[oldParents.indexOf(parent)];
        QCOMPARE(signal.at(1).toInt(), expectedRow);
        QCOMPARE(signal.at(2).toInt(), expectedRow);
    }

    QCOMPARE(rowsInserted.count(), 4);

    // first insert event should be for "Foo":
    {
        QVariantList signal = rowsInserted.takeFirst();
        QCOMPARE(signal.count(), 3);

        QModelIndex parent = signal.at(0).value<QModelIndex>();
        QVERIFY(!parent.isValid());

        int expectedRow = m_model.rowCount()-1;
        QCOMPARE(signal.at(1).toInt(), expectedRow);
        QCOMPARE(signal.at(2).toInt(), expectedRow);
    }

    // Adjust newParents and newParentsRow to include the new "Foo" category
    newParents << m_model.index(m_model.rowCount()-1, 0);
    newParentsRow << 0;

    // now check the other events for the item itself
    while (!rowsInserted.isEmpty()) {
        QVariantList signal = rowsInserted.takeFirst();
        QCOMPARE(signal.count(), 3);

        QModelIndex parent = signal.at(0).value<QModelIndex>();
        QVERIFY(newParents.contains(parent));

        int expectedRow = newParentsRow[newParents.indexOf(parent)];
        QCOMPARE(signal.at(1).toInt(), expectedRow);
        QCOMPARE(signal.at(2).toInt(), expectedRow);
    }
}

void TodoCategoriesModelTest::testSeveralReparentModification()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(4, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);
    QCOMPARE(indexes.size(), 1);
    QModelIndex index = indexes.takeFirst();

    QModelIndex oldParent = index.parent();
    QCOMPARE(oldParent, m_model.indexForCategory("Home"));

    QModelIndex newParent = m_model.indexForCategory("Errands");

    QSignalSpy rowsInserted(&m_model, SIGNAL(rowsInserted(QModelIndex, int, int)));
    QSignalSpy rowsRemoved(&m_model, SIGNAL(rowsRemoved(QModelIndex, int, int)));

    QVERIFY(m_model.setData(index, "Errands"));

    flushNotifications();

    QCOMPARE(rowsRemoved.count(), 1);
    QVariantList signal = rowsRemoved.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), oldParent);
    QCOMPARE(signal.at(1).toInt(), 0);
    QCOMPARE(signal.at(2).toInt(), 0);

    QCOMPARE(rowsInserted.count(), 1);
    signal = rowsInserted.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), newParent);
    QCOMPARE(signal.at(1).toInt(), 2);
    QCOMPARE(signal.at(2).toInt(), 2);

    indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);
    QCOMPARE(indexes.size(), 1);
    index = indexes.takeFirst();

    QVERIFY(m_model.setData(index, "Home"));

    flushNotifications();

    QCOMPARE(rowsRemoved.count(), 1);
    signal = rowsRemoved.takeFirst();
    QCOMPARE(signal.count(), 3);
    qDebug() << signal.at(0).value<QModelIndex>() << newParent;
    QCOMPARE(signal.at(0).value<QModelIndex>(), newParent);
    QCOMPARE(signal.at(1).toInt(), 2);
    QCOMPARE(signal.at(2).toInt(), 2);

    QCOMPARE(rowsInserted.count(), 1);
    signal = rowsInserted.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), oldParent);
    QCOMPARE(signal.at(1).toInt(), 1);
    QCOMPARE(signal.at(2).toInt(), 1);
}

void TodoCategoriesModelTest::testSingleRemoved()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(4, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);

    QCOMPARE(indexes.size(), 1);
    QModelIndex index = indexes.takeFirst();
    QModelIndex parent = index.parent();
    int count = m_model.rowCount(parent);

    QSignalSpy spy(&m_model, SIGNAL(rowsAboutToBeRemoved(QModelIndex, int, int)));

    Akonadi::ItemDeleteJob *job = new Akonadi::ItemDeleteJob(item);
    QVERIFY(job->exec());

    flushNotifications();

    QCOMPARE(m_model.rowCount(parent), count - 1);

    QCOMPARE(spy.count(), 1);
    QVariantList signal = spy.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), parent);
    QCOMPARE(signal.at(1).toInt(), 1);
    QCOMPARE(signal.at(1).toInt(), 1);
}

void TodoCategoriesModelTest::testMultipleRemoved()
{
    Akonadi::Item item = m_flatModel.itemForIndex(m_flatSortedModel.mapToSource(m_flatSortedModel.index(4, 0)));
    QModelIndexList indexes = m_model.indexesForItem(item, TodoFlatModel::Categories);

    QCOMPARE(indexes.size(), 3);
    QModelIndex index = indexes.takeFirst();
    QModelIndex index2 = indexes.takeFirst();
    QModelIndex index3 = indexes.takeFirst();

    QModelIndex parent = index.parent();
    QModelIndex parent2 = index2.parent();
    QModelIndex parent3 = index3.parent();
    int count = m_model.rowCount(parent);

    QSignalSpy spy(&m_model, SIGNAL(rowsAboutToBeRemoved(QModelIndex, int, int)));

    Akonadi::ItemDeleteJob *job = new Akonadi::ItemDeleteJob(item);
    QVERIFY(job->exec());

    flushNotifications();

    QCOMPARE(m_model.rowCount(parent), count - 1);

    QCOMPARE(spy.count(), 3);
    QVariantList signal = spy.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), parent);
    QCOMPARE(signal.at(1).toInt(), 0);
    QCOMPARE(signal.at(1).toInt(), 0);

    QCOMPARE(spy.count(), 2);
    signal = spy.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), parent2);
    QCOMPARE(signal.at(1).toInt(), 1);
    QCOMPARE(signal.at(1).toInt(), 1);

    QCOMPARE(spy.count(), 1);
    signal = spy.takeFirst();
    QCOMPARE(signal.count(), 3);
    QCOMPARE(signal.at(0).value<QModelIndex>(), parent3);
    QCOMPARE(signal.at(1).toInt(), 0);
    QCOMPARE(signal.at(1).toInt(), 0);
}

#include "todocategoriesmodeltest.moc"
<|endoftext|>


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

#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/io/openni_grabber.h>
#include <pcl/visualization/cloud_viewer.h>
#include <pcl/console/parse.h>
#include <pcl/common/time.h>

#include <pcl/compression/organized_pointcloud_compression.h>

#include <iostream>
#include <vector>
#include <stdio.h>
#include <sstream>
#include <stdlib.h>
#include <iostream>
#include <string>

#include <boost/asio.hpp>

using boost::asio::ip::tcp;

using namespace pcl;
using namespace pcl::io;

using namespace std;

char usage[] = "\n"
  "  PCL organized point cloud stream compression\n"
  "\n"
  "  usage: ./pcl_openni_organized_compression [mode] [parameters]\n"
  "\n"
  "  I/O: \n"
  "      -f file  : file name \n"
  "\n"
  "  file compression mode:\n"
  "      -x: encode point cloud stream to file\n"
  "      -d: decode from file and display point cloud stream\n"
  "\n"
  "  network streaming mode:\n"
  "      -s       : start server on localhost\n"
  "      -c host  : connect to server and display decoded cloud stream\n"
  "\n"
  "  optional compression parameters:\n"
  "      -a       : enable color coding\n"
  "      -t       : output statistics\n"
  "      -e       : show input cloud during encoding\n"
  "      -r       : raw encoding of disparity maps\n"
  "      -g       : gray scale conversion\n"

  "\n"
  "  example:\n"
  "      ./pcl_openni_organized_compression -x -t -f pc_compressed.pcc \n"
  "\n";

#define FPS_CALC(_WHAT_) \
do \
{ \
    static unsigned count = 0;\
    static double last = pcl::getTime ();\
    double now = pcl::getTime (); \
    ++count; \
    if (now - last >= 1.0) \
    { \
      std::cout << "Average framerate("<< _WHAT_ << "): " << double(count)/double(now - last) << " Hz" <<  std::endl; \
      count = 0; \
      last = now; \
    } \
}while(false)

void
print_usage (std::string msg)
{
  std::cerr << msg << std::endl;
  std::cout << usage << std::endl;
}

class SimpleOpenNIViewer
{
  public:
    SimpleOpenNIViewer (ostream& outputFile_arg,
                        OrganizedPointCloudCompression<PointXYZRGBA>* octreeEncoder_arg,
                        bool doColorEncoding_arg,
                        bool bShowStatistics_arg,
                        bool bRawImageEncoding_arg,
                        bool bGrayScaleConversion_arg,
                        int pngLevel_arg = -1) :
    viewer ("Input Point Cloud - PCL Compression Viewer"),
    outputFile_ (outputFile_arg),
    organizedEncoder_ (octreeEncoder_arg),
    doColorEncoding_ (doColorEncoding_arg),
    bShowStatistics_ (bShowStatistics_arg),
    bRawImageEncoding_ (bRawImageEncoding_arg),
    pngLevel_ (pngLevel_arg),
    bGrayScaleConversion_(bGrayScaleConversion_arg)
    {
    }

    void cloud_cb_ (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr &cloud)
    {
      if (!viewer.wasStopped ())
      {
        organizedEncoder_->encodePointCloud (cloud, outputFile_, doColorEncoding_,bGrayScaleConversion_,bShowStatistics_, pngLevel_);

        viewer.showCloud (cloud);
      }
    }

    void run ()
    {

      if (!bRawImageEncoding_)
      {
        // create a new grabber for OpenNI devices
        pcl::Grabber* interface = new pcl::OpenNIGrabber();

        // make callback function from member function
        boost::function<void (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr&)> f =
          boost::bind (&SimpleOpenNIViewer::cloud_cb_, this, _1);

        // connect callback function for desired signal. In this case its a point cloud with color values
        boost::signals2::connection c = interface->registerCallback (f);

        // start receiving point clouds
        interface->start ();

        while (!outputFile_.fail())
        {
          boost::this_thread::sleep(boost::posix_time::seconds(1));
        }

        interface->stop ();
      }

    }

    pcl::visualization::CloudViewer viewer;
    ostream& outputFile_;
    OrganizedPointCloudCompression<PointXYZRGBA>* organizedEncoder_;
    bool doColorEncoding_;
    bool bShowStatistics_;
    bool bRawImageEncoding_;
    int pngLevel_;
    bool bGrayScaleConversion_;
};

struct EventHelper
{
  EventHelper (ostream& outputFile_arg,
               OrganizedPointCloudCompression<PointXYZRGBA>* octreeEncoder_arg,
               bool doColorEncoding_arg,
               bool bShowStatistics_arg,
               bool bRawImageEncoding_arg,
               bool bGrayScaleConversion_arg,
               int pngLevel_arg = -1) :
  outputFile_ (outputFile_arg),
  organizedEncoder_ (octreeEncoder_arg),
  doColorEncoding_ (doColorEncoding_arg),
  bShowStatistics_ (bShowStatistics_arg),
  bRawImageEncoding_ (bRawImageEncoding_arg),
  bGrayScaleConversion_(bGrayScaleConversion_arg) ,
  pngLevel_ (pngLevel_arg)
  {
  }

  void
  cloud_cb_ (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr &cloud)
  {
    if (!outputFile_.fail ())
    {
      organizedEncoder_->encodePointCloud (cloud, outputFile_, doColorEncoding_, bGrayScaleConversion_, bShowStatistics_, pngLevel_);
    }
  }

  void
  image_callback (const boost::shared_ptr<openni_wrapper::Image> &image,
                  const boost::shared_ptr<openni_wrapper::DepthImage> &depth_image, float)
  {

    vector<uint16_t> disparity_data;
    vector<uint8_t> rgb_data;

    uint32_t width=depth_image->getWidth ();
    uint32_t height=depth_image->getHeight ();

    disparity_data.resize(width*height);
    depth_image->fillDepthImageRaw (width, height, &disparity_data[0], static_cast<unsigned int> (width * sizeof (uint16_t)));

    if (image->getEncoding() != openni_wrapper::Image::RGB)
    {
      rgb_data.resize(width*height*3);
      image->fillRGB(width, height, &rgb_data[0], static_cast<unsigned int> (width * sizeof (uint8_t) * 3));
    }

    organizedEncoder_->encodeRawDisparityMapWithColorImage (disparity_data, rgb_data, width, height, outputFile_, doColorEncoding_, bGrayScaleConversion_, bShowStatistics_, pngLevel_);

  }

  void
  run ()
  {
    if (!bRawImageEncoding_)
    {
      // create a new grabber for OpenNI devices
      pcl::Grabber* interface = new pcl::OpenNIGrabber ();

      // make callback function from member function
      boost::function<void
      (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr&)> f = boost::bind (&EventHelper::cloud_cb_, this, _1);

      // connect callback function for desired signal. In this case its a point cloud with color values
      boost::signals2::connection c = interface->registerCallback (f);

      // start receiving point clouds
      interface->start ();

      while (!outputFile_.fail ())
      {
        boost::this_thread::sleep(boost::posix_time::seconds(1));
      }

      interface->stop ();
    } else
    {
      pcl::OpenNIGrabber::Mode image_mode = pcl::OpenNIGrabber::OpenNI_Default_Mode;
      int depthformat = openni_wrapper::OpenNIDevice::OpenNI_shift_values;

      pcl::OpenNIGrabber grabber ("", pcl::OpenNIGrabber::OpenNI_Default_Mode, image_mode);

      // Set the depth output format
      grabber.getDevice ()->setDepthOutputFormat (static_cast<openni_wrapper::OpenNIDevice::DepthMode> (depthformat));

      boost::function<void (const boost::shared_ptr<openni_wrapper::Image>&, const boost::shared_ptr<openni_wrapper::DepthImage>&, float) > image_cb = boost::bind (&EventHelper::image_callback, this, _1, _2, _3);
      boost::signals2::connection image_connection = grabber.registerCallback (image_cb);

      grabber.start ();
      while (!outputFile_.fail())
      {
        boost::this_thread::sleep(boost::posix_time::seconds(1));
      }
      grabber.stop ();

    }
  }

  ostream& outputFile_;
  OrganizedPointCloudCompression<PointXYZRGBA>* organizedEncoder_;
  bool doColorEncoding_;
  bool bShowStatistics_;
  bool bRawImageEncoding_;
  bool bGrayScaleConversion_;
  int pngLevel_;
};

int
main (int argc, char **argv)
{
  OrganizedPointCloudCompression<PointXYZRGBA>* organizedCoder;

  bool showStatistics;
  bool doColorEncoding;
  bool bShowInputCloud;
  bool bRawImageEncoding;
  bool bGrayScaleConversion;

  std::string fileName = "pc_compressed.pcc";
  std::string hostName = "localhost";

  bool bServerFileMode;
  bool bEnDecode;
  bool validArguments;

  validArguments = false;
  bServerFileMode = false;
  bEnDecode = false;

  showStatistics = false;
  doColorEncoding = false;
  bShowInputCloud = false;
  bRawImageEncoding = false;
  bGrayScaleConversion = false;

  if (pcl::console::find_argument (argc, argv, "-e")>0)
    bShowInputCloud = true;

  if (pcl::console::find_argument (argc, argv, "-r")>0)
    bRawImageEncoding = true;

  if (pcl::console::find_argument (argc, argv, "-g")>0)
    bGrayScaleConversion = true;

  if (pcl::console::find_argument (argc, argv, "-s")>0)
  {
    bEnDecode = true;
    bServerFileMode = true;
    validArguments = true;
  }

  if (pcl::console::parse_argument (argc, argv, "-c", hostName)>0)
  {
    bEnDecode = false;
    bServerFileMode = true;
    validArguments = true;
  }

  if (pcl::console::find_argument (argc, argv, "-a")>0)
  {
    doColorEncoding = true;
  }

  if (pcl::console::find_argument (argc, argv, "-x")>0)
  {
    bEnDecode = true;
    bServerFileMode = false;
    validArguments = true;
  }

  if (pcl::console::find_argument (argc, argv, "-d")>0)
  {
    bEnDecode = false;
    bServerFileMode = false;
    validArguments = true;
  }

  if (pcl::console::find_argument (argc, argv, "-t")>0)
    showStatistics = true;

  pcl::console::parse_argument (argc, argv, "-f", fileName);

  if (pcl::console::find_argument (argc, argv, "-?")>0)
  {
    print_usage ("");
    return 1;
  }

  if (!validArguments)
  {
    print_usage ("Please specify compression mode..\n");
    return -1;
  }

  organizedCoder = new OrganizedPointCloudCompression<PointXYZRGBA> ();

  if (!bServerFileMode)
  {
    if (bEnDecode)
    {
      // ENCODING
      ofstream compressedPCFile;
      compressedPCFile.open (fileName.c_str(), ios::out | ios::trunc | ios::binary);

      if (!bShowInputCloud)
      {
        EventHelper v (compressedPCFile, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
        v.run ();
      }
      else
      {
        SimpleOpenNIViewer v (compressedPCFile, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
        v.run ();
      }

    } else
    {
      // DECODING
      ifstream compressedPCFile;
      compressedPCFile.open (fileName.c_str(), ios::in | ios::binary);
      compressedPCFile.seekg (0);
      compressedPCFile.unsetf (ios_base::skipws);

      pcl::visualization::CloudViewer viewer ("PCL Compression Viewer");

      while (!compressedPCFile.eof())
      {
        PointCloud<PointXYZRGBA>::Ptr cloudOut (new PointCloud<PointXYZRGBA> ());
        organizedCoder->decodePointCloud ( compressedPCFile, cloudOut );
        viewer.showCloud (cloudOut);
      }
    }
  } else
  {
    if (bEnDecode)
    {
      // ENCODING
      try
      {
        boost::asio::io_service io_service;
        tcp::endpoint endpoint (tcp::v4 (), 6666);
        tcp::acceptor acceptor (io_service, endpoint);

        tcp::iostream socketStream;

        std::cout << "Waiting for connection.." << std::endl;

        acceptor.accept (*socketStream.rdbuf ());

        std::cout << "Connected!" << std::endl;

        if (!bShowInputCloud)
        {
          EventHelper v (socketStream, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
          v.run ();
        }
        else
        {
          SimpleOpenNIViewer v (socketStream, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
          v.run ();
        }

        std::cout << "Disconnected!" << std::endl;

        boost::this_thread::sleep(boost::posix_time::seconds(3));

      }
      catch (std::exception& e)
      {
        std::cerr << e.what () << std::endl;
      }
    }
    else
    {
      // DECODING
      std::cout << "Connecting to: " << hostName << ".." << std::endl;

      try
      {
        tcp::iostream socketStream (hostName.c_str (), "6666");

        std::cout << "Connected!" << std::endl;

        pcl::visualization::CloudViewer viewer ("Decoded Point Cloud - PCL Compression Viewer");

        while (!socketStream.fail())
        {
          FPS_CALC ("drawing");
          PointCloud<PointXYZRGBA>::Ptr cloudOut (new PointCloud<PointXYZRGBA> ());
          organizedCoder->decodePointCloud (socketStream, cloudOut);
          viewer.showCloud (cloudOut);
        }

      }
      catch (std::exception& e)
      {
        std::cout << "Exception: " << e.what () << std::endl;
      }
    }
  }

  delete (organizedCoder);
  return (0);
}

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

#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/io/openni_grabber.h>
#include <pcl/visualization/cloud_viewer.h>
#include <pcl/console/parse.h>
#include <pcl/common/time.h>

#include <pcl/compression/organized_pointcloud_compression.h>

#include <iostream>
#include <vector>
#include <stdio.h>
#include <sstream>
#include <stdlib.h>
#include <iostream>
#include <string>

#include <boost/asio.hpp>

using boost::asio::ip::tcp;

using namespace pcl;
using namespace pcl::io;

using namespace std;

char usage[] = "\n"
  "  PCL organized point cloud stream compression\n"
  "\n"
  "  usage: ./pcl_openni_organized_compression [mode] [parameters]\n"
  "\n"
  "  I/O: \n"
  "      -f file  : file name \n"
  "\n"
  "  file compression mode:\n"
  "      -x: encode point cloud stream to file\n"
  "      -d: decode from file and display point cloud stream\n"
  "\n"
  "  network streaming mode:\n"
  "      -s       : start server on localhost\n"
  "      -c host  : connect to server and display decoded cloud stream\n"
  "\n"
  "  optional compression parameters:\n"
  "      -a       : enable color coding\n"
  "      -t       : output statistics\n"
  "      -e       : show input cloud during encoding\n"
  "      -r       : raw encoding of disparity maps\n"
  "      -g       : gray scale conversion\n"

  "\n"
  "  example:\n"
  "      ./pcl_openni_organized_compression -x -t -f pc_compressed.pcc \n"
  "\n";

#define FPS_CALC(_WHAT_) \
do \
{ \
    static unsigned count = 0;\
    static double last = pcl::getTime ();\
    double now = pcl::getTime (); \
    ++count; \
    if (now - last >= 1.0) \
    { \
      std::cout << "Average framerate("<< _WHAT_ << "): " << double(count)/double(now - last) << " Hz" <<  std::endl; \
      count = 0; \
      last = now; \
    } \
}while(false)

void
print_usage (std::string msg)
{
  std::cerr << msg << std::endl;
  std::cout << usage << std::endl;
}

class SimpleOpenNIViewer
{
  public:
    SimpleOpenNIViewer (ostream& outputFile_arg,
                        OrganizedPointCloudCompression<PointXYZRGBA>* octreeEncoder_arg,
                        bool doColorEncoding_arg,
                        bool bShowStatistics_arg,
                        bool bRawImageEncoding_arg,
                        bool bGrayScaleConversion_arg,
                        int pngLevel_arg = -1) :
    viewer ("Input Point Cloud - PCL Compression Viewer"),
    outputFile_ (outputFile_arg),
    organizedEncoder_ (octreeEncoder_arg),
    doColorEncoding_ (doColorEncoding_arg),
    bShowStatistics_ (bShowStatistics_arg),
    bRawImageEncoding_ (bRawImageEncoding_arg),
    bGrayScaleConversion_(bGrayScaleConversion_arg),
    pngLevel_ (pngLevel_arg)
    {
    }

    void cloud_cb_ (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr &cloud)
    {
      if (!viewer.wasStopped ())
      {
        organizedEncoder_->encodePointCloud (cloud, outputFile_, doColorEncoding_,bGrayScaleConversion_,bShowStatistics_, pngLevel_);

        viewer.showCloud (cloud);
      }
    }

    void run ()
    {

      if (!bRawImageEncoding_)
      {
        // create a new grabber for OpenNI devices
        pcl::Grabber* interface = new pcl::OpenNIGrabber();

        // make callback function from member function
        boost::function<void (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr&)> f =
          boost::bind (&SimpleOpenNIViewer::cloud_cb_, this, _1);

        // connect callback function for desired signal. In this case its a point cloud with color values
        boost::signals2::connection c = interface->registerCallback (f);

        // start receiving point clouds
        interface->start ();

        while (!outputFile_.fail())
        {
          boost::this_thread::sleep(boost::posix_time::seconds(1));
        }

        interface->stop ();
      }

    }

    pcl::visualization::CloudViewer viewer;
    ostream& outputFile_;
    OrganizedPointCloudCompression<PointXYZRGBA>* organizedEncoder_;
    bool doColorEncoding_;
    bool bShowStatistics_;
    bool bRawImageEncoding_;
    int pngLevel_;
    bool bGrayScaleConversion_;
};

struct EventHelper
{
  EventHelper (ostream& outputFile_arg,
               OrganizedPointCloudCompression<PointXYZRGBA>* octreeEncoder_arg,
               bool doColorEncoding_arg,
               bool bShowStatistics_arg,
               bool bRawImageEncoding_arg,
               bool bGrayScaleConversion_arg,
               int pngLevel_arg = -1) :
  outputFile_ (outputFile_arg),
  organizedEncoder_ (octreeEncoder_arg),
  doColorEncoding_ (doColorEncoding_arg),
  bShowStatistics_ (bShowStatistics_arg),
  bRawImageEncoding_ (bRawImageEncoding_arg),
  bGrayScaleConversion_(bGrayScaleConversion_arg) ,
  pngLevel_ (pngLevel_arg)
  {
  }

  void
  cloud_cb_ (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr &cloud)
  {
    if (!outputFile_.fail ())
    {
      organizedEncoder_->encodePointCloud (cloud, outputFile_, doColorEncoding_, bGrayScaleConversion_, bShowStatistics_, pngLevel_);
    }
  }

  void
  image_callback (const boost::shared_ptr<openni_wrapper::Image> &image,
                  const boost::shared_ptr<openni_wrapper::DepthImage> &depth_image, float)
  {

    vector<uint16_t> disparity_data;
    vector<uint8_t> rgb_data;

    uint32_t width=depth_image->getWidth ();
    uint32_t height=depth_image->getHeight ();

    disparity_data.resize(width*height);
    depth_image->fillDepthImageRaw (width, height, &disparity_data[0], static_cast<unsigned int> (width * sizeof (uint16_t)));

    if (image->getEncoding() != openni_wrapper::Image::RGB)
    {
      rgb_data.resize(width*height*3);
      image->fillRGB(width, height, &rgb_data[0], static_cast<unsigned int> (width * sizeof (uint8_t) * 3));
    }

    organizedEncoder_->encodeRawDisparityMapWithColorImage (disparity_data, rgb_data, width, height, outputFile_, doColorEncoding_, bGrayScaleConversion_, bShowStatistics_, pngLevel_);

  }

  void
  run ()
  {
    if (!bRawImageEncoding_)
    {
      // create a new grabber for OpenNI devices
      pcl::Grabber* interface = new pcl::OpenNIGrabber ();

      // make callback function from member function
      boost::function<void
      (const pcl::PointCloud<pcl::PointXYZRGBA>::ConstPtr&)> f = boost::bind (&EventHelper::cloud_cb_, this, _1);

      // connect callback function for desired signal. In this case its a point cloud with color values
      boost::signals2::connection c = interface->registerCallback (f);

      // start receiving point clouds
      interface->start ();

      while (!outputFile_.fail ())
      {
        boost::this_thread::sleep(boost::posix_time::seconds(1));
      }

      interface->stop ();
    } else
    {
      pcl::OpenNIGrabber::Mode image_mode = pcl::OpenNIGrabber::OpenNI_Default_Mode;
      int depthformat = openni_wrapper::OpenNIDevice::OpenNI_shift_values;

      pcl::OpenNIGrabber grabber ("", pcl::OpenNIGrabber::OpenNI_Default_Mode, image_mode);

      // Set the depth output format
      grabber.getDevice ()->setDepthOutputFormat (static_cast<openni_wrapper::OpenNIDevice::DepthMode> (depthformat));

      boost::function<void (const boost::shared_ptr<openni_wrapper::Image>&, const boost::shared_ptr<openni_wrapper::DepthImage>&, float) > image_cb = boost::bind (&EventHelper::image_callback, this, _1, _2, _3);
      boost::signals2::connection image_connection = grabber.registerCallback (image_cb);

      grabber.start ();
      while (!outputFile_.fail())
      {
        boost::this_thread::sleep(boost::posix_time::seconds(1));
      }
      grabber.stop ();

    }
  }

  ostream& outputFile_;
  OrganizedPointCloudCompression<PointXYZRGBA>* organizedEncoder_;
  bool doColorEncoding_;
  bool bShowStatistics_;
  bool bRawImageEncoding_;
  bool bGrayScaleConversion_;
  int pngLevel_;
};

int
main (int argc, char **argv)
{
  OrganizedPointCloudCompression<PointXYZRGBA>* organizedCoder;

  bool showStatistics;
  bool doColorEncoding;
  bool bShowInputCloud;
  bool bRawImageEncoding;
  bool bGrayScaleConversion;

  std::string fileName = "pc_compressed.pcc";
  std::string hostName = "localhost";

  bool bServerFileMode;
  bool bEnDecode;
  bool validArguments;

  validArguments = false;
  bServerFileMode = false;
  bEnDecode = false;

  showStatistics = false;
  doColorEncoding = false;
  bShowInputCloud = false;
  bRawImageEncoding = false;
  bGrayScaleConversion = false;

  if (pcl::console::find_argument (argc, argv, "-e")>0)
    bShowInputCloud = true;

  if (pcl::console::find_argument (argc, argv, "-r")>0)
    bRawImageEncoding = true;

  if (pcl::console::find_argument (argc, argv, "-g")>0)
    bGrayScaleConversion = true;

  if (pcl::console::find_argument (argc, argv, "-s")>0)
  {
    bEnDecode = true;
    bServerFileMode = true;
    validArguments = true;
  }

  if (pcl::console::parse_argument (argc, argv, "-c", hostName)>0)
  {
    bEnDecode = false;
    bServerFileMode = true;
    validArguments = true;
  }

  if (pcl::console::find_argument (argc, argv, "-a")>0)
  {
    doColorEncoding = true;
  }

  if (pcl::console::find_argument (argc, argv, "-x")>0)
  {
    bEnDecode = true;
    bServerFileMode = false;
    validArguments = true;
  }

  if (pcl::console::find_argument (argc, argv, "-d")>0)
  {
    bEnDecode = false;
    bServerFileMode = false;
    validArguments = true;
  }

  if (pcl::console::find_argument (argc, argv, "-t")>0)
    showStatistics = true;

  pcl::console::parse_argument (argc, argv, "-f", fileName);

  if (pcl::console::find_argument (argc, argv, "-?")>0)
  {
    print_usage ("");
    return 1;
  }

  if (!validArguments)
  {
    print_usage ("Please specify compression mode..\n");
    return -1;
  }

  organizedCoder = new OrganizedPointCloudCompression<PointXYZRGBA> ();

  if (!bServerFileMode)
  {
    if (bEnDecode)
    {
      // ENCODING
      ofstream compressedPCFile;
      compressedPCFile.open (fileName.c_str(), ios::out | ios::trunc | ios::binary);

      if (!bShowInputCloud)
      {
        EventHelper v (compressedPCFile, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
        v.run ();
      }
      else
      {
        SimpleOpenNIViewer v (compressedPCFile, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
        v.run ();
      }

    } else
    {
      // DECODING
      ifstream compressedPCFile;
      compressedPCFile.open (fileName.c_str(), ios::in | ios::binary);
      compressedPCFile.seekg (0);
      compressedPCFile.unsetf (ios_base::skipws);

      pcl::visualization::CloudViewer viewer ("PCL Compression Viewer");

      while (!compressedPCFile.eof())
      {
        PointCloud<PointXYZRGBA>::Ptr cloudOut (new PointCloud<PointXYZRGBA> ());
        organizedCoder->decodePointCloud ( compressedPCFile, cloudOut );
        viewer.showCloud (cloudOut);
      }
    }
  } else
  {
    if (bEnDecode)
    {
      // ENCODING
      try
      {
        boost::asio::io_service io_service;
        tcp::endpoint endpoint (tcp::v4 (), 6666);
        tcp::acceptor acceptor (io_service, endpoint);

        tcp::iostream socketStream;

        std::cout << "Waiting for connection.." << std::endl;

        acceptor.accept (*socketStream.rdbuf ());

        std::cout << "Connected!" << std::endl;

        if (!bShowInputCloud)
        {
          EventHelper v (socketStream, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
          v.run ();
        }
        else
        {
          SimpleOpenNIViewer v (socketStream, organizedCoder, doColorEncoding, showStatistics, bRawImageEncoding, bGrayScaleConversion);
          v.run ();
        }

        std::cout << "Disconnected!" << std::endl;

        boost::this_thread::sleep(boost::posix_time::seconds(3));

      }
      catch (std::exception& e)
      {
        std::cerr << e.what () << std::endl;
      }
    }
    else
    {
      // DECODING
      std::cout << "Connecting to: " << hostName << ".." << std::endl;

      try
      {
        tcp::iostream socketStream (hostName.c_str (), "6666");

        std::cout << "Connected!" << std::endl;

        pcl::visualization::CloudViewer viewer ("Decoded Point Cloud - PCL Compression Viewer");

        while (!socketStream.fail())
        {
          FPS_CALC ("drawing");
          PointCloud<PointXYZRGBA>::Ptr cloudOut (new PointCloud<PointXYZRGBA> ());
          organizedCoder->decodePointCloud (socketStream, cloudOut);
          viewer.showCloud (cloudOut);
        }

      }
      catch (std::exception& e)
      {
        std::cout << "Exception: " << e.what () << std::endl;
      }
    }
  }

  delete (organizedCoder);
  return (0);
}

<|endoftext|>


<commit_before>//
// This file is part of the Marble Virtual Globe.
//
// This program is free software licensed under the GNU LGPL. You can
// find a copy of this license in LICENSE.txt in the top directory of
// the source code.
//
// Copyright 2009      Andrew Manson  <g.real.ate@gmail.com>
// Copyright 2013      Thibaut Gridel <tgridel@free.fr>
// Copyright 2014      Calin Cruceru  <crucerucalincristian@gmail.com>
//

// Self
#include "PlacemarkTextAnnotation.h"

// Qt
#include <QImage>

// Marble
#include "AbstractProjection.h"
#include "GeoDataPlacemark.h"
#include "GeoDataStyle.h"
#include "GeoPainter.h"
#include "ViewportParams.h"
#include "MarbleDirs.h"
#include "SceneGraphicsTypes.h"

namespace Marble
{

PlacemarkTextAnnotation::PlacemarkTextAnnotation( GeoDataPlacemark *placemark ) :
    SceneGraphicsItem( placemark ),
    m_movingPlacemark( false )
{
    GeoDataStyle *newStyle = new GeoDataStyle( *placemark->style() );
    newStyle->iconStyle().setIcon( QImage() );
    newStyle->iconStyle().setIconPath( MarbleDirs::path( "bitmaps/default_location.png" ) );
    placemark->setStyle( newStyle );
}

PlacemarkTextAnnotation::~PlacemarkTextAnnotation()
{
    // nothing to do
}

void PlacemarkTextAnnotation::paint( GeoPainter *painter, const ViewportParams *viewport )
{
    m_viewport = viewport;
    painter->drawImage( placemark()->coordinate(), QImage( placemark()->style()->iconStyle().iconPath() ) );

    qreal x, y;
    viewport->currentProjection()->screenCoordinates( placemark()->coordinate(), viewport, x, y );

    m_region = QRegion( x - 10 , y - 10 , 20 , 20 );
}

bool PlacemarkTextAnnotation::containsPoint( const QPoint &eventPos ) const
{
    return m_region.contains( eventPos );
}

void PlacemarkTextAnnotation::dealWithItemChange( const SceneGraphicsItem *other )
{
    Q_UNUSED( other );
}

void PlacemarkTextAnnotation::move( const GeoDataCoordinates &source, const GeoDataCoordinates &destination )
{
    Q_UNUSED( source );
    qreal lat = destination.latitude();
    qreal lon = destination.longitude();
    GeoDataCoordinates::normalizeLonLat( lon, lat );
    placemark()->setCoordinate( lon, lat );
}

const char *PlacemarkTextAnnotation::graphicType() const
{
    return SceneGraphicsTypes::SceneGraphicPlacemark;
}

bool PlacemarkTextAnnotation::mousePressEvent( QMouseEvent *event )
{
    Q_UNUSED( event );

    setRequest( SceneGraphicsItem::NoRequest );

    if ( state() == SceneGraphicsItem::Editing ) {
        if ( event->button() == Qt::LeftButton ) {
            m_movingPlacemark = true;
        } else if ( event->button() == Qt::RightButton ) {
            setRequest( SceneGraphicsItem::ShowPlacemarkRmbMenu );
        }

        return true;
    }

    return false;
}

bool PlacemarkTextAnnotation::mouseMoveEvent( QMouseEvent *event )
{
    setRequest( SceneGraphicsItem::NoRequest );

    qreal lon, lat;
    m_viewport->geoCoordinates( event->pos().x(),
                                event->pos().y(),
                                lon, lat,
                                GeoDataCoordinates::Radian );

    if ( m_movingPlacemark ) {
        placemark()->setCoordinate( lon, lat );
        return true;
    }

    return false;
}

bool PlacemarkTextAnnotation::mouseReleaseEvent( QMouseEvent *event )
{
    Q_UNUSED( event );
    setRequest( SceneGraphicsItem::NoRequest );

    m_movingPlacemark = false;
    return true;
}

void PlacemarkTextAnnotation::dealWithStateChange( SceneGraphicsItem::ActionState previousState )
{
    Q_UNUSED( previousState );
}

}
<commit_msg>Painting the image for placemark is not needed - it is already done in placemark layout<commit_after>//
// This file is part of the Marble Virtual Globe.
//
// This program is free software licensed under the GNU LGPL. You can
// find a copy of this license in LICENSE.txt in the top directory of
// the source code.
//
// Copyright 2009      Andrew Manson  <g.real.ate@gmail.com>
// Copyright 2013      Thibaut Gridel <tgridel@free.fr>
// Copyright 2014      Calin Cruceru  <crucerucalincristian@gmail.com>
//

// Self
#include "PlacemarkTextAnnotation.h"

// Qt
#include <QImage>

// Marble
#include "AbstractProjection.h"
#include "GeoDataPlacemark.h"
#include "GeoDataStyle.h"
#include "GeoPainter.h"
#include "ViewportParams.h"
#include "MarbleDirs.h"
#include "SceneGraphicsTypes.h"

namespace Marble
{

PlacemarkTextAnnotation::PlacemarkTextAnnotation( GeoDataPlacemark *placemark ) :
    SceneGraphicsItem( placemark ),
    m_movingPlacemark( false )
{
    GeoDataStyle *newStyle = new GeoDataStyle( *placemark->style() );
    newStyle->iconStyle().setIcon( QImage() );
    newStyle->iconStyle().setIconPath( MarbleDirs::path( "bitmaps/default_location.png" ) );
    placemark->setStyle( newStyle );
}

PlacemarkTextAnnotation::~PlacemarkTextAnnotation()
{
    // nothing to do
}

void PlacemarkTextAnnotation::paint( GeoPainter *painter, const ViewportParams *viewport )
{
    Q_UNUSED( painter );
    m_viewport = viewport;

    qreal x, y;
    viewport->currentProjection()->screenCoordinates( placemark()->coordinate(), viewport, x, y );
    m_region = QRegion( x - 10 , y - 10 , 20 , 20 );
}

bool PlacemarkTextAnnotation::containsPoint( const QPoint &eventPos ) const
{
    return m_region.contains( eventPos );
}

void PlacemarkTextAnnotation::dealWithItemChange( const SceneGraphicsItem *other )
{
    Q_UNUSED( other );
}

void PlacemarkTextAnnotation::move( const GeoDataCoordinates &source, const GeoDataCoordinates &destination )
{
    Q_UNUSED( source );
    qreal lat = destination.latitude();
    qreal lon = destination.longitude();
    GeoDataCoordinates::normalizeLonLat( lon, lat );
    placemark()->setCoordinate( lon, lat );
}

const char *PlacemarkTextAnnotation::graphicType() const
{
    return SceneGraphicsTypes::SceneGraphicPlacemark;
}

bool PlacemarkTextAnnotation::mousePressEvent( QMouseEvent *event )
{
    Q_UNUSED( event );

    setRequest( SceneGraphicsItem::NoRequest );

    if ( state() == SceneGraphicsItem::Editing ) {
        if ( event->button() == Qt::LeftButton ) {
            m_movingPlacemark = true;
        } else if ( event->button() == Qt::RightButton ) {
            setRequest( SceneGraphicsItem::ShowPlacemarkRmbMenu );
        }

        return true;
    }

    return false;
}

bool PlacemarkTextAnnotation::mouseMoveEvent( QMouseEvent *event )
{
    setRequest( SceneGraphicsItem::NoRequest );

    qreal lon, lat;
    m_viewport->geoCoordinates( event->pos().x(),
                                event->pos().y(),
                                lon, lat,
                                GeoDataCoordinates::Radian );

    if ( m_movingPlacemark ) {
        placemark()->setCoordinate( lon, lat );
        return true;
    }

    return false;
}

bool PlacemarkTextAnnotation::mouseReleaseEvent( QMouseEvent *event )
{
    Q_UNUSED( event );
    setRequest( SceneGraphicsItem::NoRequest );

    m_movingPlacemark = false;
    return true;
}

void PlacemarkTextAnnotation::dealWithStateChange( SceneGraphicsItem::ActionState previousState )
{
    Q_UNUSED( previousState );
}

}
<|endoftext|>


<commit_before>#include <boost/program_options.hpp>
#include <boost/filesystem/fstream.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/trim.hpp>
#include "picpac.h"
#include "picpac-cv.h"
#include "picpac-util.h"

using namespace std;
using namespace boost;
using namespace picpac;

enum {
    FORMAT_LIST = 1,
    FORMAT_DIR = 2,
    FORMAT_ANNO_JSON = 3,
    FORMAT_ANNO_IMAGE = 4
};

class Paths: public vector<fs::path> {
public:
    Paths () {}
    Paths (fs::path const &path) {
        fs::recursive_directory_iterator it(fs::path(path), fs::symlink_option::recurse), end;
        for (; it != end; ++it) {
            if (it->status().type() == fs::regular_file) {
                push_back(it->path());
            }
        }
        CHECK(size() >= 10) << "Need at least 10 files to train: " << path;
    }
};

class Samples: public vector<Paths> {
public:
    Samples () {}
    Samples (fs::path const &root) {
        fs::directory_iterator it(root), end;
        vector<unsigned> cats;
        for (; it != end; ++it) {
            if (!fs::is_directory(it->path())) {
                LOG(ERROR) << "Not a directory: " << it->path();
                continue;
            }
            fs::path name = it->path().filename();
            try {
                unsigned c = lexical_cast<unsigned>(name.native());
                cats.push_back(c);
            }
            catch (...) {
                LOG(ERROR) << "Category directory not properly named: " << it->path();
            }
        }
        sort(cats.begin(), cats.end());
        cats.resize(unique(cats.begin(), cats.end()) - cats.begin());
        CHECK(cats.size() >= 2) << "Need at least 2 categories to train.";
        CHECK((cats.front() == 0)
                && (cats.back() == cats.size() -1 ))
            << "Subdirectories must be consecutively named from 0 to N-1.";
        for (unsigned c = 0; c < cats.size(); ++c) {
            emplace_back(root / fs::path(lexical_cast<string>(c)));
            LOG(INFO) << "Loaded " << back().size() << " paths for category " << c << ".";
        }
    }
};

int main(int argc, char **argv) {
    namespace po = boost::program_options;
    fs::path cache;
    fs::path input_path;
    fs::path output_path;
    int max_size;
    int resize;
    int format;

    po::options_description desc("Allowed options");
    desc.add_options()
    ("help,h", "produce help message.")
    ("input,i", po::value(&input_path), "")
    ("output,o", po::value(&output_path), "")
    ("max", po::value(&max_size)->default_value(-1), "")
    ("resize", po::value(&resize)->default_value(-1), "")
    ("format,f", po::value(&format)->default_value(1), "")
    ("cache", po::value(&cache)->default_value(".picpac_cache"), "")
    /*
    ("gray", "")
    ("log-level,v", po::value(&FLAGS_minloglevel)->default_value(1), "")
    */
    ;

    po::positional_options_description p;
    p.add("input", 1);
    p.add("output", 1);

    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).
                     options(desc).positional(p).run(), vm);
    po::notify(vm);

    if (vm.count("help") || input_path.empty() || output_path.empty()) {
        cerr << desc;
        cerr << endl;
        cerr << "Formats:" << endl
             << "  1: list of <image\tlabel>" << endl
             << "  2: scan a directory" << endl
             << "  3: list of <image\tjson-annotation>" << endl
             << "  4: list of <image\tannotation-image>" << endl;

        return 1;
    }
    //if (vm.count("gray")) gray = true;

    google::InitGoogleLogging(argv[0]);
    picpac::FileWriter db(output_path);
    CachedDownloader downloader(cache);
    ImageReader imreader(max_size, resize);
    int count = 0;

    if (format == FORMAT_DIR) {
        Samples all(input_path);
        for (unsigned i = 0; i < all.size(); ++i) {
            for (auto const &path: all[i]) {
                string data;
                imreader.read(path, &data);
                if (data.empty()) {
                    LOG(ERROR) << "not a image: " << path;
                    continue;
                }
                picpac::Record rec(i, data);
                db.append(rec);
                ++count;
            }
        }
    }
    else {
        fs::ifstream is(input_path.c_str());
        string line;
        while (getline(is, line)) {
            vector<string> ss;
            split(ss, line, is_any_of("\t"), token_compress_off);
            if (ss.size() != 2) {
                cerr << "Bad line: " << line << endl;
                continue;
            }
            try {
                fs::path path = downloader.download(ss[0]);
                string data;
                imreader.read(path, &data);
                if (data.empty()) {
                    LOG(ERROR) << "not a image: " << path;
                    continue;
                }
                if (format == FORMAT_LIST) {
                    float l = lexical_cast<float>(ss[1]);
                    Record record(l, data);
                    db.append(record);
                }
                else if (format == FORMAT_ANNO_JSON) {
                    Record record(0, data, ss[1]);
                    db.append(record);
                }
                else if (format == FORMAT_ANNO_IMAGE) {
                    string data2;
                    if (ss[1].size()) {
                        fs::path path2 = downloader.download(ss[1]);
                        imreader.read(path2, &data2);
                    }
                    Record record(0, data, data2);
                    db.append(record);
                }
                ++count;
            }
            catch (...) {
                LOG(ERROR) << "Fail to load " << ss[0];
            }
        }
    }
    LOG(INFO) << "Loaded " << count << " samples.";

    return 0;
}

<commit_msg>added whole directory import<commit_after>#include <boost/program_options.hpp>
#include <boost/filesystem/fstream.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/trim.hpp>
#include "picpac.h"
#include "picpac-cv.h"
#include "picpac-util.h"

using namespace std;
using namespace boost;
using namespace picpac;

enum {
    FORMAT_DIR = 0,
    FORMAT_LIST = 1,
    FORMAT_SUB_DIR = 2,
    FORMAT_ANNO_JSON = 3,
    FORMAT_ANNO_IMAGE = 4
};

class Paths: public vector<fs::path> {
public:
    Paths () {}
    Paths (fs::path const &path) {
        fs::recursive_directory_iterator it(fs::path(path), fs::symlink_option::recurse), end;
        for (; it != end; ++it) {
            if (it->status().type() == fs::regular_file) {
                push_back(it->path());
            }
        }
        CHECK(size() >= 10) << "Need at least 10 files to train: " << path;
    }
};

class Samples: public vector<Paths> {
public:
    Samples () {}
    Samples (fs::path const &root) {
        fs::directory_iterator it(root), end;
        vector<unsigned> cats;
        for (; it != end; ++it) {
            if (!fs::is_directory(it->path())) {
                LOG(ERROR) << "Not a directory: " << it->path();
                continue;
            }
            fs::path name = it->path().filename();
            try {
                unsigned c = lexical_cast<unsigned>(name.native());
                cats.push_back(c);
            }
            catch (...) {
                LOG(ERROR) << "Category directory not properly named: " << it->path();
            }
        }
        sort(cats.begin(), cats.end());
        cats.resize(unique(cats.begin(), cats.end()) - cats.begin());
        CHECK(cats.size() >= 2) << "Need at least 2 categories to train.";
        CHECK((cats.front() == 0)
                && (cats.back() == cats.size() -1 ))
            << "Subdirectories must be consecutively named from 0 to N-1.";
        for (unsigned c = 0; c < cats.size(); ++c) {
            emplace_back(root / fs::path(lexical_cast<string>(c)));
            LOG(INFO) << "Loaded " << back().size() << " paths for category " << c << ".";
        }
    }
};

int main(int argc, char **argv) {
    namespace po = boost::program_options;
    fs::path cache;
    fs::path input_path;
    fs::path output_path;
    int max_size;
    int resize;
    int format;

    po::options_description desc("Allowed options");
    desc.add_options()
    ("help,h", "produce help message.")
    ("input,i", po::value(&input_path), "")
    ("output,o", po::value(&output_path), "")
    ("max", po::value(&max_size)->default_value(-1), "")
    ("resize", po::value(&resize)->default_value(-1), "")
    ("format,f", po::value(&format)->default_value(1), "")
    ("cache", po::value(&cache)->default_value(".picpac_cache"), "")
    /*
    ("gray", "")
    ("log-level,v", po::value(&FLAGS_minloglevel)->default_value(1), "")
    */
    ;

    po::positional_options_description p;
    p.add("input", 1);
    p.add("output", 1);

    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).
                     options(desc).positional(p).run(), vm);
    po::notify(vm);

    if (vm.count("help") || input_path.empty() || output_path.empty()) {
        cerr << desc;
        cerr << endl;
        cerr << "Formats:" << endl
             << "  1: list of <image\tlabel>" << endl
             << "  2: scan a directory" << endl
             << "  3: list of <image\tjson-annotation>" << endl
             << "  4: list of <image\tannotation-image>" << endl;

        return 1;
    }
    //if (vm.count("gray")) gray = true;

    google::InitGoogleLogging(argv[0]);
    picpac::FileWriter db(output_path);
    CachedDownloader downloader(cache);
    ImageReader imreader(max_size, resize);
    int count = 0;

    if (format == FORMAT_DIR || format == FORMAT_SUB_DIR) {
        Samples all;
        if (format == FORMAT_SUB_DIR) {
            Samples tmp(input_path);
            all.swap(tmp);
        }
        else if (format == FORMAT_DIR) {
            all.emplace_back(input_path);
        }
        for (unsigned i = 0; i < all.size(); ++i) {
            for (auto const &path: all[i]) {
                string data;
                imreader.read(path, &data);
                if (data.empty()) {
                    LOG(ERROR) << "not a image: " << path;
                    continue;
                }
                picpac::Record rec(i, data);
                db.append(rec);
                ++count;
            }
        }
    }
    else {
        fs::ifstream is(input_path.c_str());
        string line;
        while (getline(is, line)) {
            vector<string> ss;
            split(ss, line, is_any_of("\t"), token_compress_off);
            if (ss.size() != 2) {
                cerr << "Bad line: " << line << endl;
                continue;
            }
            try {
                fs::path path = downloader.download(ss[0]);
                string data;
                imreader.read(path, &data);
                if (data.empty()) {
                    LOG(ERROR) << "not a image: " << path;
                    continue;
                }
                if (format == FORMAT_LIST) {
                    float l = lexical_cast<float>(ss[1]);
                    Record record(l, data);
                    db.append(record);
                }
                else if (format == FORMAT_ANNO_JSON) {
                    Record record(0, data, ss[1]);
                    db.append(record);
                }
                else if (format == FORMAT_ANNO_IMAGE) {
                    string data2;
                    if (ss[1].size()) {
                        fs::path path2 = downloader.download(ss[1]);
                        imreader.read(path2, &data2);
                    }
                    Record record(0, data, data2);
                    db.append(record);
                }
                ++count;
            }
            catch (...) {
                LOG(ERROR) << "Fail to load " << ss[0];
            }
        }
    }
    LOG(INFO) << "Loaded " << count << " samples.";

    return 0;
}

<|endoftext|>


<commit_before>/*************************************************************************
 *
 *  OpenOffice.org - a multi-platform office productivity suite
 *
 *  $RCSfile: xmlsignaturetemplateimpl.cxx,v $
 *
 *  $Revision: 1.6 $
 *
 *  last change: $Author: ihi $ $Date: 2007-04-17 10:20:13 $
 *
 *  The Contents of this file are made available subject to
 *  the terms of GNU Lesser General Public License Version 2.1.
 *
 *
 *    GNU Lesser General Public License Version 2.1
 *    =============================================
 *    Copyright 2005 by Sun Microsystems, Inc.
 *    901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *    This library is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License version 2.1, as published by the Free Software Foundation.
 *
 *    This library is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with this library; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *    MA  02111-1307  USA
 *
 ************************************************************************/

// MARKER(update_precomp.py): autogen include statement, do not remove
#include "precompiled_xmlsecurity.hxx"

#ifndef _SAL_CONFIG_H_
#include <sal/config.h>
#endif

#ifndef _RTL_USTRING_HXX_
#include <rtl/ustring.hxx>
#endif

#ifndef _RTL_UUID_H_
#include <rtl/uuid.h>
#endif

#ifndef _XMLSIGNATURETEMPLATEIMPL_HXX_
#include "xmlsignaturetemplateimpl.hxx"
#endif

using namespace ::com::sun::star::uno ;
using ::com::sun::star::lang::XMultiServiceFactory ;
using ::com::sun::star::lang::XSingleServiceFactory ;
using ::rtl::OUString ;

using ::com::sun::star::xml::wrapper::XXMLElementWrapper ;
using ::com::sun::star::xml::crypto::XXMLSignatureTemplate ;

XMLSignatureTemplateImpl :: XMLSignatureTemplateImpl( const Reference< XMultiServiceFactory >& aFactory )
    :m_xTemplate( NULL ),
     m_xServiceManager( aFactory ),
     m_nStatus ( ::com::sun::star::xml::crypto::SecurityOperationStatus_UNKNOWN )
{
}

XMLSignatureTemplateImpl :: ~XMLSignatureTemplateImpl() {
}

/* XXMLSignatureTemplate */
void SAL_CALL XMLSignatureTemplateImpl :: setTemplate( const Reference< XXMLElementWrapper >& aTemplate )
    throw( com::sun::star::uno::RuntimeException, com::sun::star::lang::IllegalArgumentException)
{
    m_xTemplate = aTemplate ;
}

/* XXMLSignatureTemplate */
Reference< XXMLElementWrapper > SAL_CALL XMLSignatureTemplateImpl :: getTemplate()
    throw (com::sun::star::uno::RuntimeException)
{
    return m_xTemplate ;
}

void SAL_CALL XMLSignatureTemplateImpl :: setTarget( const ::com::sun::star::uno::Reference< ::com::sun::star::xml::wrapper::XXMLElementWrapper >& aXmlElement )
    throw( com::sun::star::uno::RuntimeException, com::sun::star::lang::IllegalArgumentException)
{
    targets.push_back( aXmlElement );
}

::com::sun::star::uno::Sequence< ::com::sun::star::uno::Reference< ::com::sun::star::xml::wrapper::XXMLElementWrapper > > SAL_CALL XMLSignatureTemplateImpl :: getTargets()
    throw (com::sun::star::uno::RuntimeException)
{
    sal_Int32 length = targets.size();
    ::com::sun::star::uno::Sequence<
        ::com::sun::star::uno::Reference< ::com::sun::star::xml::wrapper::XXMLElementWrapper >
        > aTargets (length);

    sal_Int32 i;

    for (i=0; i<length; i++)
    {
        aTargets[i] = targets[i];
    }

    return aTargets;
}

void SAL_CALL XMLSignatureTemplateImpl::setBinding(
    const ::com::sun::star::uno::Reference<
        ::com::sun::star::xml::crypto::XUriBinding >& aUriBinding )
    throw (::com::sun::star::lang::IllegalArgumentException, ::com::sun::star::uno::RuntimeException)
{
    m_xUriBinding = aUriBinding;
}

::com::sun::star::uno::Reference< ::com::sun::star::xml::crypto::XUriBinding > SAL_CALL XMLSignatureTemplateImpl::getBinding()
    throw (::com::sun::star::uno::RuntimeException)
{
    return m_xUriBinding;
}

void SAL_CALL XMLSignatureTemplateImpl::setStatus(
    ::com::sun::star::xml::crypto::SecurityOperationStatus status )
    throw (::com::sun::star::lang::IllegalArgumentException, ::com::sun::star::uno::RuntimeException)
{
    m_nStatus = status;
}

::com::sun::star::xml::crypto::SecurityOperationStatus SAL_CALL XMLSignatureTemplateImpl::getStatus(  )
    throw (::com::sun::star::uno::RuntimeException)
{
    return m_nStatus;
}

/* XInitialization */
void SAL_CALL XMLSignatureTemplateImpl :: initialize( const Sequence< Any >& /*aArguments*/ ) throw( Exception, RuntimeException ) {
    // TBD
} ;

/* XServiceInfo */
OUString SAL_CALL XMLSignatureTemplateImpl :: getImplementationName() throw( RuntimeException ) {
    return impl_getImplementationName() ;
}

/* XServiceInfo */
sal_Bool SAL_CALL XMLSignatureTemplateImpl :: supportsService( const OUString& serviceName) throw( RuntimeException ) {
    Sequence< OUString > seqServiceNames = getSupportedServiceNames() ;
    const OUString* pArray = seqServiceNames.getConstArray() ;
    for( sal_Int32 i = 0 ; i < seqServiceNames.getLength() ; i ++ ) {
        if( *( pArray + i ) == serviceName )
            return sal_True ;
    }
    return sal_False ;
}

/* XServiceInfo */
Sequence< OUString > SAL_CALL XMLSignatureTemplateImpl :: getSupportedServiceNames() throw( RuntimeException ) {
    return impl_getSupportedServiceNames() ;
}

//Helper for XServiceInfo
Sequence< OUString > XMLSignatureTemplateImpl :: impl_getSupportedServiceNames() {
    ::osl::Guard< ::osl::Mutex > aGuard( ::osl::Mutex::getGlobalMutex() ) ;
    Sequence< OUString > seqServiceNames( 1 ) ;
    seqServiceNames.getArray()[0] = OUString::createFromAscii( "com.sun.star.xml.crypto.XMLSignatureTemplate" ) ;
    return seqServiceNames ;
}

OUString XMLSignatureTemplateImpl :: impl_getImplementationName() throw( RuntimeException ) {
    return OUString::createFromAscii( "com.sun.star.xml.security.framework.XMLSignatureTemplateImpl" ) ;
}

//Helper for registry
Reference< XInterface > SAL_CALL XMLSignatureTemplateImpl :: impl_createInstance( const Reference< XMultiServiceFactory >& aServiceManager ) throw( RuntimeException ) {
    return Reference< XInterface >( *new XMLSignatureTemplateImpl( aServiceManager ) ) ;
}

Reference< XSingleServiceFactory > XMLSignatureTemplateImpl :: impl_createFactory( const Reference< XMultiServiceFactory >& aServiceManager ) {
    //Reference< XSingleServiceFactory > xFactory ;
    //xFactory = ::cppu::createSingleFactory( aServiceManager , impl_getImplementationName , impl_createInstance , impl_getSupportedServiceNames ) ;
    //return xFactory ;
    return ::cppu::createSingleFactory( aServiceManager , impl_getImplementationName() , impl_createInstance , impl_getSupportedServiceNames() ) ;
}

<commit_msg>INTEGRATION: CWS changefileheader (1.6.66); FILE MERGED 2008/04/01 16:10:59 thb 1.6.66.3: #i85898# Stripping all external header guards 2008/04/01 13:06:22 thb 1.6.66.2: #i85898# Stripping all external header guards 2008/03/31 16:31:01 rt 1.6.66.1: #i87441# Change license header to LPGL v3.<commit_after>/*************************************************************************
 *
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * Copyright 2008 by Sun Microsystems, Inc.
 *
 * OpenOffice.org - a multi-platform office productivity suite
 *
 * $RCSfile: xmlsignaturetemplateimpl.cxx,v $
 * $Revision: 1.7 $
 *
 * This file is part of OpenOffice.org.
 *
 * OpenOffice.org is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * OpenOffice.org is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License version 3 for more details
 * (a copy is included in the LICENSE file that accompanied this code).
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with OpenOffice.org.  If not, see
 * <http://www.openoffice.org/license.html>
 * for a copy of the LGPLv3 License.
 *
 ************************************************************************/

// MARKER(update_precomp.py): autogen include statement, do not remove
#include "precompiled_xmlsecurity.hxx"
#include <sal/config.h>
#include <rtl/ustring.hxx>
#include <rtl/uuid.h>
#include "xmlsignaturetemplateimpl.hxx"

using namespace ::com::sun::star::uno ;
using ::com::sun::star::lang::XMultiServiceFactory ;
using ::com::sun::star::lang::XSingleServiceFactory ;
using ::rtl::OUString ;

using ::com::sun::star::xml::wrapper::XXMLElementWrapper ;
using ::com::sun::star::xml::crypto::XXMLSignatureTemplate ;

XMLSignatureTemplateImpl :: XMLSignatureTemplateImpl( const Reference< XMultiServiceFactory >& aFactory )
    :m_xTemplate( NULL ),
     m_xServiceManager( aFactory ),
     m_nStatus ( ::com::sun::star::xml::crypto::SecurityOperationStatus_UNKNOWN )
{
}

XMLSignatureTemplateImpl :: ~XMLSignatureTemplateImpl() {
}

/* XXMLSignatureTemplate */
void SAL_CALL XMLSignatureTemplateImpl :: setTemplate( const Reference< XXMLElementWrapper >& aTemplate )
    throw( com::sun::star::uno::RuntimeException, com::sun::star::lang::IllegalArgumentException)
{
    m_xTemplate = aTemplate ;
}

/* XXMLSignatureTemplate */
Reference< XXMLElementWrapper > SAL_CALL XMLSignatureTemplateImpl :: getTemplate()
    throw (com::sun::star::uno::RuntimeException)
{
    return m_xTemplate ;
}

void SAL_CALL XMLSignatureTemplateImpl :: setTarget( const ::com::sun::star::uno::Reference< ::com::sun::star::xml::wrapper::XXMLElementWrapper >& aXmlElement )
    throw( com::sun::star::uno::RuntimeException, com::sun::star::lang::IllegalArgumentException)
{
    targets.push_back( aXmlElement );
}

::com::sun::star::uno::Sequence< ::com::sun::star::uno::Reference< ::com::sun::star::xml::wrapper::XXMLElementWrapper > > SAL_CALL XMLSignatureTemplateImpl :: getTargets()
    throw (com::sun::star::uno::RuntimeException)
{
    sal_Int32 length = targets.size();
    ::com::sun::star::uno::Sequence<
        ::com::sun::star::uno::Reference< ::com::sun::star::xml::wrapper::XXMLElementWrapper >
        > aTargets (length);

    sal_Int32 i;

    for (i=0; i<length; i++)
    {
        aTargets[i] = targets[i];
    }

    return aTargets;
}

void SAL_CALL XMLSignatureTemplateImpl::setBinding(
    const ::com::sun::star::uno::Reference<
        ::com::sun::star::xml::crypto::XUriBinding >& aUriBinding )
    throw (::com::sun::star::lang::IllegalArgumentException, ::com::sun::star::uno::RuntimeException)
{
    m_xUriBinding = aUriBinding;
}

::com::sun::star::uno::Reference< ::com::sun::star::xml::crypto::XUriBinding > SAL_CALL XMLSignatureTemplateImpl::getBinding()
    throw (::com::sun::star::uno::RuntimeException)
{
    return m_xUriBinding;
}

void SAL_CALL XMLSignatureTemplateImpl::setStatus(
    ::com::sun::star::xml::crypto::SecurityOperationStatus status )
    throw (::com::sun::star::lang::IllegalArgumentException, ::com::sun::star::uno::RuntimeException)
{
    m_nStatus = status;
}

::com::sun::star::xml::crypto::SecurityOperationStatus SAL_CALL XMLSignatureTemplateImpl::getStatus(  )
    throw (::com::sun::star::uno::RuntimeException)
{
    return m_nStatus;
}

/* XInitialization */
void SAL_CALL XMLSignatureTemplateImpl :: initialize( const Sequence< Any >& /*aArguments*/ ) throw( Exception, RuntimeException ) {
    // TBD
} ;

/* XServiceInfo */
OUString SAL_CALL XMLSignatureTemplateImpl :: getImplementationName() throw( RuntimeException ) {
    return impl_getImplementationName() ;
}

/* XServiceInfo */
sal_Bool SAL_CALL XMLSignatureTemplateImpl :: supportsService( const OUString& serviceName) throw( RuntimeException ) {
    Sequence< OUString > seqServiceNames = getSupportedServiceNames() ;
    const OUString* pArray = seqServiceNames.getConstArray() ;
    for( sal_Int32 i = 0 ; i < seqServiceNames.getLength() ; i ++ ) {
        if( *( pArray + i ) == serviceName )
            return sal_True ;
    }
    return sal_False ;
}

/* XServiceInfo */
Sequence< OUString > SAL_CALL XMLSignatureTemplateImpl :: getSupportedServiceNames() throw( RuntimeException ) {
    return impl_getSupportedServiceNames() ;
}

//Helper for XServiceInfo
Sequence< OUString > XMLSignatureTemplateImpl :: impl_getSupportedServiceNames() {
    ::osl::Guard< ::osl::Mutex > aGuard( ::osl::Mutex::getGlobalMutex() ) ;
    Sequence< OUString > seqServiceNames( 1 ) ;
    seqServiceNames.getArray()[0] = OUString::createFromAscii( "com.sun.star.xml.crypto.XMLSignatureTemplate" ) ;
    return seqServiceNames ;
}

OUString XMLSignatureTemplateImpl :: impl_getImplementationName() throw( RuntimeException ) {
    return OUString::createFromAscii( "com.sun.star.xml.security.framework.XMLSignatureTemplateImpl" ) ;
}

//Helper for registry
Reference< XInterface > SAL_CALL XMLSignatureTemplateImpl :: impl_createInstance( const Reference< XMultiServiceFactory >& aServiceManager ) throw( RuntimeException ) {
    return Reference< XInterface >( *new XMLSignatureTemplateImpl( aServiceManager ) ) ;
}

Reference< XSingleServiceFactory > XMLSignatureTemplateImpl :: impl_createFactory( const Reference< XMultiServiceFactory >& aServiceManager ) {
    //Reference< XSingleServiceFactory > xFactory ;
    //xFactory = ::cppu::createSingleFactory( aServiceManager , impl_getImplementationName , impl_createInstance , impl_getSupportedServiceNames ) ;
    //return xFactory ;
    return ::cppu::createSingleFactory( aServiceManager , impl_getImplementationName() , impl_createInstance , impl_getSupportedServiceNames() ) ;
}

<|endoftext|>


<commit_before>/*BEGIN_LEGAL
Intel Open Source License

Copyright (c) 2002-2013 Intel Corporation. All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.  Redistributions
in binary form must reproduce the above copyright notice, this list of
conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.  Neither the name of
the Intel Corporation nor the names of its contributors may be used to
endorse or promote products derived from this software without
specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE INTEL OR
ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
END_LEGAL */
#include <iostream>
#include <fstream>
#include "pin.H"

ofstream OutFile;

// The running count of instructions is kept here
// make it static to help the compiler optimize docount
static UINT64 ccount = 0;

// This function is called before every instruction is executed
VOID docount()
{
	ccount++;
}

VOID saveCall(ADDRINT nextip)
{
	OutFile << "Return set: 0x" << hex << nextip << endl;
}

VOID checkRet(ADDRINT *rsp, UINT32 framesize)
{
	ADDRINT retval;

	ADDRINT rspval = *rsp;
	ADDRINT *psp = (ADDRINT *)rspval;
	retval = *psp;
	OutFile << "Return to: 0x" << hex << retval << endl;
}

// Pin calls this function every time a call instruction is encountered
VOID Instruction(INS ins, VOID *v)
{
	// Insert a call to docount before every call instruction, no arguments are passed
	if (INS_IsCall(ins))
	{
		INS_InsertCall(ins, IPOINT_BEFORE,
			(AFUNPTR)docount,
			IARG_END);
		INS_InsertCall(ins, IPOINT_BEFORE,
			AFUNPTR(saveCall),
			IARG_ADDRINT, INS_NextAddress(ins),
			IARG_END);
	}
	else if (INS_IsRet(ins))
	{
		UINT64 imm = 0;
		if (INS_OperandCount(ins) > 0 && INS_OperandIsImmediate(ins, 0))
			imm = INS_OperandImmediate(ins, 0);

		INS_InsertCall(ins, IPOINT_BEFORE,
			AFUNPTR(checkRet),
			IARG_CALL_ORDER, CALL_ORDER_FIRST,
			IARG_REG_REFERENCE, REG_STACK_PTR,
			IARG_ADDRINT, (ADDRINT)imm,
			IARG_END);
	}
}

KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool",
    "o", "callcount.out", "specify output file name");

// This function is called when the application exits
VOID Fini(INT32 code, VOID *v)
{
    // Write to a file since cout and cerr maybe closed by the application
    OutFile.setf(ios::showbase);
    OutFile << "Count " << dec << ccount << endl;
    OutFile.close();
}

/* ===================================================================== */
/* Print Help Message                                                    */
/* ===================================================================== */

INT32 Usage()
{
    cerr << "This tool counts the number of dynamic instructions executed" << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

/* ===================================================================== */
/* Main                                                                  */
/* ===================================================================== */
/*   argc, argv are the entire command line: pin -t <toolname> -- ...    */
/* ===================================================================== */

int main(int argc, char * argv[])
{
    // Initialize pin
    if (PIN_Init(argc, argv)) return Usage();

    OutFile.open(KnobOutputFile.Value().c_str());

    // Register Instruction to be called to instrument instructions
    INS_AddInstrumentFunction(Instruction, 0);

    // Register Fini to be called when the application exits
    PIN_AddFiniFunction(Fini, 0);

    // Start the program, never returns
    PIN_StartProgram();

    return 0;
}
<commit_msg>Saves and checks return addresses<commit_after>/*BEGIN_LEGAL
Intel Open Source License

Copyright (c) 2002-2013 Intel Corporation. All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.  Redistributions
in binary form must reproduce the above copyright notice, this list of
conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.  Neither the name of
the Intel Corporation nor the names of its contributors may be used to
endorse or promote products derived from this software without
specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE INTEL OR
ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
END_LEGAL */
#include <iostream>
#include <fstream>
#include "pin.H"
#include <stack>          // std::stack
#include <cstdlib>		  // std::exit

std::stack<ADDRINT> addrStack;
ofstream OutFile;

// The running count of instructions is kept here
// make it static to help the compiler optimize docount
static UINT64 ccount = 0;

// This function is called before every instruction is executed
VOID docount()
{
	ccount++;
}

VOID saveCall(ADDRINT nextip)
{
	addrStack.push(nextip+1);
	OutFile << "Return set: 0x" << hex << nextip << endl;
}

VOID checkRet(ADDRINT *rsp)
{
	ADDRINT retval;

	ADDRINT rspval = *rsp;
	ADDRINT *psp = (ADDRINT *)rspval;
	retval = *psp;
	ADDRINT originval = addrStack.top();
	if(!addrStack.empty()){
		addrStack.pop();
	}
	else{
		cerr << "ERROR: TOO MANY RETURNS DETECTED." << endl;
		std::exit(EXIT_FAILURE);
	}
	if(originval != retval){
		cerr << "ERROR: STACK SMASHING DETECTED: expected target 0x" << hex << originval << ", actual return target 0x" << retval << endl;
		std::exit(EXIT_FAILURE);
	}
	OutFile << "Return to: 0x" << hex << retval << endl;
}

// Pin calls this function every time a call instruction is encountered
VOID Instruction(INS ins, VOID *v)
{
	// Insert a call to docount before every call instruction, no arguments are passed
	if (INS_IsCall(ins))
	{
		INS_InsertCall(ins, IPOINT_BEFORE,
			(AFUNPTR)docount,
			IARG_END);
		INS_InsertCall(ins, IPOINT_BEFORE,
			AFUNPTR(saveCall),
			IARG_ADDRINT, INS_NextAddress(ins),
			IARG_END);
	}
	else if (INS_IsRet(ins))
	{
		INS_InsertCall(ins, IPOINT_BEFORE,
			AFUNPTR(checkRet),
			IARG_CALL_ORDER, CALL_ORDER_FIRST,
			IARG_REG_REFERENCE, REG_STACK_PTR,
			IARG_END);
	}
}

KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool",
    "o", "callcount.out", "specify output file name");

// This function is called when the application exits
VOID Fini(INT32 code, VOID *v)
{
    // Write to a file since cout and cerr maybe closed by the application
    OutFile.setf(ios::showbase);
    OutFile << "Count " << dec << ccount << endl;
    OutFile.close();
}

/* ===================================================================== */
/* Print Help Message                                                    */
/* ===================================================================== */

INT32 Usage()
{
    cerr << "This tool counts the number of dynamic instructions executed" << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

/* ===================================================================== */
/* Main                                                                  */
/* ===================================================================== */
/*   argc, argv are the entire command line: pin -t <toolname> -- ...    */
/* ===================================================================== */

int main(int argc, char * argv[])
{
    // Initialize pin
    if (PIN_Init(argc, argv)) return Usage();

    OutFile.open(KnobOutputFile.Value().c_str());

    // Register Instruction to be called to instrument instructions
    INS_AddInstrumentFunction(Instruction, 0);

    // Register Fini to be called when the application exits
    PIN_AddFiniFunction(Fini, 0);

    // Start the program, never returns
    PIN_StartProgram();

    return 0;
}
<|endoftext|>