#define __STDC_LIMIT_MACROS
#include "TokenBucket.h"
#include <string.h>

// Copyright (c) 2009 Ivan Skytte Jørgensen
// 
// This software is provided 'as-is', without any express or implied
// warranty. In no event will the authors be held liable for any damages
// arising from the use of this software.
// 
// Permission is granted to anyone to use this software for any purpose,
// including commercial applications, and to alter it and redistribute it
// freely, subject to the following restrictions:
// 
//     1. The origin of this software must not be misrepresented; you must not
//     claim that you wrote the original software. If you use this software
//     in a product, an acknowledgment in the product documentation would be
//     appreciated but is not required.
// 
//     2. Altered source versions must be plainly marked as such, and must not be
//     misrepresented as being the original software.
// 
//     3. This notice may not be removed or altered from any source
//     distribution.


// A note about the descision about what to make inline and what not:
//   Simple getters/setters are inline
//   The update() method is 28 instructions on x86-64 (102 on x86, 109 on POWER
//   32bit) so it should not be inlined.
//   Then it follows that the methods that unconditionally calls update()
//   shouldn't either.
//   The remaining methods are probably rarely called so they should not be
//   inlined either.


const TokenBucket::timeunit_t TokenBucket::never = UINT64_MAX;


bool TokenBucket::operator==(const TokenBucket &tb) const throw()
{
	//If it weren't for potential padding between members a plain memcmp() would be so much easier
	return rate==tb.rate &&
	       time_unit==tb.time_unit &&
	       limit==tb.limit &&
	       last_update==tb.last_update &&
	       tokens==tb.tokens &&
	       d==tb.d &&
	       yerr==tb.yerr &&
	       error==tb.error;
}


bool TokenBucket::operator!=(const TokenBucket &tb) const throw()
{
	return rate!=tb.rate ||
	       time_unit!=tb.time_unit ||
	       limit!=tb.limit ||
	       last_update!=tb.last_update ||
	       tokens!=tb.tokens ||
	       d!=tb.d ||
	       yerr!=tb.yerr ||
	       error!=tb.error;
}


void TokenBucket::clear() throw()
{
	tokens = 0;
	error = 0;
}


void TokenBucket::update(timeunit_t current_time) throw()
{
	timeunit_t time_passed_since_last_update = current_time-last_update;
	
	//Two simple optimizations when there is nothing to do (the algorithm
	//works fine without these optimizations)
	if(time_passed_since_last_update==0)
		return;
	if(tokens==limit) {
		last_update = current_time;
		return;
	}
	
	//Naïve calculation:
	//  uint64_t tokens_to_add = (time_passed_since_last_update/time_unit)*rate;
	//  tokens += tokens_to_add;
	//  if(tokens>limit)
	//    tokens=limit;
	//but that leads to loss of tokens due to the remainder of the division
	//thrown away. Instead a variant of Bresenham line algorithm is used.
	// variance: steep delta is ok; x is made in jumps
	
	tokencount_t tokens_to_add = d * time_passed_since_last_update;
	error += yerr * time_passed_since_last_update;
	if(error>=time_unit) {
		tokens_to_add += error/time_unit;
		error %= time_unit;
	}
	
	tokens += tokens_to_add;
	
	if(tokens>limit)
		tokens = limit;
	
	last_update = current_time;
}


bool TokenBucket::canGet(timeunit_t current_time, tokencount_t max_tokens) throw()
{
	update(current_time);
	return canGet(max_tokens);
}


TokenBucket::tokencount_t TokenBucket::getTokens(tokencount_t max_tokens) throw()
{
	if(max_tokens<=tokens) {
		tokens -= max_tokens;
		return max_tokens;
	} else {
		tokencount_t tmp = tokens;
		tokens = 0;
		return tmp;
	}
}


TokenBucket::tokencount_t TokenBucket::getTokens(timeunit_t current_time, tokencount_t max_tokens) throw()
{
	update(current_time);
	return getTokens(max_tokens);
}


void TokenBucket::putTokens(tokencount_t tokens_) throw()
{
	tokens += tokens_;
	if(tokens>limit)
		tokens = limit;
}


void TokenBucket::setTokens(tokencount_t tokens_) throw()
{
	if(tokens_<=limit)
		tokens = tokens_;
	else
		tokens = limit;
}


TokenBucket::timeunit_t TokenBucket::whenAvail(tokencount_t desired_tokens) const throw()
{
	if(desired_tokens>limit)
		return never;
	if(tokens>=desired_tokens)
		return last_update;
	if(rate==0)
		return never;
	tokencount_t tokens_needed = desired_tokens - tokens;
	
	//The calculation could take into account the remainder loss and current
	//error, but the overall estimation is ok because the estimation error is
	//not accumulative.
	//Simple (and slightly imprecise) calculation:
	timeunit_t time_needed = (tokens_needed/rate)*time_unit;
	if(time_needed==0 || tokens_needed%rate)
		time_needed += 1;
	
	return last_update + time_needed;
}


