
#ifndef CONTEXT_SWITCHER_LAMBDA_CAPTURE_H_
#define CONTEXT_SWITCHER_LAMBDA_CAPTURE_H_

#include <utility>
#include <new>
#include <type_traits>

#include "context.h"
#include "continuation.h"

#include "tsan_fiber_api.h"

/**
 * Helpers to more easily use a lambda expression as the code run on a context switch.
 * Captures lambdas by placing them in a memory region and offering a conversion form
 * the assembly_bindings callback API to the user's lambda.
 */

namespace context_switcher {

template<typename F>
struct lambda_capture {
  template<typename FARG>
  explicit lambda_capture(FARG &&lambda) : lambda_{std::forward<FARG>(lambda)} {}

#ifdef THREAD_SANITIZER
  continuation operator()(continuation_t cont) {
    return lambda_(continuation{cont, previous_tsan_fiber_});
  }
  void *previous_tsan_fiber_;
  void *my_tsan_fiber_;
#else
  continuation operator()(continuation_t cont) {
    return lambda_(continuation{cont});
  }
#endif

 private:
  F lambda_;
};

#ifdef CS_USE_CSCONTEXT
template<typename T>
cscontext::continuation_t lambda_capture_callback(cscontext::continuation_t continuation_pointer,
                                                  void *lambda_capture_param) {
  // Perform Call
  T *lambda_capture = reinterpret_cast<T *>(lambda_capture_param);
  continuation cont = (*lambda_capture)(continuation_pointer);

  // Free resources and switch to result_continuation (this execution thread is finished with the return)
  lambda_capture->~T();
  return cont.consume();
}
#else // CS_USE_CSCONTEXT
continuation switch_context(continuation &&cont);
template<typename T>
void lambda_capture_callback(fcontext::transfer_t transfer) {
  // Perform Call
  T *lambda_capture = reinterpret_cast<T *>(transfer.data);
  continuation cont = (*lambda_capture)(transfer.continuation);

  // Free resources and switch to result_continuation (this execution thread is finished with the return)
  lambda_capture->~T();

#ifdef THREAD_SANITIZER
  continuation_t cont_pointer = cont.consume();
  void *next_fiber = cont.get_tsan_fiber();
  __tsan_switch_to_fiber(next_fiber, 0);
  __tsan_destroy_fiber(lambda_capture->my_tsan_fiber_);
  fcontext::jump_fcontext(cont_pointer, (void*)0);
#else // THREAD_SANITIZER
  continuation_t cont_pointer = cont.consume();
  fcontext::jump_fcontext(cont_pointer, (void *) 0);
#endif // THREAD_SANITIZER
}
#endif // CS_USE_CSCONTEXT

template<typename F>
static lambda_capture<typename std::remove_reference<F>::type> *place_lambda_capture(F &&lambda, char *memory) {
  return new(memory) lambda_capture<typename std::remove_reference<F>::type>(std::forward<F>(lambda));
}

}

#endif //CONTEXT_SWITCHER_LAMBDA_CAPTURE_H_
