<commit_before>// @(#) $Id$

/**************************************************************************
 * This file is property of and copyright by the ALICE HLT Project        *
 * ALICE Experiment at CERN, All rights reserved.                         *
 *                                                                        *
 * Primary Authors: Matthias Richter <Matthias.Richter@ift.uib.no>        *
 *                  for The ALICE HLT Project.                            *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, provided that the above copyright notice appears in all   *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purpose. It is          *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/

/** @file   AliHLTTPCDigitPublisherComponent.cxx
    @author Matthias Richter
    @date
    @brief  TPC digit publisher component (input from offline).
*/

// see header file for class documentation
// or
// refer to README to build package
// or
// visit http://web.ift.uib.no/~kjeks/doc/alice-hlt

#include "AliHLTTPCDigitPublisherComponent.h"
#include "AliRunLoader.h"
#include "AliLog.h"
#include "TTree.h"
#include "AliHLTTPCDigitData.h"
#include "AliHLTTPCDefinitions.h"
#include "AliHLTTPCFileHandler.h"

/** ROOT macro for the implementation of ROOT specific class methods */
ClassImp(AliHLTTPCDigitPublisherComponent)

AliHLTTPCDigitPublisherComponent::AliHLTTPCDigitPublisherComponent()
  :
  fMaxSize(200000), // just a number to start with
  fMinSlice(-1),
  fMinPart(-1)
{
  // see header file for class documentation
  // or
  // refer to README to build package
  // or
  // visit http://web.ift.uib.no/~kjeks/doc/alice-hlt
}

AliHLTTPCFileHandler* AliHLTTPCDigitPublisherComponent::fgpFileHandler=NULL;
int AliHLTTPCDigitPublisherComponent::fgFileHandlerInstances=0;
int AliHLTTPCDigitPublisherComponent::fgCurrEvent=-1;

AliHLTTPCDigitPublisherComponent::~AliHLTTPCDigitPublisherComponent()
{
  // see header file for class documentation
  if (fgpFileHandler!=NULL && fgFileHandlerInstances<=0) {
    HLTWarning("improper state, de-initialization missing");
    DoDeinit();
  }
}

const char* AliHLTTPCDigitPublisherComponent::GetComponentID()
{
  // see header file for class documentation
  return "TPCDigitPublisher";
}

AliHLTComponentDataType AliHLTTPCDigitPublisherComponent::GetOutputDataType()
{
  return AliHLTTPCDefinitions::fgkUnpackedRawDataType;
}

void AliHLTTPCDigitPublisherComponent::GetOutputDataSize( unsigned long& constBase, double& inputMultiplier )
{
  constBase=fMaxSize;
  inputMultiplier=1;
}

AliHLTComponent* AliHLTTPCDigitPublisherComponent::Spawn()
{
  // see header file for class documentation
  return new AliHLTTPCDigitPublisherComponent;
}

int AliHLTTPCDigitPublisherComponent::DoInit( int argc, const char** argv )
{
  // see header file for class documentation
  int iResult=0;

  // scan arguments
  TString argument="";
  int bMissingParam=0;
  for (int i=0; i<argc && iResult>=0; i++) {
    argument=argv[i];
    if (argument.IsNull()) continue;

    // -slice
    if (argument.CompareTo("-slice")==0) {
      if ((bMissingParam=(++i>=argc))) break;
      TString parameter(argv[i]);
      parameter.Remove(TString::kLeading, ' '); // remove all blanks
      if (parameter.IsDigit()) {
	fMinSlice=parameter.Atoi();
      } else {
	HLTError("wrong parameter for argument %s, number expected", argument.Data());
	iResult=-EINVAL;
      }
      // -partition
    } else if (argument.CompareTo("-partition")==0) {
      if ((bMissingParam=(++i>=argc))) break;
      TString parameter(argv[i]);
      parameter.Remove(TString::kLeading, ' '); // remove all blanks
      if (parameter.IsDigit()) {
	fMinPart=parameter.Atoi();
      } else {
	HLTError("wrong parameter for argument %s, number expected", argument.Data());
	iResult=-EINVAL;
      }
    } else {
      HLTError("unknown argument %s", argument.Data());
      iResult=-EINVAL;
    }
  }
  if (bMissingParam) {
    HLTError("missing parameter for argument %s", argument.Data());
    iResult=-EINVAL;
  }

  if (fMinSlice<0) {
    HLTError("slice no required");
    iResult=-EINVAL;
  }

  if (fMinPart<0) {
    HLTError("partition (patch) no required");
    iResult=-EINVAL;
  }

  if (iResult<0) return iResult;

  // fetch runLoader instance from interface
  AliRunLoader* pRunLoader=GetRunLoader();
  if (pRunLoader) {
    if (fgpFileHandler==NULL) {
      fgpFileHandler=new AliHLTTPCFileHandler;
      fgCurrEvent=-1;
      fgFileHandlerInstances=1;
    } else {
      fgFileHandlerInstances++;
      //HLTDebug("publisher %p: %d references to file handler instance", this, fgFileHandlerInstances);
    }
    if (fgpFileHandler) {
      if (!fgpFileHandler->SetAliInput(pRunLoader)) {
	iResult=-EFAULT;
      }
    } else {
      AliErrorStream() << "can not allocate file handler object" << endl;
      iResult=-ENOMEM;
    }
  } else {
    AliErrorStream() << "can not get runLoader" << endl;
    iResult=-EFAULT;
  }
  return iResult;
}

int AliHLTTPCDigitPublisherComponent::DoDeinit()
{
  // see header file for class documentation
  int iResult=0;
  if (fgCurrEvent>=0) {
    fgpFileHandler->FreeDigitsTree();
    fgCurrEvent=-1;
  }
  //HLTDebug("publisher %p: %d references to file handler instance", this, fgFileHandlerInstances);
  if (--fgFileHandlerInstances==0 && fgpFileHandler!=NULL) {
    try {
      if (fgpFileHandler) {
	delete fgpFileHandler;
      }
    }
    catch (...) {
      HLTFatal("exeption during object cleanup");
      iResult=-EFAULT;
    }
    fgpFileHandler=NULL;
  }
  return iResult;
}

int AliHLTTPCDigitPublisherComponent::GetEvent(const AliHLTComponentEventData& /*evtData*/,
					       AliHLTComponentTriggerData& /*trigData*/,
					       AliHLTUInt8_t* outputPtr,
					       AliHLTUInt32_t& size,
					       vector<AliHLTComponentBlockData>& outputBlocks)
{
  // see header file for class documentation
  int iResult=0;
  if (outputPtr==NULL || size==0) {
    HLTError("no target buffer provided");
    return -EFAULT;
  }

  if (fgpFileHandler) {
    int event=GetEventCount();
    AliHLTTPCUnpackedRawData* pTgt=reinterpret_cast<AliHLTTPCUnpackedRawData*>(outputPtr);
    if (pTgt) {
      UInt_t nrow=0;
      UInt_t tgtSize=size-sizeof(AliHLTTPCUnpackedRawData);
      if (fgCurrEvent>=0 && fgCurrEvent!=event) {
	HLTDebug("new event %d, free digit tree for event %d", event, fgCurrEvent);
	fgpFileHandler->FreeDigitsTree();
      }
      fgCurrEvent=event;
      HLTDebug("converting digits for slice %d partition %d", fMinSlice, fMinPart);
      fgpFileHandler->Init(fMinSlice,fMinPart);
      AliHLTTPCDigitRowData* pData=fgpFileHandler->AliDigits2Memory(nrow, event, reinterpret_cast<Byte_t*>(pTgt->fDigits), &tgtSize);
      if (pData==NULL && tgtSize>0 && tgtSize>fMaxSize) {
	HLTDebug("target buffer too small: %d byte required, %d available", tgtSize+sizeof(AliHLTTPCUnpackedRawData), size);
	// indicate insufficient buffer size, on occasion the frameworks calls
	// again with the corrected buffer
	fMaxSize=tgtSize;
	iResult=-ENOSPC;
      } else if (pData!=pTgt->fDigits) {
	HLTError("can not read directly into output buffer");
	try {delete pData;}
	catch (...) {/* no action */}
	iResult=-EIO;
      } else {
	size=tgtSize+sizeof(AliHLTTPCUnpackedRawData);
	AliHLTComponentBlockData bd;
	FillBlockData( bd );
	bd.fOffset = 0;
	bd.fSize = size;
	bd.fSpecification = AliHLTTPCDefinitions::EncodeDataSpecification(fMinSlice, fMinSlice, fMinPart, fMinPart);
	outputBlocks.push_back( bd );
	HLTDebug("added AliHLTTPCUnpackedRawData size %d, first row %d nof digits %d", size, pTgt->fDigits->fRow, pTgt->fDigits->fNDigit);
      }
    }
  } else {
    AliErrorStream() << "component not initialized" << endl;
    iResult=-EFAULT;
  }
  return iResult;
}
<commit_msg>bugfix: digit publisher was crashing due to malfunction in SOR/EOR handling<commit_after>// @(#) $Id$

/**************************************************************************
 * This file is property of and copyright by the ALICE HLT Project        *
 * ALICE Experiment at CERN, All rights reserved.                         *
 *                                                                        *
 * Primary Authors: Matthias Richter <Matthias.Richter@ift.uib.no>        *
 *                  for The ALICE HLT Project.                            *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, provided that the above copyright notice appears in all   *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purpose. It is          *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/

/** @file   AliHLTTPCDigitPublisherComponent.cxx
    @author Matthias Richter
    @date
    @brief  TPC digit publisher component (input from offline).
*/

// see header file for class documentation
// or
// refer to README to build package
// or
// visit http://web.ift.uib.no/~kjeks/doc/alice-hlt

#include "AliHLTTPCDigitPublisherComponent.h"
#include "AliRunLoader.h"
#include "AliLog.h"
#include "TTree.h"
#include "AliHLTTPCDigitData.h"
#include "AliHLTTPCDefinitions.h"
#include "AliHLTTPCFileHandler.h"

/** ROOT macro for the implementation of ROOT specific class methods */
ClassImp(AliHLTTPCDigitPublisherComponent)

AliHLTTPCDigitPublisherComponent::AliHLTTPCDigitPublisherComponent()
  :
  fMaxSize(200000), // just a number to start with
  fMinSlice(-1),
  fMinPart(-1)
{
  // see header file for class documentation
  // or
  // refer to README to build package
  // or
  // visit http://web.ift.uib.no/~kjeks/doc/alice-hlt
}

AliHLTTPCFileHandler* AliHLTTPCDigitPublisherComponent::fgpFileHandler=NULL;
int AliHLTTPCDigitPublisherComponent::fgFileHandlerInstances=0;
int AliHLTTPCDigitPublisherComponent::fgCurrEvent=-1;

AliHLTTPCDigitPublisherComponent::~AliHLTTPCDigitPublisherComponent()
{
  // see header file for class documentation
  if (fgpFileHandler!=NULL && fgFileHandlerInstances<=0) {
    HLTWarning("improper state, de-initialization missing");
    DoDeinit();
  }
}

const char* AliHLTTPCDigitPublisherComponent::GetComponentID()
{
  // see header file for class documentation
  return "TPCDigitPublisher";
}

AliHLTComponentDataType AliHLTTPCDigitPublisherComponent::GetOutputDataType()
{
  return AliHLTTPCDefinitions::fgkUnpackedRawDataType;
}

void AliHLTTPCDigitPublisherComponent::GetOutputDataSize( unsigned long& constBase, double& inputMultiplier )
{
  constBase=fMaxSize;
  inputMultiplier=1;
}

AliHLTComponent* AliHLTTPCDigitPublisherComponent::Spawn()
{
  // see header file for class documentation
  return new AliHLTTPCDigitPublisherComponent;
}

