#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/time.h>
#include <unistd.h>   
#include <sys/sysinfo.h>
#include <asm/bitops.h>

int SEARCH;

#define spin_lock(x)
#define spin_unlock(x)
#define read_lock(x)
#define read_unlock(x)
#define unlikely(x)	(x)
#define SET_LINKS(p) do {			\
	(p)->next_task = &init_task;		\
	(p)->prev_task = init_task.prev_task;	\
	init_task.prev_task->next_task = (p);	\
	init_task.prev_task = (p);		\
} while (0)
#define REMOVE_LINKS(p) do {				\
	(p)->next_task->prev_task = (p)->prev_task;	\
	(p)->prev_task->next_task = (p)->next_task;	\
} while (0)
#define CLONE_PID	0x00001000	/* set if pid shared */
#define PID_MAX		0x8000
#define THRES_VALUE     threshold
#define MAX_TASKS       (PID_MAX-1)
#define RESERVED_PIDS   reserved

#define for_each_task(p) \
	for (p = &init_task ; (p = p->next_task) != &init_task ; )

struct task_struct {
	struct task_struct *next_task, *prev_task;
	pid_t pid;
	pid_t pgrp;
	pid_t session;
	pid_t tgid;
	int did_exec;
};

struct task_struct init_task = {
	next_task: 	&init_task,
	prev_task: 	&init_task,
	pid:		0,
	pgrp:		0,
	session:	0,
	tgid:		0,
	did_exec:	0,
};
struct task_struct *current = &init_task;
static pid_t last_pid = 0;
static int next_safe = PID_MAX;
int nr_threads;
int threshold = 0;
int reserved = 1; /* need to keep 0 */
int verbose;

#define BITS_PER_LONG  32  /* no in the kernel !! */

#ifndef SHIFT_PER_LONG
#if BITS_PER_LONG == 64
#   define SHIFT_PER_LONG 6
#elif BITS_PER_LONG == 32
#   define SHIFT_PER_LONG 5
#else
#error "SHIFT_PER_LONG"
#endif
#endif

#define PID_MAP_SIZE (PID_MAX >> SHIFT_PER_LONG)
static unsigned long pid_map[PID_MAP_SIZE];

static inline void mark_pid(int pid)
{
	pid_map[(pid) >> SHIFT_PER_LONG] |= (1 << ((pid) & (BITS_PER_LONG-1)));
}

static int get_pid_by_map(int last_pid) 
{
        struct task_struct *p;
        int i, end;
	int again = 1;
	unsigned long mask, fpos;

#if 0
	/* testing version <algo-1> ... won't pay off */
	memset(pid_map, 0, PID_MAP_SIZE * sizeof(unsigned long));
	again = 0;
	last_pid = RESERVED_PIDS;
#endif

repeat:
	for_each_task(p) {
		mark_pid(p->pid);
		mark_pid(p->pgrp);
		mark_pid(p->tgid);
		mark_pid(p->session);
	}

	/* find next free pid */
	i = last_pid >> SHIFT_PER_LONG;
	mask = pid_map[i] | (((last_pid & (BITS_PER_LONG-1)) << 1) - 1);
	
	while ((mask == ~0) && (++i < PID_MAP_SIZE)) 
		mask = pid_map[i];
	
	if (i == PID_MAP_SIZE) { 
		if (again) {
			/* we didn't find any pid , sweep and try again */
			again = 0;
			memset(pid_map, 0, PID_MAP_SIZE * sizeof(unsigned long));
			last_pid = RESERVED_PIDS;
			goto repeat;
		}
		next_safe = RESERVED_PIDS;
		return 0; 
	}

	fpos = ffz(mask);
	i &= (PID_MAX-1);
	last_pid = (i << SHIFT_PER_LONG) + fpos;

	/* find next save pid */
	mask &= ~((1 << fpos) - 1);

#if 0
	/* testing to see whether interval max size brings 
	 * benefits
	 */
	end = last_pid + 8;
	if (end > PID_MAP_SIZE) 
#endif
		end = PID_MAP_SIZE;
	while ((mask == 0) && (++i < end)) 
		mask = pid_map[i];

	if (i==PID_MAP_SIZE) 
		next_safe = PID_MAX;
	else 
		next_safe = (i << SHIFT_PER_LONG) + ffs(mask) - 1;
	return last_pid;
}

