<commit_before>/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
/*
 * This file is part of the LibreOffice project.
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * This file incorporates work covered by the following license notice:
 *
 *   Licensed to the Apache Software Foundation (ASF) under one or more
 *   contributor license agreements. See the NOTICE file distributed
 *   with this work for additional information regarding copyright
 *   ownership. The ASF licenses this file to you 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 .
 */

#include <vcl/print.hxx>
#include <svtools/prnsetup.hxx>

#include "swtypes.hxx"
#include "cmdid.h"
#include "envprt.hxx"
#include "envlop.hxx"
#include "uitool.hxx"

#include "envprt.hrc"

SwEnvPrtPage::SwEnvPrtPage(Window* pParent, const SfxItemSet& rSet)
    : SfxTabPage(pParent, "EnvPrinterPage",
        "modules/swriter/ui/envprinterpage.ui", rSet)
{
    get(m_pAlignBox, "alignbox");
    get(m_pTopButton, "top");
    get(m_pBottomButton, "bottom");
    get(m_pRightField, "right");
    get(m_pDownField, "down");
    get(m_pPrinterInfo, "printername");
    get(m_pPrtSetup, "setup");
    SetExchangeSupport();

    // Metrics
    FieldUnit eUnit = ::GetDfltMetric(sal_False);
    SetMetric(*m_pRightField, eUnit);
    SetMetric(*m_pDownField, eUnit);

    // Install handlers
    m_pTopButton->SetClickHdl(LINK(this, SwEnvPrtPage, ClickHdl));
    m_pBottomButton->SetClickHdl(LINK(this, SwEnvPrtPage, ClickHdl));

    m_pPrtSetup->SetClickHdl(LINK(this, SwEnvPrtPage, ButtonHdl));

    // Bitmaps
    m_pBottomButton->GetClickHdl().Call(m_pBottomButton);

    // ToolBox
    m_pAlignBox->SetClickHdl(LINK(this, SwEnvPrtPage, AlignHdl));
    m_aIds[ENV_HOR_LEFT] = m_pAlignBox->GetItemId("horileft");
    m_aIds[ENV_HOR_CNTR] = m_pAlignBox->GetItemId("horicenter");
    m_aIds[ENV_HOR_RGHT] = m_pAlignBox->GetItemId("horiright");
    m_aIds[ENV_VER_LEFT] = m_pAlignBox->GetItemId("vertleft");
    m_aIds[ENV_VER_CNTR] = m_pAlignBox->GetItemId("vertcenter");
    m_aIds[ENV_VER_RGHT] = m_pAlignBox->GetItemId("vertright");
}