int AliHLTTPCDigitPublisherComponent::DoInit( int argc, const char** argv )
{
  // see header file for class documentation
  int iResult=0;

  // scan arguments
  TString argument="";
  int bMissingParam=0;
  for (int i=0; i<argc && iResult>=0; i++) {
    argument=argv[i];
    if (argument.IsNull()) continue;

    // -slice
    if (argument.CompareTo("-slice")==0) {
      if ((bMissingParam=(++i>=argc))) break;
      TString parameter(argv[i]);
      parameter.Remove(TString::kLeading, ' '); // remove all blanks
      if (parameter.IsDigit()) {
	fMinSlice=parameter.Atoi();
      } else {
	HLTError("wrong parameter for argument %s, number expected", argument.Data());
	iResult=-EINVAL;
      }
      // -partition
    } else if (argument.CompareTo("-partition")==0) {
      if ((bMissingParam=(++i>=argc))) break;
      TString parameter(argv[i]);
      parameter.Remove(TString::kLeading, ' '); // remove all blanks
      if (parameter.IsDigit()) {
	fMinPart=parameter.Atoi();
      } else {
	HLTError("wrong parameter for argument %s, number expected", argument.Data());
	iResult=-EINVAL;
      }
    } else {
      HLTError("unknown argument %s", argument.Data());
      iResult=-EINVAL;
    }
  }
  if (bMissingParam) {
    HLTError("missing parameter for argument %s", argument.Data());
    iResult=-EINVAL;
  }

  if (fMinSlice<0) {
    HLTError("slice no required");
    iResult=-EINVAL;
  }

  if (fMinPart<0) {
    HLTError("partition (patch) no required");
    iResult=-EINVAL;
  }

  if (iResult<0) return iResult;

  // fetch runLoader instance from interface
  AliRunLoader* pRunLoader=GetRunLoader();
  if (pRunLoader) {
    if (fgpFileHandler==NULL) {
      fgpFileHandler=new AliHLTTPCFileHandler;
      fgCurrEvent=-1;
      fgFileHandlerInstances=1;
    } else {
      fgFileHandlerInstances++;
      //HLTDebug("publisher %p: %d references to file handler instance", this, fgFileHandlerInstances);
    }
    if (fgpFileHandler) {
      if (!fgpFileHandler->SetAliInput(pRunLoader)) {
	iResult=-EFAULT;
      }
    } else {
      AliErrorStream() << "can not allocate file handler object" << endl;
      iResult=-ENOMEM;
    }
  } else {
    AliErrorStream() << "can not get runLoader" << endl;
    iResult=-EFAULT;
  }
  return iResult;
}

int AliHLTTPCDigitPublisherComponent::DoDeinit()
{
  // see header file for class documentation
  int iResult=0;
  if (fgCurrEvent>=0) {
    fgpFileHandler->FreeDigitsTree();
    fgCurrEvent=-1;
  }
  //HLTDebug("publisher %p: %d references to file handler instance", this, fgFileHandlerInstances);
  if (--fgFileHandlerInstances==0 && fgpFileHandler!=NULL) {
    try {
      if (fgpFileHandler) {
	delete fgpFileHandler;
      }
    }
    catch (...) {
      HLTFatal("exeption during object cleanup");
      iResult=-EFAULT;
    }
    fgpFileHandler=NULL;
  }
  return iResult;
}

int AliHLTTPCDigitPublisherComponent::GetEvent(const AliHLTComponentEventData& /*evtData*/,
					       AliHLTComponentTriggerData& /*trigData*/,
					       AliHLTUInt8_t* outputPtr,
					       AliHLTUInt32_t& osize,
					       vector<AliHLTComponentBlockData>& outputBlocks)
{
  // see header file for class documentation
  int iResult=0;
  AliHLTUInt32_t capacity=osize;
  osize=0;

  // process data events only
  if (!IsDataEvent()) return 0;

  if (outputPtr==NULL || capacity==0) {
    HLTError("no target buffer provided");
    return -EFAULT;
  }

  if (fgpFileHandler) {
    int event=GetEventCount();
    AliHLTTPCUnpackedRawData* pTgt=reinterpret_cast<AliHLTTPCUnpackedRawData*>(outputPtr);
    if (pTgt) {
      UInt_t nrow=0;
      UInt_t tgtSize=capacity-sizeof(AliHLTTPCUnpackedRawData);
      if (fgCurrEvent>=0 && fgCurrEvent!=event) {
	HLTDebug("new event %d, free digit tree for event %d", event, fgCurrEvent);
	fgpFileHandler->FreeDigitsTree();
      }
      fgCurrEvent=event;
      HLTDebug("converting digits for slice %d partition %d", fMinSlice, fMinPart);
      fgpFileHandler->Init(fMinSlice,fMinPart);
      AliHLTTPCDigitRowData* pData=fgpFileHandler->AliDigits2Memory(nrow, event, reinterpret_cast<Byte_t*>(pTgt->fDigits), &tgtSize);
      if (pData==NULL && tgtSize>0 && tgtSize>fMaxSize) {
	HLTDebug("target buffer too small: %d byte required, %d available", tgtSize+sizeof(AliHLTTPCUnpackedRawData), capacity);
	// indicate insufficient buffer size, on occasion the frameworks calls
	// again with the corrected buffer
	fMaxSize=tgtSize;
	iResult=-ENOSPC;
      } else if (pData!=pTgt->fDigits) {
	HLTError("can not read directly into output buffer");
	try {delete pData;}
	catch (...) {/* no action */}
	iResult=-EIO;
      } else {
	osize=tgtSize+sizeof(AliHLTTPCUnpackedRawData);
	AliHLTComponentBlockData bd;
	FillBlockData( bd );
	bd.fOffset = 0;
	bd.fSize = osize;
	bd.fSpecification = AliHLTTPCDefinitions::EncodeDataSpecification(fMinSlice, fMinSlice, fMinPart, fMinPart);
	outputBlocks.push_back( bd );
	HLTDebug("added AliHLTTPCUnpackedRawData size %d, first row %d nof digits %d", osize, pTgt->fDigits->fRow, pTgt->fDigits->fNDigit);
      }
    }
  } else {
    AliErrorStream() << "component not initialized" << endl;
    iResult=-EFAULT;
  }
  return iResult;
}
<|endoftext|>


<commit_before>#include <ir/index_manager/index/IndexReader.h>
#include <ir/index_manager/index/TermReader.h>
#include <ir/index_manager/index/Indexer.h>
#include <ir/index_manager/index/MultiIndexBarrelReader.h>
#include <util/izene_log.h>

#include <util/ThreadModel.h>

using namespace izenelib::ir::indexmanager;

IndexReader::IndexReader(Indexer* pIndex)
        :pIndexer_(pIndex)
        ,pBarrelsInfo_(NULL)
        ,pBarrelReader_(NULL)
        ,pDocFilter_(NULL)
        ,pDocLengthReader_(NULL)
{
    pBarrelsInfo_ = pIndexer_->getBarrelsInfo();

    Directory* pDirectory = pIndexer_->getDirectory();
    if(pDirectory->fileExists(DELETED_DOCS))
    {
        pDocFilter_ = new BitVector;
        pDocFilter_->read(pDirectory, DELETED_DOCS);
    }
    ///todo since only one collection is used within indexmanager, so we only get collectionmeta for one collection to build
    ///up the DocLengthReader
    if(pIndexer_->getIndexManagerConfig()->indexStrategy_.indexDocLength_)
        pDocLengthReader_ = new DocLengthReader(
                                pIndexer_->getCollectionsMeta().begin()->second.getDocumentSchema(),
                                pIndexer_->getDirectory());
}

IndexReader::~IndexReader()
{
    if (pBarrelReader_)
    {
        delete pBarrelReader_;
        pBarrelReader_ = NULL;
    }

    flush();
    delete pDocFilter_;
    pDocFilter_ = NULL;

    if(pDocLengthReader_)
    {
        delete pDocLengthReader_;
        pDocLengthReader_ = NULL;
    }
}

void IndexReader::delDocFilter()
{
    DVLOG(2) << "=> IndexReader::delDocFilter(), pDocFilter_: " << pDocFilter_;
    if(pDocFilter_)
    {
        if(pIndexer_->getDirectory()->fileExists(DELETED_DOCS))
            pIndexer_->getDirectory()->deleteFile(DELETED_DOCS);
        pDocFilter_->clear();
    }
    DVLOG(2) << "<= IndexReader::delDocFilter()";
}

void IndexReader::flush()
{
    if(pDocFilter_)
    {
        boost::mutex::scoped_lock docFilterLock(docFilterMutex_);

        if(pDocFilter_->any())
            pDocFilter_->write(pIndexer_->getDirectory(), DELETED_DOCS);
    }
}

docid_t IndexReader::maxDoc()
{
    return pBarrelsInfo_->maxDocId();
}

size_t IndexReader::docLength(docid_t docId, fieldid_t fid)
{
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    return pDocLengthReader_->docLength(docId, fid);
}

double IndexReader::getAveragePropertyLength(fieldid_t fid)
{
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    boost::mutex::scoped_lock indexReaderLock(this->mutex_);
    return pDocLengthReader_->averagePropertyLength(fid);
}

void IndexReader::createBarrelReader()
{
    boost::mutex::scoped_lock lock(this->mutex_);

    if (pBarrelReader_)
        return;

    DVLOG(2) << "=> IndexReader::createBarrelReader()";

    if(int32_t bc = pBarrelsInfo_->getBarrelCount())
    {
        DVLOG(2) << "IndexReader::createBarrelReader() => " << bc << " barrels...";

        pBarrelReader_ = new MultiIndexBarrelReader(this, pBarrelsInfo_);

        if(pDocLengthReader_)
            pDocLengthReader_->load(pBarrelsInfo_->maxDocId());
    }

    DVLOG(2) << "<= IndexReader::createBarrelReader()";
}

TermReader* IndexReader::getTermReader(collectionid_t colID)
{
    //boost::try_mutex::scoped_try_lock lock(pIndexer_->mutex_);
    //if(!lock.owns_lock())
        //return NULL;
    izenelib::util::ScopedReadLock<izenelib::util::ReadWriteLock> lock(pIndexer_->mutex_);
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return NULL;
    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    TermReader* pTermReader = pBarrelReader_->termReader(colID);
    if (pTermReader)
    {
        pTermReader->setSkipInterval(pIndexer_->getSkipInterval());
        pTermReader->setMaxSkipLevel(pIndexer_->getMaxSkipLevel());
        if(pDocFilter_)
            pTermReader->setDocFilter(pDocFilter_);
        return pTermReader->clone();
    }
    else
        return NULL;
}

void IndexReader::reopen()
{
    {
    boost::mutex::scoped_lock lock(this->mutex_);
    if(pBarrelReader_)
    {
        delete pBarrelReader_;
        pBarrelReader_ = NULL;
    }
    }
    createBarrelReader();
}

count_t IndexReader::numDocs()
{
    return pBarrelsInfo_->getDocCount();
}

void IndexReader::delDocument(collectionid_t colID,docid_t docId)
{
    DVLOG(4) << "=> IndexReader::delDocument(), collection id: " << colID << ", doc id: " << docId << " ...";

    if(pBarrelsInfo_->deleteDocument(colID, docId))
    {
        boost::mutex::scoped_lock docFilterLock(docFilterMutex_);

        if(!pDocFilter_)
        {
            // bit 0 is reserved, so that for doc id i, we can access it just using bit i
            pDocFilter_ = new BitVector(pBarrelsInfo_->maxDocId() + 1);
        }

        pDocFilter_->set(docId);
    }

    DVLOG(4) << "<= IndexReader::delDocument()";
}

freq_t IndexReader::docFreq(collectionid_t colID, Term* term)
{
    //boost::try_mutex::scoped_try_lock lock(pIndexer_->mutex_);
    //if(!lock.owns_lock())
        //return 0;
    izenelib::util::ScopedReadLock<izenelib::util::ReadWriteLock> lock(pIndexer_->mutex_);

    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return 0;

    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    TermReader* pTermReader = pBarrelReader_->termReader(colID);
    if (pTermReader)
    {
        return pTermReader->docFreq(term);
    }

    return 0;
}

TermInfo* IndexReader::termInfo(collectionid_t colID,Term* term)
{
    //boost::try_mutex::scoped_try_lock lock(pIndexer_->mutex_);
    //if(!lock.owns_lock())
        //return NULL;
    izenelib::util::ScopedReadLock<izenelib::util::ReadWriteLock> lock(pIndexer_->mutex_);

    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return NULL;

    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    TermReader* pTermReader = pBarrelReader_->termReader(colID);
    if (pTermReader)
    {
        return pTermReader->termInfo(term);
    }

    return 0;
}

