diff --git a/evm/src/cpu/kernel/asm/account_code.asm b/evm/src/cpu/kernel/asm/account_code.asm
index 7d560ec1..d40e64ef 100644
--- a/evm/src/cpu/kernel/asm/account_code.asm
+++ b/evm/src/cpu/kernel/asm/account_code.asm
@@ -10,6 +10,7 @@
 
 
 %macro codesize
+    // stack: (empty)
     ADDRESS
     %extcodesize
 %endmacro
@@ -34,15 +35,21 @@ global extcodesize:
     %jump(extcodecopy)
 %endmacro
 
+// Pre stack: address, dest_offset, offset, size, retdest
+// Post stack: (empty)
 global extcodecopy:
     // stack: address, dest_offset, offset, size, retdest
     %stack (address, dest_offset, offset, size, retdest) -> (address, extcodecopy_contd, size, offset, dest_offset, retdest)
     %jump(load_code)
+
 extcodecopy_contd:
     // stack: code_length, size, offset, dest_offset, retdest
     SWAP1
     // stack: size, code_length, offset, dest_offset, retdest
     PUSH 0
+
+// Loop copying the `code[offset]` to `memory[dest_offset]` until `i==size`.
+// Each iteration increments `offset, dest_offset, i`.
 extcodecopy_loop:
     // stack: i, size, code_length, offset, dest_offset, retdest
     DUP2 DUP2 EQ
@@ -56,6 +63,7 @@ extcodecopy_loop:
     %mload_current(@SEGMENT_KERNEL_ACCOUNT_CODE)
     // stack: opcode, offset < code_length, offset, code_length, dest_offset, i, size, retdest
     %stack (opcode, offset_lt_code_length, offset, code_length, dest_offset, i, size, retdest) -> (offset_lt_code_length, 0, opcode, offset, code_length, dest_offset, i, size, retdest)
+    // If `offset >= code_length`, use `opcode=0`. Necessary since `SEGMENT_KERNEL_ACCOUNT_CODE` might be clobbered from previous calls.
     %select_bool
     // stack: opcode, offset, code_length, dest_offset, i, size, retdest
     DUP4
@@ -91,6 +99,8 @@ load_code:
     PROVER_INPUT(account_code::length)
     // stack: code_length, codehash, retdest
     PUSH 0
+
+// Loop non-deterministically querying `code[i]` and storing it in `SEGMENT_KERNEL_ACCOUNT_CODE` at offset `i`, until `i==code_length`.
 load_code_loop:
     // stack: i, code_length, codehash, retdest
     DUP2 DUP2 EQ
@@ -106,15 +116,12 @@ load_code_loop:
     // stack: i+1, code_length, codehash, retdest
     %jump(load_code_loop)
 
+// Check that the hash of the loaded code equals `codehash`.
 load_code_check:
     // stack: i, code_length, codehash, retdest
     POP
     // stack: code_length, codehash, retdest
-    %stack (code_length, codehash, retdest) -> (code_length, codehash, retdest, code_length)
-    PUSH 0
-    // stack: 0, code_length, codehash, retdest, code_length
-    PUSH @SEGMENT_KERNEL_ACCOUNT_CODE
-    // stack: segment, 0, code_length, codehash, retdest, code_length
+    %stack (code_length, codehash, retdest) -> (@SEGMENT_KERNEL_ACCOUNT_CODE, 0, code_length, codehash, retdest, code_length)
     GET_CONTEXT
     // stack: context, segment, 0, code_length, codehash, retdest, code_length
     KECCAK_GENERAL
diff --git a/evm/src/cpu/kernel/tests/account_code.rs b/evm/src/cpu/kernel/tests/account_code.rs
index 0576e0e4..ae9fc30d 100644
--- a/evm/src/cpu/kernel/tests/account_code.rs
+++ b/evm/src/cpu/kernel/tests/account_code.rs
@@ -13,6 +13,7 @@ use crate::cpu::kernel::tests::mpt::nibbles_64;
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::memory::segments::Segment;
 
+// Test account with a given code hash.
 fn test_account(code: &[u8]) -> AccountRlp {
     AccountRlp {
         nonce: U256::from(1111),
@@ -140,14 +141,18 @@ fn test_extcodecopy() -> Result<()> {
 
     let extcodecopy = KERNEL.global_labels["extcodecopy"];
 
-    // Put random data in main memory.
+    // Put random data in main memory and the `KernelAccountCode` segment for realism.
     let mut rng = thread_rng();
     for i in 0..2000 {
         interpreter.memory.context_memory[interpreter.context].segments
             [Segment::MainMemory as usize]
             .set(i, U256::from(rng.gen::<u8>()));
+        interpreter.memory.context_memory[interpreter.context].segments
+            [Segment::KernelAccountCode as usize]
+            .set(i, U256::from(rng.gen::<u8>()));
     }
 
+    // Random inputs
     let dest_offset = rng.gen_range(0..3000);
     let offset = rng.gen_range(0..1500);
     let size = rng.gen_range(0..1500);
diff --git a/evm/src/generation/prover_input.rs b/evm/src/generation/prover_input.rs
index 82a0c7c3..ad1cfce0 100644
--- a/evm/src/generation/prover_input.rs
+++ b/evm/src/generation/prover_input.rs
@@ -65,23 +65,26 @@ impl<F: Field> GenerationState<F> {
             .unwrap_or_else(|| panic!("Out of RLP data"))
     }
 
-    /// Account code
+    /// Account code.
     fn run_account_code(&mut self, stack: &[U256], input_fn: &ProverInputFn) -> U256 {
         match input_fn.0[1].as_str() {
             "length" => {
+                // Return length of code.
+                // stack: codehash
                 let codehash = stack.last().expect("Empty stack");
                 self.inputs.contract_code[&H256::from_uint(codehash)]
                     .len()
                     .into()
             }
             "get" => {
+                // Return `code[i]`.
+                // stack: i, code_length, codehash
                 let stacklen = stack.len();
-                // Stack looks like: i, code_length, codehash
                 let i = stack[stacklen - 1].as_usize();
                 let codehash = stack[stacklen - 3];
                 self.inputs.contract_code[&H256::from_uint(&codehash)][i].into()
             }
-            _ => panic!("sup?"),
+            _ => panic!("Invalid prover input function."),
         }
     }
 }
diff --git a/evm/src/memory/segments.rs b/evm/src/memory/segments.rs
index 021781ac..f8d536e9 100644
--- a/evm/src/memory/segments.rs
+++ b/evm/src/memory/segments.rs
@@ -95,7 +95,7 @@ impl Segment {
             Segment::ContextMetadata => 256,
             Segment::KernelGeneral => 256,
             Segment::KernelGeneral2 => 256,
-            Segment::KernelAccountCode => 256,
+            Segment::KernelAccountCode => 8,
             Segment::TxnFields => 256,
             Segment::TxnData => 256,
             Segment::RlpRaw => 8,