<commit_before>/*
 * Copyright 2018 WebAssembly Community Group participants
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//
// Optimize using the DataFlow SSA IR.
//
// This needs 'flatten' to be run before it, and you should run full
// regular opts afterwards to clean up the flattening. For example,
// you might use it like this:
//
//    --flatten --dfo -Os
//

#include "dataflow/graph.h"
#include "dataflow/node.h"
#include "dataflow/users.h"
#include "dataflow/utils.h"
#include "ir/flat.h"
#include "ir/utils.h"
#include "pass.h"
#include "wasm-builder.h"
#include "wasm.h"

namespace wasm {

struct DataFlowOpts : public WalkerPass<PostWalker<DataFlowOpts>> {
  bool isFunctionParallel() override { return true; }

  Pass* create() override { return new DataFlowOpts; }

  DataFlow::Users nodeUsers;

  // The optimization work left to do: nodes that we need to look at.
  std::unordered_set<DataFlow::Node*> workLeft;

  DataFlow::Graph graph;

  void doWalkFunction(Function* func) {
    Flat::verifyFlatness(func);
    // Build the data-flow IR.
    graph.build(func, getModule());
    nodeUsers.build(graph);
    // Propagate optimizations through the graph.
    std::unordered_set<DataFlow::Node*> optimized; // which nodes we optimized
    for (auto& node : graph.nodes) {
      workLeft.insert(node.get()); // we should try to optimize each node
    }
    while (!workLeft.empty()) {
      // std::cout << "\n\ndump before work iter\n";
      // dump(graph, std::cout);
      auto iter = workLeft.begin();
      auto* node = *iter;
      workLeft.erase(iter);
      workOn(node);
    }
    // After updating the DataFlow IR, we can update the sets in
    // the wasm.
    // TODO: we also need phis, as a phi can flow directly into say
    //       a return or a call parameter.
    for (auto* set : graph.sets) {
      auto* node = graph.setNodeMap[set];
      auto iter = optimized.find(node);
      if (iter != optimized.end()) {
        assert(node->isExpr()); // this is a set, where the node is defined
        set->value = node->expr;
      }
    }
  }

  void workOn(DataFlow::Node* node) {
    if (node->isConst()) {
      return;
    }
    // If there are no uses, there is no point to work.
    if (nodeUsers.getNumUses(node) == 0) {
      return;
    }
    // Optimize: Look for nodes that we can easily convert into
    // something simpler.
    // TODO: we can expressionify and run full normal opts on that,
    //       then copy the result if it's smaller.
    if (node->isPhi() && DataFlow::allInputsIdentical(node)) {
      // Note we don't need to check for effects when replacing, as in
      // flattened IR expression children are local.gets or consts.
      auto* value = node->getValue(1);
      if (value->isConst()) {
        replaceAllUsesWith(node, value);
      }
    } else if (node->isExpr() && DataFlow::allInputsConstant(node)) {
      assert(!node->isConst());
      // If this is a concrete value (not e.g. an eqz of unreachable),
      // it can definitely be precomputed into a constant.
      if (node->expr->type.isConcrete()) {
        // This can be precomputed.
        // TODO not just all-constant inputs? E.g. i32.mul of 0 and X.
        optimizeExprToConstant(node);
      }
    }
  }

  void optimizeExprToConstant(DataFlow::Node* node) {
    assert(node->isExpr());
    assert(!node->isConst());
    // std::cout << "will optimize an Expr of all constant inputs. before" <<
    //              '\n';
    // dump(node, std::cout);
    auto* expr = node->expr;
    // First, note that some of the expression's children may be
    // local.gets that we inferred during SSA analysis as constant.
    // We can apply those now.
    for (Index i = 0; i < node->values.size(); i++) {
      if (node->values[i]->isConst()) {
        auto* currp = getIndexPointer(expr, i);
        // Directly represent it as a constant. (Note that it may already be
        // a constant, but for now to avoid corner cases just replace them
        // all here.)
        auto* c = node->values[i]->expr->dynCast<Const>();
        *currp = Builder(*getModule()).makeConst(c->value);
      }
    }
    // Now we know that all our DataFlow inputs are constant, and all
    // our Binaryen IR representations of them are constant too. RUn
    // precompute, which will transform the expression into a constanat.
    Module temp;
    // XXX we should copy expr here, in principle, and definitely will need to
    //     when we do arbitrarily regenerated expressions
    auto* func = Builder(temp).makeFunction(
      "temp", Signature(Type::none, Type::none), {}, expr);
    PassRunner runner(&temp);
    runner.setIsNested(true);
    runner.add("precompute");
    runner.runOnFunction(func);
    // Get the optimized thing
    auto* result = func->body;
    // It may not be a constant, e.g. 0 / 0 does not optimize to 0
    if (!result->is<Const>()) {
      return;
    }
    // All good, copy it.
    node->expr = Builder(*getModule()).makeConst(result->cast<Const>()->value);
    assert(node->isConst());
    // We no longer have values, and so do not use anything.
    nodeUsers.stopUsingValues(node);
    node->values.clear();
    // Our contents changed, update our users.
    replaceAllUsesWith(node, node);
  }

  // Replaces all uses of a node with another value. This both modifies
  // the DataFlow IR to make the other users point to this one, and
  // updates the underlying Binaryen IR as well.
  // This can be used to "replace" a node with itself, which makes sense
  // when the node contents have changed and so the users must be updated.
  void replaceAllUsesWith(DataFlow::Node* node, DataFlow::Node* with) {
    // Const nodes are trivial to replace, but other stuff is trickier -
    // in particular phis.
    assert(with->isConst()); // TODO
    // All the users should be worked on later, as we will update them.
    auto& users = nodeUsers.getUsers(node);
    for (auto* user : users) {
      // Add the user to the work left to do, as we are modifying it.
      workLeft.insert(user);
      // `with` is getting another user.
      nodeUsers.addUser(with, user);
      // Replacing in the DataFlow IR is simple - just replace it,
      // in all the indexes it appears.
      std::vector<Index> indexes;
      for (Index i = 0; i < user->values.size(); i++) {
        if (user->values[i] == node) {
          user->values[i] = with;
          indexes.push_back(i);
        }
      }
      assert(!indexes.empty());
      // Replacing in the Binaryen IR requires more care
      switch (user->type) {
        case DataFlow::Node::Type::Expr: {
          auto* expr = user->expr;
          for (auto index : indexes) {
            *(getIndexPointer(expr, index)) = graph.makeUse(with);
          }
          break;
        }
        case DataFlow::Node::Type::Phi: {
          // Nothing to do: a phi is not in the Binaryen IR.
          // If the entire phi can become a constant, that will be
          // propagated when we process that phi later.
          break;
        }
        case DataFlow::Node::Type::Cond: {
          // Nothing to do: a cond is not in the Binaryen IR.
          // If the cond input is a constant, that might indicate
          // useful optimizations are possible, which perhaps we
          // should look into TODO
          break;
        }
        case DataFlow::Node::Type::Zext: {
          // Nothing to do: a zext is not in the Binaryen IR.
          // If the cond input is a constant, that might indicate
          // useful optimizations are possible, which perhaps we
          // should look into TODO
          break;
        }
        default:
          WASM_UNREACHABLE("unexpected dataflow node type");
      }
    }
    // No one is a user of this node after we replaced all the uses.
    nodeUsers.removeAllUsesOf(node);
  }

  // Gets a pointer to the expression pointer in an expression.
  // That is, given an index in the values() vector, get an
  // Expression** that we can assign to so as to modify it.
  Expression** getIndexPointer(Expression* expr, Index index) {
    if (auto* unary = expr->dynCast<Unary>()) {
      assert(index == 0);
      return &unary->value;
    } else if (auto* binary = expr->dynCast<Binary>()) {
      if (index == 0) {
        return &binary->left;
      } else if (index == 1) {
        return &binary->right;
      }
      WASM_UNREACHABLE("unexpected index");
    } else if (auto* select = expr->dynCast<Select>()) {
      if (index == 0) {
        return &select->condition;
      } else if (index == 1) {
        return &select->ifTrue;
      } else if (index == 2) {
        return &select->ifFalse;
      }
      WASM_UNREACHABLE("unexpected index");
    }
    WASM_UNREACHABLE("unexpected expression type");
  }
};

Pass* createDataFlowOptsPass() { return new DataFlowOpts(); }

} // namespace wasm
<commit_msg>Fix DataFlowOpts leaking temporary Functions (#3093)<commit_after>/*
 * Copyright 2018 WebAssembly Community Group participants
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//
// Optimize using the DataFlow SSA IR.
//
// This needs 'flatten' to be run before it, and you should run full
// regular opts afterwards to clean up the flattening. For example,
// you might use it like this:
//
//    --flatten --dfo -Os
//

#include "dataflow/graph.h"
#include "dataflow/node.h"
#include "dataflow/users.h"
#include "dataflow/utils.h"
#include "ir/flat.h"
#include "ir/utils.h"
#include "pass.h"
#include "wasm-builder.h"
#include "wasm.h"

namespace wasm {

struct DataFlowOpts : public WalkerPass<PostWalker<DataFlowOpts>> {
  bool isFunctionParallel() override { return true; }

  Pass* create() override { return new DataFlowOpts; }

  DataFlow::Users nodeUsers;

  // The optimization work left to do: nodes that we need to look at.
  std::unordered_set<DataFlow::Node*> workLeft;

  DataFlow::Graph graph;

  void doWalkFunction(Function* func) {
    Flat::verifyFlatness(func);
    // Build the data-flow IR.
    graph.build(func, getModule());
    nodeUsers.build(graph);
    // Propagate optimizations through the graph.
    std::unordered_set<DataFlow::Node*> optimized; // which nodes we optimized
    for (auto& node : graph.nodes) {
      workLeft.insert(node.get()); // we should try to optimize each node
    }
    while (!workLeft.empty()) {
      // std::cout << "\n\ndump before work iter\n";
      // dump(graph, std::cout);
      auto iter = workLeft.begin();
      auto* node = *iter;
      workLeft.erase(iter);
      workOn(node);
    }
    // After updating the DataFlow IR, we can update the sets in
    // the wasm.
    // TODO: we also need phis, as a phi can flow directly into say
    //       a return or a call parameter.
    for (auto* set : graph.sets) {
      auto* node = graph.setNodeMap[set];
      auto iter = optimized.find(node);
      if (iter != optimized.end()) {
        assert(node->isExpr()); // this is a set, where the node is defined
        set->value = node->expr;
      }
    }
  }

  void workOn(DataFlow::Node* node) {
    if (node->isConst()) {
      return;
    }
    // If there are no uses, there is no point to work.
    if (nodeUsers.getNumUses(node) == 0) {
      return;
    }
    // Optimize: Look for nodes that we can easily convert into
    // something simpler.
    // TODO: we can expressionify and run full normal opts on that,
    //       then copy the result if it's smaller.
    if (node->isPhi() && DataFlow::allInputsIdentical(node)) {
      // Note we don't need to check for effects when replacing, as in
      // flattened IR expression children are local.gets or consts.
      auto* value = node->getValue(1);
      if (value->isConst()) {
        replaceAllUsesWith(node, value);
      }
    } else if (node->isExpr() && DataFlow::allInputsConstant(node)) {
      assert(!node->isConst());
      // If this is a concrete value (not e.g. an eqz of unreachable),
      // it can definitely be precomputed into a constant.
      if (node->expr->type.isConcrete()) {
        // This can be precomputed.
        // TODO not just all-constant inputs? E.g. i32.mul of 0 and X.
        optimizeExprToConstant(node);
      }
    }
  }

  void optimizeExprToConstant(DataFlow::Node* node) {
    assert(node->isExpr());
    assert(!node->isConst());
    // std::cout << "will optimize an Expr of all constant inputs. before" <<
    //              '\n';
    // dump(node, std::cout);
    auto* expr = node->expr;
    // First, note that some of the expression's children may be
    // local.gets that we inferred during SSA analysis as constant.
    // We can apply those now.
    for (Index i = 0; i < node->values.size(); i++) {
      if (node->values[i]->isConst()) {
        auto* currp = getIndexPointer(expr, i);
        // Directly represent it as a constant. (Note that it may already be
        // a constant, but for now to avoid corner cases just replace them
        // all here.)
        auto* c = node->values[i]->expr->dynCast<Const>();
        *currp = Builder(*getModule()).makeConst(c->value);
      }
    }
    // Now we know that all our DataFlow inputs are constant, and all
    // our Binaryen IR representations of them are constant too. RUn
    // precompute, which will transform the expression into a constanat.
    Module temp;
    // XXX we should copy expr here, in principle, and definitely will need to
    //     when we do arbitrarily regenerated expressions
    std::unique_ptr<Function> tempFunc(Builder(temp).makeFunction(
      "temp", Signature(Type::none, Type::none), {}, expr));
    PassRunner runner(&temp);
    runner.setIsNested(true);
    runner.add("precompute");
    runner.runOnFunction(tempFunc.get());
    // Get the optimized thing
    auto* result = tempFunc->body;
    // It may not be a constant, e.g. 0 / 0 does not optimize to 0
    if (!result->is<Const>()) {
      return;
    }
    // All good, copy it.
    node->expr = Builder(*getModule()).makeConst(result->cast<Const>()->value);
    assert(node->isConst());
    // We no longer have values, and so do not use anything.
    nodeUsers.stopUsingValues(node);
    node->values.clear();
    // Our contents changed, update our users.
    replaceAllUsesWith(node, node);
  }

  // Replaces all uses of a node with another value. This both modifies
  // the DataFlow IR to make the other users point to this one, and
  // updates the underlying Binaryen IR as well.
  // This can be used to "replace" a node with itself, which makes sense
  // when the node contents have changed and so the users must be updated.
  void replaceAllUsesWith(DataFlow::Node* node, DataFlow::Node* with) {
    // Const nodes are trivial to replace, but other stuff is trickier -
    // in particular phis.
    assert(with->isConst()); // TODO
    // All the users should be worked on later, as we will update them.
    auto& users = nodeUsers.getUsers(node);
    for (auto* user : users) {
      // Add the user to the work left to do, as we are modifying it.
      workLeft.insert(user);
      // `with` is getting another user.
      nodeUsers.addUser(with, user);
      // Replacing in the DataFlow IR is simple - just replace it,
      // in all the indexes it appears.
      std::vector<Index> indexes;
      for (Index i = 0; i < user->values.size(); i++) {
        if (user->values[i] == node) {
          user->values[i] = with;
          indexes.push_back(i);
        }
      }
      assert(!indexes.empty());
      // Replacing in the Binaryen IR requires more care
      switch (user->type) {
        case DataFlow::Node::Type::Expr: {
          auto* expr = user->expr;
          for (auto index : indexes) {
            *(getIndexPointer(expr, index)) = graph.makeUse(with);
          }
          break;
        }
        case DataFlow::Node::Type::Phi: {
          // Nothing to do: a phi is not in the Binaryen IR.
          // If the entire phi can become a constant, that will be
          // propagated when we process that phi later.
          break;
        }
        case DataFlow::Node::Type::Cond: {
          // Nothing to do: a cond is not in the Binaryen IR.
          // If the cond input is a constant, that might indicate
          // useful optimizations are possible, which perhaps we
          // should look into TODO
          break;
        }
        case DataFlow::Node::Type::Zext: {
          // Nothing to do: a zext is not in the Binaryen IR.
          // If the cond input is a constant, that might indicate
          // useful optimizations are possible, which perhaps we
          // should look into TODO
          break;
        }
        default:
          WASM_UNREACHABLE("unexpected dataflow node type");
      }
    }
    // No one is a user of this node after we replaced all the uses.
    nodeUsers.removeAllUsesOf(node);
  }

  // Gets a pointer to the expression pointer in an expression.
  // That is, given an index in the values() vector, get an
  // Expression** that we can assign to so as to modify it.
  Expression** getIndexPointer(Expression* expr, Index index) {
    if (auto* unary = expr->dynCast<Unary>()) {
      assert(index == 0);
      return &unary->value;
    } else if (auto* binary = expr->dynCast<Binary>()) {
      if (index == 0) {
        return &binary->left;
      } else if (index == 1) {
        return &binary->right;
      }
      WASM_UNREACHABLE("unexpected index");
    } else if (auto* select = expr->dynCast<Select>()) {
      if (index == 0) {
        return &select->condition;
      } else if (index == 1) {
        return &select->ifTrue;
      } else if (index == 2) {
        return &select->ifFalse;
      }
      WASM_UNREACHABLE("unexpected index");
    }
    WASM_UNREACHABLE("unexpected expression type");
  }
};

Pass* createDataFlowOptsPass() { return new DataFlowOpts(); }

} // namespace wasm
<|endoftext|>


<commit_before>#include "pngtexture.h"
#include "detail/TextureImpl.h"
#include "ifile.h"

#include <zlib.h>
#include <png.h>
#include <boost/assert.hpp>

namespace platform {
namespace {
const double DefaultGamma = 2.2;
const double InverseGamma = 1.0 / 2.2;
const double MaxGamma = 21474.83; // Maximum gamma accepted by png library.

class PngLoadData
{
    PngLoadData( );
    bool Load( File& F );
    png_structp mPngPtr;
    png_infop mInfoPtr;
    bool mNeedConversion;
    bool mFinished;
    std::auto_ptr< detail::TextureImpl > mTexture;
    ~PngLoadData();
public:
    static detail::TextureImpl* Create( File& F );
private:
    bool ProcessData( png_bytep Buffer, png_uint_32 Length );
    static bool IsSupportedInputChannelNum( size_t Channels );
    static void InfoCallback( png_structp PngPtr, png_infop InfoPtr );
    static void RowCallback( png_structp PngPtr, png_bytep NewRow, png_uint_32 RowNum, int Pass );
    static void EndCallback( png_structp PngPtr, png_infop InfoPtr );
};

bool PngLoadData::Load( File& F )
{
    png_set_progressive_read_fn( mPngPtr, ( void* )this, &PngLoadData::InfoCallback, &PngLoadData::RowCallback, &PngLoadData::EndCallback );
    size_t RemainingSize = F.GetSize();
    while( RemainingSize )
    {
        static const size_t MaxPassSize = 4096;
        size_t CurrentPass = std::min<size_t>( MaxPassSize, RemainingSize );
        std::string Buffer;
        if( !F.Read( Buffer, CurrentPass ) )
        {
            break;
        }
        if( !ProcessData( ( png_bytep )( void* )&Buffer[0], Buffer.size() ) )
        {
            break;
        }
        RemainingSize -= CurrentPass;
    }
    if( RemainingSize || !mFinished )
    {
        mTexture->mData.clear();
        mTexture->mWidth = 0;
        mTexture->mHeight = 0;
    }
    return !mTexture->mData.empty();
}

bool PngLoadData::ProcessData( png_bytep Buffer, png_uint_32 Length )
{
    if( setjmp( png_jmpbuf( mPngPtr ) ) )
    {
        return false;
    }

    png_process_data( mPngPtr, mInfoPtr, Buffer, Length );
    return true;
}

bool PngLoadData::IsSupportedInputChannelNum( size_t Channels )
{
    return Channels == 3 || Channels == 4;
}

void PngLoadData::InfoCallback( png_structp PngPtr, png_infop InfoPtr )
{
    PngLoadData* Self = ( PngLoadData* )png_get_progressive_ptr( PngPtr );

    int BitDepth, ColorType, InterlaceType, CompressionType;
    int FilterType, Channels;
    png_uint_32 Width, Height;
    png_get_IHDR( PngPtr, InfoPtr, &Width, &Height, &BitDepth, &ColorType,
                  &InterlaceType, &CompressionType, &FilterType );

    unsigned long long TotalSize = ( unsigned long long )Width * ( unsigned long long )Height;
    if( TotalSize > ( ( 1 << 29 ) - 1 ) )
    {
        longjmp( png_jmpbuf( PngPtr ), 1 );
    }

    Self->mTexture->mWidth = ( size_t )Width;
    Self->mTexture->mHeight = ( size_t )Height;

    // Expand to ensure we use 24-bit for RGB and 32-bit for RGBA.
    if( ColorType == PNG_COLOR_TYPE_PALETTE ||
        ( ColorType == PNG_COLOR_TYPE_GRAY && BitDepth < 8 ) )
    {
        png_set_expand( PngPtr );
    }

    // Transparency for paletted images.
    if( png_get_valid( PngPtr, InfoPtr, PNG_INFO_tRNS ) )
    {
        png_set_expand( PngPtr );
    }

    // Convert 16-bit to 8-bit.
    if( BitDepth == 16 )
    {
        png_set_strip_16( PngPtr );
    }

    // Expand grayscale to RGB.
    if( ColorType == PNG_COLOR_TYPE_GRAY ||
        ColorType == PNG_COLOR_TYPE_GRAY_ALPHA )
    {
        png_set_gray_to_rgb( PngPtr );
    }

    // Deal with gamma and keep it under our control.
    double Gamma;
    if( png_get_gAMA( PngPtr, InfoPtr, &Gamma ) )
    {
        if( Gamma <= 0.0 || Gamma > MaxGamma )
        {
            Gamma = InverseGamma;
            png_set_gAMA( PngPtr, InfoPtr, Gamma );
        }
        png_set_gamma( PngPtr, DefaultGamma, Gamma );
    }
    else
    {
        png_set_gamma( PngPtr, DefaultGamma, InverseGamma );
    }

    // Tell libpng to send us rows for interlaced pngs.
    if ( InterlaceType == PNG_INTERLACE_ADAM7 )
    {
        png_set_interlace_handling( PngPtr );
    }

    // Update our info now
    png_read_update_info( PngPtr, InfoPtr );
    Channels = png_get_channels( PngPtr, InfoPtr );

    // Pick our row format converter necessary for this data.
    if( !IsSupportedInputChannelNum( Channels ) )
    {
        longjmp( png_jmpbuf( PngPtr ), 1 );
    }
    Self->mNeedConversion = ( Channels != detail::TextureImpl::mChannels );
    Self->mTexture->mData.resize( Self->mTexture->mWidth * Self->mTexture->mHeight * detail::TextureImpl::mChannels );
}

void PngLoadData::RowCallback( png_structp PngPtr, png_bytep NewRow, png_uint_32 RowNum, int Pass )
{
    PngLoadData* Self = ( PngLoadData* )png_get_progressive_ptr( PngPtr );
    BOOST_ASSERT( Pass == 0 );
    if( RowNum > Self->mTexture->mHeight )
    {
        BOOST_ASSERT( false );
        return;
    }

    uint8_t* Dst = &( ( &Self->mTexture->mData.at(0) )[Self->mTexture->mWidth * detail::TextureImpl::mChannels * RowNum] );
    if( Self->mNeedConversion )
    {
        detail::TextureImpl::ConvertRGBtoRGBA( NewRow, Self->mTexture->mWidth, Dst );
    }
    else
    {
        memcpy( Dst, NewRow, Self->mTexture->mWidth * detail::TextureImpl::mChannels );
    }
}

void PngLoadData::EndCallback( png_structp PngPtr, png_infop InfoPtr )
{
    PngLoadData* Self = ( PngLoadData* )png_get_progressive_ptr( PngPtr );
    Self->mFinished = true;
}

PngLoadData::PngLoadData()
    : mPngPtr( NULL )
    , mInfoPtr( NULL )
    , mFinished( false )
    , mNeedConversion( false )
{
    mTexture.reset( new detail::TextureImpl );
}

PngLoadData::~PngLoadData()
{
    png_destroy_read_struct( &mPngPtr, &mInfoPtr, ( png_infopp )NULL );
}

detail::TextureImpl* PngLoadData::Create( File& F )
{
    PngLoadData LoadData;
    LoadData.mPngPtr = png_create_read_struct( PNG_LIBPNG_VER_STRING, ( png_voidp )NULL, NULL, NULL );
    if( !LoadData.mPngPtr )
    {
        return NULL;
    }
    LoadData.mInfoPtr = png_create_info_struct( LoadData.mPngPtr );
    if( !LoadData.mInfoPtr )
    {
        return NULL;
    }
    if( setjmp( png_jmpbuf( LoadData.mPngPtr ) ) )
    {
        return NULL;
    }
    if( !LoadData.Load( F ) )
    {
        return NULL;
    }
    return LoadData.mTexture.release();
}

} // namespace anonymous

PngTexture::PngTexture( File& F )
{
    if( F.IsValid() )
    {
        this->mImpl.reset( PngLoadData::Create( F ) );
    }
}

} // namespace platform
<commit_msg>B: string.h is needed for memcpy (on some platforms, the includes don't pull it in)<commit_after>#include "pngtexture.h"
#include "detail/TextureImpl.h"
#include "ifile.h"

#include <zlib.h>
#include <png.h>
#include <boost/assert.hpp>
#include <string.h>

namespace platform {
namespace {
const double DefaultGamma = 2.2;
const double InverseGamma = 1.0 / 2.2;
const double MaxGamma = 21474.83; // Maximum gamma accepted by png library.

class PngLoadData
{
    PngLoadData( );
    bool Load( File& F );
    png_structp mPngPtr;
    png_infop mInfoPtr;
    bool mNeedConversion;
    bool mFinished;
    std::auto_ptr< detail::TextureImpl > mTexture;
    ~PngLoadData();
public:
    static detail::TextureImpl* Create( File& F );
private:
    bool ProcessData( png_bytep Buffer, png_uint_32 Length );
    static bool IsSupportedInputChannelNum( size_t Channels );
    static void InfoCallback( png_structp PngPtr, png_infop InfoPtr );
    static void RowCallback( png_structp PngPtr, png_bytep NewRow, png_uint_32 RowNum, int Pass );
    static void EndCallback( png_structp PngPtr, png_infop InfoPtr );
};

bool PngLoadData::Load( File& F )
{
    png_set_progressive_read_fn( mPngPtr, ( void* )this, &PngLoadData::InfoCallback, &PngLoadData::RowCallback, &PngLoadData::EndCallback );
    size_t RemainingSize = F.GetSize();
    while( RemainingSize )
    {
        static const size_t MaxPassSize = 4096;
        size_t CurrentPass = std::min<size_t>( MaxPassSize, RemainingSize );
        std::string Buffer;
        if( !F.Read( Buffer, CurrentPass ) )
        {
            break;
        }
        if( !ProcessData( ( png_bytep )( void* )&Buffer[0], Buffer.size() ) )
        {
            break;
        }
        RemainingSize -= CurrentPass;
    }
    if( RemainingSize || !mFinished )
    {
        mTexture->mData.clear();
        mTexture->mWidth = 0;
        mTexture->mHeight = 0;
    }
    return !mTexture->mData.empty();
}

bool PngLoadData::ProcessData( png_bytep Buffer, png_uint_32 Length )
{
    if( setjmp( png_jmpbuf( mPngPtr ) ) )
    {
        return false;
    }

    png_process_data( mPngPtr, mInfoPtr, Buffer, Length );
    return true;
}

bool PngLoadData::IsSupportedInputChannelNum( size_t Channels )
{
    return Channels == 3 || Channels == 4;
}

void PngLoadData::InfoCallback( png_structp PngPtr, png_infop InfoPtr )
{
    PngLoadData* Self = ( PngLoadData* )png_get_progressive_ptr( PngPtr );

    int BitDepth, ColorType, InterlaceType, CompressionType;
    int FilterType, Channels;
    png_uint_32 Width, Height;
    png_get_IHDR( PngPtr, InfoPtr, &Width, &Height, &BitDepth, &ColorType,
                  &InterlaceType, &CompressionType, &FilterType );

    unsigned long long TotalSize = ( unsigned long long )Width * ( unsigned long long )Height;
    if( TotalSize > ( ( 1 << 29 ) - 1 ) )
    {
        longjmp( png_jmpbuf( PngPtr ), 1 );
    }

    Self->mTexture->mWidth = ( size_t )Width;
    Self->mTexture->mHeight = ( size_t )Height;

    // Expand to ensure we use 24-bit for RGB and 32-bit for RGBA.
    if( ColorType == PNG_COLOR_TYPE_PALETTE ||
        ( ColorType == PNG_COLOR_TYPE_GRAY && BitDepth < 8 ) )
    {
        png_set_expand( PngPtr );
    }

    // Transparency for paletted images.
    if( png_get_valid( PngPtr, InfoPtr, PNG_INFO_tRNS ) )
    {
        png_set_expand( PngPtr );
    }

    // Convert 16-bit to 8-bit.
    if( BitDepth == 16 )
    {
        png_set_strip_16( PngPtr );
    }

    // Expand grayscale to RGB.
    if( ColorType == PNG_COLOR_TYPE_GRAY ||
        ColorType == PNG_COLOR_TYPE_GRAY_ALPHA )
    {
        png_set_gray_to_rgb( PngPtr );
    }

    // Deal with gamma and keep it under our control.
    double Gamma;
    if( png_get_gAMA( PngPtr, InfoPtr, &Gamma ) )
    {
        if( Gamma <= 0.0 || Gamma > MaxGamma )
        {
            Gamma = InverseGamma;
            png_set_gAMA( PngPtr, InfoPtr, Gamma );
        }
        png_set_gamma( PngPtr, DefaultGamma, Gamma );
    }
    else
    {
        png_set_gamma( PngPtr, DefaultGamma, InverseGamma );
    }

    // Tell libpng to send us rows for interlaced pngs.
    if ( InterlaceType == PNG_INTERLACE_ADAM7 )
    {
        png_set_interlace_handling( PngPtr );
    }

    // Update our info now
    png_read_update_info( PngPtr, InfoPtr );
    Channels = png_get_channels( PngPtr, InfoPtr );

    // Pick our row format converter necessary for this data.
    if( !IsSupportedInputChannelNum( Channels ) )
    {
        longjmp( png_jmpbuf( PngPtr ), 1 );
    }
    Self->mNeedConversion = ( Channels != detail::TextureImpl::mChannels );
    Self->mTexture->mData.resize( Self->mTexture->mWidth * Self->mTexture->mHeight * detail::TextureImpl::mChannels );
}

void PngLoadData::RowCallback( png_structp PngPtr, png_bytep NewRow, png_uint_32 RowNum, int Pass )
{
    PngLoadData* Self = ( PngLoadData* )png_get_progressive_ptr( PngPtr );
    BOOST_ASSERT( Pass == 0 );
    if( RowNum > Self->mTexture->mHeight )
    {
        BOOST_ASSERT( false );
        return;
    }

    uint8_t* Dst = &( ( &Self->mTexture->mData.at(0) )[Self->mTexture->mWidth * detail::TextureImpl::mChannels * RowNum] );
    if( Self->mNeedConversion )
    {
        detail::TextureImpl::ConvertRGBtoRGBA( NewRow, Self->mTexture->mWidth, Dst );
    }
    else
    {
        memcpy( Dst, NewRow, Self->mTexture->mWidth * detail::TextureImpl::mChannels );
    }
}

void PngLoadData::EndCallback( png_structp PngPtr, png_infop InfoPtr )
{
    PngLoadData* Self = ( PngLoadData* )png_get_progressive_ptr( PngPtr );
    Self->mFinished = true;
}

PngLoadData::PngLoadData()
    : mPngPtr( NULL )
    , mInfoPtr( NULL )
    , mFinished( false )
    , mNeedConversion( false )
{
    mTexture.reset( new detail::TextureImpl );
}

PngLoadData::~PngLoadData()
{
    png_destroy_read_struct( &mPngPtr, &mInfoPtr, ( png_infopp )NULL );
}

detail::TextureImpl* PngLoadData::Create( File& F )
{
    PngLoadData LoadData;
    LoadData.mPngPtr = png_create_read_struct( PNG_LIBPNG_VER_STRING, ( png_voidp )NULL, NULL, NULL );
    if( !LoadData.mPngPtr )
    {
        return NULL;
    }
    LoadData.mInfoPtr = png_create_info_struct( LoadData.mPngPtr );
    if( !LoadData.mInfoPtr )
    {
        return NULL;
    }
    if( setjmp( png_jmpbuf( LoadData.mPngPtr ) ) )
    {
        return NULL;
    }
    if( !LoadData.Load( F ) )
    {
        return NULL;
    }
    return LoadData.mTexture.release();
}

} // namespace anonymous

PngTexture::PngTexture( File& F )
{
    if( F.IsValid() )
    {
        this->mImpl.reset( PngLoadData::Create( F ) );
    }
}

} // namespace platform
<|endoftext|>


<commit_before>// $Id: charmm.C,v 1.2 2000/02/10 10:46:40 oliver Exp $
// Molecular Mechanics: Charmm force field class

#include <BALL/MOLMEC/CHARMM/charmm.h>
#include <BALL/MOLMEC/CHARMM/charmmStretch.h>
#include <BALL/MOLMEC/CHARMM/charmmBend.h>
#include <BALL/MOLMEC/CHARMM/charmmTorsion.h>
#include <BALL/MOLMEC/CHARMM/charmmImproperTorsion.h>
#include <BALL/MOLMEC/CHARMM/charmmNonBonded.h>
#include <BALL/MOLMEC/COMMON/assignTypes.h>
#include <BALL/MOLMEC/PARAMETER/templates.h>
#include <BALL/KERNEL/PSE.h>
#include <BALL/KERNEL/atom.h>
#include <BALL/KERNEL/bond.h>

using namespace std;

namespace BALL
{
	const char* CharmmFF::Option::FILENAME = "filename";
	const char* CharmmFF::Option::NONBONDED_CUTOFF = "nonbonded_cutoff";
	const char* CharmmFF::Option::VDW_CUTOFF = "vdw_cutoff";
	const char* CharmmFF::Option::VDW_CUTON = "vdw_cuton";
	const char* CharmmFF::Option::ELECTROSTATIC_CUTOFF = "electrostatic_cutoff";
	const char* CharmmFF::Option::SCALING_VDW_1_4 = "SCAB";
	const char* CharmmFF::Option::SCALING_ELECTROSTATIC_1_4 = "SCEE";
	const char* CharmmFF::Option::DISTANCE_DEPENDENT_DIELECTRIC = "DDDC";
	const char* CharmmFF::Option::ASSIGN_CHARGES = "assign_charges";
	const char* CharmmFF::Option::ASSIGN_TYPENAMES = "assign_type_names";
	const char* CharmmFF::Option::ASSIGN_TYPES = "assign_types";
	const char* CharmmFF::Option::OVERWRITE_CHARGES = "overwrite_non-zero_charges";
	const char* CharmmFF::Option::OVERWRITE_TYPENAMES = "overwrite_non-empty_typenames";
	const char* CharmmFF::Option::USE_EEF1 = "use_EEF1";

	const char* CharmmFF::Default::FILENAME = "CHARMM/param22.ini";
	const float CharmmFF::Default::NONBONDED_CUTOFF = 20.0;
	const float CharmmFF::Default::VDW_CUTOFF = 15.0;
	const float CharmmFF::Default::VDW_CUTON = 5.0;
	const float CharmmFF::Default::ELECTROSTATIC_CUTOFF = 15.0;
	const float CharmmFF::Default::SCALING_ELECTROSTATIC_1_4 = 2.0;
	const float CharmmFF::Default::SCALING_VDW_1_4 = 1.0;
  const bool  CharmmFF::Default::DISTANCE_DEPENDENT_DIELECTRIC = false;
	const bool	CharmmFF::Default::ASSIGN_CHARGES = true;
	const bool	CharmmFF::Default::ASSIGN_TYPENAMES = true;
	const bool	CharmmFF::Default::ASSIGN_TYPES = true;
	const bool	CharmmFF::Default::OVERWRITE_CHARGES = true;
	const bool	CharmmFF::Default::OVERWRITE_TYPENAMES = false;
	const bool	CharmmFF::Default::USE_EEF1 = true;

	// Default constructor
	CharmmFF::CharmmFF()
		: ForceField(),
			filename_(Default::FILENAME)
	{
		// set the force field name
		setName("CHARMM [" + filename_ + "]");

		// create the component list
		insertComponent(new CharmmStretch(this));
		insertComponent(new CharmmBend(this));
		insertComponent(new CharmmTorsion(this));
		insertComponent(new CharmmImproperTorsion(this));
		insertComponent(new CharmmNonBonded(this));
	}

  // Constructor initialized with a system
  CharmmFF::CharmmFF(System& system)
    : ForceField(),
			filename_(Default::FILENAME)
  {
		// create the component list
		insertComponent(new CharmmStretch(this));
		insertComponent(new CharmmBend(this));
		insertComponent(new CharmmTorsion(this));
		insertComponent(new CharmmImproperTorsion(this));
		insertComponent(new CharmmNonBonded(this));

    bool result = setup(system);

		// set the force field name
		setName("CHARMM [" + filename_ + "]");

    if (!result)
    {
      Log.level(LogStream::ERROR) << " Force Field setup failed! " << endl;
      valid_ = false;
		}
	}

  // Constructor intialized with a system and a set of options
  CharmmFF::CharmmFF(System& system, const Options& new_options)
    : ForceField(),
			filename_(Default::FILENAME)
  {
		// create the component list
		insertComponent(new CharmmStretch(this));
		insertComponent(new CharmmBend(this));
		insertComponent(new CharmmTorsion(this));
		insertComponent(new CharmmImproperTorsion(this));
		insertComponent(new CharmmNonBonded(this));

    bool result = setup(system, new_options);

		// set the force field name (this has to be done after(!) setup,
		// otherwise filename_ is not yet set if it is used in options)
		setName("CHARMM [" + filename_ + "]");

    if (!result)
    {
      Log.level(LogStream::ERROR) << " Force Field setup failed! " << endl;
      valid_ = false;
		}
	}

	// copy constructor
	CharmmFF::CharmmFF(const CharmmFF& force_field, bool clone_deep)
		:	ForceField( force_field, clone_deep),
			filename_(force_field.filename_)
	{
	}

	// destructor
	CharmmFF::~CharmmFF()
	{
	}

	// force field specific setup method BAUSTELLE
	bool CharmmFF::specificSetup()
	{
		// check whether the system is assigned
		if (getSystem() == 0)
		{
			return false;
		}

		// check whether the parameter file name
		// is set in the options
		if (options.has(Option::FILENAME))
		{
			filename_ = options[Option::FILENAME];
			setName("Charmm [" + filename_ + "]");
		} else {
			options[Option::FILENAME] = filename_;
		}

		// open parameter file
		Path    path;
		String  filename(path.find(filename_));

		if (filename == "")
		{
			throw Exception::FileNotFound(__FILE__, __LINE__, filename_);
		}

		// initialize the force field parameters
		// and retrieve the atom types
		if (parameters_.getFilename() != filename)
		{
			parameters_.setFilename(filename);
			parameters_.init();

			// this is the first time, parameters was initialized
			// tell all components about it
			parameters_initialized_ = false;
		} else {
			// parameters_ are already initialized, tell all components about it
			parameters_initialized_ = true;
		}

		// check the options whether types, type names, or charges
		// have to be (re)asigned
		options.setDefaultBool(Option::ASSIGN_CHARGES, Default::ASSIGN_CHARGES);
		bool assign_charges = options.getBool(Option::ASSIGN_CHARGES);
		options.setDefaultBool(Option::ASSIGN_TYPES, Default::ASSIGN_TYPES);
		bool assign_types = options.getBool(Option::ASSIGN_TYPES);
		options.setDefaultBool(Option::ASSIGN_TYPENAMES, Default::ASSIGN_TYPENAMES);
		bool assign_type_names = options.getBool(Option::ASSIGN_TYPENAMES);
		options.setDefaultBool(Option::OVERWRITE_TYPENAMES, Default::OVERWRITE_TYPENAMES);
		bool overwrite_type_names = options.getBool(Option::OVERWRITE_TYPENAMES);
		options.setDefaultBool(Option::OVERWRITE_CHARGES, Default::OVERWRITE_CHARGES);
		bool overwrite_charges = options.getBool(Option::OVERWRITE_CHARGES);

		bool remove_hydrogens = true;

		// extract template section (containing charges and atom types)
		if (assign_charges || assign_type_names || remove_hydrogens)
		{
			FFPSTemplates templates;
			templates.extractSection(parameters_, "ChargesAndTypeNames");

			// remove all hydrogens bound to extended atom types
			if (remove_hydrogens)
			{
				HashSet<Atom*> atoms_to_delete;
				AtomIterator it = getSystem()->beginAtom();
				for (; +it; ++it)
				{
					if (it->getElement() != PSE[Element::H])
					{
	 					if (templates.has(it->getFullName()))
						{
							String type_name = templates.getTypeName(it->getFullName());
							bool extended = getParameters().getAtomTypes().getValue(type_name, "extended").toBool();
							if (extended)
							{
								// check for hydrogen atoms and insert them into the hash set
								// We use a hash set just in case - usually a hydrogen shouldn`t be
								// bound to two atoms!
								Atom::BondIterator bond_it = it->beginBond();
								for (; +bond_it; ++bond_it)
								{
									if (bond_it->getPartner(*it)->getElement() == PSE[Element::H])
									{
										// if the atom is a hydrogen atom, store its pointer
										// for removal
										atoms_to_delete.insert(bond_it->getPartner(*it));
									}
								}
							}
						}
					}
				}

				if (atoms_to_delete.size() > 0)
				{
					Log.info() << "CharmmFF::setup: deleted " << atoms_to_delete.size() << " hydrogen atoms (using CHARMM united atoms instead)." << endl;
					HashSet<Atom*>::Iterator hydrogen_it = atoms_to_delete.begin();
					for (; hydrogen_it != atoms_to_delete.end(); ++hydrogen_it)
					{
						delete *hydrogen_it;
					}

					// tell the ForceField class to recalculate the atoms_ vector
					// - we deleted things!
					atoms_.clear();
				}
			}

			if (assign_charges && assign_type_names)
			{
				templates.assign(*getSystem(), overwrite_type_names, overwrite_charges);
			} else {
				if (assign_type_names)
				{
					templates.assignTypeNames(*getSystem(), overwrite_type_names);
				} else {
					templates.assignCharges(*getSystem(), overwrite_charges);
				}
			}
		}

		// assign the types (type names should have been set already)
		if (assign_types)
		{
			// convert the type names to types
			AssignTypeProcessor type_proc(parameters_.getAtomTypes());
			getSystem()->apply(type_proc);
		}

		return true;
	}

	float CharmmFF::getStretchEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM Stretch");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getBendEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM Bend");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getImproperTorsionEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM ImproperTorsion");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getProperTorsionEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM Torsion");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getTorsionEnergy() const
	{
		float energy = 0;

		// get the energy of the proper torsions
		ForceFieldComponent* component = getComponent("CHARMM Torsion");
		if (component != 0)
		{
			energy += component->getEnergy();
		}
		// and add the energy of the improper torsions
		component = getComponent("CHARMM ImproperTorsion");
		if (component != 0)
		{
			energy += component->getEnergy();
		}

		return energy;
	}

	float CharmmFF::getVdWEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			const CharmmNonBonded* nonbonded_component = dynamic_cast<const CharmmNonBonded*>(component);
			if (nonbonded_component != 0)
			{
				energy = nonbonded_component->getVdwEnergy();
			}
		}

		return energy;
	}

	float CharmmFF::getESEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			const CharmmNonBonded* nonbonded_component = dynamic_cast<const CharmmNonBonded*>(component);
			if (nonbonded_component != 0)
			{
				energy = nonbonded_component->getElectrostaticEnergy();
			}
		}

		return energy;
	}

	float CharmmFF::getSolvationEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			const CharmmNonBonded* nonbonded_component = dynamic_cast<const CharmmNonBonded*>(component);
			if (nonbonded_component != 0)
			{
				energy = nonbonded_component->getSolvationEnergy();
			}
		}

		return energy;
	}

	float CharmmFF::getNonbondedEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			energy += component->getEnergy();
		}

		return energy;
	}

	bool CharmmFF::hasInitializedParameters() const
	{
		return parameters_initialized_;
	}

} // namespace BALL
<commit_msg>changed: default for distance dependendend ES is ON<commit_after>// $Id: charmm.C,v 1.3 2000/02/10 16:02:32 oliver Exp $
// Molecular Mechanics: Charmm force field class

#include <BALL/MOLMEC/CHARMM/charmm.h>
#include <BALL/MOLMEC/CHARMM/charmmStretch.h>
#include <BALL/MOLMEC/CHARMM/charmmBend.h>
#include <BALL/MOLMEC/CHARMM/charmmTorsion.h>
#include <BALL/MOLMEC/CHARMM/charmmImproperTorsion.h>
#include <BALL/MOLMEC/CHARMM/charmmNonBonded.h>
#include <BALL/MOLMEC/COMMON/assignTypes.h>
#include <BALL/MOLMEC/PARAMETER/templates.h>
#include <BALL/KERNEL/PSE.h>
#include <BALL/KERNEL/atom.h>
#include <BALL/KERNEL/bond.h>

using namespace std;

namespace BALL
{
	const char* CharmmFF::Option::FILENAME = "filename";
	const char* CharmmFF::Option::NONBONDED_CUTOFF = "nonbonded_cutoff";
	const char* CharmmFF::Option::VDW_CUTOFF = "vdw_cutoff";
	const char* CharmmFF::Option::VDW_CUTON = "vdw_cuton";
	const char* CharmmFF::Option::ELECTROSTATIC_CUTOFF = "electrostatic_cutoff";
	const char* CharmmFF::Option::SCALING_VDW_1_4 = "SCAB";
	const char* CharmmFF::Option::SCALING_ELECTROSTATIC_1_4 = "SCEE";
	const char* CharmmFF::Option::DISTANCE_DEPENDENT_DIELECTRIC = "DDDC";
	const char* CharmmFF::Option::ASSIGN_CHARGES = "assign_charges";
	const char* CharmmFF::Option::ASSIGN_TYPENAMES = "assign_type_names";
	const char* CharmmFF::Option::ASSIGN_TYPES = "assign_types";
	const char* CharmmFF::Option::OVERWRITE_CHARGES = "overwrite_non-zero_charges";
	const char* CharmmFF::Option::OVERWRITE_TYPENAMES = "overwrite_non-empty_typenames";
	const char* CharmmFF::Option::USE_EEF1 = "use_EEF1";

	const char* CharmmFF::Default::FILENAME = "CHARMM/param22.ini";
	const float CharmmFF::Default::NONBONDED_CUTOFF = 20.0;
	const float CharmmFF::Default::VDW_CUTOFF = 15.0;
	const float CharmmFF::Default::VDW_CUTON = 5.0;
	const float CharmmFF::Default::ELECTROSTATIC_CUTOFF = 15.0;
	const float CharmmFF::Default::SCALING_ELECTROSTATIC_1_4 = 2.0;
	const float CharmmFF::Default::SCALING_VDW_1_4 = 1.0;
  const bool  CharmmFF::Default::DISTANCE_DEPENDENT_DIELECTRIC = true;
	const bool	CharmmFF::Default::ASSIGN_CHARGES = true;
	const bool	CharmmFF::Default::ASSIGN_TYPENAMES = true;
	const bool	CharmmFF::Default::ASSIGN_TYPES = true;
	const bool	CharmmFF::Default::OVERWRITE_CHARGES = true;
	const bool	CharmmFF::Default::OVERWRITE_TYPENAMES = false;
	const bool	CharmmFF::Default::USE_EEF1 = true;

	// Default constructor
	CharmmFF::CharmmFF()
		: ForceField(),
			filename_(Default::FILENAME)
	{
		// set the force field name
		setName("CHARMM [" + filename_ + "]");

		// create the component list
		insertComponent(new CharmmStretch(this));
		insertComponent(new CharmmBend(this));
		insertComponent(new CharmmTorsion(this));
		insertComponent(new CharmmImproperTorsion(this));
		insertComponent(new CharmmNonBonded(this));
	}

  // Constructor initialized with a system
  CharmmFF::CharmmFF(System& system)
    : ForceField(),
			filename_(Default::FILENAME)
  {
		// create the component list
		insertComponent(new CharmmStretch(this));
		insertComponent(new CharmmBend(this));
		insertComponent(new CharmmTorsion(this));
		insertComponent(new CharmmImproperTorsion(this));
		insertComponent(new CharmmNonBonded(this));

    bool result = setup(system);

		// set the force field name
		setName("CHARMM [" + filename_ + "]");

    if (!result)
    {
      Log.level(LogStream::ERROR) << " Force Field setup failed! " << endl;
      valid_ = false;
		}
	}

  // Constructor intialized with a system and a set of options
  CharmmFF::CharmmFF(System& system, const Options& new_options)
    : ForceField(),
			filename_(Default::FILENAME)
  {
		// create the component list
		insertComponent(new CharmmStretch(this));
		insertComponent(new CharmmBend(this));
		insertComponent(new CharmmTorsion(this));
		insertComponent(new CharmmImproperTorsion(this));
		insertComponent(new CharmmNonBonded(this));

    bool result = setup(system, new_options);

		// set the force field name (this has to be done after(!) setup,
		// otherwise filename_ is not yet set if it is used in options)
		setName("CHARMM [" + filename_ + "]");

    if (!result)
    {
      Log.level(LogStream::ERROR) << " Force Field setup failed! " << endl;
      valid_ = false;
		}
	}

	// copy constructor
	CharmmFF::CharmmFF(const CharmmFF& force_field, bool clone_deep)
		:	ForceField( force_field, clone_deep),
			filename_(force_field.filename_)
	{
	}

	// destructor
	CharmmFF::~CharmmFF()
	{
	}

	// force field specific setup method BAUSTELLE
	bool CharmmFF::specificSetup()
	{
		// check whether the system is assigned
		if (getSystem() == 0)
		{
			return false;
		}

		// check whether the parameter file name
		// is set in the options
		if (options.has(Option::FILENAME))
		{
			filename_ = options[Option::FILENAME];
			setName("Charmm [" + filename_ + "]");
		} else {
			options[Option::FILENAME] = filename_;
		}

		// open parameter file
		Path    path;
		String  filename(path.find(filename_));

		if (filename == "")
		{
			throw Exception::FileNotFound(__FILE__, __LINE__, filename_);
		}

		// initialize the force field parameters
		// and retrieve the atom types
		if (parameters_.getFilename() != filename)
		{
			parameters_.setFilename(filename);
			parameters_.init();

			// this is the first time, parameters was initialized
			// tell all components about it
			parameters_initialized_ = false;
		} else {
			// parameters_ are already initialized, tell all components about it
			parameters_initialized_ = true;
		}

		// check the options whether types, type names, or charges
		// have to be (re)asigned
		options.setDefaultBool(Option::ASSIGN_CHARGES, Default::ASSIGN_CHARGES);
		bool assign_charges = options.getBool(Option::ASSIGN_CHARGES);
		options.setDefaultBool(Option::ASSIGN_TYPES, Default::ASSIGN_TYPES);
		bool assign_types = options.getBool(Option::ASSIGN_TYPES);
		options.setDefaultBool(Option::ASSIGN_TYPENAMES, Default::ASSIGN_TYPENAMES);
		bool assign_type_names = options.getBool(Option::ASSIGN_TYPENAMES);
		options.setDefaultBool(Option::OVERWRITE_TYPENAMES, Default::OVERWRITE_TYPENAMES);
		bool overwrite_type_names = options.getBool(Option::OVERWRITE_TYPENAMES);
		options.setDefaultBool(Option::OVERWRITE_CHARGES, Default::OVERWRITE_CHARGES);
		bool overwrite_charges = options.getBool(Option::OVERWRITE_CHARGES);

		bool remove_hydrogens = true;

		// extract template section (containing charges and atom types)
		if (assign_charges || assign_type_names || remove_hydrogens)
		{
			FFPSTemplates templates;
			templates.extractSection(parameters_, "ChargesAndTypeNames");

			// remove all hydrogens bound to extended atom types
			if (remove_hydrogens)
			{
				HashSet<Atom*> atoms_to_delete;
				AtomIterator it = getSystem()->beginAtom();
				for (; +it; ++it)
				{
					if (it->getElement() != PSE[Element::H])
					{
	 					if (templates.has(it->getFullName()))
						{
							String type_name = templates.getTypeName(it->getFullName());
							bool extended = getParameters().getAtomTypes().getValue(type_name, "extended").toBool();
							if (extended)
							{
								// check for hydrogen atoms and insert them into the hash set
								// We use a hash set just in case - usually a hydrogen shouldn`t be
								// bound to two atoms!
								Atom::BondIterator bond_it = it->beginBond();
								for (; +bond_it; ++bond_it)
								{
									if (bond_it->getPartner(*it)->getElement() == PSE[Element::H])
									{
										// if the atom is a hydrogen atom, store its pointer
										// for removal
										atoms_to_delete.insert(bond_it->getPartner(*it));
									}
								}
							}
						}
					}
				}

				if (atoms_to_delete.size() > 0)
				{
					Log.info() << "CharmmFF::setup: deleted " << atoms_to_delete.size() << " hydrogen atoms (using CHARMM united atoms instead)." << endl;
					HashSet<Atom*>::Iterator hydrogen_it = atoms_to_delete.begin();
					for (; hydrogen_it != atoms_to_delete.end(); ++hydrogen_it)
					{
						delete *hydrogen_it;
					}

					// tell the ForceField class to recalculate the atoms_ vector
					// - we deleted things!
					atoms_.clear();
				}
			}

			if (assign_charges && assign_type_names)
			{
				templates.assign(*getSystem(), overwrite_type_names, overwrite_charges);
			} else {
				if (assign_type_names)
				{
					templates.assignTypeNames(*getSystem(), overwrite_type_names);
				} else {
					templates.assignCharges(*getSystem(), overwrite_charges);
				}
			}
		}

		// assign the types (type names should have been set already)
		if (assign_types)
		{
			// convert the type names to types
			AssignTypeProcessor type_proc(parameters_.getAtomTypes());
			getSystem()->apply(type_proc);
		}

		return true;
	}

	float CharmmFF::getStretchEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM Stretch");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getBendEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM Bend");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getImproperTorsionEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM ImproperTorsion");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getProperTorsionEnergy() const
	{
		ForceFieldComponent* component = getComponent("CHARMM Torsion");
		if (component != 0)
		{
			return component->getEnergy();
		} else {
			return 0;
		}
	}

	float CharmmFF::getTorsionEnergy() const
	{
		float energy = 0;

		// get the energy of the proper torsions
		ForceFieldComponent* component = getComponent("CHARMM Torsion");
		if (component != 0)
		{
			energy += component->getEnergy();
		}
		// and add the energy of the improper torsions
		component = getComponent("CHARMM ImproperTorsion");
		if (component != 0)
		{
			energy += component->getEnergy();
		}

		return energy;
	}

	float CharmmFF::getVdWEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			const CharmmNonBonded* nonbonded_component = dynamic_cast<const CharmmNonBonded*>(component);
			if (nonbonded_component != 0)
			{
				energy = nonbonded_component->getVdwEnergy();
			}
		}

		return energy;
	}

	float CharmmFF::getESEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			const CharmmNonBonded* nonbonded_component = dynamic_cast<const CharmmNonBonded*>(component);
			if (nonbonded_component != 0)
			{
				energy = nonbonded_component->getElectrostaticEnergy();
			}
		}

		return energy;
	}

	float CharmmFF::getSolvationEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			const CharmmNonBonded* nonbonded_component = dynamic_cast<const CharmmNonBonded*>(component);
			if (nonbonded_component != 0)
			{
				energy = nonbonded_component->getSolvationEnergy();
			}
		}

		return energy;
	}

	float CharmmFF::getNonbondedEnergy() const
	{
		float energy = 0;
		const ForceFieldComponent* component = getComponent("CHARMM NonBonded");
		if (component != 0)
		{
			energy += component->getEnergy();
		}

		return energy;
	}

	bool CharmmFF::hasInitializedParameters() const
	{
		return parameters_initialized_;
	}

} // namespace BALL
<|endoftext|>


<commit_before>//
//  iostream.hpp
//  fibio
//
//  Created by Chen Xu on 14-3-4.
//  Copyright (c) 2014 0d0a.com. All rights reserved.
//

#ifndef fibio_stream_iostream_hpp
#define fibio_stream_iostream_hpp

#include <map>
#include <boost/asio/ip/basic_resolver.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/posix/stream_descriptor.hpp>
#include <boost/asio/local/stream_protocol.hpp>
#include <boost/lexical_cast.hpp>
#include <fibio/fiber.hpp>
#include <fibio/future.hpp>
#include <fibio/stream/streambuf.hpp>

namespace fibio { namespace stream {
    namespace detail {
        template<typename Stream>
        struct iostream_base {
            typedef streambuf<Stream> streambuf_t;
            iostream_base() : sbuf_(new streambuf_t) {}
            iostream_base(iostream_base &&other)
            : sbuf_(std::move(other.sbuf_))
            {}
            // For SSL stream, construct with ssl::context
            template<typename Arg>
            iostream_base(Arg &arg)
            : sbuf_(new streambuf_t(arg))
            {}
            std::unique_ptr<streambuf_t> sbuf_;
        };
        template<typename R>
        R make_endpoint(const std::string &access_point) {}

        template<>
        inline boost::asio::ip::tcp::endpoint make_endpoint<boost::asio::ip::tcp::endpoint>(const std::string &access_point) {
            auto i=access_point.rfind(':');
            // Use 0::0 if access point contains port only
            std::string addr("0::0");
            std::string port;
            if(i==access_point.npos) {
                // ':' not found, assume arg contains only port
                port=access_point;
            } else {
                // Split access point into address and port
                addr.assign(access_point.begin(), access_point.begin()+i);
                port.assign(access_point.begin()+i+1, access_point.end());
                // Remove enclosing '[]' for IPv6 address
                if (addr[0]=='[' && *(addr.rbegin())==']') {
                    addr.assign(addr, 1, addr.size()-2);
                }
            }
            // Test if the address is numeric IP
            boost::system::error_code ec;
            boost::asio::ip::address a=boost::asio::ip::address::from_string(addr, ec);
            if(!ec) {
                try {
                    // Assume we have a numeric IP address and numeric port
                    return boost::asio::ip::tcp::endpoint(a, boost::lexical_cast<uint16_t>(port));
                } catch(boost::bad_lexical_cast &) {
                    // Port is not a number, need to resolve
                }
            }
            // Address and/or port are *not* numeric
            boost::asio::ip::tcp::resolver r(asio::get_io_service());
            boost::asio::ip::tcp::resolver::query q(addr, port);
            auto res=r.async_resolve(q, asio::yield[ec]);
            if(ec) {
                BOOST_THROW_EXCEPTION(boost::system::system_error(ec, "Resolving address failed"));
            }
            return res->endpoint();
        }

        template<>
        inline boost::asio::local::stream_protocol::endpoint make_endpoint<boost::asio::local::stream_protocol::endpoint>(const std::string &access_point) {
            return boost::asio::local::stream_protocol::endpoint(access_point);
        }
    }

    // Closable stream
    class closable_stream : public std::iostream {
    public:
        typedef std::iostream base_type;
        using base_type::base_type;
        virtual bool is_open() const=0;
        virtual void close()=0;
    };

    template<typename Stream>
    class iostream
    : public detail::iostream_base<Stream>
    , public closable_stream
    {
        typedef streambuf<Stream> streambuf_t;
        typedef detail::iostream_base<Stream> streambase_t;
    public:
        typedef typename streambuf_t::stream_type stream_type;
        typedef typename stream_type::lowest_layer_type::protocol_type protocol_type;
        typedef typename protocol_type::endpoint endpoint_type;

        iostream()
        : streambase_t()
        , closable_stream(rdbuf())
        {}

        // For SSL stream, construct with ssl::context
        template<typename Arg>
        iostream(Arg &arg)
        : streambase_t(arg)
        , closable_stream(rdbuf())
        {}

        // Movable
        iostream(iostream &&src)//=default;
        : streambase_t(std::move(src))
        , closable_stream(rdbuf())
        {}

        // Non-copyable
        iostream(const iostream&) = delete;
        iostream& operator=(const iostream&) = delete;

        template <typename... T>
        boost::system::error_code open(T... x) {
            return rdbuf()->lowest_layer().open(x...);
        }

        boost::system::error_code connect(const endpoint_type &ep) {
            return rdbuf()->connect(ep);
        }

        boost::system::error_code connect(const std::string &host, const std::string &service) {
            boost::system::error_code ec;
            typedef typename protocol_type::resolver resolver_type;
            typedef typename resolver_type::query query_type;
            resolver_type r(asio::get_io_service());
            query_type q(host, service);
            typename resolver_type::iterator i=r.async_resolve(q, asio::yield[ec]);
            if (ec) return ec;
            return rdbuf()->connect(i->endpoint());
        }

        boost::system::error_code connect(const char *access_point) {
            return rdbuf()->connect(detail::make_endpoint<endpoint_type>(access_point));
        }

        boost::system::error_code connect(const std::string &access_point) {
            return rdbuf()->connect(detail::make_endpoint<endpoint_type>(access_point));
        }

        /**
         * Close underlying stream device, flushing if necessary
         */
        inline void close() {
            if(rdbuf()->lowest_layer().is_open()) {
                flush();
            }
            rdbuf()->lowest_layer().close();
        }

        inline bool is_open() const
        { return rdbuf()->lowest_layer().is_open(); }

        inline streambuf_t *rdbuf() const
        { return this->sbuf_.get(); }

        inline stream_type &stream_descriptor()
        { return *rdbuf(); }

        void set_duplex_mode(duplex_mode dm)
        { rdbuf()->set_duplex_mode(dm); }

        duplex_mode get_duplex_mode() const
        { return rdbuf()->get_duplex_mode(); }
    };

    template<typename Stream>
    iostream<Stream> &operator<<(iostream<Stream> &is, duplex_mode dm) {
        is.set_duplex_mode(dm);
        return is;
    }

    // streams
    typedef iostream<boost::asio::ip::tcp::socket> tcp_stream;
    typedef iostream<boost::asio::posix::stream_descriptor> posix_stream;
    typedef iostream<boost::asio::local::stream_protocol::socket> local_stream;

    template<typename Stream>
    struct stream_acceptor {
        typedef Stream stream_type;
        typedef typename Stream::stream_type socket_type;
        typedef typename stream_type::protocol_type::acceptor acceptor_type;
        typedef typename stream_type::protocol_type::endpoint endpoint_type;

        stream_acceptor(const std::string &s, unsigned short port_num)
        : acc_(asio::get_io_service(),
               endpoint_type(boost::asio::ip::address::from_string(s.c_str()), port_num))
        {}

        stream_acceptor(unsigned short port_num)
        : acc_(asio::get_io_service(),
               endpoint_type(boost::asio::ip::address(), port_num))
        {}

        stream_acceptor(const endpoint_type &ep)
        : acc_(asio::get_io_service(), ep)
        {}

        stream_acceptor(const char *access_point)
        : stream_acceptor(detail::make_endpoint<endpoint_type>(access_point))
        {}

        stream_acceptor(const std::string &access_point)
        : stream_acceptor(detail::make_endpoint<endpoint_type>(access_point))
        {}

        stream_acceptor(stream_acceptor &&other)
        : acc_(std::move(other.acc_))
        {}

        stream_acceptor(const stream_acceptor &other)=delete;
        stream_acceptor &operator=(const stream_acceptor &other)=delete;

        void close()
        { acc_.close(); }

        stream_type accept() {
            stream_type s;
            boost::system::error_code ec;
            accept(s, ec);
            return s;
        }

        stream_type accept(boost::system::error_code &ec) {
            stream_type s;
            accept(s, ec);
            return s;
        }

        boost::system::error_code accept(stream_type &s) {
            boost::system::error_code ec;
            accept(s, ec);
            return ec;
        }

        void accept(stream_type &s, boost::system::error_code &ec)
        { acc_.async_accept(*(s.rdbuf()), asio::yield[ec]); }

        stream_type operator()()
        { return accept(); }

        stream_type operator()(boost::system::error_code &ec)
        { return accept(ec); }

        boost::system::error_code operator()(stream_type &s)
        { return accept(s); }

        void operator()(stream_type &s, boost::system::error_code &ec)
        { accept(s, ec); }

        acceptor_type acc_;
    };

    // acceptors
    typedef stream_acceptor<tcp_stream> tcp_stream_acceptor;
    typedef stream_acceptor<local_stream> local_stream_acceptor;

    template<typename Stream>
    struct stream_traits {
        typedef Stream stream_type;
        typedef typename stream_type::stream_type socket_type;
        typedef stream_acceptor<stream_type> acceptor_type;
        typedef typename acceptor_type::endpoint_type endpoint_type;
        // HACK:
        typedef int arg_type;
        static std::unique_ptr<stream_type> construct(arg_type *arg) {
            return std::unique_ptr<stream_type>(new stream_type());
        }
    };

    template<typename Stream>
    struct listener {
        typedef Stream stream_type;
        typedef stream_traits<Stream> traits_type;
        typedef typename traits_type::acceptor_type acceptor_type;
        typedef typename traits_type::endpoint_type endpoint_type;
        typedef typename traits_type::arg_type arg_type;

        template<typename std::enable_if<std::is_move_constructible<Stream>::value>::type* = nullptr >
        listener(const std::string &access_point)
        : ep_(detail::make_endpoint<endpoint_type>(access_point))
        {}

        // For SSL, arg type is ssl::context
        template<typename Arg>
        listener(Arg &arg, const std::string &access_point)
        : arg_(&arg)
        , ep_(detail::make_endpoint<endpoint_type>(access_point))
        {}

        endpoint_type endpoint() const {
            return ep_;
        }

        // Start and join, other fiber may stop the listener
        template<typename F>
        boost::system::error_code operator()(F f) {
            try {
                start(f);
                join();
            } catch(boost::system::system_error &e) {
                return e.code();
            }
            return boost::system::error_code();
        }

        // Start listener
        template<typename F>
        void start(F f) {
            if(!stop_signal_) {
                stop_signal_.reset(new promise<void>);
                if(!acceptor_fiber_) {
                    acceptor_fiber_.reset(new fiber(&listener::acceptor_fiber<F>,
                                                    this,
                                                    ep_,
                                                    stop_signal_.get(),
                                                    f));
                }
            }
        }

        void stop() {
            if(stop_signal_) {
                stop_signal_->set_value();
            }
        }

        void join() {
            if (acceptor_fiber_) {
                acceptor_fiber_->join();
                acceptor_fiber_.reset();
                stop_signal_.reset();
            }
        }

    private:
        template<typename F>
        void acceptor_fiber(const endpoint_type & e, promise<void> *p, F f) {
            acceptor_type acc(e);
            boost::system::error_code ec;
            fiber watchdog(fiber::attributes(fiber::attributes::stick_with_parent),
                           &listener::acceptor_watchdog_fiber,
                           this,
                           p,
                           std::ref(acc));
            while(!ec) {
                std::unique_ptr<stream_type> s(traits_type::construct(arg_));
                acc(*s, ec);
                if(ec) break;
                fiber([f](std::unique_ptr<stream_type> str){
                    f(*str);
                }, std::move(s)).detach();
            }
            watchdog.join();
        }

        void acceptor_watchdog_fiber(promise<void> *p, acceptor_type &acc) {
            p->get_future().wait();
            acc.close();
        }

        arg_type *arg_=nullptr;
        endpoint_type ep_;
        std::unique_ptr<fiber> acceptor_fiber_;
        std::unique_ptr<promise<void>> stop_signal_;
    };

    // listeners
    typedef listener<tcp_stream> tcp_listener;
    typedef listener<local_stream> local_listener;
}}  // End of namespace fibio::stream