size_t IndexReader::getDistinctNumTerms(collectionid_t colID, const std::string& property)
{
    //collection has been removed, need to rebuild the barrel reader
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return 0;

    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    return pBarrelReader_->getDistinctNumTerms(colID, property);
}

<commit_msg>undo a previous commit in 9d283a93b5825c25611eff37459ea87794d03df3, that is, in IndexReader::createBarrelReader(), instead of always using MultiIndexBarrelReader, use SingleIndexBarrelReader when there is only one barrel.<commit_after>#include <ir/index_manager/index/IndexReader.h>
#include <ir/index_manager/index/TermReader.h>
#include <ir/index_manager/index/Indexer.h>
#include <ir/index_manager/index/MultiIndexBarrelReader.h>
#include <ir/index_manager/index/IndexBarrelWriter.h>
#include <util/izene_log.h>

#include <util/ThreadModel.h>

using namespace izenelib::ir::indexmanager;

IndexReader::IndexReader(Indexer* pIndex)
        :pIndexer_(pIndex)
        ,pBarrelsInfo_(NULL)
        ,pBarrelReader_(NULL)
        ,pDocFilter_(NULL)
        ,pDocLengthReader_(NULL)
{
    pBarrelsInfo_ = pIndexer_->getBarrelsInfo();

    Directory* pDirectory = pIndexer_->getDirectory();
    if(pDirectory->fileExists(DELETED_DOCS))
    {
        pDocFilter_ = new BitVector;
        pDocFilter_->read(pDirectory, DELETED_DOCS);
    }
    ///todo since only one collection is used within indexmanager, so we only get collectionmeta for one collection to build
    ///up the DocLengthReader
    if(pIndexer_->getIndexManagerConfig()->indexStrategy_.indexDocLength_)
        pDocLengthReader_ = new DocLengthReader(
                                pIndexer_->getCollectionsMeta().begin()->second.getDocumentSchema(),
                                pIndexer_->getDirectory());
}

IndexReader::~IndexReader()
{
    if (pBarrelReader_)
    {
        delete pBarrelReader_;
        pBarrelReader_ = NULL;
    }

    flush();
    delete pDocFilter_;
    pDocFilter_ = NULL;

    if(pDocLengthReader_)
    {
        delete pDocLengthReader_;
        pDocLengthReader_ = NULL;
    }
}

void IndexReader::delDocFilter()
{
    DVLOG(2) << "=> IndexReader::delDocFilter(), pDocFilter_: " << pDocFilter_;
    if(pDocFilter_)
    {
        if(pIndexer_->getDirectory()->fileExists(DELETED_DOCS))
            pIndexer_->getDirectory()->deleteFile(DELETED_DOCS);
        pDocFilter_->clear();
    }
    DVLOG(2) << "<= IndexReader::delDocFilter()";
}

void IndexReader::flush()
{
    if(pDocFilter_)
    {
        boost::mutex::scoped_lock docFilterLock(docFilterMutex_);

        if(pDocFilter_->any())
            pDocFilter_->write(pIndexer_->getDirectory(), DELETED_DOCS);
    }
}

docid_t IndexReader::maxDoc()
{
    return pBarrelsInfo_->maxDocId();
}

size_t IndexReader::docLength(docid_t docId, fieldid_t fid)
{
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    return pDocLengthReader_->docLength(docId, fid);
}

double IndexReader::getAveragePropertyLength(fieldid_t fid)
{
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    boost::mutex::scoped_lock indexReaderLock(this->mutex_);
    return pDocLengthReader_->averagePropertyLength(fid);
}

void IndexReader::createBarrelReader()
{
    boost::mutex::scoped_lock lock(this->mutex_);

    if(pBarrelReader_)
        return;

    DVLOG(2) << "=> IndexReader::createBarrelReader()";

    if(int32_t bc = pBarrelsInfo_->getBarrelCount())
    {
        DVLOG(2) << "IndexReader::createBarrelReader() => " << bc << " barrels...";

        if(bc == 1)
        {
            boost::mutex::scoped_lock lock(pBarrelsInfo_->getMutex());

            pBarrelsInfo_->startIterator();
            if(pBarrelsInfo_->hasNext())
            {
                BarrelInfo* pBarrelInfo = pBarrelsInfo_->next();
                if(pBarrelInfo->isSearchable() && pBarrelInfo->getDocCount() > 0)
                {
                    if(IndexBarrelWriter* pBarrelWriter = pBarrelInfo->getWriter())
                        pBarrelReader_ = pBarrelWriter->inMemoryReader();
                    else
                        pBarrelReader_ = new SingleIndexBarrelReader(this, pBarrelInfo);
                }
            }
        }
        else if(bc > 1)
            pBarrelReader_ = new MultiIndexBarrelReader(this, pBarrelsInfo_);

        if(pDocLengthReader_)
            pDocLengthReader_->load(pBarrelsInfo_->maxDocId());
    }

    DVLOG(2) << "<= IndexReader::createBarrelReader()";
}

TermReader* IndexReader::getTermReader(collectionid_t colID)
{
    //boost::try_mutex::scoped_try_lock lock(pIndexer_->mutex_);
    //if(!lock.owns_lock())
        //return NULL;
    izenelib::util::ScopedReadLock<izenelib::util::ReadWriteLock> lock(pIndexer_->mutex_);
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return NULL;
    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    TermReader* pTermReader = pBarrelReader_->termReader(colID);
    if (pTermReader)
    {
        pTermReader->setSkipInterval(pIndexer_->getSkipInterval());
        pTermReader->setMaxSkipLevel(pIndexer_->getMaxSkipLevel());
        if(pDocFilter_)
            pTermReader->setDocFilter(pDocFilter_);
        return pTermReader->clone();
    }
    else
        return NULL;
}

void IndexReader::reopen()
{
    {
    boost::mutex::scoped_lock lock(this->mutex_);
    if(pBarrelReader_)
    {
        delete pBarrelReader_;
        pBarrelReader_ = NULL;
    }
    }
    createBarrelReader();
}

count_t IndexReader::numDocs()
{
    return pBarrelsInfo_->getDocCount();
}

void IndexReader::delDocument(collectionid_t colID,docid_t docId)
{
    DVLOG(4) << "=> IndexReader::delDocument(), collection id: " << colID << ", doc id: " << docId << " ...";

    if(pBarrelsInfo_->deleteDocument(colID, docId))
    {
        boost::mutex::scoped_lock docFilterLock(docFilterMutex_);

        if(!pDocFilter_)
        {
            // bit 0 is reserved, so that for doc id i, we can access it just using bit i
            pDocFilter_ = new BitVector(pBarrelsInfo_->maxDocId() + 1);
        }

        pDocFilter_->set(docId);
    }

    DVLOG(4) << "<= IndexReader::delDocument()";
}

freq_t IndexReader::docFreq(collectionid_t colID, Term* term)
{
    //boost::try_mutex::scoped_try_lock lock(pIndexer_->mutex_);
    //if(!lock.owns_lock())
        //return 0;
    izenelib::util::ScopedReadLock<izenelib::util::ReadWriteLock> lock(pIndexer_->mutex_);

    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return 0;

    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    TermReader* pTermReader = pBarrelReader_->termReader(colID);
    if (pTermReader)
    {
        return pTermReader->docFreq(term);
    }

    return 0;
}

TermInfo* IndexReader::termInfo(collectionid_t colID,Term* term)
{
    //boost::try_mutex::scoped_try_lock lock(pIndexer_->mutex_);
    //if(!lock.owns_lock())
        //return NULL;
    izenelib::util::ScopedReadLock<izenelib::util::ReadWriteLock> lock(pIndexer_->mutex_);

    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return NULL;

    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    TermReader* pTermReader = pBarrelReader_->termReader(colID);
    if (pTermReader)
    {
        return pTermReader->termInfo(term);
    }

    return 0;
}

size_t IndexReader::getDistinctNumTerms(collectionid_t colID, const std::string& property)
{
    //collection has been removed, need to rebuild the barrel reader
    if (pBarrelReader_ == NULL)
        createBarrelReader();
    if (pBarrelReader_ == NULL)
        return 0;

    boost::mutex::scoped_lock indexReaderLock(this->mutex_);

    return pBarrelReader_->getDistinctNumTerms(colID, property);
}

<|endoftext|>


<commit_before>/*
Copyright libCellML Contributors

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

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

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

#include "libcellml/model.h"

#include <algorithm>
#include <fstream>
#include <map>
#include <sstream>
#include <stack>
#include <utility>
#include <vector>

#include "libcellml/component.h"
#include "libcellml/importsource.h"
#include "libcellml/parser.h"
#include "libcellml/units.h"
#include "libcellml/variable.h"

#include "utilities.h"

namespace libcellml {

/**
 * @brief The Model::ModelImpl struct.
 *
 * This struct is the private implementation struct for the Model class.  Separating
 * the implementation from the definition allows for greater flexibility when
 * distributing the code.
 */
struct Model::ModelImpl
{
    std::vector<UnitsPtr> mUnits;

    std::vector<UnitsPtr>::iterator findUnits(const std::string &name);
    std::vector<UnitsPtr>::iterator findUnits(const UnitsPtr &units);
};

std::vector<UnitsPtr>::iterator Model::ModelImpl::findUnits(const std::string &name)
{
    return std::find_if(mUnits.begin(), mUnits.end(),
                        [=](const UnitsPtr &u) -> bool { return u->name() == name; });
}

std::vector<UnitsPtr>::iterator Model::ModelImpl::findUnits(const UnitsPtr &units)
{
    return std::find_if(mUnits.begin(), mUnits.end(),
                        [=](const UnitsPtr &u) -> bool { return units->name().empty() ? false : u->name() == units->name(); });
}

Model::Model()
    : mPimpl(new ModelImpl())
{
}

Model::Model(const std::string &name)
    : mPimpl(new ModelImpl())
{
    setName(name);
}

ModelPtr Model::create() noexcept
{
    return std::shared_ptr<Model> {new Model {}};
}

ModelPtr Model::create(const std::string &name) noexcept
{
    return std::shared_ptr<Model> {new Model {name}};
}

Model::~Model()
{
    delete mPimpl;
}

bool Model::doAddComponent(const ComponentPtr &component)
{
    if (component->hasParent()) {
        auto parent = component->parent();
        removeComponentFromEntity(parent, component);
    }
    component->setParent(shared_from_this());
    return ComponentEntity::doAddComponent(component);
}

void Model::addUnits(const UnitsPtr &units)
{
    mPimpl->mUnits.push_back(units);
    units->setParent(shared_from_this());
}

bool Model::removeUnits(size_t index)
{
    bool status = false;
    if (index < mPimpl->mUnits.size()) {
        auto units = *(mPimpl->mUnits.begin() + int64_t(index));
        units->removeParent();
        mPimpl->mUnits.erase(mPimpl->mUnits.begin() + int64_t(index));
        status = true;
    }

    return status;
}

bool Model::removeUnits(const std::string &name)
{
    bool status = false;
    auto result = mPimpl->findUnits(name);
    if (result != mPimpl->mUnits.end()) {
        (*result)->removeParent();
        mPimpl->mUnits.erase(result);
        status = true;
    }

    return status;
}

bool Model::removeUnits(const UnitsPtr &units)
{
    bool status = false;
    auto result = mPimpl->findUnits(units);
    if (result != mPimpl->mUnits.end()) {
        units->removeParent();
        mPimpl->mUnits.erase(result);
        status = true;
    }

    return status;
}

void Model::removeAllUnits()
{
    mPimpl->mUnits.clear();
}

bool Model::hasUnits(const std::string &name) const
{
    return mPimpl->findUnits(name) != mPimpl->mUnits.end();
}

bool Model::hasUnits(const UnitsPtr &units) const
{
    return mPimpl->findUnits(units) != mPimpl->mUnits.end();
}

