/*
     ntt31.c
*/
/*
** This file is placed into the public domain by its author,
** Carey Bloodworth (Carey@Bloodworth.org) on July 16, 2001
**
** This multiplication demo is not designed for high performance.
** It's a tutorial program designed to be used with the information
** on my web site at www.Bloodworth.org
*/
/*
** This file demonstrates a very basic NTT multiply.  It uses a single
** 31 bit prime.  That means you can only put a single decimal in there
** without things overflowing quickly.  It does, however, mean that
** everybody can experiment with this.
**
** To compile this using GCC:
** gcc main.c ntt31.c -o ntt31.exe
*/
#include 
#include 
#include 
#include 
#include 

#define CalcNTTLen(_NumLen) ((((_NumLen)*BASE_DIG)*2)/BASE_DIG)
/* NumLen*BaseDig*ZeroPadding/Dig_Per_FFT */

typedef short int Short;
typedef signed long    INT32; /* 32/31 bit signed int */
typedef unsigned long UINT32; /* 32 bit unsigned int */
typedef INT32 ModInt;

static ModInt *NTTNum1=NULL, *NTTNum2=NULL;
static int BASE;
static int BASE_DIG;

ModInt Prime,PrimvRoot,MulInv;
static double RecipPrime=0.0;

ModInt
ModMul(ModInt a, ModInt b)
/* Limited to 31 bits. */
{ModInt rem;
rem = a * b;
rem = rem - Prime * ((ModInt) floor(0.5+RecipPrime * ((double) a) * ((double) b)));
if (rem < 0) rem +=Prime;
return rem;
}

ModInt ModAdd(ModInt a,ModInt b)
{UINT32 Sum; /* Big enough to hold sum */
Sum=a+b;
if (Sum >= Prime) Sum-=Prime;
return (ModInt)Sum;
}

ModInt ModSub(ModInt a,ModInt b)
{INT32 Dif; /* Big enough to hold signed difference */
Dif=a-b;
if (Dif < 0) Dif+=Prime;
return (ModInt)Dif;
}


ModInt
ModPow(ModInt Base,ModInt Expon)
{ModInt prod,b;

if (Expon<=0) return 1;

b=Base;
while (!(Expon&1)) {b=ModMul(b,b);Expon>>=1;}
prod=b;

while (Expon>>=1)
  {
   b=ModMul(b,b);
   if (Expon&1) prod=ModMul(prod,b);
  }
return prod;
}

ModInt
FindInverse(ModInt Num)
{ModInt i;
i=ModPow(Num,Prime-2);
/*
** Num*3 can overflow causing the check to fail.
if (ModMul(Num*3,i) != 3)
  FatalError("Unable to find Mul inverse for %u mod %u\n",Num,Modulus);
*/
return i;
}

static void
NTTReorder(ModInt *Data, int Len)
{int Index,xednI,k;

xednI = 0;
for (Index = 0;Index < Len;Index++)
  {
   if (xednI > Index)
     {ModInt Temp;
      Temp=Data[xednI];
      Data[xednI]=Data[Index];
      Data[Index]=Temp;
     }
   k=Len/2;
   while ((k <= xednI) && (k >=1)) {xednI-=k;k/=2;}
   xednI+=k;
  }
}

void NTT(ModInt *Data, int Len, int Dir)
/* A simple minded, generic transform */
{int j,step,halfstep;
 int index,index2;
 ModInt u,w,temp;

NTTReorder(Data,Len);

step=1;
while (step < Len)
  {
   halfstep=step;
   step*=2;

   u=1;
   if (Dir > 0) w=ModPow(PrimvRoot,Prime-1-((Prime-1)/step));
   else         w=ModPow(PrimvRoot,(Prime-1)/step);

   for (j=0;j 0; x--)
    {
      Pyramid = ModMul(NTTNum1[Len2 - x],MulInv) + Carry;
      Carry       = Pyramid / BASE;
      Prod[x - 1] = Pyramid % BASE;
    }
}

void
InitFFT(unsigned long int Len,int Base,int BaseDig)
{int Bytes;

BASE=Base;
BASE_DIG=BaseDig;

Bytes=sizeof(ModInt)*CalcNTTLen(Len);

if ( (BaseDig > 1) || (BASE > 10))
  {
   printf("Error:  The ntt is hardwired for just 1 digit in the base.\n");
   printf("In 'main.c' please set BASE to 10 and BASE_DIG to 1\n");
   exit(0);
  }

NTTNum1=(ModInt*)malloc(Bytes);
NTTNum2=(ModInt*)malloc(Bytes);

if ((NTTNum1==NULL) || (NTTNum2==NULL))
   {
    printf("Unable to allocate memory for NTTNum.\n");
    printf("Len=%d Bytes=%d\n",(int)Len,(int)Bytes);
    exit(0);
   }

  Prime = 2130706433;PrimvRoot=3;
  RecipPrime=1.0/Prime;
}

void DeInitFFT(unsigned long int Len)
{
free(NTTNum1);free(NTTNum2);
}