<commit_before>/**************************************************************************
**
** This file is part of Qt Creator
**
** Copyright (c) 2011 Nokia Corporation and/or its subsidiary(-ies).
**
** Contact: Nokia Corporation (qt-info@nokia.com)
**
** No Commercial Usage
**
** This file contains pre-release code and may not be distributed.
** You may use this file in accordance with the terms and conditions
** contained in the Technology Preview License Agreement accompanying
** this package.
**
** GNU Lesser General Public License Usage
**
** Alternatively, this file may be used under the terms of the GNU Lesser
** General Public License version 2.1 as published by the Free Software
** Foundation and appearing in the file LICENSE.LGPL included in the
** packaging of this file.  Please review the following information to
** ensure the GNU Lesser General Public License version 2.1 requirements
** will be met: http://www.gnu.org/licenses/old-licenses/lgpl-2.1.html.
**
** In addition, as a special exception, Nokia gives you certain additional
** rights.  These rights are described in the Nokia Qt LGPL Exception
** version 1.1, included in the file LGPL_EXCEPTION.txt in this package.
**
** If you have questions regarding the use of this file, please contact
** Nokia at qt-info@nokia.com.
**
**************************************************************************/

#include "fancymainwindow.h"

#include <QtCore/QList>
#include <QtCore/QHash>

#include <QtGui/QAction>
#include <QtGui/QContextMenuEvent>
#include <QtGui/QMenu>
#include <QtGui/QDockWidget>
#include <QtCore/QSettings>

static const char lockedKeyC[] = "Locked";
static const char stateKeyC[] = "State";
static const int settingsVersion = 2;

namespace Utils {

struct FancyMainWindowPrivate {
    explicit FancyMainWindowPrivate(FancyMainWindow *q);

    QList<QDockWidget *> m_dockWidgets;
    QList<bool> m_dockWidgetActiveState;
    bool m_locked;
    bool m_handleDockVisibilityChanges; //todo

    QAction *m_menuSeparator1;
    QAction *m_toggleLockedAction;
    QAction *m_menuSeparator2;
    QAction *m_resetLayoutAction;
    QDockWidget *m_toolBarDockWidget;
};

FancyMainWindowPrivate::FancyMainWindowPrivate(FancyMainWindow *q) :
    m_locked(true), m_handleDockVisibilityChanges(true),
    m_menuSeparator1(new QAction(q)),
    m_toggleLockedAction(new QAction(FancyMainWindow::tr("Locked"), q)),
    m_menuSeparator2(new QAction(q)),
    m_resetLayoutAction(new QAction(FancyMainWindow::tr("Reset to Default Layout") ,q)),
    m_toolBarDockWidget(0)
{
    m_toggleLockedAction->setCheckable(true);
    m_toggleLockedAction->setChecked(m_locked);
    m_menuSeparator1->setSeparator(true);
    m_menuSeparator2->setSeparator(true);
}

FancyMainWindow::FancyMainWindow(QWidget *parent) :
    QMainWindow(parent), d(new FancyMainWindowPrivate(this))
{
    connect(d->m_toggleLockedAction, SIGNAL(toggled(bool)),
            this, SLOT(setLocked(bool)));
    connect(d->m_resetLayoutAction, SIGNAL(triggered()),
            this, SIGNAL(resetLayout()));
}

FancyMainWindow::~FancyMainWindow()
{
    delete d;
}

QDockWidget *FancyMainWindow::addDockForWidget(QWidget *widget)
{
    QDockWidget *dockWidget = new QDockWidget(widget->windowTitle(), this);
    dockWidget->setWidget(widget);
    // Set an object name to be used in settings, derive from widget name
    const QString objectName = widget->objectName();
    if (objectName.isEmpty()) {
        dockWidget->setObjectName(QLatin1String("dockWidget") + QString::number(d->m_dockWidgets.size() + 1));
    } else {
        dockWidget->setObjectName(objectName + QLatin1String("DockWidget"));
    }
    connect(dockWidget->toggleViewAction(), SIGNAL(triggered()),
        this, SLOT(onDockActionTriggered()), Qt::QueuedConnection);
    connect(dockWidget, SIGNAL(visibilityChanged(bool)),
            this, SLOT(onDockVisibilityChange(bool)));
    connect(dockWidget, SIGNAL(topLevelChanged(bool)),
            this, SLOT(onTopLevelChanged()));
    d->m_dockWidgets.append(dockWidget);
    d->m_dockWidgetActiveState.append(true);
    updateDockWidget(dockWidget);
    return dockWidget;
}

void FancyMainWindow::updateDockWidget(QDockWidget *dockWidget)
{
    const QDockWidget::DockWidgetFeatures features =
            (d->m_locked) ? QDockWidget::DockWidgetClosable | QDockWidget::DockWidgetFloatable
                       : QDockWidget::DockWidgetMovable | QDockWidget::DockWidgetClosable | QDockWidget::DockWidgetFloatable;
    QWidget *titleBarWidget = dockWidget->titleBarWidget();
    if (d->m_locked && !titleBarWidget && !dockWidget->isFloating())
        titleBarWidget = new QWidget(dockWidget);
    else if ((!d->m_locked || dockWidget->isFloating()) && titleBarWidget) {
        delete titleBarWidget;
        titleBarWidget = 0;
    }
    dockWidget->setTitleBarWidget(titleBarWidget);
    dockWidget->setFeatures(features);
}

void FancyMainWindow::onDockActionTriggered()
{
    QDockWidget *dw = qobject_cast<QDockWidget *>(sender()->parent());
    if (dw) {
        if (dw->isVisible())
            dw->raise();
    }
}

void FancyMainWindow::onDockVisibilityChange(bool visible)
{
    if (!d->m_handleDockVisibilityChanges)
        return;
    QDockWidget *dockWidget = qobject_cast<QDockWidget *>(sender());
    int index = d->m_dockWidgets.indexOf(dockWidget);
    d->m_dockWidgetActiveState[index] = visible;
}

void FancyMainWindow::onTopLevelChanged()
{
    updateDockWidget(qobject_cast<QDockWidget*>(sender()));
}

void FancyMainWindow::setTrackingEnabled(bool enabled)
{
    if (enabled) {
        d->m_handleDockVisibilityChanges = true;
        for (int i = 0; i < d->m_dockWidgets.size(); ++i)
            d->m_dockWidgetActiveState[i] = d->m_dockWidgets[i]->isVisible();
    } else {
        d->m_handleDockVisibilityChanges = false;
    }
}

void FancyMainWindow::setLocked(bool locked)
{
    d->m_locked = locked;
    foreach (QDockWidget *dockWidget, d->m_dockWidgets) {
        updateDockWidget(dockWidget);
    }
}

void FancyMainWindow::hideEvent(QHideEvent *event)
{
    Q_UNUSED(event)
    handleVisibilityChanged(false);
}

void FancyMainWindow::showEvent(QShowEvent *event)
{
    Q_UNUSED(event)
    handleVisibilityChanged(true);
}

void FancyMainWindow::contextMenuEvent(QContextMenuEvent *event)
{
    QMenu *menu = createPopupMenu();
    menu->exec(event->globalPos());
    delete menu;
}

void FancyMainWindow::handleVisibilityChanged(bool visible)
{
    d->m_handleDockVisibilityChanges = false;
    for (int i = 0; i < d->m_dockWidgets.size(); ++i) {
        QDockWidget *dockWidget = d->m_dockWidgets.at(i);
        if (dockWidget->isFloating()) {
            dockWidget->setVisible(visible && d->m_dockWidgetActiveState.at(i));
        }
    }
    if (visible)
        d->m_handleDockVisibilityChanges = true;
}

void FancyMainWindow::saveSettings(QSettings *settings) const
{
    QHash<QString, QVariant> hash = saveSettings();
    QHashIterator<QString, QVariant> it(hash);
    while (it.hasNext()) {
        it.next();
        settings->setValue(it.key(), it.value());
    }
}

void FancyMainWindow::restoreSettings(const QSettings *settings)
{
    QHash<QString, QVariant> hash;
    foreach (const QString &key, settings->childKeys()) {
        hash.insert(key, settings->value(key));
    }
    restoreSettings(hash);
}

QHash<QString, QVariant> FancyMainWindow::saveSettings() const
{
    QHash<QString, QVariant> settings;
    settings.insert(QLatin1String(stateKeyC), saveState(settingsVersion));
    settings.insert(QLatin1String(lockedKeyC), d->m_locked);
    for (int i = 0; i < d->m_dockWidgetActiveState.count(); ++i) {
        settings.insert(d->m_dockWidgets.at(i)->objectName(),
                           d->m_dockWidgetActiveState.at(i));
    }
    return settings;
}

void FancyMainWindow::restoreSettings(const QHash<QString, QVariant> &settings)
{
    QByteArray ba = settings.value(QLatin1String(stateKeyC), QByteArray()).toByteArray();
    if (!ba.isEmpty())
        restoreState(ba, settingsVersion);
    d->m_locked = settings.value(QLatin1String("Locked"), true).toBool();
    d->m_toggleLockedAction->setChecked(d->m_locked);
    for (int i = 0; i < d->m_dockWidgetActiveState.count(); ++i) {
        d->m_dockWidgetActiveState[i] = settings.value(d->m_dockWidgets.at(i)->objectName(), false).toBool();
    }
}

QList<QDockWidget *> FancyMainWindow::dockWidgets() const
{
    return d->m_dockWidgets;
}

bool FancyMainWindow::isLocked() const
{
    return d->m_locked;
}

QMenu *FancyMainWindow::createPopupMenu()
{
    QMenu *menu = QMainWindow::createPopupMenu();
    menu->addAction(d->m_menuSeparator1);
    menu->addAction(d->m_toggleLockedAction);
    menu->addAction(d->m_menuSeparator2);
    menu->addAction(d->m_resetLayoutAction);
    return menu;
}

QAction *FancyMainWindow::menuSeparator1() const
{
    return d->m_menuSeparator1;
}

QAction *FancyMainWindow::toggleLockedAction() const
{
    return d->m_toggleLockedAction;
}

QAction *FancyMainWindow::menuSeparator2() const
{
    return d->m_menuSeparator2;
}

QAction *FancyMainWindow::resetLayoutAction() const
{
    return d->m_resetLayoutAction;
}

void FancyMainWindow::setDockActionsVisible(bool v)
{
    foreach(const QDockWidget *dockWidget, d->m_dockWidgets)
        dockWidget->toggleViewAction()->setVisible(v);
    d->m_toggleLockedAction->setVisible(v);
    d->m_menuSeparator1->setVisible(v);
    d->m_menuSeparator2->setVisible(v);
    d->m_resetLayoutAction->setVisible(v);
}

QDockWidget *FancyMainWindow::toolBarDockWidget() const
{
    return d->m_toolBarDockWidget;
}

void FancyMainWindow::setToolBarDockWidget(QDockWidget *dock)
{
    d->m_toolBarDockWidget = dock;
}

} // namespace Utils
<commit_msg>fancymainwindow: remove m_dockWidgets member<commit_after>/**************************************************************************
**
** This file is part of Qt Creator
**
** Copyright (c) 2011 Nokia Corporation and/or its subsidiary(-ies).
**
** Contact: Nokia Corporation (qt-info@nokia.com)
**
** No Commercial Usage
**
** This file contains pre-release code and may not be distributed.
** You may use this file in accordance with the terms and conditions
** contained in the Technology Preview License Agreement accompanying
** this package.
**
** GNU Lesser General Public License Usage
**
** Alternatively, this file may be used under the terms of the GNU Lesser
** General Public License version 2.1 as published by the Free Software
** Foundation and appearing in the file LICENSE.LGPL included in the
** packaging of this file.  Please review the following information to
** ensure the GNU Lesser General Public License version 2.1 requirements
** will be met: http://www.gnu.org/licenses/old-licenses/lgpl-2.1.html.
**
** In addition, as a special exception, Nokia gives you certain additional
** rights.  These rights are described in the Nokia Qt LGPL Exception
** version 1.1, included in the file LGPL_EXCEPTION.txt in this package.
**
** If you have questions regarding the use of this file, please contact
** Nokia at qt-info@nokia.com.
**
**************************************************************************/