UnitsPtr Model::units(size_t index) const
{
    UnitsPtr units = nullptr;
    if (index < mPimpl->mUnits.size()) {
        units = mPimpl->mUnits.at(index);
    }

    return units;
}

UnitsPtr Model::units(const std::string &name) const
{
    UnitsPtr units = nullptr;
    auto result = mPimpl->findUnits(name);
    if (result != mPimpl->mUnits.end()) {
        units = *result;
    }

    return units;
}

UnitsPtr Model::takeUnits(size_t index)
{
    UnitsPtr units = nullptr;
    if (index < mPimpl->mUnits.size()) {
        units = mPimpl->mUnits.at(index);
        removeUnits(index);
        units->removeParent();
    }

    return units;
}

UnitsPtr Model::takeUnits(const std::string &name)
{
    return takeUnits(size_t(mPimpl->findUnits(name) - mPimpl->mUnits.begin()));
}

bool Model::replaceUnits(size_t index, const UnitsPtr &units)
{
    bool status = false;
    if (removeUnits(index)) {
        mPimpl->mUnits.insert(mPimpl->mUnits.begin() + int64_t(index), units);
        status = true;
    }

    return status;
}

bool Model::replaceUnits(const std::string &name, const UnitsPtr &units)
{
    return replaceUnits(size_t(mPimpl->findUnits(name) - mPimpl->mUnits.begin()), units);
}

bool Model::replaceUnits(const UnitsPtr &oldUnits, const UnitsPtr &newUnits)
{
    return replaceUnits(size_t(mPimpl->findUnits(oldUnits) - mPimpl->mUnits.begin()), newUnits);
}

size_t Model::unitsCount() const
{
    return mPimpl->mUnits.size();
}

/**
 * @brief Resolve the path of the given filename using the given base.
 *
 * Resolves the full path to the given @p filename using the @p base.
 *
 * This function is only intended to work with local files.  It may not
 * work with bases that use the 'file://' prefix.
 *
 * @param filename The @c std::string relative path from the base path.
 * @param base The @c std::string location on local disk for determining the full path from.
 * @return The full path from the @p base location to the @p filename
 */
std::string resolvePath(const std::string &filename, const std::string &base)
{
    // We can be naive here as we know what we are dealing with
    std::string path = base.substr(0, base.find_last_of('/') + 1) + filename;
    return path;
}

void resolveImport(const ImportedEntityPtr &importedEntity,
                   const std::string &baseFile)
{
    if (importedEntity->isImport()) {
        ImportSourcePtr importSource = importedEntity->importSource();
        if (!importSource->hasModel()) {
            std::string url = resolvePath(importSource->url(), baseFile);
            std::ifstream file(url);
            if (file.good()) {
                std::stringstream buffer;
                buffer << file.rdbuf();
                ParserPtr parser = Parser::create();
                ModelPtr model = parser->parseModel(buffer.str());
                importSource->setModel(model);
                model->resolveImports(url);
            }
        }
    }
}

void resolveComponentImports(const ComponentEntityPtr &parentComponentEntity,
                             const std::string &baseFile)
{
    for (size_t n = 0; n < parentComponentEntity->componentCount(); ++n) {
        libcellml::ComponentPtr component = parentComponentEntity->component(n);
        if (component->isImport()) {
            resolveImport(component, baseFile);
        } else {
            resolveComponentImports(component, baseFile);
        }
    }
}

void Model::resolveImports(const std::string &baseFile)
{
    for (size_t n = 0; n < unitsCount(); ++n) {
        libcellml::UnitsPtr units = Model::units(n);
        resolveImport(units, baseFile);
    }
    resolveComponentImports(shared_from_this(), baseFile);
}

bool isUnresolvedImport(const ImportedEntityPtr &importedEntity)
{
    bool unresolvedImport = false;
    if (importedEntity->isImport()) {
        ImportSourcePtr importedSource = importedEntity->importSource();
        if (!importedSource->hasModel()) {
            unresolvedImport = true;
        }
    }
    return unresolvedImport;
}

bool hasUnresolvedComponentImports(const ComponentEntityPtr &parentComponentEntity);

bool doHasUnresolvedComponentImports(const ComponentPtr &component)
{
    bool unresolvedImports = false;
    if (component->isImport()) {
        unresolvedImports = isUnresolvedImport(component);
        if (!unresolvedImports) {
            // Check that the imported component can import all it needs from its model.
            ImportSourcePtr importedSource = component->importSource();
            if (importedSource->hasModel()) {
                ModelPtr importedModel = importedSource->model();
                ComponentPtr importedComponent = importedModel->component(component->importReference());
                unresolvedImports = doHasUnresolvedComponentImports(importedComponent);
            }
        }
    } else {
        unresolvedImports = hasUnresolvedComponentImports(component);
    }
    return unresolvedImports;
}

bool hasUnresolvedComponentImports(const ComponentEntityPtr &parentComponentEntity)
{
    bool unresolvedImports = false;
    for (size_t n = 0; n < parentComponentEntity->componentCount() && !unresolvedImports; ++n) {
        libcellml::ComponentPtr component = parentComponentEntity->component(n);
        unresolvedImports = doHasUnresolvedComponentImports(component);
    }
    return unresolvedImports;
}

bool Model::hasUnresolvedImports()
{
    bool unresolvedImports = false;
    for (size_t n = 0; n < unitsCount() && !unresolvedImports; ++n) {
        libcellml::UnitsPtr units = Model::units(n);
        unresolvedImports = isUnresolvedImport(units);
    }
    if (!unresolvedImports) {
        unresolvedImports = hasUnresolvedComponentImports(shared_from_this());
    }
    return unresolvedImports;
}

ModelPtr Model::clone() const
{
    auto m = create();
    return m;
}

} // namespace libcellml
<commit_msg>Implement copy of variable equivalences when cloning models.<commit_after>/*
Copyright libCellML Contributors

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

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

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

#include "libcellml/model.h"

#include <algorithm>
#include <fstream>
#include <map>
#include <sstream>
#include <stack>
#include <utility>
#include <vector>

#include "libcellml/component.h"
#include "libcellml/importsource.h"
#include "libcellml/parser.h"
#include "libcellml/units.h"
#include "libcellml/variable.h"

#include "utilities.h"

#include "debug.h"

namespace libcellml {

/**
 * @brief The Model::ModelImpl struct.
 *
 * This struct is the private implementation struct for the Model class.  Separating
 * the implementation from the definition allows for greater flexibility when
 * distributing the code.
 */
struct Model::ModelImpl
{
    std::vector<UnitsPtr> mUnits;

    std::vector<UnitsPtr>::iterator findUnits(const std::string &name);
    std::vector<UnitsPtr>::iterator findUnits(const UnitsPtr &units);
};

std::vector<UnitsPtr>::iterator Model::ModelImpl::findUnits(const std::string &name)
{
    return std::find_if(mUnits.begin(), mUnits.end(),
                        [=](const UnitsPtr &u) -> bool { return u->name() == name; });
}

std::vector<UnitsPtr>::iterator Model::ModelImpl::findUnits(const UnitsPtr &units)
{
    return std::find_if(mUnits.begin(), mUnits.end(),
                        [=](const UnitsPtr &u) -> bool { return units->name().empty() ? false : u->name() == units->name(); });
}

Model::Model()
    : mPimpl(new ModelImpl())
{
}

Model::Model(const std::string &name)
    : mPimpl(new ModelImpl())
{
    setName(name);
}

ModelPtr Model::create() noexcept
{
    return std::shared_ptr<Model> {new Model {}};
}

ModelPtr Model::create(const std::string &name) noexcept
{
    return std::shared_ptr<Model> {new Model {name}};
}

Model::~Model()
{
    delete mPimpl;
}

bool Model::doAddComponent(const ComponentPtr &component)
{
    if (component->hasParent()) {
        auto parent = component->parent();
        removeComponentFromEntity(parent, component);
    }
    component->setParent(shared_from_this());
    return ComponentEntity::doAddComponent(component);
}

void Model::addUnits(const UnitsPtr &units)
{
    mPimpl->mUnits.push_back(units);
    units->setParent(shared_from_this());
}

bool Model::removeUnits(size_t index)
{
    bool status = false;
    if (index < mPimpl->mUnits.size()) {
        auto units = *(mPimpl->mUnits.begin() + int64_t(index));
        units->removeParent();
        mPimpl->mUnits.erase(mPimpl->mUnits.begin() + int64_t(index));
        status = true;
    }

    return status;
}

bool Model::removeUnits(const std::string &name)
{
    bool status = false;
    auto result = mPimpl->findUnits(name);
    if (result != mPimpl->mUnits.end()) {
        (*result)->removeParent();
        mPimpl->mUnits.erase(result);
        status = true;
    }

    return status;
}

bool Model::removeUnits(const UnitsPtr &units)
{
    bool status = false;
    auto result = mPimpl->findUnits(units);
    if (result != mPimpl->mUnits.end()) {
        units->removeParent();
        mPimpl->mUnits.erase(result);
        status = true;
    }

    return status;
}

void Model::removeAllUnits()
{
    mPimpl->mUnits.clear();
}

bool Model::hasUnits(const std::string &name) const
{
    return mPimpl->findUnits(name) != mPimpl->mUnits.end();
}

bool Model::hasUnits(const UnitsPtr &units) const
{
    return mPimpl->findUnits(units) != mPimpl->mUnits.end();
}

UnitsPtr Model::units(size_t index) const
{
    UnitsPtr units = nullptr;
    if (index < mPimpl->mUnits.size()) {
        units = mPimpl->mUnits.at(index);
    }

    return units;
}

UnitsPtr Model::units(const std::string &name) const
{
    UnitsPtr units = nullptr;
    auto result = mPimpl->findUnits(name);
    if (result != mPimpl->mUnits.end()) {
        units = *result;
    }

    return units;
}

UnitsPtr Model::takeUnits(size_t index)
{
    UnitsPtr units = nullptr;
    if (index < mPimpl->mUnits.size()) {
        units = mPimpl->mUnits.at(index);
        removeUnits(index);
        units->removeParent();
    }

    return units;
}

UnitsPtr Model::takeUnits(const std::string &name)
{
    return takeUnits(size_t(mPimpl->findUnits(name) - mPimpl->mUnits.begin()));
}

bool Model::replaceUnits(size_t index, const UnitsPtr &units)
{
    bool status = false;
    if (removeUnits(index)) {
        mPimpl->mUnits.insert(mPimpl->mUnits.begin() + int64_t(index), units);
        status = true;
    }

    return status;
}

bool Model::replaceUnits(const std::string &name, const UnitsPtr &units)
{
    return replaceUnits(size_t(mPimpl->findUnits(name) - mPimpl->mUnits.begin()), units);
}

bool Model::replaceUnits(const UnitsPtr &oldUnits, const UnitsPtr &newUnits)
{
    return replaceUnits(size_t(mPimpl->findUnits(oldUnits) - mPimpl->mUnits.begin()), newUnits);
}

size_t Model::unitsCount() const
{
    return mPimpl->mUnits.size();
}

/**
 * @brief Resolve the path of the given filename using the given base.
 *
 * Resolves the full path to the given @p filename using the @p base.
 *
 * This function is only intended to work with local files.  It may not
 * work with bases that use the 'file://' prefix.
 *
 * @param filename The @c std::string relative path from the base path.
 * @param base The @c std::string location on local disk for determining the full path from.
 * @return The full path from the @p base location to the @p filename
 */
std::string resolvePath(const std::string &filename, const std::string &base)
{
    // We can be naive here as we know what we are dealing with
    std::string path = base.substr(0, base.find_last_of('/') + 1) + filename;
    return path;
}

void resolveImport(const ImportedEntityPtr &importedEntity,
                   const std::string &baseFile)
{
    if (importedEntity->isImport()) {
        ImportSourcePtr importSource = importedEntity->importSource();
        if (!importSource->hasModel()) {
            std::string url = resolvePath(importSource->url(), baseFile);
            std::ifstream file(url);
            if (file.good()) {
                std::stringstream buffer;
                buffer << file.rdbuf();
                ParserPtr parser = Parser::create();
                ModelPtr model = parser->parseModel(buffer.str());
                importSource->setModel(model);
                model->resolveImports(url);
            }
        }
    }
}