namespace fibio {
    using stream::tcp_stream;
    using stream::posix_stream;
    using stream::local_stream;
    using stream::tcp_stream_acceptor;
    using stream::local_stream_acceptor;
    using stream::tcp_listener;
    using stream::local_listener;
}

#endif
<commit_msg>Windows handle stream<commit_after>//
//  iostream.hpp
//  fibio
//
//  Created by Chen Xu on 14-3-4.
//  Copyright (c) 2014 0d0a.com. All rights reserved.
//

#ifndef fibio_stream_iostream_hpp
#define fibio_stream_iostream_hpp

#include <map>
#include <boost/asio/ip/basic_resolver.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/posix/stream_descriptor.hpp>
#include <boost/asio/local/stream_protocol.hpp>
#include <boost/lexical_cast.hpp>
#include <fibio/fiber.hpp>
#include <fibio/future.hpp>
#include <fibio/stream/streambuf.hpp>

namespace fibio { namespace stream {
    namespace detail {
        template<typename Stream>
        struct iostream_base {
            typedef streambuf<Stream> streambuf_t;
            iostream_base() : sbuf_(new streambuf_t) {}
            iostream_base(iostream_base &&other)
            : sbuf_(std::move(other.sbuf_))
            {}
            // For SSL stream, construct with ssl::context
            template<typename Arg>
            iostream_base(Arg &arg)
            : sbuf_(new streambuf_t(arg))
            {}
            std::unique_ptr<streambuf_t> sbuf_;
        };
        template<typename R>
        R make_endpoint(const std::string &access_point) {}