#include "fancymainwindow.h"

#include <QtCore/QList>
#include <QtCore/QHash>

#include <QtGui/QAction>
#include <QtGui/QContextMenuEvent>
#include <QtGui/QMenu>
#include <QtGui/QDockWidget>
#include <QtCore/QSettings>

static const char lockedKeyC[] = "Locked";
static const char stateKeyC[] = "State";
static const int settingsVersion = 2;

namespace Utils {

struct FancyMainWindowPrivate {
    explicit FancyMainWindowPrivate(FancyMainWindow *q);

    QMap<QDockWidget *, bool> m_dockWidgetActiveState;
    bool m_locked;
    bool m_handleDockVisibilityChanges; //todo

    QAction *m_menuSeparator1;
    QAction *m_toggleLockedAction;
    QAction *m_menuSeparator2;
    QAction *m_resetLayoutAction;
    QDockWidget *m_toolBarDockWidget;
};

FancyMainWindowPrivate::FancyMainWindowPrivate(FancyMainWindow *q) :
    m_locked(true), m_handleDockVisibilityChanges(true),
    m_menuSeparator1(new QAction(q)),
    m_toggleLockedAction(new QAction(FancyMainWindow::tr("Locked"), q)),
    m_menuSeparator2(new QAction(q)),
    m_resetLayoutAction(new QAction(FancyMainWindow::tr("Reset to Default Layout") ,q)),
    m_toolBarDockWidget(0)
{
    m_toggleLockedAction->setCheckable(true);
    m_toggleLockedAction->setChecked(m_locked);
    m_menuSeparator1->setSeparator(true);
    m_menuSeparator2->setSeparator(true);
}

FancyMainWindow::FancyMainWindow(QWidget *parent) :
    QMainWindow(parent), d(new FancyMainWindowPrivate(this))
{
    connect(d->m_toggleLockedAction, SIGNAL(toggled(bool)),
            this, SLOT(setLocked(bool)));
    connect(d->m_resetLayoutAction, SIGNAL(triggered()),
            this, SIGNAL(resetLayout()));
}

FancyMainWindow::~FancyMainWindow()
{
    delete d;
}

QDockWidget *FancyMainWindow::addDockForWidget(QWidget *widget)
{
    QDockWidget *dockWidget = new QDockWidget(widget->windowTitle(), this);
    dockWidget->setWidget(widget);
    // Set an object name to be used in settings, derive from widget name
    const QString objectName = widget->objectName();
    if (objectName.isEmpty()) {
        dockWidget->setObjectName(QLatin1String("dockWidget") + QString::number(dockWidgets().size() + 1));
    } else {
        dockWidget->setObjectName(objectName + QLatin1String("DockWidget"));
    }
    connect(dockWidget->toggleViewAction(), SIGNAL(triggered()),
        this, SLOT(onDockActionTriggered()), Qt::QueuedConnection);
    connect(dockWidget, SIGNAL(visibilityChanged(bool)),
            this, SLOT(onDockVisibilityChange(bool)));
    connect(dockWidget, SIGNAL(topLevelChanged(bool)),
            this, SLOT(onTopLevelChanged()));
    d->m_dockWidgetActiveState[dockWidget] = true;
    updateDockWidget(dockWidget);
    return dockWidget;
}

void FancyMainWindow::updateDockWidget(QDockWidget *dockWidget)
{
    const QDockWidget::DockWidgetFeatures features =
            (d->m_locked) ? QDockWidget::DockWidgetClosable | QDockWidget::DockWidgetFloatable
                       : QDockWidget::DockWidgetMovable | QDockWidget::DockWidgetClosable | QDockWidget::DockWidgetFloatable;
    QWidget *titleBarWidget = dockWidget->titleBarWidget();
    if (d->m_locked && !titleBarWidget && !dockWidget->isFloating())
        titleBarWidget = new QWidget(dockWidget);
    else if ((!d->m_locked || dockWidget->isFloating()) && titleBarWidget) {
        delete titleBarWidget;
        titleBarWidget = 0;
    }
    dockWidget->setTitleBarWidget(titleBarWidget);
    dockWidget->setFeatures(features);
}

void FancyMainWindow::onDockActionTriggered()
{
    QDockWidget *dw = qobject_cast<QDockWidget *>(sender()->parent());
    if (dw) {
        if (dw->isVisible())
            dw->raise();
    }
}

void FancyMainWindow::onDockVisibilityChange(bool visible)
{
    if (!d->m_handleDockVisibilityChanges)
        return;
    QDockWidget *dockWidget = qobject_cast<QDockWidget *>(sender());
    d->m_dockWidgetActiveState[dockWidget] = visible;
}

void FancyMainWindow::onTopLevelChanged()
{
    updateDockWidget(qobject_cast<QDockWidget*>(sender()));
}

void FancyMainWindow::setTrackingEnabled(bool enabled)
{
    if (enabled) {
        d->m_handleDockVisibilityChanges = true;
        foreach(QDockWidget *dockWidget, dockWidgets())
            d->m_dockWidgetActiveState[dockWidget] = dockWidget->isVisible();
    } else {
        d->m_handleDockVisibilityChanges = false;
    }
}

void FancyMainWindow::setLocked(bool locked)
{
    d->m_locked = locked;
    foreach (QDockWidget *dockWidget, dockWidgets()) {
        updateDockWidget(dockWidget);
    }
}

void FancyMainWindow::hideEvent(QHideEvent *event)
{
    Q_UNUSED(event)
    handleVisibilityChanged(false);
}

void FancyMainWindow::showEvent(QShowEvent *event)
{
    Q_UNUSED(event)
    handleVisibilityChanged(true);
}

void FancyMainWindow::contextMenuEvent(QContextMenuEvent *event)
{
    QMenu *menu = createPopupMenu();
    menu->exec(event->globalPos());
    delete menu;
}

void FancyMainWindow::handleVisibilityChanged(bool visible)
{
    d->m_handleDockVisibilityChanges = false;
    foreach(QDockWidget *dockWidget, dockWidgets()) {
        if (dockWidget->isFloating()) {
            dockWidget->setVisible(visible && d->m_dockWidgetActiveState.value(dockWidget));
        }
    }
    if (visible)
        d->m_handleDockVisibilityChanges = true;
}

void FancyMainWindow::saveSettings(QSettings *settings) const
{
    QHash<QString, QVariant> hash = saveSettings();
    QHashIterator<QString, QVariant> it(hash);
    while (it.hasNext()) {
        it.next();
        settings->setValue(it.key(), it.value());
    }
}

void FancyMainWindow::restoreSettings(const QSettings *settings)
{
    QHash<QString, QVariant> hash;
    foreach (const QString &key, settings->childKeys()) {
        hash.insert(key, settings->value(key));
    }
    restoreSettings(hash);
}

QHash<QString, QVariant> FancyMainWindow::saveSettings() const
{
    QHash<QString, QVariant> settings;
    settings.insert(QLatin1String(stateKeyC), saveState(settingsVersion));
    settings.insert(QLatin1String(lockedKeyC), d->m_locked);
    foreach(QDockWidget *dockWidget, dockWidgets()) {
        settings.insert(dockWidget->objectName(),
                        d->m_dockWidgetActiveState.value(dockWidget));
    }
    return settings;
}

void FancyMainWindow::restoreSettings(const QHash<QString, QVariant> &settings)
{
    QByteArray ba = settings.value(QLatin1String(stateKeyC), QByteArray()).toByteArray();
    if (!ba.isEmpty())
        restoreState(ba, settingsVersion);
    d->m_locked = settings.value(QLatin1String("Locked"), true).toBool();
    d->m_toggleLockedAction->setChecked(d->m_locked);
    foreach(QDockWidget *widget, dockWidgets()) {
        d->m_dockWidgetActiveState[widget] = settings.value(widget->objectName(), false).toBool();
    }
}

QList<QDockWidget *> FancyMainWindow::dockWidgets() const
{
    return qFindChildren<QDockWidget *>(this);
}

bool FancyMainWindow::isLocked() const
{
    return d->m_locked;
}

QMenu *FancyMainWindow::createPopupMenu()
{
    QMenu *menu = QMainWindow::createPopupMenu();
    menu->addAction(d->m_menuSeparator1);
    menu->addAction(d->m_toggleLockedAction);
    menu->addAction(d->m_menuSeparator2);
    menu->addAction(d->m_resetLayoutAction);
    return menu;
}

QAction *FancyMainWindow::menuSeparator1() const
{
    return d->m_menuSeparator1;
}

QAction *FancyMainWindow::toggleLockedAction() const
{
    return d->m_toggleLockedAction;
}

QAction *FancyMainWindow::menuSeparator2() const
{
    return d->m_menuSeparator2;
}

QAction *FancyMainWindow::resetLayoutAction() const
{
    return d->m_resetLayoutAction;
}

void FancyMainWindow::setDockActionsVisible(bool v)
{
    foreach(const QDockWidget *dockWidget, dockWidgets())
        dockWidget->toggleViewAction()->setVisible(v);
    d->m_toggleLockedAction->setVisible(v);
    d->m_menuSeparator1->setVisible(v);
    d->m_menuSeparator2->setVisible(v);
    d->m_resetLayoutAction->setVisible(v);
}

QDockWidget *FancyMainWindow::toolBarDockWidget() const
{
    return d->m_toolBarDockWidget;
}

void FancyMainWindow::setToolBarDockWidget(QDockWidget *dock)
{
    d->m_toolBarDockWidget = dock;
}

} // namespace Utils
<|endoftext|>


