<commit_before>/*
 *  Copyright (c) 2016 The WebRTC project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <string.h>

#include <algorithm>
#include <initializer_list>
#include <iostream>  // TODO(zijiehe): Remove once flaky has been resolved.
#include <memory>
#include <utility>

// TODO(zijiehe): Remove once flaky has been resolved.
#include "webrtc/base/base64.h"
#include "webrtc/base/checks.h"
#include "webrtc/base/constructormagic.h"
#include "webrtc/base/logging.h"
#include "webrtc/modules/desktop_capture/desktop_capturer.h"
#include "webrtc/modules/desktop_capture/desktop_capture_options.h"
#include "webrtc/modules/desktop_capture/desktop_frame.h"
#include "webrtc/modules/desktop_capture/desktop_region.h"
#include "webrtc/modules/desktop_capture/mock_desktop_capturer_callback.h"
#include "webrtc/modules/desktop_capture/rgba_color.h"
#include "webrtc/modules/desktop_capture/screen_drawer.h"
#include "webrtc/test/gmock.h"
#include "webrtc/test/gtest.h"

#if defined(WEBRTC_WIN)
#include "webrtc/base/win32.h"
#include "webrtc/modules/desktop_capture/win/screen_capturer_win_directx.h"
#endif  // defined(WEBRTC_WIN)

using ::testing::_;

namespace webrtc {

namespace {

ACTION_P2(SaveCaptureResult, result, dest) {
  *result = arg0;
  *dest = std::move(*arg1);
}

// Returns true if color in |rect| of |frame| is |color|.
bool ArePixelsColoredBy(const DesktopFrame& frame,
                        DesktopRect rect,
                        RgbaColor color,
                        bool may_partially_draw) {
  if (!may_partially_draw) {
    // updated_region() should cover the painted area.
    DesktopRegion updated_region(frame.updated_region());
    updated_region.IntersectWith(rect);
    if (!updated_region.Equals(DesktopRegion(rect))) {
      return false;
    }
  }

  // Color in the |rect| should be |color|.
  uint8_t* row = frame.GetFrameDataAtPos(rect.top_left());
  for (int i = 0; i < rect.height(); i++) {
    uint8_t* column = row;
    for (int j = 0; j < rect.width(); j++) {
      if (color != RgbaColor(column)) {
        return false;
      }
      column += DesktopFrame::kBytesPerPixel;
    }
    row += frame.stride();
  }
  return true;
}

}  // namespace

class ScreenCapturerIntegrationTest : public testing::Test {
 public:
  void SetUp() override {
    capturer_ = DesktopCapturer::CreateScreenCapturer(
        DesktopCaptureOptions::CreateDefault());
  }

 protected:
  void TestCaptureUpdatedRegion(
      std::initializer_list<DesktopCapturer*> capturers) {
    RTC_DCHECK(capturers.size() > 0);
    // A large enough area for the tests, which should be able to be fulfilled
    // by most systems.
#if defined(WEBRTC_WIN)
    // On Windows, an interesting warning window may pop up randomly. The root
    // cause is still under investigation, so reduce the test area to work
    // around. Bug https://bugs.chromium.org/p/webrtc/issues/detail?id=6666.
    const int kTestArea = 416;
#else
    const int kTestArea = 512;
#endif
    const int kRectSize = 32;
    std::unique_ptr<ScreenDrawer> drawer = ScreenDrawer::Create();
    if (!drawer || drawer->DrawableRegion().is_empty()) {
      LOG(LS_WARNING) << "No ScreenDrawer implementation for current platform.";
      return;
    }
    if (drawer->DrawableRegion().width() < kTestArea ||
        drawer->DrawableRegion().height() < kTestArea) {
      LOG(LS_WARNING) << "ScreenDrawer::DrawableRegion() is too small for the "
                         "CaptureUpdatedRegion tests.";
      return;
    }

    for (DesktopCapturer* capturer : capturers) {
      capturer->Start(&callback_);
    }

    // Draw a set of |kRectSize| by |kRectSize| rectangles at (|i|, |i|), or
    // |i| by |i| rectangles at (|kRectSize|, |kRectSize|). One of (controlled
    // by |c|) its primary colors is |i|, and the other two are 0x7f. So we
    // won't draw a black or white rectangle.
    for (int c = 0; c < 3; c++) {
      // A fixed size rectangle.
      for (int i = 0; i < kTestArea - kRectSize; i += 16) {
        DesktopRect rect = DesktopRect::MakeXYWH(i, i, kRectSize, kRectSize);
        rect.Translate(drawer->DrawableRegion().top_left());
        RgbaColor color((c == 0 ? (i & 0xff) : 0x7f),
                        (c == 1 ? (i & 0xff) : 0x7f),
                        (c == 2 ? (i & 0xff) : 0x7f));
        // Fail fast.
        ASSERT_NO_FATAL_FAILURE(
            TestCaptureOneFrame(capturers, drawer.get(), rect, color));
      }

      // A variable-size rectangle.
      for (int i = 0; i < kTestArea - kRectSize; i += 16) {
        DesktopRect rect = DesktopRect::MakeXYWH(kRectSize, kRectSize, i, i);
        rect.Translate(drawer->DrawableRegion().top_left());
        RgbaColor color((c == 0 ? (i & 0xff) : 0x7f),
                        (c == 1 ? (i & 0xff) : 0x7f),
                        (c == 2 ? (i & 0xff) : 0x7f));
        // Fail fast.
        ASSERT_NO_FATAL_FAILURE(
            TestCaptureOneFrame(capturers, drawer.get(), rect, color));
      }
    }
  }

  void TestCaptureUpdatedRegion() {
    TestCaptureUpdatedRegion({capturer_.get()});
  }

#if defined(WEBRTC_WIN)
  // Enable allow_directx_capturer in DesktopCaptureOptions, but let
  // DesktopCapturer::CreateScreenCapturer() to decide whether a DirectX
  // capturer should be used.
  void MaybeCreateDirectxCapturer() {
    DesktopCaptureOptions options(DesktopCaptureOptions::CreateDefault());
    options.set_allow_directx_capturer(true);
    capturer_ = DesktopCapturer::CreateScreenCapturer(options);
  }

  bool CreateDirectxCapturer() {
    if (!ScreenCapturerWinDirectx::IsSupported()) {
      LOG(LS_WARNING) << "Directx capturer is not supported";
      return false;
    }

    MaybeCreateDirectxCapturer();
    return true;
  }

  void CreateMagnifierCapturer() {
    DesktopCaptureOptions options(DesktopCaptureOptions::CreateDefault());
    options.set_allow_use_magnification_api(true);
    capturer_ = DesktopCapturer::CreateScreenCapturer(options);
  }
#endif  // defined(WEBRTC_WIN)

  std::unique_ptr<DesktopCapturer> capturer_;
  MockDesktopCapturerCallback callback_;

 private:
  // Repeats capturing the frame by using |capturers| one-by-one for 600 times,
  // typically 30 seconds, until they succeeded captured a |color| rectangle at
  // |rect|. This function uses |drawer|->WaitForPendingDraws() between two
  // attempts to wait for the screen to update.
  void TestCaptureOneFrame(std::vector<DesktopCapturer*> capturers,
                           ScreenDrawer* drawer,
                           DesktopRect rect,
                           RgbaColor color) {
    const int wait_capture_round = 600;
    drawer->Clear();
    size_t succeeded_capturers = 0;
    for (int i = 0; i < wait_capture_round; i++) {
      drawer->DrawRectangle(rect, color);
      drawer->WaitForPendingDraws();
      for (size_t j = 0; j < capturers.size(); j++) {
        if (capturers[j] == nullptr) {
          // DesktopCapturer should return an empty updated_region() if no
          // update detected. So we won't test it again if it has captured the
          // rectangle we drew.
          continue;
        }
        std::unique_ptr<DesktopFrame> frame = CaptureFrame(capturers[j]);
        if (!frame) {
          // CaptureFrame() has triggered an assertion failure already, we only
          // need to return here.
          return;
        }

        if (ArePixelsColoredBy(
            *frame, rect, color, drawer->MayDrawIncompleteShapes())) {
          capturers[j] = nullptr;
          succeeded_capturers++;
        }
        // The following else if statement is for debugging purpose only, which
        // should be removed after flaky of ScreenCapturerIntegrationTest has
        // been resolved.
        else if (i == wait_capture_round - 1) {
          std::string result;
          rtc::Base64::EncodeFromArray(frame->data(),
                                       frame->size().height() * frame->stride(),
                                       &result);
          std::cout << frame->size().width() << " x " << frame->size().height()
                    << std::endl;
          // Split the entire string (can be over 4M) into several lines to
          // avoid browser from sticking.
          static const size_t kLineLength = 32768;
          const char* result_end = result.c_str() + result.length();
          for (const char* it = result.c_str();
               it < result_end;
               it += kLineLength) {
            const size_t max_length = result_end - it;
            std::cout << std::string(it, std::min(kLineLength, max_length))
                      << std::endl;
          }
          std::cout << "Failed to capture rectangle " << rect.left() << " x "
                    << rect.top() << " - " << rect.right() << " x "
                    << rect.bottom() << " with color ("
                    << static_cast<int>(color.red) << ", "
                    << static_cast<int>(color.green) << ", "
                    << static_cast<int>(color.blue) << ", "
                    << static_cast<int>(color.alpha) << ")" << std::endl;
          ASSERT_TRUE(false) << "ScreenCapturerIntegrationTest may be flaky. "
                                "Please kindly FYI the broken link to "
                                "zijiehe@chromium.org for investigation. If "
                                "the failure continually happens, but I have "
                                "not responded as quick as expected, disable "
                                "*all* tests in "
                                "screen_capturer_integration_test.cc to "
                                "unblock other developers.";
        }
      }

      if (succeeded_capturers == capturers.size()) {
        break;
      }
    }

    ASSERT_EQ(succeeded_capturers, capturers.size());
  }

  // Expects |capturer| to successfully capture a frame, and returns it.
  std::unique_ptr<DesktopFrame> CaptureFrame(DesktopCapturer* capturer) {
    for (int i = 0; i < 10; i++) {
      std::unique_ptr<DesktopFrame> frame;
      DesktopCapturer::Result result;
      EXPECT_CALL(callback_, OnCaptureResultPtr(_, _))
          .WillOnce(SaveCaptureResult(&result, &frame));
      capturer->CaptureFrame();
      testing::Mock::VerifyAndClearExpectations(&callback_);
      if (result == DesktopCapturer::Result::SUCCESS) {
        EXPECT_TRUE(frame);
        return frame;
      } else {
        EXPECT_FALSE(frame);
      }
    }

    EXPECT_TRUE(false);
    return nullptr;
  }
};

#if defined(WEBRTC_WIN)
// ScreenCapturerWinGdi randomly returns blank screen, the root cause is still
// unknown. Bug, https://bugs.chromium.org/p/webrtc/issues/detail?id=6843.
#define MAYBE_CaptureUpdatedRegion MANUAL_CaptureUpdatedRegion
#else
#define MAYBE_CaptureUpdatedRegion CaptureUpdatedRegion
#endif
TEST_F(ScreenCapturerIntegrationTest, MAYBE_CaptureUpdatedRegion) {
  TestCaptureUpdatedRegion();
}

#if defined(WEBRTC_WIN)
// ScreenCapturerWinGdi randomly returns blank screen, the root cause is still
// unknown. Bug, https://bugs.chromium.org/p/webrtc/issues/detail?id=6843.
#define MAYBE_TwoCapturers MANUAL_TwoCapturers
#else
#define MAYBE_TwoCapturers TwoCapturers
#endif
TEST_F(ScreenCapturerIntegrationTest, MAYBE_TwoCapturers) {
  std::unique_ptr<DesktopCapturer> capturer2 = std::move(capturer_);
  SetUp();
  TestCaptureUpdatedRegion({capturer_.get(), capturer2.get()});
}

#if defined(WEBRTC_WIN)

TEST_F(ScreenCapturerIntegrationTest, CaptureUpdatedRegionWithDirectxCapturer) {
  if (!CreateDirectxCapturer()) {
    return;
  }

  TestCaptureUpdatedRegion();
}

TEST_F(ScreenCapturerIntegrationTest, TwoDirectxCapturers) {
  if (!CreateDirectxCapturer()) {
    return;
  }

  std::unique_ptr<DesktopCapturer> capturer2 = std::move(capturer_);
  RTC_CHECK(CreateDirectxCapturer());
  TestCaptureUpdatedRegion({capturer_.get(), capturer2.get()});
}

TEST_F(ScreenCapturerIntegrationTest,
       CaptureUpdatedRegionWithMagnifierCapturer) {
  // On Windows 8 or later, magnifier APIs return a frame with a border on test
  // environment, so disable these tests.
  // Bug https://bugs.chromium.org/p/webrtc/issues/detail?id=6844
  // TODO(zijiehe): Find the root cause of the border and failure, which cannot
  // reproduce on my dev machine.
  if (rtc::IsWindows8OrLater()) {
    return;
  }
  CreateMagnifierCapturer();
  TestCaptureUpdatedRegion();
}

TEST_F(ScreenCapturerIntegrationTest, TwoMagnifierCapturers) {
  // On Windows 8 or later, magnifier APIs return a frame with a border on test
  // environment, so disable these tests.
  // Bug https://bugs.chromium.org/p/webrtc/issues/detail?id=6844
  // TODO(zijiehe): Find the root cause of the border and failure, which cannot
  // reproduce on my dev machine.
  if (rtc::IsWindows8OrLater()) {
    return;
  }
  CreateMagnifierCapturer();
  std::unique_ptr<DesktopCapturer> capturer2 = std::move(capturer_);
  CreateMagnifierCapturer();
  TestCaptureUpdatedRegion({capturer_.get(), capturer2.get()});
}

TEST_F(ScreenCapturerIntegrationTest,
       MaybeCaptureUpdatedRegionWithDirectxCapturer) {
  if (!rtc::IsWindows8OrLater()) {
    // ScreenCapturerWinGdi randomly returns blank screen, the root cause is
    // still unknown. Bug,
    // https://bugs.chromium.org/p/webrtc/issues/detail?id=6843.
    // On Windows 7 or early version, MaybeCreateDirectxCapturer() always
    // creates GDI capturer.
    return;
  }
  // Even DirectX capturer is not supported in current system, we should be able
  // to select a usable capturer.
  MaybeCreateDirectxCapturer();
  TestCaptureUpdatedRegion();
}

#endif  // defined(WEBRTC_WIN)

}  // namespace webrtc
<commit_msg>Change MANUAL -> DISABLED for ScreenCapturerIntegrationTest tests<commit_after>/*
 *  Copyright (c) 2016 The WebRTC project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <string.h>

#include <algorithm>
#include <initializer_list>
#include <iostream>  // TODO(zijiehe): Remove once flaky has been resolved.
#include <memory>
#include <utility>

// TODO(zijiehe): Remove once flaky has been resolved.
#include "webrtc/base/base64.h"
#include "webrtc/base/checks.h"
#include "webrtc/base/constructormagic.h"
#include "webrtc/base/logging.h"
#include "webrtc/modules/desktop_capture/desktop_capturer.h"
#include "webrtc/modules/desktop_capture/desktop_capture_options.h"
#include "webrtc/modules/desktop_capture/desktop_frame.h"
#include "webrtc/modules/desktop_capture/desktop_region.h"
#include "webrtc/modules/desktop_capture/mock_desktop_capturer_callback.h"
#include "webrtc/modules/desktop_capture/rgba_color.h"
#include "webrtc/modules/desktop_capture/screen_drawer.h"
#include "webrtc/test/gmock.h"
#include "webrtc/test/gtest.h"

#if defined(WEBRTC_WIN)
#include "webrtc/base/win32.h"
#include "webrtc/modules/desktop_capture/win/screen_capturer_win_directx.h"
#endif  // defined(WEBRTC_WIN)

using ::testing::_;

namespace webrtc {

namespace {

ACTION_P2(SaveCaptureResult, result, dest) {
  *result = arg0;
  *dest = std::move(*arg1);
}

// Returns true if color in |rect| of |frame| is |color|.
bool ArePixelsColoredBy(const DesktopFrame& frame,
                        DesktopRect rect,
                        RgbaColor color,
                        bool may_partially_draw) {
  if (!may_partially_draw) {
    // updated_region() should cover the painted area.
    DesktopRegion updated_region(frame.updated_region());
    updated_region.IntersectWith(rect);
    if (!updated_region.Equals(DesktopRegion(rect))) {
      return false;
    }
  }

  // Color in the |rect| should be |color|.
  uint8_t* row = frame.GetFrameDataAtPos(rect.top_left());
  for (int i = 0; i < rect.height(); i++) {
    uint8_t* column = row;
    for (int j = 0; j < rect.width(); j++) {
      if (color != RgbaColor(column)) {
        return false;
      }
      column += DesktopFrame::kBytesPerPixel;
    }
    row += frame.stride();
  }
  return true;
}

}  // namespace

class ScreenCapturerIntegrationTest : public testing::Test {
 public:
  void SetUp() override {
    capturer_ = DesktopCapturer::CreateScreenCapturer(
        DesktopCaptureOptions::CreateDefault());
  }

 protected:
  void TestCaptureUpdatedRegion(
      std::initializer_list<DesktopCapturer*> capturers) {
    RTC_DCHECK(capturers.size() > 0);
    // A large enough area for the tests, which should be able to be fulfilled
    // by most systems.
#if defined(WEBRTC_WIN)
    // On Windows, an interesting warning window may pop up randomly. The root
    // cause is still under investigation, so reduce the test area to work
    // around. Bug https://bugs.chromium.org/p/webrtc/issues/detail?id=6666.
    const int kTestArea = 416;
#else
    const int kTestArea = 512;
#endif
    const int kRectSize = 32;
    std::unique_ptr<ScreenDrawer> drawer = ScreenDrawer::Create();
    if (!drawer || drawer->DrawableRegion().is_empty()) {
      LOG(LS_WARNING) << "No ScreenDrawer implementation for current platform.";
      return;
    }
    if (drawer->DrawableRegion().width() < kTestArea ||
        drawer->DrawableRegion().height() < kTestArea) {
      LOG(LS_WARNING) << "ScreenDrawer::DrawableRegion() is too small for the "
                         "CaptureUpdatedRegion tests.";
      return;
    }

    for (DesktopCapturer* capturer : capturers) {
      capturer->Start(&callback_);
    }

    // Draw a set of |kRectSize| by |kRectSize| rectangles at (|i|, |i|), or
    // |i| by |i| rectangles at (|kRectSize|, |kRectSize|). One of (controlled
    // by |c|) its primary colors is |i|, and the other two are 0x7f. So we
    // won't draw a black or white rectangle.
    for (int c = 0; c < 3; c++) {
      // A fixed size rectangle.
      for (int i = 0; i < kTestArea - kRectSize; i += 16) {
        DesktopRect rect = DesktopRect::MakeXYWH(i, i, kRectSize, kRectSize);
        rect.Translate(drawer->DrawableRegion().top_left());
        RgbaColor color((c == 0 ? (i & 0xff) : 0x7f),
                        (c == 1 ? (i & 0xff) : 0x7f),
                        (c == 2 ? (i & 0xff) : 0x7f));
        // Fail fast.
        ASSERT_NO_FATAL_FAILURE(
            TestCaptureOneFrame(capturers, drawer.get(), rect, color));
      }

      // A variable-size rectangle.
      for (int i = 0; i < kTestArea - kRectSize; i += 16) {
        DesktopRect rect = DesktopRect::MakeXYWH(kRectSize, kRectSize, i, i);
        rect.Translate(drawer->DrawableRegion().top_left());
        RgbaColor color((c == 0 ? (i & 0xff) : 0x7f),
                        (c == 1 ? (i & 0xff) : 0x7f),
                        (c == 2 ? (i & 0xff) : 0x7f));
        // Fail fast.
        ASSERT_NO_FATAL_FAILURE(
            TestCaptureOneFrame(capturers, drawer.get(), rect, color));
      }
    }
  }

  void TestCaptureUpdatedRegion() {
    TestCaptureUpdatedRegion({capturer_.get()});
  }

#if defined(WEBRTC_WIN)
  // Enable allow_directx_capturer in DesktopCaptureOptions, but let
  // DesktopCapturer::CreateScreenCapturer() to decide whether a DirectX
  // capturer should be used.
  void MaybeCreateDirectxCapturer() {
    DesktopCaptureOptions options(DesktopCaptureOptions::CreateDefault());
    options.set_allow_directx_capturer(true);
    capturer_ = DesktopCapturer::CreateScreenCapturer(options);
  }

  bool CreateDirectxCapturer() {
    if (!ScreenCapturerWinDirectx::IsSupported()) {
      LOG(LS_WARNING) << "Directx capturer is not supported";
      return false;
    }

    MaybeCreateDirectxCapturer();
    return true;
  }

  void CreateMagnifierCapturer() {
    DesktopCaptureOptions options(DesktopCaptureOptions::CreateDefault());
    options.set_allow_use_magnification_api(true);
    capturer_ = DesktopCapturer::CreateScreenCapturer(options);
  }
#endif  // defined(WEBRTC_WIN)

  std::unique_ptr<DesktopCapturer> capturer_;
  MockDesktopCapturerCallback callback_;

 private:
  // Repeats capturing the frame by using |capturers| one-by-one for 600 times,
  // typically 30 seconds, until they succeeded captured a |color| rectangle at
  // |rect|. This function uses |drawer|->WaitForPendingDraws() between two
  // attempts to wait for the screen to update.
  void TestCaptureOneFrame(std::vector<DesktopCapturer*> capturers,
                           ScreenDrawer* drawer,
                           DesktopRect rect,
                           RgbaColor color) {
    const int wait_capture_round = 600;
    drawer->Clear();
    size_t succeeded_capturers = 0;
    for (int i = 0; i < wait_capture_round; i++) {
      drawer->DrawRectangle(rect, color);
      drawer->WaitForPendingDraws();
      for (size_t j = 0; j < capturers.size(); j++) {
        if (capturers[j] == nullptr) {
          // DesktopCapturer should return an empty updated_region() if no
          // update detected. So we won't test it again if it has captured the
          // rectangle we drew.
          continue;
        }
        std::unique_ptr<DesktopFrame> frame = CaptureFrame(capturers[j]);
        if (!frame) {
          // CaptureFrame() has triggered an assertion failure already, we only
          // need to return here.
          return;
        }

        if (ArePixelsColoredBy(
            *frame, rect, color, drawer->MayDrawIncompleteShapes())) {
          capturers[j] = nullptr;
          succeeded_capturers++;
        }
        // The following else if statement is for debugging purpose only, which
        // should be removed after flaky of ScreenCapturerIntegrationTest has
        // been resolved.
        else if (i == wait_capture_round - 1) {
          std::string result;
          rtc::Base64::EncodeFromArray(frame->data(),
                                       frame->size().height() * frame->stride(),
                                       &result);
          std::cout << frame->size().width() << " x " << frame->size().height()
                    << std::endl;
          // Split the entire string (can be over 4M) into several lines to
          // avoid browser from sticking.
          static const size_t kLineLength = 32768;
          const char* result_end = result.c_str() + result.length();
          for (const char* it = result.c_str();
               it < result_end;
               it += kLineLength) {
            const size_t max_length = result_end - it;
            std::cout << std::string(it, std::min(kLineLength, max_length))
                      << std::endl;
          }
          std::cout << "Failed to capture rectangle " << rect.left() << " x "
                    << rect.top() << " - " << rect.right() << " x "
                    << rect.bottom() << " with color ("
                    << static_cast<int>(color.red) << ", "
                    << static_cast<int>(color.green) << ", "
                    << static_cast<int>(color.blue) << ", "
                    << static_cast<int>(color.alpha) << ")" << std::endl;
          ASSERT_TRUE(false) << "ScreenCapturerIntegrationTest may be flaky. "
                                "Please kindly FYI the broken link to "
                                "zijiehe@chromium.org for investigation. If "
                                "the failure continually happens, but I have "
                                "not responded as quick as expected, disable "
                                "*all* tests in "
                                "screen_capturer_integration_test.cc to "
                                "unblock other developers.";
        }
      }

      if (succeeded_capturers == capturers.size()) {
        break;
      }
    }

    ASSERT_EQ(succeeded_capturers, capturers.size());
  }

  // Expects |capturer| to successfully capture a frame, and returns it.
  std::unique_ptr<DesktopFrame> CaptureFrame(DesktopCapturer* capturer) {
    for (int i = 0; i < 10; i++) {
      std::unique_ptr<DesktopFrame> frame;
      DesktopCapturer::Result result;
      EXPECT_CALL(callback_, OnCaptureResultPtr(_, _))
          .WillOnce(SaveCaptureResult(&result, &frame));
      capturer->CaptureFrame();
      testing::Mock::VerifyAndClearExpectations(&callback_);
      if (result == DesktopCapturer::Result::SUCCESS) {
        EXPECT_TRUE(frame);
        return frame;
      } else {
        EXPECT_FALSE(frame);
      }
    }

    EXPECT_TRUE(false);
    return nullptr;
  }
};

#if defined(WEBRTC_WIN)
// ScreenCapturerWinGdi randomly returns blank screen, the root cause is still
// unknown. Bug, https://bugs.chromium.org/p/webrtc/issues/detail?id=6843.
#define MAYBE_CaptureUpdatedRegion DISABLED_CaptureUpdatedRegion
#else
#define MAYBE_CaptureUpdatedRegion CaptureUpdatedRegion
#endif
TEST_F(ScreenCapturerIntegrationTest, MAYBE_CaptureUpdatedRegion) {
  TestCaptureUpdatedRegion();
}

#if defined(WEBRTC_WIN)
// ScreenCapturerWinGdi randomly returns blank screen, the root cause is still
// unknown. Bug, https://bugs.chromium.org/p/webrtc/issues/detail?id=6843.
#define MAYBE_TwoCapturers DISABLED_TwoCapturers
#else
#define MAYBE_TwoCapturers TwoCapturers
#endif
TEST_F(ScreenCapturerIntegrationTest, MAYBE_TwoCapturers) {
  std::unique_ptr<DesktopCapturer> capturer2 = std::move(capturer_);
  SetUp();
  TestCaptureUpdatedRegion({capturer_.get(), capturer2.get()});
}

#if defined(WEBRTC_WIN)

TEST_F(ScreenCapturerIntegrationTest, CaptureUpdatedRegionWithDirectxCapturer) {
  if (!CreateDirectxCapturer()) {
    return;
  }

  TestCaptureUpdatedRegion();
}

TEST_F(ScreenCapturerIntegrationTest, TwoDirectxCapturers) {
  if (!CreateDirectxCapturer()) {
    return;
  }

  std::unique_ptr<DesktopCapturer> capturer2 = std::move(capturer_);
  RTC_CHECK(CreateDirectxCapturer());
  TestCaptureUpdatedRegion({capturer_.get(), capturer2.get()});
}

TEST_F(ScreenCapturerIntegrationTest,
       CaptureUpdatedRegionWithMagnifierCapturer) {
  // On Windows 8 or later, magnifier APIs return a frame with a border on test
  // environment, so disable these tests.
  // Bug https://bugs.chromium.org/p/webrtc/issues/detail?id=6844
  // TODO(zijiehe): Find the root cause of the border and failure, which cannot
  // reproduce on my dev machine.
  if (rtc::IsWindows8OrLater()) {
    return;
  }
  CreateMagnifierCapturer();
  TestCaptureUpdatedRegion();
}

TEST_F(ScreenCapturerIntegrationTest, TwoMagnifierCapturers) {
  // On Windows 8 or later, magnifier APIs return a frame with a border on test
  // environment, so disable these tests.
  // Bug https://bugs.chromium.org/p/webrtc/issues/detail?id=6844
  // TODO(zijiehe): Find the root cause of the border and failure, which cannot
  // reproduce on my dev machine.
  if (rtc::IsWindows8OrLater()) {
    return;
  }
  CreateMagnifierCapturer();
  std::unique_ptr<DesktopCapturer> capturer2 = std::move(capturer_);
  CreateMagnifierCapturer();
  TestCaptureUpdatedRegion({capturer_.get(), capturer2.get()});
}

TEST_F(ScreenCapturerIntegrationTest,
       MaybeCaptureUpdatedRegionWithDirectxCapturer) {
  if (!rtc::IsWindows8OrLater()) {
    // ScreenCapturerWinGdi randomly returns blank screen, the root cause is
    // still unknown. Bug,
    // https://bugs.chromium.org/p/webrtc/issues/detail?id=6843.
    // On Windows 7 or early version, MaybeCreateDirectxCapturer() always
    // creates GDI capturer.
    return;
  }
  // Even DirectX capturer is not supported in current system, we should be able
  // to select a usable capturer.
  MaybeCreateDirectxCapturer();
  TestCaptureUpdatedRegion();
}

#endif  // defined(WEBRTC_WIN)

}  // namespace webrtc
<|endoftext|>


<commit_before>/**
 *  Copyright (C) 2015 Topology LP
 *
 *  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 <bonefish/rawsocket/tcp_connection.hpp>
#include <bonefish/trace/trace.hpp>

#include <boost/asio/read.hpp>
#include <boost/asio/write.hpp>
#include <iostream>

namespace bonefish {

tcp_connection::tcp_connection(boost::asio::ip::tcp::socket&& socket)
    : rawsocket_connection()
    , m_capabilities(0)
    , m_message_length(0)
    , m_message_buffer()
    , m_socket(std::move(socket))
{
    // Disable Nagle algorithm to get lower latency (and lower throughput).
    m_socket.set_option(boost::asio::ip::tcp::no_delay(true));
}

tcp_connection::~tcp_connection()
{
    if (m_socket.is_open()) {
        m_socket.close();
    }
}

void tcp_connection::async_handshake()
{
    auto buffer = boost::asio::buffer(&m_capabilities, sizeof(m_capabilities));

    std::weak_ptr<tcp_connection> weak_self =
            std::static_pointer_cast<tcp_connection>(shared_from_this());

    auto handler = [weak_self](
            const boost::system::error_code& error_code, size_t bytes_transferred) {
        auto shared_self = weak_self.lock();
        if (shared_self) {
            shared_self->receive_handshake_handler(error_code, bytes_transferred);
        }
    };

    boost::asio::async_read(m_socket, buffer, handler);
}

void tcp_connection::async_receive()
{
    auto buffer = boost::asio::buffer(&m_message_length, sizeof(m_message_length));

    std::weak_ptr<tcp_connection> weak_self =
            std::static_pointer_cast<tcp_connection>(shared_from_this());

    auto handler = [weak_self](
            const boost::system::error_code& error_code, size_t bytes_transferred) {
        auto shared_self = weak_self.lock();
        if (shared_self) {
            shared_self->receive_message_header_handler(error_code, bytes_transferred);
        }
    };

    // We cannot start receiving messages until the initial
    // handshake has allowed us to exchange capabilities.
    assert(m_capabilities != 0);

    boost::asio::async_read(m_socket, buffer, handler);
}

bool tcp_connection::send_handshake(uint32_t capabilities)
{
    boost::system::error_code error_code;
    boost::asio::write(m_socket,
            boost::asio::buffer(&capabilities, sizeof(capabilities)), error_code);

    if (error_code) {
        handle_system_error(error_code);
        return false;
    }

    return true;
}

bool tcp_connection::send_message(const char* message, size_t length)
{
    boost::system::error_code error_code;

    // First write the message length prefix.
    uint32_t length_prefix = htonl(length);
    boost::asio::write(m_socket, boost::asio::buffer(&length_prefix, sizeof(length_prefix)), error_code);
    if (error_code) {
        handle_system_error(error_code);
        return false;
    }

    // Then write the actual message.
    boost::asio::write(m_socket, boost::asio::buffer(message, length), error_code);
    if (error_code) {
        handle_system_error(error_code);
        return false;
    }

    return true;
}

void tcp_connection::receive_handshake_handler(
        const boost::system::error_code& error_code, size_t bytes_transferred)
{
    if (error_code) {
        return handle_system_error(error_code);
    }

    assert(bytes_transferred == sizeof(m_capabilities));
    const auto& handshake_handler = get_handshake_handler();
    handshake_handler(shared_from_this(), ntohl(m_capabilities));
}

void tcp_connection::receive_message_header_handler(
        const boost::system::error_code& error_code, size_t bytes_transferred)
{
    if (error_code) {
        return handle_system_error(error_code);
    }

    // We cannot be guaranteed that a client implementation won't accidentally
    // introduce this protocol violation. In the event that we ever encounter
    // a message that reports a zero length we fail that connection gracefully.
    assert(m_message_length != 0);
    if (m_message_length == 0) {
        BONEFISH_TRACE("invalid message length: %1%", m_message_length);
        const auto& fail_handler = get_fail_handler();
        fail_handler(shared_from_this(), "invalid message length");
    }

    m_message_buffer.reserve(ntohl(m_message_length));

    auto buffer = boost::asio::buffer(m_message_buffer.data(), ntohl(m_message_length));

    std::weak_ptr<tcp_connection> weak_self =
            std::static_pointer_cast<tcp_connection>(shared_from_this());

    auto handler = [weak_self](
            const boost::system::error_code& error_code, size_t bytes_transferred) {
        auto shared_self = weak_self.lock();
        if (shared_self) {
            shared_self->receive_message_body_handler(error_code, bytes_transferred);
        }
    };

    boost::asio::async_read(m_socket, buffer, handler);
}

void tcp_connection::receive_message_body_handler(
        const boost::system::error_code& error_code, size_t bytes_transferred)
{
    if (error_code) {
        return handle_system_error(error_code);
    }

    const auto& message_handler = get_message_handler();
    assert(message_handler);
    message_handler(shared_from_this(), m_message_buffer.data(), bytes_transferred);

    async_receive();
}

void tcp_connection::handle_system_error(const boost::system::error_code& error_code)
{
    // NOTE: The boost documentation does not indicate what all of the possible error
    //       codes are that can occur for the async receive handlers. So it will be an
    //       ongoing exercise in trying to figure this out.
    if (error_code == boost::asio::error::eof) {
        BONEFISH_TRACE("connection closed: %1%", error_code);
        const auto& close_handler = get_close_handler();
        close_handler(shared_from_this());
    } else if (error_code != boost::asio::error::operation_aborted) {
        BONEFISH_TRACE("connection failed: %1%", error_code);
        const auto& fail_handler = get_fail_handler();
        fail_handler(shared_from_this(), error_code.message().c_str());
    } else {
        BONEFISH_TRACE("unhandled system error: %1%", error_code);
        assert(0);
    }
}

} // namespace bonefish
<commit_msg>cleanly handle 0-length messages instead of letting the client crash the server on the assert<commit_after>/**
 *  Copyright (C) 2015 Topology LP
 *
 *  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 <bonefish/rawsocket/tcp_connection.hpp>
#include <bonefish/trace/trace.hpp>

#include <boost/asio/read.hpp>
#include <boost/asio/write.hpp>
#include <iostream>

namespace bonefish {

tcp_connection::tcp_connection(boost::asio::ip::tcp::socket&& socket)
    : rawsocket_connection()
    , m_capabilities(0)
    , m_message_length(0)
    , m_message_buffer()
    , m_socket(std::move(socket))
{
    // Disable Nagle algorithm to get lower latency (and lower throughput).
    m_socket.set_option(boost::asio::ip::tcp::no_delay(true));
}

tcp_connection::~tcp_connection()
{
    if (m_socket.is_open()) {
        m_socket.close();
    }
}

void tcp_connection::async_handshake()
{
    auto buffer = boost::asio::buffer(&m_capabilities, sizeof(m_capabilities));

    std::weak_ptr<tcp_connection> weak_self =
            std::static_pointer_cast<tcp_connection>(shared_from_this());

    auto handler = [weak_self](
            const boost::system::error_code& error_code, size_t bytes_transferred) {
        auto shared_self = weak_self.lock();
        if (shared_self) {
            shared_self->receive_handshake_handler(error_code, bytes_transferred);
        }
    };

    boost::asio::async_read(m_socket, buffer, handler);
}

void tcp_connection::async_receive()
{
    auto buffer = boost::asio::buffer(&m_message_length, sizeof(m_message_length));

    std::weak_ptr<tcp_connection> weak_self =
            std::static_pointer_cast<tcp_connection>(shared_from_this());

    auto handler = [weak_self](
            const boost::system::error_code& error_code, size_t bytes_transferred) {
        auto shared_self = weak_self.lock();
        if (shared_self) {
            shared_self->receive_message_header_handler(error_code, bytes_transferred);
        }
    };

    // We cannot start receiving messages until the initial
    // handshake has allowed us to exchange capabilities.
    assert(m_capabilities != 0);

    boost::asio::async_read(m_socket, buffer, handler);
}

bool tcp_connection::send_handshake(uint32_t capabilities)
{
    boost::system::error_code error_code;
    boost::asio::write(m_socket,
            boost::asio::buffer(&capabilities, sizeof(capabilities)), error_code);

    if (error_code) {
        handle_system_error(error_code);
        return false;
    }

    return true;
}

bool tcp_connection::send_message(const char* message, size_t length)
{
    boost::system::error_code error_code;

    // First write the message length prefix.
    uint32_t length_prefix = htonl(length);
    boost::asio::write(m_socket, boost::asio::buffer(&length_prefix, sizeof(length_prefix)), error_code);
    if (error_code) {
        handle_system_error(error_code);
        return false;
    }

    // Then write the actual message.
    boost::asio::write(m_socket, boost::asio::buffer(message, length), error_code);
    if (error_code) {
        handle_system_error(error_code);
        return false;
    }

    return true;
}

void tcp_connection::receive_handshake_handler(
        const boost::system::error_code& error_code, size_t bytes_transferred)
{
    if (error_code) {
        return handle_system_error(error_code);
    }

    assert(bytes_transferred == sizeof(m_capabilities));
    const auto& handshake_handler = get_handshake_handler();
    handshake_handler(shared_from_this(), ntohl(m_capabilities));
}

void tcp_connection::receive_message_header_handler(
        const boost::system::error_code& error_code, size_t bytes_transferred)
{
    if (error_code) {
        return handle_system_error(error_code);
    }

    // We cannot be guaranteed that a client implementation won't accidentally
    // introduce this protocol violation. In the event that we ever encounter
    // a message that reports a zero length we fail that connection gracefully.
    if (m_message_length == 0) {
        BONEFISH_TRACE("invalid message length: %1%", m_message_length);
        const auto& fail_handler = get_fail_handler();
        fail_handler(shared_from_this(), "invalid message length");
        return;
    }

    m_message_buffer.reserve(ntohl(m_message_length));

    auto buffer = boost::asio::buffer(m_message_buffer.data(), ntohl(m_message_length));

    std::weak_ptr<tcp_connection> weak_self =
            std::static_pointer_cast<tcp_connection>(shared_from_this());

    auto handler = [weak_self](
            const boost::system::error_code& error_code, size_t bytes_transferred) {
        auto shared_self = weak_self.lock();
        if (shared_self) {
            shared_self->receive_message_body_handler(error_code, bytes_transferred);
        }
    };

    boost::asio::async_read(m_socket, buffer, handler);
}

void tcp_connection::receive_message_body_handler(
        const boost::system::error_code& error_code, size_t bytes_transferred)
{
    if (error_code) {
        return handle_system_error(error_code);
    }

    const auto& message_handler = get_message_handler();
    assert(message_handler);
    message_handler(shared_from_this(), m_message_buffer.data(), bytes_transferred);

    async_receive();
}

void tcp_connection::handle_system_error(const boost::system::error_code& error_code)
{
    // NOTE: The boost documentation does not indicate what all of the possible error
    //       codes are that can occur for the async receive handlers. So it will be an
    //       ongoing exercise in trying to figure this out.
    if (error_code == boost::asio::error::eof) {
        BONEFISH_TRACE("connection closed: %1%", error_code);
        const auto& close_handler = get_close_handler();
        close_handler(shared_from_this());
    } else if (error_code != boost::asio::error::operation_aborted) {
        BONEFISH_TRACE("connection failed: %1%", error_code);
        const auto& fail_handler = get_fail_handler();
        fail_handler(shared_from_this(), error_code.message().c_str());
    } else {
        BONEFISH_TRACE("unhandled system error: %1%", error_code);
        assert(0);
    }
}

} // namespace bonefish
<|endoftext|>


<commit_before>// Copyright 2013 Yangqing Jia

#include <cstring>
#include <cuda_runtime.h>

#include "gtest/gtest.h"
#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/filler.hpp"
#include "caffe/vision_layers.hpp"
#include "caffe/test/test_gradient_check_util.hpp"

#include "caffe/test/test_caffe_main.hpp"

namespace caffe {

extern cudaDeviceProp CAFFE_TEST_CUDA_PROP;

template <typename Dtype>
class ConvolutionLayerTest : public ::testing::Test {
 protected:
  ConvolutionLayerTest()
      : blob_bottom_(new Blob<Dtype>()),
        blob_top_(new Blob<Dtype>()) {};
  virtual void SetUp() {
    blob_bottom_->Reshape(2, 3, 6, 5);
    // fill the values
    FillerParameter filler_param;
    filler_param.set_value(1.);
    GaussianFiller<Dtype> filler(filler_param);
    filler.Fill(this->blob_bottom_);
    blob_bottom_vec_.push_back(blob_bottom_);
    blob_top_vec_.push_back(blob_top_);
  };

  virtual ~ConvolutionLayerTest() { delete blob_bottom_; delete blob_top_; }
  Blob<Dtype>* const blob_bottom_;
  Blob<Dtype>* const blob_top_;
  vector<Blob<Dtype>*> blob_bottom_vec_;
  vector<Blob<Dtype>*> blob_top_vec_;
};

typedef ::testing::Types<float, double> Dtypes;
TYPED_TEST_CASE(ConvolutionLayerTest, Dtypes);

TYPED_TEST(ConvolutionLayerTest, TestSetup) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(4);
  shared_ptr<Layer<TypeParam> > layer(
      new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  EXPECT_EQ(this->blob_top_->num(), 2);
  EXPECT_EQ(this->blob_top_->channels(), 4);
  EXPECT_EQ(this->blob_top_->height(), 2);
  EXPECT_EQ(this->blob_top_->width(), 2);
  // setting group should not change the shape
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer.reset(new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  EXPECT_EQ(this->blob_top_->num(), 2);
  EXPECT_EQ(this->blob_top_->channels(), 3);
  EXPECT_EQ(this->blob_top_->height(), 2);
  EXPECT_EQ(this->blob_top_->width(), 2);
}

TYPED_TEST(ConvolutionLayerTest, TestSimpleConvolution) {
  // We will simply see if the convolution layer carries out averaging well.
  FillerParameter filler_param;
  filler_param.set_value(1.);
  ConstantFiller<TypeParam> filler(filler_param);
  filler.Fill(this->blob_bottom_);
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(4);
  layer_param.mutable_weight_filler()->set_type("constant");
  layer_param.mutable_weight_filler()->set_value(1);
  layer_param.mutable_bias_filler()->set_type("constant");
  layer_param.mutable_bias_filler()->set_value(0.1);
  shared_ptr<Layer<TypeParam> > layer(
      new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  Caffe::set_mode(Caffe::CPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 27.1
  const TypeParam* top_data = this->blob_top_->cpu_data();
  for (int i = 0; i < this->blob_top_->count(); ++i) {
    EXPECT_GE(top_data[i], 27.1 - 1e-4);
    EXPECT_LE(top_data[i], 27.1 + 1e-4);
  }
  // Test GPU
  Caffe::set_mode(Caffe::GPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 27.1
  top_data = this->blob_top_->cpu_data();
  for (int i = 0; i < this->blob_top_->count(); ++i) {
    EXPECT_GE(top_data[i], 27.1 - 1e-4);
    EXPECT_LE(top_data[i], 27.1 + 1e-4);
  }
}

TYPED_TEST(ConvolutionLayerTest, TestSimpleConvolutionGroup) {
  // We will simply see if the convolution layer carries out averaging well.
  FillerParameter filler_param;
  filler_param.set_value(1.);
  ConstantFiller<TypeParam> filler(filler_param);
  filler.Fill(this->blob_bottom_);
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer_param.mutable_weight_filler()->set_type("constant");
  layer_param.mutable_weight_filler()->set_value(1);
  layer_param.mutable_bias_filler()->set_type("constant");
  layer_param.mutable_bias_filler()->set_value(0.1);
  shared_ptr<Layer<TypeParam> > layer(
      new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  Caffe::set_mode(Caffe::CPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 9.1
  const TypeParam* top_data = this->blob_top_->cpu_data();
  for (int i = 0; i < this->blob_top_->count(); ++i) {
    EXPECT_GE(top_data[i], 9.1 - 1e-4);
    EXPECT_LE(top_data[i], 9.1 + 1e-4);
  }
  // Test GPU
  Caffe::set_mode(Caffe::GPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 9.1
  top_data = this->blob_top_->cpu_data();
  for (int i = 0; i < this->blob_top_->count(); ++i) {
    EXPECT_GE(top_data[i], 9.1 - 1e-4);
    EXPECT_LE(top_data[i], 9.1 + 1e-4);
  }
}

TYPED_TEST(ConvolutionLayerTest, TestCPUGradient) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(2);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::CPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

TYPED_TEST(ConvolutionLayerTest, TestCPUGradientGroup) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::CPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

TYPED_TEST(ConvolutionLayerTest, TestGPUGradient) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(2);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::GPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

TYPED_TEST(ConvolutionLayerTest, TestGPUGradientGroup) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::GPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

}
<commit_msg>test convolution more thoroughly<commit_after>// Copyright 2013 Yangqing Jia

#include <cstring>
#include <cuda_runtime.h>

#include "gtest/gtest.h"
#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/filler.hpp"
#include "caffe/vision_layers.hpp"
#include "caffe/test/test_gradient_check_util.hpp"

#include "caffe/test/test_caffe_main.hpp"

namespace caffe {

extern cudaDeviceProp CAFFE_TEST_CUDA_PROP;

template <typename Dtype>
class ConvolutionLayerTest : public ::testing::Test {
 protected:
  ConvolutionLayerTest()
      : blob_bottom_(new Blob<Dtype>()),
        blob_top_(new Blob<Dtype>()) {};
  virtual void SetUp() {
    blob_bottom_->Reshape(2, 3, 6, 5);
    // fill the values
    FillerParameter filler_param;
    filler_param.set_value(1.);
    GaussianFiller<Dtype> filler(filler_param);
    filler.Fill(this->blob_bottom_);
    blob_bottom_vec_.push_back(blob_bottom_);
    blob_top_vec_.push_back(blob_top_);
  };

  virtual ~ConvolutionLayerTest() { delete blob_bottom_; delete blob_top_; }
  Blob<Dtype>* const blob_bottom_;
  Blob<Dtype>* const blob_top_;
  vector<Blob<Dtype>*> blob_bottom_vec_;
  vector<Blob<Dtype>*> blob_top_vec_;
};

typedef ::testing::Types<float, double> Dtypes;
TYPED_TEST_CASE(ConvolutionLayerTest, Dtypes);

TYPED_TEST(ConvolutionLayerTest, TestSetup) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(4);
  shared_ptr<Layer<TypeParam> > layer(
      new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  EXPECT_EQ(this->blob_top_->num(), 2);
  EXPECT_EQ(this->blob_top_->channels(), 4);
  EXPECT_EQ(this->blob_top_->height(), 2);
  EXPECT_EQ(this->blob_top_->width(), 2);
  // setting group should not change the shape
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer.reset(new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  EXPECT_EQ(this->blob_top_->num(), 2);
  EXPECT_EQ(this->blob_top_->channels(), 3);
  EXPECT_EQ(this->blob_top_->height(), 2);
  EXPECT_EQ(this->blob_top_->width(), 2);
}

TYPED_TEST(ConvolutionLayerTest, TestSimpleConvolution) {
  // We will simply see if the convolution layer carries out averaging well.
  FillerParameter filler_param;
  filler_param.set_value(1.);
  ConstantFiller<TypeParam> filler(filler_param);
  filler.Fill(this->blob_bottom_);
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(4);
  layer_param.mutable_weight_filler()->set_type("constant");
  layer_param.mutable_weight_filler()->set_value(1);
  layer_param.mutable_bias_filler()->set_type("constant");
  layer_param.mutable_bias_filler()->set_value(0.1);
  shared_ptr<Layer<TypeParam> > layer(
      new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  Caffe::set_mode(Caffe::CPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 27.1
  const TypeParam* top_data = this->blob_top_->cpu_data();
  for (int i = 0; i < this->blob_top_->count(); ++i) {
    EXPECT_GE(top_data[i], 27.1 - 1e-4);
    EXPECT_LE(top_data[i], 27.1 + 1e-4);
  }
  // Test GPU
  Caffe::set_mode(Caffe::GPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 27.1
  top_data = this->blob_top_->cpu_data();
  for (int i = 0; i < this->blob_top_->count(); ++i) {
    EXPECT_GE(top_data[i], 27.1 - 1e-4);
    EXPECT_LE(top_data[i], 27.1 + 1e-4);
  }
}

TYPED_TEST(ConvolutionLayerTest, TestSimpleConvolutionGroup) {
  // We will simply see if the convolution layer carries out averaging well.
  FillerParameter filler_param;
  filler_param.set_value(1.);
  ConstantFiller<TypeParam> filler(filler_param);
  filler.Fill(this->blob_bottom_);
  TypeParam* bottom_data = this->blob_bottom_->mutable_cpu_data();
  for (int n = 0; n < this->blob_bottom_->num(); ++n) {
    for (int c = 0; c < this->blob_bottom_->channels(); ++c) {
      for (int h = 0; h < this->blob_bottom_->height(); ++h) {
        for (int w = 0; w < this->blob_bottom_->width(); ++w) {
          bottom_data[this->blob_bottom_->offset(n, c, h, w)] = c;
        }
      }
    }
  }
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer_param.mutable_weight_filler()->set_type("constant");
  layer_param.mutable_weight_filler()->set_value(1);
  layer_param.mutable_bias_filler()->set_type("constant");
  layer_param.mutable_bias_filler()->set_value(0.1);
  shared_ptr<Layer<TypeParam> > layer(
      new ConvolutionLayer<TypeParam>(layer_param));
  layer->SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
  Caffe::set_mode(Caffe::CPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 9.1
  const TypeParam* top_data = this->blob_top_->cpu_data();
  for (int n = 0; n < this->blob_top_->num(); ++n) {
    for (int c = 0; c < this->blob_top_->channels(); ++c) {
      for (int h = 0; h < this->blob_top_->height(); ++h) {
        for (int w = 0; w < this->blob_top_->width(); ++w) {
          TypeParam data = top_data[this->blob_top_->offset(n, c, h, w)];
          EXPECT_GE(data, c * 9 + 0.1 - 1e-4);
          EXPECT_LE(data, c * 9 + 0.1 + 1e-4);
        }
      }
    }
  }
  // Test GPU
  Caffe::set_mode(Caffe::GPU);
  layer->Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
  // After the convolution, the output should all have output values 9.1
  top_data = this->blob_top_->cpu_data();
  for (int n = 0; n < this->blob_top_->num(); ++n) {
    for (int c = 0; c < this->blob_top_->channels(); ++c) {
      for (int h = 0; h < this->blob_top_->height(); ++h) {
        for (int w = 0; w < this->blob_top_->width(); ++w) {
          TypeParam data = top_data[this->blob_top_->offset(n, c, h, w)];
          EXPECT_GE(data, c * 9 + 0.1 - 1e-4);
          EXPECT_LE(data, c * 9 + 0.1 + 1e-4);
        }
      }
    }
  }
}

TYPED_TEST(ConvolutionLayerTest, TestCPUGradient) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(2);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::CPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

TYPED_TEST(ConvolutionLayerTest, TestCPUGradientGroup) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::CPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

TYPED_TEST(ConvolutionLayerTest, TestGPUGradient) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(2);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::GPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

TYPED_TEST(ConvolutionLayerTest, TestGPUGradientGroup) {
  LayerParameter layer_param;
  layer_param.set_kernelsize(3);
  layer_param.set_stride(2);
  layer_param.set_num_output(3);
  layer_param.set_group(3);
  layer_param.mutable_weight_filler()->set_type("gaussian");
  layer_param.mutable_bias_filler()->set_type("gaussian");
  Caffe::set_mode(Caffe::GPU);
  ConvolutionLayer<TypeParam> layer(layer_param);
  GradientChecker<TypeParam> checker(1e-2, 1e-3);
  checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_);
}

}
<|endoftext|>


<commit_before>//=======================================================================
// Copyright Baptiste Wicht 2013-2016.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
//  http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================

#include <vector.hpp>
#include <string.hpp>

#include "arp_layer.hpp"
#include "arp_cache.hpp"
#include "ip_layer.hpp"
#include "logging.hpp"
#include "kernel_utils.hpp"

void network::arp::decode(network::interface_descriptor& interface, network::ethernet::packet& packet){
    header* arp_header = reinterpret_cast<header*>(packet.payload + packet.index);

    logging::logf(logging::log_level::TRACE, "arp: Start ARP packet handling\n");

    auto hw_type = switch_endian_16(arp_header->hw_type);
    auto protocol_type = switch_endian_16(arp_header->protocol_type);

    // We only support ethernet hardware frames
    if(hw_type != 0x1){
        logging::logf(logging::log_level::ERROR, "arp: Unhandled hardware type %h\n", size_t(hw_type));
        return;
    }

    // We only support IPV4 protocol
    if(protocol_type != 0x800){
        logging::logf(logging::log_level::ERROR, "arp: Unhandled protocol type %h\n", size_t(protocol_type));
        return;
    }

    auto hw_len = arp_header->hw_len;
    auto protocol_len = arp_header->protocol_len;
    auto operation = switch_endian_16(arp_header->operation);

    // This should never happen
    if(hw_len != 0x6 || protocol_len != 0x4){
        logging::logf(logging::log_level::ERROR, "arp: Unable to process the given length hw:%h prot:%h\n", size_t(hw_len), size_t(protocol_len));
        return;
    }

    if (operation != 0x1 && operation != 0x2) {
        logging::logf(logging::log_level::TRACE, "arp: Unhandled operation %h\n", size_t(operation));
        return;
    }

    size_t source_hw = 0;
    size_t target_hw = 0;

    for(size_t i = 0; i < 3; ++i){
        source_hw |= uint64_t(switch_endian_16(arp_header->source_hw_addr[i])) << ((2 - i) * 16);
        target_hw |= uint64_t(switch_endian_16(arp_header->target_hw_addr[i])) << ((2 - i) * 16);
    }

    logging::logf(logging::log_level::TRACE, "arp: Source HW Address %h \n", source_hw);
    logging::logf(logging::log_level::TRACE, "arp: Target HW Address %h \n", target_hw);

    network::ip::address source_prot;
    network::ip::address target_prot;

    for(size_t i = 0; i < 2; ++i){
        auto source = switch_endian_16(arp_header->source_protocol_addr[i]);
        auto target = switch_endian_16(arp_header->target_protocol_addr[i]);

        source_prot.set_sub(2*i, source >> 8);
        source_prot.set_sub(2*i+1, source);

        target_prot.set_sub(2*i, target >> 8);
        target_prot.set_sub(2*i+1, target);
    }

    logging::logf(logging::log_level::TRACE, "arp: Source Protocol Address %u.%u.%u.%u \n",
        uint64_t(source_prot(0)), uint64_t(source_prot(1)), uint64_t(source_prot(2)), uint64_t(source_prot(3)));
    logging::logf(logging::log_level::TRACE, "arp: Target Protocol Address %u.%u.%u.%u \n",
        uint64_t(target_prot(0)), uint64_t(target_prot(1)), uint64_t(target_prot(2)), uint64_t(target_prot(3)));

    // If not an ARP Probe, update the ARP cache
    if(source_prot.raw_address == 0x0){
        network::arp::update_cache(source_hw, source_prot);
    }

    if(operation == 0x1){
        if(target_prot == network::ip::make_address(64,65,66,67)){
            logging::logf(logging::log_level::TRACE, "arp: Reply to Request for own IP\n");

            // Ask the ethernet layer to craft a packet
            auto packet = network::ethernet::prepare_packet(interface, sizeof(header), source_hw, ethernet::ether_type::ARP);

            auto* arp_reply_header = reinterpret_cast<header*>(packet.payload + packet.index);

            arp_reply_header->hw_type = switch_endian_16(0x1); // ethernet
            arp_reply_header->protocol_type = switch_endian_16(0x800); // IPV4
            arp_reply_header->hw_len = 0x6; // MAC Address
            arp_reply_header->protocol_len = 0x4; // IP Address
            arp_reply_header->operation = switch_endian_16(0x2); //ARP Reply

            auto source_mac = interface.mac_address;

            for(size_t i = 0; i < 3; ++i){
                arp_reply_header->target_hw_addr[i] = arp_header->source_hw_addr[i];

                arp_reply_header->source_hw_addr[i] = switch_endian_16(uint16_t(source_mac >> ((2 - i) * 16)));
            }

            for(size_t i = 0; i < 2; ++i){
                arp_reply_header->source_protocol_addr[i] = arp_header->target_protocol_addr[i];
                arp_reply_header->target_protocol_addr[i] = arp_header->source_protocol_addr[i];
            }

            network::ethernet::finalize_packet(interface, packet);
        }
    } else if(operation == 0x2){
        logging::logf(logging::log_level::TRACE, "arp: Handle Reply\n");

        //TODO
    }
}
<commit_msg>Change the current IP<commit_after>//=======================================================================
// Copyright Baptiste Wicht 2013-2016.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
//  http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================

#include <vector.hpp>
#include <string.hpp>

#include "arp_layer.hpp"
#include "arp_cache.hpp"
#include "ip_layer.hpp"
#include "logging.hpp"
#include "kernel_utils.hpp"

void network::arp::decode(network::interface_descriptor& interface, network::ethernet::packet& packet){
    header* arp_header = reinterpret_cast<header*>(packet.payload + packet.index);

    logging::logf(logging::log_level::TRACE, "arp: Start ARP packet handling\n");

    auto hw_type = switch_endian_16(arp_header->hw_type);
    auto protocol_type = switch_endian_16(arp_header->protocol_type);

    // We only support ethernet hardware frames
    if(hw_type != 0x1){
        logging::logf(logging::log_level::ERROR, "arp: Unhandled hardware type %h\n", size_t(hw_type));
        return;
    }

    // We only support IPV4 protocol
    if(protocol_type != 0x800){
        logging::logf(logging::log_level::ERROR, "arp: Unhandled protocol type %h\n", size_t(protocol_type));
        return;
    }

    auto hw_len = arp_header->hw_len;
    auto protocol_len = arp_header->protocol_len;
    auto operation = switch_endian_16(arp_header->operation);

    // This should never happen
    if(hw_len != 0x6 || protocol_len != 0x4){
        logging::logf(logging::log_level::ERROR, "arp: Unable to process the given length hw:%h prot:%h\n", size_t(hw_len), size_t(protocol_len));
        return;
    }

    if (operation != 0x1 && operation != 0x2) {
        logging::logf(logging::log_level::TRACE, "arp: Unhandled operation %h\n", size_t(operation));
        return;
    }

    size_t source_hw = 0;
    size_t target_hw = 0;

    for(size_t i = 0; i < 3; ++i){
        source_hw |= uint64_t(switch_endian_16(arp_header->source_hw_addr[i])) << ((2 - i) * 16);
        target_hw |= uint64_t(switch_endian_16(arp_header->target_hw_addr[i])) << ((2 - i) * 16);
    }

    logging::logf(logging::log_level::TRACE, "arp: Source HW Address %h \n", source_hw);
    logging::logf(logging::log_level::TRACE, "arp: Target HW Address %h \n", target_hw);

    network::ip::address source_prot;
    network::ip::address target_prot;

    for(size_t i = 0; i < 2; ++i){
        auto source = switch_endian_16(arp_header->source_protocol_addr[i]);
        auto target = switch_endian_16(arp_header->target_protocol_addr[i]);

        source_prot.set_sub(2*i, source >> 8);
        source_prot.set_sub(2*i+1, source);

        target_prot.set_sub(2*i, target >> 8);
        target_prot.set_sub(2*i+1, target);
    }

    logging::logf(logging::log_level::TRACE, "arp: Source Protocol Address %u.%u.%u.%u \n",
        uint64_t(source_prot(0)), uint64_t(source_prot(1)), uint64_t(source_prot(2)), uint64_t(source_prot(3)));
    logging::logf(logging::log_level::TRACE, "arp: Target Protocol Address %u.%u.%u.%u \n",
        uint64_t(target_prot(0)), uint64_t(target_prot(1)), uint64_t(target_prot(2)), uint64_t(target_prot(3)));

    // If not an ARP Probe, update the ARP cache
    if(source_prot.raw_address == 0x0){
        network::arp::update_cache(source_hw, source_prot);
    }

    if(operation == 0x1){
        if(target_prot == network::ip::make_address(192,168,20,66)){
            logging::logf(logging::log_level::TRACE, "arp: Reply to Request for own IP\n");

            // Ask the ethernet layer to craft a packet
            auto packet = network::ethernet::prepare_packet(interface, sizeof(header), source_hw, ethernet::ether_type::ARP);

            auto* arp_reply_header = reinterpret_cast<header*>(packet.payload + packet.index);

            arp_reply_header->hw_type = switch_endian_16(0x1); // ethernet
            arp_reply_header->protocol_type = switch_endian_16(0x800); // IPV4
            arp_reply_header->hw_len = 0x6; // MAC Address
            arp_reply_header->protocol_len = 0x4; // IP Address
            arp_reply_header->operation = switch_endian_16(0x2); //ARP Reply

            auto source_mac = interface.mac_address;

            for(size_t i = 0; i < 3; ++i){
                arp_reply_header->target_hw_addr[i] = arp_header->source_hw_addr[i];

                arp_reply_header->source_hw_addr[i] = switch_endian_16(uint16_t(source_mac >> ((2 - i) * 16)));
            }

            for(size_t i = 0; i < 2; ++i){
                arp_reply_header->source_protocol_addr[i] = arp_header->target_protocol_addr[i];
                arp_reply_header->target_protocol_addr[i] = arp_header->source_protocol_addr[i];
            }

            network::ethernet::finalize_packet(interface, packet);
        }
    } else if(operation == 0x2){
        logging::logf(logging::log_level::TRACE, "arp: Handle Reply\n");

        //TODO
    }
}
<|endoftext|>


<commit_before>#include "regex_impl.hh"
#include "vector.hh"
#include "unit_tests.hh"
#include "string.hh"
#include "exception.hh"
#include "array_view.hh"

namespace Kakoune
{

namespace RegexProgram
{
enum Op : char
{
    Match,
    Literal,
    AnyChar,
    Jump,
    Split,
    LineStart,
    LineEnd,
};

using Offset = size_t;
}

namespace RegexCompiler
{
enum class Quantifier
{
    One,
    Optional,
    RepeatZeroOrMore,
    RepeatOneOrMore
};

enum class Op
{
    Literal,
    AnyChar,
    Sequence,
    Alternation,
    LineStart,
    LineEnd,
};

struct AstNode
{
    Op op;
    char value;
    Quantifier quantifier;
    Vector<std::unique_ptr<AstNode>> children;
};

using AstNodePtr = std::unique_ptr<AstNode>;

AstNodePtr make_ast_node(Op op, char value = 0,
                         Quantifier quantifier = Quantifier::One)
{
    return AstNodePtr{new AstNode{op, value, quantifier, {}}};
}

template<typename Iterator>
struct Parser
{
    static AstNodePtr parse(Iterator pos, Iterator end)
    {
        return disjunction(pos, end);
    }

private:
    static AstNodePtr disjunction(Iterator& pos, Iterator end)
    {
        AstNodePtr node = alternative(pos, end);
        if (pos == end or *pos != '|')
            return node;

        AstNodePtr res = make_ast_node(Op::Alternation);
        res->children.push_back(std::move(node));
        while (pos != end and *pos == '|')
            res->children.push_back(disjunction(++pos, end));
        return res;
    }

    static AstNodePtr alternative(Iterator& pos, Iterator end)
    {
        AstNodePtr res = make_ast_node(Op::Sequence);
        while (auto node = term(pos, end))
            res->children.push_back(std::move(node));
        return res;
    }

    static AstNodePtr term(Iterator& pos, Iterator end)
    {
        if (auto node = assertion(pos, end))
            return node;
        if (auto node = atom(pos, end))
        {
            node->quantifier = quantifier(pos, end);
            return node;
        }
        return nullptr;
    }

    static AstNodePtr assertion(Iterator& pos, Iterator end)
    {
        switch (*pos)
        {
            case '^': ++pos; return make_ast_node(Op::LineStart);
            case '$': ++pos; return make_ast_node(Op::LineEnd);
            /* TODO: \`, \', \b, \B, look ahead, look behind */
        }
        return nullptr;
    }

    static AstNodePtr atom(Iterator& pos, Iterator end)
    {
        const auto c = *pos;
        switch (c)
        {
            case '.': ++pos; return make_ast_node(Op::AnyChar);
            case '(':
            {
                ++pos;
                auto content = disjunction(pos, end);
                if (pos == end or *pos != ')')
                    throw runtime_error{"Unclosed parenthesis"};
                ++pos;
                return content;
            }
            default:
                if (contains("^$.*+?()[]{}|", c))
                    return nullptr;
                ++pos;
                return make_ast_node(Op::Literal, c);
        }
    }

    static Quantifier quantifier(Iterator& pos, Iterator end)
    {
        switch (*pos)
        {
            case '*': ++pos; return Quantifier::RepeatZeroOrMore;
            case '+': ++pos; return Quantifier::RepeatOneOrMore;
            case '?': ++pos; return Quantifier::Optional;
            default: return Quantifier::One;
        }
    }
};

