------------------------------------------------------------------------------
------------------------------------------------------------------------------
-- This file is part of 'Finite Field Arithmetic', aka 'FFA'.               --
--                                                                          --
-- (C) 2019 Stanislav Datskovskiy ( www.loper-os.org )                      --
-- http://wot.deedbot.org/17215D118B7239507FAFED98B98228A001ABFFC7.html     --
--                                                                          --
-- You do not have, nor can you ever acquire the right to use, copy or      --
-- distribute this software ; Should you use this software for any purpose, --
-- or copy and distribute it to anyone or in any manner, you are breaking   --
-- the laws of whatever soi-disant jurisdiction, and you promise to         --
-- continue doing so for the indefinite future. In any case, please         --
-- always : read and understand any software ; verify any PGP signatures    --
-- that you use - for any purpose.                                          --
--                                                                          --
-- See also http://trilema.com/2015/a-new-software-licensing-paradigm .     --
------------------------------------------------------------------------------

with Word_Ops; use Word_Ops;
with W_Pred;   use W_Pred;
with W_Shifts; use W_Shifts;

with FZ_Basic; use FZ_Basic;
with FZ_QShft; use FZ_QShft;
with FZ_Arith; use FZ_Arith;
with FZ_Divis; use FZ_Divis;
with FZ_Pred;  use FZ_Pred;
with FZ_Cmp;   use FZ_Cmp;
with FZ_Barr;  use FZ_Barr;
with FZ_ModEx; use FZ_ModEx;