        template<>
        inline boost::asio::ip::tcp::endpoint make_endpoint<boost::asio::ip::tcp::endpoint>(const std::string &access_point) {
            auto i=access_point.rfind(':');
            // Use 0::0 if access point contains port only
            std::string addr("0::0");
            std::string port;
            if(i==access_point.npos) {
                // ':' not found, assume arg contains only port
                port=access_point;
            } else {
                // Split access point into address and port
                addr.assign(access_point.begin(), access_point.begin()+i);
                port.assign(access_point.begin()+i+1, access_point.end());
                // Remove enclosing '[]' for IPv6 address
                if (addr[0]=='[' && *(addr.rbegin())==']') {
                    addr.assign(addr, 1, addr.size()-2);
                }
            }
            // Test if the address is numeric IP
            boost::system::error_code ec;
            boost::asio::ip::address a=boost::asio::ip::address::from_string(addr, ec);
            if(!ec) {
                try {
                    // Assume we have a numeric IP address and numeric port
                    return boost::asio::ip::tcp::endpoint(a, boost::lexical_cast<uint16_t>(port));
                } catch(boost::bad_lexical_cast &) {
                    // Port is not a number, need to resolve
                }
            }
            // Address and/or port are *not* numeric
            boost::asio::ip::tcp::resolver r(asio::get_io_service());
            boost::asio::ip::tcp::resolver::query q(addr, port);
            auto res=r.async_resolve(q, asio::yield[ec]);
            if(ec) {
                BOOST_THROW_EXCEPTION(boost::system::system_error(ec, "Resolving address failed"));
            }
            return res->endpoint();
        }

