1 // Copyright 2016 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "net/tools/huffman_trie/huffman/huffman_builder.h"
6
7 #include <algorithm>
8 #include <ostream>
9
10 #include "base/check.h"
11
12 namespace net::huffman_trie {
13
14 namespace {
15
16 class HuffmanNode {
17 public:
HuffmanNode(uint8_t value,uint32_t count,std::unique_ptr<HuffmanNode> left,std::unique_ptr<HuffmanNode> right)18 HuffmanNode(uint8_t value,
19 uint32_t count,
20 std::unique_ptr<HuffmanNode> left,
21 std::unique_ptr<HuffmanNode> right)
22 : value_(value),
23 count_(count),
24 left_(std::move(left)),
25 right_(std::move(right)) {}
26 ~HuffmanNode() = default;
27
IsLeaf() const28 bool IsLeaf() const {
29 return left_.get() == nullptr && right_.get() == nullptr;
30 }
31
value() const32 uint8_t value() const { return value_; }
count() const33 uint32_t count() const { return count_; }
left() const34 const std::unique_ptr<HuffmanNode>& left() const { return left_; }
right() const35 const std::unique_ptr<HuffmanNode>& right() const { return right_; }
36
37 private:
38 uint8_t value_;
39 uint32_t count_;
40 std::unique_ptr<HuffmanNode> left_;
41 std::unique_ptr<HuffmanNode> right_;
42 };
43
CompareNodes(const std::unique_ptr<HuffmanNode> & lhs,const std::unique_ptr<HuffmanNode> & rhs)44 bool CompareNodes(const std::unique_ptr<HuffmanNode>& lhs,
45 const std::unique_ptr<HuffmanNode>& rhs) {
46 return lhs->count() < rhs->count();
47 }
48
49 } // namespace
50
51 HuffmanBuilder::HuffmanBuilder() = default;
52
53 HuffmanBuilder::~HuffmanBuilder() = default;
54
RecordUsage(uint8_t character)55 void HuffmanBuilder::RecordUsage(uint8_t character) {
56 DCHECK(character < 128);
57 counts_[character & 127] += 1;
58 }
59
ToTable()60 HuffmanRepresentationTable HuffmanBuilder::ToTable() {
61 HuffmanRepresentationTable table;
62 std::unique_ptr<HuffmanNode> node(BuildTree());
63
64 TreeToTable(node.get(), 0, 0, &table);
65 return table;
66 }
67
TreeToTable(HuffmanNode * node,uint32_t bits,uint32_t number_of_bits,HuffmanRepresentationTable * table)68 void HuffmanBuilder::TreeToTable(HuffmanNode* node,
69 uint32_t bits,
70 uint32_t number_of_bits,
71 HuffmanRepresentationTable* table) {
72 if (node->IsLeaf()) {
73 HuffmanRepresentation item;
74 item.bits = bits;
75 item.number_of_bits = number_of_bits;
76
77 table->insert(HuffmanRepresentationPair(node->value(), item));
78 } else {
79 uint32_t new_bits = bits << 1;
80 TreeToTable(node->left().get(), new_bits, number_of_bits + 1, table);
81 TreeToTable(node->right().get(), new_bits | 1, number_of_bits + 1, table);
82 }
83 }
84
ToVector()85 std::vector<uint8_t> HuffmanBuilder::ToVector() {
86 std::vector<uint8_t> bytes;
87 std::unique_ptr<HuffmanNode> node(BuildTree());
88 WriteToVector(node.get(), &bytes);
89 return bytes;
90 }
91
WriteToVector(HuffmanNode * node,std::vector<uint8_t> * vector)92 uint32_t HuffmanBuilder::WriteToVector(HuffmanNode* node,
93 std::vector<uint8_t>* vector) {
94 uint8_t left_value;
95 uint8_t right_value;
96 uint32_t child_position;
97
98 if (node->left()->IsLeaf()) {
99 left_value = 128 | node->left()->value();
100 } else {
101 child_position = WriteToVector(node->left().get(), vector);
102 DCHECK(child_position < 512) << "huffman tree too large";
103 left_value = child_position / 2;
104 }
105
106 if (node->right()->IsLeaf()) {
107 right_value = 128 | node->right()->value();
108 } else {
109 child_position = WriteToVector(node->right().get(), vector);
110 DCHECK(child_position < 512) << "huffman tree to large";
111 right_value = child_position / 2;
112 }
113
114 uint32_t position = static_cast<uint32_t>(vector->size());
115 vector->push_back(left_value);
116 vector->push_back(right_value);
117 return position;
118 }
119
BuildTree()120 std::unique_ptr<HuffmanNode> HuffmanBuilder::BuildTree() {
121 std::vector<std::unique_ptr<HuffmanNode>> nodes;
122 nodes.reserve(counts_.size());
123
124 for (const auto& item : counts_) {
125 nodes.push_back(std::make_unique<HuffmanNode>(item.first, item.second,
126 nullptr, nullptr));
127 }
128
129 // At least 2 entries are required for everything to work properly. Add
130 // arbitrary values to fill the tree.
131 for (uint8_t i = 0; nodes.size() < 2 && i < 2; ++i) {
132 for (const auto& node : nodes) {
133 if (node->value() == i) {
134 break;
135 }
136 }
137
138 nodes.push_back(std::make_unique<HuffmanNode>(i, 0, nullptr, nullptr));
139 }
140
141 std::stable_sort(nodes.begin(), nodes.end(), CompareNodes);
142
143 while (nodes.size() > 1) {
144 std::unique_ptr<HuffmanNode> a = std::move(nodes[0]);
145 std::unique_ptr<HuffmanNode> b = std::move(nodes[1]);
146
147 uint32_t count_a = a->count();
148 uint32_t count_b = b->count();
149
150 auto parent = std::make_unique<HuffmanNode>(0, count_a + count_b,
151 std::move(a), std::move(b));
152
153 nodes.erase(nodes.begin());
154 nodes[0] = std::move(parent);
155
156 std::stable_sort(nodes.begin(), nodes.end(), CompareNodes);
157 }
158
159 return std::move(nodes[0]);
160 }
161
162 } // namespace net::huffman_trie
163