<commit_before>/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 * @(#)root/roofitcore:$Id$
 * Authors:                                                                  *
 *   WV, Wouter Verkerke, UC Santa Barbara, verkerke@slac.stanford.edu       *
 *   DK, David Kirkby,    UC Irvine,         dkirkby@uci.edu                 *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California          *
 *                          and Stanford University. All rights reserved.    *
 *                                                                           *
 * Redistribution and use in source and binary forms,                        *
 * with or without modification, are permitted according to the terms        *
 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)             *
 *****************************************************************************/

// -- CLASS DESCRIPTION [CAT] --
// RooMappedCategory provides a category-to-category mapping defined
// by pattern matching on their state labels
//
// The mapping function consists of a series of wild card regular expressions.
// Each expression is matched to the input categories state labels, and an associated
// output state label.

#include <cstdio>
#include <memory>
#include <cstdlib>

#include "RooFit.h"

#include "Riostream.h"
#include "TString.h"
#include "RooMappedCategory.h"
#include "RooStreamParser.h"
#include "RooMsgService.h"
#include "TBuffer.h"
#include "TString.h"
#include "RooAbsCache.h"

ClassImp(RooMappedCategory)
ClassImp(RooMappedCategory::Entry)

class RooMappedCategoryCache : public RooAbsCache {
    public:
        RooMappedCategoryCache(RooAbsArg* owner = 0) : RooAbsCache(owner)
        { initialise(); }
        RooMappedCategoryCache(const RooAbsCache& other, RooAbsArg* owner = 0) :
            RooAbsCache(other, owner)
        { initialise(); }

        // look up our parent's output based on our parent's input category index
        const RooCatType* lookup(Int_t idx) const
        { return _map[idx]; }

        virtual void wireCache()
        { _map.clear(); initialise(); }

        virtual Bool_t redirectServersHook(const RooAbsCollection& /*newServerList*/, Bool_t /*mustReplaceAll*/, Bool_t /*nameChange*/, Bool_t /*isRecursive*/)
        { _map.clear(); initialise(); return kFALSE; }

    private:
        mutable std::map<Int_t, const RooCatType*> _map;

        // pre-map categories of input category on something easily searchable
        // like the index (not the name!)
        void initialise()
        {
            const RooMappedCategory& parent = *static_cast<const RooMappedCategory*>(_owner);
            std::unique_ptr<TIterator> tit(static_cast<const RooAbsCategory&>(
                        parent._inputCat.arg()).typeIterator());
            for (const RooCatType* inCat = static_cast<const RooCatType*>(tit->Next());
                    inCat; inCat = static_cast<const RooCatType*>(tit->Next())) {
                const char* inKey = inCat->GetName();
                // Scan array of regexps
                bool found = false;
                for (std::map<std::string, RooMappedCategory::Entry>::const_iterator
                        iter = parent._mapArray.begin(),
                        end = parent._mapArray.end(); end != iter; ++iter) {
                    if (iter->second.match(inKey)) {
                        found = true;
                        _map[inCat->getVal()] = &(iter->second.outCat());
                        break;
                    }
                }
                if (!found) _map[inCat->getVal()] = parent._defCat;
            }
        }
};

RooMappedCategory::RooMappedCategory(const char *name, const char *title, RooAbsCategory& inputCat, const char* defOut, Int_t defOutIdx) :
  RooAbsCategory(name, title), _inputCat("input","Input category",this,inputCat),
  _mapcache(0)
{
  // Constructor with input category and name of default output state, which is assigned
  // to all input category states that do not follow any mapping rule.
  if (defOutIdx==NoCatIdx) {
    _defCat = (RooCatType*) defineType(defOut) ;
  } else {
    _defCat = (RooCatType*) defineType(defOut,defOutIdx) ;
  }
}

RooMappedCategory::RooMappedCategory(const RooMappedCategory& other, const char *name) :
  RooAbsCategory(other,name), _inputCat("input",this,other._inputCat), _mapArray(other._mapArray),
  _mapcache(0)
{
  _defCat = (RooCatType*) lookupType(other._defCat->GetName()) ;
}

RooMappedCategory::~RooMappedCategory()
{
    // Destructor
    delete _mapcache;
}

Bool_t RooMappedCategory::map(const char* inKeyRegExp, const char* outKey, Int_t outIdx)
{
  // Add mapping rule: any input category state label matching the 'inKeyRegExp'
  // wildcard expression will be mapped to an output state with name 'outKey'
  //
  // Rules are evaluated in the order they were added. In case an input state
  // matches more than one rule, the first rules output state will be assigned

  if (!inKeyRegExp || !outKey) return kTRUE ;

  // Check if pattern is already registered
  if (_mapArray.find(inKeyRegExp)!=_mapArray.end()) {
    coutE(InputArguments) << "RooMappedCategory::map(" << GetName() << "): ERROR expression "
                          << inKeyRegExp << " already mapped" << std::endl ;
    return kTRUE ;
  }

  // Check if output type exists, if not register
  const RooCatType* outType = lookupType(outKey) ;
  if (!outType) {
    if (outIdx==NoCatIdx) {
      outType = defineType(outKey) ;
    } else {
      outType = defineType(outKey,outIdx) ;
    }
  }
  if (!outType) {
    coutE(InputArguments) << "RooMappedCategory::map(" << GetName()
                          << "): ERROR, unable to output type " << outKey << std::endl ;
    return kTRUE ;
  }

  // Create new map entry ;
  Entry e(inKeyRegExp,outType) ;
  if (!e.ok()) {
    coutE(InputArguments) << "RooMappedCategory::map(" << GetName()
                          << "): ERROR, expression " << inKeyRegExp << " didn't compile" << std::endl ;
    return kTRUE ;
  }

  _mapArray[inKeyRegExp] = e ;
  return kFALSE ;
}

RooCatType RooMappedCategory::evaluate() const
{
    const RooMappedCategoryCache* cache = getOrCreateCache();
    return *(cache->lookup(Int_t(_inputCat)));
}

const RooMappedCategoryCache* RooMappedCategory::getOrCreateCache() const
{
    if (!_mapcache) _mapcache = new RooMappedCategoryCache(
            const_cast<RooMappedCategory*>(this));
    return _mapcache;
}

void RooMappedCategory::printMultiline(std::ostream& os, Int_t content, Bool_t verbose, TString indent) const
{
  // Print info about this mapped category to the specified stream. In addition to the info
  // from RooAbsCategory::printStream() we add:
  //
  //  Standard : input category
  //     Shape : default value
  //   Verbose : list of mapping rules

  RooAbsCategory::printMultiline(os,content,verbose,indent);

  if (verbose) {
    os << indent << "--- RooMappedCategory ---" << std::endl
       << indent << "  Maps from " ;
    _inputCat.arg().printStream(os,0,kStandard);

    os << indent << "  Default value is ";
    _defCat->printStream(os,kName|kValue,kSingleLine);

    os << indent << "  Mapping rules:" << std::endl;
    for (std::map<std::string,Entry>::const_iterator iter = _mapArray.begin() ; iter!=_mapArray.end() ; iter++) {
      os << indent << "  " << iter->first << " -> " << iter->second.outCat().GetName() << std::endl ;
    }
  }
}

Bool_t RooMappedCategory::readFromStream(std::istream& is, Bool_t compact, Bool_t /*verbose*/)
{
  // Read object contents from given stream
   if (compact) {
     coutE(InputArguments) << "RooMappedCategory::readFromSteam(" << GetName() << "): can't read in compact mode" << std::endl ;
     return kTRUE ;
   } else {

     //Clear existing definitions, but preserve default output
     TString defCatName(_defCat->GetName()) ;
     _mapArray.clear() ;
     delete _mapcache;
     _mapcache = 0;
     clearTypes() ;
     _defCat = (RooCatType*) defineType(defCatName) ;

     TString token,errorPrefix("RooMappedCategory::readFromStream(") ;
     errorPrefix.Append(GetName()) ;
     errorPrefix.Append(")") ;
     RooStreamParser parser(is,errorPrefix) ;
     parser.setPunctuation(":,") ;

     TString destKey,srcKey ;
     Bool_t readToken(kTRUE) ;

    // Loop over definition sequences
     while(1) {
       if (readToken) token=parser.readToken() ;
       if (token.IsNull()) break ;
       readToken=kTRUE ;

       destKey = token ;
       if (parser.expectToken(":",kTRUE)) return kTRUE ;

       // Loop over list of sources for this destination
       while(1) {
         srcKey = parser.readToken() ;
         token = parser.readToken() ;

         // Map a value
         if (map(srcKey,destKey)) return kTRUE ;

         // Unless next token is ',' current token
         // is destination part of next sequence
         if (token.CompareTo(",")) {
           readToken = kFALSE ;
           break ;
         }
       }
     }
     return kFALSE ;
   }
   //return kFALSE ; // statement unreachable (OSF)
}

//_____________________________________________________________________________
void RooMappedCategory::printMetaArgs(std::ostream& os) const
{
  // Customized printing of arguments of a RooMappedCategory to more intuitively reflect the contents of the
  // product operator construction

  // Scan array of regexps
  RooCatType prevOutCat ;
  Bool_t first(kTRUE) ;
  os << "map=(" ;
  for (std::map<std::string,Entry>::const_iterator iter = _mapArray.begin() ; iter!=_mapArray.end() ; iter++) {
    if (iter->second.outCat().getVal()!=prevOutCat.getVal()) {
      if (!first) { os << " " ; }
      first=kFALSE ;

      os << iter->second.outCat().GetName() << ":" << iter->first ;
      prevOutCat=iter->second.outCat() ;
    } else {
      os << "," << iter->first ;
    }
  }

  if (!first) { os << " " ; }
  os << _defCat->GetName() << ":*" ;

  os << ") " ;
}

void RooMappedCategory::writeToStream(std::ostream& os, Bool_t compact) const
{
  // Write object contents to given stream
  if (compact) {
    // Write value only
    os << getLabel() ;
  } else {
    // Write mapping expression

    // Scan array of regexps
    RooCatType prevOutCat ;
    Bool_t first(kTRUE) ;
    for (std::map<std::string,Entry>::const_iterator iter = _mapArray.begin() ; iter!=_mapArray.end() ; iter++) {
      if (iter->second.outCat().getVal()!=prevOutCat.getVal()) {
        if (!first) { os << " " ; }
        first=kFALSE ;

        os << iter->second.outCat().GetName() << "<-" << iter->first ;
        prevOutCat=iter->second.outCat() ;
      } else {
        os << "," << iter->first ;
      }
    }

    if (!first) { os << " " ; }
    os << _defCat->GetName() << ":*" ;
  }
}

//_____________________________________________________________________________
RooMappedCategory::Entry& RooMappedCategory::Entry::operator=(const RooMappedCategory::Entry& other)
{
  if (&other==this) return *this ;

  _expr = other._expr ;
  _cat = other._cat ;

  if (_regexp) {
    delete _regexp ;
  }
  _regexp = new TRegexp(_expr.Data(),kTRUE) ;

  return *this;
}

