<commit_before>// This file is part of the dune-stuff project:
//   https://users.dune-project.org/projects/dune-stuff
// Copyright holders: Rene Milk, Felix Schindler
// License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)
//
// Contributors: Sven Kaulmann

#ifndef DUNE_STUFF_DISCRETEFUNCTION_PROJECTION_HETEROGENOUS_HH
#define DUNE_STUFF_DISCRETEFUNCTION_PROJECTION_HETEROGENOUS_HH

#include <dune/common/fvector.hh>
#include <dune/grid/common/backuprestore.hh>
#include <dune/grid/common/grid.hh>
#include <dune/grid/common/entity.hh>
#include <dune/stuff/common/ranges.hh>
#include <dune/stuff/aliases.hh>
#include <dune/stuff/fem/namespace.hh>

#if HAVE_DUNE_FEM
  #include <dune/fem/function/common/discretefunction.hh>
  #include <dune/fem/quadrature/cachingquadrature.hh>
  #include <dune/fem/space/finitevolume.hh>
  #include <dune/fem/space/lagrange.hh>
#endif

#include <dune/grid/io/file/vtk/function.hh>

#include <dune/stuff/grid/search.hh>

namespace Dune {
namespace Stuff {

#if HAVE_DUNE_FEM

template< class F, class G, int p, template< class > class S >
std::vector<typename Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>::DomainType>
global_evaluation_points(const Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>& space,
         const typename Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>::EntityType& target_entity)
{
  const auto& target_lagrangepoint_set = space.lagrangePointSet(target_entity);
  const auto& target_geometry = target_entity.geometry();
  const auto quadNop = target_lagrangepoint_set.nop();
  std::vector<typename Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>::DomainType> points(quadNop);
  for(size_t qP = 0; qP < quadNop ; ++qP) {
    points[qP] = target_geometry.global(target_lagrangepoint_set.point(qP));
  }
  return points;
}

template< class F, class G, int p, template< class > class S >
std::vector<typename Dune::Fem::FiniteVolumeSpace<F,G,p,S>::DomainType>
global_evaluation_points(const Dune::Fem::FiniteVolumeSpace<F,G,p,S>& /*space*/,
         const typename Dune::Fem::FiniteVolumeSpace<F,G,p,S>::EntityType& target_entity)
{
  assert(false);
  typedef std::vector<typename Dune::Fem::FiniteVolumeSpace<F,G,p,S>::DomainType> Ret;
  return Ret(1, target_entity.geometry().center());
}

template< template< class > class SearchStrategy = Grid::EntityInlevelSearch >
class HeterogenousProjection
{
public:
  /** If your SearchStrategy only takes a leafview of the source grid, you may use this signature.
   * Otherwise you'll have to provide an instance of the SearchStrategy to the method below
   **/
  template < class SourceDFImp, class TargetDFImp >
  static void project(const Dune::Fem::DiscreteFunctionInterface<SourceDFImp>& source,
                      Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& target)
  {
    SearchStrategy<typename SourceDFImp::GridType::LeafGridView> search(source.gridPart().grid().leafView());
    project(source, target, search);
  }

  //! signature for non-default SearchStrategy constructions
  template < class SourceDFImp, class TargetDFImp, class SearchStrategyImp >
  static void project(const Dune::Fem::DiscreteFunctionInterface<SourceDFImp>& source,
                      Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& target,
                      SearchStrategyImp& search)
  {
    typedef typename TargetDFImp::DiscreteFunctionSpaceType TargetDiscreteFunctionSpaceType;
    static const int target_dimRange = TargetDiscreteFunctionSpaceType::dimRange;

    const auto& space =  target.space();

    // set all DoFs to infinity
    preprocess(target);

    const auto endit = space.end();
    for(auto it = space.begin(); it != endit ; ++it)
    {
      const auto& target_entity = *it;
      auto target_local_function = target.localFunction(target_entity);
      const auto global_quads = global_evaluation_points(space, target_entity);
      const auto evaluation_entity_ptrs = search(global_quads);
      assert(evaluation_entity_ptrs.size() >= global_quads.size());

      int k = 0;
      typename TargetDiscreteFunctionSpaceType::RangeType source_value;
      for(size_t qP = 0; qP < global_quads.size() ; ++qP)
      {
        if(std::isinf(target_local_function[ k ]))
        {
          const auto& source_entity_unique_ptr = evaluation_entity_ptrs[qP];
          if (source_entity_unique_ptr) {
            const auto& source_entity_ptr = (*source_entity_unique_ptr);
            const auto& source_geometry = source_entity_ptr->geometry();
            const auto& global_point = global_quads[qP];
            const auto& source_local_point = source_geometry.local(global_point);
            const auto& source_local_function = source.localFunction(*source_entity_ptr);
            source_local_function.evaluate(source_local_point, source_value);
            for(int i = 0; i < target_dimRange; ++i, ++k)
              setDofValue(target_local_function[k], source_value[i]);
          }
          else {
            DUNE_THROW(InvalidStateException, "Did not find the local lagrange point in the source mesh!");
          }
        }
        else
          k += target_dimRange;
      }
    }
    postprocess(target);

  } // ... project(...)

protected:
  template<class TargetDFImp>
  static void preprocess(Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& func) {
    typedef typename TargetDFImp::DofType TargetDofType;
    const auto infinity = DSC::numeric_limits< TargetDofType >::infinity();
    // set all DoFs to infinity
    const auto dend = func.dend();
    for( auto dit = func.dbegin(); dit != dend; ++dit )
      *dit = infinity;
  }

  template<class DofType, class SourceType >
  static void setDofValue(DofType& dof, const SourceType& value)
  {
    dof = value;
  }

  template<class TargetDFImp>
  static void postprocess(typename Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& /*func*/) { return; }

}; // class HeterogenousProjection

template< template< class > class SearchStrategy = Grid::EntityInlevelSearch >
class MsFEMProjection {
public:
  //! signature for non-default SearchStrategy constructions
  template < class SourceDFImp, class TargetDFImp, class SearchStrategyImp >
  static void project(const Dune::Fem::DiscreteFunctionInterface<SourceDFImp>& source,
                      Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& target,
                      SearchStrategyImp& search)
  {
    typedef typename TargetDFImp::DiscreteFunctionSpaceType TargetDiscreteFunctionSpaceType;
    static const int target_dimRange = TargetDiscreteFunctionSpaceType::dimRange;

    const auto& space =  target.space();

    // set all DoFs to infinity
    preprocess(target);

    const auto endit = space.end();
    for(auto it = space.begin(); it != endit ; ++it)
    {
      const auto& target_entity = *it;
      auto target_local_function = target.localFunction(target_entity);
      const auto global_quads = global_evaluation_points(space, target_entity);
      const auto evaluation_entity_ptrs = search(global_quads);
      assert(evaluation_entity_ptrs.size() >= global_quads.size());

      int k = 0;
      typename TargetDiscreteFunctionSpaceType::RangeType source_value;
      for(size_t qP = 0; qP < global_quads.size() ; ++qP)
      {
          const auto& source_entity_unique_ptr = evaluation_entity_ptrs[qP];
          if (source_entity_unique_ptr) {
            const auto& source_entity_ptr = (*source_entity_unique_ptr);
            const auto& source_geometry = source_entity_ptr->geometry();
            const auto& global_point = global_quads[qP];
            const auto& source_local_point = source_geometry.local(global_point);
            const auto& source_local_function = source.localFunction(*source_entity_ptr);
            source_local_function.evaluate(source_local_point, source_value);
            for(int i = 0; i < target_dimRange; ++i, ++k)
              setDofValue(target_local_function[k], source_value[i]);
          }
          else {
            DUNE_THROW(InvalidStateException, "Did not find the local lagrange point in the source mesh!");
          }
        }
      }

    postprocess(target);

  } // ... project(...)

protected:
  template<class TargetDFImp>
  static void preprocess(Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& func) {

    // set all DoFs to zero
    const auto dend = func.dend();
    for( auto dit = func.dbegin(); dit != dend; ++dit )
      *dit = 0.0;
  }

  template<class DofType, class SourceType >
  static void setDofValue(DofType& dof, const SourceType& value)
  {
    dof += value;
  }

  template<class TargetDFImp>
  static void postprocess(typename Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& func) {
    // compute node to entity relations
    typedef Dune::Fem::DiscreteFunctionInterface<TargetDFImp> DiscFuncType;
    const static int dimension = DiscFuncType::DiscreteFunctionSpaceType::GridPartType::GridType::dimension;
    std::vector<int> nodeToEntity(func.space().gridPart().grid().size(dimension), 0);
    identifySharedNodes(func.space().gridPart(), nodeToEntity);

    const auto dend = func.dend();
    auto factorsIt = nodeToEntity.begin();
    for (auto dit = func.dbegin(); dit!=dend; ++dit) {
      assert(factorsIt!=nodeToEntity.end());
      assert(*factorsIt>0);
      *dit /= *factorsIt;
      ++factorsIt;
    }
    return;
  }

  template<class GridPartType, class MapType>
  static void identifySharedNodes(const GridPartType& gridPart, MapType& map) {
    typedef typename GridPartType::GridType GridType;
    const auto& indexSet = gridPart.indexSet();

    for (auto& entity : DSC::viewRange(gridPart.grid().leafGridView())) {
      int number_of_nodes_in_entity = entity.template count<GridType::dimension>();
      for (int i = 0; i < number_of_nodes_in_entity; ++i) {
        const auto node = entity.template subEntity<GridType::dimension>(i);
        const auto global_index_node = indexSet.index(*node);

        // make sure we don't access non-existing elements
        assert(map.size() > global_index_node);
        ++map[global_index_node];
      }
    }
  }

};
#endif // HAVE_DUNE_FEM

} // namespace Stuff
} // namespace Dune

#endif // DUNE_STUFF_DISCRETEFUNCTION_PROJECTION_HETEROGENOUS_HH
<commit_msg>doc and whitespace<commit_after>// This file is part of the dune-stuff project:
//   https://users.dune-project.org/projects/dune-stuff
// Copyright holders: Rene Milk, Felix Schindler
// License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)
//
// Contributors: Sven Kaulmann

#ifndef DUNE_STUFF_DISCRETEFUNCTION_PROJECTION_HETEROGENOUS_HH
#define DUNE_STUFF_DISCRETEFUNCTION_PROJECTION_HETEROGENOUS_HH

#include <dune/common/fvector.hh>
#include <dune/grid/common/backuprestore.hh>
#include <dune/grid/common/grid.hh>
#include <dune/grid/common/entity.hh>
#include <dune/stuff/common/ranges.hh>
#include <dune/stuff/aliases.hh>
#include <dune/stuff/fem/namespace.hh>