void resolveComponentImports(const ComponentEntityPtr &parentComponentEntity,
                             const std::string &baseFile)
{
    for (size_t n = 0; n < parentComponentEntity->componentCount(); ++n) {
        libcellml::ComponentPtr component = parentComponentEntity->component(n);
        if (component->isImport()) {
            resolveImport(component, baseFile);
        } else {
            resolveComponentImports(component, baseFile);
        }
    }
}

void Model::resolveImports(const std::string &baseFile)
{
    for (size_t n = 0; n < unitsCount(); ++n) {
        libcellml::UnitsPtr units = Model::units(n);
        resolveImport(units, baseFile);
    }
    resolveComponentImports(shared_from_this(), baseFile);
}

bool isUnresolvedImport(const ImportedEntityPtr &importedEntity)
{
    bool unresolvedImport = false;
    if (importedEntity->isImport()) {
        ImportSourcePtr importedSource = importedEntity->importSource();
        if (!importedSource->hasModel()) {
            unresolvedImport = true;
        }
    }
    return unresolvedImport;
}

bool hasUnresolvedComponentImports(const ComponentEntityPtr &parentComponentEntity);

bool doHasUnresolvedComponentImports(const ComponentPtr &component)
{
    bool unresolvedImports = false;
    if (component->isImport()) {
        unresolvedImports = isUnresolvedImport(component);
        if (!unresolvedImports) {
            // Check that the imported component can import all it needs from its model.
            ImportSourcePtr importedSource = component->importSource();
            if (importedSource->hasModel()) {
                ModelPtr importedModel = importedSource->model();
                ComponentPtr importedComponent = importedModel->component(component->importReference());
                unresolvedImports = doHasUnresolvedComponentImports(importedComponent);
            }
        }
    } else {
        unresolvedImports = hasUnresolvedComponentImports(component);
    }
    return unresolvedImports;
}

bool hasUnresolvedComponentImports(const ComponentEntityPtr &parentComponentEntity)
{
    bool unresolvedImports = false;
    for (size_t n = 0; n < parentComponentEntity->componentCount() && !unresolvedImports; ++n) {
        libcellml::ComponentPtr component = parentComponentEntity->component(n);
        unresolvedImports = doHasUnresolvedComponentImports(component);
    }
    return unresolvedImports;
}

bool Model::hasUnresolvedImports()
{
    bool unresolvedImports = false;
    for (size_t n = 0; n < unitsCount() && !unresolvedImports; ++n) {
        libcellml::UnitsPtr units = Model::units(n);
        unresolvedImports = isUnresolvedImport(units);
    }
    if (!unresolvedImports) {
        unresolvedImports = hasUnresolvedComponentImports(shared_from_this());
    }
    return unresolvedImports;
}

using IndexStack = std::vector< size_t >;
using EquivalenceMap = std::map< IndexStack, std::vector< IndexStack > >; /**< Type definition for map of variable equivalences defined over model. */

size_t getVariableIndexInComponent(const ComponentPtr &component, const VariablePtr &variable)
{
    size_t index = 0;
    bool found = false;
    while (index < component->variableCount() && !found) {
        if (component->variable(index) == variable) {
            found = true;
        } else {
            ++index;
        }
    }

    if (!found) {
        Debug() << "Not found variable in parent Component!!!!";
    }
    return index;
}

size_t getComponentIndexInComponentEntity(const ComponentEntityPtr &componentParent, const ComponentEntityPtr &component)
{
    size_t index = 0;
    bool found = false;
    while (index < componentParent->componentCount() && !found) {
        if (componentParent->component(index) == component) {
            found = true;
        } else {
            ++index;
        }
    }

    if (!found) {
        Debug() << "Not found component in parent Entity!!!!";
    }
    return index;
}

IndexStack reverseEngineerIndexStack(const VariablePtr &variable)
{
    IndexStack indexStack;
    ComponentPtr component = std::dynamic_pointer_cast<Component>(variable->parent());
    indexStack.push_back(getVariableIndexInComponent(component, variable));

    ComponentEntityPtr parent = component;
    ComponentEntityPtr grandParent = std::dynamic_pointer_cast<ComponentEntity>(parent->parent());
    while (grandParent != nullptr) {
        indexStack.push_back(getComponentIndexInComponentEntity(grandParent, parent));
        parent = grandParent;
        grandParent = std::dynamic_pointer_cast<ComponentEntity>(grandParent->parent());
    }

    std::reverse(std::begin(indexStack), std::end(indexStack));

    return indexStack;
}

void recordVariableEquivalences(const ComponentPtr &component, EquivalenceMap &equivalenceMap, IndexStack &indexStack)
{
    for (size_t index = 0; index < component->variableCount(); ++index) {
        auto variable = component->variable(index);
        for (size_t j = 0; j < variable->equivalentVariableCount(); ++j) {
            if (j == 0) {
                indexStack.push_back(index);
            }
            auto equivalentVariable = variable->equivalentVariable(j);
            auto equivalentVariableIndexStack = reverseEngineerIndexStack(equivalentVariable);
            if (equivalenceMap.count(indexStack) == 0) {
                equivalenceMap[indexStack] = std::vector< IndexStack >();
            }
            equivalenceMap[indexStack].push_back(equivalentVariableIndexStack);
        }
        if (variable->equivalentVariableCount()) {
            indexStack.pop_back();
        }
    }
}

void generateEquivalenceMap(const ComponentPtr &component, EquivalenceMap &map, IndexStack &indexStack)
{
    for (size_t index = 0; index < component->componentCount(); ++index) {
        indexStack.push_back(index);
        auto c = component->component(index);
        recordVariableEquivalences(c, map, indexStack);
        generateEquivalenceMap(c, map, indexStack);
        indexStack.pop_back();
    }
}

EquivalenceMap generateEquivalenceMap(std::shared_ptr<const libcellml::Model> model)
{
    EquivalenceMap map;

    IndexStack indexStack;
    for (size_t index = 0; index < model->componentCount(); ++index) {
        indexStack.push_back(index);
        auto c = model->component(index);
        recordVariableEquivalences(c, map, indexStack);
        generateEquivalenceMap(c, map, indexStack);
        indexStack.pop_back();
    }

    return map;
}

VariablePtr getVariableLocatedAt(const IndexStack &stack, ModelPtr model)
{
    ComponentPtr component;
    for (size_t index = 0; index < stack.size() - 1; ++index) {
        if (index == 0) {
            component = model->component(stack.at(index));
        } else {
            component = component->component(stack.at(index));
        }
    }

    return component->variable(stack.back());
}

void makeEquivalence(const IndexStack &stack1, const IndexStack &stack2, ModelPtr model)
{
    auto v1 = getVariableLocatedAt(stack1, model);
    auto v2 = getVariableLocatedAt(stack2, model);
    Variable::addEquivalence(v1, v2);
}

void applyEquivalenceMapToModel(const EquivalenceMap &map, ModelPtr model)
{
    for (EquivalenceMap::const_iterator iter = map.begin(); iter != map.end(); ++iter) {
        auto key = iter->first;
        auto vector = iter->second;
        for (auto vectorIter = vector.begin(); vectorIter < vector.end(); ++vectorIter) {
            makeEquivalence(key, *vectorIter, model);
        }
    }
}

void printStack(const IndexStack &stack)
{
    Debug(false) << "[";
    for (auto iter = stack.begin(); iter < stack.end(); ++iter) {
        Debug(false) << *iter;
        if (iter + 1 < stack.end()) {
            Debug(false) << ", ";
        }
    }
    Debug() << "]";
}

void printEquivalenceMap(const EquivalenceMap &map)
{
    Debug() << "Print out of equivalence map";
    for (EquivalenceMap::const_iterator iter = map.begin(); iter != map.end(); ++iter) {
        auto key = iter->first;
        Debug(false) << "key: ";
        printStack(key);
        auto vector = iter->second;
        for (auto vectorIt = vector.begin(); vectorIt < vector.end(); ++vectorIt) {
            Debug(false) << "value: ";
            printStack(*vectorIt);
        }
    }
}

ModelPtr Model::clone() const
{
    auto m = create();

    m->setId(id());
    m->setName(name());

    m->setEncapsulationId(encapsulationId());

    for (size_t index = 0; index < mPimpl->mUnits.size(); ++index) {
        m->addUnits(units(index)->clone());
    }

    for (size_t index = 0; index < componentCount(); ++index) {
        m->addComponent(component(index)->clone());
    }

    auto map = generateEquivalenceMap(shared_from_this());
    applyEquivalenceMapToModel(map, m);
//    printEquivalenceMap(map);

    return m;
}

} // namespace libcellml
<|endoftext|>


<commit_before>//=======================================================================
// Copyright (c) 2017 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include <cuComplex.h>

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in single-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_sapxdbpy(size_t n, float alpha, const float* x, size_t incx, float beta, float* y, size_t incy);

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in double-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_dapxdbpy(size_t n, double alpha, const double* x, size_t incx, double beta, double* y, size_t incy);

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in complex single-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_capxdbpy(size_t n, cuComplex alpha, const cuComplex* x, size_t incx, cuComplex beta, cuComplex* y, size_t incy);

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in complex double-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_zapxdbpy(size_t n, cuDoubleComplex alpha, const cuDoubleComplex* x, size_t incx, cuDoubleComplex beta, cuDoubleComplex* y, size_t incy);

#define EGBLAS_HAS_SAPXDBYP true
#define EGBLAS_HAS_DAPXDBYP true
#define EGBLAS_HAS_CAPXDBYP true
#define EGBLAS_HAS_ZAPXDBPY true
<commit_msg>Fix the macros<commit_after>//=======================================================================
// Copyright (c) 2017 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include <cuComplex.h>

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in single-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_sapxdbpy(size_t n, float alpha, const float* x, size_t incx, float beta, float* y, size_t incy);

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in double-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_dapxdbpy(size_t n, double alpha, const double* x, size_t incx, double beta, double* y, size_t incy);

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in complex single-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_capxdbpy(size_t n, cuComplex alpha, const cuComplex* x, size_t incx, cuComplex beta, cuComplex* y, size_t incy);

/*!
 * \brief Compute y = (alpha + x) / (beta + y) (element wise), in complex double-precision
 * \param n The size of the two vectors
 * \param alpha The multiplicator
 * \param x The vector x (GPU memory)
 * \param incx The stride of x
 * \param beta The addition
 * \param y The vector y (GPU memory)
 * \param incy The stride of y
 */
void egblas_zapxdbpy(size_t n, cuDoubleComplex alpha, const cuDoubleComplex* x, size_t incx, cuDoubleComplex beta, cuDoubleComplex* y, size_t incy);

#define EGBLAS_HAS_SAPXDBPY true
#define EGBLAS_HAS_DAPXDBPY true
#define EGBLAS_HAS_CAPXDBPY true
#define EGBLAS_HAS_ZAPXDBPY true
<|endoftext|>


<commit_before>#include <iostream>
#include <iomanip>
#include <limits>
#include <cmath>

const int MAX_VERTICES = 1000;

typedef uint32_t Graph[MAX_VERTICES][MAX_VERTICES];

// Read graph from input stream and return number of vertices.
int readGraph(std::istream& in, Graph G) {
    // Read number of vertices.
    int N; in >> N;

    // Read vertex coordinates.
    double x[MAX_VERTICES];
    double y[MAX_VERTICES];
    for (int i = 0; i < N; ++i) {
        in >> x[i] >> y[i];
    }

    // Calculate distances.
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            G[i][j] = std::round(std::sqrt(std::pow(x[i] - x[j], 2) + pow(y[i] - y[j], 2)));
        }
    }

    return N;
}