static int get_pid(unsigned long flags)
{
	struct task_struct *p;
	int pid, beginpid;

	if (flags & CLONE_PID)
		return current->pid;

	spin_lock(&lastpid_lock);
	beginpid = last_pid;
	if((++last_pid) & 0xffff8000) {
		beginpid = last_pid = RESERVED_PIDS;		/* Skip daemons etc. */
		goto inside;
	}
	if(last_pid >= next_safe) {
inside:
		next_safe = PID_MAX;
		read_lock(&tasklist_lock);
		SEARCH++;
		if (nr_threads > THRES_VALUE) {
			last_pid = get_pid_by_map(last_pid);
		} else {
			repeat:
				for_each_task(p) {
					if(p->pid == last_pid	||
					   p->pgrp == last_pid	||
					   p->tgid == last_pid	||
					   p->session == last_pid) {
						if(++last_pid >= next_safe) {
							if(last_pid & 0xffff8000)
								last_pid = RESERVED_PIDS;
							next_safe = PID_MAX;
						}
						if(unlikely(last_pid == beginpid))
							goto nomorepids;

						goto repeat;
					}
					if(p->pid > last_pid && next_safe > p->pid)
						next_safe = p->pid;
					if(p->pgrp > last_pid && next_safe > p->pgrp)
						next_safe = p->pgrp;
					if(p->tgid > last_pid && next_safe > p->tgid)
						next_safe = p->tgid;
					if(p->session > last_pid && next_safe > p->session)
						next_safe = p->session;
				}
		}	
		read_unlock(&tasklist_lock);
	}
	pid = last_pid;
	spin_unlock(&lastpid_lock);
	return pid;

nomorepids:
	read_unlock(&tasklist_lock);
	spin_unlock(&lastpid_lock);
	return 0;
}


static void fatal(char *msg,...)
{
	char buf[256];
	va_list ap;

	va_start(ap,msg);
	vsprintf(buf,msg,ap);
	va_end(ap);
	fprintf(stderr,"%s: %s\n", buf, strerror(errno));
	exit(1);
}


static struct task_struct  tasks[PID_MAX];
static struct task_struct *tasks_by_pid[PID_MAX];
static struct task_struct *tasks_free;

static struct task_struct *find_task_by_pid(pid_t pid)
{
	if (pid & ~(PID_MAX-1)) {
		fatal("find_task_by_pid");
		return NULL;
	}
	return tasks_by_pid[pid];
}

static struct task_struct *add_task(pid_t pid, pid_t pgrp, pid_t session,
				    pid_t tgid)
{
	static int initialized = 0;
	struct task_struct *tsk;

	if (!initialized) {
		int i;
		for ( i=0 ; i<PID_MAX ; i++) {
			tasks[i].next_task = tasks_free;
			tasks_free = &tasks[i];
		}
		initialized = 1;
	}
		
	tsk = tasks_free;
	
	if (tsk == NULL)
		fatal("add_task:  pid=%d:%d",pid,nr_threads);

	tasks_free = tasks_free->next_task;
	nr_threads++;
	tasks_by_pid[pid] = tsk;

	tsk->pid = pid;
	tsk->pgrp = pgrp;
	tsk->session = session;
	tsk->tgid = tgid;
	SET_LINKS(tsk);

	return tsk;
}

static void del_task(struct task_struct *tsk)
{
	nr_threads--;
	tasks_by_pid[tsk->pid] = NULL;
	REMOVE_LINKS(tsk);
	tsk->next_task = tasks_free;
	tasks_free = tsk;
}


void populate_all(int count, int fast)
{
	int i,pid;
	int oldthreshold = 0;
	int mode;

	if (count == 0) {
		mode = 0;
		count = MAX_TASKS;
	} else {
		mode = 1;
	}
      
	if (fast) { 
		/* use the new map version */
		oldthreshold = threshold; 
		threshold = 0; 
	}

	/* first allocate all PIDs available */
	for (i = 0; i < count; i++) {
		pid = get_pid(0);
		if (pid == 0) {
			if (mode) {
				printf("failed %d\n",i);
				continue;
			} else 
				break;
		}
		add_task(pid, pid, pid, pid);
		if ((verbose >= 3) && ((i % 1000) == 0)) 
			printf("add task %d  %d %d\n",i,pid, nr_threads);
	}
	if (fast) { 
		/* return to global setting */
		threshold = oldthreshold; 
	}
}