#if HAVE_DUNE_FEM
  #include <dune/fem/function/common/discretefunction.hh>
  #include <dune/fem/quadrature/cachingquadrature.hh>
  #include <dune/fem/space/finitevolume.hh>
  #include <dune/fem/space/lagrange.hh>
#endif

#include <dune/grid/io/file/vtk/function.hh>

#include <dune/stuff/grid/search.hh>

namespace Dune {
namespace Stuff {

#if HAVE_DUNE_FEM

template< class F, class G, int p, template< class > class S >
std::vector<typename Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>::DomainType>
global_evaluation_points(const Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>& space,
         const typename Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>::EntityType& target_entity)
{
  const auto& target_lagrangepoint_set = space.lagrangePointSet(target_entity);
  const auto& target_geometry = target_entity.geometry();
  const auto quadNop = target_lagrangepoint_set.nop();
  std::vector<typename Dune::Fem::LagrangeDiscreteFunctionSpace<F,G,p,S>::DomainType> points(quadNop);
  for(size_t qP = 0; qP < quadNop ; ++qP) {
    points[qP] = target_geometry.global(target_lagrangepoint_set.point(qP));
  }
  return points;
}

template< class F, class G, int p, template< class > class S >
std::vector<typename Dune::Fem::FiniteVolumeSpace<F,G,p,S>::DomainType>
global_evaluation_points(const Dune::Fem::FiniteVolumeSpace<F,G,p,S>& /*space*/,
         const typename Dune::Fem::FiniteVolumeSpace<F,G,p,S>::EntityType& target_entity)
{
  assert(false);
  typedef std::vector<typename Dune::Fem::FiniteVolumeSpace<F,G,p,S>::DomainType> Ret;
  return Ret(1, target_entity.geometry().center());
}

template< template< class > class SearchStrategy = Grid::EntityInlevelSearch >
class HeterogenousProjection
{
public:
  /** If your SearchStrategy only takes a leafview of the source grid, you may use this signature.
   * Otherwise you'll have to provide an instance of the SearchStrategy to the method below
   **/
  template < class SourceDFImp, class TargetDFImp >
  static void project(const Dune::Fem::DiscreteFunctionInterface<SourceDFImp>& source,
                      Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& target)
  {
    SearchStrategy<typename SourceDFImp::GridType::LeafGridView> search(source.gridPart().grid().leafView());
    project(source, target, search);
  }

  //! signature for non-default SearchStrategy constructions
  template < class SourceDFImp, class TargetDFImp, class SearchStrategyImp >
  static void project(const Dune::Fem::DiscreteFunctionInterface<SourceDFImp>& source,
                      Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& target,
                      SearchStrategyImp& search)
  {
    typedef typename TargetDFImp::DiscreteFunctionSpaceType TargetDiscreteFunctionSpaceType;
    static const int target_dimRange = TargetDiscreteFunctionSpaceType::dimRange;

    const auto& space =  target.space();

    // set all DoFs to infinity
    preprocess(target);

    const auto endit = space.end();
    for(auto it = space.begin(); it != endit ; ++it)
    {
      const auto& target_entity = *it;
      auto target_local_function = target.localFunction(target_entity);
      const auto global_quads = global_evaluation_points(space, target_entity);
      const auto evaluation_entity_ptrs = search(global_quads);
      assert(evaluation_entity_ptrs.size() >= global_quads.size());

      int k = 0;
      typename TargetDiscreteFunctionSpaceType::RangeType source_value;
      for(size_t qP = 0; qP < global_quads.size() ; ++qP)
      {
        if(std::isinf(target_local_function[ k ]))
        {
          const auto& source_entity_unique_ptr = evaluation_entity_ptrs[qP];
          if (source_entity_unique_ptr) {
            const auto& source_entity_ptr = (*source_entity_unique_ptr);
            const auto& source_geometry = source_entity_ptr->geometry();
            const auto& global_point = global_quads[qP];
            const auto& source_local_point = source_geometry.local(global_point);
            const auto& source_local_function = source.localFunction(*source_entity_ptr);
            source_local_function.evaluate(source_local_point, source_value);
            for(int i = 0; i < target_dimRange; ++i, ++k)
              setDofValue(target_local_function[k], source_value[i]);
          }
          else {
            DUNE_THROW(InvalidStateException, "Did not find the local lagrange point in the source mesh!");
          }
        }
        else
          k += target_dimRange;
      }
    }
    postprocess(target);

  } // ... project(...)

protected:
  template<class TargetDFImp>
  static void preprocess(Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& func) {
    typedef typename TargetDFImp::DofType TargetDofType;
    const auto infinity = DSC::numeric_limits< TargetDofType >::infinity();
    // set all DoFs to infinity
    const auto dend = func.dend();
    for( auto dit = func.dbegin(); dit != dend; ++dit )
      *dit = infinity;
  }

  template<class DofType, class SourceType >
  static void setDofValue(DofType& dof, const SourceType& value)
  {
    dof = value;
  }

  template<class TargetDFImp>
  static void postprocess(typename Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& /*func*/) { return; }

}; // class HeterogenousProjection

template< template< class > class SearchStrategy = Grid::EntityInlevelSearch >
class MsFEMProjection {
public:
  //! signature for non-default SearchStrategy constructions
  template < class SourceDFImp, class TargetDFImp, class SearchStrategyImp >
  static void project(const Dune::Fem::DiscreteFunctionInterface<SourceDFImp>& source,
                      Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& target,
                      SearchStrategyImp& search)
  {
    typedef typename TargetDFImp::DiscreteFunctionSpaceType TargetDiscreteFunctionSpaceType;
    static const int target_dimRange = TargetDiscreteFunctionSpaceType::dimRange;

    const auto& space =  target.space();

    preprocess(target);

    const auto endit = space.end();
    for(auto it = space.begin(); it != endit ; ++it)
    {
      const auto& target_entity = *it;
      auto target_local_function = target.localFunction(target_entity);
      const auto global_quads = global_evaluation_points(space, target_entity);
      const auto evaluation_entity_ptrs = search(global_quads);
      assert(evaluation_entity_ptrs.size() >= global_quads.size());

      int k = 0;
      typename TargetDiscreteFunctionSpaceType::RangeType source_value;
      for(size_t qP = 0; qP < global_quads.size() ; ++qP)
      {
          const auto& source_entity_unique_ptr = evaluation_entity_ptrs[qP];
          if (source_entity_unique_ptr) {
            const auto& source_entity_ptr = (*source_entity_unique_ptr);
            const auto& source_geometry = source_entity_ptr->geometry();
            const auto& global_point = global_quads[qP];
            const auto& source_local_point = source_geometry.local(global_point);
            const auto& source_local_function = source.localFunction(*source_entity_ptr);
            source_local_function.evaluate(source_local_point, source_value);
            for(int i = 0; i < target_dimRange; ++i, ++k)
              setDofValue(target_local_function[k], source_value[i]);
          }
          else {
            DUNE_THROW(InvalidStateException, "Did not find the local lagrange point in the source mesh!");
          }
        }
      }

    postprocess(target);
  } // ... project(...)

protected:
  template<class TargetDFImp>
  static void preprocess(Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& func) {
    // set all DoFs to zero
    const auto dend = func.dend();
    for( auto dit = func.dbegin(); dit != dend; ++dit )
      *dit = 0.0;
  }

  template<class DofType, class SourceType >
  static void setDofValue(DofType& dof, const SourceType& value)
  {
    dof += value;
  }

  template<class TargetDFImp>
  static void postprocess(typename Dune::Fem::DiscreteFunctionInterface<TargetDFImp>& func) {
    // compute node to entity relations
    typedef Dune::Fem::DiscreteFunctionInterface<TargetDFImp> DiscFuncType;
    const static int dimension = DiscFuncType::DiscreteFunctionSpaceType::GridPartType::GridType::dimension;
    std::vector<int> nodeToEntity(func.space().gridPart().grid().size(dimension), 0);
    identifySharedNodes(func.space().gridPart(), nodeToEntity);

    const auto dend = func.dend();
    auto factorsIt = nodeToEntity.begin();
    for (auto dit = func.dbegin(); dit!=dend; ++dit) {
      assert(factorsIt!=nodeToEntity.end());
      assert(*factorsIt>0);
      *dit /= *factorsIt;
      ++factorsIt;
    }
    return;
  }

  template<class GridPartType, class MapType>
  static void identifySharedNodes(const GridPartType& gridPart, MapType& map) {
    typedef typename GridPartType::GridType GridType;
    const auto& indexSet = gridPart.indexSet();

    for (auto& entity : DSC::viewRange(gridPart.grid().leafGridView())) {
      int number_of_nodes_in_entity = entity.template count<GridType::dimension>();
      for (int i = 0; i < number_of_nodes_in_entity; ++i) {
        const auto node = entity.template subEntity<GridType::dimension>(i);
        const auto global_index_node = indexSet.index(*node);

        // make sure we don't access non-existing elements
        assert(map.size() > global_index_node);
        ++map[global_index_node];
      }
    }
  }

};
#endif // HAVE_DUNE_FEM

} // namespace Stuff
} // namespace Dune

#endif // DUNE_STUFF_DISCRETEFUNCTION_PROJECTION_HETEROGENOUS_HH
<|endoftext|>


<commit_before>#include "souper/Infer/DataflowPruning.h"

