#pragma once
#include "Function.h"
#include "Core/Array.h"
#include "Core/Hash.h"
#include "Code/Listing.h"
#include "Code/Binary.h"
#include "BuiltInSource.h"

namespace storm {
	STORM_PKG(core.lang);

	/**
	 * Type used to represent control flow lists, and to compute mappings between old and new
	 * versions of active functions.
	 *
	 * This item either represents:
	 * - A function call in a function,
	 * - A loop in a function, which in turn contains a list of ControlFlowItems, or
	 * - The start of the function,
	 */
	class ControlFlowItem {
		STORM_VALUE;
	private:
		enum {
			// Number of bits to use for status flags.
			statusBits = 1,
			// First ID used for end offset.
			loopEndShift = 3,
		};
	public:
		// Status of the item. Mainly used to add information about the match from 'diff'.
		enum Status {
			// No particular status. The default.
			none = 0,
			// The original item was removed, so the original item was matched with something else.
			removed,
		};

		// Create, represent the start of the contained function.
		STORM_CTOR ControlFlowItem()
			: data(null), offsetFlags(0), typeInfo(0) {}

		// Create, represent a call to a Function.
		STORM_CTOR ControlFlowItem(Nat offset, Function *call)
			: data(call), offsetFlags(offset << statusBits), typeInfo(1) {}

		// Create, represent a call to a built in element.
		STORM_CTOR ControlFlowItem(Nat offset, BuiltInSource *builtIn)
			: data(builtIn), offsetFlags(offset << statusBits), typeInfo(2) {}

		// Create, represent a loop.
		STORM_CTOR ControlFlowItem(Nat startOffset, Nat endOffset, Array<ControlFlowItem> *body)
			: data(body), offsetFlags(startOffset << statusBits), typeInfo(endOffset + loopEndShift) {}

		// Offset of this item in the source code:
		// - If we are representing machine code, the offset is in bytes.
		// - If we are representing uncompiled IR, the offset is in instructions.
		Nat STORM_FN offset() const {
			return offsetFlags >> statusBits;
		}

		// Set the offset.
		void STORM_FN offset(Nat v) {
			offsetFlags &= (Nat(1) << statusBits) - 1;
			offsetFlags |= v << statusBits;
		}

		// End offset. Only meaningful for loops.
		Nat STORM_FN endOffset() const {
			if (isLoop())
				return typeInfo - loopEndShift;
			else
				return offset();
		}

		// Set the end offset.
		void STORM_FN endOffset(Nat v) {
			if (isLoop())
				typeInfo = v + loopEndShift;
		}

		// Get status.
		Status STORM_FN status() const {
			return Status(offsetFlags & ((Nat(1) << statusBits) - 1));
		}

		// Set status.
		void STORM_FN status(Status v) {
			offsetFlags &= ~((Nat(1) << statusBits) - 1);
			offsetFlags |= Nat(v);
		}

		// Is this the start of the function?
		Bool STORM_FN isStart() const {
			return data == null;
		}

		// Is this a function node?
		Bool STORM_FN isCall() const {
			return typeInfo == 1 || typeInfo == 2;
		}

		// Is this a call to a Function?
		Bool STORM_FN hasFunction() const {
			return typeInfo == 1;
		}

		// Get the function to call. Assumes 'isCall' and 'hasFunction' returns true.
		Function *STORM_FN function() const {
			STORM_ASSERT(hasFunction());
			return (Function *)data;
		}

		// Is this a call to a built-in function?
		Bool STORM_FN hasBuiltIn() const {
			return typeInfo == 2;
		}

		// Get the built-in function to call. Assumes 'isCall' and 'hasBuiltIn'.
		BuiltInSource *STORM_FN builtIn() const {
			STORM_ASSERT(hasBuiltIn());
			return (BuiltInSource *)data;
		}

		// Is this a loop node?
		Bool STORM_FN isLoop() const {
			return typeInfo >= Nat(loopEndShift);
		}

		// Get the loop body. Assumes 'isLoop' returns true.
		Array<ControlFlowItem> *STORM_FN loop() const {
			STORM_ASSERT(isLoop());
			return (Array<ControlFlowItem> *)data;
		}

		// Flatten into a list. Note, loop nodes are kept in the list.
		Array<ControlFlowItem> *STORM_FN flatten() const;

		// Replace all offsets:
		void STORM_FN replaceOffsets(Map<Nat, Nat> *replace);

		// Output:
		void STORM_FN toS(StrBuf *to) const;

		// Output suffix of the status.
		static void statusSuffix(StrBuf *to, Status status);

		// Compare. Considers offset and flags.
		Bool STORM_FN operator <(const ControlFlowItem &o) const {
			if (offset() != o.offset())
				return offset() < o.offset();
			if (endOffset() != o.endOffset())
				return endOffset() < o.endOffset();
			return Nat(status()) < Nat(o.status());
		}

		// Compare. Considers offset and flags.
		Bool STORM_FN operator ==(const ControlFlowItem &o) const {
			return offset() == o.offset()
				&& endOffset() == o.endOffset()
				&& status() == o.status();
		}

		// Hash, for use in hash tables.
		Nat STORM_FN hash() const {
			return natHash(offsetFlags) ^ natHash(endOffset());
		}

	private:
		// Data in this node. Either a pointer to a Function object, or an array.
		UNKNOWN(PTR_GC) RootObject *data;

		// Start offset and flags.
		Nat offsetFlags;

		// End offset (only for loops). Also indicates if 'data' is a Function or a BuiltInSource:
		// 0 - Start.
		// 1 - Function or null
		// 2 - BuiltInSource
		// 3+- Offset + 3
		Nat typeInfo;
	};

	// Flatten an array of elements.
	Array<ControlFlowItem> *STORM_FN flatten(Array<ControlFlowItem> *items);

	// Replace all offsets in an array of control flow items.
	void STORM_FN replaceOffsets(Array<ControlFlowItem> *items, Map<Nat, Nat> *replace);

	// Generate a control flow list from a Listing.
	Array<ControlFlowItem> *STORM_FN controlFlowList(code::Listing *code);

	// Generate a control flow list from a pre-compiled function.
	Array<ControlFlowItem> *STORM_FN controlFlowList(code::Binary *code);
	Array<ControlFlowItem> *controlFlowListRaw(void *code);
}
