2 files changed, 32 insertions, 16 deletions

diff --git a/src/rtl/chacha_core.v b/src/rtl/chacha_core.v
index 32f37f5..6cb1e51 100644
--- a/src/rtl/chacha_core.v
+++ b/src/rtl/chacha_core.v
@@ -82,8 +82,10 @@ module chacha_core(
   localparam CTRL_INIT     = 3'h1;
   localparam CTRL_ROUNDS0  = 3'h2;
   localparam CTRL_ROUNDS1  = 3'h3;
-  localparam CTRL_FINALIZE = 3'h4;
-  localparam CTRL_DONE     = 3'h5;
+  localparam CTRL_ROUNDS2  = 3'h4;
+  localparam CTRL_ROUNDS3  = 3'h5;
+  localparam CTRL_FINALIZE = 3'h6;
+  localparam CTRL_DONE     = 3'h7;
 
 
   //----------------------------------------------------------------
@@ -666,6 +668,20 @@ module chacha_core(
 
         CTRL_ROUNDS1:
           begin
+            update_state    = 1;
+            chacha_ctrl_new = CTRL_ROUNDS2;
+            chacha_ctrl_we  = 1;
+          end
+
+        CTRL_ROUNDS2:
+          begin
+            update_state    = 1;
+            chacha_ctrl_new = CTRL_ROUNDS3;
+            chacha_ctrl_we  = 1;
+          end
+
+        CTRL_ROUNDS3:
+          begin
             state_we     = 1;
             update_state = 1;
             qr_ctr_inc   = 1;
diff --git a/src/rtl/chacha_qr.v b/src/rtl/chacha_qr.v
index ba68d51..ec29c70 100644
--- a/src/rtl/chacha_qr.v
+++ b/src/rtl/chacha_qr.v
@@ -59,10 +59,10 @@ module chacha_qr(
   //----------------------------------------------------------------
   reg [31 : 0] a0_reg;
   reg [31 : 0] a0_new;
-//  reg [31 : 0] a1_reg;
-//  reg [31 : 0] a1_new;
-//  reg [31 : 0] c0_reg;
-//  reg [31 : 0] c0_new;
+  reg [31 : 0] a1_reg;
+  reg [31 : 0] a1_new;
+  reg [31 : 0] c0_reg;
+  reg [31 : 0] c0_new;
 
 
   //----------------------------------------------------------------
@@ -91,15 +91,15 @@ module chacha_qr(
       if (!reset_n)
         begin
           a0_reg <= 32'h0;
-//          a1_reg <= 32'h0;
-//          c0_reg <= 32'h0;
+          a1_reg <= 32'h0;
+          c0_reg <= 32'h0;
         end
 
       else
         begin
           a0_reg <= a0_new;
-//          a1_reg <= a1_new;
-//          c0_reg <= c0_new;
+          a1_reg <= a1_new;
+          c0_reg <= c0_new;
         end
     end // reg_update
 
@@ -134,21 +134,21 @@ module chacha_qr(
       d1 = {d0[15 : 0], d0[31 : 16]};
 
       c0 = c + d1;
-//      c0_new = c0;
+      c0_new = c0;
 
-      b0 = b ^ c0;
+      b0 = b ^ c0_reg;
       b1 = {b0[19 : 0], b0[31 : 20]};
 
       a1 = a0_reg + b1;
-//      a1_new = a1;
+      a1_new = a1;
 
-      d2 = d1 ^ a1;
+      d2 = d1 ^ a1_reg;
       d3 = {d2[23 : 0], d2[31 : 24]};
-      c1 = c0 + d3;
+      c1 = c0_reg + d3;
       b2 = b1 ^ c1;
       b3 = {b2[24 : 0], b2[31 : 25]};
 
-      internal_a_prim = a1;
+      internal_a_prim = a1_reg;
       internal_b_prim = b3;
       internal_c_prim = c1;
       internal_d_prim = d3;