namespace souper {
namespace {
using souper::dataflow::EvalValue;
using souper::dataflow::ValueCache;

EvalValue evaluateInstRecursive(souper::Inst* Inst,
                                std::vector<EvalValue> args) {
  switch (Inst->K) {
    case souper::Inst::Const:
      return {Inst->Val};
    case souper::Inst::UntypedConst:
      return {Inst->Val};
    case souper::Inst::Var:
      llvm_unreachable("Should get value from cache without reaching here");
    case souper::Inst::Add:
      return {args[0].getValue() + args[1].getValue()};
    case souper::Inst::Sub:
      return {args[0].getValue() - args[1].getValue()};
    case souper::Inst::Mul:
      return {args[0].getValue() * args[1].getValue()};
    case souper::Inst::UDiv:
      return {args[0].getValue().udiv(args[1].getValue())};
    case souper::Inst::SDiv:
      return {args[0].getValue().sdiv(args[1].getValue())};
    case souper::Inst::URem:
      return {args[0].getValue().urem(args[1].getValue())};
    case souper::Inst::SRem:
      return {args[0].getValue().srem(args[1].getValue())};
    case souper::Inst::And:
      return {args[0].getValue() & args[1].getValue()};
    case souper::Inst::Or:
      return {args[0].getValue() | args[1].getValue()};
    case souper::Inst::Xor:
      return {args[0].getValue() ^ args[1].getValue()};
    case souper::Inst::Shl:
      return {args[0].getValue() << args[1].getValue()};
    case souper::Inst::LShr:
      return {args[0].getValue().lshr(args[1].getValue())};
    case souper::Inst::AShr:
      return {args[0].getValue().ashr(args[1].getValue())};
    // TODO: Handle NSW/NUW/NW variants of instructions
    case souper::Inst::Select: {
      if (!args[0].hasValue()) {
        return args[0];
      } else {
        bool cond = args[0].getValue().getBoolValue();
        return cond ? args[1] : args[2];
      }
    }
    case souper::Inst::ZExt:
      return {args[0].getValue().zext(Inst->Width)};
    case souper::Inst::SExt:
      return {args[0].getValue().sext(Inst->Width)};
    case souper::Inst::Trunc:
      return {args[0].getValue().trunc(Inst->Width)};
    case souper::Inst::Eq:
      return {{1, args[0].getValue() == args[1].getValue()}};
    case souper::Inst::Ne:
      return {{1, args[0].getValue() != args[1].getValue()}};
    case souper::Inst::Ult:
      return {{1, args[0].getValue().ult(args[1].getValue())}};
    case souper::Inst::Slt:
      return {{1, args[0].getValue().slt(args[1].getValue())}};
    case souper::Inst::Ule:
      return {{1, args[0].getValue().ule(args[1].getValue())}};
    case souper::Inst::Sle:
      return {{1, args[0].getValue().sle(args[1].getValue())}};
    case souper::Inst::CtPop:
      return {llvm::APInt(Inst->Width, args[0].getValue().countPopulation())};
    case souper::Inst::Ctlz:
      return {llvm::APInt(Inst->Width,
                          args[0].getValue().countLeadingZeros())};
    case souper::Inst::Cttz:
      return {llvm::APInt(Inst->Width,
                          args[0].getValue().countTrailingZeros())};
    case souper::Inst::BSwap:
      return {args[0].getValue().byteSwap()};
    default:
      return EvalValue(); // Indicates an 'unavailable' value
  }
}
}

// Errs on the side of an 'unavailable' result
// TODO: Some instructions unimplemented
EvalValue souper::dataflow::evaluateInst(souper::Inst* Root,
                                         ValueCache &Cache) {
  std::vector<EvalValue> EvaluatedArgs;
  EvalValue Result;

  if (Root->K == souper::Inst::Var) {
      Result = Cache[Root];
  } else {
    for (auto &&I : Root->Ops) {
      auto eval = evaluateInst(I, Cache);
      if (!eval.hasValue() && Root->K != souper::Inst::Select) {
        return Result;
        // result is `Unavailable` if args fail to evaluate
      }
      EvaluatedArgs.push_back(eval);
    }
    Result = evaluateInstRecursive(Root, EvaluatedArgs);
  }
  return Result;
}

EvalValue getValue(Inst *I, ValueCache &C) {
  if (I->K == souper::Inst::Const) {
      return {I->Val};
  } else if (I->K == souper::Inst::Var
      || I->K == souper::Inst::ReservedConst
      || I->K == souper::Inst::ReservedInst) {
    if (I->Name != "" && C.find(I) != C.end()) {
      return C[I];
    } else {
      return EvalValue(); // unavailable
    }
  }
  return EvalValue();
}

llvm::KnownBits dataflow::findKnownBits(Inst* I, ValueCache& C) {
  llvm::KnownBits Result(I->Width);
  switch(I->K) {
    case souper::Inst::Const:
    case souper::Inst::Var : {
      EvalValue V = getValue(I, C);
      if (V.hasValue()) {
        Result.One = V.getValue();
        Result.Zero = ~V.getValue();
        return Result;
      } else {
        return Result;
      }
    }
    case souper::Inst::Shl : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1V = getValue(I->Ops[1], C);
      if (Op1V.hasValue()) {
        Op0KB.One <<= Op1V.getValue();
        Op0KB.Zero <<= Op1V.getValue();
        Op0KB.Zero.setLowBits(Op1V.getValue().getLimitedValue());
        // setLowBits takes an unsiged int, so getLimitedValue is harmless
        return Op0KB;
      } else {
        return Result;
      }
    }
    case souper::Inst::And : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1KB = findKnownBits(I->Ops[1], C);

      Op0KB.One &= Op1KB.One;
      Op0KB.Zero |= Op1KB.Zero;
      return Op0KB;
    }
    case souper::Inst::Or : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1KB = findKnownBits(I->Ops[1], C);

      Op0KB.One |= Op1KB.One;
      Op0KB.Zero &= Op1KB.Zero;
      return Op0KB;
    }
    case souper::Inst::Xor : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1KB = findKnownBits(I->Ops[1], C);
      llvm::APInt KnownZeroOut =
        (Op0KB.Zero & Op1KB.Zero) | (Op0KB.One & Op1KB.One);
      Op0KB.One = (Op0KB.Zero & Op1KB.One) | (Op0KB.One & Op1KB.Zero);
      Op0KB.Zero = std::move(KnownZeroOut);
      // ^ logic copied from LLVM ValueTracking.cpp
      return Op0KB;
    }
    default :
      return Result;
  }
}

llvm::ConstantRange dataflow::findConstantRange(souper::Inst* I,
                                                souper::ValueCache& C) {
  llvm::ConstantRange result(I->Width);
  switch (I->K) {
    case souper::Inst::Const:
    case souper::Inst::Var : {
      EvalValue V = getValue(I, C);
      if (V.hasValue()) {
        return llvm::ConstantRange(V.getValue());
      } else {
        return result; // Whole range
      }
    }
    case souper::Inst::Add: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.add(R1);
    }
    case souper::Inst::AddNSW: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto V1 = getValue(I->Ops[1], C);
      if (V1.hasValue()) {
        return R0.addWithNoSignedWrap(V1.getValue());
      } else {
        return result; // full range, can we do better?
      }
    }
    case souper::Inst::Sub: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.sub(R1);
    }
    case souper::Inst::Mul: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.multiply(R1);
    }
    case souper::Inst::And: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.binaryAnd(R1);
    }
    case souper::Inst::Or: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.binaryOr(R1);
    }
    case souper::Inst::Shl: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.shl(R1);
    }
    case souper::Inst::AShr: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.ashr(R1);
    }
    case souper::Inst::LShr: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.lshr(R1);
    }
    case souper::Inst::UDiv: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.udiv(R1);
    }
//     case souper::Inst::SDiv: {
//       auto R0 = FindConstantRange(I->Ops[0], C);
//       auto R1 = FindConstantRange(I->Ops[1], C);
//       return R0.sdiv(R1); // unimplemented
//     }
    // TODO: Xor pattern for not, truncs and extends, etc
    default:
      return result;
  }
}

bool dataflow::ValueAnalysis::isInfeasible(souper::Inst* RHS) {
  for (int I = 0; I < Inputs.size(); ++I) {
    auto C = LHSValues[I];
    if (C.hasValue()) {
      auto CR = findConstantRange(RHS, Inputs[I]);
      if (!CR.contains(C.getValue())) {
        return true;
      }
      auto KB = findKnownBits(RHS, Inputs[I]);
      if ((KB.Zero & C.getValue()) != 0 || (KB.One & ~C.getValue()) != 0) {
        return true;
      }

      auto RHSV = evaluateInst(RHS, Inputs[I]);
      if (RHSV.hasValue()) {
        if (C.getValue() != RHSV.getValue()) {
          return true;
        }
      }
    }
  }
  return false;
}
dataflow::DataflowPruningManager::DataflowPruningManager(
  souper::Inst* LHS, std::vector<Inst *> &Inputs, unsigned StatsLevel)
  : VA(LHS, generateInputSets(Inputs)), NumPruned(0), TotalGuesses(0) {
  if (StatsLevel > 1) {
    DataflowPrune= [this](Inst *I, std::vector<Inst *> &RI) {
      TotalGuesses++;
      if (VA.isInfeasible(I)) {
        NumPruned++;
        llvm::outs() << "Dataflow Pruned "
          << NumPruned << "/" << TotalGuesses << "\n";
        ReplacementContext RC;
        RC.printInst(I, llvm::outs(), true);
        return false;
      }
      return true;
    };
  } else if (StatsLevel == 1) {
      DataflowPrune= [this](Inst *I, std::vector<Inst *> &RI) {
      TotalGuesses++;
      if (VA.isInfeasible(I)) {
        NumPruned++;
        return false;
      }
      return true;
    };
  } else {
    DataflowPrune= [this](Inst *I, std::vector<Inst *> &RI) {
      return !VA.isInfeasible(I);
    };
  }
}

std::vector<ValueCache> dataflow::DataflowPruningManager::generateInputSets(
  std::vector<Inst *> &Inputs) {
  std::vector<dataflow::ValueCache> InputSets;

  dataflow::ValueCache Cache;
  int64_t Current = 0;
  for (auto &&I : Inputs) {
    if (I->K == souper::Inst::Var)
      Cache[I] = {llvm::APInt(I->Width, Current++)};
  }
  InputSets.push_back(Cache);

  Current = 2*Current + 1;
  for (auto &&I : Inputs) {
    if (I->K == souper::Inst::Var)
      Cache[I] = {llvm::APInt(I->Width, Current++)};
  }
  InputSets.push_back(Cache);

  return InputSets;
}

}
<commit_msg>Support trunc, sext, zext in findConstantRange<commit_after>#include "souper/Infer/DataflowPruning.h"

