-- S.MG, 2018

package body SMG_Bit_Keccak is

  -- public function, sponge
  procedure Sponge( Input      : in Bitstream;
                    Block_Len  : in Keccak_Rate;
                    Output     : out Bitstream) is
    Internal  : State := (others => (others => (others => 0)));
  begin
    --absorb input into sponge in a loop on available blocks, including padding
    declare
      -- number of input blocks after padding (between 2 and block_len bits pad)
      Padded_Blocks : constant Positive := 1 + (Input'Length + 1) / Block_Len;
      Padded        : Bitstream ( 1 .. Padded_Blocks * Block_Len );
      Block         : Bitstream ( 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;
      Padded( Padded'First + Input'Length )                     := 1;
      Padded( Padded'Last )                                     := 1;

      -- loop through padded input and absorb block by block into sponge
      -- padded input IS a multiple of blocks, so no stray bits 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 bits from sponge in a loop as needed
    declare
      -- full blocks per output
      BPO     : constant Natural := Output'Length / Block_Len;
      -- stray bits per output
      SPO     : constant Natural := Output'Length mod Block_Len;
      Block   : Bitstream( 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 bits)
      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 bits)

    end; -- end squeeze stage

  end Sponge;

  -- helper procedures for sponge absorb/squeeze

  -- NO scramble here, this will absorb ALL given block, make sure it fits!
  procedure AbsorbBlock( Block: in Bitstream; S: in out State ) is
    X, Y                  : XYCoord;
    Z                     : ZCoord;
  begin
    -- xor current block, bit by bit, into first Block'Length bits of state
    First_Pos( X, Y, Z);
    for B of Block loop
      -- xor this bit into the state
      S( X, Y )( Z ) := S( X, Y )( Z ) + B;
      -- move to next bit of the state
      Next_Pos( X, Y, Z );
    end loop;
  end AbsorbBlock;

  -- NO scramble here, this will squeeze Block'Length bits out of *same* state S
  procedure SqueezeBlock( Block: out Bitstream; S: in State) is
    X, Y    : XYCoord;
    Z       : ZCoord;
  begin
    -- start with first position of the state
    First_Pos( X, Y, Z );
    -- squeeze bit by bit, as many bits as needed to fill Block
    for I in Block'Range loop
      -- squeeze current bit from state
      Block( I ) := S( X, Y )( Z );
      -- advance to next bit of state
      Next_Pos( X, Y, Z);
    end loop;
  end SqueezeBlock;

  -- moving one bit forwards in Keccak state
  procedure Next_Pos( X : in out XYCoord;
                      Y : in out XYCoord;
                      Z : in out ZCoord
                    ) is
  begin
    Z := Z - 1;
    if Z = ZCoord'Last then
      X := X + 1;
      if X = XYCoord'First then
        Y := Y + 1;
      end if;
    end if;
  end Next_Pos;

  -- position of first bit in Keccak state
  procedure First_Pos( X : out XYCoord;
                       Y : out XYCoord;
                       Z : out ZCoord
                     ) is
  begin
    X := XYCoord'First;
    Y := XYCoord'First;
    Z := ZCoord'Last;
  end First_Pos;

  -- operations with Bitwords
  function BWRotate_Left( Input: in Bitword;
                          Count: in Natural)
                          return Bitword is
    Output  : Bitword;
    Advance : constant ZCoord := ZCoord( Count mod Z_Length );
  begin
    for I in ZCoord loop
      Output( I ) := Input( I + Advance );
    end loop;
    return Output;
  end BWRotate_Left;

  -- Keccak transformations of the internal state
  function Theta ( Input       : in State) return State is
    Output : State;
    S1, S2 : Bit;
  begin
    for X in XYCoord loop
      for Y in XYCoord loop
        for Z in ZCoord loop
          S1 := 0;
          S2 := 0;
          for Y1 in XYCoord loop
            S1 := S1 + Input( X - 1, Y1 )( Z );
            -- Z direction is opposite to the one assumed in the ref so Z + 1 
            S2 := S2 + Input( X + 1, Y1 )( Z + 1 );
          end loop;
          Output( X, Y )(Z) := Input( X, Y )( Z ) + S1 + S2;
        end loop;
      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) := BWRotate_Left(Input(X,Y), (T+1)*(T+2)/2);
      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
        for Z in ZCoord loop
          Output(X, Y)(Z) :=   Input( X, Y )( Z ) +
                             ( Input( X + 1, Y )( Z ) + 1 ) *
                             ( Input( X + 2, Y )( Z )     );
        end loop;
      end loop;
    end loop;

    return Output;
  end Chi;

  function Iota  ( Round_Const : in Bitword; Input : in State) return State is
    Output : State;
  begin
    Output := Input;
    for Z in ZCoord loop
      Output( 0, 0 )(Z) := Input( 0, 0 )( Z ) + Round_Const( Z );
    end loop;
    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_Bit_Keccak;