<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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

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

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

#include "tensorflow/compiler/xla/service/gpu/for_thunk.h"

#include "tensorflow/compiler/xla/ptr_util.h"
#include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/lib/core/errors.h"

namespace xla {
namespace gpu {

ForThunk::ForThunk(const int64 loop_limit,
                   std::unique_ptr<ThunkSequence> body_thunk_sequence,
                   const HloInstruction* hlo)
    : Thunk(Kind::kWhile, hlo),
      loop_limit_(loop_limit),
      body_thunk_sequence_(MakeUnique<SequentialThunk>(
          // Pass nullptr as the HloInstruction* to the body_thunk_sequence_
          // constructor because this SequentialThunk is logically "part of"
          // this ForThunk, and shouldn't be profiled separately from it.
          std::move(*body_thunk_sequence), nullptr)) {}

Status ForThunk::Initialize(const GpuExecutable& executable,
                            se::StreamExecutor* executor) {
  TF_RETURN_IF_ERROR(body_thunk_sequence_->Initialize(executable, executor));
  return Status::OK();
}

Status ForThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations,
                                 se::Stream* stream,
                                 HloExecutionProfiler* profiler) {
  auto op_profiler = profiler->MakeScopedInstructionProfiler(hlo_instruction());
  for (int64 i = 0; i < loop_limit_; ++i) {
    profiler->StartHloComputation();
    // Invoke loop body thunk sequence.
    TF_RETURN_IF_ERROR(body_thunk_sequence_->ExecuteOnStream(buffer_allocations,
                                                             stream, profiler));
    profiler->FinishHloComputation(hlo_instruction()->while_body());
  }
  return Status::OK();
}

}  // namespace gpu
}  // namespace xla
<commit_msg>[XLA:GPU] Add VLOG to ForThunk.<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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

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

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

#include "tensorflow/compiler/xla/service/gpu/for_thunk.h"

#include "tensorflow/compiler/xla/ptr_util.h"
#include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/lib/core/errors.h"

namespace xla {
namespace gpu {

ForThunk::ForThunk(const int64 loop_limit,
                   std::unique_ptr<ThunkSequence> body_thunk_sequence,
                   const HloInstruction* hlo)
    : Thunk(Kind::kWhile, hlo),
      loop_limit_(loop_limit),
      body_thunk_sequence_(MakeUnique<SequentialThunk>(
          // Pass nullptr as the HloInstruction* to the body_thunk_sequence_
          // constructor because this SequentialThunk is logically "part of"
          // this ForThunk, and shouldn't be profiled separately from it.
          std::move(*body_thunk_sequence), nullptr)) {}

Status ForThunk::Initialize(const GpuExecutable& executable,
                            se::StreamExecutor* executor) {
  TF_RETURN_IF_ERROR(body_thunk_sequence_->Initialize(executable, executor));
  return Status::OK();
}

Status ForThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations,
                                 se::Stream* stream,
                                 HloExecutionProfiler* profiler) {
  VLOG(2) << "Executing ForThunk with " << loop_limit_ << " iters for "
          << (hlo_instruction() ? hlo_instruction()->ToString() : "<null>");
  auto op_profiler = profiler->MakeScopedInstructionProfiler(hlo_instruction());
  for (int64 i = 0; i < loop_limit_; ++i) {
    profiler->StartHloComputation();
    // Invoke loop body thunk sequence.
    TF_RETURN_IF_ERROR(body_thunk_sequence_->ExecuteOnStream(buffer_allocations,
                                                             stream, profiler));
    profiler->FinishHloComputation(hlo_instruction()->while_body());
  }
  return Status::OK();
}

}  // namespace gpu
}  // namespace xla
<|endoftext|>


<commit_before>#include "testparser.h"

#include "astprinter.h"
#include "lexer.h"
#include "parser.h"

void TestParser::testExamples()
{
    QDir examples(QCoreApplication::applicationDirPath() + "/../../examples");
    QVERIFY(examples.exists());
    QFileInfoList unvFiles = examples.entryInfoList(QStringList("*.unv"));
    foreach (QFileInfo f, unvFiles) {
        QFile file(f.filePath());
        QVERIFY(file.exists());
        QVERIFY(file.open(QFile::ReadOnly));
        QTextStream inFile(&file);
        QString fileText = inFile.readAll();
        file.close();

        SourceBuffer buffer(fileText, file.fileName());
        Lexer lexer;
        lexer.lex(&buffer);

        Parser parser;
        parser.parse(&buffer);

        QString astText;
        QTextStream out(&astText);
        ASTPrinter printer(&buffer, &out);
        printer.walk();

        QFile resource(":resources/" + f.baseName() + ".ast");
        QVERIFY(resource.exists());
        QVERIFY(resource.open(QFile::ReadOnly));
        QTextStream inResource(&resource);
        QString resourceText = inResource.readAll();
        resource.close();

        QCOMPARE(resourceText, astText);
    }
}
<commit_msg>Don't fail on examples where the AST has not been saved to a resource file yet<commit_after>#include "testparser.h"

#include "astprinter.h"
#include "lexer.h"
#include "parser.h"

void TestParser::testExamples()
{
    QDir examples(QCoreApplication::applicationDirPath() + "/../../examples");
    QVERIFY(examples.exists());
    QFileInfoList unvFiles = examples.entryInfoList(QStringList("*.unv"));
    foreach (QFileInfo f, unvFiles) {
        QFile file(f.filePath());
        QVERIFY(file.exists());
        QVERIFY(file.open(QFile::ReadOnly));
        QTextStream inFile(&file);
        QString fileText = inFile.readAll();
        file.close();

        SourceBuffer buffer(fileText, file.fileName());
        Lexer lexer;
        lexer.lex(&buffer);

        Parser parser;
        parser.parse(&buffer);

        QString astText;
        QTextStream out(&astText);
        ASTPrinter printer(&buffer, &out);
        printer.walk();

        QFile resource(":resources/" + f.baseName() + ".ast");
        if (!resource.exists())
            continue;

        QVERIFY(resource.open(QFile::ReadOnly));
        QTextStream inResource(&resource);
        QString resourceText = inResource.readAll();
        resource.close();

        QCOMPARE(resourceText, astText);
    }
}
<|endoftext|>


<commit_before>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

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

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/dataset.h"
#include "tensorflow/core/util/batch_util.h"

namespace tensorflow {

namespace {

// See documentation in ../ops/dataset_ops.cc for a high-level
// description of the following op.

class SlideDatasetOp : public UnaryDatasetOpKernel {
 public:
  explicit SlideDatasetOp(OpKernelConstruction* ctx)
      : UnaryDatasetOpKernel(ctx) {}

  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                   DatasetBase** output) override {
    int64 window_size = 0;
    int64 stride = 1;
    OP_REQUIRES_OK(
        ctx, ParseScalarArgument<int64>(ctx, "window_size", &window_size));
    OP_REQUIRES_OK(ctx, ParseScalarArgument<int64>(ctx, "stride", &stride));
    OP_REQUIRES(
        ctx, window_size > 0,
        errors::InvalidArgument("Window size must be greater than zero."));
    OP_REQUIRES(
        ctx, stride > 0 && stride < window_size,
        errors::InvalidArgument("Stride must be in [1, window_size)."));

    *output = new Dataset(ctx, window_size, stride, input);
  }

 private:
  class Dataset : public GraphDatasetBase {
   public:
    Dataset(OpKernelContext* ctx, int64 window_size, int64 stride,
            const DatasetBase* input)
        : GraphDatasetBase(ctx),
          window_size_(window_size),
          stride_(stride),
          input_(input) {
      input_->Ref();

      const auto& input_shapes = input_->output_shapes();
      output_shapes_.reserve(input_shapes.size());
      for (const auto& input_shape : input_shapes) {
        output_shapes_.emplace_back(
            PartialTensorShape({-1}).Concatenate(input_shape));
      }
    }

    ~Dataset() override { input_->Unref(); }

    std::unique_ptr<IteratorBase> MakeIteratorInternal(
        const string& prefix) const override {
      return std::unique_ptr<IteratorBase>(new Iterator(
          Iterator::Params{this, strings::StrCat(prefix, "::Slide")}));
    }

    const DataTypeVector& output_dtypes() const override {
      return input_->output_dtypes();
    }

    const std::vector<PartialTensorShape>& output_shapes() const override {
      return output_shapes_;
    }

    string DebugString() const override {
      return strings::StrCat("SlideDatasetOp(", window_size_, ", ", stride_,
                             ")::Dataset");
    }

