-- S.MG implementation of Keccak-f permutations

 -- (Based on The Keccak Reference, Version 3.0, January 14, 2011, by
 --   Guido Bertoni, Joan Daemen, Michael Peeters and Gilles Van Assche)

 -- S.MG, 2018

with Bits; use Bits;
package SMG_Keccak is
	pragma Pure(SMG_Keccak);  --stateless, no side effects -> can cache calls

  --knobs (can change as per keccak design but fixed here for S.MG purposes)--
  Keccak_L: constant := 6;  --gives keccak z (word) dimension of 2^6=64 and 
                            --therefore keccak function 1600 with current 
                            --constants (5*5*2^6)

  Default_Bitrate: constant := 1344; --max bits the sponge can eat/spit without
                                     --needing to scramble the state

  --constants: dimensions of keccak state and number of rounds
  XY_Length: constant := 5;
  Z_Length: constant := 2**Keccak_L;
  Width: constant := XY_Length * XY_Length * Z_Length;
  N_Rounds: constant := 12 + 2*Keccak_L;

  --types
  type XYCoord is mod XY_Length;
  type ZCoord is mod Z_Length;
  type Round_Index is mod N_Rounds;

  type ZWord is mod 2**Z_Length;	--"lane" in keccak ref
  type Plane is array(XYCoord) of ZWord; --a "horizontal slice" of keccak state
  type State is array(XYCoord, XYCoord) of ZWord; --the full keccak state

  type Round_Constants is array(Round_Index) of ZWord;  --magic keccak constants

  -- rate can be chosen by caller at each call, between 1 and width of state
  -- higher rate means sponge "eats" more bits at a time but has fewer bits in 
  --   the "secret" part of the state (i.e. lower capacity)
  subtype Keccak_Rate is Positive range 1..Width;  -- capacity = width - rate

  subtype Bitword is Bitstream( 0..Z_Length - 1 ); -- bits of one state "word"

  -- type conversions
  function BitsToWord( BWord : in Bitword ) return ZWord;
  function WordToBits( Word  : in ZWord ) return Bitword;

  -- flip input octets (i.e. groups of 8 bits)
  function FlipOctets( BWord : in Bitword ) return Bitword;

  -- public function, the sponge itself
  -- Keccak sponge structure using Keccak_Function, Pad and a given bitrate;
  -- Input - the stream of bits to hash (the message)
  -- Output - a bitstream of desired size for holding output
  -- Block_Len - the bitrate to use; this is effectively the block length 
  --             for splitting Input AND squeezing output between scrambles 
  procedure Sponge(Input      : in Bitstream;
                   Output     : out Bitstream;
                   Block_Len  : in Keccak_Rate := Default_Bitrate );

  -- state based Sponge
  type Keccak_Context (Block_Len: Keccak_Rate := Default_Bitrate) is
     record
        Internal: State := (others => (others => 0));
        Block: Bitstream(1..Block_Len) := (others => 0);
        Pos: Natural;
     end record;
 
  procedure KeccakBegin(Ctx: in out Keccak_Context);
  procedure KeccakHash(Ctx: in out Keccak_Context;
                       Input: Bitstream);
  procedure KeccakEnd(Ctx: in out Keccak_Context;
                      Output: out Bitstream);

private
  -- these are internals of the keccak implementation, not meant to be directly 
  --  accessed/used

  -- this will squeeze Block'Length bits out of state S
  -- NO scramble of state in here!
  -- NB: make SURE that Block'Length is the correct bitrate for this sponge
  -- in particular, Block'Length should be a correct bitrate aka LESS than Width 
  procedure SqueezeBlock( Block: out Bitstream; S: in State);

  -- This absorbs into sponge the given block, modifying the state accordingly
  -- NO scramble of state in here so make sure the whole Block fits in state!
  -- NB: make SURE that Block'Length is *the correct bitrate* for this sponge
  -- in particular, Block'Length should be a correct bitrate aka LESS than Width 
  procedure AbsorbBlock( Block: in Bitstream; S: in out State );

  --Keccak magic numbers
  RC : constant Round_Constants :=
    (
     16#0000_0000_0000_0001#,
     16#0000_0000_0000_8082#,
     16#8000_0000_0000_808A#,
     16#8000_0000_8000_8000#,
     16#0000_0000_0000_808B#,
     16#0000_0000_8000_0001#,
     16#8000_0000_8000_8081#,
     16#8000_0000_0000_8009#,
     16#0000_0000_0000_008A#,
     16#0000_0000_0000_0088#,
     16#0000_0000_8000_8009#,
     16#0000_0000_8000_000A#,
     16#0000_0000_8000_808B#,
     16#8000_0000_0000_008B#,
     16#8000_0000_0000_8089#,
     16#8000_0000_0000_8003#,
     16#8000_0000_0000_8002#,
     16#8000_0000_0000_0080#,
     16#0000_0000_0000_800A#,
     16#8000_0000_8000_000A#,
     16#8000_0000_8000_8081#,
     16#8000_0000_0000_8080#,
     16#0000_0000_8000_0001#,
     16#8000_0000_8000_8008#
    );

  --gnat-specific methods to have bit-ops for modular types
  function Rotate_Left( Value  : ZWord; 
	                      Amount : Natural) 
	                      return ZWord;
	pragma Import(Intrinsic, Rotate_Left); 

	function Shift_Right( Value  : ZWord; 
                        Amount : Natural) 
                        return ZWord;
  pragma Import(Intrinsic, Shift_Right); 

  function Shift_Left( Value  : ZWord;
                       Amount : Natural)
                       return ZWord;
  pragma Import(Intrinsic, Shift_Left);

  --Keccak transformations of the internal state
  function Theta ( Input       : in State) return State;
  function Rho   ( Input       : in State) return State;
  function Pi    ( Input       : in State) return State;
  function Chi   ( Input       : in State) return State;
  function Iota  ( Round_Const : in ZWord; Input : in State) return State;

  --Keccak function with block width currently 1600 (Width constant above)
  --this simply applies *all* keccak transformations in the correct order, using
  -- the keccak magic numbers (round constants) as per keccak reference
  function Keccak_Function(Input: in State) return State;

end SMG_Keccak;