1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "src/huffman_decode.h"
18 #include <queue>
19
20 namespace mindspore {
21 namespace lite {
DoHuffmanDecode(const std::string & input_str,void * decoded_data,size_t data_len)22 STATUS HuffmanDecode::DoHuffmanDecode(const std::string &input_str, void *decoded_data, size_t data_len) {
23 if (decoded_data == nullptr) {
24 MS_LOG(ERROR) << "decoded_data is nullptr.";
25 return RET_ERROR;
26 }
27
28 int status;
29 std::string huffman_decoded_str;
30 auto key_pos = input_str.find_first_of('#');
31 auto code_pos = input_str.find_first_of('#', key_pos + 1);
32 auto key = input_str.substr(0, key_pos);
33 auto code = input_str.substr(key_pos + 1, code_pos - key_pos - 1);
34 auto encoded_data = input_str.substr(code_pos + 1);
35
36 auto root = new (std::nothrow) HuffmanNode();
37 if (root == nullptr) {
38 MS_LOG(ERROR) << "new HuffmanNode failed.";
39 return RET_MEMORY_FAILED;
40 }
41 root->left = nullptr;
42 root->right = nullptr;
43 root->parent = nullptr;
44
45 status = RebuildHuffmanTree(key, code, root);
46 if (status != RET_OK) {
47 MS_LOG(ERROR) << "Rebuild huffman tree failed.";
48 delete root;
49 return status;
50 }
51
52 status = DoHuffmanDecompress(root, encoded_data, &huffman_decoded_str);
53 if (status != RET_OK) {
54 MS_LOG(ERROR) << "DoHuffmanDecompress failed.";
55 delete root;
56 return status;
57 }
58
59 size_t len = huffman_decoded_str.length();
60 if (data_len >= len) {
61 memcpy(decoded_data, huffman_decoded_str.c_str(), len);
62 } else {
63 FreeHuffmanNodeTree(root);
64 return RET_ERROR;
65 }
66 FreeHuffmanNodeTree(root);
67 return RET_OK;
68 }
69
RebuildHuffmanTree(std::string keys,std::string codes,const HuffmanNodePtr & root)70 STATUS HuffmanDecode::RebuildHuffmanTree(std::string keys, std::string codes, const HuffmanNodePtr &root) {
71 HuffmanNodePtr cur_node, tmp_node, new_node;
72
73 auto huffman_keys = Str2Vec(std::move(keys));
74 auto huffman_codes = Str2Vec(std::move(codes));
75
76 for (size_t i = 0; i < huffman_codes.size(); ++i) {
77 auto key = stoi(huffman_keys[i]);
78 auto code = huffman_codes[i];
79 auto code_len = code.length();
80 cur_node = root;
81 for (size_t j = 0; j < code_len; ++j) {
82 if (code[j] == '0') {
83 tmp_node = cur_node->left;
84 } else if (code[j] == '1') {
85 tmp_node = cur_node->right;
86 } else {
87 MS_LOG(ERROR) << "find huffman code is not 0 or 1";
88 return RET_ERROR;
89 }
90
91 if (tmp_node == nullptr) {
92 new_node = new (std::nothrow) HuffmanNode();
93 if (new_node == nullptr) {
94 MS_LOG(ERROR) << "new HuffmanNode failed.";
95 return RET_MEMORY_FAILED;
96 }
97 new_node->left = nullptr;
98 new_node->right = nullptr;
99 new_node->parent = cur_node;
100
101 if (j == code_len - 1) {
102 new_node->key = key;
103 new_node->code = code;
104 }
105
106 if (code[j] == '0') {
107 cur_node->left = new_node;
108 } else {
109 cur_node->right = new_node;
110 }
111
112 tmp_node = new_node;
113 } else if (j == code_len - 1) {
114 MS_LOG(ERROR) << "the huffman code is incomplete.";
115 return RET_ERROR;
116 } else if (tmp_node->left == nullptr && tmp_node->right == nullptr) {
117 MS_LOG(ERROR) << "the huffman code is incomplete";
118 return RET_ERROR;
119 }
120 cur_node = tmp_node;
121 }
122 }
123 return RET_OK;
124 }
125
DoHuffmanDecompress(HuffmanNodePtr root,std::string encoded_data,std::string * decoded_str)126 STATUS HuffmanDecode::DoHuffmanDecompress(HuffmanNodePtr root, std::string encoded_data, std::string *decoded_str) {
127 HuffmanNodePtr cur_node = root;
128 bool pseudo_eof = false;
129 size_t pos = 0;
130 unsigned char flag;
131
132 decoded_str->clear();
133 while (pos < encoded_data.length()) {
134 auto u_char = static_cast<unsigned char>(encoded_data[pos]);
135 flag = 0x80;
136 for (size_t i = 0; i < 8; ++i) { // traverse the 8 bit num, to find the leaf node
137 if (u_char & flag) {
138 cur_node = cur_node->right;
139 } else {
140 cur_node = cur_node->left;
141 }
142 if (cur_node->left == nullptr && cur_node->right == nullptr) {
143 auto key = cur_node->key;
144 if (key == PSEUDO_EOF) {
145 pseudo_eof = true;
146 break;
147 } else {
148 *decoded_str += static_cast<char>(cur_node->key);
149 cur_node = root;
150 }
151 }
152 flag = flag >> 1;
153 }
154 pos++;
155 if (pseudo_eof) {
156 break;
157 }
158 }
159 return RET_OK;
160 }
161
FreeHuffmanNodeTree(HuffmanNodePtr root)162 void HuffmanDecode::FreeHuffmanNodeTree(HuffmanNodePtr root) {
163 if (root == nullptr) {
164 return;
165 }
166 std::queue<HuffmanNodePtr> node_queue;
167 node_queue.push(root);
168 while (!node_queue.empty()) {
169 auto cur_node = node_queue.front();
170 node_queue.pop();
171 if (cur_node->left != nullptr) {
172 node_queue.push(cur_node->left);
173 }
174 if (cur_node->right != nullptr) {
175 node_queue.push(cur_node->right);
176 }
177 delete (cur_node);
178 }
179 }
180 } // namespace lite
181 } // namespace mindspore
182