   protected:
    Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b,
                              Node** output) const override {
      Node* input_graph_node = nullptr;
      TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_graph_node));
      Node* window_size = nullptr;
      Node* stride = nullptr;
      TF_RETURN_IF_ERROR(b->AddScalar(window_size_, &window_size));
      TF_RETURN_IF_ERROR(b->AddScalar(stride_, &stride));
      TF_RETURN_IF_ERROR(
          b->AddDataset(this, {input_graph_node, window_size, stride}, output));
      return Status::OK();
    }

   private:

    class Iterator : public DatasetIterator<Dataset> {
     public:
      explicit Iterator(const Params& params)
          : DatasetIterator<Dataset>(params) {}

      Status Initialize(IteratorContext* ctx) override {
        return dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_);
      }

      Status GetNextInternal(IteratorContext* ctx,
                             std::vector<Tensor>* out_tensors,
                             bool* end_of_sequence) override {
        const int64 window_size = dataset()->window_size_;
        const int64 stride = dataset()->stride_;
        std::vector<std::vector<Tensor>> batch_elements;
        {
          mutex_lock l(mu_);
          if (!input_impl_) {
            *end_of_sequence = true;
            return Status::OK();
          }
          batch_elements.reserve(window_size);
          const bool first_call = cache_.empty();
          if (first_call) {
            cache_.reserve(window_size);
          } else {
            // Reuse cache in the previous iteration.
            cache_.swap(batch_elements);
          }
          // Fill up with new elements.
          *end_of_sequence = false;
          for (size_t i = batch_elements.size(); i < window_size && !*end_of_sequence;
              ++i) {
            std::vector<Tensor> batch_element_tuple;
            TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, &batch_element_tuple,
                                                    end_of_sequence));
            if (!*end_of_sequence) {
              batch_elements.push_back(std::move(batch_element_tuple));
            } else {
              input_impl_.reset();
            }
          }
          // Drop the final smaller blocks.
          if (batch_elements.size() < window_size) {
            DCHECK(*end_of_sequence);
            return Status::OK();
          }
          // Cache the data used for the next iteration.
          for (size_t i = stride; i < window_size; ++i) {
            cache_.emplace_back(batch_elements[i]);
          }
        }

        // Construct output tensors.
        // Those codes below are copied from batch_dataset_op.cc.
        const size_t num_tuple_components = batch_elements[0].size();
        const int64 num_batch_elements = batch_elements.size();
        for (size_t component_index = 0; component_index < num_tuple_components;
             ++component_index) {
          const Tensor& first_element = batch_elements[0][component_index];
          TensorShape batch_component_shape({num_batch_elements});
          batch_component_shape.AppendShape(first_element.shape());
          Tensor batch_component(cpu_allocator(), first_element.dtype(),
                                 batch_component_shape);
          // Build the output tuple component by copying one slice
          // from each input element in the batch.
          for (size_t i = 0; i < num_batch_elements; ++i) {
            if (batch_elements[i][component_index].shape() !=
                first_element.shape()) {
              return errors::InvalidArgument(
                  "Cannot batch tensors with different shapes in component ",
                  component_index, ". First element had shape ",
                  first_element.shape().DebugString(), " and element ", i,
                  " had shape ",
                  batch_elements[i][component_index].shape().DebugString(),
                  ".");
            }
            TF_RETURN_IF_ERROR(batch_util::CopyElementToSlice(
                std::move(batch_elements[i][component_index]), &batch_component,
                i));
          }
          out_tensors->emplace_back(std::move(batch_component));
        }
        *end_of_sequence = false;
        return Status::OK();
      }

     protected:
      Status SaveInternal(IteratorStateWriter* writer) override {
        mutex_lock l(mu_);
        if (!input_impl_) {
          TF_RETURN_IF_ERROR(
              writer->WriteScalar(full_name("input_impl_empty"), ""));
        } else {
          TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_));
        }
        // Save cache.
        TF_RETURN_IF_ERROR(
            writer->WriteScalar(strings::StrCat("cache_size"), cache_.size()));
        for (int64 i = 0; i < cache_.size(); i++) {
          TF_RETURN_IF_ERROR(writer->WriteScalar(
              strings::StrCat("cache[", i, "]_size"), cache_[i].size()));
          for (int64 j = 0; j < cache_[i].size(); j++) {
            TF_RETURN_IF_ERROR(writer->WriteTensor(
                strings::StrCat("cache[", i, "][", j, "]"), cache_[i][j]));
          }
        }
        return Status::OK();
      }

      Status RestoreInternal(IteratorContext* ctx,
                             IteratorStateReader* reader) override {
        mutex_lock l(mu_);
        if (!reader->Contains(full_name("input_impl_empty"))) {
          TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_));
        } else {
          input_impl_.reset();
        }
        // Restore cache.
        int64 cache_size;
        TF_RETURN_IF_ERROR(
            reader->ReadScalar(strings::StrCat("cache_size"), &cache_size));
        cache_.resize(cache_size);
        for (int64 i = 0; i < cache_size; i++) {
          int64 vector_size;
          TF_RETURN_IF_ERROR(reader->ReadScalar(
              strings::StrCat("cache[", i, "]_size"), &vector_size));
          cache_[i].resize(vector_size);
          for (int64 j = 0; j < vector_size; j++) {
            TF_RETURN_IF_ERROR(reader->ReadTensor(
                strings::StrCat("cache[", i, "][", j, "]"), &cache_[i][j]));
          }
        }
        return Status::OK();
      }

     private:
      mutex mu_;
      std::vector<std::vector<Tensor>> cache_ GUARDED_BY(mu_);
      std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
    };

    const int64 window_size_;
    const int64 stride_;
    const DatasetBase* const input_;
    std::vector<PartialTensorShape> output_shapes_;
  };
};

REGISTER_KERNEL_BUILDER(Name("SlideDataset").Device(DEVICE_CPU),
                        SlideDatasetOp);

}  // namespace

}  // namespace tensorflow
<commit_msg>ENH: allow stride > window_size in c++ side<commit_after>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

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

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/dataset.h"
#include "tensorflow/core/util/batch_util.h"

namespace tensorflow {

namespace {

// See documentation in ../ops/dataset_ops.cc for a high-level
// description of the following op.

class SlideDatasetOp : public UnaryDatasetOpKernel {
 public:
  explicit SlideDatasetOp(OpKernelConstruction* ctx)
      : UnaryDatasetOpKernel(ctx) {}

  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                   DatasetBase** output) override {
    int64 window_size = 0;
    int64 stride = 0;
    OP_REQUIRES_OK(
        ctx, ParseScalarArgument<int64>(ctx, "window_size", &window_size));
    OP_REQUIRES_OK(ctx, ParseScalarArgument<int64>(ctx, "stride", &stride));
    OP_REQUIRES(
        ctx, window_size > 0,
        errors::InvalidArgument("Window size must be greater than zero."));
    OP_REQUIRES(
        ctx, stride > 0,
        errors::InvalidArgument("Stride must be greater than zero."));

    *output = new Dataset(ctx, window_size, stride, input);
  }

 private:
  class Dataset : public GraphDatasetBase {
   public:
    Dataset(OpKernelContext* ctx, int64 window_size, int64 stride,
            const DatasetBase* input)
        : GraphDatasetBase(ctx),
          window_size_(window_size),
          stride_(stride),
          input_(input) {
      input_->Ref();

      const auto& input_shapes = input_->output_shapes();
      output_shapes_.reserve(input_shapes.size());
      for (const auto& input_shape : input_shapes) {
        output_shapes_.emplace_back(
            PartialTensorShape({-1}).Concatenate(input_shape));
      }
    }

    ~Dataset() override { input_->Unref(); }

    std::unique_ptr<IteratorBase> MakeIteratorInternal(
        const string& prefix) const override {
      return std::unique_ptr<IteratorBase>(new Iterator(
          Iterator::Params{this, strings::StrCat(prefix, "::Slide")}));
    }

    const DataTypeVector& output_dtypes() const override {
      return input_->output_dtypes();
    }

    const std::vector<PartialTensorShape>& output_shapes() const override {
      return output_shapes_;
    }

    string DebugString() const override {
      return strings::StrCat("SlideDatasetOp(", window_size_, ", ", stride_,
                             ")::Dataset");
    }

   protected:
    Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b,
                              Node** output) const override {
      Node* input_graph_node = nullptr;
      TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_graph_node));
      Node* window_size = nullptr;
      Node* stride = nullptr;
      TF_RETURN_IF_ERROR(b->AddScalar(window_size_, &window_size));
      TF_RETURN_IF_ERROR(b->AddScalar(stride_, &stride));
      TF_RETURN_IF_ERROR(
          b->AddDataset(this, {input_graph_node, window_size, stride}, output));
      return Status::OK();
    }

   private:

    class Iterator : public DatasetIterator<Dataset> {
     public:
      explicit Iterator(const Params& params)
          : DatasetIterator<Dataset>(params) {}

      Status Initialize(IteratorContext* ctx) override {
        return dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_);
      }

      Status GetNextInternal(IteratorContext* ctx,
                             std::vector<Tensor>* out_tensors,
                             bool* end_of_sequence) override {
        const int64 window_size = dataset()->window_size_;
        const int64 stride = dataset()->stride_;
        std::vector<std::vector<Tensor>> batch_elements;
        {
          mutex_lock l(mu_);
          if (!input_impl_) {
            *end_of_sequence = true;
            return Status::OK();
          }
          batch_elements.reserve(window_size);
          // Use cache if stride < window_size.
          if (stride < window_size) {
            const bool first_call = cache_.empty();
            if (first_call) {
              cache_.reserve(window_size);
            } else {
              // Reuse cache in the previous iteration.
              cache_.swap(batch_elements);
            }
          }
          // Fill up with new elements.
          *end_of_sequence = false;
          for (size_t i = batch_elements.size(); i < window_size && !*end_of_sequence;
              ++i) {
            std::vector<Tensor> batch_element_tuple;
            TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, &batch_element_tuple,
                                                    end_of_sequence));
            if (!*end_of_sequence) {
              batch_elements.push_back(std::move(batch_element_tuple));
            } else {
              input_impl_.reset();
            }
          }
          // Drop the final smaller blocks.
          if (batch_elements.size() < window_size) {
            DCHECK(*end_of_sequence);
            return Status::OK();
          }

          if (stride < window_size) {
            // Cache the data used for the next iteration.
            for (size_t i = stride; i < window_size; ++i) {
              cache_.emplace_back(batch_elements[i]);
            }
          } else if (stride > window_size) {
            // Drop the data before the next iteration.
            std::vector<Tensor> batch_element_tuple;
            for (size_t i = window_size; i < stride && !*end_of_sequence; ++i) {
              TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, &batch_element_tuple,
                                                      end_of_sequence));
              if (*end_of_sequence) {
                input_impl_.reset();
              }
            }
          }
        }

        // Construct output tensors.
        // Those codes below are copied from batch_dataset_op.cc.
        const size_t num_tuple_components = batch_elements[0].size();
        const int64 num_batch_elements = batch_elements.size();
        for (size_t component_index = 0; component_index < num_tuple_components;
             ++component_index) {
          const Tensor& first_element = batch_elements[0][component_index];
          TensorShape batch_component_shape({num_batch_elements});
          batch_component_shape.AppendShape(first_element.shape());
          Tensor batch_component(cpu_allocator(), first_element.dtype(),
                                 batch_component_shape);
          // Build the output tuple component by copying one slice
          // from each input element in the batch.
          for (size_t i = 0; i < num_batch_elements; ++i) {
            if (batch_elements[i][component_index].shape() !=
                first_element.shape()) {
              return errors::InvalidArgument(
                  "Cannot batch tensors with different shapes in component ",
                  component_index, ". First element had shape ",
                  first_element.shape().DebugString(), " and element ", i,
                  " had shape ",
                  batch_elements[i][component_index].shape().DebugString(),
                  ".");
            }
            TF_RETURN_IF_ERROR(batch_util::CopyElementToSlice(
                std::move(batch_elements[i][component_index]), &batch_component,
                i));
          }
          out_tensors->emplace_back(std::move(batch_component));
        }
        *end_of_sequence = false;
        return Status::OK();
      }

     protected:
      Status SaveInternal(IteratorStateWriter* writer) override {
        mutex_lock l(mu_);
        if (!input_impl_) {
          TF_RETURN_IF_ERROR(
              writer->WriteScalar(full_name("input_impl_empty"), ""));
        } else {
          TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_));
        }
        // Save cache.
        TF_RETURN_IF_ERROR(
            writer->WriteScalar(strings::StrCat("cache_size"), cache_.size()));
        for (int64 i = 0; i < cache_.size(); i++) {
          TF_RETURN_IF_ERROR(writer->WriteScalar(
              strings::StrCat("cache[", i, "]_size"), cache_[i].size()));
          for (int64 j = 0; j < cache_[i].size(); j++) {
            TF_RETURN_IF_ERROR(writer->WriteTensor(
                strings::StrCat("cache[", i, "][", j, "]"), cache_[i][j]));
          }
        }
        return Status::OK();
      }

      Status RestoreInternal(IteratorContext* ctx,
                             IteratorStateReader* reader) override {
        mutex_lock l(mu_);
        if (!reader->Contains(full_name("input_impl_empty"))) {
          TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_));
        } else {
          input_impl_.reset();
        }
        // Restore cache.
        int64 cache_size;
        TF_RETURN_IF_ERROR(
            reader->ReadScalar(strings::StrCat("cache_size"), &cache_size));
        cache_.resize(cache_size);
        for (int64 i = 0; i < cache_size; i++) {
          int64 vector_size;
          TF_RETURN_IF_ERROR(reader->ReadScalar(
              strings::StrCat("cache[", i, "]_size"), &vector_size));
          cache_[i].resize(vector_size);
          for (int64 j = 0; j < vector_size; j++) {
            TF_RETURN_IF_ERROR(reader->ReadTensor(
                strings::StrCat("cache[", i, "][", j, "]"), &cache_[i][j]));
          }
        }
        return Status::OK();
      }

     private:
      mutex mu_;
      std::vector<std::vector<Tensor>> cache_ GUARDED_BY(mu_);
      std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
    };

    const int64 window_size_;
    const int64 stride_;
    const DatasetBase* const input_;
    std::vector<PartialTensorShape> output_shapes_;
  };
};