int delete_range(int spid, int epid)
{
	int i;
	struct task_struct *tsk;
	
	for(i=spid; i<=epid; i++) {
		if ((tsk = find_task_by_pid(i))) 
			del_task(tsk);
	}
	return epid-spid+1;
}

int delete_random(int cnt)
{
	int delcnt = 0;
	int coll=0;
	while ((delcnt < cnt) && (nr_threads > RESERVED_PIDS)) {
		struct task_struct *tsk;
		int rnums = random() & (PID_MAX-1);
		if ((rnums < RESERVED_PIDS) || ((tsk = find_task_by_pid(rnums)) == NULL))
		{
			coll++;
			continue;
		}
		del_task(tsk); 
		delcnt++; 
		if ((verbose >= 3) && ((delcnt % 1000) == 0))
			printf("\t del %d\n",delcnt);
	}
	return delcnt;
}

static inline void difftime(struct timeval *dtime, struct timeval *stime, struct timeval *etime)
{
	dtime->tv_sec = etime->tv_sec - stime->tv_sec;
	if (etime->tv_usec < stime->tv_usec) {
		dtime->tv_usec = 1000000 - stime->tv_usec + etime->tv_usec;
		dtime->tv_sec--;
	} else
		dtime->tv_usec = etime->tv_usec - stime->tv_usec;
}


static inline void addtime(struct timeval *rtime, struct timeval *dtime)
{
	rtime->tv_sec += dtime->tv_sec;
	if ((rtime->tv_usec += dtime->tv_usec) >= 1000000) {
		rtime->tv_usec -= 1000000;
		rtime->tv_sec++;
	}
}

int runtest1(int rnums[], int cnt, struct timeval *response)
{
	int i,delcnt;
	struct timeval stime, etime, ttime;
	struct task_struct *tsk;

	/* first delete */
	delcnt = 0;
	for ( i=0 ; i<cnt ; i++ ) {
		if (rnums[i] < RESERVED_PIDS) continue;
		tsk = find_task_by_pid(rnums[i]);
		if (tsk) { del_task(tsk); delcnt++; }
	}

	gettimeofday(&stime, (struct timezone *) 0); 
	for ( i=0 ; i<delcnt ; i++ ) {
		int pid = get_pid(0);
		if (pid == 0) {
			printf("failed %d\n",i);
			continue;
		}
		add_task(pid, pid, pid, pid);
	}
	gettimeofday(&etime, (struct timezone *) 0); 
	difftime(&ttime,&etime,&stime);
	addtime(response,&ttime);
	return delcnt;
}

void runtest2(int start, int end, int incr, int dels)
{
	int i,j;
	struct timeval ttime;
	int cnt;
	int aveth;

	/* first create some randomness:
	 *   a few loops delete random and refill some 
	 */	
#if 1
	populate_all(0,1);
	for (i=0;i<2;i++) {
		int n;
		if (verbose >= 2) printf("P%d\n",i);
		n = delete_random(PID_MAX/4);
		populate_all(n-1000,1);
	}
#endif	
	
	/* now delete until we get to the aveth */
	for (aveth = start ; aveth < end ; aveth += incr) {
		struct timeval response = {0,0};
		int tsearch = 0;
		int talloc  = 0;
		double avtime;

		populate_all(0,1); /* fill all remaining */
		cnt = delete_random(PID_MAX-aveth);

		for ( j=0 ; j< 10 ; j++) {
			struct timeval stime, etime;
			int delcnt = delete_random(dels);
			
			SEARCH = 0;
			gettimeofday(&stime, (struct timezone *) 0); 
			for ( i=0 ; i<delcnt ; i++ ) {
				int pid = get_pid(0);
				if (pid == 0) {
					printf("failed %d\n",i);
					continue;
				}
				add_task(pid, pid, pid, pid);
				talloc++;
			}
			gettimeofday(&etime, (struct timezone *) 0); 
			difftime(&ttime,&stime,&etime);
			addtime(&response,&ttime);
			tsearch += SEARCH;
			if (verbose >= 1) 
				printf("next_phase %d %5d %5d %5d  %ld:%06ld\n",
				       j,SEARCH,delcnt, nr_threads,
				       ttime.tv_sec, ttime.tv_usec);
		}
		avtime = ((1.0E+6*response.tv_sec) + 
			   (double)response.tv_usec) / (double)talloc;
			       
		printf("%5d %5d %5d %5d %d %ld:%06ld %10.2lf\n",
		       aveth,dels,talloc,tsearch,(threshold>0),
		       response.tv_sec,response.tv_usec, avtime
		       );
	}
}

