export module coral.stack;

import coral;

// Collections.
export namespace coral {
	/**
	 * Result codes used by [contiguous_range]-derived types when they are appended to in any way.
	 *
	 * [push_result::ok] indicates that an push operation was successful.
	 *
	 * [push_result::out_of_memory] alerts that the memory required to perform the push operation
	 * failed.
	 */
	enum class [[nodiscard]] push_result {
		ok,
		out_of_memory,
	};

	/**
	 * Base type for all stack types.
	 *
	 * Sequences are any data structure which owns a linear, non-unique set of elements which may
	 * be queried and/or mutated.
	 */
	template<typename element> struct stack {
		virtual ~stack() {};

		/**
		 * Returns a read-only [slice] of the current range values.
		 *
		 * *Note*: the behavior of retaining the returned value past the scope of the source
		 * [stack] or any subsequent modifications to it is implementation-defined.
		 */
		virtual slice<element const> as_slice() const = 0;

		/**
		 * Attempts to append `source_elements` to the stack.
		 *
		 * The returned [push_result] indicates whether the operation was successful or not.
		 *
		 * If the returned [push_result] is anything but [push_result::ok], the [stack] will be
		 * left in an implementation-defined state.
		 */
		virtual push_result push_all(slice<element const> const & source_elements) = 0;
	};

	/**
	 * Last-in-first-out contiguous sequence of `element` values optimized for small numbers of
	 * small-sized elements.
	 *
	 * [small_stack] types will default to using an inline array of `init_capacity` at first. After
	 * all local storage has been exhausted, the [small_stack] will switch to a dynamic buffer.
	 * Because of this, it is recommended to use larger `init_capacity` values for data which has a
	 * known or approximate upper bound at compile-time. Otherwise, the `init_capacity` value may
	 * be left at its default.
	 *
	 * *Note*: the [allocator] referenced in the stack must remain valid for the duration of the
	 * stack lifetime.
	 */
	template<typename element, usize init_capacity = 1> struct small_stack : public stack<element> {
		small_stack(allocator * dynamic_allocator) {
			this->dynamic_allocator = dynamic_allocator;
		}

		~small_stack() override {
			if (this->is_dynamic()) {
				for (element & e : this->elements) e.~element();

				this->dynamic_allocator->deallocate(this->elements.pointer);
			}
		}

		/**
		 * Returns a read-only [slice] of the current stack values.
		 *
		 * *Note*: the returned slice should be considered invalid if any mutable operation is
		 * performed on the source [stack] or it is no longer in scope.
		 */
		slice<element const> as_slice() const override {
			return this->elements.sliced(0, this->filled);
		}

		/**
		 * Returns `true` if the stack is backed by dynamic memory, otherwise `false`.
		 */
		bool is_dynamic() const {
			return this->elements.pointer != reinterpret_cast<element const *>(this->local_buffer);
		}

		/**
		 * Attempts to append `source_element` to the top of the stack.
		 *
		 * The returned [push_result] indicates whether the operation was successful or not.
		 *
		 * If the returned [push_result] is anything but [push_result::ok], the stack will be
		 * be left in an empty but valid state.
		 *
		 * *Note* that [push_all] is recommended when appending many values at once.
		 */
		push_result push(element const & source_element) {
			if (this->filled == this->elements.length) {
				push_result const result = this->reserve(this->elements.length);

				if (result != push_result::ok) return result;
			}

			this->elements[this->filled] = source_element;
			this->filled += 1;

			return push_result::ok;
		}

		/**
		 * Attempts to append `source_elements` to the top of the stack.
		 *
		 * The returned [push_result] indicates whether the operation was successful or not.
		 *
		 * If the returned [push_result] is anything but [push_result::ok], the stack will be left
		 * in an empty but valid state.
		 *
		 * *Note* that [push] is recommended when appending singular values.
		 */
		push_result push_all(slice<element const> const & source_elements) override {
			usize const updated_fill = this->filled + source_elements.length;

			if (updated_fill >= this->elements.length) {
				push_result const result = this->reserve(updated_fill);

				if (result != push_result::ok) return result;
			}

			for (usize i = 0; i < source_elements.length; i += 1)
				this->elements[this->filled + i] = source_elements[i];

			this->filled = updated_fill;

			return push_result::ok;
		}

		/**
		 * Attempts to reserve `capacity` number of elements additional space on the stack, forcing
		 * it to use dynamic memory _even_ if it hasn't exhausted the local buffer yet.
		 *
		 * The returned [push_result] indicates whether the operation was successful or not.
		 *
		 * If the returned [push_result] is anything but [push_result::ok], the stack will be left
		 * in an empty but valid state.
		 *
		 * *Note* that manual invocation is not recommended if the [stack] has a large
		 * `initial_capacity` argument.
		 */
		push_result reserve(usize capacity) {
			usize const requested_capacity = this->filled + capacity;

			if (this->is_dynamic()) {
				u8 * const buffer = this->dynamic_allocator->reallocate(
					reinterpret_cast<u8 *>(this->elements.pointer),
						sizeof(element) * requested_capacity);

				if (buffer == nullptr) {
					this->elements = {};

					return push_result::out_of_memory;
				}

				this->elements = {reinterpret_cast<element *>(buffer), requested_capacity};
			} else {
				usize const buffer_size = sizeof(element) * requested_capacity;
				u8 * const buffer = this->dynamic_allocator->reallocate(nullptr, buffer_size);

				if (buffer == nullptr) {
					this->elements = {};

					return push_result::out_of_memory;
				}

				copy({buffer, buffer_size}, this->elements.as_bytes());

				this->elements = {reinterpret_cast<element *>(buffer), requested_capacity};
			}

			return push_result::ok;
		}

		private:
		allocator * dynamic_allocator{nullptr};

		usize filled{0};

		slice<element> elements{reinterpret_cast<element *>(local_buffer), init_capacity};

		u8 local_buffer[init_capacity]{0};
	};
}

using byte_stack = coral::stack<coral::u8>;

// Reader / writers.
export namespace coral {
	/**
	 * Readable type for streaming data from a [stack] containing [u8] values.
	 */
	struct stack_reader : public reader {
		stack_reader(byte_stack const * stack) {
			this->cursor = 0;
			this->stack = stack;
		}

		/**
		 * Reads the data from the target stack into `buffer`, returning the number bytes read.
		 */
		expected<usize, io_error> read(slice<u8> const & buffer) override {
			slice const stack_elements {this->stack->as_slice()};
			usize const read {min(buffer.length, stack_elements.length - this->cursor)};

			copy(buffer, stack_elements.sliced(cursor, read));

			this->cursor += read;

			return read;
		}

		private:
		usize cursor {0};

		byte_stack const * stack {nullptr};
	};

	/**
	 * Writable type for appending data to a [contiguous_range] containing [u8] values.
	 */
	struct stack_writer : public writer {
		stack_writer(byte_stack * stack) {
			this->stack = stack;
		}

		/**
		 * Attempts to write `buffer` to the target stack, returning the number of bytes written or
		 * an [io_error] if it failed to commit `buffer` to the stack memory.
		 */
		expected<usize, io_error> write(slice<u8 const> const & buffer) override {
			switch (this->stack->push_all(buffer)) {
				case push_result::ok: return buffer.length;
				case push_result::out_of_memory: return io_error::unavailable;
				default: unreachable();
			}
		}

		private:
		byte_stack * stack {nullptr};
	};
}