diff --git a/circuits/aes_128.circom b/circuits/aes_128.circom
deleted file mode 100644
index e69de29..0000000
diff --git a/circuits/key_expansion.circom b/circuits/key_expansion.circom
new file mode 100644
index 0000000..d395bb9
--- /dev/null
+++ b/circuits/key_expansion.circom
@@ -0,0 +1,208 @@
+pragma circom 2.0.0;
+
+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) {
+    signal input key[nk * 4];
+    
+    var totalWords = (4 * (nr + 1));
+    var effectiveRounds = (totalWords % nk == 0 ? totalWords - nk : totalWords) \ nk;
+    var leftoverWords = totalWords - (effectiveRounds * nk);
+
+    signal output keyExpanded[totalWords][4];
+
+    for (var i = 0; i < nk; i++) {
+        for (var j = 0; j < 4; j++) {
+            keyExpanded[i][j] <== key[(4 * i) + j];
+        }
+    }
+    
+    component nextRound[effectiveRounds];
+    
+    for (var round = 1; round <= effectiveRounds; round++) {
+        var outputWordLen = round == effectiveRounds ? leftoverWords : nk;
+        nextRound[round - 1] = NextRound(nk, outputWordLen);
+
+        for (var i = 0; i < nk; i++) {
+            for (var j = 0; j < 4; j++) {
+                nextRound[round - 1].key[i][j] <== keyExpanded[(round * nk) + i - nk][j];
+            }
+        }
+
+        nextRound[round - 1].round <== round;
+
+        for (var i = 0; i < outputWordLen; i++) {
+            for (var j = 0; j < 4; j++) {
+                keyExpanded[(round * nk) + i][j] <== nextRound[round - 1].nextKey[i][j];
+            }
+        }
+    }
+}
+
+//nk, output
+template NextRound(nk, o){
+    signal input key[nk][4]; 
+    signal input round;
+    signal output nextKey[o][4];
+
+    component rotateWord = RotateWord(1);
+    for (var i = 0; i < 4; i++) {
+        rotateWord.bytes[i] <== key[nk - 1][i];
+    }
+    
+    component substituteWord[2];
+    substituteWord[0] = SubstituteWord();
+    substituteWord[0].bytes <== rotateWord.rotated;
+
+    component rcon = RCon();
+    rcon.round <== round; 
+
+    component xorWord[o + 1];
+    xorWord[0] = XorWord();
+    xorWord[0].bytes1 <== substituteWord[0].substituted;
+    xorWord[0].bytes2 <== rcon.out;
+
+    for (var i = 0; i < o; i++) {
+        xorWord[i+1] = XorWord();
+        if (i == 0) {
+            xorWord[i+1].bytes1 <== xorWord[0].out;
+        } else if(nk == 8 && i == 4) {
+            substituteWord[1] = SubstituteWord();
+            substituteWord[1].bytes <== nextKey[i - 1];
+            xorWord[i+1].bytes1 <== substituteWord[1].substituted;
+        } else {
+            xorWord[i+1].bytes1 <== nextKey[i-1];
+        }
+        xorWord[i+1].bytes2 <== key[i];
+        
+        for (var j = 0; j < 4; j++) {
+            nextKey[i][j] <== xorWord[i+1].out[j];
+        }
+    }
+}
+
+
+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/sbox128.circom b/circuits/sbox128.circom
new file mode 100644
index 0000000..e11e55e
--- /dev/null
+++ b/circuits/sbox128.circom
@@ -0,0 +1,29 @@
+pragma circom 2.0.0;
+
+include "circomlib/circuits/comparators.circom";
+
+template Sbox128() {
+    signal input in;
+    signal output out;
+
+    var sbox[16 * 16] = [
+        0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
+        0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
+        0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
+        0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
+        0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
+        0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
+        0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
+        0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
+        0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
+        0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
+        0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
+        0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
+        0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
+        0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
+        0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
+        0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
+    ];
+    
+    out <-- sbox[in];
+}
diff --git a/circuits/subbytes.circom b/circuits/subbytes.circom
deleted file mode 100644
index 32dc96f..0000000
--- a/circuits/subbytes.circom
+++ /dev/null
@@ -1,9 +0,0 @@
-pragma circom 2.0.0;
-
-template SubBytes(n) {
-    signal input key[n];
-    log("key" , key[0]);
-
-    signal output out;
-    out <==1;
-}
\ No newline at end of file
diff --git a/tests/key_expansion.test.ts b/tests/key_expansion.test.ts
new file mode 100644
index 0000000..2c47838
--- /dev/null
+++ b/tests/key_expansion.test.ts
@@ -0,0 +1,208 @@
+import { WitnessTester } from "circomkit";
+import { circomkit } from "./common";
+
+describe("KeyExpansion", () => {
+  it("should compute correctly for aes128", async () => {
+    const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
+      file: "key_expansion",
+      template: "KeyExpansion",
+      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];
+    const keyExpanded = [
+      [0x2b, 0x7e, 0x15, 0x16],
+      [0x28, 0xae, 0xd2, 0xa6],
+      [0xab, 0xf7, 0x15, 0x88],
+      [0x09, 0xcf, 0x4f, 0x3c],
+      [0xa0, 0xfa, 0xfe, 0x17],
+      [0x88, 0x54, 0x2c, 0xb1],
+      [0x23, 0xa3, 0x39, 0x39],
+      [0x2a, 0x6c, 0x76, 0x05],
+      [0xf2, 0xc2, 0x95, 0xf2],
+      [0x7a, 0x96, 0xb9, 0x43],
+      [0x59, 0x35, 0x80, 0x7a],
+      [0x73, 0x59, 0xf6, 0x7f],
+      [0x3d, 0x80, 0x47, 0x7d],
+      [0x47, 0x16, 0xfe, 0x3e],
+      [0x1e, 0x23, 0x7e, 0x44],
+      [0x6d, 0x7a, 0x88, 0x3b],
+      [0xef, 0x44, 0xa5, 0x41],
+      [0xa8, 0x52, 0x5b, 0x7f],
+      [0xb6, 0x71, 0x25, 0x3b],
+      [0xdb, 0x0b, 0xad, 0x00],
+      [0xd4, 0xd1, 0xc6, 0xf8],
+      [0x7c, 0x83, 0x9d, 0x87],
+      [0xca, 0xf2, 0xb8, 0xbc],
+      [0x11, 0xf9, 0x15, 0xbc],
+      [0x6d, 0x88, 0xa3, 0x7a],
+      [0x11, 0x0b, 0x3e, 0xfd],
+      [0xdb, 0xf9, 0x86, 0x41],
+      [0xca, 0x00, 0x93, 0xfd],
+      [0x4e, 0x54, 0xf7, 0x0e],
+      [0x5f, 0x5f, 0xc9, 0xf3],
+      [0x84, 0xa6, 0x4f, 0xb2],
+      [0x4e, 0xa6, 0xdc, 0x4f],
+      [0xea, 0xd2, 0x73, 0x21],
+      [0xb5, 0x8d, 0xba, 0xd2],
+      [0x31, 0x2b, 0xf5, 0x60],
+      [0x7f, 0x8d, 0x29, 0x2f],
+      [0xac, 0x77, 0x66, 0xf3],
+      [0x19, 0xfa, 0xdc, 0x21],
+      [0x28, 0xd1, 0x29, 0x41],
+      [0x57, 0x5c, 0x00, 0x6e],
+      [0xd0, 0x14, 0xf9, 0xa8],
+      [0xc9, 0xee, 0x25, 0x89],
+      [0xe1, 0x3f, 0x0c, 0xc8],
+      [0xb6, 0x63, 0x0c, 0xa6],
+    ];
+    await circuit.expectPass({ key }, { keyExpanded });
+  });
+  it("should compute correctly for aes 192", async () => {
+    const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
+      file: "key_expansion",
+      template: "KeyExpansion",
+      params: [6, 12],
+    });
+    console.log("#constraints:", await circuit.getConstraintCount());
+    const key = [
+      0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90, 0x79, 0xe5, 0x62, 0xf8, 0xea,
+      0xd2, 0x52, 0x2c, 0x6b, 0x7b,
+    ];
+
+    const keyExpanded = [
+      [0x8e, 0x73, 0xb0, 0xf7],
+      [0xda, 0x0e, 0x64, 0x52],
+      [0xc8, 0x10, 0xf3, 0x2b],
+      [0x80, 0x90, 0x79, 0xe5],
+      [0x62, 0xf8, 0xea, 0xd2],
+      [0x52, 0x2c, 0x6b, 0x7b],
+      [0xfe, 0x0c, 0x91, 0xf7],
+      [0x24, 0x02, 0xf5, 0xa5],
+      [0xec, 0x12, 0x06, 0x8e],
+      [0x6c, 0x82, 0x7f, 0x6b],
+      [0x0e, 0x7a, 0x95, 0xb9],
+      [0x5c, 0x56, 0xfe, 0xc2],
+      [0x4d, 0xb7, 0xb4, 0xbd],
+      [0x69, 0xb5, 0x41, 0x18],
+      [0x85, 0xa7, 0x47, 0x96],
+      [0xe9, 0x25, 0x38, 0xfd],
+      [0xe7, 0x5f, 0xad, 0x44],
+      [0xbb, 0x09, 0x53, 0x86],
+      [0x48, 0x5a, 0xf0, 0x57],
+      [0x21, 0xef, 0xb1, 0x4f],
+      [0xa4, 0x48, 0xf6, 0xd9],
+      [0x4d, 0x6d, 0xce, 0x24],
+      [0xaa, 0x32, 0x63, 0x60],
+      [0x11, 0x3b, 0x30, 0xe6],
+      [0xa2, 0x5e, 0x7e, 0xd5],
+      [0x83, 0xb1, 0xcf, 0x9a],
+      [0x27, 0xf9, 0x39, 0x43],
+      [0x6a, 0x94, 0xf7, 0x67],
+      [0xc0, 0xa6, 0x94, 0x07],
+      [0xd1, 0x9d, 0xa4, 0xe1],
+      [0xec, 0x17, 0x86, 0xeb],
+      [0x6f, 0xa6, 0x49, 0x71],
+      [0x48, 0x5f, 0x70, 0x32],
+      [0x22, 0xcb, 0x87, 0x55],
+      [0xe2, 0x6d, 0x13, 0x52],
+      [0x33, 0xf0, 0xb7, 0xb3],
+      [0x40, 0xbe, 0xeb, 0x28],
+      [0x2f, 0x18, 0xa2, 0x59],
+      [0x67, 0x47, 0xd2, 0x6b],
+      [0x45, 0x8c, 0x55, 0x3e],
+      [0xa7, 0xe1, 0x46, 0x6c],
+      [0x94, 0x11, 0xf1, 0xdf],
+      [0x82, 0x1f, 0x75, 0x0a],
+      [0xad, 0x07, 0xd7, 0x53],
+      [0xca, 0x40, 0x05, 0x38],
+      [0x8f, 0xcc, 0x50, 0x06],
+      [0x28, 0x2d, 0x16, 0x6a],
+      [0xbc, 0x3c, 0xe7, 0xb5],
+      [0xe9, 0x8b, 0xa0, 0x6f],
+      [0x44, 0x8c, 0x77, 0x3c],
+      [0x8e, 0xcc, 0x72, 0x04],
+      [0x01, 0x00, 0x22, 0x02],
+    ];
+
+    await circuit.expectPass({ key }, { keyExpanded });
+  });
+
+  it("should compute correctly for aes256", async () => {
+    const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
+      file: "key_expansion",
+      template: "KeyExpansion",
+      params: [8, 14],
+    });
+    console.log("#constraints:", await circuit.getConstraintCount());
+    const key = [
+      0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d, 0x77, 0x81, 0x1f, 0x35, 0x2c,
+      0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3, 0x09, 0x14, 0xdf, 0xf4,
+    ];
+
+    const keyExpanded = [
+      [0x60, 0x3d, 0xeb, 0x10],
+      [0x15, 0xca, 0x71, 0xbe],
+      [0x2b, 0x73, 0xae, 0xf0],
+      [0x85, 0x7d, 0x77, 0x81],
+      [0x1f, 0x35, 0x2c, 0x07],
+      [0x3b, 0x61, 0x08, 0xd7],
+      [0x2d, 0x98, 0x10, 0xa3],
+      [0x09, 0x14, 0xdf, 0xf4],
+      [0x9b, 0xa3, 0x54, 0x11],
+      [0x8e, 0x69, 0x25, 0xaf],
+      [0xa5, 0x1a, 0x8b, 0x5f],
+      [0x20, 0x67, 0xfc, 0xde],
+      [0xa8, 0xb0, 0x9c, 0x1a],
+      [0x93, 0xd1, 0x94, 0xcd],
+      [0xbe, 0x49, 0x84, 0x6e],
+      [0xb7, 0x5d, 0x5b, 0x9a],
+      [0xd5, 0x9a, 0xec, 0xb8],
+      [0x5b, 0xf3, 0xc9, 0x17],
+      [0xfe, 0xe9, 0x42, 0x48],
+      [0xde, 0x8e, 0xbe, 0x96],
+      [0xb5, 0xa9, 0x32, 0x8a],
+      [0x26, 0x78, 0xa6, 0x47],
+      [0x98, 0x31, 0x22, 0x29],
+      [0x2f, 0x6c, 0x79, 0xb3],
+      [0x81, 0x2c, 0x81, 0xad],
+      [0xda, 0xdf, 0x48, 0xba],
+      [0x24, 0x36, 0x0a, 0xf2],
+      [0xfa, 0xb8, 0xb4, 0x64],
+      [0x98, 0xc5, 0xbf, 0xc9],
+      [0xbe, 0xbd, 0x19, 0x8e],
+      [0x26, 0x8c, 0x3b, 0xa7],
+      [0x09, 0xe0, 0x42, 0x14],
+      [0x68, 0x00, 0x7b, 0xac],
+      [0xb2, 0xdf, 0x33, 0x16],
+      [0x96, 0xe9, 0x39, 0xe4],
+      [0x6c, 0x51, 0x8d, 0x80],
+      [0xc8, 0x14, 0xe2, 0x04],
+      [0x76, 0xa9, 0xfb, 0x8a],
+      [0x50, 0x25, 0xc0, 0x2d],
+      [0x59, 0xc5, 0x82, 0x39],
+      [0xde, 0x13, 0x69, 0x67],
+      [0x6c, 0xcc, 0x5a, 0x71],
+      [0xfa, 0x25, 0x63, 0x95],
+      [0x96, 0x74, 0xee, 0x15],
+      [0x58, 0x86, 0xca, 0x5d],
+      [0x2e, 0x2f, 0x31, 0xd7],
+      [0x7e, 0x0a, 0xf1, 0xfa],
+      [0x27, 0xcf, 0x73, 0xc3],
+      [0x74, 0x9c, 0x47, 0xab],
+      [0x18, 0x50, 0x1d, 0xda],
+      [0xe2, 0x75, 0x7e, 0x4f],
+      [0x74, 0x01, 0x90, 0x5a],
+      [0xca, 0xfa, 0xaa, 0xe3],
+      [0xe4, 0xd5, 0x9b, 0x34],
+      [0x9a, 0xdf, 0x6a, 0xce],
+      [0xbd, 0x10, 0x19, 0x0d],
+      [0xfe, 0x48, 0x90, 0xd1],
+      [0xe6, 0x18, 0x8d, 0x0b],
+      [0x04, 0x6d, 0xf3, 0x44],
+      [0x70, 0x6c, 0x63, 0x1e],
+    ];
+
+    await circuit.expectPass({ key }, { keyExpanded });
+  });
+});
diff --git a/tests/next_round.test.ts b/tests/next_round.test.ts
new file mode 100644
index 0000000..fe68402
--- /dev/null
+++ b/tests/next_round.test.ts
@@ -0,0 +1,64 @@
+import { WitnessTester } from "circomkit";
+import { circomkit } from "./common";
+
+describe("NextRound", () => {
+  let circuit: WitnessTester<["key", "round"], ["nextKey"]>;
+
+  describe("NextRound", () => {
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`NextRound_${4}_${4}`, {
+        file: "key_expansion",
+        template: "NextRound",
+        params: [4, 4],
+      });
+      console.log("#constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should compute correctly", async () => {
+      const key = [
+        [0x2b, 0x7e, 0x15, 0x16],
+        [0x28, 0xae, 0xd2, 0xa6],
+        [0xab, 0xf7, 0x15, 0x88],
+        [0x09, 0xcf, 0x4f, 0x3c],
+      ];
+
+      const expectedNextKey = [
+        [0xa0, 0xfa, 0xfe, 0x17],
+        [0x88, 0x54, 0x2c, 0xb1],
+        [0x23, 0xa3, 0x39, 0x39],
+        [0x2a, 0x6c, 0x76, 0x05],
+      ];
+
+      await circuit.expectPass({ key, round: 1 }, { nextKey: expectedNextKey });
+    });
+  });
+
+  // describe("NextRound", () => {
+  //   before(async () => {
+  //     circuit = await circomkit.WitnessTester(`NextRound_${6}_${6}`, {
+  //       file: "key_expansion",
+  //       template: "NextRound",
+  //       params: [6, 6],
+  //     });
+  //     console.log("#constraints:", await circuit.getConstraintCount());
+  //   });
+
+  //   it("should compute correctly for AES-192", async () => {
+  //     const key = [
+  //       [0x8e, 0x73, 0xb0, 0xf7],
+  //       [0xda, 0x0e, 0x64, 0x52],
+  //       [0xc8, 0x10, 0xf3, 0x2b],
+  //       [0x80, 0x90, 0x79, 0xe5],
+  //     ];
+
+  //     const expectedNextKey = [
+  //       [0x62, 0xf8, 0xea, 0xd2],
+  //       [0x52, 0x2c, 0x6b, 0x7b],
+  //       [0xfe, 0x0c, 0x91, 0xf7],
+  //       [0x24, 0x02, 0xf5, 0xa5],
+  //     ];
+
+  //     await circuit.expectPass({ key, round: 1 }, { nextKey: expectedNextKey });
+  //   });
+  // });
+});
diff --git a/tests/sbox128.test.ts b/tests/sbox128.test.ts
new file mode 100644
index 0000000..0c1e328
--- /dev/null
+++ b/tests/sbox128.test.ts
@@ -0,0 +1,21 @@
+import { WitnessTester } from "circomkit";
+import { circomkit } from "./common";
+
+describe("SBox128", () => {
+  let circuit: WitnessTester<["in"], ["out"]>;
+
+  describe("SubBox", () => {
+    before(async () => {
+      circuit = await circomkit.WitnessTester(`SubBytes`, {
+        file: "sbox128",
+        template: "Sbox128",
+      });
+      console.log("#constraints:", await circuit.getConstraintCount());
+    });
+
+    it("should compute correctly", async () => {
+      await circuit.expectPass({ in: 0x53 }, { out: 0xed });
+      await circuit.expectPass({ in: 0x00 }, { out: 0x63 });
+    });
+  });
+});
diff --git a/tests/subbytes.test.ts b/tests/subbytes.test.ts
deleted file mode 100644
index c71eb6b..0000000
--- a/tests/subbytes.test.ts
+++ /dev/null
@@ -1,23 +0,0 @@
-import { WitnessTester } from "circomkit";
-import { circomkit } from "./common";
-
-describe("SubBytes", () => {
-  const N = 2;
-  let circuit: WitnessTester<["key"], ["out"]>;
-
-  describe("vanilla", () => {
-    before(async () => {
-      circuit = await circomkit.WitnessTester(`SubBytes_${N}`, {
-        file: "subbytes",
-        template: "SubBytes",
-        params: [N],
-      });
-      console.log("#constraints:", await circuit.getConstraintCount());
-    });
-
-    it("should compute correctly", async () => {
-      const witness = await circuit.calculateWitness({ key: [1, 1] });
-      await circuit.expectPass({ key: [1, 1] }, { out: 1 });
-    });
-  });
-});
diff --git a/tests/utils.test.ts b/tests/utils.test.ts
new file mode 100644
index 0000000..9d37dcb
--- /dev/null
+++ b/tests/utils.test.ts
@@ -0,0 +1,111 @@
+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] }
+      );
+    });
+  });
+});