RegexProgram::Offset compile_node(Vector<char>& program, const AstNodePtr& node)
{
    RegexProgram::Offset pos = program.size();

    auto allow_none = [](Quantifier quantifier) {
        return quantifier == Quantifier::Optional or
               quantifier == Quantifier::RepeatZeroOrMore;
    };

    auto is_repeat = [](Quantifier quantifier) {
        return quantifier == Quantifier::RepeatZeroOrMore or
               quantifier == Quantifier::RepeatOneOrMore;
    };

    auto alloc_offsets = [](Vector<char>& instructions, int count) {
        auto pos = instructions.size();
        instructions.resize(instructions.size() + count * sizeof(RegexProgram::Offset));
        return pos;
    };

    auto get_offset = [](Vector<char>& instructions, RegexProgram::Offset base, int index = 0) {
        return reinterpret_cast<RegexProgram::Offset*>(&instructions[base]) + index;
    };

    RegexProgram::Offset optional_offset = -1;
    if (allow_none(node->quantifier))
    {
        program.push_back(RegexProgram::Split);
        program.push_back(2);
        auto offsets = alloc_offsets(program, 2);
        *get_offset(program, offsets) = program.size();
        optional_offset = offsets;
    }

    Vector<RegexProgram::Offset> goto_end_offsets;
    auto content_pos = program.size();
    switch (node->op)
    {
        case Op::Literal:
            program.push_back(RegexProgram::Literal);
            program.push_back(node->value);
            break;
        case Op::AnyChar:
            program.push_back(RegexProgram::AnyChar);
            break;
        case Op::Sequence:
            for (auto& child : node->children)
                compile_node(program, child);
            break;
        case Op::Alternation:
        {
            const auto count = node->children.size();
            if (count > 255)
                throw runtime_error{"More than 255 elements in an alternation is not supported"};

            program.push_back(RegexProgram::Split);
            program.push_back(count);
            auto offsets = alloc_offsets(program, count);
            auto& children = node->children;
            for (int i = 0; i < children.size(); ++i)
            {
                auto child_pos = compile_node(program, children[i]);
                *get_offset(program, offsets, i) = child_pos;
                // Jump to end after executing that children
                program.push_back(RegexProgram::Jump);
                goto_end_offsets.push_back(alloc_offsets(program, 1));
            }
            break;
        }
        case Op::LineStart:
            program.push_back(RegexProgram::LineStart);
            break;
        case Op::LineEnd:
            program.push_back(RegexProgram::LineEnd);
            break;
    }

    for (auto& offset : goto_end_offsets)
        *get_offset(program, offset) =  program.size();

    if (is_repeat(node->quantifier))
    {
        program.push_back(RegexProgram::Split);
        program.push_back(2);
        auto offsets = alloc_offsets(program, 2);
        *get_offset(program, offsets, 0) = content_pos;
        *get_offset(program, offsets, 1) = program.size();
    }

    if (optional_offset != -1)
        *get_offset(program, optional_offset, 1) = program.size();

    return pos;
}

