diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 07a957a..511d0ff 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -1,10 +1,6 @@
 name: tests
 
-on:
-  push:
-    branches:
-      - main
-  pull_request:
+on: [push, pull_request]
 
 jobs:
   test:
@@ -23,9 +19,9 @@ jobs:
             nasm \
             nlohmann-json3-dev
 
-      - name: Download Circom Binary v2.1.5
+      - name: Download Circom Binary v2.1.5\8
         run: |
-          wget -qO /home/runner/work/circom https://github.com/iden3/circom/releases/download/v2.1.5/circom-linux-amd64
+          wget -qO /home/runner/work/circom https://github.com/iden3/circom/releases/download/v2.1.8/circom-linux-amd64
           chmod +x /home/runner/work/circom
           sudo mv /home/runner/work/circom /bin/circom
 
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..730650c
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 crema-labs
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
index 04941a1..7534868 100644
--- a/README.md
+++ b/README.md
@@ -1,131 +1,43 @@
-# Circomkit Examples
+# aes-circom
 
-In this repository, we are using [Circomkit](https://github.com/erhant/circomkit) to test some example circuits using Mocha. The circuits and the statements that they prove are as follows:
-
-- **Multiplier**: "I know `n` factors that make up some number".
-- **Fibonacci**: "I know the `n`'th Fibonacci number".
-- **SHA256**: "I know the `n`-byte preimage of some SHA256 digest".
-- **Sudoku**: "I know the solution to some `(n^2)x(n^2)` Sudoku puzzle".
-- **Floating-Point Addition**: "I know two floating-point numbers that make up some number with `e` exponent and `m` mantissa bits." (adapted from [Berkeley ZKP MOOC 2023 - Lab](https://github.com/rdi-berkeley/zkp-mooc-lab)).
-
-## CLI Usage
-
-To use Circomkit CLI with a circuit, let's say for Sudoku 9x9, we follow the steps below:
-
-1. We write a circuit config in `circuits.json` with the desired parameters. In this case, we are working with the 9x9 Sudoku solution circuit, and the board size is calculated by the square of our template parameter so we should give 3. Furthermore, `puzzle` is a public input so we should specify that too.
-
-```json
-{
-  "sudoku_9x9": {
-    "file": "sudoku",
-    "template": "Sudoku",
-    "pubs": ["puzzle"],
-    "params": [3]
-  }
-}
-```
-
-2. Compile the circuit with Circomkit, providing the same circuit name as in `circuits.json`:
-
-```sh
-npx circomkit compile sudoku_9x9
-
-# print circuit info if you want to
-npx circomkit info sudoku_9x9
-```
-
-3. Commence circuit-specific setup. Normally, this requires us to download a Phase-1 PTAU file and provide it's path; however, Circomkit can determine the required PTAU and download it automatically when using `bn128` curve, thanks to [Perpetual Powers of Tau](https://github.com/privacy-scaling-explorations/perpetualpowersoftau). In this case, `sudoku_9x9` circuit has 4617 constraints, so Circomkit will download `powersOfTau28_hez_final_13.ptau` (see [here](https://github.com/iden3/snarkjs#7-prepare-phase-2)).
-
-```sh
-npx circomkit setup sudoku_9x9
+This repository contains generic implementation for AES encryption in Circom.
 
-# alternative: provide the PTAU yourself
-npx circomkit setup sudoku_9x9 <path-to-ptau>
-```
+## AES
 
-4. Prepare your input file under `./inputs/sudoku_9x9/default.json`.
+AES is a symmetric encryption algorithm that was established by the U.S. National Institute of Standards and Technology (NIST) in 2001. It is a subset of the Rijndael block cipher. AES has a fixed block size of 128 bits and a key size of 128, 192, or 256 bits. The algorithm is based on a design principle known as a substitution-permutation network (SPN).It is a symmetric block cipher that can encrypt (encipher) and decrypt (decipher) information. Encryption converts data to an unintelligible form called ciphertext; decrypting the ciphertext converts the data back into its original form, called plaintext. The symmetric key signifies that the same key is used for both encryption and decryption.
 
-```json
-{
-  "solution": [
-    [1, 9, 4, 8, 6, 5, 2, 3, 7],
-    [7, 3, 5, 4, 1, 2, 9, 6, 8],
-    [8, 6, 2, 3, 9, 7, 1, 4, 5],
-    [9, 2, 1, 7, 4, 8, 3, 5, 6],
-    [6, 7, 8, 5, 3, 1, 4, 2, 9],
-    [4, 5, 3, 9, 2, 6, 8, 7, 1],
-    [3, 8, 9, 6, 5, 4, 7, 1, 2],
-    [2, 4, 6, 1, 7, 9, 5, 8, 3],
-    [5, 1, 7, 2, 8, 3, 6, 9, 4]
-  ],
-  "puzzle": [
-    [0, 0, 0, 8, 6, 0, 2, 3, 0],
-    [7, 0, 5, 0, 0, 0, 9, 0, 8],
-    [0, 6, 0, 3, 0, 7, 0, 4, 0],
-    [0, 2, 0, 7, 0, 8, 0, 5, 0],
-    [0, 7, 8, 5, 0, 0, 0, 0, 0],
-    [4, 0, 0, 9, 0, 6, 0, 7, 0],
-    [3, 0, 9, 0, 5, 0, 7, 0, 2],
-    [0, 4, 0, 1, 0, 9, 0, 8, 0],
-    [5, 0, 7, 0, 8, 0, 0, 9, 4]
-  ]
-}
-```
-
-5. We are ready to create a proof!
-
-```sh
-npx circomkit prove sudoku_9x9 default
-```
+Read more about AES here := [FIPS 197](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197-upd1.pdf).
+Simple Rust implementation of AES can be found here := [tinyaes](https://docs.rs/crate/tinyaes/latest/source/src/aes_core.rs)
 
-6. We can then verify our proof. You can try and modify the public input at `./build/sudoku_9x9/default/public.json` and see if the proof verifies or not!
+## Circuit
 
-```sh
-npx circomkit verify sudoku_9x9 default
-```
+The circuits contain components for AES forward encryption. The implementation strictly follows the AES standard mentioned in the FIPS 197 document. The circuit is designed to be generic and can be used for any key size (128, 192, 256 bits) and block size (128 bits).
 
-## In-Code Usage
+Check the [Cipher](./circuits/aes.circom) and [KeyExpansion](./circuits/key_expansion.circom) circuits for visual representation of the design.
 
-If you would like to use Circomkit within the code itself, rather than the CLI, you can see the example at `src/index.ts`. You can `yarn start` to see it in action.
+## Design Decisions
 
-```ts
-// create circomkit
-const circomkit = new Circomkit({
-  protocol: "groth16",
-});
+The circuit only support the forward encryption of AES as we believe that the proof of computation for any proprietary use case can be refactored to use the forward encryption instead of the decryption. 
 
-// artifacts output at `build/multiplier_3` directory
-await circomkit.compile("multiplier_3", {
-  file: "multiplier",
-  template: "Multiplier",
-  params: [3],
-});
+ 💡 Create an issue if you think that the decryption circuit is necessary.
 
-// proof & public signals at `build/multiplier_3/my_input` directory
-await circomkit.prove("multiplier_3", "my_input", { in: [3, 5, 7] });
+## Circomkit
 
-// verify with proof & public signals at `build/multiplier_3/my_input`
-const ok = await circomkit.verify("multiplier_3", "my_input");
-if (ok) {
-  circomkit.log("Proof verified!", "success");
-} else {
-  circomkit.log("Verification failed.", "error");
-}
-```
+In this repository, we are using [Circomkit](https://github.com/erhant/circomkit) to test some example circuits using Mocha. The circuits and the statements that they prove are as follows:
 
-## Configuration
+### Configuration
 
 Circomkit checks for `circomkit.json` to override it's default configurations. We could for example change the target version, prime field and the proof system by setting `circomkit.json` to be:
 
 ```json
 {
-  "version": "2.1.2",
+  "version": "2.1.8",
   "protocol": "plonk",
   "prime": "bls12381"
 }
 ```
 
-## Testing
+### Testing
 
 You can use the following commands to test the circuits:
 
@@ -134,5 +46,15 @@ You can use the following commands to test the circuits:
 yarn test
 
 # test a specific circuit
-yarn test -g <circuit-name>
+yarn test -g <template-name>
 ```
+
+## Roadmap
+
+- [x] AES Forward Encryption Circuit
+- [ ] Add AES-CTR mode (priority for ECIES implementaion)
+- [ ] Add all other modes adhering to [NIST standards](https://nvlpubs.nist.gov/nistpubs/legacy/sp/nistspecialpublication800-38a.pdf)
+
+## Contribution
+
+Feel free to contribute to this repository by creating issues or pull requests. We are open to any suggestions or improvements.
\ No newline at end of file
diff --git a/circomkit.json b/circomkit.json
index 777900c..cd5fed7 100644
--- a/circomkit.json
+++ b/circomkit.json
@@ -1,5 +1,5 @@
 {
-  "version": "2.1.2",
-  "proofSystem": "plonk",
+  "version": "2.1.8",
+  "proofSystem": "groth16",
   "curve": "bn128"
 }
diff --git a/circuits.json b/circuits.json
index bd6765d..6cb94ad 100644
--- a/circuits.json
+++ b/circuits.json
@@ -1,7 +1,7 @@
 {
-  "fibonacci_11": {
-    "file": "fibonacci",
-    "template": "Fibonacci",
-    "params": [11]
+  "cipher_4": {
+    "file": "cipher",
+    "template": "Cipher",
+    "params": [4]
   }
 }
diff --git a/circuits/cipher.circom b/circuits/cipher.circom
new file mode 100644
index 0000000..f02d040
--- /dev/null
+++ b/circuits/cipher.circom
@@ -0,0 +1,106 @@
+pragma circom 2.1.8;
+
+include "key_expansion.circom";
+include "circomlib/circuits/comparators.circom";
+include "circomlib/circuits/bitify.circom";
+include "circomlib/circuits/gates.circom";
+include "transformations.circom";
+include "mix_columns.circom";
+
+// Cipher Process
+// nk: number of keys which can be 4, 6, 8
+// AES 128, 192, 256 have 10, 12, 14 rounds.
+// Input Block   Initial Round Key          Round Key             Final Round Key
+//     │                │                       │                       │
+//     ▼                ▼                       ▼                       ▼
+//  ┌─────────┐    ┌──────────┐ ┌────────┐ ┌──────────┐ ┌────────┐ ┌──────────┐
+//  │ Block   │──► │   Add    │ │  Sub   │ │   Mix    │ │  Sub   │ │   Add    │
+//  │         │    │  Round   │ │ Bytes  │ │ Columns  │ │ Bytes  │ │  Round   │
+//  │         │    │   Key    │ │        │ │          │ │        │ │   Key    │
+//  └─────────┘    └────┬─────┘ └───┬────┘ └────┬─────┘ └───┬────┘ └────┬─────┘
+//                      │           │           │           │           │
+//                      ▼           ▼           ▼           ▼           ▼
+//                 ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐
+//                 │ Round 0 │ │ Round 1 │ │ Round 2 │ │ Round   │ │  Final  │
+//                 │         │ │   to    │ │   to    │ │ Nr - 1  │ │  Round  │
+//                 │         │ │ Nr - 2  │ │ Nr - 1  │ │         │ │         │
+//                 └─────────┘ └─────────┘ └─────────┘ └─────────┘ └────┬────┘
+//                                                                      │
+//                                                                      ▼
+//                                                                 Ciphertext
+
+// @param nk: number of keys which can be 4, 6, 8
+// @inputs block: 4x4 matrix representing the input block
+// @inputs key: array of nk*4 bytes representing the key
+// @outputs cipher: 4x4 matrix representing the output block
+template Cipher(nk){
+        assert(nk == 4 || nk == 6 || nk == 8 );
+        signal input block[4][4];
+        signal output cipher[4][4];
+        signal input key[nk * 4];
+
+        var nr = Rounds(nk);
+        
+        component keyExpansion = KeyExpansion(nk,nr);
+        keyExpansion.key <== key;
+
+        component addRoundKey[nr+1]; 
+        component subBytes[nr];
+        component shiftRows[nr];
+        component mixColumns[nr-1];
+
+        signal interBlock[nr][4][4];
+
+        addRoundKey[0] = AddRoundKey();
+        addRoundKey[0].state <== block;
+        for (var i = 0; i < 4; i++) {
+                addRoundKey[0].roundKey[i] <== keyExpansion.keyExpanded[i];
+        }
+
+        interBlock[0] <== addRoundKey[0].newState;
+        for (var i = 1; i < nr; i++) {
+                subBytes[i-1] = SubBlock();
+                subBytes[i-1].state <== interBlock[i-1];
+
+                shiftRows[i-1] = ShiftRows();
+                shiftRows[i-1].state <== subBytes[i-1].newState;
+
+                mixColumns[i-1] = MixColumns();
+                mixColumns[i-1].state <== shiftRows[i-1].newState;
+
+                addRoundKey[i] = AddRoundKey();
+                addRoundKey[i].state <== mixColumns[i-1].out;
+                 for (var j = 0; j < 4; j++) {
+                        addRoundKey[i].roundKey[j] <== keyExpansion.keyExpanded[j + (i * 4)];
+                }
+
+                interBlock[i] <== addRoundKey[i].newState;
+        }
+
+        subBytes[nr-1] = SubBlock();
+        subBytes[nr-1].state <== interBlock[nr-1];
+
+        shiftRows[nr-1] = ShiftRows();
+        shiftRows[nr-1].state <== subBytes[nr-1].newState;
+
+        addRoundKey[nr] = AddRoundKey();
+        addRoundKey[nr].state <== shiftRows[nr-1].newState;
+        for (var i = 0; i < 4; i++) {
+                addRoundKey[nr].roundKey[i] <== keyExpansion.keyExpanded[i + (nr * 4)];
+        }
+
+        cipher <== addRoundKey[nr].newState;
+}
+
+// @param nk: number of keys which can be 4, 6, 8
+// @returns number of rounds
+// AES 128, 192, 256 have 10, 12, 14 rounds.
+function Rounds (nk) {
+    if (nk == 4) {
+       return 10;
+    } else if (nk == 6) {
+        return 12;
+    } else {
+        return 14;
+    }
+}
\ No newline at end of file
diff --git a/circuits/fibonacci.circom b/circuits/fibonacci.circom
deleted file mode 100644
index 19adc10..0000000
--- a/circuits/fibonacci.circom
+++ /dev/null
@@ -1,33 +0,0 @@
-pragma circom 2.0.0;
-
-// Fibonacci with custom starting numbers
-template Fibonacci(n) {
-  assert(n >= 2);
-  signal input in[2];
-  signal output out;
-
-  signal fib[n+1];
-  fib[0] <== in[0];
-  fib[1] <== in[1];
-  for (var i = 2; i <= n; i++) {
-    fib[i] <== fib[i-2] + fib[i-1];
-  }
-
-  out <== fib[n];
-}
-
-// Fibonacci with custom starting numbers, recursive & inefficient
-template FibonacciRecursive(n) {
-  signal input in[2];
-  signal output out;
-  component f1, f2;
-  if (n <= 1) {
-    out <== in[n];
-  } else {
-    f1 = FibonacciRecursive(n-1);
-    f1.in <== in;
-    f2 = FibonacciRecursive(n-2);
-    f2.in <== in;
-    out <== f1.out + f2.out;
-  }
-}
diff --git a/circuits/key_expansion.circom b/circuits/key_expansion.circom
index d395bb9..5686474 100644
--- a/circuits/key_expansion.circom
+++ b/circuits/key_expansion.circom
@@ -1,18 +1,51 @@
-pragma circom 2.0.0;
+pragma circom 2.1.8;
 
 include "sbox128.circom";
-include "../node_modules/circomlib/circuits/comparators.circom";
-include "../node_modules/circomlib/circuits/bitify.circom";
-include "../node_modules/circomlib/circuits/gates.circom";
-
-//nk is the number of keys which can be 4, 6, 8 
-//nr is the number of rounds which can be 10, 12, 14
-template KeyExpansion(nk, nr) {
+include "utils.circom";
+
+// Key Expansion Process
+//
+// Original Key (Nk words)
+// ┌───┬───┬───┬───┐
+// │W0 │W1 │W2 │W3 │  (for AES-128, Nk=4)
+// └─┬─┴─┬─┴─┬─┴─┬─┘
+//   │   │   │   │
+//   ▼   ▼   ▼   ▼
+// ┌───────────────────────────────────────────────────────┐
+// │                 Key Expansion Process                 │
+// │                                                       │
+// │  ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐      │
+// │  │RotWord  │ │SubWord  │ │  XOR    │ │  XOR    │      │
+// │  │         │ │         │ │ Rcon(i) │ │ W[i-Nk] │      │
+// │  └────┬────┘ └────┬────┘ └────┬────┘ └────┬────┘      │
+// │       │           │           │           │           │
+// │       └───────────┴───────────┴───────────┘           │
+// │                       │                               │
+// │                       ▼                               │
+// │            ┌─────────────────────┐                    │
+// │            │  New Expanded Key   │                    │
+// │            │       Word          │                    │
+// │            └─────────────────────┘                    │
+// │                       │                               │
+// └───────────────────────┼───────────────────────────────┘
+//                         │
+//                         ▼
+//               Expanded Key Words
+//        ┌───┬───┬───┬───┬───┬───┬───┬───┐
+//        │W4 │W5 │W6 │W7 │W8 │W9 │...│W43│  (for AES-128, 44 words total)
+//        └───┴───┴───┴───┴───┴───┴───┴───┘
+
+
+// @param nk: number of keys which can be 4, 6, 8
+// @param nr: number of rounds which can be 10, 12, 14 for AES 128, 192, 256
+// @inputs key: array of nk*4 bytes representing the key
+// @outputs keyExpanded: array of (nr+1)*4 words i.e for AES 128, 192, 256 it will be 44, 52, 60 words
+template KeyExpansion(nk,nr) {
+    assert(nk == 4 || nk == 6 || nk == 8 );    
     signal input key[nk * 4];
     
     var totalWords = (4 * (nr + 1));
-    var effectiveRounds = (totalWords % nk == 0 ? totalWords - nk : totalWords) \ nk;
-    var leftoverWords = totalWords - (effectiveRounds * nk);
+    var effectiveRounds = nk == 4 ? 10 : totalWords\nk;
 
     signal output keyExpanded[totalWords][4];
 
@@ -25,7 +58,7 @@ template KeyExpansion(nk, nr) {
     component nextRound[effectiveRounds];
     
     for (var round = 1; round <= effectiveRounds; round++) {
-        var outputWordLen = round == effectiveRounds ? leftoverWords : nk;
+        var outputWordLen = round == effectiveRounds ? 4 : nk; 
         nextRound[round - 1] = NextRound(nk, outputWordLen);
 
         for (var i = 0; i < nk; i++) {
@@ -44,13 +77,14 @@ template KeyExpansion(nk, nr) {
     }
 }
 
-//nk, output
+// @param nk: number of keys which can be 4, 6, 8
+// @param o: number of output words which can be 4 or nk
 template NextRound(nk, o){
     signal input key[nk][4]; 
     signal input round;
     signal output nextKey[o][4];
 
-    component rotateWord = RotateWord(1);
+    component rotateWord = Rotate(1, 4);
     for (var i = 0; i < 4; i++) {
         rotateWord.bytes[i] <== key[nk - 1][i];
     }
@@ -87,122 +121,3 @@ template NextRound(nk, o){
 }
 
 
-template BytesToWords(n) {
-    assert(n % 4 == 0);
-    signal input bytes[n];
-    signal output words[n / 4][n];
-
-    for (var i = 0; i < n / 4; i++) {
-        for(var j = 0; j < 4; j++) {
-            words[i][j] <== bytes[i * 4 + j];
-        }
-    }
-}
-
-template RotateWord(rotation) {
-    assert(rotation < 4);
-    signal input bytes[4];
-    signal output rotated[4];
-
-    signal copy[rotation];
-
-    for(var i = 0; i < rotation; i++) {
-        copy[i] <== bytes[i];
-    }
-
-    for(var i = 0; i < 4 - rotation; i++) {
-        rotated[i] <== bytes[i + rotation];
-    }
-
-    for(var i = 4 - rotation; i < 4; i++) {
-        rotated[i] <== copy[i - 4 + rotation];
-    }
-}
-
-template SubstituteWord() {
-    signal input bytes[4];
-    signal output substituted[4];
-
-    component sbox[4];
-
-    for(var i = 0; i < 4; i++) {
-        sbox[i] = Sbox128();
-        sbox[i].in <== bytes[i];
-        substituted[i] <== sbox[i].out;
-    }
-}
-
-template RCon() {
-    signal input round;
-    signal output out[4];
-
-    var rcon[10][4] = [
-        [0x01, 0x00, 0x00, 0x00],
-        [0x02, 0x00, 0x00, 0x00],
-        [0x04, 0x00, 0x00, 0x00],
-        [0x08, 0x00, 0x00, 0x00],
-        [0x10, 0x00, 0x00, 0x00],
-        [0x20, 0x00, 0x00, 0x00],
-        [0x40, 0x00, 0x00, 0x00],
-        [0x80, 0x00, 0x00, 0x00],
-        [0x1b, 0x00, 0x00, 0x00],
-        [0x36, 0x00, 0x00, 0x00]
-    ];
-
-    component isEqual[10];
-    signal sum[10][4];
-
-    for (var i = 0; i < 10; i++) {
-        isEqual[i] = IsEqual();
-        isEqual[i].in[0] <== round - 1;
-        isEqual[i].in[1] <== i;
-    }
-
-    sum[0] <== [isEqual[0].out * rcon[0][0], isEqual[1].out * rcon[0][1], isEqual[2].out * rcon[0][2], isEqual[3].out * rcon[0][3]];
-    for (var i = 1; i < 10; i++) {
-        for (var j = 0; j < 4; j++) {
-            sum[i][j] <== sum[i - 1][j] + isEqual[i].out * rcon[i][j];
-        }
-    }
-
-    out <== sum[9];
-}
-
-template XorWord() {
-    signal input bytes1[4];
-    signal input bytes2[4];
-    
-    component n2b[4 * 2];
-    component b2n[4];
-    component xor[4][8];
-
-    signal output out[4];
-
-    for(var i = 0; i < 4; i++) {
-        n2b[2 * i] = Num2Bits(8);
-        n2b[2 * i + 1] = Num2Bits(8);
-        n2b[2 * i].in <== bytes1[i];
-        n2b[2 * i + 1].in <== bytes2[i];
-        b2n[i] = Bits2Num(8);
-
-        for (var j = 0; j < 8; j++) {
-            xor[i][j] = XOR();
-            xor[i][j].a <== n2b[2 * i].out[j];
-            xor[i][j].b <== n2b[2 * i + 1].out[j];
-            b2n[i].in[j] <== xor[i][j].out;
-        }
-
-        out[i] <== b2n[i].out;
-    }
-}
-
-template WordSelector() {
-    signal input condition;
-    signal input bytes1[4];
-    signal input bytes2[4];
-    signal output out[4];
-
-    for (var i = 0; i < 4; i++) {
-        out[i] <== condition * (bytes1[i] - bytes2[i]) + bytes2[i];
-    }
-}
\ No newline at end of file
diff --git a/circuits/mix_columns.circom b/circuits/mix_columns.circom
new file mode 100644
index 0000000..817103a
--- /dev/null
+++ b/circuits/mix_columns.circom
@@ -0,0 +1,198 @@
+pragma circom 2.1.8;
+
+include "transformations.circom";
+include "circomlib/circuits/comparators.circom";
+include "circomlib/circuits/bitify.circom";
+include "circomlib/circuits/gates.circom";
+
+
+// MixColumns: Applies the equation for each column:
+// [s'0,c]   [2 3 1 1][s0,c]
+// [s'1,c] = [1 2 3 1][s1,c]
+// [s'2,c]   [1 1 2 3][s2,c]
+// [s'3,c]   [3 1 1 2][s3,c]
+// Where c is the column number, s are input bytes, s' are output bytes
+template MixColumns(){
+    signal input state[4][4];
+    signal output out[4][4];
+
+    component s0[4];
+    component s1[4];
+    component s2[4];
+    component s3[4];
+
+    for (var i = 0; i < 4; i++) {
+        s0[i] = S0();
+        s1[i] = S1();
+        s2[i] = S2();
+        s3[i] = S3();
+
+        for(var j = 0; j < 4; j++) {
+            s0[i].in[j] <== state[j][i];
+            s1[i].in[j] <== state[j][i];
+            s2[i].in[j] <== state[j][i];
+            s3[i].in[j] <== state[j][i];
+        }
+
+        out[0][i] <== s0[i].out;
+        out[1][i] <== s1[i].out;
+        out[2][i] <== s2[i].out;
+        out[3][i] <== s3[i].out;
+    }
+}
+
+// S0: Implements the equation
+// out = (2 • in[0]) ⊕ (3 • in[1]) ⊕ in[2] ⊕ in[3]
+template S0(){
+    signal input in[4];
+    signal output out;
+    component num2bits[4];
+    component xor[3];
+
+    for (var i = 0; i < 4; i++) {
+        num2bits[i] = Num2Bits(8);
+        num2bits[i].in <== in[i];
+    }
+
+    component mul = XTimes2();
+    mul.in <== num2bits[0].out;
+
+    component mul2 = XTimes(3);
+    mul2.in <== num2bits[1].out;
+
+    xor[0] = XorBits();
+    xor[0].a <== mul.out;
+    xor[0].b <== mul2.out;
+
+    xor[1] = XorBits();
+    xor[1].a <== xor[0].out;
+    xor[1].b <== num2bits[2].out;
+
+    xor[2] = XorBits();
+    xor[2].a <== xor[1].out;
+    xor[2].b <== num2bits[3].out;
+
+    component b2n = Bits2Num(8);
+    for (var i = 0; i < 8; i++) {
+        b2n.in[i] <== xor[2].out[i];
+    }
+
+    out <== b2n.out;
+}
+
+// S1: Implements the equation
+// out = in[0] ⊕ (2 • in[1]) ⊕ (3 • in[2]) ⊕ in[3]
+template S1(){
+    signal input in[4];
+    signal output out;
+    component num2bits[4];
+    component xor[3];
+
+    for (var i = 0; i < 4; i++) {
+        num2bits[i] = Num2Bits(8);
+        num2bits[i].in <== in[i];
+    }
+
+    component mul = XTimes2();
+    mul.in <== num2bits[1].out;
+
+    component mul2 = XTimes(3);
+    mul2.in <== num2bits[2].out;
+
+    xor[0] = XorBits();
+    xor[0].a <== num2bits[0].out;
+    xor[0].b <== mul.out;
+
+    xor[1] = XorBits();
+    xor[1].a <== xor[0].out;
+    xor[1].b <== mul2.out;
+
+    xor[2] = XorBits();
+    xor[2].a <== xor[1].out;
+    xor[2].b <== num2bits[3].out;
+
+    component b2n = Bits2Num(8);
+    for (var i = 0; i < 8; i++) {
+        b2n.in[i] <== xor[2].out[i];
+    }
+
+    out <== b2n.out;
+}
+
+// S2: Implements the equation
+// out = in[0] ⊕ in[1] ⊕ (2 • in[2]) ⊕ (3 • in[3])
+template S2() { 
+    signal input in[4];
+    signal output out;
+    component num2bits[4];
+    component xor[3];
+
+    for (var i = 0; i < 4; i++) {
+        num2bits[i] = Num2Bits(8);
+        num2bits[i].in <== in[i];
+    }
+
+    xor[0] = XorBits();
+    xor[0].a <== num2bits[0].out;
+    xor[0].b <== num2bits[1].out;
+
+    component mul2 = XTimes2();
+    mul2.in <== num2bits[2].out;
+
+    component mul = XTimes(3);
+    mul.in <== num2bits[3].out;
+
+    xor[1] = XorBits();
+    xor[1].a <== xor[0].out;
+    xor[1].b <== mul2.out;
+
+    xor[2] = XorBits();
+    xor[2].a <== xor[1].out;
+    xor[2].b <== mul.out;
+
+    component b2n = Bits2Num(8);
+    for (var i = 0; i < 8; i++) {
+        b2n.in[i] <== xor[2].out[i];
+    }
+
+    out <== b2n.out;
+}
+
+// S3: Implements the equation
+// out = (3 • in[0]) ⊕ in[1] ⊕ in[2] ⊕ (2 • in[3])
+template S3() {
+    signal input in[4];
+    signal output out;
+    component num2bits[4];
+    component xor[3];
+
+    for (var i = 0; i < 4; i++) {
+        num2bits[i] = Num2Bits(8);
+        num2bits[i].in <== in[i];
+    }
+
+    component mul3 = XTimes(3);
+    mul3.in <== num2bits[0].out;
+
+    xor[0] = XorBits();
+    xor[0].a <== mul3.out;
+    xor[0].b <== num2bits[1].out;
+
+    xor[1] = XorBits();
+    xor[1].a <== xor[0].out;
+    xor[1].b <== num2bits[2].out;
+
+    component mul2 = XTimes2();
+    mul2.in <== num2bits[3].out;
+
+    xor[2] = XorBits();
+    xor[2].a <== mul2.out;
+    xor[2].b <== xor[1].out;
+
+    component b2n = Bits2Num(8);
+    for (var i = 0; i < 8; i++) {
+        b2n.in[i] <== xor[2].out[i];
+    }
+
+    out <== b2n.out;
+}
\ No newline at end of file
diff --git a/circuits/sbox128.circom b/circuits/sbox128.circom
index e11e55e..b1dba34 100644
--- a/circuits/sbox128.circom
+++ b/circuits/sbox128.circom
@@ -1,8 +1,8 @@
-pragma circom 2.0.0;
+pragma circom 2.1.8;
 
 include "circomlib/circuits/comparators.circom";
 
-template Sbox128() {
+template SBox128() {
     signal input in;
     signal output out;
 
diff --git a/circuits/transformations.circom b/circuits/transformations.circom
new file mode 100644
index 0000000..83625bf
--- /dev/null
+++ b/circuits/transformations.circom
@@ -0,0 +1,118 @@
+pragma circom 2.1.8;
+
+include "circomlib/circuits/comparators.circom";
+include "circomlib/circuits/bitify.circom";
+include "circomlib/circuits/gates.circom";
+
+// ShiftRows: Performs circular left shift on each row
+// 0, 1, 2, 3 shifts for rows 0, 1, 2, 3 respectively
+template ShiftRows(){
+    signal input state[4][4];
+    signal output newState[4][4];
+
+    component shiftWord[4];
+
+    for (var i = 0; i < 4; i++) {
+        // Rotate: Performs circular left shift on each row
+        shiftWord[i] = Rotate(i, 4);
+        shiftWord[i].bytes <== state[i];
+        newState[i] <== shiftWord[i].rotated;
+    }
+}
+
+ // Applies S-box substitution to each byte
+template SubBlock(){
+        signal input state[4][4];
+        signal output newState[4][4];
+        component sbox[4];
+
+        for (var i = 0; i < 4; i++) {
+                sbox[i] = SubstituteWord();
+                sbox[i].bytes <== state[i];
+                newState[i] <== sbox[i].substituted;
+        }
+}
+
+// AddRoundKey: XORs the state with transposed the round key
+template AddRoundKey(){
+    signal input state[4][4];
+    signal input roundKey[4][4];
+    signal output newState[4][4];
+
+    component xorbyte[4][4];
+
+    for (var i = 0; i < 4; i++) {
+        for (var j = 0; j < 4; j++) {
+            xorbyte[i][j] = XorByte();
+            xorbyte[i][j].a <== state[i][j];
+            xorbyte[i][j].b <== roundKey[j][i];
+            newState[i][j] <== xorbyte[i][j].out;
+        }
+    }
+}
+
+// XTimes2: Multiplies by 2 in GF(2^8)
+template XTimes2(){
+    signal input in[8];
+    signal output out[8];
+
+    component xtimeConstant = Num2Bits(8);
+    xtimeConstant.in <== 0x1b;
+
+    component xor[7];
+
+    component isZero = IsZero();
+    isZero.in <== in[7];
+
+    out[0] <== 1-isZero.out;
+    for (var i = 0; i < 7; i++) {
+        xor[i] = XOR();
+        xor[i].a <== in[i];
+        xor[i].b <== xtimeConstant.out[i+1] * (1-isZero.out);
+        out[i+1] <== xor[i].out;
+    }
+}
+
+// XTimes: Multiplies by n in GF(2^8)
+// This uses a fast multiplication algorithm that uses the XTimes2 component
+// Number of constaints is always constant
+template XTimes(n){
+    signal input in[8];
+    signal output out[8];
+
+    component bits = Num2Bits(8);
+    bits.in <== n;
+
+    component XTimes2[7];
+
+    XTimes2[0] = XTimes2();
+    XTimes2[0].in <== in;
+
+    for (var i = 1; i < 7; i++) {
+            XTimes2[i] = XTimes2();
+            XTimes2[i].in <== XTimes2[i-1].out;
+    }
+
+    component xor[8];
+    component mul[8];
+    signal inter[8][8];
+
+    mul[0] = MulByte();
+    mul[0].a <== bits.out[0];
+    mul[0].b <== in;
+    inter[0] <== mul[0].c;
+
+    for (var i = 1; i < 8; i++) {
+        mul[i] = MulByte();
+        mul[i].a <== bits.out[i];
+        mul[i].b <== XTimes2[i-1].out;
+
+                xor[i] = XorBits();
+                xor[i].a <== inter[i-1];
+                xor[i].b <== mul[i].c;
+                inter[i] <== xor[i].out;
+        }
+
+    out <== inter[7];
+}
+
diff --git a/circuits/utils.circom b/circuits/utils.circom
new file mode 100644
index 0000000..26a4ea3
--- /dev/null
+++ b/circuits/utils.circom
@@ -0,0 +1,161 @@
+pragma circom 2.1.8;
+
+include "sbox128.circom";
+include "circomlib/circuits/comparators.circom";
+include "circomlib/circuits/bitify.circom";
+include "circomlib/circuits/gates.circom";
+
+// Converts an array of bytes to an array of words
+template BytesToWords(n) {
+    assert(n % 4 == 0);
+    signal input bytes[n];
+    signal output words[n \ 4][n];
+
+    for (var i = 0; i < n \ 4; i++) {
+        for(var j = 0; j < 4; j++) {
+            words[i][j] <== bytes[i * 4 + j];
+        }
+    }
+}
+
+// Rotates an array of bytes to the left by a specified rotation
+template Rotate(rotation, length) {
+    assert(rotation < length);
+    signal input bytes[length];
+    signal output rotated[length];
+
+    for(var i = 0; i < length - rotation; i++) {
+        rotated[i] <== bytes[i + rotation];
+    }
+
+    for(var i = length - rotation; i < length; i++) {
+        rotated[i] <== bytes[i - length + rotation];
+    }
+}
+
+// Substitutes each byte in a word using the S-Box
+template SubstituteWord() {
+    signal input bytes[4];
+    signal output substituted[4];
+
+    component sbox[4];
+
+    for(var i = 0; i < 4; i++) {
+        sbox[i] = SBox128();
+        sbox[i].in <== bytes[i];
+        substituted[i] <== sbox[i].out;
+    }
+}
+
+// Outputs a round constant for a given round number
+template RCon() {
+    signal input round;
+    signal output out[4];
+
+    assert(round > 0 && round <= 10);
+
+    var rcon[10][4] = [
+        [0x01, 0x00, 0x00, 0x00],
+        [0x02, 0x00, 0x00, 0x00],
+        [0x04, 0x00, 0x00, 0x00],
+        [0x08, 0x00, 0x00, 0x00],
+        [0x10, 0x00, 0x00, 0x00],
+        [0x20, 0x00, 0x00, 0x00],
+        [0x40, 0x00, 0x00, 0x00],
+        [0x80, 0x00, 0x00, 0x00],
+        [0x1b, 0x00, 0x00, 0x00],
+        [0x36, 0x00, 0x00, 0x00]
+    ];
+
+    out <-- rcon[round-1];
+}
+
+
+// XORs two words (arrays of 4 bytes each)
+template XorWord() {
+    signal input bytes1[4];
+    signal input bytes2[4];
+    
+    component n2b[4 * 2];
+    component b2n[4];
+    component xor[4][8];
+
+    signal output out[4];
+
+    for(var i = 0; i < 4; i++) {
+        n2b[2 * i] = Num2Bits(8);
+        n2b[2 * i + 1] = Num2Bits(8);
+        n2b[2 * i].in <== bytes1[i];
+        n2b[2 * i + 1].in <== bytes2[i];
+        b2n[i] = Bits2Num(8);
+
+        for (var j = 0; j < 8; j++) {
+            xor[i][j] = XOR();
+            xor[i][j].a <== n2b[2 * i].out[j];
+            xor[i][j].b <== n2b[2 * i + 1].out[j];
+            b2n[i].in[j] <== xor[i][j].out;
+        }
+
+        out[i] <== b2n[i].out;
+    }
+}
+
+// Selects between two words based on a condition
+template WordSelector() {
+    signal input condition;
+    signal input bytes1[4];
+    signal input bytes2[4];
+    signal output out[4];
+
+    for (var i = 0; i < 4; i++) {
+        out[i] <== condition * (bytes1[i] - bytes2[i]) + bytes2[i];
+    }
+}
+
+// Multiplies a byte by an array of bits
+template MulByte(){
+    signal input a;
+    signal input b[8];
+    signal output c[8];
+
+    for (var i = 0; i < 8; i++) {
+        c[i] <== a * b[i];
+    }
+}
+
+// XORs two bytes
+template XorByte(){
+        signal input a;
+        signal input b;
+        signal output out;
+
+        component abits = Num2Bits(8);
+        abits.in <== a;
+
+        component bbits = Num2Bits(8);
+        bbits.in <== b;
+
+        component XorBits = XorBits();
+        XorBits.a <== abits.out;
+        XorBits.b <== bbits.out;
+
+        component num = Bits2Num(8);
+        num.in <== XorBits.out;
+
+        out <== num.out;
+}
+
+// XORs two arrays of bits
+template XorBits(){
+        signal input a[8];
+        signal input b[8];
+        signal output out[8];
+
+    component xor[8];
+    for (var i = 0; i < 8; i++) {
+        xor[i] = XOR();
+        xor[i].a <== a[i];
+        xor[i].b <== b[i];
+        out[i] <== xor[i].out;
+    }
+}
\ No newline at end of file
diff --git a/inputs/cipher_4/default.json b/inputs/cipher_4/default.json
new file mode 100644
index 0000000..40b9734
--- /dev/null
+++ b/inputs/cipher_4/default.json
@@ -0,0 +1,9 @@
+{
+  "block": [
+    [50, 136, 49, 224],
+    [67, 90, 49, 55],
+    [246, 48, 152, 7],
+    [168, 141, 162, 52]
+  ],
+  "key": [43, 126, 21, 22, 40, 174, 210, 166, 171, 247, 21, 136, 9, 207, 79, 60]
+}
diff --git a/inputs/fibonacci_11/default.json b/inputs/fibonacci_11/default.json
deleted file mode 100644
index b465221..0000000
--- a/inputs/fibonacci_11/default.json
+++ /dev/null
@@ -1,3 +0,0 @@
-{
-  "in": [1, 1]
-}
diff --git a/package.json b/package.json
index 3231b60..d753bd2 100644
--- a/package.json
+++ b/package.json
@@ -3,7 +3,7 @@
   "scripts": {
     "start": "npx ts-node ./src/index.ts",
     "test": "npx mocha",
-    "compile:test": "npx circomkit compile fibonacci_11 && npx circomkit prove fibonacci_11 default && npx circomkit verify fibonacci_11 default"
+    "compile:test": "npx circomkit compile cipher_4 && npx circomkit prove cipher_4 default && npx circomkit verify cipher_4 default"
   },
   "dependencies": {
     "circomkit": "^0.0.22",
diff --git a/tests/cipher.test.ts b/tests/cipher.test.ts
new file mode 100644
index 0000000..a2dca0e
--- /dev/null
+++ b/tests/cipher.test.ts
@@ -0,0 +1,35 @@
+import { WitnessTester } from "circomkit";
+import { circomkit } from "./common";
+
+// todo: should debug cipher
+describe("Cipher", () => {
+  let circuit: WitnessTester<["block", "key"], ["cipher"]>;
+  it("should perform Cipher", async () => {
+    circuit = await circomkit.WitnessTester(`Cipher`, {
+      file: "cipher",
+      template: "Cipher",
+      params: [4],
+    });
+    console.log("@Cipher #constraints:", await circuit.getConstraintCount());
+
+    await circuit.expectPass(
+      {
+        block: [
+          [0x32, 0x88, 0x31, 0xe0],
+          [0x43, 0x5a, 0x31, 0x37],
+          [0xf6, 0x30, 0x98, 0x07],
+          [0xa8, 0x8d, 0xa2, 0x34],
+        ],
+        key: [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c],
+      },
+      {
+        cipher: [
+          [0x39, 0x02, 0xdc, 0x19],
+          [0x25, 0xdc, 0x11, 0x6a],
+          [0x84, 0x09, 0x85, 0x0b],
+          [0x1d, 0xfb, 0x97, 0x32],
+        ],
+      }
+    );
+  });
+});
diff --git a/tests/common/index.ts b/tests/common/index.ts
index cd613f6..4c2fd27 100644
--- a/tests/common/index.ts
+++ b/tests/common/index.ts
@@ -3,3 +3,11 @@ import { Circomkit } from "circomkit";
 export const circomkit = new Circomkit({
   verbose: false,
 });
+
+export const Num2Bits = (n: number) => {
+  let bits = [];
+  for (let i = 0; i < 8; i++) {
+    bits.push((n >> i) & 1);
+  }
+  return bits;
+};
diff --git a/tests/fibonacci.test.ts b/tests/fibonacci.test.ts
deleted file mode 100644
index 0fd10dc..0000000
--- a/tests/fibonacci.test.ts
+++ /dev/null
@@ -1,51 +0,0 @@
-import { WitnessTester } from "circomkit";
-import { circomkit } from "./common";
-
-describe("fibonacci", () => {
-  const N = 7;
-  let circuit: WitnessTester<["in"], ["out"]>;
-
-  describe("vanilla", () => {
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`fibonacci_${N}`, {
-        file: "fibonacci",
-        template: "Fibonacci",
-        params: [N],
-      });
-      console.log("#constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should compute correctly", async () => {
-      await circuit.expectPass({ in: [1, 1] }, { out: fibonacci([1, 1], N) });
-    });
-  });
-
-  describe("recursive", () => {
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`fibonacci_${N}_recursive`, {
-        file: "fibonacci",
-        template: "FibonacciRecursive",
-        params: [N],
-      });
-      console.log("#constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should compute correctly", async () => {
-      await circuit.expectPass({ in: [1, 1] }, { out: fibonacci([1, 1], N) });
-    });
-  });
-});
-
-// simple fibonacci with 2 variables
-function fibonacci(init: [number, number], n: number): number {
-  if (n < 0) {
-    throw new Error("N must be positive");
-  }
-
-  let [a, b] = init;
-  for (let i = 2; i <= n; i++) {
-    b = a + b;
-    a = b - a;
-  }
-  return n === 0 ? a : b;
-}
diff --git a/tests/key_expansion.test.ts b/tests/key_expansion.test.ts
index 2c47838..d1ed1f0 100644
--- a/tests/key_expansion.test.ts
+++ b/tests/key_expansion.test.ts
@@ -6,7 +6,7 @@ describe("KeyExpansion", () => {
     const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
       file: "key_expansion",
       template: "KeyExpansion",
-      params: [4, 10],
+      params: [4,10],
     });
     console.log("#constraints:", await circuit.getConstraintCount());
     const key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c];
@@ -62,7 +62,7 @@ describe("KeyExpansion", () => {
     const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
       file: "key_expansion",
       template: "KeyExpansion",
-      params: [6, 12],
+      params: [6,12],
     });
     console.log("#constraints:", await circuit.getConstraintCount());
     const key = [
@@ -132,7 +132,7 @@ describe("KeyExpansion", () => {
     const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
       file: "key_expansion",
       template: "KeyExpansion",
-      params: [8, 14],
+      params: [8,14],
     });
     console.log("#constraints:", await circuit.getConstraintCount());
     const key = [
diff --git a/tests/sbox128.test.ts b/tests/sbox128.test.ts
index 0c1e328..7582b81 100644
--- a/tests/sbox128.test.ts
+++ b/tests/sbox128.test.ts
@@ -8,7 +8,7 @@ describe("SBox128", () => {
     before(async () => {
       circuit = await circomkit.WitnessTester(`SubBytes`, {
         file: "sbox128",
-        template: "Sbox128",
+        template: "SBox128",
       });
       console.log("#constraints:", await circuit.getConstraintCount());
     });
diff --git a/tests/transformations.test.ts b/tests/transformations.test.ts
new file mode 100644
index 0000000..6db652e
--- /dev/null
+++ b/tests/transformations.test.ts
@@ -0,0 +1,360 @@
+import { WitnessTester } from "circomkit";
+import { circomkit } from "./common";
+
+describe("AES Key Expansion Components", () => {
+  describe("BytesToWords", () => {
+    let circuit: WitnessTester<["bytes"], ["words"]>;
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`BytesToWords`, {
+        file: "key_expansion",
+        template: "BytesToWords",
+        params: [4],
+      });
+      console.log("BytesToWords #constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should compute correctly", async () => {
+      await circuit.expectPass({ bytes: [0x01, 0x12, 0x02, 0x30] }, { words: [[0x01, 0x12, 0x02, 0x30]] });
+    });
+  });
+
+  describe("Rotate", () => {
+    let circuit: WitnessTester<["bytes"], ["rotated"]>;
+
+    it("should rotate correctly", async () => {
+      circuit = await circomkit.WitnessTester(`Rotate`, {
+        file: "key_expansion",
+        template: "Rotate",
+        params: [1, 4],
+      });
+      console.log("RotateWord #constraints:", await circuit.getConstraintCount());
+      await circuit.expectPass({ bytes: [0x01, 0x12, 0x02, 0x30] }, { rotated: [0x12, 0x02, 0x30, 0x01] });
+    });
+  });
+
+  describe("SubstituteWord", () => {
+    let circuit: WitnessTester<["bytes"], ["substituted"]>;
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`SubstituteWord`, {
+        file: "key_expansion",
+        template: "SubstituteWord",
+      });
+      console.log("SubstituteWord #constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should substitute correctly", async () => {
+      await circuit.expectPass({ bytes: [0x00, 0x10, 0x20, 0x30] }, { substituted: [0x63, 0xca, 0xb7, 0x04] });
+    });
+  });
+
+  describe("RCon", () => {
+    let circuit: WitnessTester<["round"], ["out"]>;
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`RCon`, {
+        file: "key_expansion",
+        template: "RCon",
+      });
+      console.log("RCon #constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should compute round constant correctly", async () => {
+      await circuit.expectPass({ round: 1 }, { out: [0x01, 0x00, 0x00, 0x00] });
+      await circuit.expectPass({ round: 2 }, { out: [0x02, 0x00, 0x00, 0x00] });
+      await circuit.expectPass({ round: 10 }, { out: [0x36, 0x00, 0x00, 0x00] });
+    });
+  });
+
+  describe("XorWord", () => {
+    let circuit: WitnessTester<["bytes1", "bytes2"], ["out"]>;
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`XorWord`, {
+        file: "key_expansion",
+        template: "XorWord",
+      });
+      console.log("XorWord #constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should XOR correctly", async () => {
+      await circuit.expectPass(
+        { bytes1: [0x0a, 0x0b, 0x0c, 0x0d], bytes2: [0x01, 0x02, 0x03, 0x04] },
+        { out: [0x0b, 0x09, 0x0f, 0x09] }
+      );
+    });
+  });
+
+  describe("WordSelector", () => {
+    let circuit: WitnessTester<["condition", "bytes1", "bytes2"], ["out"]>;
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`WordSelector`, {
+        file: "key_expansion",
+        template: "WordSelector",
+      });
+      console.log("WordSelector #constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should select first word when condition is 1", async () => {
+      await circuit.expectPass(
+        { condition: 1, bytes1: [0x0a, 0x0b, 0x0c, 0x0d], bytes2: [0x01, 0x02, 0x03, 0x04] },
+        { out: [0x0a, 0x0b, 0x0c, 0x0d] }
+      );
+    });
+
+    it("should select second word when condition is 0", async () => {
+      await circuit.expectPass(
+        { condition: 0, bytes1: [0x0a, 0x0b, 0x0c, 0x0d], bytes2: [0x01, 0x02, 0x03, 0x04] },
+        { out: [0x01, 0x02, 0x03, 0x04] }
+      );
+    });
+  });
+});
+
+describe("XTimes2", () => {
+  let circuit: WitnessTester<["in"], ["out"]>;
+  it("should perform 2 times", async () => {
+    circuit = await circomkit.WitnessTester(`XTimes2`, {
+      file: "cipher",
+      template: "XTimes2",
+    });
+    console.log("@XTimes2 #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 2 = 0xae
+    await circuit.expectPass({ in: [1, 1, 1, 0, 1, 0, 1, 0] }, { out: [0, 1, 1, 1, 0, 1, 0, 1] });
+    // 0x54 . 2 = 0xa8
+    await circuit.expectPass({ in: [0, 0, 1, 0, 1, 0, 1, 0] }, { out: [0, 0, 0, 1, 0, 1, 0, 1] });
+    // 0xae . 2 = 0x47
+    await circuit.expectPass({ in: [0, 1, 1, 1, 0, 1, 0, 1] }, { out: [1, 1, 1, 0, 0, 0, 1, 0] });
+    // 0x47 . 2 = 0x8e
+    await circuit.expectPass({ in: [1, 1, 1, 0, 0, 0, 1, 0] }, { out: [0, 1, 1, 1, 0, 0, 0, 1] });
+  });
+});
+describe("XTimes", () => {
+  let circuit: WitnessTester<["in"], ["out"]>;
+  it("should perform  xtimes", async () => {
+    circuit = await circomkit.WitnessTester(`XTimes`, {
+      file: "cipher",
+      template: "XTimes",
+      params: [0x13],
+    });
+    console.log("@XTimes2 #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 0x13 = 0xfe
+    await circuit.expectPass({ in: [1, 1, 1, 0, 1, 0, 1, 0] }, { out: [0, 1, 1, 1, 1, 1, 1, 1] });
+  });
+});
+
+describe("XTimes2 with XTimes", () => {
+  let circuit: WitnessTester<["in"], ["out"]>;
+  it("should perform 2 times with XTERMS", async () => {
+    circuit = await circomkit.WitnessTester(`XTimes`, {
+      file: "cipher",
+      template: "XTimes",
+      params: [0x2],
+    });
+    console.log("@XTimes2 #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 2 = 0xae
+    await circuit.expectPass({ in: [1, 1, 1, 0, 1, 0, 1, 0] }, { out: [0, 1, 1, 1, 0, 1, 0, 1] });
+    // 0x54 . 2 = 0xa8
+    await circuit.expectPass({ in: [0, 0, 1, 0, 1, 0, 1, 0] }, { out: [0, 0, 0, 1, 0, 1, 0, 1] });
+    // 0xae . 2 = 0x47
+    await circuit.expectPass({ in: [0, 1, 1, 1, 0, 1, 0, 1] }, { out: [1, 1, 1, 0, 0, 0, 1, 0] });
+    // 0x47 . 2 = 0x8e
+    await circuit.expectPass({ in: [1, 1, 1, 0, 0, 0, 1, 0] }, { out: [0, 1, 1, 1, 0, 0, 0, 1] });
+  });
+});
+
+describe("XTimes1 with XTimes", () => {
+  let circuit: WitnessTester<["in"], ["out"]>;
+  it("should perform 1 times with XTERMS", async () => {
+    circuit = await circomkit.WitnessTester(`XTimes`, {
+      file: "cipher",
+      template: "XTimes",
+      params: [0x1],
+    });
+    console.log("@XTimes2 #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 2 = 0xae
+    await circuit.expectPass({ in: [1, 1, 1, 0, 1, 0, 1, 0] }, { out: [1, 1, 1, 0, 1, 0, 1, 0] });
+    // 0x54 . 2 = 0xa8
+    await circuit.expectPass({ in: [0, 0, 1, 0, 1, 0, 1, 0] }, { out: [0, 0, 1, 0, 1, 0, 1, 0] });
+    // 0xae . 2 = 0x47
+    await circuit.expectPass({ in: [0, 1, 1, 1, 0, 1, 0, 1] }, { out: [0, 1, 1, 1, 0, 1, 0, 1] });
+    // 0x47 . 2 = 0x8e
+    await circuit.expectPass({ in: [1, 1, 1, 0, 1, 0, 1, 0] }, { out: [1, 1, 1, 0, 1, 0, 1, 0] });
+  });
+});
+
+describe("MixColumns", () => {
+  it("s0 should compute correctly", async () => {
+    let circuit: WitnessTester<["in"], ["out"]>;
+    circuit = await circomkit.WitnessTester(`s0`, {
+      file: "cipher",
+      template: "S0",
+      params: [],
+    });
+    console.log("@S0 #constraints:", await circuit.getConstraintCount());
+
+    await circuit.expectPass({ in: [0xd4, 0xbf, 0x5d, 0x30] }, { out: 0x04 });
+  });
+
+  it("s1 should compute correctly", async () => {
+    let circuit: WitnessTester<["in"], ["out"]>;
+    circuit = await circomkit.WitnessTester(`s1`, {
+      file: "cipher",
+      template: "S1",
+      params: [],
+    });
+    console.log("@S1 #constraints:", await circuit.getConstraintCount());
+
+    await circuit.expectPass({ in: [0xd4, 0xbf, 0x5d, 0x30] }, { out: 0x66 });
+  });
+
+  it("s2 should compute correctly", async () => {
+    let circuit: WitnessTester<["in"], ["out"]>;
+    circuit = await circomkit.WitnessTester(`s2`, {
+      file: "cipher",
+      template: "S2",
+      params: [],
+    });
+    console.log("@S2 #constraints:", await circuit.getConstraintCount());
+
+    await circuit.expectPass({ in: [0xd4, 0xbf, 0x5d, 0x30] }, { out: 0x81 });
+  });
+
+  it("s3 should compute correctly", async () => {
+    let circuit: WitnessTester<["in"], ["out"]>;
+    circuit = await circomkit.WitnessTester(`s3`, {
+      file: "cipher",
+      template: "S3",
+      params: [],
+    });
+    console.log("@S3 #constraints:", await circuit.getConstraintCount());
+
+    await circuit.expectPass({ in: [0xd4, 0xbf, 0x5d, 0x30] }, { out: 0xe5 });
+  });
+
+  it("s4 should compute correctly", async () => {
+    let circuit: WitnessTester<["state"], ["out"]>;
+    circuit = await circomkit.WitnessTester(`MixColumns`, {
+      file: "cipher",
+      template: "MixColumns",
+      params: [],
+    });
+    console.log("@MixColumns #constraints:", await circuit.getConstraintCount());
+    const state = [
+      [0xd4, 0xe0, 0xb8, 0x1e],
+      [0xbf, 0xb4, 0x41, 0x27],
+      [0x5d, 0x52, 0x11, 0x98],
+      [0x30, 0xae, 0xf1, 0xe5],
+    ];
+
+    const out = [
+      [0x04, 0xe0, 0x48, 0x28],
+      [0x66, 0xcb, 0xf8, 0x06],
+      [0x81, 0x19, 0xd3, 0x26],
+      [0xe5, 0x9a, 0x7a, 0x4c],
+    ];
+
+    await circuit.expectPass({ state }, { out });
+  });
+});
+
+describe("AddRoundKey", () => {
+  let circuit: WitnessTester<["state", "roundKey"], ["newState"]>;
+  it("should perform AddRoundKey", async () => {
+    circuit = await circomkit.WitnessTester(`AddRoundKey`, {
+      file: "cipher",
+      template: "AddRoundKey",
+    });
+    console.log("@AddRoundKey #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 2 = 0xae
+    await circuit.expectPass(
+      {
+        state: [
+          [4, 224, 72, 40],
+          [102, 203, 248, 6],
+          [129, 25, 211, 38],
+          [229, 154, 122, 76],
+        ],
+        roundKey: [
+          [160, 250, 254, 23],
+          [136, 84, 44, 177],
+          [35, 163, 57, 57],
+          [42, 108, 118, 5],
+        ],
+      },
+      {
+        newState: [
+          [164, 104, 107, 2],
+          [156, 159, 91, 106],
+          [127, 53, 234, 80],
+          [242, 43, 67, 73],
+        ],
+      }
+    );
+  });
+});
+
+describe("SubBlock", () => {
+  let circuit: WitnessTester<["state"], ["newState"]>;
+  it("should perform SubBlock", async () => {
+    circuit = await circomkit.WitnessTester(`SubBlock`, {
+      file: "cipher",
+      template: "SubBlock",
+    });
+    console.log("@SubBlock #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 2 = 0xae
+    await circuit.expectPass(
+      {
+        state: [
+          [25, 160, 154, 233],
+          [61, 244, 198, 248],
+          [227, 226, 141, 72],
+          [190, 43, 42, 8],
+        ],
+      },
+      {
+        newState: [
+          [212, 224, 184, 30],
+          [39, 191, 180, 65],
+          [17, 152, 93, 82],
+          [174, 241, 229, 48],
+        ],
+      }
+    );
+  });
+});
+
+describe("ShiftRows", () => {
+  let circuit: WitnessTester<["state"], ["newState"]>;
+  it("should perform ShiftRows", async () => {
+    circuit = await circomkit.WitnessTester(`ShiftRows`, {
+      file: "cipher",
+      template: "ShiftRows",
+      params: [],
+    });
+    console.log("@ShiftRows #constraints:", await circuit.getConstraintCount());
+
+    // 0x57 . 2 = 0xae
+    await circuit.expectPass(
+      {
+        state: [
+          [212, 224, 184, 30],
+          [39, 191, 180, 65],
+          [17, 152, 93, 82],
+          [174, 241, 229, 48],
+        ],
+      },
+      {
+        newState: [
+          [212, 224, 184, 30],
+          [191, 180, 65, 39],
+          [93, 82, 17, 152],
+          [48, 174, 241, 229],
+        ],
+      }
+    );
+  });
+});
diff --git a/tests/utils.test.ts b/tests/utils.test.ts
deleted file mode 100644
index 9d37dcb..0000000
--- a/tests/utils.test.ts
+++ /dev/null
@@ -1,111 +0,0 @@
-import { WitnessTester } from "circomkit";
-import { circomkit } from "./common";
-
-describe("AES Key Expansion Components", () => {
-  describe("BytesToWords", () => {
-    let circuit: WitnessTester<["bytes"], ["words"]>;
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`BytesToWords`, {
-        file: "key_expansion",
-        template: "BytesToWords",
-        params: [4],
-      });
-      console.log("BytesToWords #constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should compute correctly", async () => {
-      await circuit.expectPass({ bytes: [0x01, 0x12, 0x02, 0x30] }, { words: [[0x01, 0x12, 0x02, 0x30]] });
-    });
-  });
-
-  describe("RotateWord", () => {
-    let circuit: WitnessTester<["bytes"], ["rotated"]>;
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`RotateWord`, {
-        file: "key_expansion",
-        template: "RotateWord",
-        params: [1],
-      });
-      console.log("RotateWord #constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should rotate correctly", async () => {
-      await circuit.expectPass({ bytes: [0x01, 0x12, 0x02, 0x30] }, { rotated: [0x12, 0x02, 0x30, 0x01] });
-    });
-  });
-
-  describe("SubstituteWord", () => {
-    let circuit: WitnessTester<["bytes"], ["substituted"]>;
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`SubstituteWord`, {
-        file: "key_expansion",
-        template: "SubstituteWord",
-      });
-      console.log("SubstituteWord #constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should substitute correctly", async () => {
-      await circuit.expectPass({ bytes: [0x00, 0x10, 0x20, 0x30] }, { substituted: [0x63, 0xca, 0xb7, 0x04] });
-    });
-  });
-
-  describe("RCon", () => {
-    let circuit: WitnessTester<["round"], ["out"]>;
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`RCon`, {
-        file: "key_expansion",
-        template: "RCon",
-      });
-      console.log("RCon #constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should compute round constant correctly", async () => {
-      await circuit.expectPass({ round: 1 }, { out: [0x01, 0x00, 0x00, 0x00] });
-      await circuit.expectPass({ round: 2 }, { out: [0x02, 0x00, 0x00, 0x00] });
-      await circuit.expectPass({ round: 10 }, { out: [0x36, 0x00, 0x00, 0x00] });
-    });
-  });
-
-  describe("XorWord", () => {
-    let circuit: WitnessTester<["bytes1", "bytes2"], ["out"]>;
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`XorWord`, {
-        file: "key_expansion",
-        template: "XorWord",
-      });
-      console.log("XorWord #constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should XOR correctly", async () => {
-      await circuit.expectPass(
-        { bytes1: [0x0a, 0x0b, 0x0c, 0x0d], bytes2: [0x01, 0x02, 0x03, 0x04] },
-        { out: [0x0b, 0x09, 0x0f, 0x09] }
-      );
-    });
-  });
-
-  describe("WordSelector", () => {
-    let circuit: WitnessTester<["condition", "bytes1", "bytes2"], ["out"]>;
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`WordSelector`, {
-        file: "key_expansion",
-        template: "WordSelector",
-      });
-      console.log("WordSelector #constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should select first word when condition is 1", async () => {
-      await circuit.expectPass(
-        { condition: 1, bytes1: [0x0a, 0x0b, 0x0c, 0x0d], bytes2: [0x01, 0x02, 0x03, 0x04] },
-        { out: [0x0a, 0x0b, 0x0c, 0x0d] }
-      );
-    });
-
-    it("should select second word when condition is 0", async () => {
-      await circuit.expectPass(
-        { condition: 0, bytes1: [0x0a, 0x0b, 0x0c, 0x0d], bytes2: [0x01, 0x02, 0x03, 0x04] },
-        { out: [0x01, 0x02, 0x03, 0x04] }
-      );
-    });
-  });
-});