namespace souper {
namespace {
using souper::dataflow::EvalValue;
using souper::dataflow::ValueCache;

EvalValue evaluateInstRecursive(souper::Inst* Inst,
                                std::vector<EvalValue> args) {
  switch (Inst->K) {
    case souper::Inst::Const:
      return {Inst->Val};
    case souper::Inst::UntypedConst:
      return {Inst->Val};
    case souper::Inst::Var:
      llvm_unreachable("Should get value from cache without reaching here");
    case souper::Inst::Add:
      return {args[0].getValue() + args[1].getValue()};
    case souper::Inst::Sub:
      return {args[0].getValue() - args[1].getValue()};
    case souper::Inst::Mul:
      return {args[0].getValue() * args[1].getValue()};
    case souper::Inst::UDiv:
      return {args[0].getValue().udiv(args[1].getValue())};
    case souper::Inst::SDiv:
      return {args[0].getValue().sdiv(args[1].getValue())};
    case souper::Inst::URem:
      return {args[0].getValue().urem(args[1].getValue())};
    case souper::Inst::SRem:
      return {args[0].getValue().srem(args[1].getValue())};
    case souper::Inst::And:
      return {args[0].getValue() & args[1].getValue()};
    case souper::Inst::Or:
      return {args[0].getValue() | args[1].getValue()};
    case souper::Inst::Xor:
      return {args[0].getValue() ^ args[1].getValue()};
    case souper::Inst::Shl:
      return {args[0].getValue() << args[1].getValue()};
    case souper::Inst::LShr:
      return {args[0].getValue().lshr(args[1].getValue())};
    case souper::Inst::AShr:
      return {args[0].getValue().ashr(args[1].getValue())};
    // TODO: Handle NSW/NUW/NW variants of instructions
    case souper::Inst::Select: {
      if (!args[0].hasValue()) {
        return args[0];
      } else {
        bool cond = args[0].getValue().getBoolValue();
        return cond ? args[1] : args[2];
      }
    }
    case souper::Inst::ZExt:
      return {args[0].getValue().zext(Inst->Width)};
    case souper::Inst::SExt:
      return {args[0].getValue().sext(Inst->Width)};
    case souper::Inst::Trunc:
      return {args[0].getValue().trunc(Inst->Width)};
    case souper::Inst::Eq:
      return {{1, args[0].getValue() == args[1].getValue()}};
    case souper::Inst::Ne:
      return {{1, args[0].getValue() != args[1].getValue()}};
    case souper::Inst::Ult:
      return {{1, args[0].getValue().ult(args[1].getValue())}};
    case souper::Inst::Slt:
      return {{1, args[0].getValue().slt(args[1].getValue())}};
    case souper::Inst::Ule:
      return {{1, args[0].getValue().ule(args[1].getValue())}};
    case souper::Inst::Sle:
      return {{1, args[0].getValue().sle(args[1].getValue())}};
    case souper::Inst::CtPop:
      return {llvm::APInt(Inst->Width, args[0].getValue().countPopulation())};
    case souper::Inst::Ctlz:
      return {llvm::APInt(Inst->Width,
                          args[0].getValue().countLeadingZeros())};
    case souper::Inst::Cttz:
      return {llvm::APInt(Inst->Width,
                          args[0].getValue().countTrailingZeros())};
    case souper::Inst::BSwap:
      return {args[0].getValue().byteSwap()};
    default:
      return EvalValue(); // Indicates an 'unavailable' value
  }
}
}

// Errs on the side of an 'unavailable' result
// TODO: Some instructions unimplemented
EvalValue souper::dataflow::evaluateInst(souper::Inst* Root,
                                         ValueCache &Cache) {
  std::vector<EvalValue> EvaluatedArgs;
  EvalValue Result;

  if (Root->K == souper::Inst::Var) {
      Result = Cache[Root];
  } else {
    for (auto &&I : Root->Ops) {
      auto eval = evaluateInst(I, Cache);
      if (!eval.hasValue() && Root->K != souper::Inst::Select) {
        return Result;
        // result is `Unavailable` if args fail to evaluate
      }
      EvaluatedArgs.push_back(eval);
    }
    Result = evaluateInstRecursive(Root, EvaluatedArgs);
  }
  return Result;
}

EvalValue getValue(Inst *I, ValueCache &C) {
  if (I->K == souper::Inst::Const) {
      return {I->Val};
  } else if (I->K == souper::Inst::Var
      || I->K == souper::Inst::ReservedConst
      || I->K == souper::Inst::ReservedInst) {
    if (I->Name != "" && C.find(I) != C.end()) {
      return C[I];
    } else {
      return EvalValue(); // unavailable
    }
  }
  return EvalValue();
}

llvm::KnownBits dataflow::findKnownBits(Inst* I, ValueCache& C) {
  llvm::KnownBits Result(I->Width);
  switch(I->K) {
    case souper::Inst::Const:
    case souper::Inst::Var : {
      EvalValue V = getValue(I, C);
      if (V.hasValue()) {
        Result.One = V.getValue();
        Result.Zero = ~V.getValue();
        return Result;
      } else {
        return Result;
      }
    }
    case souper::Inst::Shl : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1V = getValue(I->Ops[1], C);
      if (Op1V.hasValue()) {
        Op0KB.One <<= Op1V.getValue();
        Op0KB.Zero <<= Op1V.getValue();
        Op0KB.Zero.setLowBits(Op1V.getValue().getLimitedValue());
        // setLowBits takes an unsiged int, so getLimitedValue is harmless
        return Op0KB;
      } else {
        return Result;
      }
    }
    case souper::Inst::And : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1KB = findKnownBits(I->Ops[1], C);

      Op0KB.One &= Op1KB.One;
      Op0KB.Zero |= Op1KB.Zero;
      return Op0KB;
    }
    case souper::Inst::Or : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1KB = findKnownBits(I->Ops[1], C);

      Op0KB.One |= Op1KB.One;
      Op0KB.Zero &= Op1KB.Zero;
      return Op0KB;
    }
    case souper::Inst::Xor : {
      auto Op0KB = findKnownBits(I->Ops[0], C);
      auto Op1KB = findKnownBits(I->Ops[1], C);
      llvm::APInt KnownZeroOut =
        (Op0KB.Zero & Op1KB.Zero) | (Op0KB.One & Op1KB.One);
      Op0KB.One = (Op0KB.Zero & Op1KB.One) | (Op0KB.One & Op1KB.Zero);
      Op0KB.Zero = std::move(KnownZeroOut);
      // ^ logic copied from LLVM ValueTracking.cpp
      return Op0KB;
    }
    default :
      return Result;
  }
}

llvm::ConstantRange dataflow::findConstantRange(souper::Inst* I,
                                                souper::ValueCache& C) {
  llvm::ConstantRange result(I->Width);
  switch (I->K) {
    case souper::Inst::Const:
    case souper::Inst::Var : {
      EvalValue V = getValue(I, C);
      if (V.hasValue()) {
        return llvm::ConstantRange(V.getValue());
      } else {
        return result; // Whole range
      }
    }
    case souper::Inst::Trunc: {
      auto R0 = findConstantRange(I->Ops[0], C);
      return R0.truncate(I->Width);
    }
    case souper::Inst::SExt: {
      auto R0 = findConstantRange(I->Ops[0], C);
      return R0.signExtend(I->Width);
    }
    case souper::Inst::ZExt: {
      auto R0 = findConstantRange(I->Ops[0], C);
      return R0.zeroExtend(I->Width);
    }
    case souper::Inst::Add: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.add(R1);
    }
    case souper::Inst::AddNSW: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto V1 = getValue(I->Ops[1], C);
      if (V1.hasValue()) {
        return R0.addWithNoSignedWrap(V1.getValue());
      } else {
        return result; // full range, can we do better?
      }
    }
    case souper::Inst::Sub: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.sub(R1);
    }
    case souper::Inst::Mul: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.multiply(R1);
    }
    case souper::Inst::And: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.binaryAnd(R1);
    }
    case souper::Inst::Or: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.binaryOr(R1);
    }
    case souper::Inst::Shl: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.shl(R1);
    }
    case souper::Inst::AShr: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.ashr(R1);
    }
    case souper::Inst::LShr: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.lshr(R1);
    }
    case souper::Inst::UDiv: {
      auto R0 = findConstantRange(I->Ops[0], C);
      auto R1 = findConstantRange(I->Ops[1], C);
      return R0.udiv(R1);
    }
//     case souper::Inst::SDiv: {
//       auto R0 = FindConstantRange(I->Ops[0], C);
//       auto R1 = FindConstantRange(I->Ops[1], C);
//       return R0.sdiv(R1); // unimplemented
//     }
    // TODO: Xor pattern for not, truncs and extends, etc
    default:
      return result;
  }
}

bool dataflow::ValueAnalysis::isInfeasible(souper::Inst* RHS) {
  for (int I = 0; I < Inputs.size(); ++I) {
    auto C = LHSValues[I];
    if (C.hasValue()) {
      auto CR = findConstantRange(RHS, Inputs[I]);
      if (!CR.contains(C.getValue())) {
        return true;
      }
      auto KB = findKnownBits(RHS, Inputs[I]);
      if ((KB.Zero & C.getValue()) != 0 || (KB.One & ~C.getValue()) != 0) {
        return true;
      }

      auto RHSV = evaluateInst(RHS, Inputs[I]);
      if (RHSV.hasValue()) {
        if (C.getValue() != RHSV.getValue()) {
          return true;
        }
      }
    }
  }
  return false;
}
dataflow::DataflowPruningManager::DataflowPruningManager(
  souper::Inst* LHS, std::vector<Inst *> &Inputs, unsigned StatsLevel)
  : VA(LHS, generateInputSets(Inputs)), NumPruned(0), TotalGuesses(0) {
  if (StatsLevel > 1) {
    DataflowPrune= [this](Inst *I, std::vector<Inst *> &RI) {
      TotalGuesses++;
      if (VA.isInfeasible(I)) {
        NumPruned++;
        llvm::outs() << "Dataflow Pruned "
          << NumPruned << "/" << TotalGuesses << "\n";
        ReplacementContext RC;
        RC.printInst(I, llvm::outs(), true);
        return false;
      }
      return true;
    };
  } else if (StatsLevel == 1) {
      DataflowPrune= [this](Inst *I, std::vector<Inst *> &RI) {
      TotalGuesses++;
      if (VA.isInfeasible(I)) {
        NumPruned++;
        return false;
      }
      return true;
    };
  } else {
    DataflowPrune= [this](Inst *I, std::vector<Inst *> &RI) {
      return !VA.isInfeasible(I);
    };
  }
}

std::vector<ValueCache> dataflow::DataflowPruningManager::generateInputSets(
  std::vector<Inst *> &Inputs) {
  std::vector<dataflow::ValueCache> InputSets;

  dataflow::ValueCache Cache;
  int64_t Current = 0;
  for (auto &&I : Inputs) {
    if (I->K == souper::Inst::Var)
      Cache[I] = {llvm::APInt(I->Width, Current++)};
  }
  InputSets.push_back(Cache);

  Current = 2*Current + 1;
  for (auto &&I : Inputs) {
    if (I->K == souper::Inst::Var)
      Cache[I] = {llvm::APInt(I->Width, Current++)};
  }
  InputSets.push_back(Cache);

  return InputSets;
}

}
<|endoftext|>


<commit_before>#include "HotkeyManager.h"
#include "Logger.h"

HotkeyManager *HotkeyManager::instance = NULL;

HotkeyManager *HotkeyManager::Instance() {
    return instance;
}

HotkeyManager *HotkeyManager::Instance(HWND notifyWnd) {
    if (instance == NULL) {
        instance = new HotkeyManager();
        instance->Hook();
        instance->_notifyWnd = notifyWnd;
    }

    return instance;
}