Vector<char> compile(const AstNodePtr& node)
{
    Vector<char> res;
    compile_node(res, node);
    res.push_back(RegexProgram::Match);
    return res;
}

template<typename Iterator>
Vector<char> compile(Iterator begin, Iterator end)
{
    return compile(Parser<Iterator>::parse(begin, end));
}
}

namespace RegexProgram
{
void dump(ConstArrayView<char> program)
{
    for (size_t pos = 0; pos < program.size(); )
    {
        printf("%4zd    ", pos);
        switch ((RegexProgram::Op)program[pos++])
        {
            case RegexProgram::Literal:
                printf("literal %c\n", program[pos++]);
                break;
            case RegexProgram::AnyChar:
                printf("any char\n");
                break;
            case RegexProgram::Jump:
                printf("jump %zd\n", *reinterpret_cast<const RegexProgram::Offset*>(&program[pos]));
                pos += sizeof(RegexProgram::Offset);
                break;
            case RegexProgram::Split:
            {
                int count = program[pos++];
                printf("split [");
                for (int i = 0; i < count; ++i)
                    printf("%zd%s", reinterpret_cast<const RegexProgram::Offset*>(&program[pos])[i],
                           (i == count - 1) ? "]\n" : ", ");
                pos += count * sizeof(RegexProgram::Offset);
                break;
            }
            case RegexProgram::LineStart:
                printf("line start\n");
                break;
            case RegexProgram::LineEnd:
                printf("line end\n");
                break;
            case RegexProgram::Match:
                printf("match\n");
        }
    }
}

struct StepResult
{
    enum Result { Consumed, Matched, Failed } result;
    const char* next = nullptr;
};

StepResult step_thread(const char* inst, char c, const char* start, Vector<const char*>& threads)
{
    while (true)
    {
        const RegexProgram::Op op = (RegexProgram::Op)*inst++;
        switch (op)
        {
            case RegexProgram::Literal:
                if (*inst++ == c)
                    return { StepResult::Consumed, inst };
                return { StepResult::Failed };
            case RegexProgram::AnyChar:
                return { StepResult::Consumed, inst };
            case RegexProgram::Jump:
                inst = start + *reinterpret_cast<const RegexProgram::Offset*>(inst);
                break;
            case RegexProgram::Split:
            {
                const int count = *inst++;
                auto* offsets = reinterpret_cast<const RegexProgram::Offset*>(inst);
                for (int o = 1; o < count; ++o)
                    threads.push_back(start + offsets[o]);
                inst = start + offsets[0];
                break;
            }
            case RegexProgram::LineStart:
                // TODO
                break;
            case RegexProgram::LineEnd:
                // TODO
                break;
            case RegexProgram::Match:
                return { StepResult::Matched };
        }
    }
    return { StepResult::Failed };
}

bool match(ConstArrayView<char> program, StringView data)
{
    const char* start = program.begin();
    Vector<const char*> threads = { start };

    for (auto c : data)
    {
        for (int i = 0; i < threads.size(); ++i)
        {
            auto res = step_thread(threads[i], c, start, threads);
            threads[i] = res.next;
            if (res.result == StepResult::Matched)
                return true;
        }
        threads.erase(std::remove(threads.begin(), threads.end(), nullptr), threads.end());
        if (threads.empty())
            break;
    }

    // Step remaining threads to see if they match without consuming anything else
    for (int i = 0; i < threads.size(); ++i)
    {
        if (step_thread(threads[i], 0, start, threads).result == StepResult::Matched)
            return true;
    }
    return false;
}
}