//_____________________________________________________________________________
TString RooMappedCategory::Entry::mangle(const char* exp) const
{
  // Mangle name : escape regexp character '+'
  TString t ;
  const char *c = exp ;
  while(*c) {
    if (*c=='+') t.Append('\\') ;
    t.Append(*c) ;
    c++ ;
  }
  return t ;
}

//_____________________________________________________________________________
void RooMappedCategory::Entry::Streamer(TBuffer &R__b)
{
  typedef ::RooMappedCategory::Entry ThisClass;

   // Stream an object of class RooWorkspace::CodeRepo.
   if (R__b.IsReading()) {

     UInt_t R__s, R__c;
     R__b.ReadVersion(&R__s, &R__c);

     // Stream contents of ClassFiles map
     R__b >> _expr ;
     _cat.Streamer(R__b) ;
     _regexp = new TRegexp(_expr.Data(),kTRUE) ;
     R__b.CheckByteCount(R__s, R__c, ThisClass::IsA());

   } else {

     UInt_t R__c;
     R__c = R__b.WriteVersion(ThisClass::IsA(), kTRUE);

     // Stream contents of ClassRelInfo map
     R__b << _expr ;
     _cat.Streamer(R__b) ;

     R__b.SetByteCount(R__c, kTRUE);

   }
}
<commit_msg>v5.34 can not require C++11<commit_after>/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 * @(#)root/roofitcore:$Id$
 * Authors:                                                                  *
 *   WV, Wouter Verkerke, UC Santa Barbara, verkerke@slac.stanford.edu       *
 *   DK, David Kirkby,    UC Irvine,         dkirkby@uci.edu                 *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California          *
 *                          and Stanford University. All rights reserved.    *
 *                                                                           *
 * Redistribution and use in source and binary forms,                        *
 * with or without modification, are permitted according to the terms        *
 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)             *
 *****************************************************************************/

// -- CLASS DESCRIPTION [CAT] --
// RooMappedCategory provides a category-to-category mapping defined
// by pattern matching on their state labels
//
// The mapping function consists of a series of wild card regular expressions.
// Each expression is matched to the input categories state labels, and an associated
// output state label.

#include <cstdio>
#include <memory>
#include <cstdlib>

#include "RooFit.h"

#include "Riostream.h"
#include "TString.h"
#include "RooMappedCategory.h"
#include "RooStreamParser.h"
#include "RooMsgService.h"
#include "TBuffer.h"
#include "TString.h"
#include "RooAbsCache.h"

ClassImp(RooMappedCategory)
ClassImp(RooMappedCategory::Entry)

class RooMappedCategoryCache : public RooAbsCache {
    public:
        RooMappedCategoryCache(RooAbsArg* owner = 0) : RooAbsCache(owner)
        { initialise(); }
        RooMappedCategoryCache(const RooAbsCache& other, RooAbsArg* owner = 0) :
            RooAbsCache(other, owner)
        { initialise(); }

        // look up our parent's output based on our parent's input category index
        const RooCatType* lookup(Int_t idx) const
        { return _map[idx]; }

        virtual void wireCache()
        { _map.clear(); initialise(); }

        virtual Bool_t redirectServersHook(const RooAbsCollection& /*newServerList*/, Bool_t /*mustReplaceAll*/, Bool_t /*nameChange*/, Bool_t /*isRecursive*/)
        { _map.clear(); initialise(); return kFALSE; }

    private:
        mutable std::map<Int_t, const RooCatType*> _map;

        // pre-map categories of input category on something easily searchable
        // like the index (not the name!)
        void initialise()
        {
            const RooMappedCategory& parent = *static_cast<const RooMappedCategory*>(_owner);
            TIterator* tit(static_cast<const RooAbsCategory&>(
                        parent._inputCat.arg()).typeIterator());
            for (const RooCatType* inCat = static_cast<const RooCatType*>(tit->Next());
                    inCat; inCat = static_cast<const RooCatType*>(tit->Next())) {
                const char* inKey = inCat->GetName();
                // Scan array of regexps
                bool found = false;
                for (std::map<std::string, RooMappedCategory::Entry>::const_iterator
                        iter = parent._mapArray.begin(),
                        end = parent._mapArray.end(); end != iter; ++iter) {
                    if (iter->second.match(inKey)) {
                        found = true;
                        _map[inCat->getVal()] = &(iter->second.outCat());
                        break;
                    }
                }
                if (!found) _map[inCat->getVal()] = parent._defCat;
            }
            delete tit;
        }
};

RooMappedCategory::RooMappedCategory(const char *name, const char *title, RooAbsCategory& inputCat, const char* defOut, Int_t defOutIdx) :
  RooAbsCategory(name, title), _inputCat("input","Input category",this,inputCat),
  _mapcache(0)
{
  // Constructor with input category and name of default output state, which is assigned
  // to all input category states that do not follow any mapping rule.
  if (defOutIdx==NoCatIdx) {
    _defCat = (RooCatType*) defineType(defOut) ;
  } else {
    _defCat = (RooCatType*) defineType(defOut,defOutIdx) ;
  }
}

RooMappedCategory::RooMappedCategory(const RooMappedCategory& other, const char *name) :
  RooAbsCategory(other,name), _inputCat("input",this,other._inputCat), _mapArray(other._mapArray),
  _mapcache(0)
{
  _defCat = (RooCatType*) lookupType(other._defCat->GetName()) ;
}

RooMappedCategory::~RooMappedCategory()
{
    // Destructor
    delete _mapcache;
}

Bool_t RooMappedCategory::map(const char* inKeyRegExp, const char* outKey, Int_t outIdx)
{
  // Add mapping rule: any input category state label matching the 'inKeyRegExp'
  // wildcard expression will be mapped to an output state with name 'outKey'
  //
  // Rules are evaluated in the order they were added. In case an input state
  // matches more than one rule, the first rules output state will be assigned

  if (!inKeyRegExp || !outKey) return kTRUE ;

  // Check if pattern is already registered
  if (_mapArray.find(inKeyRegExp)!=_mapArray.end()) {
    coutE(InputArguments) << "RooMappedCategory::map(" << GetName() << "): ERROR expression "
                          << inKeyRegExp << " already mapped" << std::endl ;
    return kTRUE ;
  }

  // Check if output type exists, if not register
  const RooCatType* outType = lookupType(outKey) ;
  if (!outType) {
    if (outIdx==NoCatIdx) {
      outType = defineType(outKey) ;
    } else {
      outType = defineType(outKey,outIdx) ;
    }
  }
  if (!outType) {
    coutE(InputArguments) << "RooMappedCategory::map(" << GetName()
                          << "): ERROR, unable to output type " << outKey << std::endl ;
    return kTRUE ;
  }

  // Create new map entry ;
  Entry e(inKeyRegExp,outType) ;
  if (!e.ok()) {
    coutE(InputArguments) << "RooMappedCategory::map(" << GetName()
                          << "): ERROR, expression " << inKeyRegExp << " didn't compile" << std::endl ;
    return kTRUE ;
  }

  _mapArray[inKeyRegExp] = e ;
  return kFALSE ;
}

RooCatType RooMappedCategory::evaluate() const
{
    const RooMappedCategoryCache* cache = getOrCreateCache();
    return *(cache->lookup(Int_t(_inputCat)));
}

const RooMappedCategoryCache* RooMappedCategory::getOrCreateCache() const
{
    if (!_mapcache) _mapcache = new RooMappedCategoryCache(
            const_cast<RooMappedCategory*>(this));
    return _mapcache;
}

void RooMappedCategory::printMultiline(std::ostream& os, Int_t content, Bool_t verbose, TString indent) const
{
  // Print info about this mapped category to the specified stream. In addition to the info
  // from RooAbsCategory::printStream() we add:
  //
  //  Standard : input category
  //     Shape : default value
  //   Verbose : list of mapping rules

  RooAbsCategory::printMultiline(os,content,verbose,indent);

  if (verbose) {
    os << indent << "--- RooMappedCategory ---" << std::endl
       << indent << "  Maps from " ;
    _inputCat.arg().printStream(os,0,kStandard);

    os << indent << "  Default value is ";
    _defCat->printStream(os,kName|kValue,kSingleLine);

    os << indent << "  Mapping rules:" << std::endl;
    for (std::map<std::string,Entry>::const_iterator iter = _mapArray.begin() ; iter!=_mapArray.end() ; iter++) {
      os << indent << "  " << iter->first << " -> " << iter->second.outCat().GetName() << std::endl ;
    }
  }
}

Bool_t RooMappedCategory::readFromStream(std::istream& is, Bool_t compact, Bool_t /*verbose*/)
{
  // Read object contents from given stream
   if (compact) {
     coutE(InputArguments) << "RooMappedCategory::readFromSteam(" << GetName() << "): can't read in compact mode" << std::endl ;
     return kTRUE ;
   } else {

     //Clear existing definitions, but preserve default output
     TString defCatName(_defCat->GetName()) ;
     _mapArray.clear() ;
     delete _mapcache;
     _mapcache = 0;
     clearTypes() ;
     _defCat = (RooCatType*) defineType(defCatName) ;

     TString token,errorPrefix("RooMappedCategory::readFromStream(") ;
     errorPrefix.Append(GetName()) ;
     errorPrefix.Append(")") ;
     RooStreamParser parser(is,errorPrefix) ;
     parser.setPunctuation(":,") ;

     TString destKey,srcKey ;
     Bool_t readToken(kTRUE) ;

    // Loop over definition sequences
     while(1) {
       if (readToken) token=parser.readToken() ;
       if (token.IsNull()) break ;
       readToken=kTRUE ;

       destKey = token ;
       if (parser.expectToken(":",kTRUE)) return kTRUE ;

       // Loop over list of sources for this destination
       while(1) {
         srcKey = parser.readToken() ;
         token = parser.readToken() ;

         // Map a value
         if (map(srcKey,destKey)) return kTRUE ;

         // Unless next token is ',' current token
         // is destination part of next sequence
         if (token.CompareTo(",")) {
           readToken = kFALSE ;
           break ;
         }
       }
     }
     return kFALSE ;
   }
   //return kFALSE ; // statement unreachable (OSF)
}

//_____________________________________________________________________________
void RooMappedCategory::printMetaArgs(std::ostream& os) const
{
  // Customized printing of arguments of a RooMappedCategory to more intuitively reflect the contents of the
  // product operator construction

  // Scan array of regexps
  RooCatType prevOutCat ;
  Bool_t first(kTRUE) ;
  os << "map=(" ;
  for (std::map<std::string,Entry>::const_iterator iter = _mapArray.begin() ; iter!=_mapArray.end() ; iter++) {
    if (iter->second.outCat().getVal()!=prevOutCat.getVal()) {
      if (!first) { os << " " ; }
      first=kFALSE ;

      os << iter->second.outCat().GetName() << ":" << iter->first ;
      prevOutCat=iter->second.outCat() ;
    } else {
      os << "," << iter->first ;
    }
  }

  if (!first) { os << " " ; }
  os << _defCat->GetName() << ":*" ;

  os << ") " ;
}

