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
16namespace taichi {
17
18namespace zip {
19
20inline 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
28void 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 s_pComment = "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
49void write(const std::string &fn, const std::string &data) {
50 write(fn, reinterpret_cast<const uint8 *>(data.c_str()), data.size() + 1);
51}
52
53std::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