auto test_regex = UnitTest{[]{
    {
        StringView re = "a*b";
        auto program = RegexCompiler::compile(re.begin(), re.end());
        RegexProgram::dump(program);
        kak_assert(RegexProgram::match(program, "b"));
        kak_assert(RegexProgram::match(program, "ab"));
        kak_assert(RegexProgram::match(program, "aaab"));
        kak_assert(not RegexProgram::match(program, "acb"));
        kak_assert(not RegexProgram::match(program, ""));
    }
    {
        StringView re = "^(foo|qux)+(bar)?baz$";
        auto program = RegexCompiler::compile(re.begin(), re.end());
        RegexProgram::dump(program);
        kak_assert(RegexProgram::match(program, "fooquxbarbaz"));
        kak_assert(not RegexProgram::match(program, "quxbar"));
        kak_assert(not RegexProgram::match(program, "blahblah"));
        kak_assert(RegexProgram::match(program, "foobaz"));
        kak_assert(RegexProgram::match(program, "quxbaz"));
    }
}};

}
<commit_msg>Regex: Introduce RegexProgram::ThreadedExecutor and add line end/begin impl<commit_after>#include "regex_impl.hh"
#include "vector.hh"
#include "unit_tests.hh"
#include "string.hh"
#include "exception.hh"
#include "array_view.hh"