void RooMappedCategory::writeToStream(std::ostream& os, Bool_t compact) const
{
  // Write object contents to given stream
  if (compact) {
    // Write value only
    os << getLabel() ;
  } else {
    // Write mapping expression

    // Scan array of regexps
    RooCatType prevOutCat ;
    Bool_t first(kTRUE) ;
    for (std::map<std::string,Entry>::const_iterator iter = _mapArray.begin() ; iter!=_mapArray.end() ; iter++) {
      if (iter->second.outCat().getVal()!=prevOutCat.getVal()) {
        if (!first) { os << " " ; }
        first=kFALSE ;

        os << iter->second.outCat().GetName() << "<-" << iter->first ;
        prevOutCat=iter->second.outCat() ;
      } else {
        os << "," << iter->first ;
      }
    }

    if (!first) { os << " " ; }
    os << _defCat->GetName() << ":*" ;
  }
}

//_____________________________________________________________________________
RooMappedCategory::Entry& RooMappedCategory::Entry::operator=(const RooMappedCategory::Entry& other)
{
  if (&other==this) return *this ;

  _expr = other._expr ;
  _cat = other._cat ;

  if (_regexp) {
    delete _regexp ;
  }
  _regexp = new TRegexp(_expr.Data(),kTRUE) ;

  return *this;
}

//_____________________________________________________________________________
TString RooMappedCategory::Entry::mangle(const char* exp) const
{
  // Mangle name : escape regexp character '+'
  TString t ;
  const char *c = exp ;
  while(*c) {
    if (*c=='+') t.Append('\\') ;
    t.Append(*c) ;
    c++ ;
  }
  return t ;
}

//_____________________________________________________________________________
void RooMappedCategory::Entry::Streamer(TBuffer &R__b)
{
  typedef ::RooMappedCategory::Entry ThisClass;

   // Stream an object of class RooWorkspace::CodeRepo.
   if (R__b.IsReading()) {

     UInt_t R__s, R__c;
     R__b.ReadVersion(&R__s, &R__c);

     // Stream contents of ClassFiles map
     R__b >> _expr ;
     _cat.Streamer(R__b) ;
     _regexp = new TRegexp(_expr.Data(),kTRUE) ;
     R__b.CheckByteCount(R__s, R__c, ThisClass::IsA());

   } else {

     UInt_t R__c;
     R__c = R__b.WriteVersion(ThisClass::IsA(), kTRUE);

     // Stream contents of ClassRelInfo map
     R__b << _expr ;
     _cat.Streamer(R__b) ;

     R__b.SetByteCount(R__c, kTRUE);

   }
}
<|endoftext|>


<commit_before>//------------------------------------------------------------------------------
//  oryol-shdc
//------------------------------------------------------------------------------
#include "ExportUtil/CmdLineArgs.h"
#include "ExportUtil/Log.h"
#include "spirv_hlsl.hpp"
#include "spirv_msl.hpp"
#include "pystring.h"
#include "cJSON.h"
#include <stdio.h>

using namespace OryolTools;
using namespace spv;
using namespace spirv_cross;
using namespace std;

//------------------------------------------------------------------------------
vector<uint32_t> read_spirv_file(const string& path) {
    FILE* fp = fopen(path.c_str(), "rb");
    if (!fp) {
        Log::Fatal("Failed to open SPIRV file '%s'\n", path.c_str());
    }
    fseek(fp, 0, SEEK_END);
    long len = ftell(fp) / sizeof(uint32_t);
    fseek(fp, 0, SEEK_SET);
    vector<uint32_t> spirv(len);
    if (fread(spirv.data(), sizeof(uint32_t), len, fp) != size_t(len)) {
        Log::Fatal("Error reading SPIRV file '%s'\n", path.c_str());
    }
    fclose(fp);
    return spirv;
}

//------------------------------------------------------------------------------
void write_source_file(const string& path, const string& content) {
    FILE* fp = fopen(path.c_str(), "w");
    if (!fp) {
        Log::Fatal("Failed to open '%s' for writing\n", path.c_str());
    }
    fwrite(content.c_str(), 1, content.length(), fp);
    fclose(fp);
}

//------------------------------------------------------------------------------
const char* type_to_uniform_type(const SPIRType& type) {
    if (type.basetype == SPIRType::Float) {
        if (type.columns == 1) {
            // scalar or vec
            switch (type.vecsize) {
                case 1: return "float";
                case 2: return "vec2";
                case 3: return "vec3";
                case 4: return "vec4";
            }
        }
        else {
            // a matrix
            if ((type.vecsize == 2) && (type.columns == 2)) {
                return "mat2";
            }
            else if ((type.vecsize == 3) && (type.columns == 3)) {
                return "mat3";
            }
            else if ((type.vecsize == 4) && (type.columns == 4)) {
                return "mat4";
            }
        }
    }
    else if (type.basetype == SPIRType::Int) {
        return "int";
    }
    else if (type.basetype == SPIRType::Boolean) {
        return "bool";
    }
    Log::Fatal("Invalid member type in uniform block! (expected: float, vec2, vec3, vec4, mat2, mat3, mat4, int, bool)\n");
    return nullptr;
}

//------------------------------------------------------------------------------
const char* type_to_attr_format(const SPIRType& type) {
    if (type.basetype == SPIRType::Float && type.columns == 1) {
        switch (type.vecsize) {
            case 1: return "float";
            case 2: return "vec2";
            case 3: return "vec3";
            case 4: return "vec4";
        }
    }
    Log::Fatal("Invalid vertex attribute type! (expected: float, vec2, vec3, vec4)\n");
    return nullptr;
}

//------------------------------------------------------------------------------
cJSON* extract_resource_info(Compiler* compiler) {
    cJSON* root = cJSON_CreateObject();

    // shader stage
    const char* stage_str = nullptr;
    switch (compiler->get_execution_model()) {
        case ExecutionModelVertex: stage_str = "vs"; break;
        case ExecutionModelFragment: stage_str = "fs"; break;
        default: break;
    }
    if (stage_str) {
        cJSON_AddItemToObject(root, "stage", cJSON_CreateString(stage_str));
    }
    else {
        Log::Fatal("only vertex- or fragment-shaders allowed!\n");
    }

    // uniform blocks
    int ub_slot = 0;
    ShaderResources res = compiler->get_shader_resources();
    cJSON* ub_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "uniform_blocks", ub_array);
    for (const Resource& ub_res : res.uniform_buffers) {
        const SPIRType& ub_type = compiler->get_type(ub_res.base_type_id);
        string ub_name = compiler->get_name(ub_res.id);
        if (ub_name.empty()) {
            char buf[64];
            snprintf(buf, sizeof(buf), "_%d", ub_res.id);
            ub_name = buf;
        }
        cJSON* ub = cJSON_CreateObject();
        cJSON_AddItemToArray(ub_array, ub);
        cJSON_AddItemToObject(ub, "type", cJSON_CreateString(ub_res.name.c_str()));
        cJSON_AddItemToObject(ub, "name", cJSON_CreateString(ub_name.c_str()));
        cJSON_AddItemToObject(ub, "slot", cJSON_CreateNumber(ub_slot++));
        cJSON* ub_members = cJSON_CreateArray();
        cJSON_AddItemToObject(ub, "members", ub_members);
        for (int m_index = 0; m_index < int(ub_type.member_types.size()); m_index++) {
            cJSON* ub_member = cJSON_CreateObject();
            cJSON_AddItemToArray(ub_members, ub_member);
            string m_name = compiler->get_member_name(ub_res.base_type_id, m_index);
            const SPIRType& m_type = compiler->get_type(ub_type.member_types[m_index]);
            const char* m_type_str = type_to_uniform_type(m_type);
            cJSON_AddItemToObject(ub_member, "name", cJSON_CreateString(m_name.c_str()));
            cJSON_AddItemToObject(ub_member, "type", cJSON_CreateString(m_type_str));
            int num = 1;
            if (m_type.array.size() > 0) {
                num = m_type.array[0];
            }
            cJSON_AddItemToObject(ub_member, "num", cJSON_CreateNumber(num));
        }
    }

    // textures
    int tex_slot = 0;
    cJSON* tex_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "textures", tex_array);
    for (const Resource& img_res : res.sampled_images) {
        const SPIRType& img_type = compiler->get_type(img_res.type_id);
        cJSON* tex = cJSON_CreateObject();
        cJSON_AddItemToArray(tex_array, tex);
        const char* tex_type_str = nullptr;
        if (img_type.image.arrayed) {
            if (img_type.image.dim == Dim2D) {
                tex_type_str = "sampler2DArray";
            }
        }
        else {
            switch (img_type.image.dim) {
                case Dim2D:     tex_type_str = "sampler2D"; break;
                case DimCube:   tex_type_str = "samplerCube"; break;
                case Dim3D:     tex_type_str = "sampler3D"; break;
                default:        break;
            }
        }
        if (!tex_type_str) {
            Log::Fatal("Invalid texture type! (expected: 2D, Cube, 3D or 2D-array)\n");
        }
        cJSON_AddItemToObject(tex, "name", cJSON_CreateString(img_res.name.c_str()));
        cJSON_AddItemToObject(tex, "type", cJSON_CreateString(tex_type_str));
        cJSON_AddItemToObject(tex, "slot", cJSON_CreateNumber(tex_slot++));
    }

    // stage inputs
    cJSON* inputs_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "inputs", inputs_array);
    for (const Resource& stage_input : res.stage_inputs) {
        cJSON* input = cJSON_CreateObject();
        cJSON_AddItemToArray(inputs_array, input);
        const SPIRType& type = compiler->get_type(stage_input.base_type_id);
        const char* type_str = type_to_attr_format(type);
        cJSON_AddItemToObject(input, "name", cJSON_CreateString(stage_input.name.c_str()));
        cJSON_AddItemToObject(input, "type", cJSON_CreateString(type_str));
        if (compiler->get_execution_model() == ExecutionModelVertex) {
            uint32_t loc = compiler->get_decoration(stage_input.id, DecorationLocation);
            cJSON_AddItemToObject(input, "slot", cJSON_CreateNumber(loc));
        }
    }

    // stage outputs
    cJSON* outputs_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "outputs", outputs_array);
    for (const Resource& stage_output : res.stage_outputs) {
        cJSON* output = cJSON_CreateObject();
        cJSON_AddItemToArray(outputs_array, output);
        const SPIRType& type = compiler->get_type(stage_output.base_type_id);
        const char* type_str = type_to_attr_format(type);
        cJSON_AddItemToObject(output, "name", cJSON_CreateString(stage_output.name.c_str()));
        cJSON_AddItemToObject(output, "type", cJSON_CreateString(type_str));
    }

    return root;
}