REGISTER_KERNEL_BUILDER(Name("SlideDataset").Device(DEVICE_CPU),
                        SlideDatasetOp);

}  // namespace

}  // namespace tensorflow
<|endoftext|>


<commit_before>/*
Implement parallel vectorize workqueue on top of Intel TBB.
*/

#define TBB_PREVIEW_WAITING_FOR_WORKERS 1
/* tbb.h redefines these */
#include "../_pymodule.h"
#ifdef _POSIX_C_SOURCE
#undef _POSIX_C_SOURCE
#endif
#ifdef _XOPEN_SOURCE
#undef _XOPEN_SOURCE
#endif

#include <tbb/tbb.h>
#include <string.h>
#include <stdio.h>
#include "workqueue.h"

#include "gufunc_scheduler.h"

#if TBB_INTERFACE_VERSION >= 9106
#define TSI_INIT(count) tbb::task_scheduler_init(count)
#define TSI_TERMINATE(tsi) tsi->blocking_terminate(std::nothrow)
#else
#if __TBB_SUPPORTS_WORKERS_WAITING_IN_TERMINATE
#define TSI_INIT(count) tbb::task_scheduler_init(count, 0, /*blocking termination*/true)
#define TSI_TERMINATE(tsi) tsi->terminate()
#else
#error This version of TBB does not support blocking terminate
#endif
#endif

#define _DEBUG 0
#define _TRACE_SPLIT 0

static tbb::task_group *tg = NULL;
static tbb::task_scheduler_init *tsi = NULL;
static int tsi_count = 0;

static void
add_task(void *fn, void *args, void *dims, void *steps, void *data)
{
    tg->run([=]
    {
        auto func = reinterpret_cast<void (*)(void *args, void *dims, void *steps, void *data)>(fn);
        func(args, dims, steps, data);
    });
}

static void
parallel_for(void *fn, char **args, size_t *dimensions, size_t *steps, void *data,
             size_t inner_ndim, size_t array_count)
{
    static bool printed = false;
    if(!printed && _DEBUG)
    {
        puts("Using parallel_for");
        printed = true;
    }

    //     args = <ir.Argument '.1' of type i8**>,
    //     dimensions = <ir.Argument '.2' of type i64*>
    //     steps = <ir.Argument '.3' of type i64*>
    //     data = <ir.Argument '.4' of type i8*>

    const size_t arg_len = (inner_ndim + 1);

    if(_DEBUG && _TRACE_SPLIT)
    {
        printf("inner_ndim: %lu\n",inner_ndim);
        printf("arg_len: %lu\n", arg_len);
        printf("total: %lu\n", dimensions[0]);
        printf("dimensions: ");
        for(size_t j = 0; j < arg_len; j++)
            printf("%lu, ", ((size_t *)dimensions)[j]);
        printf("\nsteps: ");
        for(size_t j = 0; j < array_count; j++)
            printf("%lu, ", steps[j]);
        printf("\n*args: ");
        for(size_t j = 0; j < array_count; j++)
            printf("%p, ", (void *)args[j]);
        printf("\n");
    }

    using range_t = tbb::blocked_range<size_t>;
    tbb::parallel_for(range_t(0, dimensions[0]), [=](const range_t &range)
    {
        size_t * count_space = (size_t *)alloca(sizeof(size_t) * arg_len);
        char ** array_arg_space = (char**)alloca(sizeof(char*) * array_count);
        memcpy(count_space, dimensions, arg_len * sizeof(size_t));
        count_space[0] = range.size();

        if(_DEBUG && _TRACE_SPLIT > 1)
        {
            printf("THREAD %p:", count_space);
            printf("count_space: ");
            for(size_t j = 0; j < arg_len; j++)
                printf("%lu, ", count_space[j]);
            printf("\n");
        }
        for(size_t j = 0; j < array_count; j++)
        {
            char * base = args[j];
            size_t step = steps[j];
            ptrdiff_t offset = step * range.begin();
            array_arg_space[j] = base + offset;

            if(_DEBUG && _TRACE_SPLIT > 2)
            {
                printf("Index %ld\n", j);
                printf("-->Got base %p\n", (void *)base);
                printf("-->Got step %lu\n", step);
                printf("-->Got offset %ld\n", offset);
                printf("-->Got addr %p\n", (void *)array_arg_space[j]);
            }
        }

        if(_DEBUG && _TRACE_SPLIT > 2)
        {
            printf("array_arg_space: ");
            for(size_t j = 0; j < array_count; j++)
                printf("%p, ", (void *)array_arg_space[j]);
            printf("\n");
        }
        auto func = reinterpret_cast<void (*)(char **args, size_t *dims, size_t *steps, void *data)>(fn);
        func(array_arg_space, count_space, steps, data);
    });
}

void ignore_blocking_terminate_assertion( const char*, int, const char*, const char * )
{
    tbb::internal::runtime_warning("Unable to wait for threads to shut down before fork(). It can break multithreading in child process\n");
}

void ignore_assertion( const char*, int, const char*, const char * ) {}

static void prepare_fork(void)
{
    if(_DEBUG)
    {
        puts("Suspending TBB: prepare fork");
    }
    if(tsi)
    {
        assertion_handler_type orig = tbb::set_assertion_handler(ignore_blocking_terminate_assertion);
        TSI_TERMINATE(tsi);
        tbb::set_assertion_handler(orig);
    }
}

static void reset_after_fork(void)
{
    if(_DEBUG)
    {
        puts("Resuming TBB: after fork");
    }
    if(tsi)
        tsi->initialize(tsi_count);
}

#if PY_MAJOR_VERSION >= 3
static void unload_tbb(void)
{
    if(tsi)
    {
        if(_DEBUG)
        {
            puts("Unloading TBB");
        }
        tg->wait();
        delete tg;
        tg = NULL;
        assertion_handler_type orig = tbb::set_assertion_handler(ignore_assertion);
        tsi->terminate(); // no blocking terminate is needed here
        tbb::set_assertion_handler(orig);
        delete tsi;
        tsi = NULL;
    }
}
#endif

static void launch_threads(int count)
{
    if(tsi)
        return;
    if(_DEBUG)
        puts("Using TBB");
    if(count < 1)
        count = tbb::task_scheduler_init::automatic;
    tsi = new TSI_INIT(tsi_count = count);
    tg = new tbb::task_group;
    tg->run([] {}); // start creating threads asynchronously

#ifndef _MSC_VER
    pthread_atfork(prepare_fork, reset_after_fork, reset_after_fork);
#endif
}

static void synchronize(void)
{
    tg->wait();
}

static void ready(void)
{
}

MOD_INIT(tbbpool)
{
    PyObject *m;
    MOD_DEF(m, "tbbpool", "No docs", NULL)
    if (m == NULL)
        return MOD_ERROR_VAL;
#if PY_MAJOR_VERSION >= 3
    PyModuleDef *md = PyModule_GetDef(m);
    if (md)
    {
        md->m_free = (freefunc)unload_tbb;
    }
#endif

    PyObject_SetAttrString(m, "launch_threads",
                           PyLong_FromVoidPtr((void*)&launch_threads));
    PyObject_SetAttrString(m, "synchronize",
                           PyLong_FromVoidPtr((void*)&synchronize));
    PyObject_SetAttrString(m, "ready",
                           PyLong_FromVoidPtr((void*)&ready));
    PyObject_SetAttrString(m, "add_task",
                           PyLong_FromVoidPtr((void*)&add_task));
    PyObject_SetAttrString(m, "parallel_for",
                           PyLong_FromVoidPtr((void*)&parallel_for));
    PyObject_SetAttrString(m, "do_scheduling_signed",
                           PyLong_FromVoidPtr((void*)&do_scheduling_signed));
    PyObject_SetAttrString(m, "do_scheduling_unsigned",
                           PyLong_FromVoidPtr((void*)&do_scheduling_unsigned));

    return MOD_SUCCESS_VAL(m);
}
<commit_msg>Add get_num_threads and set_num_threads functions to tbbpool.cpp<commit_after>/*
Implement parallel vectorize workqueue on top of Intel TBB.
*/

