1 | #include "taichi/common/core.h" |
2 | |
3 | #if defined(__GNUC__) |
4 | // Ensure we get the 64-bit variants of the CRT's file I/O calls |
5 | #ifndef _FILE_OFFSET_BITS |
6 | #define _FILE_OFFSET_BITS 64 |
7 | #endif |
8 | #if !defined(_LARGEFILE64_SOURCE) && defined(TI_PLATFORM_LINUX) |
9 | // Only Linux has large file extension |
10 | #define _LARGEFILE64_SOURCE 1 |
11 | #endif |
12 | #endif |
13 | |
14 | #include "taichi/common/miniz.h" |
15 | |
16 | namespace taichi { |
17 | |
18 | namespace zip { |
19 | |
20 | inline std::string get_file_name_from_whole_path(const std::string &fn) { |
21 | std::size_t start_index = 0; |
22 | if (fn.rfind("/" ) != std::string::npos) { |
23 | start_index = fn.rfind("/" ) + 1; |
24 | } |
25 | return std::string(fn.begin() + start_index, fn.end()); |
26 | } |
27 | |
28 | void write(std::string fn, const uint8 *data, std::size_t len) { |
29 | mz_bool status; |
30 | TI_ERROR_UNLESS(taichi::ends_with(fn, ".tcb.zip" ), |
31 | "Filename must end with .tcb.zip" ); |
32 | |
33 | std::string fn_uncompressed = get_file_name_from_whole_path(fn); |
34 | fn_uncompressed = std::string(fn_uncompressed.begin(), |
35 | fn_uncompressed.end() - 4); // remove .zip |
36 | std::remove(fn.c_str()); |
37 | |
38 | auto = "Taichi Binary File" ; |
39 | |
40 | status = mz_zip_add_mem_to_archive_file_in_place( |
41 | fn.c_str(), fn_uncompressed.c_str(), |
42 | reinterpret_cast<char *>(const_cast<uint8 *>(data)), len, s_pComment, |
43 | (uint16)strlen(s_pComment), MZ_BEST_COMPRESSION); |
44 | if (!status) { |
45 | TI_ERROR("mz_zip_add_mem_to_archive_file_in_place failed!\n" ); |
46 | } |
47 | } |
48 | |
49 | void write(const std::string &fn, const std::string &data) { |
50 | write(fn, reinterpret_cast<const uint8 *>(data.c_str()), data.size() + 1); |
51 | } |
52 | |
53 | std::vector<uint8> read(const std::string fn, bool verbose) { |
54 | TI_ERROR_UNLESS(taichi::ends_with(fn, ".tcb.zip" ), |
55 | "Filename must end with .tcb.zip" ); |
56 | |
57 | mz_zip_archive zip_archive; |
58 | mz_zip_archive_file_stat file_stat; |
59 | mz_bool status; |
60 | |
61 | memset(&zip_archive, 0, sizeof(zip_archive)); |
62 | status = mz_zip_reader_init_file(&zip_archive, fn.c_str(), 0); |
63 | if (!status) { |
64 | TI_ERROR("mz_zip_reader_init_file() failed!\n" ); |
65 | } |
66 | if (!mz_zip_reader_file_stat(&zip_archive, 0, &file_stat)) { |
67 | mz_zip_reader_end(&zip_archive); |
68 | TI_ERROR("mz_zip_reader_file_stat() failed!\n" ); |
69 | } |
70 | |
71 | if (verbose) { |
72 | TI_TRACE( |
73 | "Filename: {}, Comment: {}, Uncompressed size: {}, Compressed size: " |
74 | "{}, Is Dir: {}\n" , |
75 | file_stat.m_filename, file_stat.m_comment, |
76 | (uint)file_stat.m_uncomp_size, (uint)file_stat.m_comp_size, |
77 | mz_zip_reader_is_file_a_directory(&zip_archive, 0)); |
78 | } |
79 | |
80 | // Close the archive, freeing any resources it was using |
81 | mz_zip_reader_end(&zip_archive); |
82 | |
83 | size_t uncomp_size; |
84 | memset(&zip_archive, 0, sizeof(zip_archive)); |
85 | status = mz_zip_reader_init_file(&zip_archive, fn.c_str(), 0); |
86 | if (!status) { |
87 | TI_ERROR("mz_zip_reader_init_file() failed!\n" ); |
88 | } |
89 | |
90 | std::string fn_uncompressed = get_file_name_from_whole_path(fn); |
91 | fn_uncompressed = std::string(fn_uncompressed.begin(), |
92 | fn_uncompressed.end() - 4); // remove .zip |
93 | auto archive_filename = fn_uncompressed; |
94 | auto p = reinterpret_cast<const uint8 *>(mz_zip_reader_extract_file_to_heap( |
95 | &zip_archive, archive_filename.c_str(), &uncomp_size, 0)); |
96 | |
97 | if (!p) { |
98 | mz_zip_reader_end(&zip_archive); |
99 | TI_ERROR("mz_zip_reader_extract_file_to_heap() failed!" ); |
100 | } |
101 | |
102 | if (verbose) { |
103 | TI_TRACE("Successfully extracted file {}, size {}" , archive_filename, |
104 | (uint)uncomp_size); |
105 | TI_TRACE("File data: {}" , (const char *)p); |
106 | } |
107 | |
108 | std::vector<uint8> ret(p, p + file_stat.m_uncomp_size); |
109 | mz_free((void *)p); |
110 | return ret; |
111 | } |
112 | |
113 | } // namespace zip |
114 | |
115 | } // namespace taichi |
116 | |