TokenBucket::timeunit_t TokenBucket::whenAvail(timeunit_t current_time, tokencount_t desired_tokens) throw()
{
	update(current_time);
	return whenAvail(desired_tokens);
}


void TokenBucket::swap(TokenBucket &tb) throw()
{
	//due to const members we do a raw memcpy
	TokenBucket tmp(0,1,0);
	memcpy(&tmp,this,sizeof(TokenBucket));
	memcpy(this,&tb,sizeof(TokenBucket));
	memcpy(&tb,&tmp,sizeof(TokenBucket));
}


#ifdef UNITTEST
#include <assert.h>
int main(void) {
	{
		TokenBucket tb(10,1,5);
		assert(tb.capacity()==5);
		assert(tb.tokensLeft()==0);
		assert(tb.empty());
		tb.update(1);
		assert(!tb.empty());
	}
	{
		TokenBucket tb(10,1,5);
		assert(tb.getTokens(1,1)==1);
		assert(tb.tokensLeft()==4);
		assert(tb.getTokens(1,1)==1);
		assert(tb.tokensLeft()==3);
		assert(tb.getTokens(1,1)==1);
		assert(tb.getTokens(1,1)==1);
		assert(tb.getTokens(1,1)==1);
		assert(tb.tokensLeft()==0);
		assert(tb.getTokens(1,1)==0);
		//make a leap in time
		assert(tb.getTokens(1000,1)==1);
		assert(tb.getTokens(1000,8)==4);
		assert(tb.getTokens(1000,1)==0);
	}
	{
		//make sure that the rounding error correction works
		TokenBucket tb(3,2,3);
		assert(tb.getTokens(1000,1)==1);
		assert(tb.tokensLeft()==2);
		assert(tb.getTokens(1000,1)==1);
		assert(tb.tokensLeft()==1);
		assert(tb.getTokens(1000,1)==1);
		assert(tb.tokensLeft()==0);
		//leap half a time unit
		assert(tb.getTokens(1001,1)==1);
		assert(tb.getTokens(1001,1)==0);
		//half a time unit again
		assert(tb.getTokens(1002,1)==1);
		assert(tb.getTokens(1002,1)==1);
		//a whole time unit
		assert(tb.getTokens(1004,1)==1);
		assert(tb.getTokens(1004,1)==1);
		assert(tb.getTokens(1004,1)==1);
		assert(tb.getTokens(1004,1)==0);
	}
	{
		//verify that canGet() works
		TokenBucket tb(10,2,50);
		assert(!tb.canGet(0,1));
		tb.update(2);
		assert(tb.canGet(2,10));
	}
	{
		//verify that putTokens() works
		TokenBucket tb(10,2,50);
		assert(tb.tokensLeft()==0);
		tb.putTokens(40);
		assert(tb.tokensLeft()==40);
		tb.putTokens(40);
		assert(tb.tokensLeft()==50);
	}
	{
		//verify that whenAvail() works
		TokenBucket tb(10,2,50);
		assert(tb.tokensLeft()==0);
		assert(tb.whenAvail(0,10)==2);
		tb.putTokens(5);
		assert(tb.whenAvail(0,10)==1);
	}
	{
		//verify that clear() works
		TokenBucket tb(10,2,50);
		tb.update(2);
		assert(tb.tokensLeft()==10);
		tb.clear();
		assert(tb.tokensLeft()==0);
	}
	{
		//verify that setTokens() works
		TokenBucket tb(10,2,50);
		assert(tb.tokensLeft()==0);
		tb.setTokens(40);
		assert(tb.tokensLeft()==40);
		tb.setTokens(60);
		assert(tb.tokensLeft()==50);
	}
	{
		//verify that comparison functions work
		TokenBucket tb0(10,2,50);
		TokenBucket tb1(10,2,50);
		TokenBucket tb2(9,2,50);
		assert(tb0==tb1);
		assert(tb0!=tb2);
		tb0.update(2);
		assert(tb0!=tb1);
		tb1.update(2);
		assert(tb0==tb1);
	}
	return 0;
}
#endif

#ifdef SPEEDTEST
int main(void) {
	TokenBucket tb0(10,2,50);
	TokenBucket tb1(20,2,50);
	TokenBucket tb2(30,2,50);
	TokenBucket tb3(40,2,50);
	TokenBucket tb4(50,5,50);
	TokenBucket tb5(60,5,50);
	TokenBucket tb6(70,5,50);
	TokenBucket tb7(80,8,50);
	TokenBucket tb8(90,8,50);
	uint64_t x=0;
	for(int r=0; r<10000000; r++) {
		x+=tb0.getTokens(r,20);
		x+=tb1.getTokens(r,20);
		x+=tb2.getTokens(r,20);
		x+=tb3.getTokens(r,20);
		x+=tb4.getTokens(r,20);
		x+=tb5.getTokens(r,20);
		x+=tb6.getTokens(r,20);
		x+=tb7.getTokens(r,20);
		x+=tb8.getTokens(r,20);
	}
	return (int)x;
}
#endif