diff --git a/aes/rtl/vhdl/aes_enc.vhd b/aes/rtl/vhdl/aes_enc.vhd
index 21d5258..6b3b443 100644
--- a/aes/rtl/vhdl/aes_enc.vhd
+++ b/aes/rtl/vhdl/aes_enc.vhd
@@ -52,11 +52,11 @@ begin
 
 
   IterG : if design_type = "ITER" generate
-  
-  
+
+
     signal s_round : t_enc_rounds;
-  
-  
+
+
   begin
 
 
@@ -79,7 +79,7 @@ begin
             if (accept_o = '1' and valid_i = '1') then
               accept_o <= '0';
               v_state  := set_state(data_i);
-              v_key    := (key_i(0 to 31), key_i(32 to 63), key_i(64 to 95), key_i(96 to 127));
+              v_key    := set_key(key_i);
               s_round  <= s_round + 1;
             end if;
 
@@ -133,7 +133,7 @@ begin
 
       cover {accept_o};
       assert always (accept_o -> s_round = 0);
-    
+
       cover {valid_i and accept_o};
       assert always (valid_i and accept_o -> next not accept_o);
 
@@ -152,7 +152,7 @@ begin
 
 
   end generate IterG;
-  
+
 
 
 end architecture rtl;
diff --git a/aes/rtl/vhdl/aes_pkg.vhd b/aes/rtl/vhdl/aes_pkg.vhd
index f148052..338eab6 100644
--- a/aes/rtl/vhdl/aes_pkg.vhd
+++ b/aes/rtl/vhdl/aes_pkg.vhd
@@ -130,8 +130,6 @@ package aes_pkg is
 
   constant c_rcon : t_rcon := (x"01", x"02", x"04", x"08", x"10", x"20", x"40", x"80", x"1B", x"36");
 
-  type t_mode is (ENCRYPT, DECRYPT);
-
 
   function bytesub    (input : std_logic_vector(7 downto 0)) return std_logic_vector;
   function invbytesub (input : std_logic_vector(7 downto 0)) return std_logic_vector;
@@ -145,8 +143,6 @@ package aes_pkg is
   function mixcolumns (input : t_datatable2d) return t_datatable2d;
   function invmixcolumns (input : t_datatable2d) return t_datatable2d;
 
-  function sortdata (input : std_logic_vector(127 downto 0)) return t_datatable2d;
-
   function gmul (a : std_logic_vector(7 downto 0); b : std_logic_vector(7 downto 0)) return std_logic_vector;
 
   function addroundkey (input : in t_datatable2d; key : in t_key) return t_datatable2d;
@@ -161,6 +157,8 @@ package aes_pkg is
 
   function get_state (input : in t_datatable2d) return std_logic_vector;
 
+  function set_key (input : in std_logic_vector(0 to 127)) return t_key;
+
   function to_string(input : t_datatable2d) return string;
 
 
@@ -171,27 +169,15 @@ end package aes_pkg;
 package body aes_pkg is
 
 
-  function sortdata (input : std_logic_vector(127 downto 0)) return t_datatable2d is
-    variable v_datamatrix : t_datatable2d;
-  begin
-    for outdex in 0 to 3 loop
-      for index in 0 to 3 loop
-        v_datamatrix(index)(outdex) := input(outdex*32+(index+1)*7 downto outdex*32+index*8);
-      end loop;
-    end loop;
-    return v_datamatrix;
-  end function sortdata;
-
-
   function bytesub (input : std_logic_vector(7 downto 0)) return std_logic_vector is
   begin
-    return(c_sbox(to_integer(unsigned(input(7 downto 4))))(to_integer(unsigned(input(3 downto 0)))));
+    return c_sbox(to_integer(unsigned(input(7 downto 4))))(to_integer(unsigned(input(3 downto 0))));
   end function bytesub;
 
 
   function invbytesub (input : std_logic_vector(7 downto 0)) return std_logic_vector is
   begin
-    return(c_sbox_invers(to_integer(unsigned(input(7 downto 4))))(to_integer(unsigned(input(3 downto 0)))));
+    return c_sbox_invers(to_integer(unsigned(input(7 downto 4))))(to_integer(unsigned(input(3 downto 0))));
   end function invbytesub;
 
 
@@ -384,6 +370,12 @@ package body aes_pkg is
   end function get_state;
 
 
+  function set_key (input : in std_logic_vector(0 to 127)) return t_key is
+  begin
+    return (input(0 to 31), input(32 to 63), input(64 to 95), input(96 to 127));
+  end function set_key;
+
+
   function to_string(input : t_datatable2d) return string is
   begin
     return '(' & to_hstring(input(0)(0)) & ',' & to_hstring(input(0)(1)) & ',' & to_hstring(input(0)(2)) & ',' & to_hstring(input(0)(3)) & ')' & LF &