// Calculate and print MST using Prim's algorithm.
void primMST(Graph G, int N) {
    bool inMST[MAX_VERTICES];
    uint32_t cost[MAX_VERTICES];
    int parent[MAX_VERTICES];

    std::fill(inMST, inMST + N, false);
    std::fill(parent, parent + N, -1);
    std::fill(cost, cost + N, std::numeric_limits<uint32_t>::max());
    cost[0] = 0;

    for (int added = 0; added < N - 1; ++added) {
        // Find lowest cost vertex not in MST and add it to MST.
        int u = -1;
        uint32_t c = std::numeric_limits<int>::max();
        for (int v = 0; v < N; ++v) {
            if (cost[v] < c && !inMST[v]) {
                u = v;
                c = cost[v];
            }
        }
        inMST[u] = true;

        // For each neighbor v of u not already in MST.
        for (int v = 0; v < N; ++v) {
            if (0 < G[u][v] && G[u][v] < cost[v] && !inMST[v]) {
                // Distance from u to v is lower than current cost of v, so mark u
                // as parent of v and set cost of v to the distance from u to v.
                parent[v] = u;
                cost[v] = G[u][v];
            }
        }
    }

    // Print MST.
    for (int v = 1; v < N; ++v) {
        std::cout << v << " " << parent[v] << std::endl;
    }
}

int main(void) {

    // Read graph.
    Graph G;
    int N = readGraph(std::cin, G);

    // Calculate and print MST.
    primMST(G, N);

    return 0;
}
<commit_msg>Add a small test case.<commit_after>#include <iostream>
#include <iomanip>
#include <limits>
#include <cmath>
#include <cassert>

const int MAX_VERTICES = 1000;
const int NO_PARENT = -1;

typedef uint32_t Graph[MAX_VERTICES][MAX_VERTICES];

// Read graph from input stream and return number of vertices.
int readGraph(std::istream& in, Graph G) {
    // Read number of vertices.
    int N; in >> N;

    // Read vertex coordinates.
    double x[MAX_VERTICES];
    double y[MAX_VERTICES];
    for (int i = 0; i < N; ++i) {
        in >> x[i] >> y[i];
    }

    // Calculate distances.
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            G[i][j] = std::round(std::sqrt(std::pow(x[i] - x[j], 2) + pow(y[i] - y[j], 2)));
        }
    }

    return N;
}

// Calculate MST rooted at 0 using Prim's algorithm.
void primMST(Graph G, int parent[MAX_VERTICES], int N) {
    bool inMST[MAX_VERTICES];
    uint32_t cost[MAX_VERTICES];

    std::fill(inMST, inMST + N, false);
    std::fill(parent, parent + N, NO_PARENT);
    std::fill(cost, cost + N, std::numeric_limits<uint32_t>::max());
    cost[0] = 0;

    for (int added = 0; added < N - 1; ++added) {
        // Find lowest cost vertex not in MST and add it to MST.
        int u = -1;
        uint32_t c = std::numeric_limits<int>::max();
        for (int v = 0; v < N; ++v) {
            if (cost[v] < c && !inMST[v]) {
                u = v;
                c = cost[v];
            }
        }
        inMST[u] = true;

        // For each neighbor v of u not already in MST.
        for (int v = 0; v < N; ++v) {
            if (0 < G[u][v] && G[u][v] < cost[v] && !inMST[v]) {
                // Distance from u to v is lower than current cost of v, so mark u
                // as parent of v and set cost of v to the distance from u to v.
                parent[v] = u;
                cost[v] = G[u][v];
            }
        }
    }
}

// Test function for the primMST function.
void testPrimMST() {
    // Run the function on a small test case.
    int N = 5;
    Graph G = {{0, 2, 0, 6, 0},
               {2, 0, 3, 8, 5},
               {0, 3, 0, 0, 7},
               {6, 8, 0, 0, 9},
               {0, 5, 7, 9, 0}};
    int parent[MAX_VERTICES];
    std::fill(parent, parent + N, 42); // Just in case.

    primMST(G, parent, N);

    // Check results.
    assert(parent[0] == NO_PARENT);
    assert(parent[1] == 0);
    assert(parent[2] == 1);
    assert(parent[3] == 0);
    assert(parent[4] == 1);
}

int main(void) {
    // Run a small test case first.
    testPrimMST();

    // Read graph from standard input.
    Graph G;
    int N = readGraph(std::cin, G);

    // Calculate and print MST.
    int parent[MAX_VERTICES];
    primMST(G, parent, N);

    // Print MST.
    for (int v = 1; v < N; ++v) {
        std::cout << v << " " << parent[v] << std::endl;
    }

    return 0;
}
<|endoftext|>


<commit_before>

#include <iostream>
#include <string.h>
#include <stdlib.h>
#include <errno.h>

#include <openssl/md5.h>
#include <openssl/sha.h>

#include "include/types.h"
#include "objclass/objclass.h"

#include "include/rbd_types.h"

CLS_VER(1,0)
CLS_NAME(rbd)

cls_handle_t h_class;
cls_method_handle_t h_snapshots_list;
cls_method_handle_t h_snapshot_add;
cls_method_handle_t h_snapshot_revert;

static int snap_read_header(cls_method_context_t hctx, bufferlist& bl)
{
  int snap_count = 0;
  uint64_t snap_names_len = 0;
  int rc;
  struct rbd_obj_header_ondisk *header;

  cls_log("snapshots_list");

  while (1) {
    int len = sizeof(*header) +
      snap_count * sizeof(struct rbd_obj_snap_ondisk) +
      snap_names_len;

    rc = cls_cxx_read(hctx, 0, len, &bl);
    if (rc < 0)
      return rc;

    header = (struct rbd_obj_header_ondisk *)bl.c_str();

    if ((snap_count != header->snap_count) ||
        (snap_names_len != header->snap_names_len)) {
      snap_count = header->snap_count;
      snap_names_len = header->snap_names_len;
      bl.clear();
      continue;
    }
    break;
  }

  return 0;
}

int snapshots_list(cls_method_context_t hctx, bufferlist *in, bufferlist *out)
{
  bufferlist bl;
  struct rbd_obj_header_ondisk *header;
  int rc = snap_read_header(hctx, bl);
  if (rc < 0)
    return rc;

  header = (struct rbd_obj_header_ondisk *)bl.c_str();
  bufferptr p(header->snap_names_len);
  char *buf = (char *)header;
  char *name = buf + sizeof(*header) + header->snap_count * sizeof(struct rbd_obj_snap_ondisk);
  char *end = name + header->snap_names_len;
  memcpy(p.c_str(),
         buf + sizeof(*header) + header->snap_count * sizeof(struct rbd_obj_snap_ondisk),
         header->snap_names_len);

  ::encode(header->snap_count, *out);

  for (int i = 0; i < header->snap_count; i++) {
    string s = name;
    ::encode(header->snaps[i].id, *out);
    ::encode(header->snaps[i].image_size, *out);
    ::encode(s, *out);

    name += strlen(name) + 1;
    if (name > end)
      return -EIO;
  }

  return 0;
}

int snapshot_add(cls_method_context_t hctx, bufferlist *in, bufferlist *out)
{
  bufferlist bl;
  struct rbd_obj_header_ondisk *header;
  bufferlist newbl;
  bufferptr header_bp(sizeof(*header));
  struct rbd_obj_snap_ondisk *new_snaps;

  int rc = snap_read_header(hctx, bl);
  if (rc < 0)
    return rc;

  header = (struct rbd_obj_header_ondisk *)bl.c_str();

  int snaps_id_ofs = sizeof(*header);
  int len = snaps_id_ofs;
  int names_ofs = snaps_id_ofs + sizeof(*new_snaps) * header->snap_count;
  const char *snap_name;
  const char *snap_names = ((char *)header) + names_ofs;
  const char *end = snap_names + header->snap_names_len;
  bufferlist::iterator iter = in->begin();
  string s;
  uint64_t snap_id;

  try {
    ::decode(s, iter);
    ::decode(snap_id, iter);
  } catch (buffer::error *err) {
    return -EINVAL;
  }
  snap_name = s.c_str();

  const char *cur_snap_name;
  for (cur_snap_name = snap_names; cur_snap_name < end; cur_snap_name += strlen(cur_snap_name) + 1) {
    if (strncmp(cur_snap_name, snap_name, end - cur_snap_name) == 0)
      return -EEXIST;
  }
  if (cur_snap_name > end)
    return -EIO;

  int snap_name_len = strlen(snap_name);

  bufferptr new_names_bp(header->snap_names_len + snap_name_len + 1);
  bufferptr new_snaps_bp(sizeof(*new_snaps) * (header->snap_count + 1));

  /* copy snap names and append to new snap name */
  char *new_snap_names = new_names_bp.c_str();
  strcpy(new_snap_names, snap_name);
  memcpy(new_snap_names + snap_name_len + 1, snap_names, header->snap_names_len);

  /* append new snap id */
  new_snaps = (struct rbd_obj_snap_ondisk *)new_snaps_bp.c_str();
  memcpy(new_snaps + 1, header->snaps, sizeof(*new_snaps) * header->snap_count);

  header->snap_count = header->snap_count + 1;
  header->snap_names_len = header->snap_names_len + snap_name_len + 1;
  header->snap_seq = header->snap_seq + 1;

  new_snaps[0].id = snap_id;
  new_snaps[0].image_size = header->image_size;

  len += sizeof(*new_snaps) * header->snap_count + header->snap_names_len;

  memcpy(header_bp.c_str(), header, sizeof(*header));

  newbl.push_back(header_bp);
  newbl.push_back(new_snaps_bp);
  newbl.push_back(new_names_bp);

  rc = cls_cxx_write(hctx, 0, len, &newbl);
  if (rc < 0)
    return rc;

  return 0;
}

int snapshot_revert(cls_method_context_t hctx, bufferlist *in, bufferlist *out)
{
  bufferlist bl;
  struct rbd_obj_header_ondisk *header;
  bufferlist newbl;
  bufferptr header_bp(sizeof(*header));
  struct rbd_obj_snap_ondisk *new_snaps;

  int rc = snap_read_header(hctx, bl);
  if (rc < 0)
    return rc;

  header = (struct rbd_obj_header_ondisk *)bl.c_str();

  int snaps_id_ofs = sizeof(*header);
  int len = snaps_id_ofs;
  int names_ofs = snaps_id_ofs + sizeof(*new_snaps) * header->snap_count;
  const char *snap_name;
  const char *snap_names = ((char *)header) + names_ofs;
  const char *end = snap_names + header->snap_names_len;
  bufferlist::iterator iter = in->begin();
  string s;
  uint64_t snap_id;
  int i;
  bool found = false;
  struct rbd_obj_snap_ondisk snap;

  try {
    ::decode(s, iter);
  } catch (buffer::error *err) {
    return -EINVAL;
  }
  snap_name = s.c_str();

  for (i = 0; snap_names < end; i++) {
    if (strcmp(snap_names, snap_name) == 0) {
      snap = header->snaps[i];
      found = true;
      break;
    }
    snap_names += strlen(snap_names);
  }
  if (!found)
    return -ENOENT;

  header->image_size = snap.image_size;
  header->snap_seq = header->snap_seq + 1;

  snap_names += strlen(snap_names);
  i++;

  header->snap_count = header->snap_count - i;
  bufferptr new_names_bp(end - snap_names);
  bufferptr new_snaps_bp(sizeof(header->snaps[0]) * header->snap_count);

  memcpy(header_bp.c_str(), header, sizeof(*header));
  if (header->snap_count) {
    char *new_snap_names;
    memcpy(new_snaps_bp.c_str(), header->snaps + i, sizeof(header->snaps[0]) * header->snap_count);
    memcpy(new_names_bp.c_str(), snap_names, end - snap_names);
    newbl.push_back(new_snaps_bp);
    newbl.push_back(new_names_bp);
  }

  rc = cls_cxx_write(hctx, 0, len, &newbl);
  if (rc < 0)
    return rc;

  ::encode(snap.id, *out);

  return 0;
}

void class_init()
{
  CLS_LOG("Loaded rbd class!");

  cls_register("rbd", &h_class);
  cls_register_cxx_method(h_class, "snap_list", CLS_METHOD_RD, snapshots_list, &h_snapshots_list);
  cls_register_cxx_method(h_class, "snap_add", CLS_METHOD_RD | CLS_METHOD_WR, snapshot_add, &h_snapshot_add);
  cls_register_cxx_method(h_class, "snap_revert", CLS_METHOD_RD | CLS_METHOD_WR, snapshot_revert, &h_snapshot_revert);

  return;
}