int main(int argc, char *argv[])
{
	int i,cnt,c;
	pid_t pid=0;
	pid_t pgrp=0;
	struct task_struct *tsk = NULL;
	int si,sv;
	int testcase = 2;
	struct timeval response,stime,etime,ttime;
	int startth, endth, incth, dels;

	startth = endth = incth = dels = 1000;
	
        while ((c = getopt(argc, argv, "s:e:i:d:t:c:v:r:")) != -1) {
		switch (c) {
		case 's': startth   = atoi(optarg); break;
		case 'e': endth     = atoi(optarg); break;
		case 'i': incth     = atoi(optarg); break;
		case 'd': dels      = atoi(optarg); break;
		case 't': threshold = atoi(optarg); break;
		case 'c': testcase  = atoi(optarg); break;
		case 'v': verbose   = atoi(optarg); break;
		case 'r': reserved  = atoi(optarg); break;
		}
	}
	if (threshold) threshold = (1<<20);
	if ((reserved < 1) || (reserved > 1000)) 
		fatal("bad reserved %d\n",reserved);

	response.tv_sec = response.tv_usec = 0;
	switch (testcase) {
	case 0:
		populate_all(MAX_TASKS,1);
		cnt = 0;
		cnt += delete_range(300,337);
		cnt += delete_range(1037,1097);
		cnt += delete_range(32723,32767);
		
		/* fork() */
		printf("allocate more %d\n",cnt);
		si = -1; sv=-1;
		for (i=0; i<cnt; i++) {
			pid = get_pid(0);
			if (pid == 0) {
				printf("failed %d\n",i);
				continue;
			}
			if (si == -1) {
				si = pid;
				sv = pid;
			} else {
				if (pid != sv+1) {
					printf("%d - %d\n",si,sv);
					si = pid;
				}
				sv = pid;
			}
			add_task(pid, pid, pid, pid);
		}
		printf("%d - %d\n",si,sv);
		break;

	case 1:
		threshold = 0;
		while (1) {
			const int NUM=100;
			int rnums[NUM];
			int num;
			for (i=0;i<NUM;i++)
				rnums[i] = random() & (PID_MAX-1);
			
			num = runtest1(rnums,NUM,&response);
			printf("%5d %ld:%06ld\n",num,
			       response.tv_sec,response.tv_usec);
		}
		break;

	case 2:
		SEARCH=0;
		runtest2(startth,endth,incth,dels);
		break;

	case 3:
		gettimeofday(&stime, (struct timezone *) 0); 
		pid = get_pid(0);
		add_task(pid, pid, pid, pid);
		printf("new pid: %d, 300: pid %d, pgrp %d\n",
		       pid, tsk->pid, tsk->pgrp);
		/* exit() */
		tsk = find_task_by_pid(300);
		del_task(tsk);
		/* fork() */
		pid = get_pid(0);
		add_task(pid, pgrp, pid, pid);
		printf("new pid: %d, 300: pid %d, pgrp %d\n",
		       pid, tsk->pid, tsk->pgrp);
		/* exit() */
		tsk = find_task_by_pid(301);
		del_task(tsk);
		
		/* fork() */
		pid = get_pid(0);
		add_task(pid, pgrp, pid, pid);
		printf("new pid: %d, 300: pid %d, pgrp %d\n",
		       pid, tsk->pid, tsk->pgrp);
		
		tsk = find_task_by_pid(300);
		gettimeofday(&etime, (struct timezone *) 0); 
		difftime(&ttime,&stime,&etime);
		addtime(&response,&ttime);
		printf("new pid: %d, 300: pid %d, pgrp %d, time:%ld:%ld sec\n",
		       pid, tsk->pid, tsk->pgrp, response.tv_sec,response.tv_usec);
		break;
		
	default: 
		printf("Unknown testcase %d\n",testcase);
		break;
	}
	return 0;
}