HotkeyManager::HotkeyManager() :
_fixWin(false) {

}

HotkeyManager::~HotkeyManager() {
    while (true) {
        auto i = _keyCombinations.begin();
        if (_keyCombinations.size() > 0) {
            Unregister(*i);
        } else {
            break;
        }
    }

    Unhook();
}

void HotkeyManager::Shutdown() {
    delete instance;
    instance = NULL;
}

bool HotkeyManager::Hook() {
    _mouseHook = SetWindowsHookEx(WH_MOUSE_LL,
        LowLevelMouseProc, NULL, NULL);

    _keyHook = SetWindowsHookEx(WH_KEYBOARD_LL,
        LowLevelKeyboardProc, NULL, NULL);

    return _mouseHook && _keyHook;
}

bool HotkeyManager::Unhook() {
    BOOL unMouse = UnhookWindowsHookEx(_mouseHook);
    BOOL unKey = UnhookWindowsHookEx(_keyHook);
    return unMouse && unKey;
}

void HotkeyManager::Register(int keyCombination) {
    if (_keyCombinations.count(keyCombination) > 0) {
        CLOG(L"Hotkey combination [%d] already registered", keyCombination);
        return;
    } else {
        _keyCombinations.insert(keyCombination);
    }

    /* get VK_* value; includes unused bits */
    int vk = 0xFFFF & keyCombination;

    if ((keyCombination >> 20) > 0 /* uses a HKM_MOUSE_* flag */
        || vk == VK_LBUTTON
        || vk == VK_RBUTTON
        || vk == VK_MBUTTON) {

        /* mouse-based hotkeys; we are done */
        CLOG(L"Registered new mouse-based hotkey: %d", keyCombination);
        return;
    }

    /* keyboard-only hotkeys; use WinAPI */
    int mods = (0xF0000 & keyCombination) >> 16;
    if (!RegisterHotKey(_notifyWnd, keyCombination, mods, vk)) {
        CLOG(L"Failed to register hotkey [%d]\n"
            L"Mods: %d, VK: %d", keyCombination, mods, vk);
        return;
    }

    CLOG(L"Registered new keyboard-based hotkey: %d", keyCombination);
}

bool HotkeyManager::Unregister(int keyCombination) {
    CLOG(L"Unregistering hotkey combination: %d", keyCombination);

    if (_keyCombinations.count(keyCombination) <= 0) {
        QCLOG(L"Hotkey combination [%d] was not previously registered",
            keyCombination);
        return false;
    }

    _keyCombinations.erase(keyCombination);
    if ((keyCombination >> 20) == 0) {
        /* This hotkey isn't mouse-based; unregister with Windows */
        if (!UnregisterHotKey(_notifyWnd, keyCombination)) {
            CLOG(L"Failed to unregister hotkey: %d", keyCombination);
            return false;
        }
    }
    return true;
}

bool HotkeyManager::IsModifier(DWORD vk) {
    switch (vk) {
    case VK_MENU:
    case VK_LMENU:
    case VK_RMENU:
    case VK_CONTROL:
    case VK_LCONTROL:
    case VK_RCONTROL:
    case VK_SHIFT:
    case VK_LSHIFT:
    case VK_RSHIFT:
    case VK_LWIN:
    case VK_RWIN:
        return true;
    }
    return false;
}

int HotkeyManager::Modifiers() {
    int mods = 0;
    mods += (GetKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

int HotkeyManager::ModifiersAsync() {
    int mods = 0;
    mods += (GetAsyncKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetAsyncKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetAsyncKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetAsyncKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetAsyncKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

std::wstring HotkeyManager::HotkeysToModString(int combination,
    std::wstring separator) {

    std::wstring str = L"";
    if (combination & HKM_MOD_ALT) {
        str += VKToString(VK_MENU) + separator;
    }
    if (combination & HKM_MOD_CTRL) {
        str += VKToString(VK_CONTROL) + separator;
    }
    if (combination & HKM_MOD_SHF) {
        str += VKToString(VK_SHIFT) + separator;
    }
    if (combination & HKM_MOD_WIN) {
        str += L"Win" + separator;
    }

    return str;
}

std::wstring HotkeyManager::HotkeysToString(int combination,
    std::wstring separator) {

    std::wstring mods = HotkeysToModString(combination, separator);
    int vk = combination & 0xFF;
    std::wstring str = VKToString(vk);

    return mods + str;
}

std::wstring HotkeyManager::VKToString(unsigned int vk, bool extendedKey) {
    int extended = extendedKey ? 0x1 : 0x0;

    unsigned int scanCode = MapVirtualKey(vk, MAPVK_VK_TO_VSC);
    scanCode = scanCode << 16;
    scanCode |= extended << 24;
    wchar_t buf[256] = {};
    GetKeyNameText(scanCode, buf, 256);
    return std::wstring(buf);
}

LRESULT CALLBACK
HotkeyManager::KeyProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0) {
        if (wParam == WM_KEYUP) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT*) lParam;
            if ((kbInfo->vkCode == VK_LWIN || kbInfo->vkCode == VK_RWIN)
                && _fixWin) {
                /* WIN+Mouse combination used; we need to prevent the
                 * system from only seeing a WIN keypress (and usually
                 * popping up the start menu).  We simulate WIN+VK_NONAME. */

                INPUT input = { 0 };
                input.type = INPUT_KEYBOARD;

                input.ki.wVk = VK_NONAME;
                input.ki.wScan = 0;
                input.ki.dwFlags = 0;
                input.ki.time = 1;
                input.ki.dwExtraInfo = GetMessageExtraInfo();

                /* key down: */
                SendInput(1, &input, sizeof(INPUT));

                input.ki.dwFlags = KEYEVENTF_KEYUP;
                /* key up: */
                SendInput(1, &input, sizeof(INPUT));

                _fixWin = false;
            }
        }
    }

    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::MouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0) {
        int mouseState = 0;
        MSLLHOOKSTRUCT *msInfo;

        switch (wParam) {

            case WM_LBUTTONDOWN:
                mouseState += VK_LBUTTON;
                break;

            case WM_MBUTTONDOWN:
                mouseState += VK_MBUTTON;
                break;

            case WM_RBUTTONDOWN:
                mouseState += VK_RBUTTON;
                break;

            case WM_XBUTTONDOWN: {
                msInfo = (MSLLHOOKSTRUCT*) lParam;

                int button = msInfo->mouseData >> 16 & 0xFFFF;
                if (button == 1)
                    mouseState += HKM_MOUSE_XB1;
                else if (button == 2)
                    mouseState += HKM_MOUSE_XB2;

                break;
            }

            case WM_MOUSEWHEEL: {
                msInfo = (MSLLHOOKSTRUCT*) lParam;

                if ((int) msInfo->mouseData > 0) {
                    mouseState += HKM_MOUSE_WHUP;
                } else {
                    mouseState += HKM_MOUSE_WHDN;
                }

                break;
            }
        }

        if (mouseState > 0) {
            mouseState += Modifiers();

            if (_keyCombinations.count(mouseState) > 0) {
                PostMessage(_notifyWnd, WM_HOTKEY,
                    mouseState, mouseState & 0xF0000);

                if (mouseState & HKM_MOD_WIN) {
                    _fixWin = true; /* enable fix right before WIN goes up */
                }

                return 1; /* processed the message; eat it */
            }
        }
    }
    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelMouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->MouseProc(nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->KeyProc(nCode, wParam, lParam);
}<commit_msg>Reorder comments/statements to clarify<commit_after>#include "HotkeyManager.h"
#include "Logger.h"

HotkeyManager *HotkeyManager::instance = NULL;

HotkeyManager *HotkeyManager::Instance() {
    return instance;
}

HotkeyManager *HotkeyManager::Instance(HWND notifyWnd) {
    if (instance == NULL) {
        instance = new HotkeyManager();
        instance->Hook();
        instance->_notifyWnd = notifyWnd;
    }

    return instance;
}

HotkeyManager::HotkeyManager() :
_fixWin(false) {

}

HotkeyManager::~HotkeyManager() {
    while (true) {
        auto i = _keyCombinations.begin();
        if (_keyCombinations.size() > 0) {
            Unregister(*i);
        } else {
            break;
        }
    }

    Unhook();
}

void HotkeyManager::Shutdown() {
    delete instance;
    instance = NULL;
}

bool HotkeyManager::Hook() {
    _mouseHook = SetWindowsHookEx(WH_MOUSE_LL,
        LowLevelMouseProc, NULL, NULL);

    _keyHook = SetWindowsHookEx(WH_KEYBOARD_LL,
        LowLevelKeyboardProc, NULL, NULL);

    return _mouseHook && _keyHook;
}

bool HotkeyManager::Unhook() {
    BOOL unMouse = UnhookWindowsHookEx(_mouseHook);
    BOOL unKey = UnhookWindowsHookEx(_keyHook);
    return unMouse && unKey;
}

void HotkeyManager::Register(int keyCombination) {
    if (_keyCombinations.count(keyCombination) > 0) {
        CLOG(L"Hotkey combination [%d] already registered", keyCombination);
        return;
    } else {
        _keyCombinations.insert(keyCombination);
    }

    /* get VK_* value; includes unused bits */
    int vk = 0xFFFF & keyCombination;

    if ((keyCombination >> 20) > 0 /* uses a HKM_MOUSE_* flag */
        || vk == VK_LBUTTON
        || vk == VK_RBUTTON
        || vk == VK_MBUTTON) {

        /* mouse-based hotkeys; we are done */
        CLOG(L"Registered new mouse-based hotkey: %d", keyCombination);
        return;
    }

    /* keyboard-only hotkeys; use WinAPI */
    int mods = (0xF0000 & keyCombination) >> 16;
    if (!RegisterHotKey(_notifyWnd, keyCombination, mods, vk)) {
        CLOG(L"Failed to register hotkey [%d]\n"
            L"Mods: %d, VK: %d", keyCombination, mods, vk);
        return;
    }

    CLOG(L"Registered new keyboard-based hotkey: %d", keyCombination);
}

bool HotkeyManager::Unregister(int keyCombination) {
    CLOG(L"Unregistering hotkey combination: %d", keyCombination);

    if (_keyCombinations.count(keyCombination) <= 0) {
        QCLOG(L"Hotkey combination [%d] was not previously registered",
            keyCombination);
        return false;
    }

    _keyCombinations.erase(keyCombination);
    if ((keyCombination >> 20) == 0) {
        /* This hotkey isn't mouse-based; unregister with Windows */
        if (!UnregisterHotKey(_notifyWnd, keyCombination)) {
            CLOG(L"Failed to unregister hotkey: %d", keyCombination);
            return false;
        }
    }
    return true;
}