//------------------------------------------------------------------------------
void fix_vertex_attr_locations(Compiler* compiler) {
    if (compiler->get_execution_model() == ExecutionModelVertex) {
        ShaderResources res = compiler->get_shader_resources();
        for (const auto& input : res.stage_inputs) {
            int loc = -1;
            if (input.name == "position")       loc = 0;
            else if (input.name == "normal")    loc = 1;
            else if (input.name == "texcoord0") loc = 2;
            else if (input.name == "texcoord1") loc = 3;
            else if (input.name == "texcoord2") loc = 4;
            else if (input.name == "texcoord3") loc = 5;
            else if (input.name == "tangent")   loc = 6;
            else if (input.name == "binormal")  loc = 7;
            else if (input.name == "weights")   loc = 8;
            else if (input.name == "indices")   loc = 9;
            else if (input.name == "color0")    loc = 10;
            else if (input.name == "color1")    loc = 11;
            else if (input.name == "instance0") loc = 12;
            else if (input.name == "instance1") loc = 13;
            else if (input.name == "instance2") loc = 14;
            else if (input.name == "instance3") loc = 15;
            if (-1 != loc) {
                compiler->set_decoration(input.id, DecorationLocation, (uint32_t)loc);
            }
        }
    }
}

//------------------------------------------------------------------------------
void to_reflection_json(const vector<uint32_t>& spirv, const string& base_path) {
    CompilerGLSL compiler(spirv);
    fix_vertex_attr_locations(&compiler);
    cJSON* json = extract_resource_info(&compiler);
    char* json_raw_str = cJSON_Print(json);
    string json_str(json_raw_str);
    std::free(json_raw_str);
    cJSON_Delete(json);
    write_source_file(base_path + ".json", json_str);
}

//------------------------------------------------------------------------------
void to_glsl(const vector<uint32_t>& spirv, const string& base_path, const string& file_ext, uint32_t version, bool is_es) {
    CompilerGLSL compiler(spirv);
    auto opts = compiler.get_options();
    opts.version = version;
    opts.es = is_es;
    opts.vertex.fixup_clipspace = false;
    compiler.set_options(opts);
    fix_vertex_attr_locations(&compiler);
    string src = compiler.compile();
    if (src.empty()) {
        Log::Fatal("Failed to compile GLSL v100 source for '%s'!\n", base_path.c_str());
    }
    else {
        write_source_file(base_path + file_ext, src);
    }
}

//------------------------------------------------------------------------------
void to_hlsl_sm5(const vector<uint32_t>& spirv, const string& base_path) {
    CompilerHLSL compiler(spirv);
    auto opts = compiler.get_options();
    opts.shader_model = 50;
    opts.fixup_clipspace = true;
    compiler.set_options(opts);
    fix_vertex_attr_locations(&compiler);
    string src = compiler.compile();
    if (src.empty()) {
        Log::Fatal("Failed to compile HLSL5 source for '%s'!\n", base_path.c_str());
    }
    else {
        write_source_file(base_path + ".hlsl", src);
    }
}

//------------------------------------------------------------------------------
void to_mlsl(const vector<uint32_t>& spirv, const string& base_path) {
    CompilerMSL compiler(spirv);
    fix_vertex_attr_locations(&compiler);
    string src = compiler.compile();
    if (src.empty()) {
        Log::Fatal("Failed to compile MetalSL source for '%s'!\n", base_path.c_str());
    }
    else {
        write_source_file(base_path + ".metal", src);
    }
}

//------------------------------------------------------------------------------
int main(int argc, const char** argv) {
    CmdLineArgs args;
    args.AddBool("-help", "show help");
    args.AddString("-spirv", "SPIR-V input file", "");
    if (!args.Parse(argc, argv)) {
        Log::Warn("Failed to parse args!\n");
        return 10;
    }

    if (args.HasArg("-help")) {
        Log::Info("Oryol SPIR-V to GLSL/HLSL/MSL cross-compiler\n"
                  "Based on SPIRV-Cross: https://github.com/KhronosGroup/SPIRV-Cross\n");
        args.ShowHelp();
        return 0;
    }
    string spirv_path = args.GetString("-spirv");
    if (spirv_path.empty()) {
        Log::Fatal("-spirv arg expected");
    }

    // load SPIRV byte code
    auto spirv = read_spirv_file(spirv_path);

    // ...translate and write to output files...
    string base_path, ext;
    pystring::os::path::splitext(base_path, ext, spirv_path);
    to_reflection_json(spirv, base_path);
    to_glsl(spirv, base_path, ".glsl100.glsl", 100, true);
    to_glsl(spirv, base_path, ".glsl120.glsl", 120, false);
    to_glsl(spirv, base_path, ".glsles3.glsl", 300, true);
    to_glsl(spirv, base_path, ".glsl330.glsl", 330, false);
    to_hlsl_sm5(spirv, base_path);
    to_mlsl(spirv, base_path);

    return 0;
}
<commit_msg>Force ub matrices to ColMajor<commit_after>//------------------------------------------------------------------------------
//  oryol-shdc
//------------------------------------------------------------------------------
#include "ExportUtil/CmdLineArgs.h"
#include "ExportUtil/Log.h"
#include "spirv_hlsl.hpp"
#include "spirv_msl.hpp"
#include "pystring.h"
#include "cJSON.h"
#include <stdio.h>

using namespace OryolTools;
using namespace spv;
using namespace spirv_cross;
using namespace std;

//------------------------------------------------------------------------------
vector<uint32_t> read_spirv_file(const string& path) {
    FILE* fp = fopen(path.c_str(), "rb");
    if (!fp) {
        Log::Fatal("Failed to open SPIRV file '%s'\n", path.c_str());
    }
    fseek(fp, 0, SEEK_END);
    long len = ftell(fp) / sizeof(uint32_t);
    fseek(fp, 0, SEEK_SET);
    vector<uint32_t> spirv(len);
    if (fread(spirv.data(), sizeof(uint32_t), len, fp) != size_t(len)) {
        Log::Fatal("Error reading SPIRV file '%s'\n", path.c_str());
    }
    fclose(fp);
    return spirv;
}

//------------------------------------------------------------------------------
void write_source_file(const string& path, const string& content) {
    FILE* fp = fopen(path.c_str(), "w");
    if (!fp) {
        Log::Fatal("Failed to open '%s' for writing\n", path.c_str());
    }
    fwrite(content.c_str(), 1, content.length(), fp);
    fclose(fp);
}

//------------------------------------------------------------------------------
const char* type_to_uniform_type(const SPIRType& type) {
    if (type.basetype == SPIRType::Float) {
        if (type.columns == 1) {
            // scalar or vec
            switch (type.vecsize) {
                case 1: return "float";
                case 2: return "vec2";
                case 3: return "vec3";
                case 4: return "vec4";
            }
        }
        else {
            // a matrix
            if ((type.vecsize == 2) && (type.columns == 2)) {
                return "mat2";
            }
            else if ((type.vecsize == 3) && (type.columns == 3)) {
                return "mat3";
            }
            else if ((type.vecsize == 4) && (type.columns == 4)) {
                return "mat4";
            }
        }
    }
    else if (type.basetype == SPIRType::Int) {
        return "int";
    }
    else if (type.basetype == SPIRType::Boolean) {
        return "bool";
    }
    Log::Fatal("Invalid member type in uniform block! (expected: float, vec2, vec3, vec4, mat2, mat3, mat4, int, bool)\n");
    return nullptr;
}

//------------------------------------------------------------------------------
const char* type_to_attr_format(const SPIRType& type) {
    if (type.basetype == SPIRType::Float && type.columns == 1) {
        switch (type.vecsize) {
            case 1: return "float";
            case 2: return "vec2";
            case 3: return "vec3";
            case 4: return "vec4";
        }
    }
    Log::Fatal("Invalid vertex attribute type! (expected: float, vec2, vec3, vec4)\n");
    return nullptr;
}

//------------------------------------------------------------------------------
cJSON* extract_resource_info(Compiler* compiler) {
    cJSON* root = cJSON_CreateObject();

    // shader stage
    const char* stage_str = nullptr;
    switch (compiler->get_execution_model()) {
        case ExecutionModelVertex: stage_str = "vs"; break;
        case ExecutionModelFragment: stage_str = "fs"; break;
        default: break;
    }
    if (stage_str) {
        cJSON_AddItemToObject(root, "stage", cJSON_CreateString(stage_str));
    }
    else {
        Log::Fatal("only vertex- or fragment-shaders allowed!\n");
    }

    // uniform blocks
    int ub_slot = 0;
    ShaderResources res = compiler->get_shader_resources();
    cJSON* ub_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "uniform_blocks", ub_array);
    for (const Resource& ub_res : res.uniform_buffers) {
        const SPIRType& ub_type = compiler->get_type(ub_res.base_type_id);
        string ub_name = compiler->get_name(ub_res.id);
        if (ub_name.empty()) {
            char buf[64];
            snprintf(buf, sizeof(buf), "_%d", ub_res.id);
            ub_name = buf;
        }
        cJSON* ub = cJSON_CreateObject();
        cJSON_AddItemToArray(ub_array, ub);
        cJSON_AddItemToObject(ub, "type", cJSON_CreateString(ub_res.name.c_str()));
        cJSON_AddItemToObject(ub, "name", cJSON_CreateString(ub_name.c_str()));
        cJSON_AddItemToObject(ub, "slot", cJSON_CreateNumber(ub_slot++));
        cJSON* ub_members = cJSON_CreateArray();
        cJSON_AddItemToObject(ub, "members", ub_members);
        for (int m_index = 0; m_index < int(ub_type.member_types.size()); m_index++) {
            cJSON* ub_member = cJSON_CreateObject();
            cJSON_AddItemToArray(ub_members, ub_member);
            string m_name = compiler->get_member_name(ub_res.base_type_id, m_index);
            const SPIRType& m_type = compiler->get_type(ub_type.member_types[m_index]);
            const char* m_type_str = type_to_uniform_type(m_type);
            cJSON_AddItemToObject(ub_member, "name", cJSON_CreateString(m_name.c_str()));
            cJSON_AddItemToObject(ub_member, "type", cJSON_CreateString(m_type_str));
            int num = 1;
            if (m_type.array.size() > 0) {
                num = m_type.array[0];
            }
            cJSON_AddItemToObject(ub_member, "num", cJSON_CreateNumber(num));
        }
    }

    // textures
    int tex_slot = 0;
    cJSON* tex_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "textures", tex_array);
    for (const Resource& img_res : res.sampled_images) {
        const SPIRType& img_type = compiler->get_type(img_res.type_id);
        cJSON* tex = cJSON_CreateObject();
        cJSON_AddItemToArray(tex_array, tex);
        const char* tex_type_str = nullptr;
        if (img_type.image.arrayed) {
            if (img_type.image.dim == Dim2D) {
                tex_type_str = "sampler2DArray";
            }
        }
        else {
            switch (img_type.image.dim) {
                case Dim2D:     tex_type_str = "sampler2D"; break;
                case DimCube:   tex_type_str = "samplerCube"; break;
                case Dim3D:     tex_type_str = "sampler3D"; break;
                default:        break;
            }
        }
        if (!tex_type_str) {
            Log::Fatal("Invalid texture type! (expected: 2D, Cube, 3D or 2D-array)\n");
        }
        cJSON_AddItemToObject(tex, "name", cJSON_CreateString(img_res.name.c_str()));
        cJSON_AddItemToObject(tex, "type", cJSON_CreateString(tex_type_str));
        cJSON_AddItemToObject(tex, "slot", cJSON_CreateNumber(tex_slot++));
    }

    // stage inputs
    cJSON* inputs_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "inputs", inputs_array);
    for (const Resource& stage_input : res.stage_inputs) {
        cJSON* input = cJSON_CreateObject();
        cJSON_AddItemToArray(inputs_array, input);
        const SPIRType& type = compiler->get_type(stage_input.base_type_id);
        const char* type_str = type_to_attr_format(type);
        cJSON_AddItemToObject(input, "name", cJSON_CreateString(stage_input.name.c_str()));
        cJSON_AddItemToObject(input, "type", cJSON_CreateString(type_str));
        if (compiler->get_execution_model() == ExecutionModelVertex) {
            uint32_t loc = compiler->get_decoration(stage_input.id, DecorationLocation);
            cJSON_AddItemToObject(input, "slot", cJSON_CreateNumber(loc));
        }
    }

    // stage outputs
    cJSON* outputs_array = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "outputs", outputs_array);
    for (const Resource& stage_output : res.stage_outputs) {
        cJSON* output = cJSON_CreateObject();
        cJSON_AddItemToArray(outputs_array, output);
        const SPIRType& type = compiler->get_type(stage_output.base_type_id);
        const char* type_str = type_to_attr_format(type);
        cJSON_AddItemToObject(output, "name", cJSON_CreateString(stage_output.name.c_str()));
        cJSON_AddItemToObject(output, "type", cJSON_CreateString(type_str));
    }

    return root;
}