        template<>
        inline boost::asio::local::stream_protocol::endpoint make_endpoint<boost::asio::local::stream_protocol::endpoint>(const std::string &access_point) {
            return boost::asio::local::stream_protocol::endpoint(access_point);
        }
    }

    // Closable stream
    class closable_stream : public std::iostream {
    public:
        typedef std::iostream base_type;
        using base_type::base_type;
        virtual bool is_open() const=0;
        virtual void close()=0;
    };

    template<typename Stream>
    class iostream
    : public detail::iostream_base<Stream>
    , public closable_stream
    {
        typedef streambuf<Stream> streambuf_t;
        typedef detail::iostream_base<Stream> streambase_t;
    public:
        typedef typename streambuf_t::stream_type stream_type;
        typedef typename stream_type::lowest_layer_type::protocol_type protocol_type;
        typedef typename protocol_type::endpoint endpoint_type;

        iostream()
        : streambase_t()
        , closable_stream(rdbuf())
        {}

        // For SSL stream, construct with ssl::context
        template<typename Arg>
        iostream(Arg &arg)
        : streambase_t(arg)
        , closable_stream(rdbuf())
        {}

        // Movable
        iostream(iostream &&src)//=default;
        : streambase_t(std::move(src))
        , closable_stream(rdbuf())
        {}

        // Non-copyable
        iostream(const iostream&) = delete;
        iostream& operator=(const iostream&) = delete;

        template <typename... T>
        boost::system::error_code open(T... x) {
            return rdbuf()->lowest_layer().open(x...);
        }

        boost::system::error_code connect(const endpoint_type &ep) {
            return rdbuf()->connect(ep);
        }

        boost::system::error_code connect(const std::string &host, const std::string &service) {
            boost::system::error_code ec;
            typedef typename protocol_type::resolver resolver_type;
            typedef typename resolver_type::query query_type;
            resolver_type r(asio::get_io_service());
            query_type q(host, service);
            typename resolver_type::iterator i=r.async_resolve(q, asio::yield[ec]);
            if (ec) return ec;
            return rdbuf()->connect(i->endpoint());
        }

        boost::system::error_code connect(const char *access_point) {
            return rdbuf()->connect(detail::make_endpoint<endpoint_type>(access_point));
        }

        boost::system::error_code connect(const std::string &access_point) {
            return rdbuf()->connect(detail::make_endpoint<endpoint_type>(access_point));
        }

        /**
         * Close underlying stream device, flushing if necessary
         */
        inline void close() {
            if(rdbuf()->lowest_layer().is_open()) {
                flush();
            }
            rdbuf()->lowest_layer().close();
        }

        inline bool is_open() const
        { return rdbuf()->lowest_layer().is_open(); }

        inline streambuf_t *rdbuf() const
        { return this->sbuf_.get(); }

        inline stream_type &stream_descriptor()
        { return *rdbuf(); }

        void set_duplex_mode(duplex_mode dm)
        { rdbuf()->set_duplex_mode(dm); }

        duplex_mode get_duplex_mode() const
        { return rdbuf()->get_duplex_mode(); }
    };

