#pragma once
#include "Thread.h"
#include "Value.h"
#include "ActiveFunctions.h"
#include "Gc/ObjMap.h"

namespace storm {
	STORM_PKG(core.lang);

	class Named;
	class Function;
	class VTable;
	class ReplaceContext;
	class TypeTransform;
	class ReplaceTfmWalker;

	/**
	 * An object encapsulating remaining tasks to do after replacing one or more objects, and other
	 * state that is needed while replacing objects.
	 *
	 * These tasks are things that cannot be made by the replacing objects themselves, for example
	 * since they require the GC to be paused or other similar things. Operations that are
	 * beneficial to perform in batch whenever all replacements have been performed are also queued
	 * up in this object.
	 *
	 * This encompasses:
	 * - Replacing all references of one object with another.
	 * - Modifying the layout of objects in memory.
	 *
	 * Note: If replacing object A with B, and both A and B are keys in a Map<>, future lookups and
	 * modifications of B might return inconsistent results until a full rehash is made (due to
	 * moved objects, not due to capacity). This does not seem to happen very often in the system
	 * currently, as we only allow replacing Named and Handle.
	 */
	class ReplaceTasks : public ObjectOn<Compiler> {
		STORM_CLASS;
	public:
		// Create, without information about active functions (mainly for unit-tests).
		ReplaceTasks();

		// Create, provide information about active functions.
		ReplaceTasks(PauseThreads &from);

		// Create, provide an array to store exceptions in.
		ReplaceTasks(PauseThreads &from, Array<Exception *> *exceptions);

		// Destroy.
		~ReplaceTasks();

		/**
		 * Replacing references.
		 */

		// Schedule a reference to be replaced.
		void STORM_FN replace(Named *old, Named *with);
		void STORM_FN replace(const Handle *old, const Handle *with);

		// Schedule a GcType reference to be replaced. Note: We don't touch 'myGcType' inside Type objects.
		void replace(const GcType *old, const GcType *with);

		// VTable replacement. This refers to the content of the vtable rather than the vtable itself.
		// VTables are not replaced globally, only the 'vtable' field of the type is actually modified.
		void STORM_FN replace(VTable *old, VTable *with);

		// Transform a type that has been updated. Done instead of 'replace' of VTables and GcTypes.
		void STORM_FN transform(TypeTransform *transform);

		// Check if a type is transformed.
		Bool STORM_FN hasTransformFor(Type *type);

		// Schedule an update of a potentially active function.
		void STORM_FN replaceActive(Function *newCode);

		// Apply changes requested. We don't allow doing this from Storm. It is not necessarily
		// good to do while other threads are running.
		void apply();


		/**
		 * Error reporting.
		 */

		// Report errors during the replacement process. Note that errors during the replacement
		// process must not be fatal - it is often not possible to roll back the replacement process
		// once it has been started. Note, that all involved entities have the chance to check for
		// obvious issues beforehand, so this type of errors should be comparatively rare. These
		// errors will all be replaced as one big lump at the end of the update process, so that the
		// user at least knows about them.
		void STORM_FN error(Exception *exception);

		// Throw the stored errors if there are any.
		void STORM_FN throwErrors();


		/**
		 * Inspecting/replacing active functions. Mainly used from ReplaceActive.cpp.
		 */

		// Check if a function is currently active. Returns one element for each active instance of
		// the function, each indicating the current offset of the return location into the
		// function.
		vector<ActiveOffset> findActive(const void *function) const;

		// Relace an active function by replacing the return address on the stack. Replaces all
		// instances of 'offset' inside 'function' with 'replace'. 'replace' is typically the start
		// of a newly generated thunk that migrates control flow to the new version of the function.
		// The system does, however, not make assumptions about this except that 'function' and
		// 'replace' need to be allocated code blocks on the GC heap.
		size_t replaceActive(const void *function, size_t fOffset, const void *replace, size_t rOffset) const;

	private:
		// References to replace.
		UNKNOWN(PTR_NOGC) RawObjMap *replaceMap;

		// VTables to replace.
		UNKNOWN(PTR_NOGC) RawObjMap *vtableMap;

		// Currently active functions.
		UNKNOWN(PTR_NOGC) ActiveFunctions *activeFunctions;

		// Store functions that should be updated.
		Array<Function *> *functionsToUpdate;

		// All types to transform.
		Map<Type *, TypeTransform *> *transforms;

		// All stored exceptions.
		Array<Exception *> *exceptions;

		// Helper to transform objects on the heap.
		void applyTransforms(ReplaceTfmWalker &tfm);
	};

}
