// Copyright 2019 The SwiftShader Authors. All Rights Reserved.
//
// 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 "Reactor.hpp"

#include <memory>

#ifndef rr_ReactorCoroutine_hpp
#define rr_ReactorCoroutine_hpp

namespace rr
{
	// Base class for the template Stream<T>
	class StreamBase
	{
	protected:
		StreamBase(const std::shared_ptr<Routine> &routine, Nucleus::CoroutineHandle handle)
			: routine(routine), handle(handle) {}

		~StreamBase()
		{
			auto pfn = (Nucleus::CoroutineDestroy*)routine->getEntry(Nucleus::CoroutineEntryDestroy);
			pfn(handle);
		}

		bool await(void* out)
		{
			auto pfn = (Nucleus::CoroutineAwait*)routine->getEntry(Nucleus::CoroutineEntryAwait);
			return pfn(handle, out);
		}

private:
		std::shared_ptr<Routine> routine;
		Nucleus::CoroutineHandle handle;
	};

	// Stream is the interface to a running Coroutine instance.
	// A Coroutine may Yield() values of type T, which can be retrieved with
	// await().
	template<typename T>
	class Stream : public StreamBase
	{
	public:
		inline Stream(const std::shared_ptr<Routine> &routine, Nucleus::CoroutineHandle handle)
			: StreamBase(routine, handle) {}

		// await() retrieves the next yielded value from the coroutine.
		// Returns true if the coroutine yieled a value and out was assigned a
		// new value. If await() returns false, the coroutine has finished
		// execution and await() will return false for all future calls.
		inline bool await(T& out) { return StreamBase::await(&out); }
	};

	template<typename FunctionType>
	class Coroutine;

	// Coroutine constructs a reactor Coroutine function.
	// rr::Coroutine is similar to rr::Function in that it builds a new
	// executable function, but Coroutines have the following differences:
	//  (1) Coroutines do not support Return() statements.
	//  (2) Coroutines support Yield() statements to suspend execution of the
	//      coroutine and pass a value up to the caller. Yield can be called
	//      multiple times in a single execution of a coroutine.
	//  (3) The template argument T to Coroutine<T> is a C-style function
	//      signature.
	//  (4) Coroutine::operator() returns a rr::Stream<T> instead of an
	//      rr::Routine.
	//  (5) operator() starts execution of the coroutine immediately.
	//  (6) operator() uses the Coroutine's template function signature to
	//      ensure the argument types match the generated function signature.
	//
	// Example usage:
	//
	//   // Build the coroutine function
	//   Coroutine<int()> coroutine;
	//   {
	//       Yield(Int(0));
	//       Yield(Int(1));
	//       Int current = 1;
	//       Int next = 1;
	//       While (true) {
	//           Yield(next);
	//           auto tmp = current + next;
	//           current = next;
	//           next = tmp;
	//       }
	//   }
	//
	//   // Start the execution of the coroutine.
	//   auto s = coroutine();
	//
	//   // Grab the first 20 yielded values and print them.
	//   for (int i = 0; i < 20; i++)
	//   {
	//       int val = 0;
	//       s->await(val);
	//       printf("Fibonacci(%d): %d", i, val);
	//   }
	//
	template<typename Return, typename... Arguments>
	class Coroutine<Return(Arguments...)>
	{
	public:
		Coroutine();

		template<int index>
		using CArgumentType = typename std::tuple_element<index, std::tuple<Arguments...>>::type;

		template<int index>
		using RArgumentType = CToReactor<CArgumentType<index>>;

		// Return the argument value with the given index.
		template<int index>
		Argument<RArgumentType<index>> Arg() const
		{
			Value *arg = Nucleus::getArgument(index);
			return Argument<RArgumentType<index>>(arg);
		}

		// Completes building of the coroutine and generates the coroutine's
		// executable code. After calling, no more reactor functions may be
		// called without building a new rr::Function or rr::Coroutine.
		// While automatically called by operator(), finalize() should be called
		// as early as possible to release the global Reactor mutex lock.
		inline void finalize(OptimizationLevel optLevel = OptimizationLevel::Default);

		// Starts execution of the coroutine and returns a unique_ptr to a
		// Stream<> that exposes the await() function for obtaining yielded
		// values.
		std::unique_ptr<Stream<Return>> operator()(Arguments...);

	protected:
		std::unique_ptr<Nucleus> core;
		std::shared_ptr<Routine> routine;
		std::vector<Type*> arguments;
	};

	template<typename Return, typename... Arguments>
	Coroutine<Return(Arguments...)>::Coroutine()
	{
		core.reset(new Nucleus());

		std::vector<Type*> types = {CToReactor<Arguments>::getType()...};
		for(auto type : types)
		{
			if(type != Void::getType())
			{
				arguments.push_back(type);
			}
		}

		Nucleus::createCoroutine(CToReactor<Return>::getType(), arguments);
	}

	template<typename Return, typename... Arguments>
	void Coroutine<Return(Arguments...)>::finalize(OptimizationLevel optLevel /* = OptimizationLevel::Default */)
	{
		if(core != nullptr)
		{
			routine.reset(core->acquireCoroutine("coroutine", optLevel));
			core.reset(nullptr);
		}
	}

	template<typename Return, typename... Arguments>
	std::unique_ptr<Stream<Return>>
	Coroutine<Return(Arguments...)>::operator()(Arguments... args)
	{
		finalize();

		using Sig = Nucleus::CoroutineBegin<Arguments...>;
		auto pfn = (Sig*)routine->getEntry(Nucleus::CoroutineEntryBegin);
		auto handle = pfn(args...);
		return std::unique_ptr<Stream<Return>>(new Stream<Return>(routine, handle));
	}

#ifdef Yield // Defined in WinBase.h
#undef Yield
#endif

	// Suspends execution of the coroutine and yields val to the caller.
	// Execution of the coroutine will resume after val is retrieved.
	template<typename T>
	inline void Yield(const T &val) { Nucleus::yield(ValueOf(val)); }

} // namespace rr

#endif // rr_ReactorCoroutine_hpp