    template<typename Stream>
    iostream<Stream> &operator<<(iostream<Stream> &is, duplex_mode dm) {
        is.set_duplex_mode(dm);
        return is;
    }

    template<typename Stream>
    struct stream_acceptor {
        typedef Stream stream_type;
        typedef typename Stream::stream_type socket_type;
        typedef typename stream_type::protocol_type::acceptor acceptor_type;
        typedef typename stream_type::protocol_type::endpoint endpoint_type;

        stream_acceptor(const std::string &s, unsigned short port_num)
        : acc_(asio::get_io_service(),
               endpoint_type(boost::asio::ip::address::from_string(s.c_str()), port_num))
        {}

        stream_acceptor(unsigned short port_num)
        : acc_(asio::get_io_service(),
               endpoint_type(boost::asio::ip::address(), port_num))
        {}

        stream_acceptor(const endpoint_type &ep)
        : acc_(asio::get_io_service(), ep)
        {}

        stream_acceptor(const char *access_point)
        : stream_acceptor(detail::make_endpoint<endpoint_type>(access_point))
        {}

        stream_acceptor(const std::string &access_point)
        : stream_acceptor(detail::make_endpoint<endpoint_type>(access_point))
        {}

        stream_acceptor(stream_acceptor &&other)
        : acc_(std::move(other.acc_))
        {}

