1 files changed, 5 insertions, 5 deletions
diff --git a/src/shift_unit.v b/src/shift_unit.v
index 55b6add..eb931e2 100644
--- a/src/shift_unit.v
+++ b/src/shift_unit.v
@@ -1,17 +1,17 @@
 module shift_unit #(
   parameter N = 32
 )(
-  input [N-1:0] A,
-  input [clog2(N):0] SHAMT,
+  input signed [N-1:0] A,
+  input unsigned [N-1:0] SHAMT,
   input [1:0] OP, // 00: SLL, 01: SRL, 11: SRA 
   output reg [N-1:0] RESULT
 );
 
 always @ (*) begin
   case (OP)
-    2'b00: RESULT <= A << SHAMT;
-    2'b01: RESULT <= A >> SHAMT;
-    2'b11: RESULT <= A >>> SHAMT;
+    2'b00: RESULT <= A << SHAMT % N;
+    2'b01: RESULT <= A >> SHAMT % N;
+    2'b11: RESULT <= A >>> SHAMT % N;
   endcase
 end