#define TBB_PREVIEW_WAITING_FOR_WORKERS 1
/* tbb.h redefines these */
#include "../_pymodule.h"
#ifdef _POSIX_C_SOURCE
#undef _POSIX_C_SOURCE
#endif
#ifdef _XOPEN_SOURCE
#undef _XOPEN_SOURCE
#endif

#include <tbb/tbb.h>
#include <string.h>
#include <stdio.h>
#include "workqueue.h"

#include "gufunc_scheduler.h"

#if TBB_INTERFACE_VERSION >= 9106
#define TSI_INIT(count) tbb::task_scheduler_init(count)
#define TSI_TERMINATE(tsi) tsi->blocking_terminate(std::nothrow)
#else
#if __TBB_SUPPORTS_WORKERS_WAITING_IN_TERMINATE
#define TSI_INIT(count) tbb::task_scheduler_init(count, 0, /*blocking termination*/true)
#define TSI_TERMINATE(tsi) tsi->terminate()
#else
#error This version of TBB does not support blocking terminate
#endif
#endif

#define _DEBUG 0
#define _TRACE_SPLIT 0

static tbb::task_group *tg = NULL;
static tbb::task_scheduler_init *tsi = NULL;
static int tsi_count = 0;
int num_threads = 0;

static void
add_task(void *fn, void *args, void *dims, void *steps, void *data)
{
    tg->run([=]
    {
        auto func = reinterpret_cast<void (*)(void *args, void *dims, void *steps, void *data)>(fn);
        func(args, dims, steps, data);
    });
}

static void
parallel_for(void *fn, char **args, size_t *dimensions, size_t *steps, void *data,
             size_t inner_ndim, size_t array_count)
{
    static bool printed = false;
    if(!printed && _DEBUG)
    {
        puts("Using parallel_for");
        printed = true;
    }

    //     args = <ir.Argument '.1' of type i8**>,
    //     dimensions = <ir.Argument '.2' of type i64*>
    //     steps = <ir.Argument '.3' of type i64*>
    //     data = <ir.Argument '.4' of type i8*>

    const size_t arg_len = (inner_ndim + 1);

    if(_DEBUG && _TRACE_SPLIT)
    {
        printf("inner_ndim: %lu\n",inner_ndim);
        printf("arg_len: %lu\n", arg_len);
        printf("total: %lu\n", dimensions[0]);
        printf("dimensions: ");
        for(size_t j = 0; j < arg_len; j++)
            printf("%lu, ", ((size_t *)dimensions)[j]);
        printf("\nsteps: ");
        for(size_t j = 0; j < array_count; j++)
            printf("%lu, ", steps[j]);
        printf("\n*args: ");
        for(size_t j = 0; j < array_count; j++)
            printf("%p, ", (void *)args[j]);
        printf("\n");
    }

    using range_t = tbb::blocked_range<size_t>;
    tbb::parallel_for(range_t(0, dimensions[0]), [=](const range_t &range)
    {
        size_t * count_space = (size_t *)alloca(sizeof(size_t) * arg_len);
        char ** array_arg_space = (char**)alloca(sizeof(char*) * array_count);
        memcpy(count_space, dimensions, arg_len * sizeof(size_t));
        count_space[0] = range.size();

        if(_DEBUG && _TRACE_SPLIT > 1)
        {
            printf("THREAD %p:", count_space);
            printf("count_space: ");
            for(size_t j = 0; j < arg_len; j++)
                printf("%lu, ", count_space[j]);
            printf("\n");
        }
        for(size_t j = 0; j < array_count; j++)
        {
            char * base = args[j];
            size_t step = steps[j];
            ptrdiff_t offset = step * range.begin();
            array_arg_space[j] = base + offset;

            if(_DEBUG && _TRACE_SPLIT > 2)
            {
                printf("Index %ld\n", j);
                printf("-->Got base %p\n", (void *)base);
                printf("-->Got step %lu\n", step);
                printf("-->Got offset %ld\n", offset);
                printf("-->Got addr %p\n", (void *)array_arg_space[j]);
            }
        }

        if(_DEBUG && _TRACE_SPLIT > 2)
        {
            printf("array_arg_space: ");
            for(size_t j = 0; j < array_count; j++)
                printf("%p, ", (void *)array_arg_space[j]);
            printf("\n");
        }
        auto func = reinterpret_cast<void (*)(char **args, size_t *dims, size_t *steps, void *data)>(fn);
        func(array_arg_space, count_space, steps, data);
    });
}

void ignore_blocking_terminate_assertion( const char*, int, const char*, const char * )
{
    tbb::internal::runtime_warning("Unable to wait for threads to shut down before fork(). It can break multithreading in child process\n");
}

void ignore_assertion( const char*, int, const char*, const char * ) {}

static void prepare_fork(void)
{
    if(_DEBUG)
    {
        puts("Suspending TBB: prepare fork");
    }
    if(tsi)
    {
        assertion_handler_type orig = tbb::set_assertion_handler(ignore_blocking_terminate_assertion);
        TSI_TERMINATE(tsi);
        tbb::set_assertion_handler(orig);
    }
}

static void reset_after_fork(void)
{
    if(_DEBUG)
    {
        puts("Resuming TBB: after fork");
    }
    if(tsi)
        tsi->initialize(tsi_count);
}

#if PY_MAJOR_VERSION >= 3
static void unload_tbb(void)
{
    if(tsi)
    {
        if(_DEBUG)
        {
            puts("Unloading TBB");
        }
        tg->wait();
        delete tg;
        tg = NULL;
        assertion_handler_type orig = tbb::set_assertion_handler(ignore_assertion);
        tsi->terminate(); // no blocking terminate is needed here
        tbb::set_assertion_handler(orig);
        delete tsi;
        tsi = NULL;
    }
}
#endif

static void launch_threads(int count)
{
    if(tsi)
        return;
    if(_DEBUG)
        puts("Using TBB");
    if(count < 1)
        count = tbb::task_scheduler_init::automatic;
    tsi = new TSI_INIT(tsi_count = count);
    tg = new tbb::task_group;
    tg->run([] {}); // start creating threads asynchronously

#ifndef _MSC_VER
    pthread_atfork(prepare_fork, reset_after_fork, reset_after_fork);
#endif
}

static void synchronize(void)
{
    tg->wait();
}

static void ready(void)
{
}

static void set_num_threads(int count)
{
    num_threads = count;
}

static int get_num_threads(void)
{
    return num_threads;
}

MOD_INIT(tbbpool)
{
    PyObject *m;
    MOD_DEF(m, "tbbpool", "No docs", NULL)
    if (m == NULL)
        return MOD_ERROR_VAL;
#if PY_MAJOR_VERSION >= 3
    PyModuleDef *md = PyModule_GetDef(m);
    if (md)
    {
        md->m_free = (freefunc)unload_tbb;
    }
#endif

    PyObject_SetAttrString(m, "launch_threads",
                           PyLong_FromVoidPtr((void*)&launch_threads));
    PyObject_SetAttrString(m, "synchronize",
                           PyLong_FromVoidPtr((void*)&synchronize));
    PyObject_SetAttrString(m, "ready",
                           PyLong_FromVoidPtr((void*)&ready));
    PyObject_SetAttrString(m, "add_task",
                           PyLong_FromVoidPtr((void*)&add_task));
    PyObject_SetAttrString(m, "parallel_for",
                           PyLong_FromVoidPtr((void*)&parallel_for));
    PyObject_SetAttrString(m, "do_scheduling_signed",
                           PyLong_FromVoidPtr((void*)&do_scheduling_signed));
    PyObject_SetAttrString(m, "do_scheduling_unsigned",
                           PyLong_FromVoidPtr((void*)&do_scheduling_unsigned));
    PyObject_SetAttrString(m, "set_num_threads",
                           PyLong_FromVoidPtr((void*)&set_num_threads));
    PyObject_SetAttrString(m, "get_num_threads",
                           PyLong_FromVoidPtr((void*)&get_num_threads));

    return MOD_SUCCESS_VAL(m);
}
<|endoftext|>


<commit_before>#include <exception>
//#if (__GLIBCXX__ / 10000) == 2014
namespace std {
inline bool uncaught_exception() noexcept(true)
{ return current_exception() != nullptr; }
} // namespace std
//#endif

#define CATCH_CONFIG_MAIN
#include <catch.hpp>
<commit_msg>Remove GCC fix that is not needed<commit_after>#define CATCH_CONFIG_MAIN
#include <catch.hpp>
<|endoftext|>


<commit_before>#include "archive.hh"
#include "binary-cache-store.hh"
#include "compression.hh"
#include "derivations.hh"
#include "fs-accessor.hh"
#include "globals.hh"
#include "nar-info.hh"
#include "sync.hh"
#include "remote-fs-accessor.hh"
#include "nar-info-disk-cache.hh"
#include "nar-accessor.hh"
#include "json.hh"

#include <chrono>

#include <future>