        stream_acceptor(const stream_acceptor &other)=delete;
        stream_acceptor &operator=(const stream_acceptor &other)=delete;

        void close()
        { acc_.close(); }

        stream_type accept() {
            stream_type s;
            boost::system::error_code ec;
            accept(s, ec);
            return s;
        }

        stream_type accept(boost::system::error_code &ec) {
            stream_type s;
            accept(s, ec);
            return s;
        }

        boost::system::error_code accept(stream_type &s) {
            boost::system::error_code ec;
            accept(s, ec);
            return ec;
        }

        void accept(stream_type &s, boost::system::error_code &ec)
        { acc_.async_accept(*(s.rdbuf()), asio::yield[ec]); }

        stream_type operator()()
        { return accept(); }

        stream_type operator()(boost::system::error_code &ec)
        { return accept(ec); }

        boost::system::error_code operator()(stream_type &s)
        { return accept(s); }

        void operator()(stream_type &s, boost::system::error_code &ec)
        { accept(s, ec); }

        acceptor_type acc_;
    };

    template<typename Stream>
    struct stream_traits {
        typedef Stream stream_type;
        typedef typename stream_type::stream_type socket_type;
        typedef stream_acceptor<stream_type> acceptor_type;
        typedef typename acceptor_type::endpoint_type endpoint_type;
        // HACK:
        typedef int arg_type;
        static std::unique_ptr<stream_type> construct(arg_type *arg) {
            return std::unique_ptr<stream_type>(new stream_type());
        }
    };