package body FZ_Prime is
   
   -- Find the highest power of 2 which divides N. ( or 0 if N is zero. )
   function FZ_Count_Bottom_Zeros(N : in FZ) return FZBit_Index is
      
      -- The result (default : 0, will remain 0 if N is in fact zero)
      Index     : Word := 0;
      
      -- The currently-examined Word, starting from the highest
      Ni        : Word;
      
      -- The most recently-seen nonzero Word, if indeed any exist
      W         : Word := 0;
      
      -- 1 if currently-examined Word is zero, otherwise 0
      NiZ       : WBool;
      
   begin
      
      -- Find lowest non-zero Word (or simply stop at lowest, if N = 0)
      for i in reverse 0 .. Indices(N'Length - 1) loop
         Ni    := N(N'First + i);              -- Ni := current Word;
         NiZ   := W_ZeroP(Ni);                 -- ... is this Word = 0?
         Index := W_Mux(Word(i), Index, NiZ);  -- If NO, save its Index,
         W     := W_Mux(Ni,      W,     NiZ);  -- ... and save the Word.
      end loop;
      
      -- Set Index to be the bit-position of the lowest non-zero Word:
      Index := Shift_Left(Index, BitnessLog2); -- Index := Index * Bitness
      
      -- Handle degenerate case: if W is equal to 0, Index is not changed;
      -- Otherwise, start by advancing Index by an entire Word Bitness:
      Index := Index + ((0 - W_NZeroP(W)) and Word(Bitness));
      
      -- Now crank back the Index by the number of high bits of W (i.e.
      -- starting from its top) that must be discarded for W to become zero:
      for b in 1 .. Bitness loop
         
         -- If W is non-zero at this time, decrement the Index:
         Index := Index - W_NZeroP(W);
         
         -- Advance W left, i.e. discard the top bit of it:
         W     := Shift_Left(W, 1);
         
      end loop;
      
      -- If N = 0, result will be 0; otherwise: length of bottom zeros in N.
      return FZBit_Index(Index);
      
   end FZ_Count_Bottom_Zeros;
   
   
   -- Constant-Time Miller-Rabin Test on N using the given Witness.
   -- Witness will be used as-is if it conforms to the valid bounds,
   -- i.e. 2 <= Witness <= N - 2; otherwise will be transformed into a
   -- valid Witness via modular arithmetic.
   -- Outputs ONE if N WAS FOUND composite; ZERO if NOT FOUND.
   -- Handles degenerate cases of N that M-R per se cannot eat:
   -- N=0, N=1: ALWAYS 'FOUND COMPOSITE'; 2, 3 - ALWAYS 'NOT FOUND'.
   -- If N is Even and not equal to 2, N is ALWAYS 'FOUND COMPOSITE.'
   -- For ALL other N, the output is equal to that of the M-R test.
   function FZ_MR_Composite_On_Witness(N       : in  FZ;
                                       Witness : in  FZ) return WBool is
      
      -- Widths of N, Witness, and Result are equal :
      subtype Width is Word_Index range N'Range;
      
      -- Whether N is 'small' (i.e. 1 Word) and therefore possibly degenerate :
      N_Is_Small       : constant WBool := FZ_OneWordP(N);
      
      -- Head of N, for detecting (and handling) the degenerate-N case :
      N_Head           : constant Word  := FZ_Get_Head(N);
      
      -- Zero and One are defined as COMPOSITE degenerate cases of N :
      N_Is_Zero_Or_One : constant WBool
        := N_Is_Small and (W_EqP(N_Head, 0) or W_EqP(N_Head, 1));
      
      -- Two and Three are PRIME degenerate cases of N :
      N_Is_Two         : constant WBool := N_Is_Small and W_EqP(N_Head, 2);
      N_Is_Three       : constant WBool := N_Is_Small and W_EqP(N_Head, 3);
      
      -- The COMPOSITE degenerate case of N where N != 2 and N is Even :
      N_Ne_2_Is_Even   : constant WBool :=
        (1 - N_Is_Two) and (1 - W_OddP(N_Head));
      
      -- Degeneracy latch. If N is Zero or One, then set immediately :
      Degen_Composite  : constant WBool := N_Is_Zero_Or_One or N_Ne_2_Is_Even;
      
      -- Possible-primality latch. If N is Two or Three, then set immediately :
      Possibly_Prime   : WBool := N_Is_Two or N_Is_Three;
      
      -- The writable copy of N that we will operate on :
      X                : FZ(Width) := N;
      
      -- Potentially-replaced head of X (if degenerate N) :
      X_Head           : Word := N_Head;
      
      -- Will be Barrettoid(X), for operations modulo X :
      XBar             : Barretoid(ZXMLength       => X'Length + 1,
                                   BarretoidLength => 2 * X'Length);
      
      -- The Bound Witness, constrained to valid range 2 <= BW <= X - 2 :
      BW               : FZ(Width);
      
      -- Head of BW, for detecting crossing of the lower bound :
      BW_Head          : Word;
      
      -- X - 1 (for M-R proper, and to constrain BW) :
      X_Minus_One      : FZ(Width);
      
      -- X - 1 = 2^R * K, with K odd :
      K                : FZ(Width);
      R                : FZBit_Index;
      
      -- Working register for all M-R modular operations :
      T                : FZ(Width);
      
      -- For subtraction where no underflow can happen, barring cosmic ray:
      NoCarry          : WZeroOrDie := 0;
      
   begin
      
      -- First: X, which will remain equal to N unless N is degenerate:
      
      -- If N is one of the two prohibited small primes (2,3) X will become 5:
      X_Head := W_Mux(A => X_Head, B => 5, Sel => Possibly_Prime);
      
      -- If one of the two prohibited small composites (0,1), or even, then 9:
      X_Head := W_Mux(A => X_Head, B => 9, Sel => Degen_Composite);
      
      -- Given as we're stuck carrying out M-R even if N is degenerate case,
      -- then let the M-R do The Right Thing, for added cosmic ray resistance.
      
      -- X gets a new head, if N was a degenerate case; else keeps old head:
      FZ_Set_Head(X, X_Head);
      
      -- Compute X - 1. As now X > 3, underflow is not permitted:
      FZ_Sub_W(X => X, W => 1, Difference => X_Minus_One,
               Underflow => NoCarry);
      
      -- Now, compute BW (Bound Witness), which satisfies 2 <= BW <= X - 2:
      
      -- First, BW := Witness mod (X - 1). After this, 0 <= BW <= X - 2:
      FZ_Mod(Dividend => Witness, Divisor => X_Minus_One, Remainder => BW);
      
      -- Now, adjust BW if it is found to be below Two (the lower bound) :
      
      -- Get head of BW:
      BW_Head := FZ_Get_Head(BW);
      
      -- If BW is equal to Zero or One, it will be forced to equal Two:
      BW_Head := W_Mux(A   => BW_Head,
                       B   => 2,
                       Sel => FZ_OneWordP(BW) and
                         (W_EqP(BW_Head, 0) or W_EqP(BW_Head, 1)));
      
      -- BW gets the new head, if it must; otherwise keeps its old head:
      FZ_Set_Head(BW, BW_Head);
      
      -- We finished adjusting X and BW for degeneracies, and can now M-R:
      
      -- Generate Barrettoid(X) to use in all of the modulo-X operations:
      FZ_Make_Barrettoid(Modulus => X, Result => XBar);
      
      -- Find R >= 1, and odd K, where X − 1 = 2^R * K :
      
      -- ... first, find R, the largest power of two which divides X - 1 :
      R := FZ_Count_Bottom_Zeros(X_Minus_One);
      
      -- ... and now obtain K := X / 2^R, i.e., K := X >> R :
      FZ_Quiet_ShiftRight(N => X_Minus_One, ShiftedN => K, Count => R);
      
      -- T := BW^K mod X
      FZ_Mod_Exp_Barrett(Base => BW, Exponent => K, Bar => XBar, Result => T);
      
      -- if T = 1 or T = X - 1, then X is possibly-prime:
      Possibly_Prime := Possibly_Prime or
        FZ_EqP_W(T, 1) or FZ_EqP(T, X_Minus_One);
      
      -- Needs R - 1 times, but perform for max possible count (gated):
      for i in 1 .. FZ_Bitness(N) - 1 loop
         
         -- T := T^2 mod X
         FZ_Mod_Sqr_Barrett(X => T, Bar => XBar, Product => T);
         
         -- if T = X - 1, and i < R, then X is possibly-prime:
         Possibly_Prime := Possibly_Prime or 
           (FZ_EqP(T, X_Minus_One) and W_LtP(Word(i), Word(R)));
         
      end loop;
      
      -- The output, which will be 1 iff X WAS FOUND to be composite via BW,
      -- ... or if X was found to equal Zero or One (without regard to BW.)
      return (1 - Possibly_Prime) or Degen_Composite;
      -- If X was found to equal Two or Three, output will be 0 (under any BW).
      
   end FZ_MR_Composite_On_Witness;
   
end FZ_Prime;