namespace nix {

BinaryCacheStore::BinaryCacheStore(const Params & params)
    : Store(params)
{
    if (secretKeyFile != "")
        secretKey = std::unique_ptr<SecretKey>(new SecretKey(readFile(secretKeyFile)));

    StringSink sink;
    sink << narVersionMagic1;
    narMagic = *sink.s;
}

void BinaryCacheStore::init()
{
    std::string cacheInfoFile = "nix-cache-info";

    auto cacheInfo = getFile(cacheInfoFile);
    if (!cacheInfo) {
        upsertFile(cacheInfoFile, "StoreDir: " + storeDir + "\n", "text/x-nix-cache-info");
    } else {
        for (auto & line : tokenizeString<Strings>(*cacheInfo, "\n")) {
            size_t colon = line.find(':');
            if (colon == std::string::npos) continue;
            auto name = line.substr(0, colon);
            auto value = trim(line.substr(colon + 1, std::string::npos));
            if (name == "StoreDir") {
                if (value != storeDir)
                    throw Error(format("binary cache '%s' is for Nix stores with prefix '%s', not '%s'")
                        % getUri() % value % storeDir);
            } else if (name == "WantMassQuery") {
                wantMassQuery_ = value == "1";
            } else if (name == "Priority") {
                string2Int(value, priority);
            }
        }
    }
}

std::shared_ptr<std::string> BinaryCacheStore::getFile(const std::string & path)
{
    std::promise<std::shared_ptr<std::string>> promise;
    getFile(path,
        [&](std::shared_ptr<std::string> result) {
            promise.set_value(result);
        },
        [&](std::exception_ptr exc) {
            promise.set_exception(exc);
        });
    return promise.get_future().get();
}

Path BinaryCacheStore::narInfoFileFor(const Path & storePath)
{
    assertStorePath(storePath);
    return storePathToHash(storePath) + ".narinfo";
}

void BinaryCacheStore::addToStore(const ValidPathInfo & info, const ref<std::string> & nar,
    RepairFlag repair, CheckSigsFlag checkSigs, std::shared_ptr<FSAccessor> accessor)
{
    if (!repair && isValidPath(info.path)) return;

    /* Verify that all references are valid. This may do some .narinfo
       reads, but typically they'll already be cached. */
    for (auto & ref : info.references)
        try {
            if (ref != info.path)
                queryPathInfo(ref);
        } catch (InvalidPath &) {
            throw Error(format("cannot add '%s' to the binary cache because the reference '%s' is not valid")
                % info.path % ref);
        }

    auto narInfoFile = narInfoFileFor(info.path);

    assert(nar->compare(0, narMagic.size(), narMagic) == 0);

    auto narInfo = make_ref<NarInfo>(info);

    narInfo->narSize = nar->size();
    narInfo->narHash = hashString(htSHA256, *nar);

    if (info.narHash && info.narHash != narInfo->narHash)
        throw Error(format("refusing to copy corrupted path '%1%' to binary cache") % info.path);

    auto accessor_ = std::dynamic_pointer_cast<RemoteFSAccessor>(accessor);

    /* Optionally write a JSON file containing a listing of the
       contents of the NAR. */
    if (writeNARListing) {
        std::ostringstream jsonOut;

        {
            JSONObject jsonRoot(jsonOut);
            jsonRoot.attr("version", 1);

            auto narAccessor = makeNarAccessor(nar);

            if (accessor_) {
                accessor_->nars.emplace(info.path, narAccessor);
                //accessor_->addToCache(info.path, *nar);
            }

            std::function<void(const Path &, JSONPlaceholder &)> recurse;

            recurse = [&](const Path & path, JSONPlaceholder & res) {
                auto st = narAccessor->stat(path);

                auto obj = res.object();

                switch (st.type) {
                case FSAccessor::Type::tRegular:
                    obj.attr("type", "regular");
                    obj.attr("size", st.fileSize);
                    if (st.isExecutable)
                        obj.attr("executable", true);
                    break;
                case FSAccessor::Type::tDirectory:
                    obj.attr("type", "directory");
                    {
                        auto res2 = obj.object("entries");
                        for (auto & name : narAccessor->readDirectory(path)) {
                            auto res3 = res2.placeholder(name);
                            recurse(path + "/" + name, res3);
                        }
                    }
                    break;
                case FSAccessor::Type::tSymlink:
                    obj.attr("type", "symlink");
                    obj.attr("target", narAccessor->readLink(path));
                    break;
                default:
                    abort();
                }
            };

            {
                auto res = jsonRoot.placeholder("root");
                recurse("", res);
            }
        }

        upsertFile(storePathToHash(info.path) + ".ls", jsonOut.str(), "application/json");
    }

    else {
        if (accessor_) {
            accessor_->nars.emplace(info.path, makeNarAccessor(nar));
            //accessor_->addToCache(info.path, *nar);
        }
    }

    /* Compress the NAR. */
    narInfo->compression = compression;
    auto now1 = std::chrono::steady_clock::now();
    auto narCompressed = compress(compression, *nar);
    auto now2 = std::chrono::steady_clock::now();
    narInfo->fileHash = hashString(htSHA256, *narCompressed);
    narInfo->fileSize = narCompressed->size();

    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();
    printMsg(lvlTalkative, format("copying path '%1%' (%2% bytes, compressed %3$.1f%% in %4% ms) to binary cache")
        % narInfo->path % narInfo->narSize
        % ((1.0 - (double) narCompressed->size() / nar->size()) * 100.0)
        % duration);

    /* Atomically write the NAR file. */
    narInfo->url = "nar/" + narInfo->fileHash.to_string(Base32, false) + ".nar"
        + (compression == "xz" ? ".xz" :
           compression == "bzip2" ? ".bz2" :
           compression == "br" ? ".br" :
           "");
    if (repair || !fileExists(narInfo->url)) {
        stats.narWrite++;
        upsertFile(narInfo->url, *narCompressed, "application/x-nix-nar");
    } else
        stats.narWriteAverted++;

    stats.narWriteBytes += nar->size();
    stats.narWriteCompressedBytes += narCompressed->size();
    stats.narWriteCompressionTimeMs += duration;

    /* Atomically write the NAR info file.*/
    if (secretKey) narInfo->sign(*secretKey);

    upsertFile(narInfoFile, narInfo->to_string(), "text/x-nix-narinfo");

    auto hashPart = storePathToHash(narInfo->path);

    {
        auto state_(state.lock());
        state_->pathInfoCache.upsert(hashPart, std::shared_ptr<NarInfo>(narInfo));
    }

    if (diskCache)
        diskCache->upsertNarInfo(getUri(), hashPart, std::shared_ptr<NarInfo>(narInfo));

    stats.narInfoWrite++;
}

bool BinaryCacheStore::isValidPathUncached(const Path & storePath)
{
    // FIXME: this only checks whether a .narinfo with a matching hash
    // part exists. So ‘f4kb...-foo’ matches ‘f4kb...-bar’, even
    // though they shouldn't. Not easily fixed.
    return fileExists(narInfoFileFor(storePath));
}

void BinaryCacheStore::narFromPath(const Path & storePath, Sink & sink)
{
    auto info = queryPathInfo(storePath).cast<const NarInfo>();

    auto nar = getFile(info->url);

    if (!nar) throw Error(format("file '%s' missing from binary cache") % info->url);

    stats.narRead++;
    stats.narReadCompressedBytes += nar->size();

    /* Decompress the NAR. FIXME: would be nice to have the remote
       side do this. */
    try {
        nar = decompress(info->compression, *nar);
    } catch (UnknownCompressionMethod &) {
        throw Error(format("binary cache path '%s' uses unknown compression method '%s'")
            % storePath % info->compression);
    }

    stats.narReadBytes += nar->size();

    printMsg(lvlTalkative, format("exporting path '%1%' (%2% bytes)") % storePath % nar->size());

    assert(nar->size() % 8 == 0);

    sink((unsigned char *) nar->c_str(), nar->size());
}

void BinaryCacheStore::queryPathInfoUncached(const Path & storePath,
        std::function<void(std::shared_ptr<ValidPathInfo>)> success,
        std::function<void(std::exception_ptr exc)> failure)
{
    auto uri = getUri();
    auto act = std::make_shared<Activity>(*logger, lvlTalkative, actQueryPathInfo,
        fmt("querying info about '%s' on '%s'", storePath, uri), Logger::Fields{storePath, uri});
    PushActivity pact(act->id);

    auto narInfoFile = narInfoFileFor(storePath);

    getFile(narInfoFile,
        [=](std::shared_ptr<std::string> data) {
            if (!data) return success(0);

            stats.narInfoRead++;

            callSuccess(success, failure, (std::shared_ptr<ValidPathInfo>)
                std::make_shared<NarInfo>(*this, *data, narInfoFile));

            (void) act; // force Activity into this lambda to ensure it stays alive
        },
        failure);
}

Path BinaryCacheStore::addToStore(const string & name, const Path & srcPath,
    bool recursive, HashType hashAlgo, PathFilter & filter, RepairFlag repair)
{
    // FIXME: some cut&paste from LocalStore::addToStore().

    /* Read the whole path into memory. This is not a very scalable
       method for very large paths, but `copyPath' is mainly used for
       small files. */
    StringSink sink;
    Hash h;
    if (recursive) {
        dumpPath(srcPath, sink, filter);
        h = hashString(hashAlgo, *sink.s);
    } else {
        auto s = readFile(srcPath);
        dumpString(s, sink);
        h = hashString(hashAlgo, s);
    }

    ValidPathInfo info;
    info.path = makeFixedOutputPath(recursive, h, name);

    addToStore(info, sink.s, repair, CheckSigs, nullptr);

    return info.path;
}

Path BinaryCacheStore::addTextToStore(const string & name, const string & s,
    const PathSet & references, RepairFlag repair)
{
    ValidPathInfo info;
    info.path = computeStorePathForText(name, s, references);
    info.references = references;

    if (repair || !isValidPath(info.path)) {
        StringSink sink;
        dumpString(s, sink);
        addToStore(info, sink.s, repair, CheckSigs, nullptr);
    }

    return info.path;
}

ref<FSAccessor> BinaryCacheStore::getFSAccessor()
{
    return make_ref<RemoteFSAccessor>(ref<Store>(shared_from_this()), localNarCache);
}

std::shared_ptr<std::string> BinaryCacheStore::getBuildLog(const Path & path)
{
    Path drvPath;

    if (isDerivation(path))
        drvPath = path;
    else {
        try {
            auto info = queryPathInfo(path);
            // FIXME: add a "Log" field to .narinfo
            if (info->deriver == "") return nullptr;
            drvPath = info->deriver;
        } catch (InvalidPath &) {
            return nullptr;
        }
    }

    auto logPath = "log/" + baseNameOf(drvPath);

    debug("fetching build log from binary cache '%s/%s'", getUri(), logPath);

    return getFile(logPath);
}

}
<commit_msg>Revert "Let's not populate the NAR cache from hydra-queue-runner for now"<commit_after>#include "archive.hh"
#include "binary-cache-store.hh"
#include "compression.hh"
#include "derivations.hh"
#include "fs-accessor.hh"
#include "globals.hh"
#include "nar-info.hh"
#include "sync.hh"
#include "remote-fs-accessor.hh"
#include "nar-info-disk-cache.hh"
#include "nar-accessor.hh"
#include "json.hh"

#include <chrono>

#include <future>