namespace Kakoune
{

namespace RegexProgram
{
enum Op : char
{
    Match,
    Literal,
    AnyChar,
    Jump,
    Split,
    LineStart,
    LineEnd,
};

using Offset = size_t;
}

namespace RegexCompiler
{
enum class Quantifier
{
    One,
    Optional,
    RepeatZeroOrMore,
    RepeatOneOrMore
};

enum class Op
{
    Literal,
    AnyChar,
    Sequence,
    Alternation,
    LineStart,
    LineEnd,
};

struct AstNode
{
    Op op;
    char value;
    Quantifier quantifier;
    Vector<std::unique_ptr<AstNode>> children;
};

using AstNodePtr = std::unique_ptr<AstNode>;

AstNodePtr make_ast_node(Op op, char value = 0,
                         Quantifier quantifier = Quantifier::One)
{
    return AstNodePtr{new AstNode{op, value, quantifier, {}}};
}

template<typename Iterator>
struct Parser
{
    static AstNodePtr parse(Iterator pos, Iterator end)
    {
        return disjunction(pos, end);
    }

private:
    static AstNodePtr disjunction(Iterator& pos, Iterator end)
    {
        AstNodePtr node = alternative(pos, end);
        if (pos == end or *pos != '|')
            return node;

        AstNodePtr res = make_ast_node(Op::Alternation);
        res->children.push_back(std::move(node));
        while (pos != end and *pos == '|')
            res->children.push_back(disjunction(++pos, end));
        return res;
    }

