<commit_before>#pragma once
/*
**  Copyright (C) 2013 Aldebaran Robotics
**  See COPYING for the license
*/

#ifndef _QITYPE_DETAIL_PROPERTY_HXX_
#define _QITYPE_DETAIL_PROPERTY_HXX_

#include <boost/thread/locks.hpp>
#include <qi/future.hpp>
#include <qi/property.hpp>
#include <ka/errorhandling.hpp>

namespace qi
{
  inline FutureSync<void> GenericProperty::set(const AnyValue& v)
  {
    auto conv = v.convert(_type);
    if (!conv->type())
      throw std::runtime_error(std::string("Failed converting ") + v.type()->infoString() + " to " + _type->infoString());

    Property<AnyValue>::set(AnyValue(*conv, false, conv.ownsReference()));
    conv.release();

    return FutureSync<void>(0);
  }

  template<typename T>
  SignalingProperty<T>::SignalingProperty(SignalBase::OnSubscribers onsubscribe)
    : SignalF<void(const T&)>(std::move(onsubscribe))
  {
  }

  template<typename T>
  SignalingProperty<T>::SignalingProperty(ExecutionContext* execContext,
                                          SignalBase::OnSubscribers onsubscribe)
    : SignalF<void(const T&)>(execContext, std::move(onsubscribe))
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(ExecutionContext* execContext,
                                Getter getter,
                                Setter setter,
                                SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(execContext, std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(Getter getter, Setter setter, SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(AutoAnyReference defaultValue,
                                Getter getter,
                                Setter setter,
                                SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
    , _value(defaultValue.to<T>())
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(AutoAnyReference defaultValue,
                                ExecutionContext* execContext,
                                Getter getter,
                                Setter setter,
                                SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(execContext, std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
    , _value(defaultValue.to<T>())
  {
  }

  template<typename T>
  T PropertyImpl<T>::getImpl() const
  {
    if (_getter)
      return _getter(boost::ref(_value));
    else
      return _value;
  }
  template<typename T>
  void PropertyImpl<T>::setImpl(const T& v)
  {
    qiLogDebug("qitype.property") << "set " << this << " " << (!!_setter);
    if (_setter)
    {
      const bool ok = _setter(boost::ref(_value), v);
      if (ok)
        (*this)(_value);
    }
    else
    {
      _value = v;
      (*this)(_value);
    }
  }

  template<typename T>
  ReadOnlyProperty<T>::ReadOnlyProperty(SignalingProperty<T>& s)
    : _source{ s }
    , _subscription{ s.connect(track([=](const T& val){ (*this)(val); }, this)) }
  {
  }

  template<typename T>
  ReadOnlyProperty<T>::~ReadOnlyProperty()
  {
    ka::invoke_catch(qi::exceptionLogError("qi.type.readonlyproperty",
                                           "Error while destroying a read only property"),
                     [&] {
                       Trackable<ReadOnlyProperty>::destroy();
                       _source.disconnectAsync(_subscription);
                     });
  }

  template<typename T>
  FutureSync<T> ReadOnlyProperty<T>::get() const
  {
    return _source.get();
  }

  template<typename T>
  FutureSync<AnyValue> ReadOnlyProperty<T>::value() const
  {
    return _source.get().async().then(FutureCallbackType_Sync, [](const T& v) {
      return AnyValue::from(v);
    });
  }

  template<typename T>
  FutureSync<void> ReadOnlyProperty<T>::setValue(AutoAnyReference)
  {
    return makeFutureError<void>("Property is read-only, it cannot be set from this interface.");
  }

  template<typename T>
  FutureSync<T> UnsafeProperty<T>::get() const
  {
    return FutureSync<T>(this->getImpl());
  }

  template<typename T>
  FutureSync<void> UnsafeProperty<T>::set(const T& v)
  {
    this->setImpl(v);
    return FutureSync<void>(0);
  }

  template<typename T>
  FutureSync<AnyValue> UnsafeProperty<T>::value() const
  {
    return FutureSync<AnyValue>(AnyValue::from(this->getImpl()));
  }

  template<typename T>
  FutureSync<void> UnsafeProperty<T>::setValue(AutoAnyReference value)
  {
    this->setImpl(value.to<T>());
    return FutureSync<void>(0);
  }

  template<typename T>
  Property<T>::~Property()
  {
    _tracked.destroy();
    joinStrand();
    SignalBase::clearExecutionContext();
  }

  template<typename T>
  FutureSync<T> Property<T>::get() const
  {
    return strand().async(track([this]{ return this->getImpl(); }, &_tracked));
  }

  template<typename T>
  FutureSync<void> Property<T>::set(const T& v)
  {
    return strand().async(track([this, v]{ this->setImpl(v); }, &_tracked));
  }

  template<typename T>
  FutureSync<AnyValue> Property<T>::value() const
  {
    return strand().async(track([this]{ return AnyValue::from(this->getImpl()); }, &_tracked));
  }

  template<typename T>
  FutureSync<void> Property<T>::setValue(AutoAnyReference value)
  {
    const auto v = value.to<T>();
    return strand().async(track([this, v]{ this->setImpl(v); }, &_tracked));
  }

  template<typename T>
  Strand& Property<T>::strand() const
  {
    struct Src : boost::static_visitor<Strand&>, ka::src_t {};
    return boost::apply_visitor(Src{}, _strand);
  }

  template<typename T>
  void Property<T>::joinStrand() QI_NOEXCEPT(true)
  {
    struct JoinStrand : boost::static_visitor<void>
    {
      void operator()(Strand*) const
      {
        // Do nothing, we do not own the strand, we have no right to join it.
      }

      void operator()(Strand& strand) const
      {
        strand.join();
      }
    };
    boost::apply_visitor(JoinStrand{}, _strand);
  }
}

#endif  // _QITYPE_DETAIL_PROPERTY_HXX_
<commit_msg>property: Makes GenericProperty::set asynchronous.<commit_after>#pragma once
/*
**  Copyright (C) 2013 Aldebaran Robotics
**  See COPYING for the license
*/

#ifndef _QITYPE_DETAIL_PROPERTY_HXX_
#define _QITYPE_DETAIL_PROPERTY_HXX_

#include <boost/thread/locks.hpp>
#include <qi/property.hpp>
#include <qi/future.hpp>
#include <qi/property.hpp>
#include <ka/errorhandling.hpp>

namespace qi
{
  inline FutureSync<void> GenericProperty::set(const AnyValue& v)
  {
    auto conv = v.convert(_type);
    if (!conv->type())
      throw std::runtime_error(std::string("Failed converting ") + v.type()->infoString() + " to " + _type->infoString());

    const auto convOwnedRef = conv.ownsReference();
    return Property<AnyValue>::set(AnyValue(conv.release(), false, convOwnedRef));
  }

  template<typename T>
  SignalingProperty<T>::SignalingProperty(SignalBase::OnSubscribers onsubscribe)
    : SignalF<void(const T&)>(std::move(onsubscribe))
  {
  }

  template<typename T>
  SignalingProperty<T>::SignalingProperty(ExecutionContext* execContext,
                                          SignalBase::OnSubscribers onsubscribe)
    : SignalF<void(const T&)>(execContext, std::move(onsubscribe))
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(ExecutionContext* execContext,
                                Getter getter,
                                Setter setter,
                                SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(execContext, std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(Getter getter, Setter setter, SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(AutoAnyReference defaultValue,
                                Getter getter,
                                Setter setter,
                                SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
    , _value(defaultValue.to<T>())
  {
  }

  template <typename T>
  PropertyImpl<T>::PropertyImpl(AutoAnyReference defaultValue,
                                ExecutionContext* execContext,
                                Getter getter,
                                Setter setter,
                                SignalBase::OnSubscribers onsubscribe)
    : SignalingProperty<T>(execContext, std::move(onsubscribe))
    , _getter(std::move(getter))
    , _setter(std::move(setter))
    , _value(defaultValue.to<T>())
  {
  }

  template<typename T>
  T PropertyImpl<T>::getImpl() const
  {
    if (_getter)
      return _getter(boost::ref(_value));
    else
      return _value;
  }
  template<typename T>
  void PropertyImpl<T>::setImpl(const T& v)
  {
    qiLogDebug("qitype.property") << "set " << this << " " << (!!_setter);
    if (_setter)
    {
      const bool ok = _setter(boost::ref(_value), v);
      if (ok)
        (*this)(_value);
    }
    else
    {
      _value = v;
      (*this)(_value);
    }
  }

  template<typename T>
  ReadOnlyProperty<T>::ReadOnlyProperty(SignalingProperty<T>& s)
    : _source{ s }
    , _subscription{ s.connect(track([=](const T& val){ (*this)(val); }, this)) }
  {
  }

  template<typename T>
  ReadOnlyProperty<T>::~ReadOnlyProperty()
  {
    ka::invoke_catch(qi::exceptionLogError("qi.type.readonlyproperty",
                                           "Error while destroying a read only property"),
                     [&] {
                       Trackable<ReadOnlyProperty>::destroy();
                       _source.disconnectAsync(_subscription);
                     });
  }

  template<typename T>
  FutureSync<T> ReadOnlyProperty<T>::get() const
  {
    return _source.get();
  }

  template<typename T>
  FutureSync<AnyValue> ReadOnlyProperty<T>::value() const
  {
    return _source.get().async().then(FutureCallbackType_Sync, [](const T& v) {
      return AnyValue::from(v);
    });
  }

  template<typename T>
  FutureSync<void> ReadOnlyProperty<T>::setValue(AutoAnyReference)
  {
    return makeFutureError<void>("Property is read-only, it cannot be set from this interface.");
  }

  template<typename T>
  FutureSync<T> UnsafeProperty<T>::get() const
  {
    return FutureSync<T>(this->getImpl());
  }

  template<typename T>
  FutureSync<void> UnsafeProperty<T>::set(const T& v)
  {
    this->setImpl(v);
    return FutureSync<void>(0);
  }

  template<typename T>
  FutureSync<AnyValue> UnsafeProperty<T>::value() const
  {
    return FutureSync<AnyValue>(AnyValue::from(this->getImpl()));
  }

  template<typename T>
  FutureSync<void> UnsafeProperty<T>::setValue(AutoAnyReference value)
  {
    this->setImpl(value.to<T>());
    return FutureSync<void>(0);
  }

  template<typename T>
  Property<T>::~Property()
  {
    _tracked.destroy();
    joinStrand();
    SignalBase::clearExecutionContext();
  }

  template<typename T>
  FutureSync<T> Property<T>::get() const
  {
    return strand().async(track([this]{ return this->getImpl(); }, &_tracked));
  }

  template<typename T>
  FutureSync<void> Property<T>::set(const T& v)
  {
    return strand().async(track([this, v]{ this->setImpl(v); }, &_tracked));
  }

  template<typename T>
  FutureSync<AnyValue> Property<T>::value() const
  {
    return strand().async(track([this]{ return AnyValue::from(this->getImpl()); }, &_tracked));
  }

  template<typename T>
  FutureSync<void> Property<T>::setValue(AutoAnyReference value)
  {
    const auto v = value.to<T>();
    return strand().async(track([this, v]{ this->setImpl(v); }, &_tracked));
  }

  template<typename T>
  Strand& Property<T>::strand() const
  {
    struct Src : boost::static_visitor<Strand&>, ka::src_t {};
    return boost::apply_visitor(Src{}, _strand);
  }

  template<typename T>
  void Property<T>::joinStrand() QI_NOEXCEPT(true)
  {
    struct JoinStrand : boost::static_visitor<void>
    {
      void operator()(Strand*) const
      {
        // Do nothing, we do not own the strand, we have no right to join it.
      }

      void operator()(Strand& strand) const
      {
        strand.join();
      }
    };
    boost::apply_visitor(JoinStrand{}, _strand);
  }
}

#endif  // _QITYPE_DETAIL_PROPERTY_HXX_
<|endoftext|>


<commit_before>// Licensed GNU LGPL v3 or later: http://www.gnu.org/licenses/lgpl.html

#include "uithread.hh"
#include "internal.hh"
#include <semaphore.h>
#include <stdlib.h>
#include <sys/eventfd.h>
#include <deque>
#include <set>

namespace { // Anon
static void wrap_test_runner  (void);
} // Anon

namespace Rapicorn {

class Channel { // Channel for cross-thread Aida::FieldBuffer IO
  pthread_spinlock_t              msg_spinlock;
  std::vector<Aida::FieldBuffer*> msg_vector;
  std::vector<Aida::FieldBuffer*> msg_queue;
  size_t                          msg_index, msg_last_size;
  virtual void  data_notify () = 0;
  virtual void  data_wait   () = 0;
  virtual void  flush_waits () = 0;
public:
  Aida::FieldBuffer*
  fetch_msg (const int blockpop)
  {
    do
      {
        // fetch new messages
        if (msg_index >= msg_queue.size())      // no buffered messages left
          {
            if (msg_index)
              msg_queue.resize (0);             // get rid of stale pointers
            msg_index = 0;
            flush_waits();
            pthread_spin_lock (&msg_spinlock);
            msg_vector.swap (msg_queue);        // actual cross-thread fetching
            pthread_spin_unlock (&msg_spinlock);
          }
        // hand out message
        if (msg_index < msg_queue.size())       // have buffered messages
          {
            const size_t index = msg_index;
            if (blockpop > 0) // advance
              msg_index++;
            return msg_queue[index];
          }
        // no messages available
        if (blockpop < 0)
          data_wait();
      }
    while (blockpop < 0);
    return NULL;
  }
  Channel () :
    msg_index (0), msg_last_size (0)
  {
    pthread_spin_init (&msg_spinlock, 0 /* pshared */);
  }
  ~Channel ()
  {
    pthread_spin_destroy (&msg_spinlock);
  }
  void
  push_msg (Aida::FieldBuffer *fb) // takes fb ownership
  {
    pthread_spin_lock (&msg_spinlock);
    msg_vector.push_back (fb);
    msg_last_size = msg_vector.size();
    pthread_spin_unlock (&msg_spinlock);
    data_notify();
  }
  bool               has_msg            () { return fetch_msg (0) != NULL; }
  Aida::FieldBuffer* pop_msg            () { wait_msg(); return fetch_msg (1); } // passes fbmsg ownership
  void               wait_msg           () { while (!fetch_msg (-1)); }
  Aida::FieldBuffer* pop_msg_ndelay     () { return fetch_msg (1); } // passes fbmsg ownership
};

class ChannelS : public Channel { // Channel with semaphore for syncronization
  sem_t         msg_sem;
  virtual void  data_notify ()  { sem_post (&msg_sem); }
  virtual void  data_wait   ()  { sem_wait (&msg_sem); }
  virtual void  flush_waits ()  { /*sem_trywait (&msg_sem);*/ }
public:
  explicit      ChannelS    ()  { sem_init (&msg_sem, 0 /* unshared */, 0 /* init */); }
  /*dtor*/     ~ChannelS    ()  { sem_destroy (&msg_sem); }
};

class ChannelE : public Channel, public EventFd { // Channel with EventFd for syncronization
  virtual void  data_notify ()  { wakeup(); }
  virtual void  data_wait   ()  { pollin(); }
  virtual void  flush_waits ()  { flush(); }
public:
  ChannelE ()
  {
    if (open() < 0)
      fatal ("failed to open pipe for thread communication: %s", strerror());
  }
};

class ServerConnectionSource : public virtual EventLoop::Source {
  const char            *WHERE;
  Aida::ServerConnection m_connection;
  PollFD                 pollfd;
  bool                   last_seen_primary, need_check_primary;
public:
  ServerConnectionSource (EventLoop &loop, Aida::ServerConnection scon) :
    WHERE ("Rapicorn::UIThread::ServerConnection"),
    m_connection (scon), last_seen_primary (false), need_check_primary (false)
  {
    primary (false);
    loop.add (this, EventLoop::PRIORITY_NORMAL);
    pollfd.fd = m_connection.notify_fd();
    pollfd.events = PollFD::IN;
    pollfd.revents = 0;
    add_poll (&pollfd);
  }
  void
  wakeup () // allow external wakeups
  {
    // evil kludge, we're assuming Aida::ServerConnection.notify_fd() is an eventfd
    eventfd_write (pollfd.fd, 1);
  }
private:
  ~ServerConnectionSource ()
  {
    remove_poll (&pollfd);
    loop_remove();
  }
  virtual bool
  prepare (const EventLoop::State &state, int64*)
  {
    if (UNLIKELY (last_seen_primary && !state.seen_primary))
      need_check_primary = true;
    return need_check_primary || m_connection.pending();
  }
  virtual bool
  check (const EventLoop::State &state)
  {
    if (UNLIKELY (last_seen_primary && !state.seen_primary))
      need_check_primary = true;
    last_seen_primary = state.seen_primary;
    return need_check_primary || m_connection.pending();
  }
  virtual bool
  dispatch (const EventLoop::State &state)
  {
    m_connection.dispatch();
    if (need_check_primary)
      {
        need_check_primary = false;
        m_loop->exec_background (check_primaries);
      }
    return true;
  }
  static void
  check_primaries()
  {
    // seen_primary is merely a hint, this handler checks the real loop state
    const bool uithread_finishable = uithread_main_loop()->finishable();
    if (uithread_finishable)
      ApplicationImpl::the().lost_primaries();
  }
};

struct Initializer {
  int *argcp; char **argv; const StringVector *args;
  Aida::ServerConnection server_connection;
  Mutex mutex; Cond cond; uint64 app_id;
};

static Atomic<ThreadInfo*> uithread_threadinfo = NULL;

class UIThread {
  std::thread            m_thread;
  pthread_mutex_t        m_thread_mutex;
  volatile bool          m_running;
  Aida::ServerConnection m_server_connection;
  Initializer           *m_idata;
  MainLoop              &m_main_loop; // FIXME: non-NULL only while running
public:
  Aida::ClientConnection m_client_connection;
  UIThread (Initializer *idata) :
    m_thread_mutex (PTHREAD_MUTEX_INITIALIZER), m_running (0), m_idata (idata),
    m_main_loop (*ref_sink (MainLoop::_new()))
  {
    m_main_loop.set_lock_hooks (rapicorn_thread_entered, rapicorn_thread_enter, rapicorn_thread_leave);
    m_server_connection = Aida::ServerConnection::create_threaded();
    m_client_connection = Aida::ClientConnection (m_server_connection);
  }
  bool  running() const { return m_running; }
  void
  start()
  {
    pthread_mutex_lock (&m_thread_mutex);
    if (m_thread.get_id() == std::thread::id())
      {
        assert (m_running == false);
        m_thread = std::thread (std::ref (*this));
      }
    pthread_mutex_unlock (&m_thread_mutex);
  }
  void
  join()
  {
    pthread_mutex_lock (&m_thread_mutex);
    if (m_thread.joinable())
      m_thread.join();
    pthread_mutex_unlock (&m_thread_mutex);
    assert (m_running == false);
  }
  void
  queue_stop()
  {
    pthread_mutex_lock (&m_thread_mutex);
    if (&m_main_loop)
      m_main_loop.quit();
    pthread_mutex_unlock (&m_thread_mutex);
  }
  MainLoop*         main_loop()   { return &m_main_loop; }
private:
  ~UIThread ()
  {
    fatal ("UIThread singleton in dtor");
  }
  void
  initialize ()
  {
    assert_return (m_idata != NULL);
    // stay inside rapicorn_thread_enter/rapicorn_thread_leave while not polling
    assert (rapicorn_thread_entered() == true);
    // idata_core() already called
    ThisThread::affinity (string_to_int (string_vector_find (*m_idata->args, "cpu-affinity=", "-1")));
    // initialize ui_thread loop before components
    ServerConnectionSource *server_source = ref_sink (new ServerConnectionSource (m_main_loop, m_server_connection));
    (void) server_source;
    // initialize sub systems
    struct InitHookCaller : public InitHook {
      static void  invoke (const String &kind, int *argcp, char **argv, const StringVector &args)
      { invoke_hooks (kind, argcp, argv, args); }
    };
    // UI library core parts
    InitHookCaller::invoke ("ui-core/", m_idata->argcp, m_idata->argv, *m_idata->args);
    // Application Singleton
    InitHookCaller::invoke ("ui-thread/", m_idata->argcp, m_idata->argv, *m_idata->args);
    assert_return (NULL != &ApplicationImpl::the());
    // Initializations after Application Singleton
    InitHookCaller::invoke ("ui-app/", m_idata->argcp, m_idata->argv, *m_idata->args);
    // initialize uithread connection handling
    uithread_serverglue (m_server_connection);
    // Complete initialization by signalling caller
    m_idata->mutex.lock();
    m_idata->app_id = connection_object2id (ApplicationImpl::the());
    m_idata->cond.signal();
    m_idata->mutex.unlock();
    m_idata = NULL;
  }
public:
  void
  operator() ()
  {
    assert (uithread_threadinfo == NULL);
    uithread_threadinfo = &ThreadInfo::self();
    ThreadInfo::self().name ("RapicornUIThread");
    const bool running_twice = __sync_fetch_and_add (&m_running, +1);
    assert (running_twice == false);

    rapicorn_thread_enter();
    initialize();
    assert_return (m_idata == NULL);
    m_main_loop.run();
    WindowImpl::forcefully_close_all();
    ScreenDriver::forcefully_close_all();
    while (!m_main_loop.finishable())
      if (!m_main_loop.iterate (false))
        break;  // handle primary idle handlers like exec_now
    m_main_loop.kill_loops();
    rapicorn_thread_leave();

    assert (m_running == true);
    const bool stopped_twice = !__sync_fetch_and_sub (&m_running, +1);
    assert (stopped_twice == false);

    assert (m_running == false);
    assert (uithread_threadinfo == &ThreadInfo::self());
    uithread_threadinfo = NULL;
  }
};
static UIThread *the_uithread = NULL;

Aida::ClientConnection
uithread_connection (void) // prototype in ui/internal.hh
{
  return the_uithread && the_uithread->running() ? the_uithread->m_client_connection : Aida::ClientConnection();
}

MainLoop*
uithread_main_loop ()
{
  return the_uithread ? the_uithread->main_loop() : NULL;
}

bool
uithread_is_current ()
{
  return uithread_threadinfo == &ThreadInfo::self();
}

void
uithread_shutdown (void)
{
  if (the_uithread && the_uithread->running())
    {
      the_uithread->queue_stop(); // stops ui thread main loop
      the_uithread->join();
    }
}

static void
uithread_uncancelled_atexit()
{
  if (the_uithread && the_uithread->running())
    {
      /* For proper shutdown, the ui-thread needs to stop running before global
       * dtors or any atexit() handlers are being executed. C9x and C++03 leave
       * this unsolved, so we provide explicit API for the user, like
       * Rapicorn::exit() and Application::shutdown(). In case these are omitted,
       * we're not 100% safe, some earlier atexit() handler could have shot some
       * of our required resources already, however we do our best to start a
       * graceful shutdown at this point.
       */
      critical ("FIX""ME: UI-Thread still running during exit(), call Application::shutdown()");
      uithread_shutdown();
    }
}

uint64
uithread_bootup (int *argcp, char **argv, const StringVector &args)
{
  assert_return (the_uithread == NULL, 0);
  // catch exit() while UIThread is still running
  atexit (uithread_uncancelled_atexit);
  // setup client/server connection pair
  Initializer idata;
  // initialize and create UIThread
  idata.argcp = argcp; idata.argv = argv; idata.args = &args; idata.app_id = 0;
  // start and syncronize with thread
  idata.mutex.lock();
  the_uithread = new UIThread (&idata);
  the_uithread->start();
  while (idata.app_id == 0)
    idata.cond.wait (idata.mutex);
  uint64 app_id = idata.app_id;
  idata.mutex.unlock();
  assert (the_uithread->running());
  // install handler for UIThread test cases
  wrap_test_runner();
  return app_id;
}

} // Rapicorn

#include <rcore/testutils.hh>

namespace { // Anon
using namespace Rapicorn;

// === UI-Thread Syscalls ===
struct Callable {
  virtual int64 operator() () = 0;
  virtual      ~Callable   () {}
};
static std::deque<Callable*> syscall_queue;
static Mutex                 syscall_mutex;

static Aida::FieldBuffer*
ui_thread_syscall_twoway (Aida::FieldReader &fbr)
{
  Aida::FieldBuffer &rb = *Aida::FieldBuffer::new_result();
  int64 result = -1;
  syscall_mutex.lock();
  while (syscall_queue.size())
    {
      Callable &callable = *syscall_queue.front();
      syscall_queue.pop_front();
      syscall_mutex.unlock();
      result = callable ();
      delete &callable;
      syscall_mutex.lock();
    }
  syscall_mutex.unlock();
  rb.add_int64 (result);
  return &rb;
}

static const Aida::ServerConnection::MethodEntry ui_thread_call_entries[] = {
  { Aida::MSGID_TWOWAY | 0x0c0ffee01, 0x52617069636f726eULL, ui_thread_syscall_twoway, },
};
static Aida::ServerConnection::MethodRegistry ui_thread_call_registry (ui_thread_call_entries);

static int64
ui_thread_syscall (Callable *callable)
{
  Aida::FieldBuffer *fb = Aida::FieldBuffer::_new (2);
  fb->add_msgid (Aida::MSGID_TWOWAY | 0x0c0ffee01, 0x52617069636f726eULL); // ui_thread_syscall_twoway
  syscall_mutex.lock();
  syscall_queue.push_back (callable);
  syscall_mutex.unlock();
  Aida::FieldBuffer *fr = uithread_connection().call_remote (fb); // deletes fb
  Aida::FieldReader frr (*fr);
  const Aida::MessageId msgid = Aida::MessageId (frr.pop_int64());
  frr.skip(); // FIXME: check full msgid
  int64 result = 0;
  if (Aida::msgid_is_result (msgid))
    result = frr.pop_int64();
  delete fr;
  return result;
}

// === UI-Thread Test Runs ===
class SyscallTestRunner : public Callable {
  void (*test_runner) (void);
public:
  SyscallTestRunner (void (*func) (void)) : test_runner (func) {}
  int64 operator()  ()                    { test_runner(); return 0; }
};

static void
wrap_test_runner (void)
{
  Test::RegisterTest::test_set_trigger (uithread_test_trigger);
}

} // Anon

// === UI-Thread Test Trigger ===
namespace Rapicorn {
class SyscallTestTrigger : public Callable {
  void (*test_func) (void);
public:
  SyscallTestTrigger (void (*tfunc) (void)) : test_func (tfunc) {}
  int64 operator()  ()                      { test_func(); return 0; }
};
void
uithread_test_trigger (void (*test_func) ())
{
  assert_return (test_func != NULL);
  assert_return (the_uithread != NULL);
  // run tests from ui-thread
  ui_thread_syscall (new SyscallTestTrigger (test_func));
  // ensure ui-thread shutdown
  uithread_shutdown();
}
} // Rapicorn
<commit_msg>UI: removed unused outdated code<commit_after>// Licensed GNU LGPL v3 or later: http://www.gnu.org/licenses/lgpl.html

#include "uithread.hh"
#include "internal.hh"
#include <semaphore.h>
#include <stdlib.h>
#include <sys/eventfd.h>
#include <deque>
#include <set>

namespace { // Anon
static void wrap_test_runner  (void);
} // Anon

namespace Rapicorn {

class ServerConnectionSource : public virtual EventLoop::Source {
  const char            *WHERE;
  Aida::ServerConnection m_connection;
  PollFD                 pollfd;
  bool                   last_seen_primary, need_check_primary;
public:
  ServerConnectionSource (EventLoop &loop, Aida::ServerConnection scon) :
    WHERE ("Rapicorn::UIThread::ServerConnection"),
    m_connection (scon), last_seen_primary (false), need_check_primary (false)
  {
    primary (false);
    loop.add (this, EventLoop::PRIORITY_NORMAL);
    pollfd.fd = m_connection.notify_fd();
    pollfd.events = PollFD::IN;
    pollfd.revents = 0;
    add_poll (&pollfd);
  }
  void
  wakeup () // allow external wakeups
  {
    // evil kludge, we're assuming Aida::ServerConnection.notify_fd() is an eventfd
    eventfd_write (pollfd.fd, 1);
  }
private:
  ~ServerConnectionSource ()
  {
    remove_poll (&pollfd);
    loop_remove();
  }
  virtual bool
  prepare (const EventLoop::State &state, int64*)
  {
    if (UNLIKELY (last_seen_primary && !state.seen_primary))
      need_check_primary = true;
    return need_check_primary || m_connection.pending();
  }
  virtual bool
  check (const EventLoop::State &state)
  {
    if (UNLIKELY (last_seen_primary && !state.seen_primary))
      need_check_primary = true;
    last_seen_primary = state.seen_primary;
    return need_check_primary || m_connection.pending();
  }
  virtual bool
  dispatch (const EventLoop::State &state)
  {
    m_connection.dispatch();
    if (need_check_primary)
      {
        need_check_primary = false;
        m_loop->exec_background (check_primaries);
      }
    return true;
  }
  static void
  check_primaries()
  {
    // seen_primary is merely a hint, this handler checks the real loop state
    const bool uithread_finishable = uithread_main_loop()->finishable();
    if (uithread_finishable)
      ApplicationImpl::the().lost_primaries();
  }
};

struct Initializer {
  int *argcp; char **argv; const StringVector *args;
  Aida::ServerConnection server_connection;
  Mutex mutex; Cond cond; uint64 app_id;
};

static Atomic<ThreadInfo*> uithread_threadinfo = NULL;

class UIThread {
  std::thread            m_thread;
  pthread_mutex_t        m_thread_mutex;
  volatile bool          m_running;
  Aida::ServerConnection m_server_connection;
  Initializer           *m_idata;
  MainLoop              &m_main_loop; // FIXME: non-NULL only while running
public:
  Aida::ClientConnection m_client_connection;
  UIThread (Initializer *idata) :
    m_thread_mutex (PTHREAD_MUTEX_INITIALIZER), m_running (0), m_idata (idata),
    m_main_loop (*ref_sink (MainLoop::_new()))
  {
    m_main_loop.set_lock_hooks (rapicorn_thread_entered, rapicorn_thread_enter, rapicorn_thread_leave);
    m_server_connection = Aida::ServerConnection::create_threaded();
    m_client_connection = Aida::ClientConnection (m_server_connection);
  }
  bool  running() const { return m_running; }
  void
  start()
  {
    pthread_mutex_lock (&m_thread_mutex);
    if (m_thread.get_id() == std::thread::id())
      {
        assert (m_running == false);
        m_thread = std::thread (std::ref (*this));
      }
    pthread_mutex_unlock (&m_thread_mutex);
  }
  void
  join()
  {
    pthread_mutex_lock (&m_thread_mutex);
    if (m_thread.joinable())
      m_thread.join();
    pthread_mutex_unlock (&m_thread_mutex);
    assert (m_running == false);
  }
  void
  queue_stop()
  {
    pthread_mutex_lock (&m_thread_mutex);
    if (&m_main_loop)
      m_main_loop.quit();
    pthread_mutex_unlock (&m_thread_mutex);
  }
  MainLoop*         main_loop()   { return &m_main_loop; }
private:
  ~UIThread ()
  {
    fatal ("UIThread singleton in dtor");
  }
  void
  initialize ()
  {
    assert_return (m_idata != NULL);
    // stay inside rapicorn_thread_enter/rapicorn_thread_leave while not polling
    assert (rapicorn_thread_entered() == true);
    // idata_core() already called
    ThisThread::affinity (string_to_int (string_vector_find (*m_idata->args, "cpu-affinity=", "-1")));
    // initialize ui_thread loop before components
    ServerConnectionSource *server_source = ref_sink (new ServerConnectionSource (m_main_loop, m_server_connection));
    (void) server_source;
    // initialize sub systems
    struct InitHookCaller : public InitHook {
      static void  invoke (const String &kind, int *argcp, char **argv, const StringVector &args)
      { invoke_hooks (kind, argcp, argv, args); }
    };
    // UI library core parts
    InitHookCaller::invoke ("ui-core/", m_idata->argcp, m_idata->argv, *m_idata->args);
    // Application Singleton
    InitHookCaller::invoke ("ui-thread/", m_idata->argcp, m_idata->argv, *m_idata->args);
    assert_return (NULL != &ApplicationImpl::the());
    // Initializations after Application Singleton
    InitHookCaller::invoke ("ui-app/", m_idata->argcp, m_idata->argv, *m_idata->args);
    // initialize uithread connection handling
    uithread_serverglue (m_server_connection);
    // Complete initialization by signalling caller
    m_idata->mutex.lock();
    m_idata->app_id = connection_object2id (ApplicationImpl::the());
    m_idata->cond.signal();
    m_idata->mutex.unlock();
    m_idata = NULL;
  }
public:
  void
  operator() ()
  {
    assert (uithread_threadinfo == NULL);
    uithread_threadinfo = &ThreadInfo::self();
    ThreadInfo::self().name ("RapicornUIThread");
    const bool running_twice = __sync_fetch_and_add (&m_running, +1);
    assert (running_twice == false);

    rapicorn_thread_enter();
    initialize();
    assert_return (m_idata == NULL);
    m_main_loop.run();
    WindowImpl::forcefully_close_all();
    ScreenDriver::forcefully_close_all();
    while (!m_main_loop.finishable())
      if (!m_main_loop.iterate (false))
        break;  // handle primary idle handlers like exec_now
    m_main_loop.kill_loops();
    rapicorn_thread_leave();

    assert (m_running == true);
    const bool stopped_twice = !__sync_fetch_and_sub (&m_running, +1);
    assert (stopped_twice == false);

    assert (m_running == false);
    assert (uithread_threadinfo == &ThreadInfo::self());
    uithread_threadinfo = NULL;
  }
};
static UIThread *the_uithread = NULL;

Aida::ClientConnection
uithread_connection (void) // prototype in ui/internal.hh
{
  return the_uithread && the_uithread->running() ? the_uithread->m_client_connection : Aida::ClientConnection();
}

MainLoop*
uithread_main_loop ()
{
  return the_uithread ? the_uithread->main_loop() : NULL;
}

bool
uithread_is_current ()
{
  return uithread_threadinfo == &ThreadInfo::self();
}

void
uithread_shutdown (void)
{
  if (the_uithread && the_uithread->running())
    {
      the_uithread->queue_stop(); // stops ui thread main loop
      the_uithread->join();
    }
}

static void
uithread_uncancelled_atexit()
{
  if (the_uithread && the_uithread->running())
    {
      /* For proper shutdown, the ui-thread needs to stop running before global
       * dtors or any atexit() handlers are being executed. C9x and C++03 leave
       * this unsolved, so we provide explicit API for the user, like
       * Rapicorn::exit() and Application::shutdown(). In case these are omitted,
       * we're not 100% safe, some earlier atexit() handler could have shot some
       * of our required resources already, however we do our best to start a
       * graceful shutdown at this point.
       */
      critical ("FIX""ME: UI-Thread still running during exit(), call Application::shutdown()");
      uithread_shutdown();
    }
}

uint64
uithread_bootup (int *argcp, char **argv, const StringVector &args)
{
  assert_return (the_uithread == NULL, 0);
  // catch exit() while UIThread is still running
  atexit (uithread_uncancelled_atexit);
  // setup client/server connection pair
  Initializer idata;
  // initialize and create UIThread
  idata.argcp = argcp; idata.argv = argv; idata.args = &args; idata.app_id = 0;
  // start and syncronize with thread
  idata.mutex.lock();
  the_uithread = new UIThread (&idata);
  the_uithread->start();
  while (idata.app_id == 0)
    idata.cond.wait (idata.mutex);
  uint64 app_id = idata.app_id;
  idata.mutex.unlock();
  assert (the_uithread->running());
  // install handler for UIThread test cases
  wrap_test_runner();
  return app_id;
}

} // Rapicorn

#include <rcore/testutils.hh>

namespace { // Anon
using namespace Rapicorn;

// === UI-Thread Syscalls ===
struct Callable {
  virtual int64 operator() () = 0;
  virtual      ~Callable   () {}
};
static std::deque<Callable*> syscall_queue;
static Mutex                 syscall_mutex;

static Aida::FieldBuffer*
ui_thread_syscall_twoway (Aida::FieldReader &fbr)
{
  Aida::FieldBuffer &rb = *Aida::FieldBuffer::new_result();
  int64 result = -1;
  syscall_mutex.lock();
  while (syscall_queue.size())
    {
      Callable &callable = *syscall_queue.front();
      syscall_queue.pop_front();
      syscall_mutex.unlock();
      result = callable ();
      delete &callable;
      syscall_mutex.lock();
    }
  syscall_mutex.unlock();
  rb.add_int64 (result);
  return &rb;
}

static const Aida::ServerConnection::MethodEntry ui_thread_call_entries[] = {
  { Aida::MSGID_TWOWAY | 0x0c0ffee01, 0x52617069636f726eULL, ui_thread_syscall_twoway, },
};
static Aida::ServerConnection::MethodRegistry ui_thread_call_registry (ui_thread_call_entries);

static int64
ui_thread_syscall (Callable *callable)
{
  Aida::FieldBuffer *fb = Aida::FieldBuffer::_new (2);
  fb->add_msgid (Aida::MSGID_TWOWAY | 0x0c0ffee01, 0x52617069636f726eULL); // ui_thread_syscall_twoway
  syscall_mutex.lock();
  syscall_queue.push_back (callable);
  syscall_mutex.unlock();
  Aida::FieldBuffer *fr = uithread_connection().call_remote (fb); // deletes fb
  Aida::FieldReader frr (*fr);
  const Aida::MessageId msgid = Aida::MessageId (frr.pop_int64());
  frr.skip(); // FIXME: check full msgid
  int64 result = 0;
  if (Aida::msgid_is_result (msgid))
    result = frr.pop_int64();
  delete fr;
  return result;
}

// === UI-Thread Test Runs ===
class SyscallTestRunner : public Callable {
  void (*test_runner) (void);
public:
  SyscallTestRunner (void (*func) (void)) : test_runner (func) {}
  int64 operator()  ()                    { test_runner(); return 0; }
};

static void
wrap_test_runner (void)
{
  Test::RegisterTest::test_set_trigger (uithread_test_trigger);
}

} // Anon

// === UI-Thread Test Trigger ===
namespace Rapicorn {
class SyscallTestTrigger : public Callable {
  void (*test_func) (void);
public:
  SyscallTestTrigger (void (*tfunc) (void)) : test_func (tfunc) {}
  int64 operator()  ()                      { test_func(); return 0; }
};
void
uithread_test_trigger (void (*test_func) ())
{
  assert_return (test_func != NULL);
  assert_return (the_uithread != NULL);
  // run tests from ui-thread
  ui_thread_syscall (new SyscallTestTrigger (test_func));
  // ensure ui-thread shutdown
  uithread_shutdown();
}
} // Rapicorn
<|endoftext|>


<commit_before>#ifndef ITER_UNIQUE_EVERSEEN_HPP_
#define ITER_UNIQUE_EVERSEEN_HPP_

#include "filter.hpp"
#include "internal/iterbase.hpp"

#include <functional>
#include <iterator>
#include <type_traits>
#include <unordered_set>
#include <utility>

namespace iter {
  namespace impl {
    struct UniqueEverseenFn : Pipeable<UniqueEverseenFn> {
      template <typename Container>
      auto operator()(Container&& container) const {
        using elem_type = impl::iterator_deref<Container>;
        auto func = [elem_seen = std::unordered_set<std::decay_t<elem_type>>()](
            const elem_type& e) mutable {
          return elem_seen.insert(e).second;
        };
        return filter(func, std::forward<Container>(container));
      }
    };
  }

  constexpr impl::UniqueEverseenFn unique_everseen{};
}

#endif
<commit_msg>Fixes const reference build<commit_after>#ifndef ITER_UNIQUE_EVERSEEN_HPP_
#define ITER_UNIQUE_EVERSEEN_HPP_

#include "filter.hpp"
#include "internal/iterbase.hpp"

#include <functional>
#include <iterator>
#include <type_traits>
#include <unordered_set>
#include <utility>

namespace iter {
  namespace impl {
    struct UniqueEverseenFn : Pipeable<UniqueEverseenFn> {
      template <typename Container>
      auto operator()(Container&& container) const {
        using elem_type = impl::iterator_deref<Container>;
        auto func = [elem_seen = std::unordered_set<std::decay_t<elem_type>>()](
            const std::remove_reference_t<elem_type>& e) mutable {
          return elem_seen.insert(e).second;
        };
        return filter(func, std::forward<Container>(container));
      }
    };
  }

  constexpr impl::UniqueEverseenFn unique_everseen{};
}

#endif
<|endoftext|>


<commit_before>/*
    RawSpeed - RAW file decoder.

    Copyright (C) 2009-2014 Klaus Post
    Copyright (C) 2017 Axel Waggershauser
    Copyright (C) 2017 Roman Lebedev

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public
    License as published by the Free Software Foundation; either
    version 2 of the License, or (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/

#include "decompressors/AbstractLJpegDecompressor.h"
#include "common/Point.h"                       // for iPoint2D
#include "decoders/RawDecoderException.h"       // for ThrowRDE
#include "decompressors/AbstractHuffmanTable.h" // for AbstractHuffmanTable
#include "decompressors/HuffmanTable.h"         // for HuffmanTable, Huffma...
#include "io/ByteStream.h"                      // for ByteStream
#include "io/Endianness.h"                      // for Endianness, Endianne...
#include <array>                                // for array
#include <cassert>                              // for assert
#include <memory>                               // for unique_ptr, make_unique
#include <utility>                              // for move
#include <vector>                               // for vector

namespace rawspeed {

AbstractLJpegDecompressor::AbstractLJpegDecompressor(ByteStream bs,
                                                     const RawImage& img)
    : input(std::move(bs)), mRaw(img) {
  input.setByteOrder(Endianness::big);

  if (mRaw->dim.x == 0 || mRaw->dim.y == 0)
    ThrowRDE("Image has zero size");

#ifdef FUZZING_BUILD_MODE_UNSAFE_FOR_PRODUCTION
  // Yeah, sure, here it would be just dumb to leave this for production :)
  if (mRaw->dim.x > 8896 || mRaw->dim.y > 6304) {
    ThrowRDE("Unexpected image dimensions found: (%u; %u)", mRaw->dim.x,
             mRaw->dim.y);
  }
#endif
}

void AbstractLJpegDecompressor::decode() {
  if (getNextMarker(false) != M_SOI)
    ThrowRDE("Image did not start with SOI. Probably not an LJPEG");

  struct FoundMarkers {
    bool DHT = false;
    bool SOF = false;
    bool SOS = false;
  } FoundMarkers;

  JpegMarker m;
  do {
    m = getNextMarker(true);

    if (m == M_EOI)
      break;

    ByteStream data(input.getStream(input.peekU16()));
    data.skipBytes(2); // headerLength

    switch (m) {
    case M_DHT:
      if (FoundMarkers.SOS)
        ThrowRDE("Found second DHT marker after SOS");
      // there can be more than one DHT markers.
      // FIXME: do we really want to reparse and use the last one?
      parseDHT(data);
      FoundMarkers.DHT = true;
      break;
    case M_SOF3:
      if (FoundMarkers.SOS)
        ThrowRDE("Found second SOF marker after SOS");
      if (FoundMarkers.SOF)
        ThrowRDE("Found second SOF marker");
      // SOF is not required to be after DHT
      parseSOF(data, &frame);
      FoundMarkers.SOF = true;
      break;
    case M_SOS:
      if (FoundMarkers.SOS)
        ThrowRDE("Found second SOS marker");
      if (!FoundMarkers.DHT)
        ThrowRDE("Did not find DHT marker before SOS.");
      if (!FoundMarkers.SOF)
        ThrowRDE("Did not find SOF marker before SOS.");
      parseSOS(data);
      FoundMarkers.SOS = true;
      break;
    case M_DQT:
      ThrowRDE("Not a valid RAW file.");
    default: // Just let it skip to next marker
      break;
    }
  } while (m != M_EOI);

  if (!FoundMarkers.SOS)
    ThrowRDE("Did not find SOS marker.");
}

void AbstractLJpegDecompressor::parseSOF(ByteStream sofInput, SOFInfo* sof) {
  sof->prec = sofInput.getByte();
  sof->h = sofInput.getU16();
  sof->w = sofInput.getU16();
  sof->cps = sofInput.getByte();

  if (sof->prec < 2 || sof->prec > 16)
    ThrowRDE("Invalid precision (%u).", sof->prec);

  if (sof->h == 0 || sof->w == 0)
    ThrowRDE("Frame width or height set to zero");

  if (sof->cps > 4 || sof->cps < 1)
    ThrowRDE("Only from 1 to 4 components are supported.");

  if (sof->cps < mRaw->getCpp()) {
    ThrowRDE("Component count should be no less than sample count (%u vs %u).",
             sof->cps, mRaw->getCpp());
  }

  if (sof->cps > static_cast<uint32_t>(mRaw->dim.x)) {
    ThrowRDE("Component count should be no greater than row length (%u vs %u).",
             sof->cps, mRaw->dim.x);
  }

  if (sofInput.getRemainSize() != 3 * sof->cps)
    ThrowRDE("Header size mismatch.");

  for (uint32_t i = 0; i < sof->cps; i++) {
    sof->compInfo[i].componentId = sofInput.getByte();

    uint32_t subs = sofInput.getByte();
    frame.compInfo[i].superV = subs & 0xf;
    frame.compInfo[i].superH = subs >> 4;

    if (frame.compInfo[i].superV < 1 || frame.compInfo[i].superV > 4)
      ThrowRDE("Horizontal sampling factor is invalid.");

    if (frame.compInfo[i].superH < 1 || frame.compInfo[i].superH > 4)
      ThrowRDE("Horizontal sampling factor is invalid.");

    uint32_t Tq = sofInput.getByte();
    if (Tq != 0)
      ThrowRDE("Quantized components not supported.");
  }

  sof->initialized = true;

  mRaw->metadata.subsampling.x = sof->compInfo[0].superH;
  mRaw->metadata.subsampling.y = sof->compInfo[0].superV;
}

void AbstractLJpegDecompressor::parseSOS(ByteStream sos) {
  assert(frame.initialized);

  if (sos.getRemainSize() != 1 + 2 * frame.cps + 3)
    ThrowRDE("Invalid SOS header length.");

  uint32_t soscps = sos.getByte();
  if (frame.cps != soscps)
    ThrowRDE("Component number mismatch.");

  for (uint32_t i = 0; i < frame.cps; i++) {
    uint32_t cs = sos.getByte();
    uint32_t td = sos.getByte() >> 4;

    if (td >= huff.size() || !huff[td])
      ThrowRDE("Invalid Huffman table selection.");

    int ciIndex = -1;
    for (uint32_t j = 0; j < frame.cps; ++j) {
      if (frame.compInfo[j].componentId == cs)
        ciIndex = j;
    }

    if (ciIndex == -1)
      ThrowRDE("Invalid Component Selector");

    frame.compInfo[ciIndex].dcTblNo = td;
  }

  // Get predictor, see table H.1 from the JPEG spec
  predictorMode = sos.getByte();
  // The spec says predictoreMode is in [0..7], but Hasselblad uses '8'.
  if (predictorMode > 8)
    ThrowRDE("Invalid predictor mode.");

  // Se + Ah Not used in LJPEG
  if (sos.getByte() != 0)
    ThrowRDE("Se/Ah not zero.");

  Pt = sos.getByte(); // Point Transform
  if (Pt > 15)
    ThrowRDE("Invalid Point transform.");

  decodeScan();
}

void AbstractLJpegDecompressor::parseDHT(ByteStream dht) {
  while (dht.getRemainSize() > 0) {
    uint32_t b = dht.getByte();

    uint32_t htClass = b >> 4;
    if (htClass != 0)
      ThrowRDE("Unsupported Table class.");

    uint32_t htIndex = b & 0xf;
    if (htIndex >= huff.size())
      ThrowRDE("Invalid huffman table destination id.");

    if (huff[htIndex] != nullptr)
      ThrowRDE("Duplicate table definition");

    // copy 16 bytes from input stream to number of codes per length table
    uint32_t nCodes = ht_.setNCodesPerLength(dht.getBuffer(16));

    // spec says 16 different codes is max but Hasselblad violates that -> 17
    if (nCodes > 17)
      ThrowRDE("Invalid DHT table.");

    // copy nCodes bytes from input stream to code values table
    ht_.setCodeValues(dht.getBuffer(nCodes));

    // see if we already have a HuffmanTable with the same codes
    for (const auto& i : huffmanTableStore)
      if (*i == ht_)
        huff[htIndex] = i.get();

    if (!huff[htIndex]) {
      // setup new ht_ and put it into the store
      auto dHT = std::make_unique<HuffmanTable>(ht_);
      dHT->setup(fullDecodeHT, fixDng16Bug);
      huff[htIndex] = dHT.get();
      huffmanTableStore.emplace_back(std::move(dHT));
    }
  }
}

JpegMarker AbstractLJpegDecompressor::getNextMarker(bool allowskip) {
  uint8_t c0;
  uint8_t c1 = input.getByte();
  do {
    c0 = c1;
    c1 = input.getByte();
  } while (allowskip && !(c0 == 0xFF && c1 != 0 && c1 != 0xFF));

  if (!(c0 == 0xFF && c1 != 0 && c1 != 0xFF))
    ThrowRDE("(Noskip) Expected marker not found. Probably corrupt file.");

  return static_cast<JpegMarker>(c1);
}

} // namespace rawspeed
<commit_msg>AbstractLJpegDecompressor::parseSOF(): don't write mRaw->metadata.subsampling, but check against it<commit_after>/*
    RawSpeed - RAW file decoder.

    Copyright (C) 2009-2014 Klaus Post
    Copyright (C) 2017 Axel Waggershauser
    Copyright (C) 2017 Roman Lebedev

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public
    License as published by the Free Software Foundation; either
    version 2 of the License, or (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/

#include "decompressors/AbstractLJpegDecompressor.h"
#include "common/Point.h"                       // for iPoint2D
#include "decoders/RawDecoderException.h"       // for ThrowRDE
#include "decompressors/AbstractHuffmanTable.h" // for AbstractHuffmanTable
#include "decompressors/HuffmanTable.h"         // for HuffmanTable, Huffma...
#include "io/ByteStream.h"                      // for ByteStream
#include "io/Endianness.h"                      // for Endianness, Endianne...
#include <array>                                // for array
#include <cassert>                              // for assert
#include <memory>                               // for unique_ptr, make_unique
#include <utility>                              // for move
#include <vector>                               // for vector

namespace rawspeed {

AbstractLJpegDecompressor::AbstractLJpegDecompressor(ByteStream bs,
                                                     const RawImage& img)
    : input(std::move(bs)), mRaw(img) {
  input.setByteOrder(Endianness::big);

  if (mRaw->dim.x == 0 || mRaw->dim.y == 0)
    ThrowRDE("Image has zero size");

#ifdef FUZZING_BUILD_MODE_UNSAFE_FOR_PRODUCTION
  // Yeah, sure, here it would be just dumb to leave this for production :)
  if (mRaw->dim.x > 8896 || mRaw->dim.y > 6304) {
    ThrowRDE("Unexpected image dimensions found: (%u; %u)", mRaw->dim.x,
             mRaw->dim.y);
  }
#endif
}

void AbstractLJpegDecompressor::decode() {
  if (getNextMarker(false) != M_SOI)
    ThrowRDE("Image did not start with SOI. Probably not an LJPEG");

  struct FoundMarkers {
    bool DHT = false;
    bool SOF = false;
    bool SOS = false;
  } FoundMarkers;

  JpegMarker m;
  do {
    m = getNextMarker(true);

    if (m == M_EOI)
      break;

    ByteStream data(input.getStream(input.peekU16()));
    data.skipBytes(2); // headerLength

    switch (m) {
    case M_DHT:
      if (FoundMarkers.SOS)
        ThrowRDE("Found second DHT marker after SOS");
      // there can be more than one DHT markers.
      // FIXME: do we really want to reparse and use the last one?
      parseDHT(data);
      FoundMarkers.DHT = true;
      break;
    case M_SOF3:
      if (FoundMarkers.SOS)
        ThrowRDE("Found second SOF marker after SOS");
      if (FoundMarkers.SOF)
        ThrowRDE("Found second SOF marker");
      // SOF is not required to be after DHT
      parseSOF(data, &frame);
      FoundMarkers.SOF = true;
      break;
    case M_SOS:
      if (FoundMarkers.SOS)
        ThrowRDE("Found second SOS marker");
      if (!FoundMarkers.DHT)
        ThrowRDE("Did not find DHT marker before SOS.");
      if (!FoundMarkers.SOF)
        ThrowRDE("Did not find SOF marker before SOS.");
      parseSOS(data);
      FoundMarkers.SOS = true;
      break;
    case M_DQT:
      ThrowRDE("Not a valid RAW file.");
    default: // Just let it skip to next marker
      break;
    }
  } while (m != M_EOI);

  if (!FoundMarkers.SOS)
    ThrowRDE("Did not find SOS marker.");
}

void AbstractLJpegDecompressor::parseSOF(ByteStream sofInput, SOFInfo* sof) {
  sof->prec = sofInput.getByte();
  sof->h = sofInput.getU16();
  sof->w = sofInput.getU16();
  sof->cps = sofInput.getByte();

  if (sof->prec < 2 || sof->prec > 16)
    ThrowRDE("Invalid precision (%u).", sof->prec);

  if (sof->h == 0 || sof->w == 0)
    ThrowRDE("Frame width or height set to zero");

  if (sof->cps > 4 || sof->cps < 1)
    ThrowRDE("Only from 1 to 4 components are supported.");

  if (sof->cps < mRaw->getCpp()) {
    ThrowRDE("Component count should be no less than sample count (%u vs %u).",
             sof->cps, mRaw->getCpp());
  }

  if (sof->cps > static_cast<uint32_t>(mRaw->dim.x)) {
    ThrowRDE("Component count should be no greater than row length (%u vs %u).",
             sof->cps, mRaw->dim.x);
  }

  if (sofInput.getRemainSize() != 3 * sof->cps)
    ThrowRDE("Header size mismatch.");

  for (uint32_t i = 0; i < sof->cps; i++) {
    sof->compInfo[i].componentId = sofInput.getByte();

    uint32_t subs = sofInput.getByte();
    frame.compInfo[i].superV = subs & 0xf;
    frame.compInfo[i].superH = subs >> 4;

    if (frame.compInfo[i].superV < 1 || frame.compInfo[i].superV > 4)
      ThrowRDE("Horizontal sampling factor is invalid.");

    if (frame.compInfo[i].superH < 1 || frame.compInfo[i].superH > 4)
      ThrowRDE("Horizontal sampling factor is invalid.");

    uint32_t Tq = sofInput.getByte();
    if (Tq != 0)
      ThrowRDE("Quantized components not supported.");
  }

  if (static_cast<int>(sof->compInfo[0].superH) !=
          mRaw->metadata.subsampling.x ||
      static_cast<int>(sof->compInfo[0].superV) != mRaw->metadata.subsampling.y)
    ThrowRDE("LJpeg's subsampling does not match image's subsampling.");

  sof->initialized = true;
}

void AbstractLJpegDecompressor::parseSOS(ByteStream sos) {
  assert(frame.initialized);

  if (sos.getRemainSize() != 1 + 2 * frame.cps + 3)
    ThrowRDE("Invalid SOS header length.");

  uint32_t soscps = sos.getByte();
  if (frame.cps != soscps)
    ThrowRDE("Component number mismatch.");

  for (uint32_t i = 0; i < frame.cps; i++) {
    uint32_t cs = sos.getByte();
    uint32_t td = sos.getByte() >> 4;

    if (td >= huff.size() || !huff[td])
      ThrowRDE("Invalid Huffman table selection.");

    int ciIndex = -1;
    for (uint32_t j = 0; j < frame.cps; ++j) {
      if (frame.compInfo[j].componentId == cs)
        ciIndex = j;
    }

    if (ciIndex == -1)
      ThrowRDE("Invalid Component Selector");

    frame.compInfo[ciIndex].dcTblNo = td;
  }

  // Get predictor, see table H.1 from the JPEG spec
  predictorMode = sos.getByte();
  // The spec says predictoreMode is in [0..7], but Hasselblad uses '8'.
  if (predictorMode > 8)
    ThrowRDE("Invalid predictor mode.");

  // Se + Ah Not used in LJPEG
  if (sos.getByte() != 0)
    ThrowRDE("Se/Ah not zero.");

  Pt = sos.getByte(); // Point Transform
  if (Pt > 15)
    ThrowRDE("Invalid Point transform.");

  decodeScan();
}

void AbstractLJpegDecompressor::parseDHT(ByteStream dht) {
  while (dht.getRemainSize() > 0) {
    uint32_t b = dht.getByte();

    uint32_t htClass = b >> 4;
    if (htClass != 0)
      ThrowRDE("Unsupported Table class.");

    uint32_t htIndex = b & 0xf;
    if (htIndex >= huff.size())
      ThrowRDE("Invalid huffman table destination id.");

    if (huff[htIndex] != nullptr)
      ThrowRDE("Duplicate table definition");

    // copy 16 bytes from input stream to number of codes per length table
    uint32_t nCodes = ht_.setNCodesPerLength(dht.getBuffer(16));

    // spec says 16 different codes is max but Hasselblad violates that -> 17
    if (nCodes > 17)
      ThrowRDE("Invalid DHT table.");

    // copy nCodes bytes from input stream to code values table
    ht_.setCodeValues(dht.getBuffer(nCodes));

    // see if we already have a HuffmanTable with the same codes
    for (const auto& i : huffmanTableStore)
      if (*i == ht_)
        huff[htIndex] = i.get();

    if (!huff[htIndex]) {
      // setup new ht_ and put it into the store
      auto dHT = std::make_unique<HuffmanTable>(ht_);
      dHT->setup(fullDecodeHT, fixDng16Bug);
      huff[htIndex] = dHT.get();
      huffmanTableStore.emplace_back(std::move(dHT));
    }
  }
}

JpegMarker AbstractLJpegDecompressor::getNextMarker(bool allowskip) {
  uint8_t c0;
  uint8_t c1 = input.getByte();
  do {
    c0 = c1;
    c1 = input.getByte();
  } while (allowskip && !(c0 == 0xFF && c1 != 0 && c1 != 0xFF));

  if (!(c0 == 0xFF && c1 != 0 && c1 != 0xFF))
    ThrowRDE("(Noskip) Expected marker not found. Probably corrupt file.");

  return static_cast<JpegMarker>(c1);
}

} // namespace rawspeed
<|endoftext|>


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

#include "wfa/virtual_people/core/labeler/labeler.h"

#include "absl/container/flat_hash_map.h"
#include "gmock/gmock.h"
#include "google/protobuf/text_format.h"
#include "gtest/gtest.h"
#include "src/main/proto/wfa/virtual_people/common/event.pb.h"
#include "src/main/proto/wfa/virtual_people/common/label.pb.h"
#include "src/main/proto/wfa/virtual_people/common/model.pb.h"
#include "src/test/cc/testutil/matchers.h"
#include "src/test/cc/testutil/status_macros.h"

namespace wfa_virtual_people {
namespace {

using ::testing::DoubleNear;
using ::testing::Pair;
using ::testing::UnorderedElementsAre;
using ::wfa::IsOk;
using ::wfa::StatusIs;

constexpr int kEventIdNumber = 10000;

TEST(LabelerTest, TestBuildFromRoot) {
  // A model with
  // * 40% probability to assign virtual person id 10
  // * 60% probability to assign virtual person id 20
  CompiledNode root;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node {
            population_node {
              pools {
                population_offset: 10
                total_population: 1
              }
              random_seed: "TestPopulationNodeSeed1"
            }
          }
          chance: 0.4
        }
        branches {
          node {
            population_node {
              pools {
                population_offset: 20
                total_population: 1
              }
              random_seed: "TestPopulationNodeSeed2"
            }
          }
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &root));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(root));

  absl::flat_hash_map<int64_t, double> vpid_counts;
  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    ++vpid_counts[output.people(0).virtual_person_id()];
  }
  for (auto& [key, value] : vpid_counts) {
    value /= static_cast<double>(kEventIdNumber);
  }
  // Absolute error more than 2% is very unlikely.
  EXPECT_THAT(vpid_counts, UnorderedElementsAre(
      Pair(10, DoubleNear(0.4, 0.02)),
      Pair(20, DoubleNear(0.6, 0.02))));
}

TEST(LabelerTest, TestBuildFromNodesRootWithIndex) {
  // All nodes, including root node, have index set.
  // A model with
  // * 40% probability to assign virtual person id 10
  // * 60% probability to assign virtual person id 20
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 0.4
        }
        branches {
          node_index: 3
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(nodes));

  absl::flat_hash_map<int64_t, double> vpid_counts;
  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    ++vpid_counts[output.people(0).virtual_person_id()];
  }
  for (auto& [key, value] : vpid_counts) {
    value /= static_cast<double>(kEventIdNumber);
  }
  // Absolute error more than 2% is very unlikely.
  EXPECT_THAT(vpid_counts, UnorderedElementsAre(
      Pair(10, DoubleNear(0.4, 0.02)),
      Pair(20, DoubleNear(0.6, 0.02))));
}

TEST(LabelerTest, TestBuildFromNodesRootWithoutIndex) {
  // All nodes, except root node, have index set.
  // A model with
  // * 40% probability to assign virtual person id 10
  // * 60% probability to assign virtual person id 20
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 0.4
        }
        branches {
          node_index: 3
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(nodes));

  absl::flat_hash_map<int64_t, double> vpid_counts;
  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    ++vpid_counts[output.people(0).virtual_person_id()];
  }
  for (auto& [key, value] : vpid_counts) {
    value /= static_cast<double>(kEventIdNumber);
  }
  // Absolute error more than 2% is very unlikely.
  EXPECT_THAT(vpid_counts, UnorderedElementsAre(
      Pair(10, DoubleNear(0.4, 0.02)),
      Pair(20, DoubleNear(0.6, 0.02))));
}

TEST(LabelerTest, TestBuildFromNodesNodeAfterRoot) {
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 3
          chance: 1
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesMultipleRoots) {
  // 2 trees exist:
  // 1 -> 3, 4
  // 2 -> 5
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 4
      name: "TestNode4"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 5
      name: "TestNode5"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed3"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 3
          chance: 0.4
        }
        branches {
          node_index: 4
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      branch_node {
        branches {
          node_index: 5
          chance: 1
        }
        random_seed: "TestBranchNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesNoRoot) {
  // Nodes 1 and 2 reference each other as child node. No root node can be
  // recognized.
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 1
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      branch_node {
        branches {
          node_index: 1
          chance: 1
        }
        random_seed: "TestBranchNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesNoNodeForIndex) {
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 1
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesDuplicatedIndexes) {
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 1
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesMultipleParents) {
  // This is invalid as node 3 is the child node of both node 1 and node 2.
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      branch_node {
        branches {
          node_index: 3
          chance: 1
        }
        random_seed: "TestBranchNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 0.5
        }
        branches {
          node_index: 3
          chance: 0.5
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

}  // namespace
}  // namespace wfa_virtual_people
<commit_msg>Add a test case to build Labeler with a vector containing single node.<commit_after>// Copyright 2021 The Cross-Media Measurement Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "wfa/virtual_people/core/labeler/labeler.h"

#include "absl/container/flat_hash_map.h"
#include "gmock/gmock.h"
#include "google/protobuf/text_format.h"
#include "gtest/gtest.h"
#include "src/main/proto/wfa/virtual_people/common/event.pb.h"
#include "src/main/proto/wfa/virtual_people/common/label.pb.h"
#include "src/main/proto/wfa/virtual_people/common/model.pb.h"
#include "src/test/cc/testutil/matchers.h"
#include "src/test/cc/testutil/status_macros.h"

namespace wfa_virtual_people {
namespace {

using ::testing::DoubleNear;
using ::testing::Pair;
using ::testing::UnorderedElementsAre;
using ::wfa::IsOk;
using ::wfa::StatusIs;

constexpr int kEventIdNumber = 10000;

TEST(LabelerTest, TestBuildFromRoot) {
  // A model with
  // * 40% probability to assign virtual person id 10
  // * 60% probability to assign virtual person id 20
  CompiledNode root;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node {
            population_node {
              pools {
                population_offset: 10
                total_population: 1
              }
              random_seed: "TestPopulationNodeSeed1"
            }
          }
          chance: 0.4
        }
        branches {
          node {
            population_node {
              pools {
                population_offset: 20
                total_population: 1
              }
              random_seed: "TestPopulationNodeSeed2"
            }
          }
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &root));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(root));

  absl::flat_hash_map<int64_t, double> vpid_counts;
  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    ++vpid_counts[output.people(0).virtual_person_id()];
  }
  for (auto& [key, value] : vpid_counts) {
    value /= static_cast<double>(kEventIdNumber);
  }
  // Absolute error more than 2% is very unlikely.
  EXPECT_THAT(vpid_counts, UnorderedElementsAre(
      Pair(10, DoubleNear(0.4, 0.02)),
      Pair(20, DoubleNear(0.6, 0.02))));
}

TEST(LabelerTest, TestBuildFromNodesRootWithIndex) {
  // All nodes, including root node, have index set.
  // A model with
  // * 40% probability to assign virtual person id 10
  // * 60% probability to assign virtual person id 20
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 0.4
        }
        branches {
          node_index: 3
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(nodes));

  absl::flat_hash_map<int64_t, double> vpid_counts;
  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    ++vpid_counts[output.people(0).virtual_person_id()];
  }
  for (auto& [key, value] : vpid_counts) {
    value /= static_cast<double>(kEventIdNumber);
  }
  // Absolute error more than 2% is very unlikely.
  EXPECT_THAT(vpid_counts, UnorderedElementsAre(
      Pair(10, DoubleNear(0.4, 0.02)),
      Pair(20, DoubleNear(0.6, 0.02))));
}

TEST(LabelerTest, TestBuildFromNodesRootWithoutIndex) {
  // All nodes, except root node, have index set.
  // A model with
  // * 40% probability to assign virtual person id 10
  // * 60% probability to assign virtual person id 20
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 0.4
        }
        branches {
          node_index: 3
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(nodes));

  absl::flat_hash_map<int64_t, double> vpid_counts;
  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    ++vpid_counts[output.people(0).virtual_person_id()];
  }
  for (auto& [key, value] : vpid_counts) {
    value /= static_cast<double>(kEventIdNumber);
  }
  // Absolute error more than 2% is very unlikely.
  EXPECT_THAT(vpid_counts, UnorderedElementsAre(
      Pair(10, DoubleNear(0.4, 0.02)),
      Pair(20, DoubleNear(0.6, 0.02))));
}

TEST(LabelerTest, TestBuildFromListWithSingleNode) {
  // A model with single node, which always assigns virtual person id 10.
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<Labeler> labeler, Labeler::Build(nodes));

  for (int event_id = 0; event_id < kEventIdNumber; ++event_id) {
    LabelerInput input;
    input.mutable_event_id()->set_id(std::to_string(event_id));
    LabelerOutput output;
    EXPECT_THAT(labeler->Label(input, output), IsOk());
    EXPECT_EQ(output.people(0).virtual_person_id(), 10);
  }
}

TEST(LabelerTest, TestBuildFromNodesNodeAfterRoot) {
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 3
          chance: 1
        }
        random_seed: "TestBranchNodeSeed"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesMultipleRoots) {
  // 2 trees exist:
  // 1 -> 3, 4
  // 2 -> 5
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 4
      name: "TestNode4"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 5
      name: "TestNode5"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed3"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 3
          chance: 0.4
        }
        branches {
          node_index: 4
          chance: 0.6
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      branch_node {
        branches {
          node_index: 5
          chance: 1
        }
        random_seed: "TestBranchNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesNoRoot) {
  // Nodes 1 and 2 reference each other as child node. No root node can be
  // recognized.
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 1
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      branch_node {
        branches {
          node_index: 1
          chance: 1
        }
        random_seed: "TestBranchNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesNoNodeForIndex) {
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 1
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 1
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesDuplicatedIndexes) {
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 20
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 1
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

TEST(LabelerTest, TestBuildFromNodesMultipleParents) {
  // This is invalid as node 3 is the child node of both node 1 and node 2.
  std::vector<CompiledNode> nodes;
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 3
      name: "TestNode3"
      population_node {
        pools {
          population_offset: 10
          total_population: 1
        }
        random_seed: "TestPopulationNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      index: 2
      name: "TestNode2"
      branch_node {
        branches {
          node_index: 3
          chance: 1
        }
        random_seed: "TestBranchNodeSeed2"
      }
  )pb", &nodes.emplace_back()));
  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(R"pb(
      name: "TestNode1"
      branch_node {
        branches {
          node_index: 2
          chance: 0.5
        }
        branches {
          node_index: 3
          chance: 0.5
        }
        random_seed: "TestBranchNodeSeed1"
      }
  )pb", &nodes.emplace_back()));
  EXPECT_THAT(
      Labeler::Build(nodes).status(),
      StatusIs(absl::StatusCode::kInvalidArgument, ""));
}

}  // namespace
}  // namespace wfa_virtual_people
<|endoftext|>


<commit_before>#include "stdafx.h"
#include "SettingsUI.h"
#include "Hotkeys.h"
#include "afxdialogex.h"

#include <unordered_map>

#include "HotkeyInfo.h"
#include "HotkeyPrompt.h"
#include "KeyGrabber.h"
#include "Settings.h"

IMPLEMENT_DYNAMIC(Hotkeys, CPropertyPage)

Hotkeys::Hotkeys() :
CPropertyPage(Hotkeys::IDD) {

}

Hotkeys::~Hotkeys() {
}

void Hotkeys::DoDataExchange(CDataExchange* pDX)
{
    CPropertyPage::DoDataExchange(pDX);
    DDX_Control(pDX, LST_KEYS, _list);
    DDX_Control(pDX, BTN_KEYS, _keys);
    DDX_Control(pDX, CMB_ACTION, _action);
    DDX_Control(pDX, LBL_ARG, _argument);
    DDX_Control(pDX, CMB_DRIVE, _argCombo);
}

BOOL Hotkeys::OnApply() {
    OutputDebugString(L"-> Hotkeys\n");
    Settings *settings = Settings::Instance();
    settings->Hotkeys(_keyInfo);
    UIUtils::SaveSettings(*this);
    return CPropertyPage::OnApply();
}

BOOL Hotkeys::OnCommand(WPARAM wParam, LPARAM lParam) {
    SetModified();
    return CPropertyPage::OnCommand(wParam, lParam);
}

BOOL Hotkeys::OnInitDialog() {
    CPropertyPage::OnInitDialog();

    Settings *settings = Settings::Instance();

    DWORD exStyle = _list.GetExtendedStyle();
    exStyle |= LVS_EX_FULLROWSELECT;
    _list.SetExtendedStyle(exStyle);

    /* Set up listview and editor area */
    RECT r;
    _list.GetWindowRect(&r);
    int width = r.right - r.left;

    _list.InsertColumn(0, L"Hotkeys", NULL, (int) (width * .485));
    _list.InsertColumn(1, L"Action", NULL, (int) (width * .445));

    for (std::wstring action : HotkeyInfo::ActionNames) {
        _action.AddString(action.c_str());
    }

    std::unordered_map<int, HotkeyInfo> hotkeys = settings->Hotkeys();
    for (auto it = hotkeys.begin(); it != hotkeys.end(); ++it) {
        _keyInfo.push_back(it->second);
    }

    for (unsigned int i = 0; i < _keyInfo.size(); ++i) {
        LoadSelection(i);
        HotkeyInfo hki = _keyInfo[i];
        std::wstring hkStr = HotkeyManager::HotkeysToString(hki.keyCombination);
        _list.InsertItem(i, hkStr.c_str());
        _list.SetItemText(i, 1, HotkeyInfo::ActionNames[hki.action].c_str());
    }

    SelectItem(0);

    return TRUE;
}

void Hotkeys::SelectItem(int idx) {
    /* First unselect the selected item (not strictly necessary in our case) */
    int sel = _list.GetSelectionMark();
    if (sel >= 0) {
        _list.SetItemState(sel, ~LVIS_SELECTED, LVIS_SELECTED);
    }

    /* Make sure we're not selecting something out of range */
    if (idx < 0) {
        idx = 0;
    }
    int numItems = _list.GetItemCount();
    if (idx > numItems - 1) {
        idx = numItems - 1;
    }

    /* Select the item */
    _list.SetItemState(idx, LVIS_SELECTED, LVIS_SELECTED);
    _list.SetSelectionMark(idx);

    /* Scroll down to the item, if necessary */
    _list.EnsureVisible(idx, FALSE);
}

void Hotkeys::LoadSelection(int idx) {
    HotkeyInfo current = _keyInfo[_selIdx];

    _keys.SetWindowText(L"");

    if (current.keyCombination > 0) {
        std::wstring keyStr
            = HotkeyManager::HotkeysToString(current.keyCombination);
        _keys.SetWindowText(keyStr.c_str());
        _list.SetItemText(_selIdx, 0, keyStr.c_str());
    }

    if (current.action >= 0) {
        _list.SetItemText(_selIdx, 1,
            HotkeyInfo::ActionNames[current.action].c_str());

        switch ((HotkeyInfo::HotkeyActions) current.action) {
        case HotkeyInfo::EjectDrive:
            _argument.SetWindowTextW(L"Drive:");
            _argCombo.ResetContent();
            for (int i = 0; i < 26; ++i) {
                wchar_t ch = (wchar_t) i + 65;
                _argCombo.InsertString(i, CString(ch));
            }
            if (current.args.size() > 0) {
                _argCombo.SelectString(0, current.args[0].c_str());
            }

            _argument.ShowWindow(SW_SHOW);
            _argCombo.ShowWindow(SW_SHOW);
            break;

        default:
            _argument.ShowWindow(SW_HIDE);
            _argCombo.ShowWindow(SW_HIDE);
        }
    } else {
        /* TODO fix these nasty repeated hides */
        _argument.ShowWindow(SW_HIDE);
        _argCombo.ShowWindow(SW_HIDE);
    }

    _action.SetCurSel(current.action);
}

BEGIN_MESSAGE_MAP(Hotkeys, CPropertyPage)
    ON_BN_CLICKED(BTN_ADD, &Hotkeys::OnBnClickedAdd)
    ON_BN_CLICKED(BTN_REMOVE, &Hotkeys::OnBnClickedRemove)
    ON_NOTIFY(LVN_ITEMCHANGED, LST_KEYS, &Hotkeys::OnLvnItemchangedKeys)
    ON_BN_CLICKED(BTN_KEYS, &Hotkeys::OnBnClickedKeys)
    ON_CBN_SELCHANGE(CMB_ACTION, &Hotkeys::OnCbnSelchangeAction)
    ON_CBN_SELCHANGE(CMB_ARG, &Hotkeys::OnCbnSelchangeArg)
END_MESSAGE_MAP()

void Hotkeys::OnBnClickedAdd() {
    int items = _list.GetItemCount();
    for (int i = 0; i < items; ++i) {
        if (_list.GetItemText(i, 0) == L""
            && _list.GetItemText(i, 1) == L"") {
            /* We found a blank item already in the list */
            SelectItem(i);
            return;
        }
    }

    HotkeyInfo hi;
    _keyInfo.push_back(hi);
    int idx = _list.InsertItem(items, L"");
    _list.SetItemText(idx, 1, L"");
    SelectItem(idx);
}

void Hotkeys::OnBnClickedRemove() {
    _keyInfo.erase(_keyInfo.begin() + _selIdx);
    _list.DeleteItem(_selIdx);

    /* Select the item closest to the previous selection: */
    SelectItem(_selIdx);
}

void Hotkeys::OnLvnItemchangedKeys(NMHDR *pNMHDR, LRESULT *pResult) {
    LPNMLISTVIEW pNMLV = reinterpret_cast<LPNMLISTVIEW>(pNMHDR);

    if (pNMLV->uChanged & LVIF_STATE) {
        if (pNMLV->uNewState & LVIS_SELECTED) {
            int sel = pNMLV->iItem;
            _selIdx = sel;
            OutputDebugString(std::to_wstring(sel).c_str());
            OutputDebugString(L"\n");
            LoadSelection(sel);
        }
    }

    *pResult = 0;
}

void Hotkeys::OnBnClickedKeys() {
    HotkeyPrompt hkp;
    hkp.DoModal();
    KeyGrabber::Instance()->Unhook();
    int keyCombo = KeyGrabber::Instance()->KeyCombination();
    if (keyCombo > 0) {
        std::wstring keyStr = HotkeyManager::HotkeysToString(keyCombo);
        _keys.SetWindowText(keyStr.c_str());
        _keyInfo[_selIdx].keyCombination = keyCombo;
        LoadSelection(_selIdx);
    }
}

void Hotkeys::OnCbnSelchangeAction() {
    int sel = _action.GetCurSel();
    _keyInfo[_selIdx].action = sel;
    LoadSelection(_selIdx);
}

void Hotkeys::OnCbnSelchangeArg() {
    HotkeyInfo *current = &_keyInfo[_selIdx];

    HotkeyInfo::HotkeyActions action
        = (HotkeyInfo::HotkeyActions) _action.GetCurSel();

    switch (action) {
    case HotkeyInfo::EjectDrive:
        wchar_t buf[2];
        _argCombo.GetWindowText(buf, 2);
        if (current->args.size() <= 0) {
            current->args.resize(1);
        }
        current->args[0] = buf;
        break;
    }
    LoadSelection(_selIdx);
}
<commit_msg>Implement media key chooser in hotkeys tab<commit_after>#include "stdafx.h"
#include "SettingsUI.h"
#include "Hotkeys.h"
#include "afxdialogex.h"

#include <unordered_map>

#include "HotkeyInfo.h"
#include "HotkeyPrompt.h"
#include "KeyGrabber.h"
#include "Settings.h"

IMPLEMENT_DYNAMIC(Hotkeys, CPropertyPage)

Hotkeys::Hotkeys() :
CPropertyPage(Hotkeys::IDD) {

}

Hotkeys::~Hotkeys() {
}

void Hotkeys::DoDataExchange(CDataExchange* pDX)
{
    CPropertyPage::DoDataExchange(pDX);
    DDX_Control(pDX, LST_KEYS, _list);
    DDX_Control(pDX, BTN_KEYS, _keys);
    DDX_Control(pDX, CMB_ACTION, _action);
    DDX_Control(pDX, LBL_ARG, _argument);
    DDX_Control(pDX, CMB_DRIVE, _argCombo);
}

BOOL Hotkeys::OnApply() {
    OutputDebugString(L"-> Hotkeys\n");
    Settings *settings = Settings::Instance();
    settings->Hotkeys(_keyInfo);
    UIUtils::SaveSettings(*this);
    return CPropertyPage::OnApply();
}

BOOL Hotkeys::OnCommand(WPARAM wParam, LPARAM lParam) {
    SetModified();
    return CPropertyPage::OnCommand(wParam, lParam);
}

BOOL Hotkeys::OnInitDialog() {
    CPropertyPage::OnInitDialog();

    Settings *settings = Settings::Instance();

    DWORD exStyle = _list.GetExtendedStyle();
    exStyle |= LVS_EX_FULLROWSELECT;
    _list.SetExtendedStyle(exStyle);

    /* Set up listview and editor area */
    RECT r;
    _list.GetWindowRect(&r);
    int width = r.right - r.left;

    _list.InsertColumn(0, L"Hotkeys", NULL, (int) (width * .485));
    _list.InsertColumn(1, L"Action", NULL, (int) (width * .445));

    for (std::wstring action : HotkeyInfo::ActionNames) {
        _action.AddString(action.c_str());
    }

    std::unordered_map<int, HotkeyInfo> hotkeys = settings->Hotkeys();
    for (auto it = hotkeys.begin(); it != hotkeys.end(); ++it) {
        _keyInfo.push_back(it->second);
    }

    for (unsigned int i = 0; i < _keyInfo.size(); ++i) {
        LoadSelection(i);
        HotkeyInfo hki = _keyInfo[i];
        std::wstring hkStr = HotkeyManager::HotkeysToString(hki.keyCombination);
        _list.InsertItem(i, hkStr.c_str());
        _list.SetItemText(i, 1, HotkeyInfo::ActionNames[hki.action].c_str());
    }

    SelectItem(0);

    return TRUE;
}

void Hotkeys::SelectItem(int idx) {
    /* First unselect the selected item (not strictly necessary in our case) */
    int sel = _list.GetSelectionMark();
    if (sel >= 0) {
        _list.SetItemState(sel, ~LVIS_SELECTED, LVIS_SELECTED);
    }

    /* Make sure we're not selecting something out of range */
    if (idx < 0) {
        idx = 0;
    }
    int numItems = _list.GetItemCount();
    if (idx > numItems - 1) {
        idx = numItems - 1;
    }

    /* Select the item */
    _list.SetItemState(idx, LVIS_SELECTED, LVIS_SELECTED);
    _list.SetSelectionMark(idx);

    /* Scroll down to the item, if necessary */
    _list.EnsureVisible(idx, FALSE);
}

void Hotkeys::LoadSelection(int idx) {
    HotkeyInfo current = _keyInfo[_selIdx];

    _keys.SetWindowText(L"");

    if (current.keyCombination > 0) {
        std::wstring keyStr
            = HotkeyManager::HotkeysToString(current.keyCombination);
        _keys.SetWindowText(keyStr.c_str());
        _list.SetItemText(_selIdx, 0, keyStr.c_str());
    }

    if (current.action >= 0) {
        _list.SetItemText(_selIdx, 1,
            HotkeyInfo::ActionNames[current.action].c_str());

        switch ((HotkeyInfo::HotkeyActions) current.action) {
        case HotkeyInfo::EjectDrive:
            _argument.SetWindowTextW(L"Drive:");
            _argCombo.ResetContent();
            for (int i = 0; i < 26; ++i) {
                wchar_t ch = (wchar_t) i + 65;
                _argCombo.InsertString(i, CString(ch));
            }
            if (current.args.size() > 0) {
                _argCombo.SelectString(0, current.args[0].c_str());
            }

            _argument.ShowWindow(SW_SHOW);
            _argCombo.ShowWindow(SW_SHOW);
            break;

        case HotkeyInfo::MediaKey: {
            _argument.SetWindowTextW(L"Action:");
            _argCombo.ResetContent();
            unsigned int numMediaKeys = HotkeyInfo::MediaKeyNames.size();
            for (unsigned int i = 0; i < numMediaKeys; ++i) {
                _argCombo.InsertString(i, HotkeyInfo::MediaKeyNames[i].c_str());
            }
            if (current.args.size() > 0) {
                //TODO: this needs to be translated again
                _argCombo.SelectString(0, current.args[0].c_str());
            }

            _argument.ShowWindow(SW_SHOW);
            _argCombo.ShowWindow(SW_SHOW);
        }
            break;

        default:
            _argument.ShowWindow(SW_HIDE);
            _argCombo.ShowWindow(SW_HIDE);
        }
    } else {
        /* TODO fix these nasty repeated hides */
        _argument.ShowWindow(SW_HIDE);
        _argCombo.ShowWindow(SW_HIDE);
    }

    _action.SetCurSel(current.action);
}

BEGIN_MESSAGE_MAP(Hotkeys, CPropertyPage)
    ON_BN_CLICKED(BTN_ADD, &Hotkeys::OnBnClickedAdd)
    ON_BN_CLICKED(BTN_REMOVE, &Hotkeys::OnBnClickedRemove)
    ON_NOTIFY(LVN_ITEMCHANGED, LST_KEYS, &Hotkeys::OnLvnItemchangedKeys)
    ON_BN_CLICKED(BTN_KEYS, &Hotkeys::OnBnClickedKeys)
    ON_CBN_SELCHANGE(CMB_ACTION, &Hotkeys::OnCbnSelchangeAction)
    ON_CBN_SELCHANGE(CMB_ARG, &Hotkeys::OnCbnSelchangeArg)
END_MESSAGE_MAP()

void Hotkeys::OnBnClickedAdd() {
    int items = _list.GetItemCount();
    for (int i = 0; i < items; ++i) {
        if (_list.GetItemText(i, 0) == L""
            && _list.GetItemText(i, 1) == L"") {
            /* We found a blank item already in the list */
            SelectItem(i);
            return;
        }
    }

    HotkeyInfo hi;
    _keyInfo.push_back(hi);
    int idx = _list.InsertItem(items, L"");
    _list.SetItemText(idx, 1, L"");
    SelectItem(idx);
}

void Hotkeys::OnBnClickedRemove() {
    _keyInfo.erase(_keyInfo.begin() + _selIdx);
    _list.DeleteItem(_selIdx);

    /* Select the item closest to the previous selection: */
    SelectItem(_selIdx);
}

void Hotkeys::OnLvnItemchangedKeys(NMHDR *pNMHDR, LRESULT *pResult) {
    LPNMLISTVIEW pNMLV = reinterpret_cast<LPNMLISTVIEW>(pNMHDR);

    if (pNMLV->uChanged & LVIF_STATE) {
        if (pNMLV->uNewState & LVIS_SELECTED) {
            int sel = pNMLV->iItem;
            _selIdx = sel;
            OutputDebugString(std::to_wstring(sel).c_str());
            OutputDebugString(L"\n");
            LoadSelection(sel);
        }
    }

    *pResult = 0;
}

void Hotkeys::OnBnClickedKeys() {
    HotkeyPrompt hkp;
    hkp.DoModal();
    KeyGrabber::Instance()->Unhook();
    int keyCombo = KeyGrabber::Instance()->KeyCombination();
    if (keyCombo > 0) {
        std::wstring keyStr = HotkeyManager::HotkeysToString(keyCombo);
        _keys.SetWindowText(keyStr.c_str());
        _keyInfo[_selIdx].keyCombination = keyCombo;
        LoadSelection(_selIdx);
    }
}

void Hotkeys::OnCbnSelchangeAction() {
    int sel = _action.GetCurSel();
    _keyInfo[_selIdx].action = sel;
    LoadSelection(_selIdx);
}

void Hotkeys::OnCbnSelchangeArg() {
    HotkeyInfo *current = &_keyInfo[_selIdx];

    HotkeyInfo::HotkeyActions action
        = (HotkeyInfo::HotkeyActions) _action.GetCurSel();

    switch (action) {
    case HotkeyInfo::EjectDrive: {
        wchar_t buf[2];
        _argCombo.GetWindowText(buf, 2);
        if (current->args.size() <= 0) {
            current->args.resize(1);
        }
        current->args[0] = buf;
    }
        break;

    case HotkeyInfo::MediaKey: {
        wchar_t buf[128];
        _argCombo.GetWindowText(buf, 128);
        if (current->args.size() <= 0) {
            current->args.resize(1);
        }
        current->args[0] = buf;
    }
        break;

    }
    LoadSelection(_selIdx);
}
<|endoftext|>