-- S.MG, 2018
with System; use System;  -- for Bit_Order
with Interfaces; use Interfaces;

package body SMG_Keccak is

-- public function, sponge
  procedure Sponge( Input      : in Bytestream;
                    Output     : out Bytestream;
                    Block_Len  : in Keccak_Rate := Default_Byterate ) is
    Internal  : State := (others => (others => 0));
  begin
    --absorb input into sponge in a loop on available blocks, including padding
    declare
      -- number of input blocks after padding (pad between 1 and block_len)
      Padded_Blocks : constant Positive := 1 + Input'Length / Block_Len;
      Padded        : Bytestream ( 1 .. Padded_Blocks * Block_Len );
      Block         : Bytestream ( 1 .. Block_Len );
    begin
      -- initialise Padded with 0 everywhere
      Padded := ( others => 0 );
      -- copy and pad input with rule 10*1
      Padded( Padded'First .. Padded'First + Input'Length - 1 ) := Input;
      -- padding is 10*1 so start and end with an 1 but LSB order hence 16#80#
      Padded( Padded'First + Input'Length )                     := 1;
      Padded( Padded'Last ) := Padded( Padded'Last ) + 16#80#;

      -- loop through padded input and absorb block by block into sponge
      -- padded input IS a multiple of blocks, so no stray octets left
      for B in 0 .. Padded_Blocks - 1 loop
        -- first get the current block to absorb
        Block   := Padded( Padded'First + B * Block_Len .. 
                           Padded'First + (B+1) * Block_Len - 1 );
        AbsorbBlock( Block, Internal );
        -- scramble state with Keccak function
        Internal := Keccak_Function( Internal );

      end loop; -- end absorb loop for blocks
    end; -- end absorb stage

    --squeeze required octets from sponge in a loop as needed
    declare
      -- full blocks per output
      BPO     : constant Natural := Output'Length / Block_Len;
      -- stray octets per output
      SPO     : constant Natural := Output'Length mod Block_Len;
      Block   : Bytestream( 1 .. Block_Len );
    begin
      -- squeeze block by block (if at least one full block is needed)
      for I in 0 .. BPO - 1 loop
        SqueezeBlock( Block, Internal );
        Output( Output'First + I * Block_Len .. 
                Output'First + (I + 1) * Block_Len -1) := Block;
 
        -- scramble state
        Internal := Keccak_Function( Internal );
      end loop;  -- end squeezing full blocks

      -- squeeze any partial block needed (stray octets)
      if SPO > 0 then
        SqueezeBlock( Block, Internal );
        Output( Output'Last - SPO + 1 .. Output'Last ) := 
                Block( Block'First .. Block'First + SPO - 1 );
      end if; -- end squeezing partial last block (stray octets)

    end; -- end squeeze stage
  end Sponge;

  -- convert from a bytestream of ZWord/8 size to an actual ZWord number
  -- NB: this will FLIP bits on big endian because keccak expects input LSB
  -- NOT exact opposite of WordToBytes
  function BytesToWordLE( BWord: in Byteword ) return ZWord is
    W    : ZWord;
    B    : Byteword;
  begin
    -- just copy octets if machine is little endian
    -- flip octets AND bits if machine is big endian
    if Default_Bit_Order = Low_Order_First then
      B := BWord;
    else
      B := FlipOctets( BWord );
      for I in B'First..B'Last loop
        B(I) := Reverse_Table(Natural(B(I)));
      end loop;
    end if;
    -- actual bytes to word conversion
    W := Cast(B);
    return W;
  end BytesToWordLE;

  -- convert from a ZWord (lane of state) to a bytestream of ZWord size
  -- NOT exact oppositve of BytesToWordLE
  -- Keccak sponge spits out MSB so bits are flipped on LITTLE Endian iron.
  function WordToBytesBE( Word: in ZWord ) return Byteword is
    B: Byteword;
  begin
    B := Cast( Word );

    -- flip octets if machine is big endian
    if Default_Bit_Order = High_Order_First then
      B := FlipOctets( B );
    else
    -- onth flip bits if machine is little endian....
      for I in B'First..B'Last loop
        B(I) := Reverse_Table(Natural(B(I)));
      end loop;
    end if;

    return B;
  end WordToBytesBE;

  -- flip given octets (i.e. groups of 8 bits)
  function FlipOctets( BWord : in Byteword ) return Byteword is
    B : Byteword;
  begin
    -- copy octets changing their order in the array 
    -- i.e. 1st octet in BWord becomes last octet in B and so on
    for I in 0 .. BWord'Length-1 loop
      B(B'First + I) := BWord(BWord'Last-I);
    end loop;
    return B;
  end FlipOctets;

-- helper procedures for sponge absorb/squeeze

  -- NO scramble here, this will absorb ALL given block, make sure it fits!
  procedure AbsorbBlock( Block: in Bytestream; S: in out State ) is
    WPB: constant Natural := Block'Length / Byteword'Length; -- words per block
    SBB: constant Natural := Block'Length mod Byteword'Length; -- stray octets
    FromPos, ToPos        : Natural;
    X, Y                  : XYCoord;
    Word                  : ZWord;
    BWord                 : Byteword;
  begin
    -- xor current block into first Block'Length octets of state
    -- a block can consist in more than one word
    X := 0;
    Y := 0;
    for I in 0..WPB-1 loop
      FromPos := Block'First + I * Byteword'Length;
      ToPos   := FromPos + Byteword'Length - 1;
      Word := BytesToWordLE( Block( FromPos .. ToPos ) );
      S( X, Y ) := S( X, Y ) xor Word;
      -- move on to next word in state
      X := X + 1;
      if X = 0 then
        Y := Y + 1;
      end if;
    end loop;
    -- absorb also any remaining bytes from block
    if SBB > 0 then
      ToPos := Block'Last;
      FromPos := ToPos - SBB + 1;
      BWord := (others => 0);
      BWord(Byteword'First .. Byteword'First + SBB - 1) := Block(FromPos..ToPos);
      Word := BytesToWordLE( BWord );
      S( X, Y ) := S( X, Y ) xor Word;
    end if;
  end AbsorbBlock;

  -- NO scramble here, this will squeeze Block'Length bits out of *same* state S
  procedure SqueezeBlock( Block: out Bytestream; S: in State) is
    X, Y    : XYCoord;
    BWord   : Byteword;
    FromPos : Natural;
    Len     : Natural;
  begin
    X := 0;
    Y := 0;
    FromPos := Block'First;

    while FromPos <= Block'Last loop
      BWord := WordToBytesBE( S(X, Y) );

      X := X + 1;
      if X = 0 then
        Y := Y + 1;
      end if;

      -- copy full word if it fits or
      --   only as many bytes as are still needed to fill the block
      Len := Block'Last - FromPos + 1;
      if Len > BWord'Length then
        Len := BWord'Length;
      end if;

      Block(FromPos..FromPos+Len-1) := BWord(BWord'First..BWord'First+Len-1);
      FromPos := FromPos + Len;
    end loop;
  end SqueezeBlock;


-- private, internal transformations
  function Theta(Input : in State) return State is
    Output : State;
    C      : Plane;
    W      : ZWord;
  begin
    for X in XYCoord loop
      C(X) := Input(X, 0);
      for Y in 1..XYCoord'Last loop
        C(X) := C(X) xor Input(X, Y);
      end loop;
    end loop;

    for X in XYCoord loop
      W := C(X-1) xor Rotate_Left(C(X+1), 1);
      for Y in XYCoord loop
        Output(X,Y) := Input(X,Y) xor W;
      end loop;
    end loop;

    return Output;
  end Theta;

  function Rho(Input : in State) return State is
    Output      : State;
    X, Y, Old_Y : XYCoord;
  begin
    Output(0,0) := Input(0,0);
    X           := 1;
    Y           := 0;

    for T in 0..23 loop
      Output(X, Y) := Rotate_Left(Input(X,Y), ((T+1)*(T+2)/2) mod Z_Length);
      Old_Y := Y;
      Y := 2*X + 3*Y;
      X := Old_Y;
    end loop;
    return Output;
  end rho;

  function Pi(Input : in State) return State is
    Output: State;
  begin
    for X in XYCoord loop
      for Y in XYCoord loop
        Output(Y, 2*X + 3*Y) := Input(X, Y);
      end loop;
    end loop;
    return Output;
  end pi;

  function Chi(Input : in State) return State is
    Output: State;
  begin
    for Y in XYCoord loop
      for X in XYCoord loop
        Output(X, Y) := Input(X, Y) xor 
                        ( (not Input(X + 1, Y)) and Input(X + 2, Y) );
      end loop;
    end loop;
    return Output;
  end chi;

  function Iota(Round_Const : in ZWord; Input : in State) return State is
    Output: State;
  begin
    Output := Input;
    Output(0,0) := Input(0,0) xor Round_Const;
    return Output;
  end iota;

  function Keccak_Function(Input: in State) return State is
    Output: State;
  begin
    Output := Input;
    for I in Round_Index loop
      Output := Iota(RC(I), Chi(Pi(Rho(Theta(Output)))));
    end loop;

    return Output;
  end Keccak_Function;

end SMG_Keccak;