bool HotkeyManager::IsModifier(DWORD vk) {
    switch (vk) {
    case VK_MENU:
    case VK_LMENU:
    case VK_RMENU:
    case VK_CONTROL:
    case VK_LCONTROL:
    case VK_RCONTROL:
    case VK_SHIFT:
    case VK_LSHIFT:
    case VK_RSHIFT:
    case VK_LWIN:
    case VK_RWIN:
        return true;
    }
    return false;
}

int HotkeyManager::Modifiers() {
    int mods = 0;
    mods += (GetKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

int HotkeyManager::ModifiersAsync() {
    int mods = 0;
    mods += (GetAsyncKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetAsyncKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetAsyncKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetAsyncKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetAsyncKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

std::wstring HotkeyManager::HotkeysToModString(int combination,
    std::wstring separator) {

    std::wstring str = L"";
    if (combination & HKM_MOD_ALT) {
        str += VKToString(VK_MENU) + separator;
    }
    if (combination & HKM_MOD_CTRL) {
        str += VKToString(VK_CONTROL) + separator;
    }
    if (combination & HKM_MOD_SHF) {
        str += VKToString(VK_SHIFT) + separator;
    }
    if (combination & HKM_MOD_WIN) {
        str += L"Win" + separator;
    }

    return str;
}

std::wstring HotkeyManager::HotkeysToString(int combination,
    std::wstring separator) {

    std::wstring mods = HotkeysToModString(combination, separator);
    int vk = combination & 0xFF;
    std::wstring str = VKToString(vk);

    return mods + str;
}

std::wstring HotkeyManager::VKToString(unsigned int vk, bool extendedKey) {
    int extended = extendedKey ? 0x1 : 0x0;

    unsigned int scanCode = MapVirtualKey(vk, MAPVK_VK_TO_VSC);
    scanCode = scanCode << 16;
    scanCode |= extended << 24;
    wchar_t buf[256] = {};
    GetKeyNameText(scanCode, buf, 256);
    return std::wstring(buf);
}

LRESULT CALLBACK
HotkeyManager::KeyProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0) {
        if (wParam == WM_KEYUP) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT*) lParam;
            if ((kbInfo->vkCode == VK_LWIN || kbInfo->vkCode == VK_RWIN)
                && _fixWin) {
                /* WIN+Mouse combination used; we need to prevent the
                 * system from only seeing a WIN keypress (and usually
                 * popping up the start menu).  We simulate WIN+VK_NONAME. */

                INPUT input = { 0 };
                input.type = INPUT_KEYBOARD;

                input.ki.wVk = VK_NONAME;
                input.ki.wScan = 0;
                input.ki.dwFlags = 0;
                input.ki.time = 1;
                input.ki.dwExtraInfo = GetMessageExtraInfo();

                /* key down: */
                SendInput(1, &input, sizeof(INPUT));

                /* key up: */
                input.ki.dwFlags = KEYEVENTF_KEYUP;
                SendInput(1, &input, sizeof(INPUT));

                _fixWin = false;
            }
        }
    }

    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::MouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0) {
        int mouseState = 0;
        MSLLHOOKSTRUCT *msInfo;

        switch (wParam) {

            case WM_LBUTTONDOWN:
                mouseState += VK_LBUTTON;
                break;

            case WM_MBUTTONDOWN:
                mouseState += VK_MBUTTON;
                break;

            case WM_RBUTTONDOWN:
                mouseState += VK_RBUTTON;
                break;

            case WM_XBUTTONDOWN: {
                msInfo = (MSLLHOOKSTRUCT*) lParam;

                int button = msInfo->mouseData >> 16 & 0xFFFF;
                if (button == 1)
                    mouseState += HKM_MOUSE_XB1;
                else if (button == 2)
                    mouseState += HKM_MOUSE_XB2;

                break;
            }

            case WM_MOUSEWHEEL: {
                msInfo = (MSLLHOOKSTRUCT*) lParam;

                if ((int) msInfo->mouseData > 0) {
                    mouseState += HKM_MOUSE_WHUP;
                } else {
                    mouseState += HKM_MOUSE_WHDN;
                }

                break;
            }
        }

        if (mouseState > 0) {
            mouseState += Modifiers();

            if (_keyCombinations.count(mouseState) > 0) {
                PostMessage(_notifyWnd, WM_HOTKEY,
                    mouseState, mouseState & 0xF0000);

                if (mouseState & HKM_MOD_WIN) {
                    _fixWin = true; /* enable fix right before WIN goes up */
                }

                return 1; /* processed the message; eat it */
            }
        }
    }
    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelMouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->MouseProc(nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->KeyProc(nCode, wParam, lParam);
}<|endoftext|>


<commit_before>// Copyright (c) 2015, Matthew Malensek.
// Distributed under the BSD 2-Clause License (see LICENSE.txt for details)

#include "HotkeyManager.h"

#include "Logger.h"
#include "SyntheticKeyboard.h"

HotkeyManager *HotkeyManager::instance = NULL;

std::unordered_map<UINT, std::wstring> HotkeyManager::_vkStringMap = {
    { VK_SELECT, L"Select" },
    { VK_PRINT, L"Print" },
    { VK_EXECUTE, L"Execute" },
    { VK_SNAPSHOT, L"Print Screen" },
    { VK_HELP, L"Help" },
    { VK_SLEEP, L"Sleep" },
    { VK_BROWSER_BACK, L"Browser Back" },
    { VK_BROWSER_FORWARD, L"Browser Forward" },
    { VK_BROWSER_REFRESH, L"Browser Refresh" },
    { VK_BROWSER_STOP, L"Browser Stop" },
    { VK_BROWSER_SEARCH, L"Browser Search" },
    { VK_BROWSER_FAVORITES, L"Browser Favorites" },
    { VK_BROWSER_HOME, L"Browser Home" },
    { VK_VOLUME_MUTE, L"Mute" },
    { VK_VOLUME_DOWN, L"Volume Down" },
    { VK_VOLUME_UP, L"Volume Up" },
    { VK_MEDIA_NEXT_TRACK, L"Next Track" },
    { VK_MEDIA_PREV_TRACK, L"Previous Track" },
    { VK_MEDIA_STOP, L"Stop" },
    { VK_MEDIA_PLAY_PAUSE, L"Play/Pause" },
    { VK_LAUNCH_MAIL, L"Mail" },
    { VK_LAUNCH_MEDIA_SELECT, L"Select Media" },
    { VK_LAUNCH_APP1, L"App 1" },
    { VK_LAUNCH_APP2, L"App 2" },
    { VK_PLAY, L"Play" },
    { VK_ZOOM, L"Zoom" },
    { VK_OEM_CLEAR, L"Clear" },
};

HotkeyManager *HotkeyManager::Instance() {
    return instance;
}

HotkeyManager *HotkeyManager::Instance(HWND notifyWnd) {
    if (instance == NULL) {
        instance = new HotkeyManager();
        instance->Hook();
        instance->_notifyWnd = notifyWnd;
    }

    return instance;
}

HotkeyManager::HotkeyManager() :
_fixWin(false) {

}

HotkeyManager::~HotkeyManager() {
    while (true) {
        auto i = _keyCombinations.begin();
        if (_keyCombinations.size() > 0) {
            Unregister(*i);
        } else {
            break;
        }
    }
    _hookCombinations.clear();

    Unhook();
}

void HotkeyManager::Shutdown() {
    delete instance;
    instance = NULL;
}

bool HotkeyManager::Hook() {
    _mouseHook = SetWindowsHookEx(WH_MOUSE_LL,
        LowLevelMouseProc, NULL, NULL);

    _keyHook = SetWindowsHookEx(WH_KEYBOARD_LL,
        LowLevelKeyboardProc, NULL, NULL);

    return _mouseHook && _keyHook;
}

bool HotkeyManager::Unhook() {
    BOOL unMouse = UnhookWindowsHookEx(_mouseHook);
    BOOL unKey = UnhookWindowsHookEx(_keyHook);
    return unMouse && unKey;
}

void HotkeyManager::Register(int keyCombination) {
    if (_keyCombinations.count(keyCombination) > 0) {
        CLOG(L"Hotkey combination [%d] already registered", keyCombination);
        return;
    } else {
        _keyCombinations.insert(keyCombination);
    }

    /* get VK_* value; */
    int vk = 0xFF & keyCombination;

    if ((keyCombination >> 20) > 0 /* uses a HKM_MOUSE_* flag */
        || vk == VK_LBUTTON
        || vk == VK_RBUTTON
        || vk == VK_MBUTTON) {

        /* mouse-based hotkeys; we are done */
        CLOG(L"Registered new mouse-based hotkey: %d", keyCombination);
        return;
    }

    /* keyboard-only hotkeys; use WinAPI */
    int mods = (0xF0000 & keyCombination) >> 16;
    if (!RegisterHotKey(_notifyWnd, keyCombination, mods, vk)) {
        CLOG(L"Failed to register hotkey [%d]\n"
            L"Mods: %d, VK: %d\n"
            L"Placing in hook list", keyCombination, mods, vk);
        _hookCombinations.insert(keyCombination);
        return;
    }

    CLOG(L"Registered new keyboard-based hotkey: %d", keyCombination);
}

bool HotkeyManager::Unregister(int keyCombination) {
    CLOG(L"Unregistering hotkey combination: %d", keyCombination);

    if (_keyCombinations.count(keyCombination) <= 0) {
        QCLOG(L"Hotkey combination [%d] was not previously registered",
            keyCombination);
        return false;
    }

    _keyCombinations.erase(keyCombination);
    if ((keyCombination >> 20) == 0) {
        /* This hotkey isn't mouse-based; unregister with Windows */
        if (!UnregisterHotKey(_notifyWnd, keyCombination)) {
            CLOG(L"Failed to unregister hotkey: %d", keyCombination);
            return false;
        }
    }
    return true;
}