//------------------------------------------------------------------------------
void fix_vertex_attr_locations(Compiler* compiler) {
    if (compiler->get_execution_model() == ExecutionModelVertex) {
        ShaderResources res = compiler->get_shader_resources();
        for (const auto& input : res.stage_inputs) {
            int loc = -1;
            if (input.name == "position")       loc = 0;
            else if (input.name == "normal")    loc = 1;
            else if (input.name == "texcoord0") loc = 2;
            else if (input.name == "texcoord1") loc = 3;
            else if (input.name == "texcoord2") loc = 4;
            else if (input.name == "texcoord3") loc = 5;
            else if (input.name == "tangent")   loc = 6;
            else if (input.name == "binormal")  loc = 7;
            else if (input.name == "weights")   loc = 8;
            else if (input.name == "indices")   loc = 9;
            else if (input.name == "color0")    loc = 10;
            else if (input.name == "color1")    loc = 11;
            else if (input.name == "instance0") loc = 12;
            else if (input.name == "instance1") loc = 13;
            else if (input.name == "instance2") loc = 14;
            else if (input.name == "instance3") loc = 15;
            if (-1 != loc) {
                compiler->set_decoration(input.id, DecorationLocation, (uint32_t)loc);
            }
        }
    }
}

//------------------------------------------------------------------------------
void fix_ub_matrix_force_colmajor(Compiler* compiler) {
    // go though all uniform block matrixes and decorate them with
    // column-major, this is needed in the HLSL backend to fix the
    // multiplication order
    ShaderResources res = compiler->get_shader_resources();
    for (const Resource& ub_res : res.uniform_buffers) {
        const SPIRType& ub_type = compiler->get_type(ub_res.base_type_id);
        for (int m_index = 0; m_index < int(ub_type.member_types.size()); m_index++) {
            const SPIRType& m_type = compiler->get_type(ub_type.member_types[m_index]);
            if ((m_type.basetype == SPIRType::Float) && (m_type.vecsize > 1) && (m_type.columns > 1)) {
                compiler->set_member_decoration(ub_res.base_type_id, m_index, DecorationColMajor);
            }
        }
    }
}

//------------------------------------------------------------------------------
void to_reflection_json(const vector<uint32_t>& spirv, const string& base_path) {
    CompilerGLSL compiler(spirv);
    fix_vertex_attr_locations(&compiler);
    cJSON* json = extract_resource_info(&compiler);
    char* json_raw_str = cJSON_Print(json);
    string json_str(json_raw_str);
    std::free(json_raw_str);
    cJSON_Delete(json);
    write_source_file(base_path + ".json", json_str);
}

//------------------------------------------------------------------------------
void to_glsl(const vector<uint32_t>& spirv, const string& base_path, const string& file_ext, uint32_t version, bool is_es) {
    CompilerGLSL compiler(spirv);
    auto opts = compiler.get_options();
    opts.version = version;
    opts.es = is_es;
    opts.vertex.fixup_clipspace = false;
    compiler.set_options(opts);
    fix_vertex_attr_locations(&compiler);
    fix_ub_matrix_force_colmajor(&compiler);
    string src = compiler.compile();
    if (src.empty()) {
        Log::Fatal("Failed to compile GLSL v100 source for '%s'!\n", base_path.c_str());
    }
    else {
        write_source_file(base_path + file_ext, src);
    }
}

//------------------------------------------------------------------------------
void to_hlsl_sm5(const vector<uint32_t>& spirv, const string& base_path) {
    CompilerHLSL compiler(spirv);
    auto opts = compiler.get_options();
    opts.shader_model = 50;
    opts.fixup_clipspace = true;
    compiler.set_options(opts);
    fix_vertex_attr_locations(&compiler);
    fix_ub_matrix_force_colmajor(&compiler);
    string src = compiler.compile();
    if (src.empty()) {
        Log::Fatal("Failed to compile HLSL5 source for '%s'!\n", base_path.c_str());
    }
    else {
        write_source_file(base_path + ".hlsl", src);
    }
}

//------------------------------------------------------------------------------
void to_mlsl(const vector<uint32_t>& spirv, const string& base_path) {
    CompilerMSL compiler(spirv);
    fix_vertex_attr_locations(&compiler);
    string src = compiler.compile();
    if (src.empty()) {
        Log::Fatal("Failed to compile MetalSL source for '%s'!\n", base_path.c_str());
    }
    else {
        write_source_file(base_path + ".metal", src);
    }
}

//------------------------------------------------------------------------------
int main(int argc, const char** argv) {
    CmdLineArgs args;
    args.AddBool("-help", "show help");
    args.AddString("-spirv", "SPIR-V input file", "");
    if (!args.Parse(argc, argv)) {
        Log::Warn("Failed to parse args!\n");
        return 10;
    }

    if (args.HasArg("-help")) {
        Log::Info("Oryol SPIR-V to GLSL/HLSL/MSL cross-compiler\n"
                  "Based on SPIRV-Cross: https://github.com/KhronosGroup/SPIRV-Cross\n");
        args.ShowHelp();
        return 0;
    }
    string spirv_path = args.GetString("-spirv");
    if (spirv_path.empty()) {
        Log::Fatal("-spirv arg expected");
    }

    // load SPIRV byte code
    auto spirv = read_spirv_file(spirv_path);

    // ...translate and write to output files...
    string base_path, ext;
    pystring::os::path::splitext(base_path, ext, spirv_path);
    to_reflection_json(spirv, base_path);
    to_glsl(spirv, base_path, ".glsl100.glsl", 100, true);
    to_glsl(spirv, base_path, ".glsl120.glsl", 120, false);
    to_glsl(spirv, base_path, ".glsles3.glsl", 300, true);
    to_glsl(spirv, base_path, ".glsl330.glsl", 330, false);
    to_hlsl_sm5(spirv, base_path);
    to_mlsl(spirv, base_path);

    return 0;
}
<|endoftext|>


<commit_before>/*
 +----------------------------------------------------------------------+
 | Swoole                                                               |
 +----------------------------------------------------------------------+
 | This source file is subject to version 2.0 of the Apache license,    |
 | that is bundled with this package in the file LICENSE, and is        |
 | available through the world-wide-web at the following url:           |
 | http://www.apache.org/licenses/LICENSE-2.0.html                      |
 | If you did not receive a copy of the Apache2.0 license and are unable|
 | to obtain it through the world-wide-web, please send a note to       |
 | license@swoole.com so we can mail you a copy immediately.            |
 +----------------------------------------------------------------------+
 | Author: Tianfeng Han  <mikan.tenny@gmail.com>                        |
 +----------------------------------------------------------------------+
 */

#include "swoole_api.h"
#include "async.h"

#include <thread>
#include <atomic>
#include <unordered_map>
#include <condition_variable>
#include <mutex>
#include <queue>

using namespace std;

typedef swAio_event async_event;

swAsyncIO SwooleAIO;

static void swAio_free(void *private_data);

int swAio_callback(swReactor *reactor, swEvent *_event)
{
    int i;
    async_event *events[SW_AIO_EVENT_NUM];
    ssize_t n = read(_event->fd, events, sizeof(async_event*) * SW_AIO_EVENT_NUM);
    if (n < 0)
    {
        swSysWarn("read() failed");
        return SW_ERR;
    }
    for (i = 0; i < n / (int) sizeof(async_event*); i++)
    {
        if (!events[i]->canceled)
        {
            events[i]->callback(events[i]);
        }
        SwooleAIO.task_num--;
        delete events[i];
    }
    return SW_OK;
}

class async_event_queue
{
public:
    inline bool push(async_event *event)
    {
        unique_lock<mutex> lock(_mutex);
        _queue.push(event);
        return true;
    }
    inline async_event* pop()
    {
        unique_lock<mutex> lock(_mutex);
        if (_queue.empty())
        {
            return nullptr;
        }
        async_event* retval = _queue.front();
        _queue.pop();
        return retval;
    }
    inline bool empty()
    {
        unique_lock<mutex> lock(_mutex);
        return _queue.empty();
    }
    inline size_t count()
    {
        return _queue.size();
    }
private:
    queue<async_event*> _queue;
    mutex _mutex;
};

class async_thread_pool
{
public:
    async_thread_pool(int _min_threads, int _max_threads)
    {
        n_waiting = 0;
        running = false;
        min_threads = _min_threads;
        max_threads = _max_threads;
        current_task_id = 0;
        current_pid = getpid();

        if (swPipeBase_create(&_aio_pipe, 0) < 0)
        {
            swoole_throw_error(SW_ERROR_SYSTEM_CALL_FAIL);
        }
        _pipe_read = _aio_pipe.getFd(&_aio_pipe, 0);
        _pipe_write = _aio_pipe.getFd(&_aio_pipe, 1);
        swoole_event_add(_pipe_read, SW_EVENT_READ, SW_FD_AIO);
    }

    ~async_thread_pool()
    {
        shutdown();
        if (SwooleTG.reactor)
        {
            swoole_event_del(_pipe_read);
        }
        _aio_pipe.close(&_aio_pipe);
    }

    void schedule()
    {
        int i = (int) threads.size();
        //++
        if (n_waiting == 0 && i < max_threads)
        {
            create_thread(i);
        }
        //--
        else if (n_waiting > min_threads)
        {
            i -= 1;
            *exit_flags[i] = true;
            threads[i]->detach();
            threads.erase(i);
            exit_flags.erase(i);
        }
    }