<commit_msg>rbd: snap revert header manipulation fixes<commit_after>

#include <iostream>
#include <string.h>
#include <stdlib.h>
#include <errno.h>

#include <openssl/md5.h>
#include <openssl/sha.h>

#include "include/types.h"
#include "objclass/objclass.h"

#include "include/rbd_types.h"

CLS_VER(1,0)
CLS_NAME(rbd)

cls_handle_t h_class;
cls_method_handle_t h_snapshots_list;
cls_method_handle_t h_snapshot_add;
cls_method_handle_t h_snapshot_revert;

static int snap_read_header(cls_method_context_t hctx, bufferlist& bl)
{
  int snap_count = 0;
  uint64_t snap_names_len = 0;
  int rc;
  struct rbd_obj_header_ondisk *header;

  cls_log("snapshots_list");

  while (1) {
    int len = sizeof(*header) +
      snap_count * sizeof(struct rbd_obj_snap_ondisk) +
      snap_names_len;

    rc = cls_cxx_read(hctx, 0, len, &bl);
    if (rc < 0)
      return rc;

    header = (struct rbd_obj_header_ondisk *)bl.c_str();

    if ((snap_count != header->snap_count) ||
        (snap_names_len != header->snap_names_len)) {
      snap_count = header->snap_count;
      snap_names_len = header->snap_names_len;
      bl.clear();
      continue;
    }
    break;
  }

  return 0;
}

int snapshots_list(cls_method_context_t hctx, bufferlist *in, bufferlist *out)
{
  bufferlist bl;
  struct rbd_obj_header_ondisk *header;
  int rc = snap_read_header(hctx, bl);
  if (rc < 0)
    return rc;

  header = (struct rbd_obj_header_ondisk *)bl.c_str();
  bufferptr p(header->snap_names_len);
  char *buf = (char *)header;
  char *name = buf + sizeof(*header) + header->snap_count * sizeof(struct rbd_obj_snap_ondisk);
  char *end = name + header->snap_names_len;
  memcpy(p.c_str(),
         buf + sizeof(*header) + header->snap_count * sizeof(struct rbd_obj_snap_ondisk),
         header->snap_names_len);

  ::encode(header->snap_count, *out);

  for (int i = 0; i < header->snap_count; i++) {
    string s = name;
    ::encode(header->snaps[i].id, *out);
    ::encode(header->snaps[i].image_size, *out);
    ::encode(s, *out);

    name += strlen(name) + 1;
    if (name > end)
      return -EIO;
  }

  return 0;
}

int snapshot_add(cls_method_context_t hctx, bufferlist *in, bufferlist *out)
{
  bufferlist bl;
  struct rbd_obj_header_ondisk *header;
  bufferlist newbl;
  bufferptr header_bp(sizeof(*header));
  struct rbd_obj_snap_ondisk *new_snaps;

  int rc = snap_read_header(hctx, bl);
  if (rc < 0)
    return rc;

  header = (struct rbd_obj_header_ondisk *)bl.c_str();

  int snaps_id_ofs = sizeof(*header);
  int len = snaps_id_ofs;
  int names_ofs = snaps_id_ofs + sizeof(*new_snaps) * header->snap_count;
  const char *snap_name;
  const char *snap_names = ((char *)header) + names_ofs;
  const char *end = snap_names + header->snap_names_len;
  bufferlist::iterator iter = in->begin();
  string s;
  uint64_t snap_id;

  try {
    ::decode(s, iter);
    ::decode(snap_id, iter);
  } catch (buffer::error *err) {
    return -EINVAL;
  }
  snap_name = s.c_str();

  const char *cur_snap_name;
  for (cur_snap_name = snap_names; cur_snap_name < end; cur_snap_name += strlen(cur_snap_name) + 1) {
    if (strncmp(cur_snap_name, snap_name, end - cur_snap_name) == 0)
      return -EEXIST;
  }
  if (cur_snap_name > end)
    return -EIO;

  int snap_name_len = strlen(snap_name);

  bufferptr new_names_bp(header->snap_names_len + snap_name_len + 1);
  bufferptr new_snaps_bp(sizeof(*new_snaps) * (header->snap_count + 1));

  /* copy snap names and append to new snap name */
  char *new_snap_names = new_names_bp.c_str();
  strcpy(new_snap_names, snap_name);
  memcpy(new_snap_names + snap_name_len + 1, snap_names, header->snap_names_len);

  /* append new snap id */
  new_snaps = (struct rbd_obj_snap_ondisk *)new_snaps_bp.c_str();
  memcpy(new_snaps + 1, header->snaps, sizeof(*new_snaps) * header->snap_count);

  header->snap_count = header->snap_count + 1;
  header->snap_names_len = header->snap_names_len + snap_name_len + 1;
  header->snap_seq = header->snap_seq + 1;

  new_snaps[0].id = snap_id;
  new_snaps[0].image_size = header->image_size;

  len += sizeof(*new_snaps) * header->snap_count + header->snap_names_len;

  memcpy(header_bp.c_str(), header, sizeof(*header));

  newbl.push_back(header_bp);
  newbl.push_back(new_snaps_bp);
  newbl.push_back(new_names_bp);

  rc = cls_cxx_write(hctx, 0, len, &newbl);
  if (rc < 0)
    return rc;

  return 0;
}

int snapshot_revert(cls_method_context_t hctx, bufferlist *in, bufferlist *out)
{
  bufferlist bl;
  struct rbd_obj_header_ondisk *header;
  bufferlist newbl;
  bufferptr header_bp(sizeof(*header));
  struct rbd_obj_snap_ondisk *new_snaps;

  int rc = snap_read_header(hctx, bl);
  if (rc < 0)
    return rc;

  header = (struct rbd_obj_header_ondisk *)bl.c_str();

  int snaps_id_ofs = sizeof(*header);
  int len = snaps_id_ofs;
  int names_ofs = snaps_id_ofs + sizeof(*new_snaps) * header->snap_count;
  const char *snap_name;
  const char *snap_names = ((char *)header) + names_ofs;
  const char *end = snap_names + header->snap_names_len;
  bufferlist::iterator iter = in->begin();
  string s;
  uint64_t snap_id;
  int i;
  bool found = false;
  struct rbd_obj_snap_ondisk snap;

  try {
    ::decode(s, iter);
  } catch (buffer::error *err) {
    return -EINVAL;
  }
  snap_name = s.c_str();

  for (i = 0; snap_names < end; i++) {
    if (strcmp(snap_names, snap_name) == 0) {
      snap = header->snaps[i];
      found = true;
      break;
    }
    snap_names += strlen(snap_names) + 1;
  }
  if (!found)
    return -ENOENT;

  header->image_size = snap.image_size;
  header->snap_seq = header->snap_seq + 1;

  snap_names += strlen(snap_names) + 1;
  i++;

  header->snap_count = header->snap_count - i;
  CLS_LOG("snap_count=%d\n", header->snap_count);
  bufferptr new_names_bp(end - snap_names);
  bufferptr new_snaps_bp(sizeof(header->snaps[0]) * header->snap_count);

  memcpy(header_bp.c_str(), header, sizeof(*header));
  newbl.push_back(header_bp);

  if (header->snap_count) {
    char *new_snap_names;
    memcpy(new_snaps_bp.c_str(), header->snaps + i, sizeof(header->snaps[0]) * header->snap_count);
    memcpy(new_names_bp.c_str(), snap_names, end - snap_names);
    newbl.push_back(new_snaps_bp);
    newbl.push_back(new_names_bp);
  }

  rc = cls_cxx_write(hctx, 0, newbl.length(), &newbl);
  if (rc < 0)
    return rc;

  ::encode(snap.id, *out);

  return 0;
}

void class_init()
{
  CLS_LOG("Loaded rbd class!");

  cls_register("rbd", &h_class);
  cls_register_cxx_method(h_class, "snap_list", CLS_METHOD_RD, snapshots_list, &h_snapshots_list);
  cls_register_cxx_method(h_class, "snap_add", CLS_METHOD_RD | CLS_METHOD_WR, snapshot_add, &h_snapshot_add);
  cls_register_cxx_method(h_class, "snap_revert", CLS_METHOD_RD | CLS_METHOD_WR, snapshot_revert, &h_snapshot_revert);

  return;
}

<|endoftext|>


<commit_before>//========= Copyright  1996-2005, Valve Corporation, All rights reserved. ============//
//
// Purpose:
//
// $NoKeywords: $
//=============================================================================//

#include "cbase.h"
#include "hud.h"
#include "sdk_hud_stamina.h"
#include "hud_macros.h"
#include "c_sdk_player.h"
#include "iclientmode.h"
#include <vgui_controls/AnimationController.h>
#include <vgui/ISurface.h>
#include <vgui/ILocalize.h>
#include <igameresources.h>
#include "c_baseplayer.h"
#include "in_buttons.h"
#include "sdk_gamerules.h"
#include "c_ios_replaymanager.h"
#include "c_match_ball.h"
#include "view.h"
#include "ios_camera.h"
#include "iinput.h"

// memdbgon must be the last include file in a .cpp file!!!
#include "tier0/memdbgon.h"

extern ConVar centeredstaminabar;

class CHudChargedshotBar : public CHudElement, public vgui::Panel
{
	DECLARE_CLASS_SIMPLE( CHudChargedshotBar, vgui::Panel );

public:
	CHudChargedshotBar( const char *pElementName );
	virtual void	Init( void );
	virtual void	Reset( void );
	virtual void	OnThink( void );
	bool			ShouldDraw( void );

protected:
	virtual void	ApplySchemeSettings( vgui::IScheme *scheme );
	virtual void	Paint();

	float m_flNextUpdate;
};

DECLARE_HUDELEMENT( CHudChargedshotBar );

using namespace vgui;

enum { BAR_WIDTH = 80, BAR_HEIGHT = 14, BAR_HPADDING = 1, BAR_VPADDING = 1 };
enum { INDICATOR_WIDTH = 9, INDICATOR_OFFSET = 7, INDICATOR_BORDER = 1 };

//-----------------------------------------------------------------------------
// Purpose: Constructor
//-----------------------------------------------------------------------------
CHudChargedshotBar::CHudChargedshotBar( const char *pElementName ) : CHudElement( pElementName ), BaseClass( NULL, "HudChargedshotBar" )
{
	SetHiddenBits(HIDEHUD_POWERSHOTBAR);

	vgui::Panel *pParent = g_pClientMode->GetViewport();
	SetParent( pParent );

	m_flNextUpdate = gpGlobals->curtime;
}

void CHudChargedshotBar::ApplySchemeSettings( IScheme *scheme )
{
	BaseClass::ApplySchemeSettings(scheme);

	SetBounds(0, 0, ScreenWidth(), ScreenHeight());
}

//-----------------------------------------------------------------------------
// Purpose:
//-----------------------------------------------------------------------------
void CHudChargedshotBar::Init( void )
{
}

//-----------------------------------------------------------------------------
// Purpose:
//-----------------------------------------------------------------------------
void CHudChargedshotBar::Reset( void )
{
	Init();
}

//-----------------------------------------------------------------------------
// Purpose: Save CPU cycles by letting the HUD system early cull
// costly traversal.  Called per frame, return true if thinking and
// painting need to occur.
//-----------------------------------------------------------------------------
bool CHudChargedshotBar::ShouldDraw()
{
	if (!CHudElement::ShouldDraw())
		return false;

	C_SDKPlayer *pPlayer = C_SDKPlayer::GetLocalSDKPlayer();
	if (!pPlayer || pPlayer->GetTeamNumber() != TEAM_A && pPlayer->GetTeamNumber() != TEAM_B)
		return false;

	if (GetMatchBall() && GetMatchBall()->m_eBallState == BALL_STATE_GOAL)
		return false;

	if (GetReplayManager() && GetReplayManager()->IsReplaying())
		return false;

	return true;
}

//-----------------------------------------------------------------------------
// Purpose:
//-----------------------------------------------------------------------------
void CHudChargedshotBar::OnThink( void )
{
	BaseClass::OnThink();
}