    static AstNodePtr alternative(Iterator& pos, Iterator end)
    {
        AstNodePtr res = make_ast_node(Op::Sequence);
        while (auto node = term(pos, end))
            res->children.push_back(std::move(node));
        return res;
    }

    static AstNodePtr term(Iterator& pos, Iterator end)
    {
        if (auto node = assertion(pos, end))
            return node;
        if (auto node = atom(pos, end))
        {
            node->quantifier = quantifier(pos, end);
            return node;
        }
        return nullptr;
    }

    static AstNodePtr assertion(Iterator& pos, Iterator end)
    {
        switch (*pos)
        {
            case '^': ++pos; return make_ast_node(Op::LineStart);
            case '$': ++pos; return make_ast_node(Op::LineEnd);
            /* TODO: \`, \', \b, \B, look ahead, look behind */
        }
        return nullptr;
    }

    static AstNodePtr atom(Iterator& pos, Iterator end)
    {
        const auto c = *pos;
        switch (c)
        {
            case '.': ++pos; return make_ast_node(Op::AnyChar);
            case '(':
            {
                ++pos;
                auto content = disjunction(pos, end);
                if (pos == end or *pos != ')')
                    throw runtime_error{"Unclosed parenthesis"};
                ++pos;
                return content;
            }
            default:
                if (contains("^$.*+?()[]{}|", c))
                    return nullptr;
                ++pos;
                return make_ast_node(Op::Literal, c);
        }
    }