namespace nix {

BinaryCacheStore::BinaryCacheStore(const Params & params)
    : Store(params)
{
    if (secretKeyFile != "")
        secretKey = std::unique_ptr<SecretKey>(new SecretKey(readFile(secretKeyFile)));

    StringSink sink;
    sink << narVersionMagic1;
    narMagic = *sink.s;
}

void BinaryCacheStore::init()
{
    std::string cacheInfoFile = "nix-cache-info";

    auto cacheInfo = getFile(cacheInfoFile);
    if (!cacheInfo) {
        upsertFile(cacheInfoFile, "StoreDir: " + storeDir + "\n", "text/x-nix-cache-info");
    } else {
        for (auto & line : tokenizeString<Strings>(*cacheInfo, "\n")) {
            size_t colon = line.find(':');
            if (colon == std::string::npos) continue;
            auto name = line.substr(0, colon);
            auto value = trim(line.substr(colon + 1, std::string::npos));
            if (name == "StoreDir") {
                if (value != storeDir)
                    throw Error(format("binary cache '%s' is for Nix stores with prefix '%s', not '%s'")
                        % getUri() % value % storeDir);
            } else if (name == "WantMassQuery") {
                wantMassQuery_ = value == "1";
            } else if (name == "Priority") {
                string2Int(value, priority);
            }
        }
    }
}

std::shared_ptr<std::string> BinaryCacheStore::getFile(const std::string & path)
{
    std::promise<std::shared_ptr<std::string>> promise;
    getFile(path,
        [&](std::shared_ptr<std::string> result) {
            promise.set_value(result);
        },
        [&](std::exception_ptr exc) {
            promise.set_exception(exc);
        });
    return promise.get_future().get();
}

Path BinaryCacheStore::narInfoFileFor(const Path & storePath)
{
    assertStorePath(storePath);
    return storePathToHash(storePath) + ".narinfo";
}

void BinaryCacheStore::addToStore(const ValidPathInfo & info, const ref<std::string> & nar,
    RepairFlag repair, CheckSigsFlag checkSigs, std::shared_ptr<FSAccessor> accessor)
{
    if (!repair && isValidPath(info.path)) return;

    /* Verify that all references are valid. This may do some .narinfo
       reads, but typically they'll already be cached. */
    for (auto & ref : info.references)
        try {
            if (ref != info.path)
                queryPathInfo(ref);
        } catch (InvalidPath &) {
            throw Error(format("cannot add '%s' to the binary cache because the reference '%s' is not valid")
                % info.path % ref);
        }

    auto narInfoFile = narInfoFileFor(info.path);

    assert(nar->compare(0, narMagic.size(), narMagic) == 0);

    auto narInfo = make_ref<NarInfo>(info);

    narInfo->narSize = nar->size();
    narInfo->narHash = hashString(htSHA256, *nar);

    if (info.narHash && info.narHash != narInfo->narHash)
        throw Error(format("refusing to copy corrupted path '%1%' to binary cache") % info.path);

    auto accessor_ = std::dynamic_pointer_cast<RemoteFSAccessor>(accessor);

    /* Optionally write a JSON file containing a listing of the
       contents of the NAR. */
    if (writeNARListing) {
        std::ostringstream jsonOut;

        {
            JSONObject jsonRoot(jsonOut);
            jsonRoot.attr("version", 1);

            auto narAccessor = makeNarAccessor(nar);

            if (accessor_) {
                accessor_->nars.emplace(info.path, narAccessor);
                accessor_->addToCache(info.path, *nar);
            }

            std::function<void(const Path &, JSONPlaceholder &)> recurse;

            recurse = [&](const Path & path, JSONPlaceholder & res) {
                auto st = narAccessor->stat(path);

                auto obj = res.object();

                switch (st.type) {
                case FSAccessor::Type::tRegular:
                    obj.attr("type", "regular");
                    obj.attr("size", st.fileSize);
                    if (st.isExecutable)
                        obj.attr("executable", true);
                    break;
                case FSAccessor::Type::tDirectory:
                    obj.attr("type", "directory");
                    {
                        auto res2 = obj.object("entries");
                        for (auto & name : narAccessor->readDirectory(path)) {
                            auto res3 = res2.placeholder(name);
                            recurse(path + "/" + name, res3);
                        }
                    }
                    break;
                case FSAccessor::Type::tSymlink:
                    obj.attr("type", "symlink");
                    obj.attr("target", narAccessor->readLink(path));
                    break;
                default:
                    abort();
                }
            };

            {
                auto res = jsonRoot.placeholder("root");
                recurse("", res);
            }
        }

        upsertFile(storePathToHash(info.path) + ".ls", jsonOut.str(), "application/json");
    }

    else {
        if (accessor_) {
            accessor_->nars.emplace(info.path, makeNarAccessor(nar));
            accessor_->addToCache(info.path, *nar);
        }
    }

    /* Compress the NAR. */
    narInfo->compression = compression;
    auto now1 = std::chrono::steady_clock::now();
    auto narCompressed = compress(compression, *nar);
    auto now2 = std::chrono::steady_clock::now();
    narInfo->fileHash = hashString(htSHA256, *narCompressed);
    narInfo->fileSize = narCompressed->size();

    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();
    printMsg(lvlTalkative, format("copying path '%1%' (%2% bytes, compressed %3$.1f%% in %4% ms) to binary cache")
        % narInfo->path % narInfo->narSize
        % ((1.0 - (double) narCompressed->size() / nar->size()) * 100.0)
        % duration);

    /* Atomically write the NAR file. */
    narInfo->url = "nar/" + narInfo->fileHash.to_string(Base32, false) + ".nar"
        + (compression == "xz" ? ".xz" :
           compression == "bzip2" ? ".bz2" :
           compression == "br" ? ".br" :
           "");
    if (repair || !fileExists(narInfo->url)) {
        stats.narWrite++;
        upsertFile(narInfo->url, *narCompressed, "application/x-nix-nar");
    } else
        stats.narWriteAverted++;

    stats.narWriteBytes += nar->size();
    stats.narWriteCompressedBytes += narCompressed->size();
    stats.narWriteCompressionTimeMs += duration;

    /* Atomically write the NAR info file.*/
    if (secretKey) narInfo->sign(*secretKey);

    upsertFile(narInfoFile, narInfo->to_string(), "text/x-nix-narinfo");

    auto hashPart = storePathToHash(narInfo->path);

    {
        auto state_(state.lock());
        state_->pathInfoCache.upsert(hashPart, std::shared_ptr<NarInfo>(narInfo));
    }

    if (diskCache)
        diskCache->upsertNarInfo(getUri(), hashPart, std::shared_ptr<NarInfo>(narInfo));

    stats.narInfoWrite++;
}

bool BinaryCacheStore::isValidPathUncached(const Path & storePath)
{
    // FIXME: this only checks whether a .narinfo with a matching hash
    // part exists. So ‘f4kb...-foo’ matches ‘f4kb...-bar’, even
    // though they shouldn't. Not easily fixed.
    return fileExists(narInfoFileFor(storePath));
}

void BinaryCacheStore::narFromPath(const Path & storePath, Sink & sink)
{
    auto info = queryPathInfo(storePath).cast<const NarInfo>();

    auto nar = getFile(info->url);

    if (!nar) throw Error(format("file '%s' missing from binary cache") % info->url);

    stats.narRead++;
    stats.narReadCompressedBytes += nar->size();

    /* Decompress the NAR. FIXME: would be nice to have the remote
       side do this. */
    try {
        nar = decompress(info->compression, *nar);
    } catch (UnknownCompressionMethod &) {
        throw Error(format("binary cache path '%s' uses unknown compression method '%s'")
            % storePath % info->compression);
    }

    stats.narReadBytes += nar->size();

    printMsg(lvlTalkative, format("exporting path '%1%' (%2% bytes)") % storePath % nar->size());

    assert(nar->size() % 8 == 0);

    sink((unsigned char *) nar->c_str(), nar->size());
}

void BinaryCacheStore::queryPathInfoUncached(const Path & storePath,
        std::function<void(std::shared_ptr<ValidPathInfo>)> success,
        std::function<void(std::exception_ptr exc)> failure)
{
    auto uri = getUri();
    auto act = std::make_shared<Activity>(*logger, lvlTalkative, actQueryPathInfo,
        fmt("querying info about '%s' on '%s'", storePath, uri), Logger::Fields{storePath, uri});
    PushActivity pact(act->id);

    auto narInfoFile = narInfoFileFor(storePath);

    getFile(narInfoFile,
        [=](std::shared_ptr<std::string> data) {
            if (!data) return success(0);

            stats.narInfoRead++;

            callSuccess(success, failure, (std::shared_ptr<ValidPathInfo>)
                std::make_shared<NarInfo>(*this, *data, narInfoFile));

            (void) act; // force Activity into this lambda to ensure it stays alive
        },
        failure);
}

Path BinaryCacheStore::addToStore(const string & name, const Path & srcPath,
    bool recursive, HashType hashAlgo, PathFilter & filter, RepairFlag repair)
{
    // FIXME: some cut&paste from LocalStore::addToStore().

    /* Read the whole path into memory. This is not a very scalable
       method for very large paths, but `copyPath' is mainly used for
       small files. */
    StringSink sink;
    Hash h;
    if (recursive) {
        dumpPath(srcPath, sink, filter);
        h = hashString(hashAlgo, *sink.s);
    } else {
        auto s = readFile(srcPath);
        dumpString(s, sink);
        h = hashString(hashAlgo, s);
    }

    ValidPathInfo info;
    info.path = makeFixedOutputPath(recursive, h, name);

    addToStore(info, sink.s, repair, CheckSigs, nullptr);

    return info.path;
}

Path BinaryCacheStore::addTextToStore(const string & name, const string & s,
    const PathSet & references, RepairFlag repair)
{
    ValidPathInfo info;
    info.path = computeStorePathForText(name, s, references);
    info.references = references;

    if (repair || !isValidPath(info.path)) {
        StringSink sink;
        dumpString(s, sink);
        addToStore(info, sink.s, repair, CheckSigs, nullptr);
    }

    return info.path;
}

ref<FSAccessor> BinaryCacheStore::getFSAccessor()
{
    return make_ref<RemoteFSAccessor>(ref<Store>(shared_from_this()), localNarCache);
}

std::shared_ptr<std::string> BinaryCacheStore::getBuildLog(const Path & path)
{
    Path drvPath;

    if (isDerivation(path))
        drvPath = path;
    else {
        try {
            auto info = queryPathInfo(path);
            // FIXME: add a "Log" field to .narinfo
            if (info->deriver == "") return nullptr;
            drvPath = info->deriver;
        } catch (InvalidPath &) {
            return nullptr;
        }
    }

    auto logPath = "log/" + baseNameOf(drvPath);

    debug("fetching build log from binary cache '%s/%s'", getUri(), logPath);

    return getFile(logPath);
}

}
<|endoftext|>