diff --git a/enc/brotli_bit_stream.cc b/enc/brotli_bit_stream.cc
index 031a715..314eb15 100644
--- a/enc/brotli_bit_stream.cc
+++ b/enc/brotli_bit_stream.cc
@@ -18,8 +18,11 @@
 
 #include "./brotli_bit_stream.h"
 
+#include <algorithm>
+#include <limits>
 #include <vector>
 
+#include "./bit_cost.h"
 #include "./entropy_encode.h"
 #include "./fast_log.h"
 #include "./prefix.h"
@@ -322,6 +325,156 @@ void BuildAndStoreHuffmanTree(const int *histogram,
   }
 }
 
+int IndexOf(const std::vector<int>& v, int value) {
+  for (int i = 0; i < v.size(); ++i) {
+    if (v[i] == value) return i;
+  }
+  return -1;
+}
+
+void MoveToFront(std::vector<int>* v, int index) {
+  int value = (*v)[index];
+  for (int i = index; i > 0; --i) {
+    (*v)[i] = (*v)[i - 1];
+  }
+  (*v)[0] = value;
+}
+
+std::vector<int> MoveToFrontTransform(const std::vector<int>& v) {
+  if (v.empty()) return v;
+  std::vector<int> mtf(*max_element(v.begin(), v.end()) + 1);
+  for (int i = 0; i < mtf.size(); ++i) mtf[i] = i;
+  std::vector<int> result(v.size());
+  for (int i = 0; i < v.size(); ++i) {
+    int index = IndexOf(mtf, v[i]);
+    result[i] = index;
+    MoveToFront(&mtf, index);
+  }
+  return result;
+}
+
+// Finds runs of zeros in v_in and replaces them with a prefix code of the run
+// length plus extra bits in *v_out and *extra_bits. Non-zero values in v_in are
+// shifted by *max_length_prefix. Will not create prefix codes bigger than the
+// initial value of *max_run_length_prefix. The prefix code of run length L is
+// simply Log2Floor(L) and the number of extra bits is the same as the prefix
+// code.
+void RunLengthCodeZeros(const std::vector<int>& v_in,
+                        int* max_run_length_prefix,
+                        std::vector<int>* v_out,
+                        std::vector<int>* extra_bits) {
+  int max_reps = 0;
+  for (int i = 0; i < v_in.size();) {
+    for (; i < v_in.size() && v_in[i] != 0; ++i) ;
+    int reps = 0;
+    for (; i < v_in.size() && v_in[i] == 0; ++i) {
+      ++reps;
+    }
+    max_reps = std::max(reps, max_reps);
+  }
+  int max_prefix = max_reps > 0 ? Log2Floor(max_reps) : 0;
+  *max_run_length_prefix = std::min(max_prefix, *max_run_length_prefix);
+  for (int i = 0; i < v_in.size();) {
+    if (v_in[i] != 0) {
+      v_out->push_back(v_in[i] + *max_run_length_prefix);
+      extra_bits->push_back(0);
+      ++i;
+    } else {
+      int reps = 1;
+      for (uint32_t k = i + 1; k < v_in.size() && v_in[k] == 0; ++k) {
+        ++reps;
+      }
+      i += reps;
+      while (reps) {
+        if (reps < (2 << *max_run_length_prefix)) {
+          int run_length_prefix = Log2Floor(reps);
+          v_out->push_back(run_length_prefix);
+          extra_bits->push_back(reps - (1 << run_length_prefix));
+          break;
+        } else {
+          v_out->push_back(*max_run_length_prefix);
+          extra_bits->push_back((1 << *max_run_length_prefix) - 1);
+          reps -= (2 << *max_run_length_prefix) - 1;
+        }
+      }
+    }
+  }
+}
+
+// Returns a maximum zero-run-length-prefix value such that run-length coding
+// zeros in v with this maximum prefix value and then encoding the resulting
+// histogram and entropy-coding v produces the least amount of bits.
+int BestMaxZeroRunLengthPrefix(const std::vector<int>& v) {
+  int min_cost = std::numeric_limits<int>::max();
+  int best_max_prefix = 0;
+  for (int max_prefix = 0; max_prefix <= 16; ++max_prefix) {
+    std::vector<int> rle_symbols;
+    std::vector<int> extra_bits;
+    int max_run_length_prefix = max_prefix;
+    RunLengthCodeZeros(v, &max_run_length_prefix, &rle_symbols, &extra_bits);
+    if (max_run_length_prefix < max_prefix) break;
+    HistogramContextMap histogram;
+    for (int i = 0; i < rle_symbols.size(); ++i) {
+      histogram.Add(rle_symbols[i]);
+    }
+    int bit_cost = PopulationCost(histogram);
+    if (max_prefix > 0) {
+      bit_cost += 4;
+    }
+    for (int i = 1; i <= max_prefix; ++i) {
+      bit_cost += histogram.data_[i] * i;  // extra bits
+    }
+    if (bit_cost < min_cost) {
+      min_cost = bit_cost;
+      best_max_prefix = max_prefix;
+    }
+  }
+  return best_max_prefix;
+}
+
+void EncodeContextMap(const std::vector<int>& context_map,
+                      int num_clusters,
+                      int* storage_ix, uint8_t* storage) {
+  StoreVarLenUint8(num_clusters - 1, storage_ix, storage);
+
+  if (num_clusters == 1) {
+    return;
+  }
+
+  std::vector<int> transformed_symbols = MoveToFrontTransform(context_map);
+  std::vector<int> rle_symbols;
+  std::vector<int> extra_bits;
+  int max_run_length_prefix = BestMaxZeroRunLengthPrefix(transformed_symbols);
+  RunLengthCodeZeros(transformed_symbols, &max_run_length_prefix,
+                     &rle_symbols, &extra_bits);
+  HistogramContextMap symbol_histogram;
+  for (int i = 0; i < rle_symbols.size(); ++i) {
+    symbol_histogram.Add(rle_symbols[i]);
+  }
+  bool use_rle = max_run_length_prefix > 0;
+  WriteBits(1, use_rle, storage_ix, storage);
+  if (use_rle) {
+    WriteBits(4, max_run_length_prefix - 1, storage_ix, storage);
+  }
+  EntropyCodeContextMap symbol_code;
+  memset(symbol_code.depth_, 0, sizeof(symbol_code.depth_));
+  memset(symbol_code.bits_, 0, sizeof(symbol_code.bits_));
+  BuildAndStoreHuffmanTree(symbol_histogram.data_,
+                           num_clusters + max_run_length_prefix,
+                           9,  // quality
+                           symbol_code.depth_, symbol_code.bits_,
+                           storage_ix, storage);
+  for (int i = 0; i < rle_symbols.size(); ++i) {
+    WriteBits(symbol_code.depth_[rle_symbols[i]],
+              symbol_code.bits_[rle_symbols[i]],
+              storage_ix, storage);
+    if (rle_symbols[i] > 0 && rle_symbols[i] <= max_run_length_prefix) {
+      WriteBits(rle_symbols[i], extra_bits[i], storage_ix, storage);
+    }
+  }
+  WriteBits(1, 1, storage_ix, storage);  // use move-to-front
+}
+
 void StoreBlockSwitch(const BlockSplitCode& code,
                       const int block_ix,
                       int* storage_ix,
diff --git a/enc/brotli_bit_stream.h b/enc/brotli_bit_stream.h
index de68dfd..252edf8 100644
--- a/enc/brotli_bit_stream.h
+++ b/enc/brotli_bit_stream.h
@@ -53,6 +53,12 @@ void StoreTrivialContextMap(int num_types,
                             int* storage_ix,
                             uint8_t* storage);
 
+void StoreHuffmanTreeOfHuffmanTreeToBitMask(
+    const int num_codes,
+    const uint8_t *code_length_bitdepth,
+    int *storage_ix,
+    uint8_t *storage);
+
 // Builds a Huffman tree from histogram[0:length] into depth[0:length] and
 // bits[0:length] and stores the encoded tree to the bit stream.
 void BuildAndStoreHuffmanTree(const int *histogram,
@@ -63,6 +69,12 @@ void BuildAndStoreHuffmanTree(const int *histogram,
                               int* storage_ix,
                               uint8_t* storage);
 
+// Encodes the given context map to the bit stream. The number of different
+// histogram ids is given by num_clusters.
+void EncodeContextMap(const std::vector<int>& context_map,
+                      int num_clusters,
+                      int* storage_ix, uint8_t* storage);
+
 // Data structure that stores everything that is needed to encode each block
 // switch command.
 struct BlockSplitCode {
diff --git a/enc/encode.cc b/enc/encode.cc
index 7fb77fd..c819389 100644
--- a/enc/encode.cc
+++ b/enc/encode.cc
@@ -224,153 +224,6 @@ void RecomputeDistancePrefixes(std::vector<Command>* cmds,
   }
 }
 
-int IndexOf(const std::vector<int>& v, int value) {
-  for (int i = 0; i < v.size(); ++i) {
-    if (v[i] == value) return i;
-  }
-  return -1;
-}
-
-void MoveToFront(std::vector<int>* v, int index) {
-  int value = (*v)[index];
-  for (int i = index; i > 0; --i) {
-    (*v)[i] = (*v)[i - 1];
-  }
-  (*v)[0] = value;
-}
-
-std::vector<int> MoveToFrontTransform(const std::vector<int>& v) {
-  if (v.empty()) return v;
-  std::vector<int> mtf(*max_element(v.begin(), v.end()) + 1);
-  for (int i = 0; i < mtf.size(); ++i) mtf[i] = i;
-  std::vector<int> result(v.size());
-  for (int i = 0; i < v.size(); ++i) {
-    int index = IndexOf(mtf, v[i]);
-    result[i] = index;
-    MoveToFront(&mtf, index);
-  }
-  return result;
-}
-
-// Finds runs of zeros in v_in and replaces them with a prefix code of the run
-// length plus extra bits in *v_out and *extra_bits. Non-zero values in v_in are
-// shifted by *max_length_prefix. Will not create prefix codes bigger than the
-// initial value of *max_run_length_prefix. The prefix code of run length L is
-// simply Log2Floor(L) and the number of extra bits is the same as the prefix
-// code.
-void RunLengthCodeZeros(const std::vector<int>& v_in,
-                        int* max_run_length_prefix,
-                        std::vector<int>* v_out,
-                        std::vector<int>* extra_bits) {
-  int max_reps = 0;
-  for (int i = 0; i < v_in.size();) {
-    for (; i < v_in.size() && v_in[i] != 0; ++i) ;
-    int reps = 0;
-    for (; i < v_in.size() && v_in[i] == 0; ++i) {
-      ++reps;
-    }
-    max_reps = std::max(reps, max_reps);
-  }
-  int max_prefix = max_reps > 0 ? Log2Floor(max_reps) : 0;
-  *max_run_length_prefix = std::min(max_prefix, *max_run_length_prefix);
-  for (int i = 0; i < v_in.size();) {
-    if (v_in[i] != 0) {
-      v_out->push_back(v_in[i] + *max_run_length_prefix);
-      extra_bits->push_back(0);
-      ++i;
-    } else {
-      int reps = 1;
-      for (uint32_t k = i + 1; k < v_in.size() && v_in[k] == 0; ++k) {
-        ++reps;
-      }
-      i += reps;
-      while (reps) {
-        if (reps < (2 << *max_run_length_prefix)) {
-          int run_length_prefix = Log2Floor(reps);
-          v_out->push_back(run_length_prefix);
-          extra_bits->push_back(reps - (1 << run_length_prefix));
-          break;
-        } else {
-          v_out->push_back(*max_run_length_prefix);
-          extra_bits->push_back((1 << *max_run_length_prefix) - 1);
-          reps -= (2 << *max_run_length_prefix) - 1;
-        }
-      }
-    }
-  }
-}
-
-// Returns a maximum zero-run-length-prefix value such that run-length coding
-// zeros in v with this maximum prefix value and then encoding the resulting
-// histogram and entropy-coding v produces the least amount of bits.
-int BestMaxZeroRunLengthPrefix(const std::vector<int>& v) {
-  int min_cost = std::numeric_limits<int>::max();
-  int best_max_prefix = 0;
-  for (int max_prefix = 0; max_prefix <= 16; ++max_prefix) {
-    std::vector<int> rle_symbols;
-    std::vector<int> extra_bits;
-    int max_run_length_prefix = max_prefix;
-    RunLengthCodeZeros(v, &max_run_length_prefix, &rle_symbols, &extra_bits);
-    if (max_run_length_prefix < max_prefix) break;
-    HistogramContextMap histogram;
-    for (int i = 0; i < rle_symbols.size(); ++i) {
-      histogram.Add(rle_symbols[i]);
-    }
-    int bit_cost = PopulationCost(histogram);
-    if (max_prefix > 0) {
-      bit_cost += 4;
-    }
-    for (int i = 1; i <= max_prefix; ++i) {
-      bit_cost += histogram.data_[i] * i;  // extra bits
-    }
-    if (bit_cost < min_cost) {
-      min_cost = bit_cost;
-      best_max_prefix = max_prefix;
-    }
-  }
-  return best_max_prefix;
-}
-
-void EncodeContextMap(const std::vector<int>& context_map,
-                      int num_clusters,
-                      int* storage_ix, uint8_t* storage) {
-  StoreVarLenUint8(num_clusters - 1, storage_ix, storage);
-
-  if (num_clusters == 1) {
-    return;
-  }
-
-  std::vector<int> transformed_symbols = MoveToFrontTransform(context_map);
-  std::vector<int> rle_symbols;
-  std::vector<int> extra_bits;
-  int max_run_length_prefix = BestMaxZeroRunLengthPrefix(transformed_symbols);
-  RunLengthCodeZeros(transformed_symbols, &max_run_length_prefix,
-                     &rle_symbols, &extra_bits);
-  HistogramContextMap symbol_histogram;
-  for (int i = 0; i < rle_symbols.size(); ++i) {
-    symbol_histogram.Add(rle_symbols[i]);
-  }
-  bool use_rle = max_run_length_prefix > 0;
-  WriteBits(1, use_rle, storage_ix, storage);
-  if (use_rle) {
-    WriteBits(4, max_run_length_prefix - 1, storage_ix, storage);
-  }
-  EntropyCodeContextMap symbol_code;
-  BuildAndStoreEntropyCode(symbol_histogram, 15,
-                           num_clusters + max_run_length_prefix,
-                           &symbol_code,
-                           storage_ix, storage);
-  for (int i = 0; i < rle_symbols.size(); ++i) {
-    WriteBits(symbol_code.depth_[rle_symbols[i]],
-              symbol_code.bits_[rle_symbols[i]],
-              storage_ix, storage);
-    if (rle_symbols[i] > 0 && rle_symbols[i] <= max_run_length_prefix) {
-      WriteBits(rle_symbols[i], extra_bits[i], storage_ix, storage);
-    }
-  }
-  WriteBits(1, 1, storage_ix, storage);  // use move-to-front
-}
-
 void MoveAndEncode(const BlockSplitCode& code,
                    BlockSplitIterator* it,
                    int* storage_ix, uint8_t* storage) {