    static Quantifier quantifier(Iterator& pos, Iterator end)
    {
        switch (*pos)
        {
            case '*': ++pos; return Quantifier::RepeatZeroOrMore;
            case '+': ++pos; return Quantifier::RepeatOneOrMore;
            case '?': ++pos; return Quantifier::Optional;
            default: return Quantifier::One;
        }
    }
};

RegexProgram::Offset compile_node(Vector<char>& program, const AstNodePtr& node)
{
    RegexProgram::Offset pos = program.size();

    auto allow_none = [](Quantifier quantifier) {
        return quantifier == Quantifier::Optional or
               quantifier == Quantifier::RepeatZeroOrMore;
    };

    auto is_repeat = [](Quantifier quantifier) {
        return quantifier == Quantifier::RepeatZeroOrMore or
               quantifier == Quantifier::RepeatOneOrMore;
    };

    auto alloc_offsets = [](Vector<char>& instructions, int count) {
        auto pos = instructions.size();
        instructions.resize(instructions.size() + count * sizeof(RegexProgram::Offset));
        return pos;
    };

    auto get_offset = [](Vector<char>& instructions, RegexProgram::Offset base, int index = 0) {
        return reinterpret_cast<RegexProgram::Offset*>(&instructions[base]) + index;
    };

    RegexProgram::Offset optional_offset = -1;
    if (allow_none(node->quantifier))
    {
        program.push_back(RegexProgram::Split);
        program.push_back(2);
        auto offsets = alloc_offsets(program, 2);
        *get_offset(program, offsets) = program.size();
        optional_offset = offsets;
    }

    Vector<RegexProgram::Offset> goto_end_offsets;
    auto content_pos = program.size();
    switch (node->op)
    {
        case Op::Literal:
            program.push_back(RegexProgram::Literal);
            program.push_back(node->value);
            break;
        case Op::AnyChar:
            program.push_back(RegexProgram::AnyChar);
            break;
        case Op::Sequence:
            for (auto& child : node->children)
                compile_node(program, child);
            break;
        case Op::Alternation:
        {
            const auto count = node->children.size();
            if (count > 255)
                throw runtime_error{"More than 255 elements in an alternation is not supported"};

            program.push_back(RegexProgram::Split);
            program.push_back(count);
            auto offsets = alloc_offsets(program, count);
            auto& children = node->children;
            for (int i = 0; i < children.size(); ++i)
            {
                auto child_pos = compile_node(program, children[i]);
                *get_offset(program, offsets, i) = child_pos;
                // Jump to end after executing that children
                program.push_back(RegexProgram::Jump);
                goto_end_offsets.push_back(alloc_offsets(program, 1));
            }
            break;
        }
        case Op::LineStart:
            program.push_back(RegexProgram::LineStart);
            break;
        case Op::LineEnd:
            program.push_back(RegexProgram::LineEnd);
            break;
    }

    for (auto& offset : goto_end_offsets)
        *get_offset(program, offset) =  program.size();

    if (is_repeat(node->quantifier))
    {
        program.push_back(RegexProgram::Split);
        program.push_back(2);
        auto offsets = alloc_offsets(program, 2);
        *get_offset(program, offsets, 0) = content_pos;
        *get_offset(program, offsets, 1) = program.size();
    }

    if (optional_offset != -1)
        *get_offset(program, optional_offset, 1) = program.size();

    return pos;
}

Vector<char> compile(const AstNodePtr& node)
{
    Vector<char> res;
    compile_node(res, node);
    res.push_back(RegexProgram::Match);
    return res;
}

template<typename Iterator>
Vector<char> compile(Iterator begin, Iterator end)
{
    return compile(Parser<Iterator>::parse(begin, end));
}
}

namespace RegexProgram
{
void dump(ConstArrayView<char> program)
{
    for (size_t pos = 0; pos < program.size(); )
    {
        printf("%4zd    ", pos);
        switch ((RegexProgram::Op)program[pos++])
        {
            case RegexProgram::Literal:
                printf("literal %c\n", program[pos++]);
                break;
            case RegexProgram::AnyChar:
                printf("any char\n");
                break;
            case RegexProgram::Jump:
                printf("jump %zd\n", *reinterpret_cast<const RegexProgram::Offset*>(&program[pos]));
                pos += sizeof(RegexProgram::Offset);
                break;
            case RegexProgram::Split:
            {
                int count = program[pos++];
                printf("split [");
                for (int i = 0; i < count; ++i)
                    printf("%zd%s", reinterpret_cast<const RegexProgram::Offset*>(&program[pos])[i],
                           (i == count - 1) ? "]\n" : ", ");
                pos += count * sizeof(RegexProgram::Offset);
                break;
            }
            case RegexProgram::LineStart:
                printf("line start\n");
                break;
            case RegexProgram::LineEnd:
                printf("line end\n");
                break;
            case RegexProgram::Match:
                printf("match\n");
        }
    }
}

struct ThreadedExecutor
{
    ThreadedExecutor(ConstArrayView<char> program) : m_program{program} {}