    template<typename Stream>
    struct listener {
        typedef Stream stream_type;
        typedef stream_traits<Stream> traits_type;
        typedef typename traits_type::acceptor_type acceptor_type;
        typedef typename traits_type::endpoint_type endpoint_type;
        typedef typename traits_type::arg_type arg_type;

        template<typename std::enable_if<std::is_move_constructible<Stream>::value>::type* = nullptr >
        listener(const std::string &access_point)
        : ep_(detail::make_endpoint<endpoint_type>(access_point))
        {}

        // For SSL, arg type is ssl::context
        template<typename Arg>
        listener(Arg &arg, const std::string &access_point)
        : arg_(&arg)
        , ep_(detail::make_endpoint<endpoint_type>(access_point))
        {}

        endpoint_type endpoint() const {
            return ep_;
        }

        // Start and join, other fiber may stop the listener
        template<typename F>
        boost::system::error_code operator()(F f) {
            try {
                start(f);
                join();
            } catch(boost::system::system_error &e) {
                return e.code();
            }
            return boost::system::error_code();
        }

        // Start listener
        template<typename F>
        void start(F f) {
            if(!stop_signal_) {
                stop_signal_.reset(new promise<void>);
                if(!acceptor_fiber_) {
                    acceptor_fiber_.reset(new fiber(&listener::acceptor_fiber<F>,
                                                    this,
                                                    ep_,
                                                    stop_signal_.get(),
                                                    f));
                }
            }
        }