LRESULT CALLBACK
HotkeyManager::KeyProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0) {
        if (wParam == WM_KEYUP) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT *) lParam;
            if ((kbInfo->vkCode == VK_LWIN || kbInfo->vkCode == VK_RWIN)
                && _fixWin) {
                /* WIN+Mouse combination used; we need to prevent the
                 * system from only seeing a WIN keypress (and usually
                 * popping up the start menu).  We simulate WIN+VK_NONAME. */
                SyntheticKeyboard::SimulateKeypress(VK_NONAME, false);
                _fixWin = false;
                return CallNextHookEx(NULL, nCode, wParam, lParam);
            }
        }

        if (_hookCombinations.size() <= 0) {
            return CallNextHookEx(NULL, nCode, wParam, lParam);
        }

        if (wParam == WM_KEYDOWN || wParam == WM_SYSKEYDOWN) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT *) lParam;

            DWORD vk = kbInfo->vkCode;
            int mods = HotkeyManager::IsModifier(vk);

            if (mods) {
                _modifiers |= mods;
                return CallNextHookEx(NULL, nCode, wParam, lParam);
            } else {
                /* Is this an extended key? */
                int ext = (kbInfo->flags & 0x1) << EXT_OFFSET;
                int keys = _modifiers | ext | vk;
                if (_hookCombinations.count(keys) > 0) {
                    SendMessage(_notifyWnd, WM_HOTKEY,
                        keys, _modifiers >> MOD_OFFSET);
                    return (LRESULT) 1;
                }
            }
        }

        if (wParam == WM_KEYUP || wParam == WM_SYSKEYUP) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT *) lParam;
            int m = HotkeyManager::IsModifier(kbInfo->vkCode);
            if (m) {
                _modifiers ^= m;
            }
            return CallNextHookEx(NULL, nCode, wParam, lParam);
        }
    }

    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::MouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0 && wParam != WM_MOUSEMOVE) {
        int mouseState = 0;
        MSLLHOOKSTRUCT *msInfo;

        switch (wParam) {

            case WM_LBUTTONDOWN:
                mouseState += VK_LBUTTON;
                break;

            case WM_MBUTTONDOWN:
                mouseState += VK_MBUTTON;
                break;

            case WM_RBUTTONDOWN:
                mouseState += VK_RBUTTON;
                break;

            case WM_XBUTTONDOWN: {
                msInfo = (MSLLHOOKSTRUCT *) lParam;

                int button = msInfo->mouseData >> 16 & 0xFFFF;
                if (button == 1)
                    mouseState += HKM_MOUSE_XB1;
                else if (button == 2)
                    mouseState += HKM_MOUSE_XB2;

                break;
            }

            case WM_MOUSEWHEEL: {
                /* Note: WM_MOUSEWHEEL on a hook is a little different, so we
                 * need to grab the wheel delta info from the lParam. */
                msInfo = (MSLLHOOKSTRUCT *) lParam;

                if ((int) msInfo->mouseData > 0) {
                    mouseState += HKM_MOUSE_WHUP;
                } else {
                    mouseState += HKM_MOUSE_WHDN;
                }

                break;
            }
        }

        if (mouseState > 0) {
            mouseState += Modifiers();

            if (_keyCombinations.count(mouseState) > 0) {
                PostMessage(_notifyWnd, WM_HOTKEY,
                    mouseState, mouseState & 0xF0000);

                if (mouseState & HKM_MOD_WIN) {
                    _fixWin = true; /* enable fix right before WIN goes up */
                }

                return 1; /* processed the message; eat it */
            }
        }
    }
    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelMouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->MouseProc(nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->KeyProc(nCode, wParam, lParam);
}

int HotkeyManager::IsModifier(int vk) {
    switch (vk) {
    case VK_MENU:
    case VK_LMENU:
    case VK_RMENU:
        return HKM_MOD_ALT;

    case VK_CONTROL:
    case VK_LCONTROL:
    case VK_RCONTROL:
        return HKM_MOD_CTRL;

    case VK_SHIFT:
    case VK_LSHIFT:
    case VK_RSHIFT:
        return HKM_MOD_SHF;

    case VK_LWIN:
    case VK_RWIN:
        return HKM_MOD_WIN;
    }

    return 0;
}

bool HotkeyManager::IsMouseKey(int vk) {
    if (vk < 0x07 && vk != VK_CANCEL) {
        return true;
    }

    if (vk & (0xF << MOUSE_OFFSET)) {
        /* Has wheel or xbutton flags */
        return true;
    }

    return false;
}

int HotkeyManager::Modifiers() {
    int mods = 0;
    mods += (GetKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

int HotkeyManager::ModifiersAsync() {
    int mods = 0;
    mods += (GetAsyncKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetAsyncKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetAsyncKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetAsyncKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetAsyncKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

std::wstring HotkeyManager::HotkeysToModString(int combination,
    std::wstring separator) {

    std::wstring str = L"";
    if (combination & HKM_MOD_ALT) {
        str += VKToString(VK_MENU) + separator;
    }
    if (combination & HKM_MOD_CTRL) {
        str += VKToString(VK_CONTROL) + separator;
    }
    if (combination & HKM_MOD_SHF) {
        str += VKToString(VK_SHIFT) + separator;
    }
    if (combination & HKM_MOD_WIN) {
        str += L"Win" + separator;
    }

    return str;
}

std::wstring HotkeyManager::HotkeysToString(int combination,
    std::wstring separator) {

    std::wstring mods = HotkeysToModString(combination, separator);
    int vk = combination & 0xFF;

    std::wstring str;
    if (IsMouseKey(vk)) {
        str = MouseString(combination);
    } else {
        bool ext = (combination & 0x100) > 0;
        str = VKToString(vk, ext);
    }

    return mods + str;
}

std::wstring HotkeyManager::MouseString(int combination) {
    int vk = combination & 0xFF;
    if (vk > 0) {
        switch (vk) {
        case VK_LBUTTON:
            return L"Mouse 1";
        case VK_RBUTTON:
            return L"Mouse 2";
        case VK_MBUTTON:
            return L"Mouse 3";
        }
    }

    int flags = combination & (0xF << MOUSE_OFFSET);
    if (flags == HKM_MOUSE_XB1) {
        return L"Mouse 4";
    } else if (flags == HKM_MOUSE_XB2) {
        return L"Mouse 5";
    } else if (flags == HKM_MOUSE_WHUP) {
        return L"Mouse Wheel Up";
    } else if (flags == HKM_MOUSE_WHDN) {
        return L"Mouse Wheel Down";
    }

    return L"";
}

std::wstring HotkeyManager::VKToString(unsigned int vk, bool extendedKey) {
    if (_vkStringMap.find(vk) != _vkStringMap.end()) {
        return _vkStringMap[vk];
    }

    if (
        vk == 0x03 /* break */
        || (vk >= 0x21 && vk <= 0x2f) /* arrow keys, home/insert/del, etc */
        || (vk >= 0x5b && vk <= 0x5d) /* win, app keys (natural keyboard) */
        || (vk == 0x90) /* num lock */
        ) {
        extendedKey = true;
    }

    /* GetKeyNameText expects the following:
     * 16-23: scan code
     *    24: extended key flag
     *    25: 'do not care' bit (don't distinguish between L/R keys) */
    unsigned int scanCode = MapVirtualKey(vk, MAPVK_VK_TO_VSC);
    scanCode = scanCode << 16;

    if (vk == VK_RSHIFT) {
        /* For some reason, the right shift key ends up having its extended
         * key flag set and confuses GetKeyNameText. */
        extendedKey = false;
    }

    int extended = extendedKey ? 0x1 : 0x0;
    scanCode |= extended << 24;

    wchar_t buf[256] = {};
    GetKeyNameText(scanCode, buf, 256);
    return std::wstring(buf);
}

void HotkeyManager::VKStringTest() {
    for (unsigned int i = 0; i < 0xFF; ++i) {
        CLOG(L"%02x - %s", i, VKToString(i).c_str());
    }
}
<commit_msg>Allow using pause as a hotkey (super useful on Logitech G513)<commit_after>// Copyright (c) 2015, Matthew Malensek.
// Distributed under the BSD 2-Clause License (see LICENSE.txt for details)

#include "HotkeyManager.h"

#include "Logger.h"
#include "SyntheticKeyboard.h"

HotkeyManager *HotkeyManager::instance = NULL;

std::unordered_map<UINT, std::wstring> HotkeyManager::_vkStringMap = {
    { VK_SELECT, L"Select" },
    { VK_PRINT, L"Print" },
    { VK_EXECUTE, L"Execute" },
    { VK_SNAPSHOT, L"Print Screen" },
    { VK_HELP, L"Help" },
    { VK_SLEEP, L"Sleep" },
    { VK_BROWSER_BACK, L"Browser Back" },
    { VK_BROWSER_FORWARD, L"Browser Forward" },
    { VK_BROWSER_REFRESH, L"Browser Refresh" },
    { VK_BROWSER_STOP, L"Browser Stop" },
    { VK_BROWSER_SEARCH, L"Browser Search" },
    { VK_BROWSER_FAVORITES, L"Browser Favorites" },
    { VK_BROWSER_HOME, L"Browser Home" },
    { VK_VOLUME_MUTE, L"Mute" },
    { VK_VOLUME_DOWN, L"Volume Down" },
    { VK_VOLUME_UP, L"Volume Up" },
    { VK_MEDIA_NEXT_TRACK, L"Next Track" },
    { VK_MEDIA_PREV_TRACK, L"Previous Track" },
    { VK_MEDIA_STOP, L"Stop" },
    { VK_MEDIA_PLAY_PAUSE, L"Play/Pause" },
    { VK_LAUNCH_MAIL, L"Mail" },
    { VK_LAUNCH_MEDIA_SELECT, L"Select Media" },
    { VK_LAUNCH_APP1, L"App 1" },
    { VK_LAUNCH_APP2, L"App 2" },
    { VK_PLAY, L"Play" },
    { VK_ZOOM, L"Zoom" },
    { VK_OEM_CLEAR, L"Clear" },
};

HotkeyManager *HotkeyManager::Instance() {
    return instance;
}

HotkeyManager *HotkeyManager::Instance(HWND notifyWnd) {
    if (instance == NULL) {
        instance = new HotkeyManager();
        instance->Hook();
        instance->_notifyWnd = notifyWnd;
    }

    return instance;
}

HotkeyManager::HotkeyManager() :
_fixWin(false) {

}

HotkeyManager::~HotkeyManager() {
    while (true) {
        auto i = _keyCombinations.begin();
        if (_keyCombinations.size() > 0) {
            Unregister(*i);
        } else {
            break;
        }
    }
    _hookCombinations.clear();

    Unhook();
}

void HotkeyManager::Shutdown() {
    delete instance;
    instance = NULL;
}

bool HotkeyManager::Hook() {
    _mouseHook = SetWindowsHookEx(WH_MOUSE_LL,
        LowLevelMouseProc, NULL, NULL);

    _keyHook = SetWindowsHookEx(WH_KEYBOARD_LL,
        LowLevelKeyboardProc, NULL, NULL);

    return _mouseHook && _keyHook;
}

bool HotkeyManager::Unhook() {
    BOOL unMouse = UnhookWindowsHookEx(_mouseHook);
    BOOL unKey = UnhookWindowsHookEx(_keyHook);
    return unMouse && unKey;
}

void HotkeyManager::Register(int keyCombination) {
    if (_keyCombinations.count(keyCombination) > 0) {
        CLOG(L"Hotkey combination [%d] already registered", keyCombination);
        return;
    } else {
        _keyCombinations.insert(keyCombination);
    }

    /* get VK_* value; */
    int vk = 0xFF & keyCombination;

    if ((keyCombination >> 20) > 0 /* uses a HKM_MOUSE_* flag */
        || vk == VK_LBUTTON
        || vk == VK_RBUTTON
        || vk == VK_MBUTTON) {

        /* mouse-based hotkeys; we are done */
        CLOG(L"Registered new mouse-based hotkey: %d", keyCombination);
        return;
    }

    /* keyboard-only hotkeys; use WinAPI */
    int mods = (0xF0000 & keyCombination) >> 16;
    if (!RegisterHotKey(_notifyWnd, keyCombination, mods, vk)) {
        CLOG(L"Failed to register hotkey [%d]\n"
            L"Mods: %d, VK: %d\n"
            L"Placing in hook list", keyCombination, mods, vk);
        _hookCombinations.insert(keyCombination);
        return;
    }

    CLOG(L"Registered new keyboard-based hotkey: %d", keyCombination);
}

bool HotkeyManager::Unregister(int keyCombination) {
    CLOG(L"Unregistering hotkey combination: %d", keyCombination);

    if (_keyCombinations.count(keyCombination) <= 0) {
        QCLOG(L"Hotkey combination [%d] was not previously registered",
            keyCombination);
        return false;
    }

    _keyCombinations.erase(keyCombination);
    if ((keyCombination >> 20) == 0) {
        /* This hotkey isn't mouse-based; unregister with Windows */
        if (!UnregisterHotKey(_notifyWnd, keyCombination)) {
            CLOG(L"Failed to unregister hotkey: %d", keyCombination);
            return false;
        }
    }
    return true;
}

LRESULT CALLBACK
HotkeyManager::KeyProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0) {
        if (wParam == WM_KEYUP) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT *) lParam;
            if ((kbInfo->vkCode == VK_LWIN || kbInfo->vkCode == VK_RWIN)
                && _fixWin) {
                /* WIN+Mouse combination used; we need to prevent the
                 * system from only seeing a WIN keypress (and usually
                 * popping up the start menu).  We simulate WIN+VK_NONAME. */
                SyntheticKeyboard::SimulateKeypress(VK_NONAME, false);
                _fixWin = false;
                return CallNextHookEx(NULL, nCode, wParam, lParam);
            }
        }

        if (_hookCombinations.size() <= 0) {
            return CallNextHookEx(NULL, nCode, wParam, lParam);
        }

        if (wParam == WM_KEYDOWN || wParam == WM_SYSKEYDOWN) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT *) lParam;

            DWORD vk = kbInfo->vkCode;
            int mods = HotkeyManager::IsModifier(vk);

            if (mods) {
                _modifiers |= mods;
                return CallNextHookEx(NULL, nCode, wParam, lParam);
            } else {
                /* Is this an extended key? */
                int ext = (kbInfo->flags & 0x1) << EXT_OFFSET;
                int keys = _modifiers | ext | vk;
                if (_hookCombinations.count(keys) > 0) {
                    SendMessage(_notifyWnd, WM_HOTKEY,
                        keys, _modifiers >> MOD_OFFSET);
                    return (LRESULT) 1;
                }
            }
        }

        if (wParam == WM_KEYUP || wParam == WM_SYSKEYUP) {
            KBDLLHOOKSTRUCT *kbInfo = (KBDLLHOOKSTRUCT *) lParam;
            int m = HotkeyManager::IsModifier(kbInfo->vkCode);
            if (m) {
                _modifiers ^= m;
            }
            return CallNextHookEx(NULL, nCode, wParam, lParam);
        }
    }

    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::MouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode >= 0 && wParam != WM_MOUSEMOVE) {
        int mouseState = 0;
        MSLLHOOKSTRUCT *msInfo;

        switch (wParam) {

            case WM_LBUTTONDOWN:
                mouseState += VK_LBUTTON;
                break;

            case WM_MBUTTONDOWN:
                mouseState += VK_MBUTTON;
                break;

            case WM_RBUTTONDOWN:
                mouseState += VK_RBUTTON;
                break;

            case WM_XBUTTONDOWN: {
                msInfo = (MSLLHOOKSTRUCT *) lParam;

                int button = msInfo->mouseData >> 16 & 0xFFFF;
                if (button == 1)
                    mouseState += HKM_MOUSE_XB1;
                else if (button == 2)
                    mouseState += HKM_MOUSE_XB2;

                break;
            }

            case WM_MOUSEWHEEL: {
                /* Note: WM_MOUSEWHEEL on a hook is a little different, so we
                 * need to grab the wheel delta info from the lParam. */
                msInfo = (MSLLHOOKSTRUCT *) lParam;

                if ((int) msInfo->mouseData > 0) {
                    mouseState += HKM_MOUSE_WHUP;
                } else {
                    mouseState += HKM_MOUSE_WHDN;
                }

                break;
            }
        }

        if (mouseState > 0) {
            mouseState += Modifiers();

            if (_keyCombinations.count(mouseState) > 0) {
                PostMessage(_notifyWnd, WM_HOTKEY,
                    mouseState, mouseState & 0xF0000);

                if (mouseState & HKM_MOD_WIN) {
                    _fixWin = true; /* enable fix right before WIN goes up */
                }

                return 1; /* processed the message; eat it */
            }
        }
    }
    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelMouseProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->MouseProc(nCode, wParam, lParam);
}