//-----------------------------------------------------------------------------
// Purpose: draws the stamina bar
//-----------------------------------------------------------------------------
void CHudChargedshotBar::Paint()
{
	C_SDKPlayer *pPlayer = C_SDKPlayer::GetLocalSDKPlayer();

	if ( !pPlayer )
		return;

	float stamina = pPlayer->m_Shared.GetStamina();
	float relStamina = stamina / 100.0f;

	Color fgColor, bgColor;

	Color missingMaxStaminaColor = Color(100, 100, 100, 255);

	if (pPlayer->GetFlags() & FL_REMOTECONTROLLED)
	{
		fgColor = Color(255, 255, 255, 255);
		bgColor = Color(100, 100, 100, 255);
	}
	else if (GetMatchBall() && GetMatchBall()->m_bShotsBlocked)
	{
		fgColor = Color(255, 0, 0, 255);
		bgColor = Color(139, 0, 0, 255);
	}
	else if (GetMatchBall() && GetMatchBall()->m_bChargedshotBlocked)
	{
		fgColor = Color(255, 140, 0, 255);
		bgColor = Color(255, 69, 0, 255);
	}
	else
	{
		fgColor = Color(255 * (1 - relStamina), 255 * relStamina, 0, 255);
		bgColor = Color(0, 0, 0, 255);
	}

	float shotStrength;

	bool drawChargedshotIndicator = false;

	if (pPlayer->m_Shared.m_bDoChargedShot || pPlayer->m_Shared.m_bIsShotCharging)
	{
		shotStrength = pPlayer->GetChargedShotStrength();

		// Flash
		if (shotStrength > 0.9f)
		{
			//relStamina = 1;
			bgColor = Color(255, 255, 255, 255);
		}

		drawChargedshotIndicator = true;
	}

	const int CenterX = GetWide() / 2;
	const int CenterY = min(GetTall() - 30, GetTall() / 2 + (100 -::input->GetCameraAngles()[PITCH]) * (2 * GetTall() / 480.0f));

	// Draw stamina bar back
	surface()->DrawSetColor(bgColor);
	surface()->DrawFilledRect(
		CenterX - (BAR_WIDTH / 2),
		CenterY - (BAR_HEIGHT / 2),
		CenterX + (BAR_WIDTH / 2),
		CenterY + (BAR_HEIGHT / 2));

	// Draw missing max stamina
	surface()->DrawSetColor(missingMaxStaminaColor);
	surface()->DrawFilledRect(
		CenterX + (BAR_WIDTH / 2 - BAR_HPADDING) - (1.0f - pPlayer->m_Shared.GetMaxStamina() / 100.0f) * (BAR_WIDTH - 2 * BAR_VPADDING),
		CenterY - (BAR_HEIGHT / 2 - BAR_VPADDING),
		CenterX + (BAR_WIDTH / 2 - BAR_HPADDING),
		CenterY + (BAR_HEIGHT / 2 - BAR_VPADDING));

	// Draw stamina bar front
	surface()->DrawSetColor(fgColor);
	surface()->DrawFilledRect(
		CenterX - (BAR_WIDTH / 2 - BAR_HPADDING),
		CenterY - (BAR_HEIGHT / 2 - BAR_VPADDING),
		CenterX - (BAR_WIDTH / 2 - BAR_HPADDING) + relStamina * (BAR_WIDTH - 2 * BAR_VPADDING),
		CenterY + (BAR_HEIGHT / 2 - BAR_VPADDING));

	const int partCount = 4;
	const int hMargin = BAR_WIDTH / partCount;

	surface()->DrawSetColor(bgColor);

	for (int i = 1; i < partCount; i++)
	{
		surface()->DrawFilledRect(
			CenterX - BAR_WIDTH / 2 + i * hMargin,
			CenterY - BAR_HEIGHT / 2,
			CenterX - BAR_WIDTH / 2 + i * hMargin + 1,
			CenterY + BAR_HEIGHT / 2);
	}

	if (drawChargedshotIndicator)
	{
		// Draw chargedshot indicator back
		surface()->DrawSetColor(0, 0, 0, 255);
		surface()->DrawFilledRect(
			CenterX - BAR_WIDTH / 2 + shotStrength * BAR_WIDTH - INDICATOR_WIDTH / 2,
			CenterY - BAR_HEIGHT / 2 - INDICATOR_OFFSET,
			CenterX - BAR_WIDTH / 2 + shotStrength * BAR_WIDTH + INDICATOR_WIDTH / 2,
			CenterY + BAR_HEIGHT / 2 + INDICATOR_OFFSET);

		// Draw chargedshot indicator front
		surface()->DrawSetColor(255, 255, 255, 255);
		surface()->DrawFilledRect(
			CenterX - BAR_WIDTH / 2 + shotStrength * BAR_WIDTH - INDICATOR_WIDTH / 2 + INDICATOR_BORDER,
			CenterY - BAR_HEIGHT / 2 - INDICATOR_OFFSET + INDICATOR_BORDER,
			CenterX - BAR_WIDTH / 2 + shotStrength * BAR_WIDTH + INDICATOR_WIDTH / 2 - INDICATOR_BORDER,
			CenterY + BAR_HEIGHT / 2 + INDICATOR_OFFSET - INDICATOR_BORDER);

	}
}<commit_msg>Improve the chargedshot and stamina bar<commit_after>//========= Copyright  1996-2005, Valve Corporation, All rights reserved. ============//
//
// Purpose:
//
// $NoKeywords: $
//=============================================================================//

#include "cbase.h"
#include "hud.h"
#include "sdk_hud_stamina.h"
#include "hud_macros.h"
#include "c_sdk_player.h"
#include "iclientmode.h"
#include <vgui_controls/AnimationController.h>
#include <vgui/ISurface.h>
#include <vgui/ILocalize.h>
#include <igameresources.h>
#include "c_baseplayer.h"
#include "in_buttons.h"
#include "sdk_gamerules.h"
#include "c_ios_replaymanager.h"
#include "c_match_ball.h"
#include "view.h"
#include "ios_camera.h"
#include "iinput.h"

// memdbgon must be the last include file in a .cpp file!!!
#include "tier0/memdbgon.h"

extern ConVar centeredstaminabar;

class CHudChargedshotBar : public CHudElement, public vgui::Panel
{
	DECLARE_CLASS_SIMPLE( CHudChargedshotBar, vgui::Panel );

public:
	CHudChargedshotBar( const char *pElementName );
	virtual void	Init( void );
	virtual void	Reset( void );
	virtual void	OnThink( void );
	bool			ShouldDraw( void );

protected:
	virtual void	ApplySchemeSettings( vgui::IScheme *scheme );
	virtual void	Paint();

	float m_flNextUpdate;
};

DECLARE_HUDELEMENT( CHudChargedshotBar );

using namespace vgui;

enum { BAR_WIDTH = 80, BAR_HEIGHT = 14, BAR_HPADDING = 2, BAR_VPADDING = 6, SHOTBAR_HEIGHT = 6 };
enum { INDICATOR_WIDTH = 9, INDICATOR_OFFSET = 7, INDICATOR_BORDER = 1 };

//-----------------------------------------------------------------------------
// Purpose: Constructor
//-----------------------------------------------------------------------------
CHudChargedshotBar::CHudChargedshotBar( const char *pElementName ) : CHudElement( pElementName ), BaseClass( NULL, "HudChargedshotBar" )
{
	SetHiddenBits(HIDEHUD_POWERSHOTBAR);

	vgui::Panel *pParent = g_pClientMode->GetViewport();
	SetParent( pParent );

	m_flNextUpdate = gpGlobals->curtime;
}

void CHudChargedshotBar::ApplySchemeSettings( IScheme *scheme )
{
	BaseClass::ApplySchemeSettings(scheme);

	SetBounds(0, 0, ScreenWidth(), ScreenHeight());
}

//-----------------------------------------------------------------------------
// Purpose:
//-----------------------------------------------------------------------------
void CHudChargedshotBar::Init( void )
{
}

//-----------------------------------------------------------------------------
// Purpose:
//-----------------------------------------------------------------------------
void CHudChargedshotBar::Reset( void )
{
	Init();
}

//-----------------------------------------------------------------------------
// Purpose: Save CPU cycles by letting the HUD system early cull
// costly traversal.  Called per frame, return true if thinking and
// painting need to occur.
//-----------------------------------------------------------------------------
bool CHudChargedshotBar::ShouldDraw()
{
	if (!CHudElement::ShouldDraw())
		return false;

	C_SDKPlayer *pPlayer = C_SDKPlayer::GetLocalSDKPlayer();
	if (!pPlayer || pPlayer->GetTeamNumber() != TEAM_A && pPlayer->GetTeamNumber() != TEAM_B)
		return false;

	if (GetMatchBall() && GetMatchBall()->m_eBallState == BALL_STATE_GOAL)
		return false;

	if (GetReplayManager() && GetReplayManager()->IsReplaying())
		return false;

	return true;
}

//-----------------------------------------------------------------------------
// Purpose:
//-----------------------------------------------------------------------------
void CHudChargedshotBar::OnThink( void )
{
	BaseClass::OnThink();
}

//-----------------------------------------------------------------------------
// Purpose: draws the stamina bar
//-----------------------------------------------------------------------------
void CHudChargedshotBar::Paint()
{
	C_SDKPlayer *pPlayer = C_SDKPlayer::GetLocalSDKPlayer();

	if ( !pPlayer )
		return;

	float stamina = pPlayer->m_Shared.GetStamina() / 100.0f;
	float shotStrength = pPlayer->m_Shared.m_bDoChargedShot || pPlayer->m_Shared.m_bIsShotCharging ? pPlayer->GetChargedShotStrength() : 0;

	Color staminaFgColor, staminaBgColor, shotFgColor;

	if (pPlayer->GetFlags() & FL_REMOTECONTROLLED)
		staminaBgColor = Color(100, 100, 100, 255);
	else if (GetMatchBall() && GetMatchBall()->m_bShotsBlocked)
		staminaBgColor = Color(139, 0, 0, 255);
	else if (GetMatchBall() && GetMatchBall()->m_bChargedshotBlocked)
		staminaBgColor = Color(255, 69, 0, 255);
	else
		staminaBgColor = Color(0, 0, 0, 255);

	staminaFgColor = Color(255 * (1 - stamina), 255 * stamina, 0, 255);
	staminaBgColor = Color(0, 0, 0, 255);

	shotFgColor = Color(255, 255, 255, 255);

	int centerX = GetWide() / 2;
	int centerY = min(GetTall() - 30, GetTall() / 2 + (100 -::input->GetCameraAngles()[PITCH]) * (2 * GetTall() / 480.0f));

	// Draw stamina bar back
	surface()->DrawSetColor(staminaBgColor);
	surface()->DrawFilledRect(
		centerX - (BAR_WIDTH / 2 + BAR_HPADDING),
		centerY - (BAR_HEIGHT / 2 + BAR_VPADDING),
		centerX + (BAR_WIDTH / 2 + BAR_HPADDING),
		centerY + (BAR_HEIGHT / 2 + BAR_VPADDING));

	// Draw stamina bar front
	surface()->DrawSetColor(staminaFgColor);
	surface()->DrawFilledRect(
		centerX - (BAR_WIDTH / 2),
		centerY - (BAR_HEIGHT / 2),
		centerX - (BAR_WIDTH / 2) + stamina * (BAR_WIDTH),
		centerY + (BAR_HEIGHT / 2));

	// Draw shot bars
	surface()->DrawSetColor(shotFgColor);
	surface()->DrawFilledRect(
		centerX - (BAR_WIDTH / 2),
		centerY - (BAR_HEIGHT / 2 + 4),
		centerX - (BAR_WIDTH / 2) + shotStrength * (BAR_WIDTH),
		centerY - (BAR_HEIGHT / 2 - 2));

	surface()->DrawSetColor(shotFgColor);
	surface()->DrawFilledRect(
		centerX - (BAR_WIDTH / 2),
		centerY + (BAR_HEIGHT / 2 - 2),
		centerX - (BAR_WIDTH / 2) + shotStrength * (BAR_WIDTH),
		centerY + (BAR_HEIGHT / 2 + 4));
}<|endoftext|>