• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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