    bool start()
    {
        running = true;
        for (int i = 0; i < min_threads; i++)
        {
            create_thread(i);
        }
        return true;
    }

    bool shutdown()
    {
        if (!running)
        {
            return false;
        }
        running = false;

        _mutex.lock();
        _cv.notify_all();
        _mutex.unlock();

        for (int i = 0; i < static_cast<int>(threads.size()); ++i)
        {
            if (threads[i]->joinable())
            {
                threads[i]->join();
            }
        }

        threads.clear();
        exit_flags.clear();
        return true;
    }

    async_event* dispatch(const async_event *request)
    {
        auto _event_copy = new async_event(*request);
        schedule();
        _event_copy->task_id = current_task_id++;
        queue.push(_event_copy);
        _cv.notify_one();
        return _event_copy;
    }

    inline size_t thread_count()
    {
        return threads.size();
    }

    inline size_t queue_count()
    {
        return queue.count();
    }

    pid_t current_pid;

private:
    void create_thread(int i)
    {
        exit_flags[i] = make_shared<atomic<bool>>(false);
        shared_ptr<atomic<bool>> flag(exit_flags[i]);
        thread *_thread;

        try
        {
            _thread = new thread([this, flag]()
            {
                SwooleTG.buffer_stack = swString_new(SW_STACK_BUFFER_SIZE);
                if (SwooleTG.buffer_stack == nullptr)
                {
                    return;
                }

                swSignal_none();

                atomic<bool> &_flag = *flag;
                async_event *event;
                _accept:
                event = queue.pop();
                if (event)
                {
                    if (sw_unlikely(event->handler == nullptr))
                    {
                        event->error = SW_ERROR_AIO_BAD_REQUEST;
                        event->ret = -1;
                        goto _error;
                    }
                    else if (sw_unlikely(event->canceled))
                    {
                        event->error = SW_ERROR_AIO_BAD_REQUEST;
                        event->ret = -1;
                        goto _error;
                    }
                    else
                    {
                        event->handler(event);
                    }

                    swTrace("aio_thread ok. ret=%d, error=%d", event->ret, event->error);

                    _error:
                    while (true)
                    {
                        SwooleAIO.lock.lock(&SwooleAIO.lock);
                        int ret = write(_pipe_write, &event, sizeof(event));
                        SwooleAIO.lock.unlock(&SwooleAIO.lock);
                        if (ret < 0)
                        {
                            if (errno == EAGAIN)
                            {
                                swSocket_wait(_pipe_write, 1000, SW_EVENT_WRITE);
                                continue;
                            }
                            else if (errno == EINTR)
                            {
                                continue;
                            }
                            else
                            {
                                swSysWarn("sendto swoole_aio_pipe_write failed");
                            }
                        }
                        break;
                    }
                    //exit
                    if (_flag)
                    {
                        return;
                    }
                }
                else
                {
                    unique_lock<mutex> lock(_mutex);
                    if (running)
                    {
                        ++n_waiting;
                        _cv.wait(lock);
                        --n_waiting;
                    }
                }
                if (running)
                {
                    goto _accept;
                }
            });
        }
        catch (const std::system_error& e)
        {
            swSysNotice("create aio thread#%d failed, please check your system configuration or adjust max_thread_count", i);
            return;
        }

        threads[i] = unique_ptr<thread>(_thread);
    }

    swPipe _aio_pipe;
    int _pipe_read;
    int _pipe_write;
    int current_task_id;

    unordered_map<int, unique_ptr<thread>> threads;
    unordered_map<int, shared_ptr<atomic<bool>>> exit_flags;

    async_event_queue queue;
    bool running;
    atomic<int> n_waiting;
    int min_threads;
    int max_threads;
    mutex _mutex;
    condition_variable _cv;
};

static async_thread_pool *pool = nullptr;

static int swAio_init()
{
    if (SwooleAIO.init)
    {
        swWarn("AIO has already been initialized");
        return SW_ERR;
    }
    if (!SwooleTG.reactor)
    {
        swWarn("no event loop, cannot initialized");
        return SW_ERR;
    }

    if (swMutex_create(&SwooleAIO.lock, 0) < 0)
    {
        swWarn("create mutex lock error");
        return SW_ERR;
    }

    if (SwooleAIO.min_thread_count == 0)
    {
        SwooleAIO.min_thread_count = SW_AIO_THREAD_DEFAULT_NUM;
    }
    if (SwooleAIO.max_thread_count == 0)
    {
        SwooleAIO.max_thread_count = (SW_CPU_NUM * 2) * SW_AIO_THREAD_NUM_MULTIPLE;
    }
    if (SwooleAIO.min_thread_count > SwooleAIO.max_thread_count)
    {
        SwooleAIO.max_thread_count = SwooleAIO.min_thread_count;
    }

    swReactor_add_destroy_callback(SwooleTG.reactor, swAio_free, nullptr);

    pool = new async_thread_pool(SwooleAIO.min_thread_count, SwooleAIO.max_thread_count);
    pool->start();
    SwooleAIO.init = 1;

    return SW_OK;
}

size_t swAio_thread_count()
{
    return pool ? pool->thread_count() : 0;
}

int swAio_dispatch(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    async_event *event = pool->dispatch(request);
    return event->task_id;
}

swAio_event* swAio_dispatch2(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    return pool->dispatch(request);
}

static void swAio_free(void *private_data)
{
    if (!SwooleAIO.init)
    {
        return;
    }
    if (pool->current_pid == getpid())
    {
        delete pool;
    }
    pool = nullptr;
    SwooleAIO.init = 0;
}
<commit_msg>Remove shared ptr<commit_after>/*
 +----------------------------------------------------------------------+
 | Swoole                                                               |
 +----------------------------------------------------------------------+
 | This source file is subject to version 2.0 of the Apache license,    |
 | that is bundled with this package in the file LICENSE, and is        |
 | available through the world-wide-web at the following url:           |
 | http://www.apache.org/licenses/LICENSE-2.0.html                      |
 | If you did not receive a copy of the Apache2.0 license and are unable|
 | to obtain it through the world-wide-web, please send a note to       |
 | license@swoole.com so we can mail you a copy immediately.            |
 +----------------------------------------------------------------------+
 | Author: Tianfeng Han  <mikan.tenny@gmail.com>                        |
 +----------------------------------------------------------------------+
 */

#include "swoole_api.h"
#include "async.h"

#include <thread>
#include <atomic>
#include <unordered_map>
#include <condition_variable>
#include <mutex>
#include <queue>

using namespace std;

typedef swAio_event async_event;

swAsyncIO SwooleAIO;

static void swAio_free(void *private_data);

int swAio_callback(swReactor *reactor, swEvent *_event)
{
    int i;
    async_event *events[SW_AIO_EVENT_NUM];
    ssize_t n = read(_event->fd, events, sizeof(async_event*) * SW_AIO_EVENT_NUM);
    if (n < 0)
    {
        swSysWarn("read() failed");
        return SW_ERR;
    }
    for (i = 0; i < n / (int) sizeof(async_event*); i++)
    {
        if (!events[i]->canceled)
        {
            events[i]->callback(events[i]);
        }
        SwooleAIO.task_num--;
        delete events[i];
    }
    return SW_OK;
}

class async_event_queue
{
public:
    inline bool push(async_event *event)
    {
        unique_lock<mutex> lock(_mutex);
        _queue.push(event);
        return true;
    }
    inline async_event* pop()
    {
        unique_lock<mutex> lock(_mutex);
        if (_queue.empty())
        {
            return nullptr;
        }
        async_event* retval = _queue.front();
        _queue.pop();
        return retval;
    }
    inline bool empty()
    {
        unique_lock<mutex> lock(_mutex);
        return _queue.empty();
    }
    inline size_t count()
    {
        return _queue.size();
    }
private:
    queue<async_event*> _queue;
    mutex _mutex;
};

class async_thread_pool
{
public:
    async_thread_pool(int _min_threads, int _max_threads)
    {
        n_waiting = 0;
        running = false;
        min_threads = _min_threads;
        max_threads = _max_threads;
        current_task_id = 0;
        current_pid = getpid();

        if (swPipeBase_create(&_aio_pipe, 0) < 0)
        {
            swoole_throw_error(SW_ERROR_SYSTEM_CALL_FAIL);
        }
        _pipe_read = _aio_pipe.getFd(&_aio_pipe, 0);
        _pipe_write = _aio_pipe.getFd(&_aio_pipe, 1);
        swoole_event_add(_pipe_read, SW_EVENT_READ, SW_FD_AIO);
    }

    ~async_thread_pool()
    {
        shutdown();
        if (SwooleTG.reactor)
        {
            swoole_event_del(_pipe_read);
        }
        _aio_pipe.close(&_aio_pipe);
    }

    void schedule()
    {
        int i = (int) threads.size();
        //++
        if (n_waiting == 0 && i < max_threads)
        {
            create_thread(i);
        }
        //--
        else if (n_waiting > min_threads)
        {
            i -= 1;
            exit_flags[i] = true;
            threads[i]->detach();
            threads.erase(i);
            exit_flags.erase(i);
        }
    }

    bool start()
    {
        running = true;
        for (int i = 0; i < min_threads; i++)
        {
            create_thread(i);
        }
        return true;
    }

    bool shutdown()
    {
        if (!running)
        {
            return false;
        }
        running = false;

        _mutex.lock();
        _cv.notify_all();
        _mutex.unlock();

        for (int i = 0; i < static_cast<int>(threads.size()); ++i)
        {
            if (threads[i]->joinable())
            {
                threads[i]->join();
            }
        }

        threads.clear();
        exit_flags.clear();
        return true;
    }

    async_event* dispatch(const async_event *request)
    {
        auto _event_copy = new async_event(*request);
        schedule();
        _event_copy->task_id = current_task_id++;
        queue.push(_event_copy);
        _cv.notify_one();
        return _event_copy;
    }

    inline size_t thread_count()
    {
        return threads.size();
    }

    inline size_t queue_count()
    {
        return queue.count();
    }

