#include <iostream>
#include <functional>
#include <limits>
#include <memory>
#include <vector>
#include <cstddef>
#include <memory>

#include "generator_stack.h"

void* operator new([[maybe_unused]]std::size_t n)
{
	std::cout << "new\n";
	throw std::bad_alloc();
}

void operator delete(void*) throw()
{
	std::cout << "delete\n";
	// Do nothing. actual memory is allocated on the stack
}

void operator delete(void*, std::size_t) throw()
{
	std::cout << "delete[]\n";
	// Do nothing. actual memory is allocated on the stack
}

// Code heavily inspired by:
//  https://stackoverflow.com/questions/66891368/template-parametric-type-allocator-in-c
template <typename T, std::size_t S>
struct StackAllocator
{
	using value_type = T;
	using size_type = size_t;
	using pointer = value_type*;

	template<class U> struct rebind { using other = StackAllocator<U, S>; };

	StackAllocator() = default;
	StackAllocator(const StackAllocator&) = default;
	template <typename U> StackAllocator(const StackAllocator<U, S>&) {}

	pointer allocate(size_type n)
	{
		if(n > S) throw std::bad_alloc();
		return buf;
	}

	void deallocate(void*, size_type)
	{
	}

private:
	T buf[S];
};

template<typename T>
Generator<T, 64> iota(T start)
{
	while(true)
	{
		co_yield start++;
	}
}

int main()
{
	std::cout << " ** std::string:\n";
	{
		std::basic_string<char, std::char_traits<char>, StackAllocator<char, 32>> str(31, 'a');
		std::cout << str << " " << sizeof(str) << '\n';

		std::basic_string<char, std::char_traits<char>, StackAllocator<char, 64>> str2(31, 'a');
		std::cout << str2 << " " << sizeof(str2) << '\n';
	}
	std::cout << '\n';

	std::cout << " ** std::vector:\n";
	{
		try
		{
			std::vector<int, StackAllocator<int, 10>> vec{42};
		}
		catch(std::bad_alloc &e)
		{
			std::cout << "No SBO for std::vector!\n";
		}
	}
	std::cout << '\n';


	std::cout << " ** std::function:\n";
	{
//		{ // Support for allocators with std::function removed in C++17
//			std::function<std::allocator_arg, StackAllocator<int, 10>, int()> f;
//			char foo[16]{};
//			f = [foo]() // capture up to 16 bytes is ok for std::function
//			    {
//				    int i = 0;
//				    for(auto v : foo) i += v;
//				    return i;
//			    };
//		}
	}
	std::cout << '\n';


	std::cout << " ** co-routines:\n";
	{
		try
		{
			auto gen = iota(0);
			while(true)
			{
				auto i = gen();
				if(i > 10) break;
				std::cout << i << " ";
			}
			std::cout << '\n';
		}
		catch(std::bad_alloc &e)
		{
			std::cout << "Co-routines always allocate!\n";
		}
	}
	std::cout << '\n';
}