#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/shm.h>
#include <sys/sem.h>

void genarray ( int *A, int N )
{
   int i;

   for (i=0; i<N; ++i) A[i] = rand() % 49 - 24;
}

void prnarray ( int *A, int N )
{
   int i;

   for (i=0; i<N; ++i) {
      if (i % 15 == 0) printf("\n    ");
      printf("%5d", A[i]);
   }
   printf("\n");
}

void prefixsum1 ( int *A, int N, int L, int l )
{
   int i, N1, N2;

   if (l == L) {
      for (i=1; i<N; ++i) A[i] += A[i-1];
      exit(0);
   }
   N1 = N / 2; N2 = (N+1) / 2;
   if (!fork()) prefixsum1(A,N1,L,l+1);
   if (!fork()) prefixsum1(A+N1,N2,L,l+1);
   wait(NULL); wait(NULL);
   for (i=N1; i<N; ++i) A[i] += A[N1-1];
   if (l > 0) exit(0);
}

void P ( int semid, int i )
{
   struct sembuf pop;

   pop.sem_num = i;
   pop.sem_flg = 0;
   pop.sem_op = -1 ;
   semop(semid, &pop, 1);
}

void V ( int semid, int i )
{
   struct sembuf vop;

   vop.sem_num = i;
   vop.sem_flg = 0;
   vop.sem_op = 1 ; 
   semop(semid, &vop, 1);
}

void prefixsum2 ( int *A, int N, int L, int l, int semidL, int serno, int semidLP, int semidR )
{
   int i, N1, N2;

   if (l == L) {
      for (i=1; i<N; ++i) A[i] += A[i-1];
      V(semidLP,serno/2);
      P(semidL,serno);
      printf("+++ Leaf process %d: Final array segment", serno);
      prnarray(A,N);
      V(semidR,0);
      exit(0);
   }
   N1 = N / 2; N2 = (N+1) / 2;
   if (!fork()) prefixsum2(A,N1,L,l+1,semidL,2*serno,semidLP,semidR);
   if (!fork()) prefixsum2(A+N1,N2,L,l+1,semidL,2*serno+1,semidLP,semidR);
   if (l == L-1) {
      P(semidLP,serno);
      P(semidLP,serno);
   } else {
      wait(NULL); wait(NULL);
   }
   for (i=N1; i<N; ++i) A[i] += A[N1-1];
   if (l > 0) exit(0);
}

int main ( int argc, char *argv[] )
{
   int N, L;
   int shmid, *A, *B;
   int semidL, semidLP, semidR;
   int i;

   srand((unsigned int)time(NULL));

   if (argc >= 3) {
      N = atoi(argv[1]);
      L = atoi(argv[2]);
   } else {
      scanf("%d%d", &N, &L);
   }
   if ((L <= 0) || (L >= 6)) L = 3;
   if (N < (1 << L)) N = (2 << L);
   printf("+++ N = %d\n+++ L = %d\n", N, L);

   shmid = shmget(IPC_PRIVATE, N * sizeof(int), 0600 | IPC_CREAT);
   A = (int *)shmat(shmid, NULL, 0);
   genarray(A,N);
   B = (int *)malloc(N * sizeof(int));
   for (i=0; i<N; ++i) B[i] = A[i];
   printf("\n+++ Initial array"); prnarray(A,N);

   prefixsum1(A,N,L,0);

   printf("\n*** PART 1\n+++ Root process: Final array"); prnarray(A,N);

   for (i=0; i<N; ++i) A[i] = B[i];

   /* Set of semaphores for the leaf processes to wait on */
   semidL = semget(IPC_PRIVATE, (1 << L), 0600 | IPC_CREAT);
   for (i=0; i<(1<<L); ++i) semctl(semidL, 0, SETVAL, 0);

   /* Set of semaphores for the parents of leaf processes to wait on */
   semidLP = semget(IPC_PRIVATE, (1 << (L-1)), 0600 | IPC_CREAT);
   for (i=0; i<(1<<(L-1)); ++i) semctl(semidLP, 0, SETVAL, 0);

   /* Semaphore for the root to wait on */
   semidR = semget(IPC_PRIVATE, 1, 0600 | IPC_CREAT);

   printf("\n*** PART 2\n");
   prefixsum2(A,N,L,0,semidL,0,semidLP,semidR);

   for (i=0; i<(1 <<L); ++i) {
      V(semidL,i);
      P(semidR,0);
   }

   free(B);
   shmdt(A);
   shmctl(shmid, IPC_RMID, NULL);
   semctl(semidL, -1, IPC_RMID, 0);
   semctl(semidLP, -1, IPC_RMID, 0);
   semctl(semidR, -1, IPC_RMID, 0);
   exit(0);
}