    pid_t current_pid;

private:
    void create_thread(int i)
    {
        exit_flags[i] = false;
        atomic<bool> &exit_flag = exit_flags[i];
        thread *_thread;

        try
        {
            _thread = new thread([this, &exit_flag]()
            {
                SwooleTG.buffer_stack = swString_new(SW_STACK_BUFFER_SIZE);
                if (SwooleTG.buffer_stack == nullptr)
                {
                    return;
                }

                swSignal_none();

                async_event *event;
                _accept:
                event = queue.pop();
                if (event)
                {
                    if (sw_unlikely(event->handler == nullptr))
                    {
                        event->error = SW_ERROR_AIO_BAD_REQUEST;
                        event->ret = -1;
                        goto _error;
                    }
                    else if (sw_unlikely(event->canceled))
                    {
                        event->error = SW_ERROR_AIO_BAD_REQUEST;
                        event->ret = -1;
                        goto _error;
                    }
                    else
                    {
                        event->handler(event);
                    }

                    swTrace("aio_thread ok. ret=%d, error=%d", event->ret, event->error);

                    _error:
                    while (true)
                    {
                        SwooleAIO.lock.lock(&SwooleAIO.lock);
                        int ret = write(_pipe_write, &event, sizeof(event));
                        SwooleAIO.lock.unlock(&SwooleAIO.lock);
                        if (ret < 0)
                        {
                            if (errno == EAGAIN)
                            {
                                swSocket_wait(_pipe_write, 1000, SW_EVENT_WRITE);
                                continue;
                            }
                            else if (errno == EINTR)
                            {
                                continue;
                            }
                            else
                            {
                                swSysWarn("sendto swoole_aio_pipe_write failed");
                            }
                        }
                        break;
                    }
                    //exit
                    if (exit_flag)
                    {
                        return;
                    }
                }
                else
                {
                    unique_lock<mutex> lock(_mutex);
                    if (running)
                    {
                        ++n_waiting;
                        _cv.wait(lock);
                        --n_waiting;
                    }
                }
                if (running)
                {
                    goto _accept;
                }
            });
        }
        catch (const std::system_error& e)
        {
            swSysNotice("create aio thread#%d failed, please check your system configuration or adjust max_thread_count", i);
            return;
        }

        threads[i] = unique_ptr<thread>(_thread);
    }

    swPipe _aio_pipe;
    int _pipe_read;
    int _pipe_write;
    int current_task_id;

    unordered_map<int, unique_ptr<thread>> threads;
    unordered_map<int, atomic<bool>> exit_flags;

    async_event_queue queue;
    bool running;
    atomic<int> n_waiting;
    int min_threads;
    int max_threads;
    mutex _mutex;
    condition_variable _cv;
};

static async_thread_pool *pool = nullptr;

static int swAio_init()
{
    if (SwooleAIO.init)
    {
        swWarn("AIO has already been initialized");
        return SW_ERR;
    }
    if (!SwooleTG.reactor)
    {
        swWarn("no event loop, cannot initialized");
        return SW_ERR;
    }

    if (swMutex_create(&SwooleAIO.lock, 0) < 0)
    {
        swWarn("create mutex lock error");
        return SW_ERR;
    }

    if (SwooleAIO.min_thread_count == 0)
    {
        SwooleAIO.min_thread_count = SW_AIO_THREAD_DEFAULT_NUM;
    }
    if (SwooleAIO.max_thread_count == 0)
    {
        SwooleAIO.max_thread_count = (SW_CPU_NUM * 2) * SW_AIO_THREAD_NUM_MULTIPLE;
    }
    if (SwooleAIO.min_thread_count > SwooleAIO.max_thread_count)
    {
        SwooleAIO.max_thread_count = SwooleAIO.min_thread_count;
    }

    swReactor_add_destroy_callback(SwooleTG.reactor, swAio_free, nullptr);

    pool = new async_thread_pool(SwooleAIO.min_thread_count, SwooleAIO.max_thread_count);
    pool->start();
    SwooleAIO.init = 1;

    return SW_OK;
}

size_t swAio_thread_count()
{
    return pool ? pool->thread_count() : 0;
}

int swAio_dispatch(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    async_event *event = pool->dispatch(request);
    return event->task_id;
}

swAio_event* swAio_dispatch2(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    return pool->dispatch(request);
}

static void swAio_free(void *private_data)
{
    if (!SwooleAIO.init)
    {
        return;
    }
    if (pool->current_pid == getpid())
    {
        delete pool;
    }
    pool = nullptr;
    SwooleAIO.init = 0;
}
<|endoftext|>


<commit_before>/*
 +----------------------------------------------------------------------+
 | Swoole                                                               |
 +----------------------------------------------------------------------+
 | This source file is subject to version 2.0 of the Apache license,    |
 | that is bundled with this package in the file LICENSE, and is        |
 | available through the world-wide-web at the following url:           |
 | http://www.apache.org/licenses/LICENSE-2.0.html                      |
 | If you did not receive a copy of the Apache2.0 license and are unable|
 | to obtain it through the world-wide-web, please send a note to       |
 | license@swoole.com so we can mail you a copy immediately.            |
 +----------------------------------------------------------------------+
 | Author: Tianfeng Han  <mikan.tenny@gmail.com>                        |
 +----------------------------------------------------------------------+
 */

#include "swoole_api.h"
#include "async.h"

#include <thread>
#include <atomic>
#include <unordered_map>
#include <condition_variable>
#include <mutex>
#include <queue>

using namespace std;

typedef swAio_event async_event;

swAsyncIO SwooleAIO;

static void swAio_free(void *private_data);

int swAio_callback(swReactor *reactor, swEvent *_event)
{
    async_event *events[SW_AIO_EVENT_NUM];
    ssize_t n = read(_event->fd, events, sizeof(async_event*) * SW_AIO_EVENT_NUM);
    if (n < 0)
    {
        swSysWarn("read() failed");
        return SW_ERR;
    }
    for (size_t i = 0; i < n / sizeof(async_event *); i++)
    {
        if (!events[i]->canceled)
        {
            events[i]->callback(events[i]);
        }
        SwooleAIO.task_num--;
        delete events[i];
    }
    return SW_OK;
}

struct thread_context
{
    thread *_thread;
    atomic<bool> *_exit_flag;
    thread_context(thread *_thread, atomic<bool> *_exit_flag) : _thread(_thread), _exit_flag(_exit_flag) { }
};

class async_event_queue
{
public:
    inline bool push(async_event *event)
    {
        unique_lock<mutex> lock(_mutex);
        _queue.push(event);
        return true;
    }
    inline async_event* pop()
    {
        unique_lock<mutex> lock(_mutex);
        if (_queue.empty())
        {
            return nullptr;
        }
        async_event* retval = _queue.front();
        _queue.pop();
        return retval;
    }
    inline bool empty()
    {
        unique_lock<mutex> lock(_mutex);
        return _queue.empty();
    }
    inline size_t count()
    {
        return _queue.size();
    }
private:
    queue<async_event*> _queue;
    mutex _mutex;
};

class async_thread_pool
{
public:
    async_thread_pool(size_t _min_threads, size_t _max_threads)
    {
        n_waiting = 0;
        running = false;
        min_threads = SW_MAX(SW_AIO_THREAD_DEFAULT_NUM, _min_threads);
        max_threads = SW_MAX(min_threads, _max_threads);
        current_task_id = 0;
        current_pid = getpid();

        if (swPipeBase_create(&_aio_pipe, 0) < 0)
        {
            swoole_throw_error(SW_ERROR_SYSTEM_CALL_FAIL);
        }
        _pipe_read = _aio_pipe.getFd(&_aio_pipe, 0);
        _pipe_write = _aio_pipe.getFd(&_aio_pipe, 1);
        swoole_event_add(_pipe_read, SW_EVENT_READ, SW_FD_AIO);
    }

    ~async_thread_pool()
    {
        shutdown();
        if (SwooleTG.reactor)
        {
            swoole_event_del(_pipe_read);
        }
        _aio_pipe.close(&_aio_pipe);
    }

    void schedule()
    {
        //++
        if (n_waiting == 0 && threads.size() < max_threads)
        {
            create_thread();
        }
        //--
        else if (n_waiting > min_threads)
        {
            thread_context *tc = &threads.front();
            *tc->_exit_flag = false;
            tc->_thread->detach();
            delete tc->_thread;
            threads.pop();
        }
    }

    bool start()
    {
        running = true;
        for (size_t i = 0; i < min_threads; i++)
        {
            create_thread();
        }
        return true;
    }

    bool shutdown()
    {
        if (!running)
        {
            return false;
        }
        running = false;

        _mutex.lock();
        _cv.notify_all();
        _mutex.unlock();

        while (!threads.empty())
        {
            thread_context *tc = &threads.front();
            if (tc->_thread->joinable())
            {
                tc->_thread->join();
            }
            threads.pop();
        }

        return true;
    }

    async_event* dispatch(const async_event *request)
    {
        auto _event_copy = new async_event(*request);
        schedule();
        _event_copy->task_id = current_task_id++;
        _queue.push(_event_copy);
        _cv.notify_one();
        return _event_copy;
    }

    inline size_t thread_count()
    {
        return threads.size();
    }

    inline size_t queue_count()
    {
        return _queue.count();
    }

    pid_t current_pid;

private:
    void create_thread()
    {
        atomic<bool> *_exit_flag = new atomic<bool>(false);
        try
        {
            thread *_thread = new thread([this, _exit_flag]()
            {
                SwooleTG.buffer_stack = swString_new(SW_STACK_BUFFER_SIZE);
                if (SwooleTG.buffer_stack == nullptr)
                {
                    return;
                }

                swSignal_none();

                while (running)
                {
                    async_event *event = _queue.pop();
                    if (event)
                    {
                        if (sw_unlikely(event->handler == nullptr))
                        {
                            event->error = SW_ERROR_AIO_BAD_REQUEST;
                            event->ret = -1;
                            goto _error;
                        }
                        else if (sw_unlikely(event->canceled))
                        {
                            event->error = SW_ERROR_AIO_BAD_REQUEST;
                            event->ret = -1;
                            goto _error;
                        }
                        else
                        {
                            event->handler(event);
                        }

                        swTrace("aio_thread ok. ret=%d, error=%d", event->ret, event->error);

                        _error:
                        while (true)
                        {
                            SwooleAIO.lock.lock(&SwooleAIO.lock);
                            int ret = write(_pipe_write, &event, sizeof(event));
                            SwooleAIO.lock.unlock(&SwooleAIO.lock);
                            if (ret < 0)
                            {
                                if (errno == EAGAIN)
                                {
                                    swSocket_wait(_pipe_write, 1000, SW_EVENT_WRITE);
                                    continue;
                                }
                                else if (errno == EINTR)
                                {
                                    continue;
                                }
                                else
                                {
                                    swSysWarn("sendto swoole_aio_pipe_write failed");
                                }
                            }
                            break;
                        }

                        // exit
                        if (*_exit_flag)
                        {
                            break;
                        }
                    }
                    else
                    {
                        unique_lock<mutex> lock(_mutex);
                        if (running)
                        {
                            ++n_waiting;
                            _cv.wait(lock);
                            --n_waiting;
                        }
                    }
                }

                delete _exit_flag;
            });
            threads.push(thread_context(_thread, _exit_flag));
        }
        catch (const std::system_error& e)
        {
            swSysNotice("create aio thread failed, please check your system configuration or adjust max_thread_count");
            delete _exit_flag;
            return;
        }
    }

    size_t min_threads;
    size_t max_threads;

    swPipe _aio_pipe;
    int _pipe_read;
    int _pipe_write;
    int current_task_id;

    queue<thread_context> threads;
    async_event_queue _queue;
    bool running;
    atomic<size_t> n_waiting;
    mutex _mutex;
    condition_variable _cv;
};

static async_thread_pool *pool = nullptr;