    struct StepResult
    {
        enum Result { Consumed, Matched, Failed } result;
        const char* next = nullptr;
    };

    StepResult step(const char* inst)
    {
        while (true)
        {
            auto c = m_pos == m_subject.end() ? 0 : *m_pos;
            const RegexProgram::Op op = (RegexProgram::Op)*inst++;
            switch (op)
            {
                case RegexProgram::Literal:
                    if (*inst++ == c)
                        return { StepResult::Consumed, inst };
                    return { StepResult::Failed };
                case RegexProgram::AnyChar:
                    return { StepResult::Consumed, inst };
                case RegexProgram::Jump:
                    inst = m_program.begin() + *reinterpret_cast<const RegexProgram::Offset*>(inst);
                    break;
                case RegexProgram::Split:
                {
                    const int count = *inst++;
                    auto* offsets = reinterpret_cast<const RegexProgram::Offset*>(inst);
                    for (int o = 1; o < count; ++o)
                        m_threads.push_back(m_program.begin() + offsets[o]);
                    inst = m_program.begin() + offsets[0];
                    break;
                }
                case RegexProgram::LineStart:
                    if (not is_line_start())
                        return { StepResult::Failed };
                    break;
                case RegexProgram::LineEnd:
                    if (not is_line_end())
                        return { StepResult::Failed };
                    break;
                case RegexProgram::Match:
                    return { StepResult::Matched };
            }
        }
        return { StepResult::Failed };
    }

    bool match(ConstArrayView<char> program, StringView data)
    {
        m_threads = Vector<const char*>{program.begin()};
        m_subject = data;
        m_pos = data.begin();

        for (m_pos = m_subject.begin(); m_pos != m_subject.end(); ++m_pos)
        {
            for (int i = 0; i < m_threads.size(); ++i)
            {
                auto res = step(m_threads[i]);
                m_threads[i] = res.next;
                if (res.result == StepResult::Matched)
                    return true;
            }
            m_threads.erase(std::remove(m_threads.begin(), m_threads.end(), nullptr), m_threads.end());
            if (m_threads.empty())
                break;
        }

        // Step remaining threads to see if they match without consuming anything else
        for (int i = 0; i < m_threads.size(); ++i)
        {
            if (step(m_threads[i]).result == StepResult::Matched)
                return true;
        }
        return false;
    }

    bool is_line_start() const
    {
        return m_pos == m_subject.begin() or *(m_pos-1) == '\n';
    }

    bool is_line_end() const
    {
        return m_pos == m_subject.end() or *m_pos == '\n';
    }

    ConstArrayView<char> m_program;
    Vector<const char*> m_threads;
    StringView m_subject;
    const char* m_pos;
};
}

auto test_regex = UnitTest{[]{
    using Exec = RegexProgram::ThreadedExecutor;
    {
        StringView re = "a*b";
        auto program = RegexCompiler::compile(re.begin(), re.end());
        RegexProgram::dump(program);
        Exec exec{program};
        kak_assert(exec.match(program, "b"));
        kak_assert(exec.match(program, "ab"));
        kak_assert(exec.match(program, "aaab"));
        kak_assert(not exec.match(program, "acb"));
        kak_assert(not exec.match(program, ""));
    }
    {
        StringView re = "^(foo|qux)+(bar)?baz$";
        auto program = RegexCompiler::compile(re.begin(), re.end());
        RegexProgram::dump(program);
        Exec exec{program};
        kak_assert(exec.match(program, "fooquxbarbaz"));
        kak_assert(not exec.match(program, "quxbar"));
        kak_assert(not exec.match(program, "blahblah"));
        kak_assert(exec.match(program, "foobaz"));
        kak_assert(exec.match(program, "quxbaz"));
    }
}};

}
<|endoftext|>