------------------------------------------------------------
-- Copyright: 2010 Integrated Sytems Laboratory, ETH Zurich
--            http://www.iis.ee.ethz.ch/~sha3
------------------------------------------------------------
library ieee;
use ieee.numeric_std.all;
use ieee.std_logic_1164.all;
use work.shavitepkg.all;

entity expansion is

  port (
    ClkxCI     : in  std_logic;
    RstxRBI    : in  std_logic;
    InputEnxEI : in  std_logic;
    DataBlkxDI : in  wordmat512;
    BitCntxDI  : in  wordmat64;
    SubKeyxDO  : out wordmat128);

end expansion;

architecture rtl of expansion is

  component aes
    port (
      ClkxCI   : in  std_logic;
      RstxRBI  : in  std_logic;
      PlainxDI : in  wordmat128;
      OutxDO   : out wordmat128);
  end component;

  signal AESInxD, AESOutxD, AESXorxD : wordmat128;
  signal StatexDP, StatexDN          : unsigned(5 downto 0);  --cycle counter:
  -- 0 = input fetch/AES, mod8=4 = XOR
  signal StateMatxDP, StateMatxDN    : wordmat512;
  
begin  -- rtl

  u_aes : aes
    port map (
      ClkxCI   => ClkxCI,
      RstxRBI  => RstxRBI,
      PlainxDI => AESInxD,
      OutxDO   => AESOutxD);


  -- Output
  output : for i in 3 downto 0 generate
    SubKeyxDO(i) <= StateMatxDP(i) when InputEnxEI = '0' else DataBlkxDI(i);
  end generate output;


  -- purpose: Control the expansion function
  -- type   : combinational
  -- inputs : AESOutxD, AESXorxD, BitCntxDI, DataBlkxDI, InputEnxEI, StateMatxDP, StatexDP
  -- outputs: AESInxD, AESXorxD, StateMatxDN, StatexDN
  p_fsm : process (AESOutxD, AESXorxD, BitCntxDI, DataBlkxDI, InputEnxEI,
                   StateMatxDP, StatexDP)
  begin  -- process p_fsm

    -- default assignments
    AESInxD(2) <= StateMatxDP(3);
    AESInxD(1) <= StateMatxDP(2);
    AESInxD(0) <= StateMatxDP(1);
    AESInxD(3) <= StateMatxDP(0);

    for i in 3 downto 0 loop
      AESXorxD(i) <= AESOutxD(i) xor StateMatxDP(12+i);
    end loop;  -- i
    StateMatxDN <= StateMatxDP;
    StatexDN    <= StatexDP+1;


    -- Walk through states 0...35
    StatexDN <= StatexDP+1;
    if to_integer(StatexDP) = 35 or (to_integer(StatexDP) = 0 and InputEnxEI = '0') then
      StatexDN <= (others => '0');
    end if;

    -- main operations: AES round, XOR of previous values
    if to_integer(StatexDP) mod 8 < 4 then  --AES
      for i in 0 to 3 loop
        StateMatxDN(i+12) <= AESXorxD(i);
      end loop;  -- i
    else                                    --XOR
      for i in 0 to 2 loop
        StateMatxDN(i+12) <= StateMatxDP(i) xor StateMatxDP(i+13);
      end loop;
      StateMatxDN(15) <= StateMatxDP(3) xor StateMatxDP(0) xor StateMatxDP(13);
    end if;

    -- Shift the state matrix
    for i in 11 downto 0 loop
      StateMatxDN(i) <= StateMatxDP(i+4);
    end loop;  -- i

    -- special cases
    case to_integer(StatexDP) is
      
      when 0 =>
        AESInxD(2)  <= DataBlkxDI(3);
        AESInxD(1)  <= DataBlkxDI(2);
        AESInxD(0)  <= DataBlkxDI(1);
        AESInxD(3)  <= DataBlkxDI(0);
        AESXorxD(0) <= AESOutxD(0) xor DataBlkxDI(12) xor BitCntxDI(0);
        AESXorxD(1) <= AESOutxD(1) xor DataBlkxDI(13) xor (not BitCntxDI(1));
        AESXorxD(2) <= AESOutxD(2) xor DataBlkxDI(14);
        AESXorxD(3) <= AESOutxD(3) xor DataBlkxDI(15);
        for i in 11 downto 0 loop
          StateMatxDN(i) <= DataBlkxDI(i+4);
        end loop;  -- i
        
      when 10 =>
        AESXorxD(0) <= AESOutxD(0) xor StateMatxDP(12);
        AESXorxD(1) <= AESOutxD(1) xor StateMatxDP(13) xor BitCntxDI(1);
        AESXorxD(2) <= AESOutxD(2) xor StateMatxDP(14) xor (not BitCntxDI(0));
        AESXorxD(3) <= AESOutxD(3) xor StateMatxDP(15);

      when 17 =>
        AESXorxD(0) <= AESOutxD(0) xor StateMatxDP(12);
        AESXorxD(1) <= AESOutxD(1) xor StateMatxDP(13);
        AESXorxD(2) <= AESOutxD(2) xor StateMatxDP(14) xor BitCntxDI(1);
        AESXorxD(3) <= AESOutxD(3) xor StateMatxDP(15) xor (not BitCntxDI(0));

      when 27 =>
        AESXorxD(0) <= AESOutxD(0) xor StateMatxDP(12) xor BitCntxDI(0);
        AESXorxD(1) <= AESOutxD(1) xor StateMatxDP(13);
        AESXorxD(2) <= AESOutxD(2) xor StateMatxDP(14);
        AESXorxD(3) <= AESOutxD(3) xor StateMatxDP(15) xor (not BitCntxDI(1));

      when others => null;

    end case;

  end process;


  p_mem : process (ClkxCI, RstxRBI)
  begin  -- process p_mem
    if RstxRBI = '0' then               -- asynchronous reset (active low)
      StatexDP    <= (others => '0');
      StateMatxDP <= (others => (others => '0'));

      
    elsif ClkxCI'event and ClkxCI = '1' then  -- rising clock edge
      StatexDP    <= StatexDN;
      StateMatxDP <= StateMatxDN;
      
    end if;
  end process p_mem;

end rtl;