        void stop() {
            if(stop_signal_) {
                stop_signal_->set_value();
            }
        }

        void join() {
            if (acceptor_fiber_) {
                acceptor_fiber_->join();
                acceptor_fiber_.reset();
                stop_signal_.reset();
            }
        }

    private:
        template<typename F>
        void acceptor_fiber(const endpoint_type & e, promise<void> *p, F f) {
            acceptor_type acc(e);
            boost::system::error_code ec;
            fiber watchdog(fiber::attributes(fiber::attributes::stick_with_parent),
                           &listener::acceptor_watchdog_fiber,
                           this,
                           p,
                           std::ref(acc));
            while(!ec) {
                std::unique_ptr<stream_type> s(traits_type::construct(arg_));
                acc(*s, ec);
                if(ec) break;
                fiber([f](std::unique_ptr<stream_type> str){
                    f(*str);
                }, std::move(s)).detach();
            }
            watchdog.join();
        }

        void acceptor_watchdog_fiber(promise<void> *p, acceptor_type &acc) {
            p->get_future().wait();
            acc.close();
        }

        arg_type *arg_=nullptr;
        endpoint_type ep_;
        std::unique_ptr<fiber> acceptor_fiber_;
        std::unique_ptr<promise<void>> stop_signal_;
    };

    // TCP
    typedef iostream<boost::asio::ip::tcp::socket> tcp_stream;
    typedef stream_acceptor<tcp_stream> tcp_stream_acceptor;
    typedef listener<tcp_stream> tcp_listener;

#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
    // Unix-domain socket
    typedef iostream<boost::asio::local::stream_protocol::socket> local_stream;
    typedef stream_acceptor<local_stream> local_stream_acceptor;
    typedef listener<local_stream> local_listener;
#endif

#ifdef BOOST_ASIO_HAS_POSIX_STREAM_DESCRIPTOR
    // POSIX native stream
    typedef iostream<boost::asio::posix::stream_descriptor> posix_stream;
#endif

#ifdef BOOST_ASIO_HAS_WINDOWS_STREAM_HANDLE
    // Windows stream handle
    typedef iostream<boost::asio::windows::stream_handle> handle_stream;
#endif

}}  // End of namespace fibio::stream

namespace fibio {
    using stream::tcp_stream;
    using stream::tcp_stream_acceptor;
    using stream::tcp_listener;

#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
    using stream::local_stream;
    using stream::local_stream_acceptor;
    using stream::local_listener;
#endif

#ifdef BOOST_ASIO_HAS_POSIX_STREAM_DESCRIPTOR
    using stream::posix_stream;
#endif

#ifdef BOOST_ASIO_HAS_WINDOWS_STREAM_HANDLE
    using stream::handle_stream;
#endif
}

#endif
<|endoftext|>