LRESULT CALLBACK
HotkeyManager::LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam) {
    return HotkeyManager::instance->KeyProc(nCode, wParam, lParam);
}

int HotkeyManager::IsModifier(int vk) {
    switch (vk) {
    case VK_MENU:
    case VK_LMENU:
    case VK_RMENU:
        return HKM_MOD_ALT;

    case VK_CONTROL:
    case VK_LCONTROL:
    case VK_RCONTROL:
        return HKM_MOD_CTRL;

    case VK_SHIFT:
    case VK_LSHIFT:
    case VK_RSHIFT:
        return HKM_MOD_SHF;

    case VK_LWIN:
    case VK_RWIN:
        return HKM_MOD_WIN;
    }

    return 0;
}

bool HotkeyManager::IsMouseKey(int vk) {
    if (vk < 0x07 && vk != VK_CANCEL) {
        return true;
    }

    if (vk & (0xF << MOUSE_OFFSET)) {
        /* Has wheel or xbutton flags */
        return true;
    }

    return false;
}

int HotkeyManager::Modifiers() {
    int mods = 0;
    mods += (GetKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

int HotkeyManager::ModifiersAsync() {
    int mods = 0;
    mods += (GetAsyncKeyState(VK_MENU) & 0x8000) << 1;
    mods += (GetAsyncKeyState(VK_CONTROL) & 0x8000) << 2;
    mods += (GetAsyncKeyState(VK_SHIFT) & 0x8000) << 3;
    mods += (GetAsyncKeyState(VK_LWIN) & 0x8000) << 4;
    mods += (GetAsyncKeyState(VK_RWIN) & 0x8000) << 4;
    return mods;
}

std::wstring HotkeyManager::HotkeysToModString(int combination,
    std::wstring separator) {

    std::wstring str = L"";
    if (combination & HKM_MOD_ALT) {
        str += VKToString(VK_MENU) + separator;
    }
    if (combination & HKM_MOD_CTRL) {
        str += VKToString(VK_CONTROL) + separator;
    }
    if (combination & HKM_MOD_SHF) {
        str += VKToString(VK_SHIFT) + separator;
    }
    if (combination & HKM_MOD_WIN) {
        str += L"Win" + separator;
    }

    return str;
}

std::wstring HotkeyManager::HotkeysToString(int combination,
    std::wstring separator) {

    std::wstring mods = HotkeysToModString(combination, separator);
    int vk = combination & 0xFF;

    std::wstring str;
    if (IsMouseKey(vk)) {
        str = MouseString(combination);
    } else {
        bool ext = (combination & 0x100) > 0;
        str = VKToString(vk, ext);
    }

    return mods + str;
}

std::wstring HotkeyManager::MouseString(int combination) {
    int vk = combination & 0xFF;
    if (vk > 0) {
        switch (vk) {
        case VK_LBUTTON:
            return L"Mouse 1";
        case VK_RBUTTON:
            return L"Mouse 2";
        case VK_MBUTTON:
            return L"Mouse 3";
        }
    }

    int flags = combination & (0xF << MOUSE_OFFSET);
    if (flags == HKM_MOUSE_XB1) {
        return L"Mouse 4";
    } else if (flags == HKM_MOUSE_XB2) {
        return L"Mouse 5";
    } else if (flags == HKM_MOUSE_WHUP) {
        return L"Mouse Wheel Up";
    } else if (flags == HKM_MOUSE_WHDN) {
        return L"Mouse Wheel Down";
    }

    return L"";
}

std::wstring HotkeyManager::VKToString(unsigned int vk, bool extendedKey) {
    if (_vkStringMap.find(vk) != _vkStringMap.end()) {
        return _vkStringMap[vk];
    }

    if (
        vk == 0x03 /* break */
        || (vk >= 0x21 && vk <= 0x2f) /* arrow keys, home/insert/del, etc */
        || (vk >= 0x5b && vk <= 0x5d) /* win, app keys (natural keyboard) */
        || (vk == 0x90) /* num lock */
        ) {
        extendedKey = true;
    }

    /* GetKeyNameText expects the following:
     * 16-23: scan code
     *    24: extended key flag
     *    25: 'do not care' bit (don't distinguish between L/R keys) */
    unsigned int scanCode;
    if (vk == VK_PAUSE)
        scanCode = 0x45; /* MapVirtualKey is unable to map VK_PAUSE (this is a known bug), hence we map that by hand. */
    else
        scanCode = MapVirtualKey(vk, MAPVK_VK_TO_VSC);
    scanCode = scanCode << 16;

    if (vk == VK_RSHIFT) {
        /* For some reason, the right shift key ends up having its extended
         * key flag set and confuses GetKeyNameText. */
        extendedKey = false;
    }

    int extended = extendedKey ? 0x1 : 0x0;
    scanCode |= extended << 24;

    wchar_t buf[256] = {};
    GetKeyNameText(scanCode, buf, 256);
    return std::wstring(buf);
}

void HotkeyManager::VKStringTest() {
    for (unsigned int i = 0; i < 0xFF; ++i) {
        CLOG(L"%02x - %s", i, VKToString(i).c_str());
    }
}
<|endoftext|>