IMPL_LINK_NOARG(SwEnvPrtPage, ClickHdl)
{
    if (m_pBottomButton->IsChecked())
    {
        // Envelope from botton
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_LEFT], Image(Bitmap(SW_RES(BMP_HOR_LEFT_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_CNTR], Image(Bitmap(SW_RES(BMP_HOR_CNTR_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_RGHT], Image(Bitmap(SW_RES(BMP_HOR_RGHT_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_LEFT], Image(Bitmap(SW_RES(BMP_VER_LEFT_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_CNTR], Image(Bitmap(SW_RES(BMP_VER_CNTR_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_RGHT], Image(Bitmap(SW_RES(BMP_VER_RGHT_LOWER))));
    }
    else
    {
        // Envelope from top
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_LEFT], Image(Bitmap(SW_RES(BMP_HOR_LEFT_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_CNTR], Image(Bitmap(SW_RES(BMP_HOR_CNTR_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_RGHT], Image(Bitmap(SW_RES(BMP_HOR_RGHT_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_LEFT], Image(Bitmap(SW_RES(BMP_VER_LEFT_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_CNTR], Image(Bitmap(SW_RES(BMP_VER_CNTR_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_RGHT], Image(Bitmap(SW_RES(BMP_VER_RGHT_UPPER))));
    }
    return 0;
}

IMPL_LINK( SwEnvPrtPage, ButtonHdl, Button *, pBtn )
{
    if (pBtn == m_pPrtSetup)
    {
        // Call printer setup
        if (pPrt)
        {
            PrinterSetupDialog* pDlg = new PrinterSetupDialog(this );
            pDlg->SetPrinter(pPrt);
            pDlg->Execute();
            delete pDlg;
            GrabFocus();
            m_pPrinterInfo->SetText(pPrt->GetName());
        }
    }
    return 0;
}

IMPL_LINK_NOARG(SwEnvPrtPage, AlignHdl)
{
    if (m_pAlignBox->GetCurItemId())
    {
        for (sal_uInt16 i = ENV_HOR_LEFT; i <= ENV_VER_RGHT; ++i)
            m_pAlignBox->CheckItem(m_aIds[i], false);
        m_pAlignBox->CheckItem(m_pAlignBox->GetCurItemId(), true);
    }
    else
    {
        // GetCurItemId() == 0 is possible!
        const SwEnvItem& rItem = (const SwEnvItem&) GetItemSet().Get(FN_ENVELOP);
        m_pAlignBox->CheckItem(m_aIds[rItem.eAlign], true);
    }
    return 0;
}

SfxTabPage* SwEnvPrtPage::Create(Window* pParent, const SfxItemSet& rSet)
{
    return new SwEnvPrtPage(pParent, rSet);
}

void SwEnvPrtPage::ActivatePage(const SfxItemSet&)
{
    if (pPrt)
        m_pPrinterInfo->SetText(pPrt->GetName());
}

int SwEnvPrtPage::DeactivatePage(SfxItemSet* _pSet)
{
    if( _pSet )
        FillItemSet(*_pSet);
    return SfxTabPage::LEAVE_PAGE;
}

void SwEnvPrtPage::FillItem(SwEnvItem& rItem)
{
    sal_uInt16 nOrient = 0;
    for (sal_uInt16 i = ENV_HOR_LEFT; i <= ENV_VER_RGHT; ++i)
    {
        if (m_pAlignBox->IsItemChecked(m_aIds[i]))
        {
            nOrient = i;
            break;
        }
    }

    rItem.eAlign          = (SwEnvAlign) (nOrient);
    rItem.bPrintFromAbove = m_pTopButton->IsChecked();
    rItem.lShiftRight     = static_cast< sal_Int32 >(GetFldVal(*m_pRightField));
    rItem.lShiftDown      = static_cast< sal_Int32 >(GetFldVal(*m_pDownField ));
}

bool SwEnvPrtPage::FillItemSet(SfxItemSet& rSet)
{
    FillItem(GetParentSwEnvDlg()->aEnvItem);
    rSet.Put(GetParentSwEnvDlg()->aEnvItem);
    return true;
}

void SwEnvPrtPage::Reset(const SfxItemSet& rSet)
{
    // Read item
    const SwEnvItem& rItem = (const SwEnvItem&) rSet.Get(FN_ENVELOP);
    m_pAlignBox->CheckItem(m_aIds[rItem.eAlign]);

    if (rItem.bPrintFromAbove)
        m_pTopButton->Check();
    else
        m_pBottomButton->Check();

    SetFldVal(*m_pRightField, rItem.lShiftRight);
    SetFldVal(*m_pDownField , rItem.lShiftDown );

    ActivatePage(rSet);
    ClickHdl(m_pTopButton);
}

/* vim:set shiftwidth=4 softtabstop=4 expandtab: */
<commit_msg>sal_uInt16 to int<commit_after>/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
/*
 * This file is part of the LibreOffice project.
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * This file incorporates work covered by the following license notice:
 *
 *   Licensed to the Apache Software Foundation (ASF) under one or more
 *   contributor license agreements. See the NOTICE file distributed
 *   with this work for additional information regarding copyright
 *   ownership. The ASF licenses this file to you 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 .
 */

#include <vcl/print.hxx>
#include <svtools/prnsetup.hxx>

#include "swtypes.hxx"
#include "cmdid.h"
#include "envprt.hxx"
#include "envlop.hxx"
#include "uitool.hxx"

#include "envprt.hrc"

SwEnvPrtPage::SwEnvPrtPage(Window* pParent, const SfxItemSet& rSet)
    : SfxTabPage(pParent, "EnvPrinterPage",
        "modules/swriter/ui/envprinterpage.ui", rSet)
{
    get(m_pAlignBox, "alignbox");
    get(m_pTopButton, "top");
    get(m_pBottomButton, "bottom");
    get(m_pRightField, "right");
    get(m_pDownField, "down");
    get(m_pPrinterInfo, "printername");
    get(m_pPrtSetup, "setup");
    SetExchangeSupport();

    // Metrics
    FieldUnit eUnit = ::GetDfltMetric(sal_False);
    SetMetric(*m_pRightField, eUnit);
    SetMetric(*m_pDownField, eUnit);

    // Install handlers
    m_pTopButton->SetClickHdl(LINK(this, SwEnvPrtPage, ClickHdl));
    m_pBottomButton->SetClickHdl(LINK(this, SwEnvPrtPage, ClickHdl));

    m_pPrtSetup->SetClickHdl(LINK(this, SwEnvPrtPage, ButtonHdl));

    // Bitmaps
    m_pBottomButton->GetClickHdl().Call(m_pBottomButton);

    // ToolBox
    m_pAlignBox->SetClickHdl(LINK(this, SwEnvPrtPage, AlignHdl));
    m_aIds[ENV_HOR_LEFT] = m_pAlignBox->GetItemId("horileft");
    m_aIds[ENV_HOR_CNTR] = m_pAlignBox->GetItemId("horicenter");
    m_aIds[ENV_HOR_RGHT] = m_pAlignBox->GetItemId("horiright");
    m_aIds[ENV_VER_LEFT] = m_pAlignBox->GetItemId("vertleft");
    m_aIds[ENV_VER_CNTR] = m_pAlignBox->GetItemId("vertcenter");
    m_aIds[ENV_VER_RGHT] = m_pAlignBox->GetItemId("vertright");
}

IMPL_LINK_NOARG(SwEnvPrtPage, ClickHdl)
{
    if (m_pBottomButton->IsChecked())
    {
        // Envelope from botton
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_LEFT], Image(Bitmap(SW_RES(BMP_HOR_LEFT_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_CNTR], Image(Bitmap(SW_RES(BMP_HOR_CNTR_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_RGHT], Image(Bitmap(SW_RES(BMP_HOR_RGHT_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_LEFT], Image(Bitmap(SW_RES(BMP_VER_LEFT_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_CNTR], Image(Bitmap(SW_RES(BMP_VER_CNTR_LOWER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_RGHT], Image(Bitmap(SW_RES(BMP_VER_RGHT_LOWER))));
    }
    else
    {
        // Envelope from top
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_LEFT], Image(Bitmap(SW_RES(BMP_HOR_LEFT_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_CNTR], Image(Bitmap(SW_RES(BMP_HOR_CNTR_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_HOR_RGHT], Image(Bitmap(SW_RES(BMP_HOR_RGHT_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_LEFT], Image(Bitmap(SW_RES(BMP_VER_LEFT_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_CNTR], Image(Bitmap(SW_RES(BMP_VER_CNTR_UPPER))));
        m_pAlignBox->SetItemImage(m_aIds[ENV_VER_RGHT], Image(Bitmap(SW_RES(BMP_VER_RGHT_UPPER))));
    }
    return 0;
}

IMPL_LINK( SwEnvPrtPage, ButtonHdl, Button *, pBtn )
{
    if (pBtn == m_pPrtSetup)
    {
        // Call printer setup
        if (pPrt)
        {
            PrinterSetupDialog* pDlg = new PrinterSetupDialog(this );
            pDlg->SetPrinter(pPrt);
            pDlg->Execute();
            delete pDlg;
            GrabFocus();
            m_pPrinterInfo->SetText(pPrt->GetName());
        }
    }
    return 0;
}

IMPL_LINK_NOARG(SwEnvPrtPage, AlignHdl)
{
    if (m_pAlignBox->GetCurItemId())
    {
        for (int i = ENV_HOR_LEFT; i <= ENV_VER_RGHT; ++i)
            m_pAlignBox->CheckItem(m_aIds[i], false);
        m_pAlignBox->CheckItem(m_pAlignBox->GetCurItemId(), true);
    }
    else
    {
        // GetCurItemId() == 0 is possible!
        const SwEnvItem& rItem = (const SwEnvItem&) GetItemSet().Get(FN_ENVELOP);
        m_pAlignBox->CheckItem(m_aIds[rItem.eAlign], true);
    }
    return 0;
}

SfxTabPage* SwEnvPrtPage::Create(Window* pParent, const SfxItemSet& rSet)
{
    return new SwEnvPrtPage(pParent, rSet);
}

void SwEnvPrtPage::ActivatePage(const SfxItemSet&)
{
    if (pPrt)
        m_pPrinterInfo->SetText(pPrt->GetName());
}

int SwEnvPrtPage::DeactivatePage(SfxItemSet* _pSet)
{
    if( _pSet )
        FillItemSet(*_pSet);
    return SfxTabPage::LEAVE_PAGE;
}

void SwEnvPrtPage::FillItem(SwEnvItem& rItem)
{
    int nOrient = 0;
    for (int i = ENV_HOR_LEFT; i <= ENV_VER_RGHT; ++i)
    {
        if (m_pAlignBox->IsItemChecked(m_aIds[i]))
        {
            nOrient = i;
            break;
        }
    }

    rItem.eAlign          = (SwEnvAlign) (nOrient);
    rItem.bPrintFromAbove = m_pTopButton->IsChecked();
    rItem.lShiftRight     = static_cast< sal_Int32 >(GetFldVal(*m_pRightField));
    rItem.lShiftDown      = static_cast< sal_Int32 >(GetFldVal(*m_pDownField ));
}

bool SwEnvPrtPage::FillItemSet(SfxItemSet& rSet)
{
    FillItem(GetParentSwEnvDlg()->aEnvItem);
    rSet.Put(GetParentSwEnvDlg()->aEnvItem);
    return true;
}

void SwEnvPrtPage::Reset(const SfxItemSet& rSet)
{
    // Read item
    const SwEnvItem& rItem = (const SwEnvItem&) rSet.Get(FN_ENVELOP);
    m_pAlignBox->CheckItem(m_aIds[rItem.eAlign]);

    if (rItem.bPrintFromAbove)
        m_pTopButton->Check();
    else
        m_pBottomButton->Check();

    SetFldVal(*m_pRightField, rItem.lShiftRight);
    SetFldVal(*m_pDownField , rItem.lShiftDown );

    ActivatePage(rSet);
    ClickHdl(m_pTopButton);
}

/* vim:set shiftwidth=4 softtabstop=4 expandtab: */
<|endoftext|>


<commit_before>// Copyright 2019 The TCMalloc Authors
//
// 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
//
//     https://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 <stddef.h>
#include <stdio.h>
#include <stdlib.h>

#include <limits>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "tcmalloc/common.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/parameter_accessors.h"
#include "tcmalloc/malloc_extension.h"
#include "tcmalloc/testing/testutil.h"

namespace tcmalloc {
namespace {

using ::testing::ContainsRegex;
using ::testing::HasSubstr;

void DumpHeapStats(absl::string_view label) {
  std::string buffer = MallocExtension::GetStats();
  absl::FPrintF(stderr, "%s\n%s\n", label, buffer);
}

// Fixture for friend access to MallocExtension.
class LimitTest : public ::testing::Test {
 protected:
  void SetLimit(size_t limit, bool is_hard) {
    MallocExtension::MemoryLimit v;
    v.limit = limit;
    v.hard = is_hard;
    MallocExtension::SetMemoryLimit(v);
  }

  size_t GetLimit(bool is_hard) {
    MallocExtension::MemoryLimit v = MallocExtension::GetMemoryLimit();
    if (v.hard == is_hard) {
      return v.limit;
    } else {
      // Return no limit, as we do not have a limit matching is_hard.
      return std::numeric_limits<size_t>::max();
    }
  }

  // avoid fragmentation in local caches
  void *malloc_pages(size_t bytes) {
    CHECK_CONDITION(bytes % kPageSize == 0);
    void *ptr;
    CHECK_CONDITION(posix_memalign(&ptr, kPageSize, bytes) == 0);
    return ptr;
  }

  size_t physical_memory_used() {
    std::map<std::string, MallocExtension::Property> m =
        MallocExtension::GetProperties();
    auto i = m.find("generic.physical_memory_used");
    CHECK_CONDITION(i != m.end());
    return i->second.value;
  }
};

TEST_F(LimitTest, LimitRespected) {
  static const size_t kLim = 4ul * 1024 * 1024 * 1024;
  SetLimit(kLim, false);

  std::string statsBuf = MallocExtension::GetStats();
  std::string statsPbtxt = GetStatsInPbTxt();
  EXPECT_THAT(statsBuf, HasSubstr(absl::StrFormat(
                            "PARAMETER desired_usage_limit_bytes %u", kLim)));
  EXPECT_THAT(statsBuf, HasSubstr("Number of times limit was hit: 0"));
  EXPECT_THAT(statsPbtxt, HasSubstr(absl::StrFormat(
                              "desired_usage_limit_bytes: %u", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr("hard_limit: false"));
  EXPECT_THAT(statsPbtxt, HasSubstr("limit_hits: 0"));

  // Avoid failing due to usage by test itself.
  static const size_t kLimForUse = kLim * 9 / 10;
  // First allocate many small objects...
  size_t used = 0;
  std::vector<void *> ptrs;
  while (used < kLimForUse) {
    ptrs.push_back(malloc_pages(kPageSize));
    used += kPageSize;
  }
  DumpHeapStats("after allocating small objects");
  // return much of the space, fragmented...
  bool ret = false;
  for (auto &p : ptrs) {
    if (ret) {
      free(p);
      p = nullptr;
      used -= kPageSize;
    }
    ret = !ret;
  }
  DumpHeapStats("after freeing many small objects");
  // Now ensure we can re use it for large allocations.

  while (used < kLimForUse) {
    const size_t large = kPageSize * 10;
    ptrs.push_back(malloc_pages(large));
    used += large;
  }
  DumpHeapStats("after allocating large objects");
  EXPECT_LE(physical_memory_used(), kLim);

  statsBuf = MallocExtension::GetStats();
  statsPbtxt = GetStatsInPbTxt();
  // The HugePageAwareAllocator hits the limit more than once.
  EXPECT_THAT(statsBuf,
              ContainsRegex(R"(Number of times limit was hit: [1-9]\d*)"));
  EXPECT_THAT(statsPbtxt, ContainsRegex(R"(limit_hits: [1-9]\d*)"));

  for (auto p : ptrs) {
    free(p);
  }
}

  TEST_F(LimitTest, DISABLED_HardLimitRespected) {  // b/143314400
  static const size_t kLim = 300 << 20;
  SetLimit(kLim, true);

  std::string statsBuf = MallocExtension::GetStats();
  std::string statsPbtxt = GetStatsInPbTxt();
  EXPECT_THAT(statsBuf,
              HasSubstr(absl::StrFormat(
                  "PARAMETER desired_usage_limit_bytes %u (hard)", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr(absl::StrFormat(
                              "desired_usage_limit_bytes: %u", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr("hard_limit: true"));

  void *ptr = malloc_pages(50 << 20);
  ASSERT_DEATH(malloc_pages(400 << 20), "limit");
  free(ptr);
  SetLimit(std::numeric_limits<size_t>::max(), false);
}

TEST_F(LimitTest, HardLimitRespectsNoSubrelease) {
  static const size_t kLim = 300 << 20;
  SetLimit(kLim, true);
  TCMalloc_Internal_SetHPAASubrelease(false);
  EXPECT_FALSE(TCMalloc_Internal_GetHPAASubrelease());

  std::string statsBuf = MallocExtension::GetStats();
  std::string statsPbtxt = GetStatsInPbTxt();
  EXPECT_THAT(statsBuf,
              HasSubstr(absl::StrFormat(
                  "PARAMETER desired_usage_limit_bytes %u (hard)", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr(absl::StrFormat(
                              "desired_usage_limit_bytes: %u", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr("hard_limit: true"));

  ASSERT_DEATH(
      []() {
        // Allocate a bunch of medium objects, free half of them to cause some
        // fragmentation, then allocate some large objects. If we subrelease we
        // could stay under our hard limit, but if we don't then we should go
        // over.
        std::vector<void *> ptrs;
        constexpr size_t kNumMediumObjs = 400;
        constexpr size_t kNumLargeObjs = 200;
        for (size_t i = 0; i < kNumMediumObjs; i++) {
          ptrs.push_back(::operator new(512 << 10));
        }
        DumpHeapStats("after allocating medium objects");
        for (size_t i = 0; i < ptrs.size(); i++) {
          if (i % 2) continue;
          ::operator delete(ptrs[i]);
          ptrs[i] = static_cast<void *>(0);
        }
        DumpHeapStats("after freeing half of medium objects");
        for (size_t i = 0; i < kNumLargeObjs; i++) {
          ptrs.push_back(::operator new(1 << 20));
        }
        DumpHeapStats("after allocating large objects");
        while (!ptrs.empty()) {
          ::operator delete(ptrs.back());
          ptrs.pop_back();
        }
        DumpHeapStats("after freeing all objects");
      }(),
      "limit");
  SetLimit(std::numeric_limits<size_t>::max(), false);
}

}  // namespace
}  // namespace tcmalloc
<commit_msg>Enable hard limit test.<commit_after>// Copyright 2019 The TCMalloc Authors
//
// 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
//
//     https://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 <stddef.h>
#include <stdio.h>
#include <stdlib.h>

#include <limits>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "tcmalloc/common.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/parameter_accessors.h"
#include "tcmalloc/malloc_extension.h"
#include "tcmalloc/testing/testutil.h"

namespace tcmalloc {
namespace {

using ::testing::ContainsRegex;
using ::testing::HasSubstr;

void DumpHeapStats(absl::string_view label) {
  std::string buffer = MallocExtension::GetStats();
  absl::FPrintF(stderr, "%s\n%s\n", label, buffer);
}

// Fixture for friend access to MallocExtension.
class LimitTest : public ::testing::Test {
 protected:
  void SetLimit(size_t limit, bool is_hard) {
    MallocExtension::MemoryLimit v;
    v.limit = limit;
    v.hard = is_hard;
    MallocExtension::SetMemoryLimit(v);
  }

  size_t GetLimit(bool is_hard) {
    MallocExtension::MemoryLimit v = MallocExtension::GetMemoryLimit();
    if (v.hard == is_hard) {
      return v.limit;
    } else {
      // Return no limit, as we do not have a limit matching is_hard.
      return std::numeric_limits<size_t>::max();
    }
  }

  // avoid fragmentation in local caches
  void *malloc_pages(size_t bytes) {
    CHECK_CONDITION(bytes % kPageSize == 0);
    void *ptr;
    CHECK_CONDITION(posix_memalign(&ptr, kPageSize, bytes) == 0);
    return ptr;
  }

  size_t physical_memory_used() {
    std::map<std::string, MallocExtension::Property> m =
        MallocExtension::GetProperties();
    auto i = m.find("generic.physical_memory_used");
    CHECK_CONDITION(i != m.end());
    return i->second.value;
  }
};

TEST_F(LimitTest, LimitRespected) {
  static const size_t kLim = 4ul * 1024 * 1024 * 1024;
  SetLimit(kLim, false);

  std::string statsBuf = MallocExtension::GetStats();
  std::string statsPbtxt = GetStatsInPbTxt();
  EXPECT_THAT(statsBuf, HasSubstr(absl::StrFormat(
                            "PARAMETER desired_usage_limit_bytes %u", kLim)));
  EXPECT_THAT(statsBuf, HasSubstr("Number of times limit was hit: 0"));
  EXPECT_THAT(statsPbtxt, HasSubstr(absl::StrFormat(
                              "desired_usage_limit_bytes: %u", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr("hard_limit: false"));
  EXPECT_THAT(statsPbtxt, HasSubstr("limit_hits: 0"));

  // Avoid failing due to usage by test itself.
  static const size_t kLimForUse = kLim * 9 / 10;
  // First allocate many small objects...
  size_t used = 0;
  std::vector<void *> ptrs;
  while (used < kLimForUse) {
    ptrs.push_back(malloc_pages(kPageSize));
    used += kPageSize;
  }
  DumpHeapStats("after allocating small objects");
  // return much of the space, fragmented...
  bool ret = false;
  for (auto &p : ptrs) {
    if (ret) {
      free(p);
      p = nullptr;
      used -= kPageSize;
    }
    ret = !ret;
  }
  DumpHeapStats("after freeing many small objects");
  // Now ensure we can re use it for large allocations.

  while (used < kLimForUse) {
    const size_t large = kPageSize * 10;
    ptrs.push_back(malloc_pages(large));
    used += large;
  }
  DumpHeapStats("after allocating large objects");
  EXPECT_LE(physical_memory_used(), kLim);

  statsBuf = MallocExtension::GetStats();
  statsPbtxt = GetStatsInPbTxt();
  // The HugePageAwareAllocator hits the limit more than once.
  EXPECT_THAT(statsBuf,
              ContainsRegex(R"(Number of times limit was hit: [1-9]\d*)"));
  EXPECT_THAT(statsPbtxt, ContainsRegex(R"(limit_hits: [1-9]\d*)"));

  for (auto p : ptrs) {
    free(p);
  }
}

TEST_F(LimitTest, HardLimitRespected) {
  static const size_t kLim = 300 << 20;
  SetLimit(kLim, true);

  std::string statsBuf = MallocExtension::GetStats();
  std::string statsPbtxt = GetStatsInPbTxt();
  EXPECT_THAT(statsBuf,
              HasSubstr(absl::StrFormat(
                  "PARAMETER desired_usage_limit_bytes %u (hard)", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr(absl::StrFormat(
                              "desired_usage_limit_bytes: %u", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr("hard_limit: true"));

  void *ptr = malloc_pages(50 << 20);
  ASSERT_DEATH(malloc_pages(400 << 20), "limit");
  free(ptr);
  SetLimit(std::numeric_limits<size_t>::max(), false);
}

TEST_F(LimitTest, HardLimitRespectsNoSubrelease) {
  static const size_t kLim = 300 << 20;
  SetLimit(kLim, true);
  TCMalloc_Internal_SetHPAASubrelease(false);
  EXPECT_FALSE(TCMalloc_Internal_GetHPAASubrelease());

  std::string statsBuf = MallocExtension::GetStats();
  std::string statsPbtxt = GetStatsInPbTxt();
  EXPECT_THAT(statsBuf,
              HasSubstr(absl::StrFormat(
                  "PARAMETER desired_usage_limit_bytes %u (hard)", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr(absl::StrFormat(
                              "desired_usage_limit_bytes: %u", kLim)));
  EXPECT_THAT(statsPbtxt, HasSubstr("hard_limit: true"));

  ASSERT_DEATH(
      []() {
        // Allocate a bunch of medium objects, free half of them to cause some
        // fragmentation, then allocate some large objects. If we subrelease we
        // could stay under our hard limit, but if we don't then we should go
        // over.
        std::vector<void *> ptrs;
        constexpr size_t kNumMediumObjs = 400;
        constexpr size_t kNumLargeObjs = 200;
        for (size_t i = 0; i < kNumMediumObjs; i++) {
          ptrs.push_back(::operator new(512 << 10));
        }
        DumpHeapStats("after allocating medium objects");
        for (size_t i = 0; i < ptrs.size(); i++) {
          if (i % 2) continue;
          ::operator delete(ptrs[i]);
          ptrs[i] = static_cast<void *>(0);
        }
        DumpHeapStats("after freeing half of medium objects");
        for (size_t i = 0; i < kNumLargeObjs; i++) {
          ptrs.push_back(::operator new(1 << 20));
        }
        DumpHeapStats("after allocating large objects");
        while (!ptrs.empty()) {
          ::operator delete(ptrs.back());
          ptrs.pop_back();
        }
        DumpHeapStats("after freeing all objects");
      }(),
      "limit");
  SetLimit(std::numeric_limits<size_t>::max(), false);
}

}  // namespace
}  // namespace tcmalloc
<|endoftext|>


<commit_before>/**
 * @file lcc_main.cpp
 * @author Nishant Mehta
 *
 * Executable for Local Coordinate Coding.
 */
#include <mlpack/core.hpp>
#include "lcc.hpp"

PROGRAM_INFO("Local Coordinate Coding",
    "An implementation of Local Coordinate Coding (LCC), which "
    "codes data that approximately lives on a manifold using a variation of l1-"
    "norm regularized sparse coding.  Given a dense data matrix X with n points"
    " and d dimensions, LCC seeks to find a dense dictionary matrix D with k "
    "atoms in d dimensions, and a coding matrix Z with n points in k "
    "dimensions.  Because of the regularization method used, the atoms in D "
    "should lie close to the manifold on which the data points lie."
    "\n\n"
    "The original data matrix X can then be reconstructed as D * Z.  Therefore,"
    " this program finds a representation of each point in X as a sparse linear"
    " combination of atoms in the dictionary D."
    "\n\n"
    "The coding is found with an algorithm which alternates between a "
    "dictionary step, which updates the dictionary D, and a coding step, which "
    "updates the coding matrix Z."
    "\n\n"
    "To run this program, the input matrix X must be specified (with -i), along"
    " with the number of atoms in the dictionary (-k).  An initial dictionary "
    "may also be specified with the --initial_dictionary option.  The l1-norm "
    "regularization parameter is specified with -l.  For example, to run LCC on"
    " the dataset in data.csv using 200 atoms and an l1-regularization "
    "parameter of 0.1, saving the dictionary into dict.csv and the codes into "
    "codes.csv, use "
    "\n\n"
    "$ local_coordinate_coding -i data.csv -k 200 -l 0.1 -d dict.csv -c "
    "codes.csv"
    "\n\n"
    "The maximum number of iterations may be specified with the -n option. "
    "Optionally, the input data matrix X can be normalized before coding with "
    "the -N option.");

PARAM_STRING_REQ("input_file", "Filename of the input data.", "i");
PARAM_INT_REQ("atoms", "Number of atoms in the dictionary.", "k");

PARAM_DOUBLE("lambda", "Weighted l1-norm regularization parameter.", "l", 0.0);

PARAM_INT("max_iterations", "Maximum number of iterations for LCC (0 indicates "
    "no limit).", "n", 0);

PARAM_STRING("initial_dictionary", "Filename for optional initial dictionary.",
    "D", "");

PARAM_STRING("dictionary_file", "Filename to save the output dictionary to.",
    "d", "dictionary.csv");
PARAM_STRING("codes_file", "Filename to save the output codes to.", "c",
    "codes.csv");

PARAM_FLAG("normalize", "If set, the input data matrix will be normalized "
    "before coding.", "N");

PARAM_INT("seed", "Random seed.  If 0, 'std::time(NULL)' is used.", "s", 0);

PARAM_DOUBLE("objective_tolerance", "Tolerance for objective function.", "o",
    0.01);

using namespace arma;
using namespace std;
using namespace mlpack;
using namespace mlpack::math;
using namespace mlpack::lcc;
using namespace mlpack::sparse_coding; // For NothingInitializer.

int main(int argc, char* argv[])
{
  CLI::ParseCommandLine(argc, argv);

  if (CLI::GetParam<int>("seed") != 0)
    RandomSeed((size_t) CLI::GetParam<int>("seed"));
  else
    RandomSeed((size_t) std::time(NULL));

  const double lambda = CLI::GetParam<double>("lambda");

  const string inputFile = CLI::GetParam<string>("input_file");
  const string dictionaryFile = CLI::GetParam<string>("dictionary_file");
  const string codesFile = CLI::GetParam<string>("codes_file");
  const string initialDictionaryFile =
      CLI::GetParam<string>("initial_dictionary");

  const size_t maxIterations = CLI::GetParam<int>("max_iterations");
  const size_t atoms = CLI::GetParam<int>("atoms");

  const bool normalize = CLI::HasParam("normalize");

  const double objTolerance = CLI::GetParam<double>("objective_tolerance");

  mat input;
  data::Load(inputFile, input, true);

  Log::Info << "Loaded " << input.n_cols << " point in " << input.n_rows
      << " dimensions." << endl;

  // Normalize each point if the user asked for it.
  if (normalize)
  {
    Log::Info << "Normalizing data before coding..." << endl;
    for (size_t i = 0; i < input.n_cols; ++i)
      input.col(i) /= norm(input.col(i), 2);
  }

  // If there is an initial dictionary, be sure we do not initialize one.
  if (initialDictionaryFile != "")
  {
    LocalCoordinateCoding lcc(atoms, lambda, maxIterations, objTolerance);

    // Load initial dictionary directly into LCC object.
    data::Load(initialDictionaryFile, lcc.Dictionary(), true);

    // Validate size of initial dictionary.
    if (lcc.Dictionary().n_cols != atoms)
    {
      Log::Fatal << "The initial dictionary has " << lcc.Dictionary().n_cols
          << " atoms, but the number of atoms was specified to be " << atoms
          << "!" << endl;
    }

    if (lcc.Dictionary().n_rows != input.n_rows)
    {
      Log::Fatal << "The initial dictionary has " << lcc.Dictionary().n_rows
          << " dimensions, but the data has " << input.n_rows << " dimensions!"
          << endl;
    }

    // Run LCC.
    lcc.Train<NothingInitializer>(input);

    // Save the results.
    Log::Info << "Saving dictionary matrix to '" << dictionaryFile << "'.\n";
    data::Save(dictionaryFile, lcc.Dictionary());

    mat codes;
    lcc.Encode(input, codes);

    Log::Info << "Saving sparse codes to '" << codesFile << "'.\n";
    data::Save(codesFile, codes);
  }
  else
  {
    // No initial dictionary.
    LocalCoordinateCoding lcc(input, atoms, lambda, maxIterations,
        objTolerance);

    // Save the results.
    Log::Info << "Saving dictionary matrix to '" << dictionaryFile << "'.\n";
    data::Save(dictionaryFile, lcc.Dictionary());

    mat codes;
    lcc.Encode(input, codes);

    Log::Info << "Saving sparse codes to '" << codesFile << "'.\n";
    data::Save(codesFile, codes);
  }
}
<commit_msg>Refactor lcc program.<commit_after>/**
 * @file lcc_main.cpp
 * @author Nishant Mehta
 *
 * Executable for Local Coordinate Coding.
 */
#include <mlpack/core.hpp>
#include "lcc.hpp"

PROGRAM_INFO("Local Coordinate Coding",
    "An implementation of Local Coordinate Coding (LCC), which "
    "codes data that approximately lives on a manifold using a variation of l1-"
    "norm regularized sparse coding.  Given a dense data matrix X with n points"
    " and d dimensions, LCC seeks to find a dense dictionary matrix D with k "
    "atoms in d dimensions, and a coding matrix Z with n points in k "
    "dimensions.  Because of the regularization method used, the atoms in D "
    "should lie close to the manifold on which the data points lie."
    "\n\n"
    "The original data matrix X can then be reconstructed as D * Z.  Therefore,"
    " this program finds a representation of each point in X as a sparse linear"
    " combination of atoms in the dictionary D."
    "\n\n"
    "The coding is found with an algorithm which alternates between a "
    "dictionary step, which updates the dictionary D, and a coding step, which "
    "updates the coding matrix Z."
    "\n\n"
    "To run this program, the input matrix X must be specified (with -i), along"
    " with the number of atoms in the dictionary (-k).  An initial dictionary "
    "may also be specified with the --initial_dictionary option.  The l1-norm "
    "regularization parameter is specified with -l.  For example, to run LCC on"
    " the dataset in data.csv using 200 atoms and an l1-regularization "
    "parameter of 0.1, saving the dictionary into dict.csv and the codes into "
    "codes.csv, use "
    "\n\n"
    "$ local_coordinate_coding -i data.csv -k 200 -l 0.1 -d dict.csv -c "
    "codes.csv"
    "\n\n"
    "The maximum number of iterations may be specified with the -n option. "
    "Optionally, the input data matrix X can be normalized before coding with "
    "the -N option.");

// Training parameters.
PARAM_STRING("training_file", "Filename of the training data (X).", "t", "");
PARAM_INT("atoms", "Number of atoms in the dictionary.", "k", 0);
PARAM_DOUBLE("lambda", "Weighted l1-norm regularization parameter.", "l", 0.0);
PARAM_INT("max_iterations", "Maximum number of iterations for LCC (0 indicates "
    "no limit).", "n", 0);
PARAM_STRING("initial_dictionary", "Filename for optional initial dictionary.",
    "i", "");
PARAM_FLAG("normalize", "If set, the input data matrix will be normalized "
    "before coding.", "N");
PARAM_DOUBLE("tolerance", "Tolerance for objective function.", "o", 0.01);

// Load/save a model.
PARAM_STRING("input_model_file", "File containing input LCC model.", "m", "");
PARAM_STRING("output_model_file", "File to save trained LCC model to.", "M",
    "");

// Test on another dataset.
PARAM_STRING("test_file", "File of test points to encode.", "T", "");
PARAM_STRING("dictionary_file", "Filename to save the output dictionary to.",
    "d", "");
PARAM_STRING("codes_file", "Filename to save the output codes to.", "c", "");

PARAM_INT("seed", "Random seed.  If 0, 'std::time(NULL)' is used.", "s", 0);

using namespace arma;
using namespace std;
using namespace mlpack;
using namespace mlpack::math;
using namespace mlpack::lcc;
using namespace mlpack::sparse_coding; // For NothingInitializer.

int main(int argc, char* argv[])
{
  CLI::ParseCommandLine(argc, argv);

  if (CLI::GetParam<int>("seed") != 0)
    RandomSeed((size_t) CLI::GetParam<int>("seed"));
  else
    RandomSeed((size_t) std::time(NULL));

  // Check for parameter validity.
  if (CLI::HasParam("input_model_file") && CLI::HasParam("initial_dictionary"))
    Log::Fatal << "Cannot specify both --input_model_file (-m) and "
        << "--initial_dictionary (-i)!" << endl;

  if (CLI::HasParam("training_file") && !CLI::HasParam("atoms"))
    Log::Fatal << "If --training_file is specified, the number of atoms in the "
        << "dictionary must be specified with --atoms (-k)!" << endl;

  if (!CLI::HasParam("training_file") && !CLI::HasParam("input_model_file"))
    Log::Fatal << "One of --training_file (-t) or --input_model_file (-m) must "
        << "be specified!" << endl;

  if (!CLI::HasParam("codes_file") && !CLI::HasParam("dictionary_file") &&
      !CLI::HasParam("output_model_file"))
    Log::Warn << "Neither --codes_file (-c), --dictionary_file (-d), nor "
        << "--output_model_file (-M) are specified; no output will be saved."
        << endl;

  if (CLI::HasParam("codes_file") && !CLI::HasParam("test_file"))
    Log::Fatal << "--codes_file (-c) is specified, but no test matrix ("
        << "specified with --test_file or -T) is given to encode!" << endl;

  if (!CLI::HasParam("training_file"))
  {
    if (CLI::HasParam("atoms"))
      Log::Warn << "--atoms (-k) ignored because --training_file (-t) is not "
          << "specified." << endl;
    if (CLI::HasParam("lambda"))
      Log::Warn << "--lambda (-l) ignored because --training_file (-t) is not "
          << "specified." << endl;
    if (CLI::HasParam("initial_dictionary"))
      Log::Warn << "--initial_dictionary (-i) ignored because --training_file "
          << "(-t) is not specified." << endl;
    if (CLI::HasParam("max_iterations"))
      Log::Warn << "--max_iterations (-n) ignored because --training_file (-t) "
          << "is not specified." << endl;
    if (CLI::HasParam("normalize"))
      Log::Warn << "--normalize (-N) ignored because --training_file (-t) is "
          << "not specified." << endl;
    if (CLI::HasParam("tolerance"))
      Log::Warn << "--tolerance (-o) ignored because --training_file (-t) is "
          << "not specified." << endl;
  }

  // Do we have an existing model?
  LocalCoordinateCoding lcc(0, 0.0);
  if (CLI::HasParam("input_model_file"))
  {
    data::Load(CLI::GetParam<string>("input_model_file"), "lcc_model", lcc,
        true);
  }

  if (CLI::HasParam("training_file"))
  {
    mat matX;
    data::Load(CLI::GetParam<string>("training_file"), matX, true);

    // Normalize each point if the user asked for it.
    if (CLI::HasParam("normalize"))
    {
      Log::Info << "Normalizing data before coding..." << endl;
      for (size_t i = 0; i < matX.n_cols; ++i)
        matX.col(i) /= norm(matX.col(i), 2);
    }

    lcc.Lambda() = CLI::GetParam<double>("lambda");
    lcc.Atoms() = (size_t) CLI::GetParam<int>("atoms");
    lcc.MaxIterations() = (size_t) CLI::GetParam<int>("max_iterations");
    lcc.Tolerance() = CLI::GetParam<double>("tolerance");

    // Inform the user if we are overwriting their model.
    if (CLI::HasParam("input_model_file"))
    {
      Log::Info << "Using dictionary from existing model in '"
          << CLI::GetParam<string>("input_model_file") << "' as initial "
          << "dictionary for training." << endl;
      lcc.Train<NothingInitializer>(matX);
    }
    else if (CLI::HasParam("initial_dictionary"))
    {
      // Load initial dictionary directly into LCC object.
      data::Load(CLI::GetParam<string>("initial_dictionary"), lcc.Dictionary(),
          true);

      // Validate the size of the initial dictionary.
      if (lcc.Dictionary().n_cols != lcc.Atoms())
      {
        Log::Fatal << "The initial dictionary has " << lcc.Dictionary().n_cols
            << " atoms, but the number of atoms was specified to be "
            << lcc.Atoms() << "!" << endl;
      }

      if (lcc.Dictionary().n_rows != matX.n_rows)
      {
        Log::Fatal << "The initial dictionary has " << lcc.Dictionary().n_rows
            << " dimensions, but the data has " << matX.n_rows << " dimensions!"
            << endl;
      }

      // Train the model.
      lcc.Train<NothingInitializer>(matX);
    }
    else
    {
      // Run with the default initialization.
      lcc.Train(matX);
    }
  }

  // Now, do we have any matrix to encode?
  if (CLI::HasParam("test_file"))
  {
    mat matY;
    data::Load(CLI::GetParam<string>("test_file"), matY, true);

    if (matY.n_rows != lcc.Dictionary().n_rows)
      Log::Fatal << "Model was trained with a dimensionality of "
          << lcc.Dictionary().n_rows << ", but data in test file "
          << CLI::GetParam<string>("test_file") << " has a dimensionality of "
          << matY.n_rows << "!" << endl;

    // Normalize each point if the user asked for it.
    if (CLI::HasParam("normalize"))
    {
      Log::Info << "Normalizing test data before coding..." << endl;
      for (size_t i = 0; i < matY.n_cols; ++i)
        matY.col(i) /= norm(matY.col(i), 2);
    }

    mat codes;
    lcc.Encode(matY, codes);

    if (CLI::HasParam("codes_file"))
      data::Save(CLI::GetParam<string>("codes_file"), codes);
  }

  // Did the user want to save the dictionary?
  if (CLI::HasParam("dictionary_file"))
    data::Save(CLI::GetParam<string>("dictionary_file"), lcc.Dictionary());

  // Did the user want to save the model?
  if (CLI::HasParam("output_model_file"))
    data::Save(CLI::GetParam<string>("output_model_file"), "lcc_model", lcc,
        false); // Non-fatal on failure.
}
<|endoftext|>


<commit_before>/**
 * @file sparse_coding_impl.hpp
 * @author Nishant Mehta
 *
 * Implementation of Sparse Coding with Dictionary Learning using l1 (LASSO) or
 * l1+l2 (Elastic Net) regularization.
 */
#ifndef __MLPACK_METHODS_SPARSE_CODING_SPARSE_CODING_IMPL_HPP
#define __MLPACK_METHODS_SPARSE_CODING_SPARSE_CODING_IMPL_HPP

// In case it hasn't already been included.
#include "sparse_coding.hpp"

namespace mlpack {
namespace sparse_coding {

// TODO: parameterizable; options to methods?
#define OBJ_TOL 1e-2 // 1E-9
#define NEWTON_TOL 1e-6 // 1E-9

template<typename DictionaryInitializer>
SparseCoding<DictionaryInitializer>::SparseCoding(const arma::mat& data,
                                                  const size_t atoms,
                                                  const double lambda1,
                                                  const double lambda2) :
    atoms(atoms),
    data(data),
    codes(atoms, data.n_cols),
    lambda1(lambda1),
    lambda2(lambda2)
{
  // Initialize the dictionary.
  DictionaryInitializer::Initialize(data, atoms, dictionary);
}

template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::Encode(const size_t maxIterations)
{
  double lastObjVal = DBL_MAX;

  // Take the initial coding step, which has to happen before entering the main
  // optimization loop.
  Log::Info << "Initial Coding Step." << std::endl;

  OptimizeCode();
  arma::uvec adjacencies = find(codes);

  Log::Info << "  Sparsity level: " << 100.0 * ((double) (adjacencies.n_elem))
      / ((double) (atoms * data.n_cols)) << "%." << std::endl;
  Log::Info << "  Objective value: " << Objective() << "." << std::endl;

  for (size_t t = 1; t != maxIterations; ++t)
  {
    Log::Info << "Iteration " << t << " of " << maxIterations << "."
        << std::endl;

    // First step: optimize the dictionary.
    Log::Info << "Performing dictionary step... " << std::endl;
    OptimizeDictionary(adjacencies);
    Log::Info << "  Objective value: " << Objective() << "." << std::endl;

    // Second step: perform the coding.
    Log::Info << "Performing coding step..." << std::endl;
    OptimizeCode();
    // Get the indices of all the nonzero elements in the codes.
    adjacencies = find(codes);
    Log::Info << "  Sparsity level: " << 100.0 * ((double) (adjacencies.n_elem))
        / ((double) (atoms * data.n_cols)) << "%." << std::endl;

    // Find the new objective value and improvement so we can check for
    // convergence.
    double curObjVal = Objective();
    double improvement = lastObjVal - curObjVal;
    Log::Info << "  Objective value: " << curObjVal << " (improvement "
        << std::scientific << improvement << ")." << std::endl;

    // Have we converged?
    if (improvement < OBJ_TOL)
    {
      Log::Info << "Converged within tolerance " << OBJ_TOL << ".\n";
      break;
    }

    lastObjVal = curObjVal;
  }
}

template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::OptimizeCode()
{
  // When using the Cholesky version of LARS, this is correct even if
  // lambda2 > 0.
  arma::mat matGram = trans(dictionary) * dictionary;

  for (size_t i = 0; i < data.n_cols; ++i)
  {
    // Report progress.
    if ((i % 100) == 0)
      Log::Debug << "Optimization at point " << i << "." << std::endl;

    bool useCholesky = true;
    regression::LARS lars(useCholesky, matGram, lambda1, lambda2);

    // Create an alias of the code (using the same memory), and then LARS will
    // place the result directly into that; then we will not need to have an
    // extra copy.
    arma::vec code = codes.unsafe_col(i);
    lars.Regress(dictionary, data.unsafe_col(i), code, true);
  }
}

// Dictionary step for optimization.
template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::OptimizeDictionary(
    const arma::uvec& adjacencies)
{
  // Count the number of atomic neighbors for each point x^i.
  arma::uvec neighborCounts = arma::zeros<arma::uvec>(data.n_cols, 1);

  if (adjacencies.n_elem > 0)
  {
    // This gets the column index.  Intentional integer division.
    size_t curPointInd = (size_t) (adjacencies(0) / atoms);

    size_t nextColIndex = (curPointInd + 1) * atoms;
    for (size_t l = 1; l < adjacencies.n_elem; ++l)
    {
      // If l no longer refers to an element in this column, advance the column
      // number accordingly.
      if (adjacencies(l) >= nextColIndex)
      {
        curPointInd = (size_t) (adjacencies(l) / atoms);
        nextColIndex = (curPointInd + 1) * atoms;
      }

      ++neighborCounts(curPointInd);
    }
  }

  // Handle the case of inactive atoms (atoms not used in the given coding).
  std::vector<size_t> inactiveAtoms;
  std::vector<size_t> activeAtoms;
  activeAtoms.reserve(atoms);

  for (size_t j = 0; j < atoms; ++j)
  {
    if (accu(codes.row(j) != 0) == 0)
      inactiveAtoms.push_back(j);
  }

  const size_t nInactiveAtoms = inactiveAtoms.size();
  const size_t nActiveAtoms = atoms - nInactiveAtoms;

  // Efficient construction of Z restricted to active atoms.
  arma::mat matActiveZ;
  if (inactiveAtoms.empty())
  {
    matActiveZ = codes;
  }
  else
  {
    arma::uvec inactiveAtomsVec =
        arma::conv_to<arma::uvec>::from(inactiveAtoms);
    RemoveRows(codes, inactiveAtomsVec, matActiveZ);
  }

  if (nInactiveAtoms > 0)
  {
    Log::Warn << "There are " << nInactiveAtoms
        << " inactive atoms. They will be re-initialized randomly.\n";
  }

  Log::Debug << "Solving Dual via Newton's Method.\n";

  // Solve using Newton's method in the dual - note that the final dot
  // multiplication with inv(A) seems to be unavoidable. Although more
  // expensive, the code written this way (we use solve()) should be more
  // numerically stable than just using inv(A) for everything.
  arma::vec dualVars = arma::zeros<arma::vec>(nActiveAtoms);

  //vec dualVars = 1e-14 * ones<vec>(nActiveAtoms);

  // Method used by feature sign code - fails miserably here.  Perhaps the
  // MATLAB optimizer fmincon does something clever?
  //vec dualVars = 10.0 * randu(nActiveAtoms, 1);

  //vec dualVars = diagvec(solve(dictionary, data * trans(codes))
  //    - codes * trans(codes));
  //for (size_t i = 0; i < dualVars.n_elem; i++)
  //  if (dualVars(i) < 0)
  //    dualVars(i) = 0;

  bool converged = false;
  arma::mat codesXT = matActiveZ * trans(data);
  arma::mat codesZT = matActiveZ * trans(matActiveZ);

  double improvement;
  for (size_t t = 1; !converged; ++t)
  {
    arma::mat A = codesZT + diagmat(dualVars);

    arma::mat matAInvZXT = solve(A, codesXT);

    arma::vec gradient = -(arma::sum(arma::square(matAInvZXT), 1) -
        arma::ones<arma::vec>(nActiveAtoms));

    arma::mat hessian = -(-2 * (matAInvZXT * trans(matAInvZXT)) % inv(A));

    arma::vec searchDirection = -solve(hessian, gradient);

    // Armijo line search.
    const double c = 1e-4;
    double alpha = 1.0;
    const double rho = 0.9;
    double sufficientDecrease = c * dot(gradient, searchDirection);

    while (true)
    {
      // Calculate objective.
      double sumDualVars = sum(dualVars);
      double fOld = -(-trace(trans(codesXT) * matAInvZXT) - sumDualVars);
      double fNew = -(-trace(trans(codesXT) * solve(codesZT +
          diagmat(dualVars + alpha * searchDirection), codesXT)) -
          (sumDualVars + alpha * sum(searchDirection)));

      if (fNew <= fOld + alpha * sufficientDecrease)
      {
        searchDirection = alpha * searchDirection;
        improvement = fOld - fNew;
        break;
      }

      alpha *= rho;
    }

    // Take step and print useful information.
    dualVars += searchDirection;
    double normGradient = norm(gradient, 2);
    Log::Debug << "Newton Method iteration " << t << ":" << std::endl;
    Log::Debug << "  Gradient norm: " << std::scientific << normGradient
        << "." << std::endl;
    Log::Debug << "  Improvement: " << std::scientific << improvement << ".\n";

    if (improvement < NEWTON_TOL)
      converged = true;
  }

  if (inactiveAtoms.empty())
  {
    // Directly update dictionary.
    dictionary = trans(solve(codesZT + diagmat(dualVars), codesXT));
  }
  else
  {
    arma::mat activeDictionary = trans(solve(codesZT +
        diagmat(dualVars), codesXT));

    // Update all atoms.
    size_t currentInactiveIndex = 0;
    for (size_t i = 0; i < atoms; ++i)
    {
      if (inactiveAtoms[currentInactiveIndex] == i)
      {
        // This atom is inactive.  Reinitialize it randomly.
        dictionary.col(i) = (data.col(math::RandInt(data.n_cols)) +
                             data.col(math::RandInt(data.n_cols)) +
                             data.col(math::RandInt(data.n_cols)));

        dictionary.col(i) /= norm(dictionary.col(i), 2);

        // Increment inactive index counter.
        ++currentInactiveIndex;
      }
      else
      {
        // Update estimate.
        dictionary.col(i) = activeDictionary.col(i - currentInactiveIndex);
      }
    }
  }
}

// Project each atom of the dictionary back into the unit ball (if necessary).
template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::ProjectDictionary()
{
  for (size_t j = 0; j < atoms; j++)
  {
    double atomNorm = norm(dictionary.col(j), 2);
    if (atomNorm > 1)
    {
      Log::Info << "Norm of atom " << j << " exceeds 1 (" << std::scientific
          << atomNorm << ").  Shrinking...\n";
      dictionary.col(j) /= atomNorm;
    }
  }
}

// Compute the objective function.
template<typename DictionaryInitializer>
double SparseCoding<DictionaryInitializer>::Objective()
{
  double l11NormZ = sum(sum(abs(codes)));
  double froNormResidual = norm(data - (dictionary * codes), "fro");

  if (lambda2 > 0)
  {
    double froNormZ = norm(codes, "fro");
    return 0.5 * (std::pow(froNormResidual, 2.0) + (lambda2 *
        std::pow(froNormZ, 2.0))) + (lambda1 * l11NormZ);
  }
  else // It can be simpler.
  {
    return 0.5 * std::pow(froNormResidual, 2.0) + lambda1 * l11NormZ;
  }
}

void RemoveRows(const arma::mat& X,
                const arma::uvec& rowsToRemove,
                arma::mat& modX)
{
  const size_t cols = X.n_cols;
  const size_t rows = X.n_rows;
  const size_t nRemove = rowsToRemove.n_elem;
  const size_t nKeep = rows - nRemove;

  if (nRemove == 0)
  {
    modX = X;
  }
  else
  {
    modX.set_size(nKeep, cols);

    size_t curRow = 0;
    size_t removeInd = 0;
    // First, check 0 to first row to remove.
    if (rowsToRemove(0) > 0)
    {
      // Note that this implies that n_rows > 1.
      size_t height = rowsToRemove(0);
      modX(arma::span(curRow, curRow + height - 1), arma::span::all) =
          X(arma::span(0, rowsToRemove(0) - 1), arma::span::all);
      curRow += height;
    }
    // Now, check i'th row to remove to (i + 1)'th row to remove, until i is the
    // penultimate row.
    while (removeInd < nRemove - 1)
    {
      size_t height = rowsToRemove[removeInd + 1] -
          rowsToRemove[removeInd] - 1;

      if (height > 0)
      {
        modX(arma::span(curRow, curRow + height - 1), arma::span::all) =
            X(arma::span(rowsToRemove[removeInd] + 1,
            rowsToRemove[removeInd + 1] - 1), arma::span::all);
        curRow += height;
      }

      removeInd++;
    }

    // Now that i is the last row to remove, check last row to remove to last
    // row.
    if (rowsToRemove[removeInd] < rows - 1)
    {
      modX(arma::span(curRow, nKeep - 1), arma::span::all) =
          X(arma::span(rowsToRemove[removeInd] + 1, rows - 1),
          arma::span::all);
    }
  }
}

}; // namespace sparse_coding
}; // namespace mlpack

#endif
<commit_msg>Avoid a matrix copy for matActiveZ and potentially speed up gradient calculation by avoiding instantiation of arma::ones<arma::vec> (Armadillo may have already been smart enough to avoid that anyway).<commit_after>/**
 * @file sparse_coding_impl.hpp
 * @author Nishant Mehta
 *
 * Implementation of Sparse Coding with Dictionary Learning using l1 (LASSO) or
 * l1+l2 (Elastic Net) regularization.
 */
#ifndef __MLPACK_METHODS_SPARSE_CODING_SPARSE_CODING_IMPL_HPP
#define __MLPACK_METHODS_SPARSE_CODING_SPARSE_CODING_IMPL_HPP

// In case it hasn't already been included.
#include "sparse_coding.hpp"

namespace mlpack {
namespace sparse_coding {

// TODO: parameterizable; options to methods?
#define OBJ_TOL 1e-2 // 1E-9
#define NEWTON_TOL 1e-6 // 1E-9

template<typename DictionaryInitializer>
SparseCoding<DictionaryInitializer>::SparseCoding(const arma::mat& data,
                                                  const size_t atoms,
                                                  const double lambda1,
                                                  const double lambda2) :
    atoms(atoms),
    data(data),
    codes(atoms, data.n_cols),
    lambda1(lambda1),
    lambda2(lambda2)
{
  // Initialize the dictionary.
  DictionaryInitializer::Initialize(data, atoms, dictionary);
}

template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::Encode(const size_t maxIterations)
{
  double lastObjVal = DBL_MAX;

  // Take the initial coding step, which has to happen before entering the main
  // optimization loop.
  Log::Info << "Initial Coding Step." << std::endl;

  OptimizeCode();
  arma::uvec adjacencies = find(codes);

  Log::Info << "  Sparsity level: " << 100.0 * ((double) (adjacencies.n_elem))
      / ((double) (atoms * data.n_cols)) << "%." << std::endl;
  Log::Info << "  Objective value: " << Objective() << "." << std::endl;

  for (size_t t = 1; t != maxIterations; ++t)
  {
    Log::Info << "Iteration " << t << " of " << maxIterations << "."
        << std::endl;

    // First step: optimize the dictionary.
    Log::Info << "Performing dictionary step... " << std::endl;
    OptimizeDictionary(adjacencies);
    Log::Info << "  Objective value: " << Objective() << "." << std::endl;

    // Second step: perform the coding.
    Log::Info << "Performing coding step..." << std::endl;
    OptimizeCode();
    // Get the indices of all the nonzero elements in the codes.
    adjacencies = find(codes);
    Log::Info << "  Sparsity level: " << 100.0 * ((double) (adjacencies.n_elem))
        / ((double) (atoms * data.n_cols)) << "%." << std::endl;

    // Find the new objective value and improvement so we can check for
    // convergence.
    double curObjVal = Objective();
    double improvement = lastObjVal - curObjVal;
    Log::Info << "  Objective value: " << curObjVal << " (improvement "
        << std::scientific << improvement << ")." << std::endl;

    // Have we converged?
    if (improvement < OBJ_TOL)
    {
      Log::Info << "Converged within tolerance " << OBJ_TOL << ".\n";
      break;
    }

    lastObjVal = curObjVal;
  }
}

template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::OptimizeCode()
{
  // When using the Cholesky version of LARS, this is correct even if
  // lambda2 > 0.
  arma::mat matGram = trans(dictionary) * dictionary;

  for (size_t i = 0; i < data.n_cols; ++i)
  {
    // Report progress.
    if ((i % 100) == 0)
      Log::Debug << "Optimization at point " << i << "." << std::endl;

    bool useCholesky = true;
    regression::LARS lars(useCholesky, matGram, lambda1, lambda2);

    // Create an alias of the code (using the same memory), and then LARS will
    // place the result directly into that; then we will not need to have an
    // extra copy.
    arma::vec code = codes.unsafe_col(i);
    lars.Regress(dictionary, data.unsafe_col(i), code, true);
  }
}

// Dictionary step for optimization.
template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::OptimizeDictionary(
    const arma::uvec& adjacencies)
{
  // Count the number of atomic neighbors for each point x^i.
  arma::uvec neighborCounts = arma::zeros<arma::uvec>(data.n_cols, 1);

  if (adjacencies.n_elem > 0)
  {
    // This gets the column index.  Intentional integer division.
    size_t curPointInd = (size_t) (adjacencies(0) / atoms);

    size_t nextColIndex = (curPointInd + 1) * atoms;
    for (size_t l = 1; l < adjacencies.n_elem; ++l)
    {
      // If l no longer refers to an element in this column, advance the column
      // number accordingly.
      if (adjacencies(l) >= nextColIndex)
      {
        curPointInd = (size_t) (adjacencies(l) / atoms);
        nextColIndex = (curPointInd + 1) * atoms;
      }

      ++neighborCounts(curPointInd);
    }
  }

  // Handle the case of inactive atoms (atoms not used in the given coding).
  std::vector<size_t> inactiveAtoms;
  std::vector<size_t> activeAtoms;
  activeAtoms.reserve(atoms);

  for (size_t j = 0; j < atoms; ++j)
  {
    if (accu(codes.row(j) != 0) == 0)
      inactiveAtoms.push_back(j);
  }

  const size_t nInactiveAtoms = inactiveAtoms.size();
  const size_t nActiveAtoms = atoms - nInactiveAtoms;

  // Efficient construction of Z restricted to active atoms.
  arma::mat matActiveZ;
  if (!inactiveAtoms.empty())
  {
    arma::uvec inactiveAtomsVec =
        arma::conv_to<arma::uvec>::from(inactiveAtoms);
    RemoveRows(codes, inactiveAtomsVec, matActiveZ);
  }

  if (nInactiveAtoms > 0)
  {
    Log::Warn << "There are " << nInactiveAtoms
        << " inactive atoms. They will be re-initialized randomly.\n";
  }

  Log::Debug << "Solving Dual via Newton's Method.\n";

  // Solve using Newton's method in the dual - note that the final dot
  // multiplication with inv(A) seems to be unavoidable. Although more
  // expensive, the code written this way (we use solve()) should be more
  // numerically stable than just using inv(A) for everything.
  arma::vec dualVars = arma::zeros<arma::vec>(nActiveAtoms);

  //vec dualVars = 1e-14 * ones<vec>(nActiveAtoms);

  // Method used by feature sign code - fails miserably here.  Perhaps the
  // MATLAB optimizer fmincon does something clever?
  //vec dualVars = 10.0 * randu(nActiveAtoms, 1);

  //vec dualVars = diagvec(solve(dictionary, data * trans(codes))
  //    - codes * trans(codes));
  //for (size_t i = 0; i < dualVars.n_elem; i++)
  //  if (dualVars(i) < 0)
  //    dualVars(i) = 0;

  bool converged = false;

  // If we have any inactive atoms, we must construct these differently.
  arma::mat codesXT;
  arma::mat codesZT;

  if (inactiveAtoms.empty())
  {
    codesXT = codes * trans(data);
    codesZT = codes * trans(codes);
  }
  else
  {
    codesXT = matActiveZ * trans(data);
    codesZT = matActiveZ * trans(matActiveZ);
  }

  double improvement;
  for (size_t t = 1; !converged; ++t)
  {
    arma::mat A = codesZT + diagmat(dualVars);

    arma::mat matAInvZXT = solve(A, codesXT);

    arma::vec gradient = -arma::sum(arma::square(matAInvZXT), 1);
    gradient += 1;

    arma::mat hessian = -(-2 * (matAInvZXT * trans(matAInvZXT)) % inv(A));

    arma::vec searchDirection = -solve(hessian, gradient);

    // Armijo line search.
    const double c = 1e-4;
    double alpha = 1.0;
    const double rho = 0.9;
    double sufficientDecrease = c * dot(gradient, searchDirection);

    while (true)
    {
      // Calculate objective.
      double sumDualVars = sum(dualVars);
      double fOld = -(-trace(trans(codesXT) * matAInvZXT) - sumDualVars);
      double fNew = -(-trace(trans(codesXT) * solve(codesZT +
          diagmat(dualVars + alpha * searchDirection), codesXT)) -
          (sumDualVars + alpha * sum(searchDirection)));

      if (fNew <= fOld + alpha * sufficientDecrease)
      {
        searchDirection = alpha * searchDirection;
        improvement = fOld - fNew;
        break;
      }

      alpha *= rho;
    }

    // Take step and print useful information.
    dualVars += searchDirection;
    double normGradient = norm(gradient, 2);
    Log::Debug << "Newton Method iteration " << t << ":" << std::endl;
    Log::Debug << "  Gradient norm: " << std::scientific << normGradient
        << "." << std::endl;
    Log::Debug << "  Improvement: " << std::scientific << improvement << ".\n";

    if (improvement < NEWTON_TOL)
      converged = true;
  }

  if (inactiveAtoms.empty())
  {
    // Directly update dictionary.
    dictionary = trans(solve(codesZT + diagmat(dualVars), codesXT));
  }
  else
  {
    arma::mat activeDictionary = trans(solve(codesZT +
        diagmat(dualVars), codesXT));

    // Update all atoms.
    size_t currentInactiveIndex = 0;
    for (size_t i = 0; i < atoms; ++i)
    {
      if (inactiveAtoms[currentInactiveIndex] == i)
      {
        // This atom is inactive.  Reinitialize it randomly.
        dictionary.col(i) = (data.col(math::RandInt(data.n_cols)) +
                             data.col(math::RandInt(data.n_cols)) +
                             data.col(math::RandInt(data.n_cols)));

        dictionary.col(i) /= norm(dictionary.col(i), 2);

        // Increment inactive index counter.
        ++currentInactiveIndex;
      }
      else
      {
        // Update estimate.
        dictionary.col(i) = activeDictionary.col(i - currentInactiveIndex);
      }
    }
  }
}

// Project each atom of the dictionary back into the unit ball (if necessary).
template<typename DictionaryInitializer>
void SparseCoding<DictionaryInitializer>::ProjectDictionary()
{
  for (size_t j = 0; j < atoms; j++)
  {
    double atomNorm = norm(dictionary.col(j), 2);
    if (atomNorm > 1)
    {
      Log::Info << "Norm of atom " << j << " exceeds 1 (" << std::scientific
          << atomNorm << ").  Shrinking...\n";
      dictionary.col(j) /= atomNorm;
    }
  }
}

// Compute the objective function.
template<typename DictionaryInitializer>
double SparseCoding<DictionaryInitializer>::Objective()
{
  double l11NormZ = sum(sum(abs(codes)));
  double froNormResidual = norm(data - (dictionary * codes), "fro");

  if (lambda2 > 0)
  {
    double froNormZ = norm(codes, "fro");
    return 0.5 * (std::pow(froNormResidual, 2.0) + (lambda2 *
        std::pow(froNormZ, 2.0))) + (lambda1 * l11NormZ);
  }
  else // It can be simpler.
  {
    return 0.5 * std::pow(froNormResidual, 2.0) + lambda1 * l11NormZ;
  }
}

void RemoveRows(const arma::mat& X,
                const arma::uvec& rowsToRemove,
                arma::mat& modX)
{
  const size_t cols = X.n_cols;
  const size_t rows = X.n_rows;
  const size_t nRemove = rowsToRemove.n_elem;
  const size_t nKeep = rows - nRemove;

  if (nRemove == 0)
  {
    modX = X;
  }
  else
  {
    modX.set_size(nKeep, cols);

    size_t curRow = 0;
    size_t removeInd = 0;
    // First, check 0 to first row to remove.
    if (rowsToRemove(0) > 0)
    {
      // Note that this implies that n_rows > 1.
      size_t height = rowsToRemove(0);
      modX(arma::span(curRow, curRow + height - 1), arma::span::all) =
          X(arma::span(0, rowsToRemove(0) - 1), arma::span::all);
      curRow += height;
    }
    // Now, check i'th row to remove to (i + 1)'th row to remove, until i is the
    // penultimate row.
    while (removeInd < nRemove - 1)
    {
      size_t height = rowsToRemove[removeInd + 1] -
          rowsToRemove[removeInd] - 1;

      if (height > 0)
      {
        modX(arma::span(curRow, curRow + height - 1), arma::span::all) =
            X(arma::span(rowsToRemove[removeInd] + 1,
            rowsToRemove[removeInd + 1] - 1), arma::span::all);
        curRow += height;
      }

      removeInd++;
    }

    // Now that i is the last row to remove, check last row to remove to last
    // row.
    if (rowsToRemove[removeInd] < rows - 1)
    {
      modX(arma::span(curRow, nKeep - 1), arma::span::all) =
          X(arma::span(rowsToRemove[removeInd] + 1, rows - 1),
          arma::span::all);
    }
  }
}

}; // namespace sparse_coding
}; // namespace mlpack

#endif
<|endoftext|>


<commit_before>/**
 * @file
 * @brief Implementation of ROOT data file reader module
 * @copyright Copyright (c) 2017 CERN and the Allpix Squared authors.
 * This software is distributed under the terms of the MIT License, copied verbatim in the file "LICENSE.md".
 * In applying this license, CERN does not waive the privileges and immunities granted to it by virtue of its status as an
 * Intergovernmental Organization or submit itself to any jurisdiction.
 */

#include "ROOTObjectReaderModule.hpp"

#include <climits>
#include <string>
#include <utility>

#include <TBranch.h>
#include <TKey.h>
#include <TObjArray.h>
#include <TTree.h>

#include "core/messenger/Messenger.hpp"
#include "core/utils/log.h"
#include "core/utils/type.h"

#include "objects/Object.hpp"
#include "objects/objects.h"

#include "core/utils/type.h"

using namespace allpix;

ROOTObjectReaderModule::ROOTObjectReaderModule(Configuration& config, Messenger* messenger, GeometryManager* geo_mgr)
    : Module(config), messenger_(messenger), geo_mgr_(geo_mgr) {}

/**
 * @note Objects cannot be stored in smart pointers due to internal ROOT logic
 */
ROOTObjectReaderModule::~ROOTObjectReaderModule() {
    for(auto message_inf : message_info_array_) {
        delete message_inf.objects;
    }
}

/**
 * Adds lambda function map to convert a vector of generic objects to a templated message containing this particular type of
 * object from its typeid.
 */
template <typename T> static void add_creator(ROOTObjectReaderModule::MessageCreatorMap& map) {
    map[typeid(T)] = [&](std::vector<Object*> objects, std::shared_ptr<Detector> detector = nullptr) {
        std::vector<T> data;
        for(auto& object : objects) {
            data.emplace_back(std::move(*static_cast<T*>(object)));
        }

        if(detector == nullptr) {
            return std::make_shared<Message<T>>(data);
        }
        return std::make_shared<Message<T>>(data, detector);
    };
}

/**
 * Uses SFINAE trick to call the add_creator function for all template arguments of a container class. Used to add creators
 * for every object in a tuple of objects.
 */
template <template <typename...> class T, typename... Args>
static void gen_creator_map_from_tag(ROOTObjectReaderModule::MessageCreatorMap& map, type_tag<T<Args...>>) {
    std::initializer_list<int> value{(add_creator<Args>(map), 0)...};
    (void)value;
}

/**
 * Wrapper function to make the SFINAE trick in \ref gen_creator_map_from_tag work.
 */
template <typename T> static ROOTObjectReaderModule::MessageCreatorMap gen_creator_map() {
    ROOTObjectReaderModule::MessageCreatorMap ret_map;
    gen_creator_map_from_tag(ret_map, type_tag<T>());
    return ret_map;
}

void ROOTObjectReaderModule::init() {
    // Read include and exclude list
    if(config_.has("include") && config_.has("exclude")) {
        throw InvalidCombinationError(
            config_, {"exclude", "include"}, "include and exclude parameter are mutually exclusive");
    } else if(config_.has("include")) {
        auto inc_arr = config_.getArray<std::string>("include");
        include_.insert(inc_arr.begin(), inc_arr.end());
    } else if(config_.has("exclude")) {
        auto exc_arr = config_.getArray<std::string>("exclude");
        exclude_.insert(exc_arr.begin(), exc_arr.end());
    }

    // Initialize the call map from the tuple of available objects
    message_creator_map_ = gen_creator_map<allpix::OBJECTS>();

    // Open the file with the objects
    input_file_ = std::make_unique<TFile>(config_.getPath("file_name", true).c_str());

    // Read all the trees in the file
    TList* keys = input_file_->GetListOfKeys();
    std::set<std::string> tree_names;

    for(auto&& object : *keys) {
        auto& key = dynamic_cast<TKey&>(*object);
        if(std::string(key.GetClassName()) == "TTree") {
            auto tree = static_cast<TTree*>(key.ReadObjectAny(nullptr));

            // Check if a version of this tree has already been read
            if(tree_names.find(tree->GetName()) != tree_names.end()) {
                LOG(TRACE) << "Skipping copy of tree with name " << tree->GetName()
                           << " because one with identical name has already been processed";
                continue;
            }
            tree_names.insert(tree->GetName());

            // Check if this tree should be used
            if((!include_.empty() && include_.find(tree->GetName()) == include_.end()) ||
               (!exclude_.empty() && exclude_.find(tree->GetName()) != exclude_.end())) {
                LOG(TRACE) << "Ignoring tree with " << tree->GetName()
                           << " objects because it has been excluded or not explicitly included";
                continue;
            }

            trees_.push_back(tree);
        }
    }

    if(trees_.empty()) {
        LOG(ERROR) << "Provided ROOT file does not contain any trees, module will not read any data";
    }

    // Cross-check the core random seed stored in the file with the one configured:
    auto& global_config = getConfigManager()->getGlobalConfiguration();
    auto config_seed = global_config.get<uint64_t>("random_seed_core");

    std::string* str = nullptr;
    input_file_->GetObject("config/Allpix/random_seed_core", str);
    uint64_t file_seed;

    if(str == nullptr) {
        // check if missing random seed core in config file should be ignored
        if(config_.get<bool>("ignore_seed_mismatch", false)) {
            LOG(WARNING) << "No random seed for core set in the input data file, cross-check with configured value - "
                         << "this might lead to unexpected behavior. Random seed core from the input data is be used.";
            file_seed = config_seed;
        } else {
            throw InvalidValueError(global_config,
                                    "random_seed_core",
                                    "no random seed for core set in the input data file, cross-check with configured value "
                                    "impossible - this might lead to unexpected behavior.");
        }
    } else {
        file_seed = allpix::from_string<uint64_t>(*str);
    }

    if(config_seed != file_seed) {
        // check if mismatch between random seed core in config and in input file should be ignored
        if(config_.get<bool>("ignore_seed_mismatch", false)) {
            LOG(WARNING) << "Mismatch between core random seed in configuration file and input data"
                         << " - this might lead to unexpected behavior.";
        } else {
            throw InvalidValueError(global_config,
                                    "random_seed_core",
                                    "mismatch between core random seed in configuration file and input data - this "
                                    "might lead to unexpected behavior. Set to value configured in the input data file: " +
                                        (*str));
        }
    }

    // Cross-check version, print warning only in case of a mismatch:
    std::string* version_str = nullptr;
    input_file_->GetObject("config/Allpix/version", version_str);
    if(version_str != nullptr && (*version_str) != ALLPIX_PROJECT_VERSION) {
        LOG(WARNING) << "Reading data produced with different version " << (*version_str)
                     << " - this might lead to unexpected behavior.";
    }

    // Loop over all found trees
    for(auto& tree : trees_) {
        // Loop over the list of branches and create the set of receiver objects
        TObjArray* branches = tree->GetListOfBranches();
        for(int i = 0; i < branches->GetEntries(); i++) {
            auto* branch = static_cast<TBranch*>(branches->At(i));

            // Add a new vector of objects and bind it to the branch
            message_info message_inf;
            message_inf.objects = new std::vector<Object*>;
            message_info_array_.emplace_back(message_inf);
            branch->SetAddress(&(message_info_array_.back().objects));

            // Fill the rest of the message information
            // FIXME: we want to index this in a different way
            std::string branch_name = branch->GetName();
            auto split = allpix::split<std::string>(branch_name, "_");

            // Fetch information from the tree name
            size_t expected_size = 2;
            size_t det_idx = 0;
            size_t name_idx = 1;
            if(branch_name.front() == '_' || branch_name.empty()) {
                --expected_size;
                det_idx = INT_MAX;
                --name_idx;
            }
            if(branch_name.find('_') == std::string::npos) {
                --expected_size;
                name_idx = INT_MAX;
            }

            // Check tree structure and if object type matches name
            auto split_type = allpix::split<std::string>(branch->GetClassName(), "<>");
            if(expected_size != split.size() || split_type.size() != 2 || split_type[1].size() <= 2) {
                throw ModuleError("Tree is malformed and cannot be used for creating messages");
            }
            std::string class_name = split_type[1].substr(0, split_type[1].size() - 1);
            std::string apx_namespace = "allpix::";
            size_t ap_idx = class_name.find(apx_namespace);
            if(ap_idx != std::string::npos) {
                class_name.replace(ap_idx, apx_namespace.size(), "");
            }
            if(class_name != tree->GetName()) {
                throw ModuleError("Tree contains objects of the wrong type");
            }

            if(name_idx != INT_MAX) {
                message_info_array_.back().name = split[name_idx];
            }
            if(det_idx != INT_MAX) {
                if(split[det_idx] != "global") {
                    message_info_array_.back().detector = geo_mgr_->getDetector(split[det_idx]);
                }
            }
        }
    }
}

void ROOTObjectReaderModule::run(unsigned int event_num) {
    --event_num;
    for(auto& tree : trees_) {
        if(event_num >= tree->GetEntries()) {
            throw EndOfRunException("Requesting end of run because TTree only contains data for " +
                                    std::to_string(event_num) + " events");
        }
        tree->GetEntry(event_num);
    }
    LOG(TRACE) << "Building messages from stored objects";

    // Loop through all branches
    for(auto message_inf : message_info_array_) {
        auto objects = message_inf.objects;

        // Skip empty objects in current event
        if(objects->empty()) {
            continue;
        }

        // Check if a pointer to a dispatcher method exist
        auto first_object = (*objects)[0];
        auto iter = message_creator_map_.find(typeid(*first_object));
        if(iter == message_creator_map_.end()) {
            LOG(INFO) << "Cannot dispatch message with object " << allpix::demangle(typeid(*first_object).name())
                      << " because it not registered for messaging";
            continue;
        }

        // Update statistics
        read_cnt_ += objects->size();

        // Create a message
        std::shared_ptr<BaseMessage> message = iter->second(*objects, message_inf.detector);

        // Dispatch the message
        messenger_->dispatchMessage(this, message, message_inf.name);
    }
}

void ROOTObjectReaderModule::finalize() {
    int branch_count = 0;
    for(auto& tree : trees_) {
        branch_count += tree->GetListOfBranches()->GetEntries();
    }

    // Print statistics
    LOG(INFO) << "Read " << read_cnt_ << " objects from " << branch_count << " branches";

    // Close the file
    input_file_->Close();
}
<commit_msg>Instead of using two different if-statements to test if random seed core from config file is missing and if there is a mismatch, one if statement is used.<commit_after>/**
 * @file
 * @brief Implementation of ROOT data file reader module
 * @copyright Copyright (c) 2017 CERN and the Allpix Squared authors.
 * This software is distributed under the terms of the MIT License, copied verbatim in the file "LICENSE.md".
 * In applying this license, CERN does not waive the privileges and immunities granted to it by virtue of its status as an
 * Intergovernmental Organization or submit itself to any jurisdiction.
 */

#include "ROOTObjectReaderModule.hpp"

#include <climits>
#include <string>
#include <utility>

#include <TBranch.h>
#include <TKey.h>
#include <TObjArray.h>
#include <TTree.h>

#include "core/messenger/Messenger.hpp"
#include "core/utils/log.h"
#include "core/utils/type.h"

#include "objects/Object.hpp"
#include "objects/objects.h"

#include "core/utils/type.h"

using namespace allpix;

ROOTObjectReaderModule::ROOTObjectReaderModule(Configuration& config, Messenger* messenger, GeometryManager* geo_mgr)
    : Module(config), messenger_(messenger), geo_mgr_(geo_mgr) {}

/**
 * @note Objects cannot be stored in smart pointers due to internal ROOT logic
 */
ROOTObjectReaderModule::~ROOTObjectReaderModule() {
    for(auto message_inf : message_info_array_) {
        delete message_inf.objects;
    }
}

/**
 * Adds lambda function map to convert a vector of generic objects to a templated message containing this particular type of
 * object from its typeid.
 */
template <typename T> static void add_creator(ROOTObjectReaderModule::MessageCreatorMap& map) {
    map[typeid(T)] = [&](std::vector<Object*> objects, std::shared_ptr<Detector> detector = nullptr) {
        std::vector<T> data;
        for(auto& object : objects) {
            data.emplace_back(std::move(*static_cast<T*>(object)));
        }

        if(detector == nullptr) {
            return std::make_shared<Message<T>>(data);
        }
        return std::make_shared<Message<T>>(data, detector);
    };
}

/**
 * Uses SFINAE trick to call the add_creator function for all template arguments of a container class. Used to add creators
 * for every object in a tuple of objects.
 */
template <template <typename...> class T, typename... Args>
static void gen_creator_map_from_tag(ROOTObjectReaderModule::MessageCreatorMap& map, type_tag<T<Args...>>) {
    std::initializer_list<int> value{(add_creator<Args>(map), 0)...};
    (void)value;
}

/**
 * Wrapper function to make the SFINAE trick in \ref gen_creator_map_from_tag work.
 */
template <typename T> static ROOTObjectReaderModule::MessageCreatorMap gen_creator_map() {
    ROOTObjectReaderModule::MessageCreatorMap ret_map;
    gen_creator_map_from_tag(ret_map, type_tag<T>());
    return ret_map;
}

void ROOTObjectReaderModule::init() {
    // Read include and exclude list
    if(config_.has("include") && config_.has("exclude")) {
        throw InvalidCombinationError(
            config_, {"exclude", "include"}, "include and exclude parameter are mutually exclusive");
    } else if(config_.has("include")) {
        auto inc_arr = config_.getArray<std::string>("include");
        include_.insert(inc_arr.begin(), inc_arr.end());
    } else if(config_.has("exclude")) {
        auto exc_arr = config_.getArray<std::string>("exclude");
        exclude_.insert(exc_arr.begin(), exc_arr.end());
    }

    // Initialize the call map from the tuple of available objects
    message_creator_map_ = gen_creator_map<allpix::OBJECTS>();

    // Open the file with the objects
    input_file_ = std::make_unique<TFile>(config_.getPath("file_name", true).c_str());

    // Read all the trees in the file
    TList* keys = input_file_->GetListOfKeys();
    std::set<std::string> tree_names;

    for(auto&& object : *keys) {
        auto& key = dynamic_cast<TKey&>(*object);
        if(std::string(key.GetClassName()) == "TTree") {
            auto tree = static_cast<TTree*>(key.ReadObjectAny(nullptr));

            // Check if a version of this tree has already been read
            if(tree_names.find(tree->GetName()) != tree_names.end()) {
                LOG(TRACE) << "Skipping copy of tree with name " << tree->GetName()
                           << " because one with identical name has already been processed";
                continue;
            }
            tree_names.insert(tree->GetName());

            // Check if this tree should be used
            if((!include_.empty() && include_.find(tree->GetName()) == include_.end()) ||
               (!exclude_.empty() && exclude_.find(tree->GetName()) != exclude_.end())) {
                LOG(TRACE) << "Ignoring tree with " << tree->GetName()
                           << " objects because it has been excluded or not explicitly included";
                continue;
            }

            trees_.push_back(tree);
        }
    }

    if(trees_.empty()) {
        LOG(ERROR) << "Provided ROOT file does not contain any trees, module will not read any data";
    }

    // Cross-check the core random seed stored in the file with the one configured:
    auto& global_config = getConfigManager()->getGlobalConfiguration();
    auto config_seed = global_config.get<uint64_t>("random_seed_core");

    std::string* str = nullptr;
    input_file_->GetObject("config/Allpix/random_seed_core", str);

    if(str == nullptr) {
        // check if missing random seed core in config file should be ignored
        if(config_.get<bool>("ignore_seed_mismatch", false)) {
            LOG(WARNING) << "No random seed for core set in the input data file, cross-check with configured value - "
                         << "this might lead to unexpected behavior. Random seed core from the input data is used.";
        } else {
            throw InvalidValueError(global_config,
                                    "random_seed_core",
                                    "no random seed for core set in the input data file, cross-check with configured value "
                                    "impossible - this might lead to unexpected behavior.");
        }
    } else if(config_seed != allpix::from_string<uint64_t>(*str)) {
        // check if mismatch between random seed core in config and in input file should be ignored
        if(config_.get<bool>("ignore_seed_mismatch", false)) {
            LOG(WARNING) << "Mismatch between core random seed in configuration file and input data"
                         << " - this might lead to unexpected behavior.";
        } else {
            throw InvalidValueError(global_config,
                                    "random_seed_core",
                                    "mismatch between core random seed in configuration file and input data - this "
                                    "might lead to unexpected behavior. Set to value configured in the input data file: " +
                                        (*str));
        }
    }

    // Cross-check version, print warning only in case of a mismatch:
    std::string* version_str = nullptr;
    input_file_->GetObject("config/Allpix/version", version_str);
    if(version_str != nullptr && (*version_str) != ALLPIX_PROJECT_VERSION) {
        LOG(WARNING) << "Reading data produced with different version " << (*version_str)
                     << " - this might lead to unexpected behavior.";
    }

    // Loop over all found trees
    for(auto& tree : trees_) {
        // Loop over the list of branches and create the set of receiver objects
        TObjArray* branches = tree->GetListOfBranches();
        for(int i = 0; i < branches->GetEntries(); i++) {
            auto* branch = static_cast<TBranch*>(branches->At(i));

            // Add a new vector of objects and bind it to the branch
            message_info message_inf;
            message_inf.objects = new std::vector<Object*>;
            message_info_array_.emplace_back(message_inf);
            branch->SetAddress(&(message_info_array_.back().objects));

            // Fill the rest of the message information
            // FIXME: we want to index this in a different way
            std::string branch_name = branch->GetName();
            auto split = allpix::split<std::string>(branch_name, "_");

            // Fetch information from the tree name
            size_t expected_size = 2;
            size_t det_idx = 0;
            size_t name_idx = 1;
            if(branch_name.front() == '_' || branch_name.empty()) {
                --expected_size;
                det_idx = INT_MAX;
                --name_idx;
            }
            if(branch_name.find('_') == std::string::npos) {
                --expected_size;
                name_idx = INT_MAX;
            }

            // Check tree structure and if object type matches name
            auto split_type = allpix::split<std::string>(branch->GetClassName(), "<>");
            if(expected_size != split.size() || split_type.size() != 2 || split_type[1].size() <= 2) {
                throw ModuleError("Tree is malformed and cannot be used for creating messages");
            }
            std::string class_name = split_type[1].substr(0, split_type[1].size() - 1);
            std::string apx_namespace = "allpix::";
            size_t ap_idx = class_name.find(apx_namespace);
            if(ap_idx != std::string::npos) {
                class_name.replace(ap_idx, apx_namespace.size(), "");
            }
            if(class_name != tree->GetName()) {
                throw ModuleError("Tree contains objects of the wrong type");
            }

            if(name_idx != INT_MAX) {
                message_info_array_.back().name = split[name_idx];
            }
            if(det_idx != INT_MAX) {
                if(split[det_idx] != "global") {
                    message_info_array_.back().detector = geo_mgr_->getDetector(split[det_idx]);
                }
            }
        }
    }
}

void ROOTObjectReaderModule::run(unsigned int event_num) {
    --event_num;
    for(auto& tree : trees_) {
        if(event_num >= tree->GetEntries()) {
            throw EndOfRunException("Requesting end of run because TTree only contains data for " +
                                    std::to_string(event_num) + " events");
        }
        tree->GetEntry(event_num);
    }
    LOG(TRACE) << "Building messages from stored objects";

    // Loop through all branches
    for(auto message_inf : message_info_array_) {
        auto objects = message_inf.objects;

        // Skip empty objects in current event
        if(objects->empty()) {
            continue;
        }

        // Check if a pointer to a dispatcher method exist
        auto first_object = (*objects)[0];
        auto iter = message_creator_map_.find(typeid(*first_object));
        if(iter == message_creator_map_.end()) {
            LOG(INFO) << "Cannot dispatch message with object " << allpix::demangle(typeid(*first_object).name())
                      << " because it not registered for messaging";
            continue;
        }

        // Update statistics
        read_cnt_ += objects->size();

        // Create a message
        std::shared_ptr<BaseMessage> message = iter->second(*objects, message_inf.detector);

        // Dispatch the message
        messenger_->dispatchMessage(this, message, message_inf.name);
    }
}

void ROOTObjectReaderModule::finalize() {
    int branch_count = 0;
    for(auto& tree : trees_) {
        branch_count += tree->GetListOfBranches()->GetEntries();
    }

    // Print statistics
    LOG(INFO) << "Read " << read_cnt_ << " objects from " << branch_count << " branches";

    // Close the file
    input_file_->Close();
}
<|endoftext|>


<commit_before><commit_msg>return nullptr in indirect symbol if indirect symbol could not resolve<commit_after><|endoftext|>


<commit_before>/*
 * Copyright (C) 2011-2014 Intel Corporation. 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.
 */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "decodeinputavformat.h"
#include "common/common_def.h"
#include "common/log.h"
#include <Yami.h>

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(55, 39, 100)
#define av_packet_unref av_free_packet
#endif

DecodeInputAvFormat::DecodeInputAvFormat()
:m_format(NULL),m_videoId(-1), m_codecId(AV_CODEC_ID_NONE), m_isEos(true)
{
    av_register_all();

    av_init_packet(&m_packet);
}

bool DecodeInputAvFormat::initInput(const char* fileName)
{
    uint16_t width = 0, height = 0;
    int ret = avformat_open_input(&m_format, fileName, NULL, NULL);
    if (ret)
        goto error;
    uint32_t i;
    for (i = 0; i < m_format->nb_streams; ++i) {
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 33, 100)
        AVCodecParameters* codec = m_format->streams[i]->codecpar;
#else
        AVCodecContext* codec = m_format->streams[i]->codec;
#endif
        //VP9: width and height of the IVF header,VP8: width and height is zero
        if (AVMEDIA_TYPE_VIDEO == codec->codec_type) {
            width = codec->width;
            height = codec->height;
            break;
        }
    }

    ret = avformat_find_stream_info(m_format, NULL);
    if (ret < 0)
        goto error;
    for (i = 0; i < m_format->nb_streams; i++)
    {
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 33, 100)
        AVCodecParameters* codec = m_format->streams[i]->codecpar;
#else
        AVCodecContext* codec = m_format->streams[i]->codec;
#endif
        if (AVMEDIA_TYPE_VIDEO == codec->codec_type) {
            m_codecId = codec->codec_id;
            if (codec->extradata && codec->extradata_size)
                m_codecData.append((char*)codec->extradata, codec->extradata_size);
            m_videoId = i;
            //VP9: display_width and display_height of the first frame
            if (codec->width > width)
                width = codec->width;
            if (codec->height > height)
                height = codec->height;
            setResolution(width, height);
            break;
        }
    }
    if (i == m_format->nb_streams) {
        ERROR("no video stream");
        goto error;
    }
    m_isEos = false;
    return true;
error:
    if (m_format)
        avformat_close_input(&m_format);
    return false;

}

struct MimeEntry
{
    AVCodecID id;
    const char* mime;
};

static const MimeEntry MimeEntrys[] = {
    AV_CODEC_ID_MPEG2VIDEO, YAMI_MIME_MPEG2,
    AV_CODEC_ID_VP8, YAMI_MIME_VP8,

#if LIBAVCODEC_VERSION_INT > AV_VERSION_INT(54, 40, 0)
    AV_CODEC_ID_VP9, YAMI_MIME_VP9,
#endif

    AV_CODEC_ID_H264, YAMI_MIME_H264,
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 39, 100)
    AV_CODEC_ID_H265, YAMI_MIME_H265,
#endif
    AV_CODEC_ID_WMV3, YAMI_MIME_VC1,
    AV_CODEC_ID_VC1, YAMI_MIME_VC1
};

const char * DecodeInputAvFormat::getMimeType()
{
    for (size_t i = 0; i < N_ELEMENTS(MimeEntrys); i++) {
        if (MimeEntrys[i].id == m_codecId)
            return MimeEntrys[i].mime;
    }
    return "unknow";
}

bool DecodeInputAvFormat::getNextDecodeUnit(VideoDecodeBuffer &inputBuffer)
{
    if (!m_format || m_isEos)
        return false;
    int ret;
    while (1) {
        //free old packet
        av_packet_unref(&m_packet);

        ret = av_read_frame(m_format, &m_packet);
        if (ret) {
            m_isEos = true;
            return false;
        }
        if (m_packet.stream_index == m_videoId) {
            memset(&inputBuffer, 0, sizeof(inputBuffer));
            inputBuffer.data = m_packet.data;
            inputBuffer.size = m_packet.size;
            inputBuffer.timeStamp = m_packet.dts;
            return true;
        }
    }
    return false;
}

const string& DecodeInputAvFormat::getCodecData()
{
    return m_codecData;
}

DecodeInputAvFormat::~DecodeInputAvFormat()
{
    if (m_format) {
        av_packet_unref(&m_packet);
        avformat_close_input(&m_format);
    }

}
<commit_msg>decode avformat: av_register_all is deprecated since 58.9.100<commit_after>/*
 * Copyright (C) 2011-2014 Intel Corporation. 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.
 */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "decodeinputavformat.h"
#include "common/common_def.h"
#include "common/log.h"
#include <Yami.h>

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(55, 39, 100)
#define av_packet_unref av_free_packet
#endif

DecodeInputAvFormat::DecodeInputAvFormat()
:m_format(NULL),m_videoId(-1), m_codecId(AV_CODEC_ID_NONE), m_isEos(true)
{
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(58, 9, 100)
    av_register_all();
#endif

    av_init_packet(&m_packet);
}

bool DecodeInputAvFormat::initInput(const char* fileName)
{
    uint16_t width = 0, height = 0;
    int ret = avformat_open_input(&m_format, fileName, NULL, NULL);
    if (ret)
        goto error;
    uint32_t i;
    for (i = 0; i < m_format->nb_streams; ++i) {
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 33, 100)
        AVCodecParameters* codec = m_format->streams[i]->codecpar;
#else
        AVCodecContext* codec = m_format->streams[i]->codec;
#endif
        //VP9: width and height of the IVF header,VP8: width and height is zero
        if (AVMEDIA_TYPE_VIDEO == codec->codec_type) {
            width = codec->width;
            height = codec->height;
            break;
        }
    }

    ret = avformat_find_stream_info(m_format, NULL);
    if (ret < 0)
        goto error;
    for (i = 0; i < m_format->nb_streams; i++)
    {
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 33, 100)
        AVCodecParameters* codec = m_format->streams[i]->codecpar;
#else
        AVCodecContext* codec = m_format->streams[i]->codec;
#endif
        if (AVMEDIA_TYPE_VIDEO == codec->codec_type) {
            m_codecId = codec->codec_id;
            if (codec->extradata && codec->extradata_size)
                m_codecData.append((char*)codec->extradata, codec->extradata_size);
            m_videoId = i;
            //VP9: display_width and display_height of the first frame
            if (codec->width > width)
                width = codec->width;
            if (codec->height > height)
                height = codec->height;
            setResolution(width, height);
            break;
        }
    }
    if (i == m_format->nb_streams) {
        ERROR("no video stream");
        goto error;
    }
    m_isEos = false;
    return true;
error:
    if (m_format)
        avformat_close_input(&m_format);
    return false;

}

struct MimeEntry
{
    AVCodecID id;
    const char* mime;
};

static const MimeEntry MimeEntrys[] = {
    AV_CODEC_ID_MPEG2VIDEO, YAMI_MIME_MPEG2,
    AV_CODEC_ID_VP8, YAMI_MIME_VP8,

#if LIBAVCODEC_VERSION_INT > AV_VERSION_INT(54, 40, 0)
    AV_CODEC_ID_VP9, YAMI_MIME_VP9,
#endif

    AV_CODEC_ID_H264, YAMI_MIME_H264,
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 39, 100)
    AV_CODEC_ID_H265, YAMI_MIME_H265,
#endif
    AV_CODEC_ID_WMV3, YAMI_MIME_VC1,
    AV_CODEC_ID_VC1, YAMI_MIME_VC1
};

const char * DecodeInputAvFormat::getMimeType()
{
    for (size_t i = 0; i < N_ELEMENTS(MimeEntrys); i++) {
        if (MimeEntrys[i].id == m_codecId)
            return MimeEntrys[i].mime;
    }
    return "unknow";
}

bool DecodeInputAvFormat::getNextDecodeUnit(VideoDecodeBuffer &inputBuffer)
{
    if (!m_format || m_isEos)
        return false;
    int ret;
    while (1) {
        //free old packet
        av_packet_unref(&m_packet);

        ret = av_read_frame(m_format, &m_packet);
        if (ret) {
            m_isEos = true;
            return false;
        }
        if (m_packet.stream_index == m_videoId) {
            memset(&inputBuffer, 0, sizeof(inputBuffer));
            inputBuffer.data = m_packet.data;
            inputBuffer.size = m_packet.size;
            inputBuffer.timeStamp = m_packet.dts;
            return true;
        }
    }
    return false;
}

const string& DecodeInputAvFormat::getCodecData()
{
    return m_codecData;
}

DecodeInputAvFormat::~DecodeInputAvFormat()
{
    if (m_format) {
        av_packet_unref(&m_packet);
        avformat_close_input(&m_format);
    }

}
<|endoftext|>