static int swAio_init()
{
    if (SwooleAIO.init)
    {
        swWarn("AIO has already been initialized");
        return SW_ERR;
    }
    if (!SwooleTG.reactor)
    {
        swWarn("no event loop, cannot initialized");
        return SW_ERR;
    }

    if (swMutex_create(&SwooleAIO.lock, 0) < 0)
    {
        swWarn("create mutex lock error");
        return SW_ERR;
    }

    if (SwooleAIO.min_thread_num == 0)
    {
        SwooleAIO.min_thread_num = SW_AIO_THREAD_DEFAULT_NUM;
    }
    if (SwooleAIO.max_thread_num == 0)
    {
        SwooleAIO.max_thread_num = (SW_CPU_NUM * 2) * SW_AIO_THREAD_NUM_MULTIPLE;
    }
    if (SwooleAIO.min_thread_num > SwooleAIO.max_thread_num)
    {
        SwooleAIO.max_thread_num = SwooleAIO.min_thread_num;
    }

    swReactor_add_destroy_callback(SwooleTG.reactor, swAio_free, nullptr);

    pool = new async_thread_pool(SwooleAIO.min_thread_num, SwooleAIO.max_thread_num);
    pool->start();
    SwooleAIO.init = 1;

    return SW_OK;
}

size_t swAio_thread_count()
{
    return pool ? pool->thread_count() : 0;
}

int swAio_dispatch(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    async_event *event = pool->dispatch(request);
    return event->task_id;
}

swAio_event* swAio_dispatch2(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    return pool->dispatch(request);
}

static void swAio_free(void *private_data)
{
    if (!SwooleAIO.init)
    {
        return;
    }
    if (pool->current_pid == getpid())
    {
        delete pool;
    }
    pool = nullptr;
    SwooleAIO.init = 0;
}
<commit_msg>Fix warning<commit_after>/*
 +----------------------------------------------------------------------+
 | Swoole                                                               |
 +----------------------------------------------------------------------+
 | This source file is subject to version 2.0 of the Apache license,    |
 | that is bundled with this package in the file LICENSE, and is        |
 | available through the world-wide-web at the following url:           |
 | http://www.apache.org/licenses/LICENSE-2.0.html                      |
 | If you did not receive a copy of the Apache2.0 license and are unable|
 | to obtain it through the world-wide-web, please send a note to       |
 | license@swoole.com so we can mail you a copy immediately.            |
 +----------------------------------------------------------------------+
 | Author: Tianfeng Han  <mikan.tenny@gmail.com>                        |
 +----------------------------------------------------------------------+
 */

#include "swoole_api.h"
#include "async.h"

#include <thread>
#include <atomic>
#include <unordered_map>
#include <condition_variable>
#include <mutex>
#include <queue>

using namespace std;

typedef swAio_event async_event;

swAsyncIO SwooleAIO;

static void swAio_free(void *private_data);

int swAio_callback(swReactor *reactor, swEvent *_event)
{
    async_event *events[SW_AIO_EVENT_NUM];
    ssize_t n = read(_event->fd, events, sizeof(async_event*) * SW_AIO_EVENT_NUM);
    if (n < 0)
    {
        swSysWarn("read() failed");
        return SW_ERR;
    }
    for (size_t i = 0; i < n / sizeof(async_event *); i++)
    {
        if (!events[i]->canceled)
        {
            events[i]->callback(events[i]);
        }
        SwooleAIO.task_num--;
        delete events[i];
    }
    return SW_OK;
}

struct thread_context
{
    thread *_thread;
    atomic<bool> *_exit_flag;
    thread_context(thread *_thread, atomic<bool> *_exit_flag) : _thread(_thread), _exit_flag(_exit_flag) { }
};

class async_event_queue
{
public:
    inline bool push(async_event *event)
    {
        unique_lock<mutex> lock(_mutex);
        _queue.push(event);
        return true;
    }
    inline async_event* pop()
    {
        unique_lock<mutex> lock(_mutex);
        if (_queue.empty())
        {
            return nullptr;
        }
        async_event* retval = _queue.front();
        _queue.pop();
        return retval;
    }
    inline bool empty()
    {
        unique_lock<mutex> lock(_mutex);
        return _queue.empty();
    }
    inline size_t count()
    {
        return _queue.size();
    }
private:
    queue<async_event*> _queue;
    mutex _mutex;
};

class async_thread_pool
{
public:
    async_thread_pool(size_t _min_threads, size_t _max_threads)
    {
        n_waiting = 0;
        running = false;
        min_threads = SW_MAX(SW_AIO_THREAD_DEFAULT_NUM, _min_threads);
        max_threads = SW_MAX(min_threads, _max_threads);
        current_task_id = 0;
        current_pid = getpid();

        if (swPipeBase_create(&_aio_pipe, 0) < 0)
        {
            swoole_throw_error(SW_ERROR_SYSTEM_CALL_FAIL);
        }
        _pipe_read = _aio_pipe.getFd(&_aio_pipe, 0);
        _pipe_write = _aio_pipe.getFd(&_aio_pipe, 1);
        swoole_event_add(_pipe_read, SW_EVENT_READ, SW_FD_AIO);
    }

    ~async_thread_pool()
    {
        shutdown();
        if (SwooleTG.reactor)
        {
            swoole_event_del(_pipe_read);
        }
        _aio_pipe.close(&_aio_pipe);
    }

    void schedule()
    {
        //++
        if (n_waiting == 0 && threads.size() < max_threads)
        {
            create_thread();
        }
        //--
        else if (n_waiting > min_threads && threads.size() > min_threads)
        {
            thread_context *tc = &threads.front();
            *tc->_exit_flag = false;
            tc->_thread->detach();
            delete tc->_thread;
            threads.pop();
        }
    }

    bool start()
    {
        running = true;
        for (size_t i = 0; i < min_threads; i++)
        {
            create_thread();
        }
        return true;
    }

    bool shutdown()
    {
        if (!running)
        {
            return false;
        }
        running = false;

        _mutex.lock();
        _cv.notify_all();
        _mutex.unlock();

        while (!threads.empty())
        {
            thread_context *tc = &threads.front();
            if (tc->_thread->joinable())
            {
                tc->_thread->join();
            }
            threads.pop();
        }

        return true;
    }

    async_event* dispatch(const async_event *request)
    {
        auto _event_copy = new async_event(*request);
        schedule();
        _event_copy->task_id = current_task_id++;
        _queue.push(_event_copy);
        _cv.notify_one();
        return _event_copy;
    }

    inline size_t thread_count()
    {
        return threads.size();
    }

    inline size_t queue_count()
    {
        return _queue.count();
    }

    pid_t current_pid;

private:
    void create_thread()
    {
        atomic<bool> *_exit_flag = new atomic<bool>(false);
        try
        {
            thread *_thread = new thread([this, _exit_flag]()
            {
                SwooleTG.buffer_stack = swString_new(SW_STACK_BUFFER_SIZE);
                if (SwooleTG.buffer_stack == nullptr)
                {
                    return;
                }

                swSignal_none();

                while (running)
                {
                    async_event *event = _queue.pop();
                    if (event)
                    {
                        if (sw_unlikely(event->handler == nullptr))
                        {
                            event->error = SW_ERROR_AIO_BAD_REQUEST;
                            event->ret = -1;
                            goto _error;
                        }
                        else if (sw_unlikely(event->canceled))
                        {
                            event->error = SW_ERROR_AIO_BAD_REQUEST;
                            event->ret = -1;
                            goto _error;
                        }
                        else
                        {
                            event->handler(event);
                        }

                        swTrace("aio_thread ok. ret=%d, error=%d", event->ret, event->error);

                        _error:
                        while (true)
                        {
                            SwooleAIO.lock.lock(&SwooleAIO.lock);
                            int ret = write(_pipe_write, &event, sizeof(event));
                            SwooleAIO.lock.unlock(&SwooleAIO.lock);
                            if (ret < 0)
                            {
                                if (errno == EAGAIN)
                                {
                                    swSocket_wait(_pipe_write, 1000, SW_EVENT_WRITE);
                                    continue;
                                }
                                else if (errno == EINTR)
                                {
                                    continue;
                                }
                                else
                                {
                                    swSysWarn("sendto swoole_aio_pipe_write failed");
                                }
                            }
                            break;
                        }

                        // exit
                        if (*_exit_flag)
                        {
                            break;
                        }
                    }
                    else
                    {
                        unique_lock<mutex> lock(_mutex);
                        if (running)
                        {
                            ++n_waiting;
                            _cv.wait(lock);
                            --n_waiting;
                        }
                    }
                }

                delete _exit_flag;
            });
            threads.push(thread_context(_thread, _exit_flag));
        }
        catch (const std::system_error& e)
        {
            swSysNotice("create aio thread failed, please check your system configuration or adjust max_thread_count");
            delete _exit_flag;
            return;
        }
    }

    size_t min_threads;
    size_t max_threads;

    swPipe _aio_pipe;
    int _pipe_read;
    int _pipe_write;
    int current_task_id;

    queue<thread_context> threads;
    async_event_queue _queue;
    bool running;
    atomic<size_t> n_waiting;
    mutex _mutex;
    condition_variable _cv;
};

static async_thread_pool *pool = nullptr;

static int swAio_init()
{
    if (SwooleAIO.init)
    {
        swWarn("AIO has already been initialized");
        return SW_ERR;
    }
    if (!SwooleTG.reactor)
    {
        swWarn("no event loop, cannot initialized");
        return SW_ERR;
    }

    if (swMutex_create(&SwooleAIO.lock, 0) < 0)
    {
        swWarn("create mutex lock error");
        return SW_ERR;
    }

    if (SwooleAIO.min_thread_num == 0)
    {
        SwooleAIO.min_thread_num = SW_AIO_THREAD_DEFAULT_NUM;
    }
    if (SwooleAIO.max_thread_num == 0)
    {
        SwooleAIO.max_thread_num = (SW_CPU_NUM * 2) * SW_AIO_THREAD_NUM_MULTIPLE;
    }
    if (SwooleAIO.min_thread_num > SwooleAIO.max_thread_num)
    {
        SwooleAIO.max_thread_num = SwooleAIO.min_thread_num;
    }

    swReactor_add_destroy_callback(SwooleTG.reactor, swAio_free, nullptr);

    pool = new async_thread_pool(SwooleAIO.min_thread_num, SwooleAIO.max_thread_num);
    pool->start();
    SwooleAIO.init = 1;

    return SW_OK;
}

size_t swAio_thread_count()
{
    return pool ? pool->thread_count() : 0;
}

int swAio_dispatch(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    async_event *event = pool->dispatch(request);
    return event->task_id;
}

swAio_event* swAio_dispatch2(const swAio_event *request)
{
    if (sw_unlikely(!SwooleAIO.init))
    {
        swAio_init();
    }
    SwooleAIO.task_num++;
    return pool->dispatch(request);
}

static void swAio_free(void *private_data)
{
    if (!SwooleAIO.init)
    {
        return;
    }
    if (pool->current_pid == getpid())
    {
        delete pool;
    }
    pool = nullptr;
    SwooleAIO.init = 0;
}
<|endoftext|>