1 | #include "taichi/common/virtual_dir.h" |
2 | #include "taichi/common/zip.h" |
3 | |
4 | namespace taichi { |
5 | namespace io { |
6 | |
7 | struct FilesystemVirtualDir : public VirtualDir { |
8 | std::string base_dir_; |
9 | |
10 | explicit FilesystemVirtualDir(const std::string &base_dir) |
11 | : base_dir_(base_dir) { |
12 | } |
13 | |
14 | static std::unique_ptr<VirtualDir> create(const std::string &base_dir) { |
15 | std::string base_dir2; |
16 | if (base_dir.empty()) { |
17 | base_dir2 = "./" ; |
18 | } else if (base_dir.back() != '/') { |
19 | base_dir2 = base_dir + "/" ; |
20 | } else { |
21 | base_dir2 = base_dir; |
22 | } |
23 | |
24 | return std::unique_ptr<VirtualDir>(new FilesystemVirtualDir(base_dir2)); |
25 | } |
26 | |
27 | bool get_file_size(const std::string &path, size_t &size) const override { |
28 | std::fstream f(base_dir_ + path, |
29 | std::ios::in | std::ios::binary | std::ios::ate); |
30 | if (!f.is_open()) { |
31 | return false; |
32 | } |
33 | size = f.tellg(); |
34 | return true; |
35 | } |
36 | size_t load_file(const std::string &path, |
37 | void *data, |
38 | size_t size) const override { |
39 | std::fstream f(base_dir_ + path, std::ios::in | std::ios::binary); |
40 | if (!f.is_open()) { |
41 | return false; |
42 | } |
43 | |
44 | f.read((char *)data, size); |
45 | size_t n = f.gcount(); |
46 | return n; |
47 | } |
48 | }; |
49 | struct ZipArchiveVirtualDir : public VirtualDir { |
50 | zip::ZipArchive archive_; |
51 | |
52 | explicit ZipArchiveVirtualDir(zip::ZipArchive &&archive) |
53 | : archive_(std::move(archive)) { |
54 | } |
55 | |
56 | static std::unique_ptr<VirtualDir> create(const std::string &archive_path) { |
57 | std::fstream f(archive_path, |
58 | std::ios::in | std::ios::binary | std::ios::ate); |
59 | std::vector<uint8_t> archive_data(f.tellg()); |
60 | f.seekg(std::ios::beg); |
61 | f.read((char *)archive_data.data(), archive_data.size()); |
62 | |
63 | return from_zip(archive_data.data(), archive_data.size()); |
64 | } |
65 | static std::unique_ptr<VirtualDir> from_zip(const void *data, size_t size) { |
66 | zip::ZipArchive archive; |
67 | bool succ = zip::ZipArchive::try_from_bytes(data, size, archive); |
68 | if (!succ) { |
69 | return nullptr; |
70 | } |
71 | |
72 | return std::unique_ptr<VirtualDir>( |
73 | new ZipArchiveVirtualDir(std::move(archive))); |
74 | } |
75 | |
76 | bool get_file_size(const std::string &path, size_t &size) const override { |
77 | auto it = archive_.file_dict.find(path); |
78 | if (it == archive_.file_dict.end()) { |
79 | return false; |
80 | } |
81 | |
82 | size = it->second.size(); |
83 | return true; |
84 | } |
85 | size_t load_file(const std::string &path, |
86 | void *data, |
87 | size_t size) const override { |
88 | auto it = archive_.file_dict.find(path); |
89 | if (it == archive_.file_dict.end()) { |
90 | return 0; |
91 | } |
92 | |
93 | size_t n = std::min<size_t>(size, it->second.size()); |
94 | std::memcpy(data, it->second.data(), n); |
95 | return n; |
96 | } |
97 | }; |
98 | |
99 | inline bool is_zip_file(const std::string &path) { |
100 | std::fstream f(path, std::ios::in | std::ios::binary); |
101 | if (!f.is_open()) { |
102 | return false; |
103 | } |
104 | |
105 | // Ensure the file magic matches the Zip format. |
106 | char magic[2]; |
107 | f.read(magic, 2); |
108 | size_t n = f.gcount(); |
109 | if (n == 2 && magic[0] == 'P' && magic[1] == 'K') { |
110 | return true; |
111 | } |
112 | |
113 | return false; |
114 | } |
115 | |
116 | std::unique_ptr<VirtualDir> VirtualDir::open(const std::string &path) { |
117 | if (is_zip_file(path)) { |
118 | return ZipArchiveVirtualDir::create(path); |
119 | } else { |
120 | // (penguinliong) I wanted to use `std::filesyste::is_directory`. But it |
121 | // seems `<filesystem>` is only supported in MSVC. |
122 | return FilesystemVirtualDir::create(path); |
123 | } |
124 | } |
125 | |
126 | std::unique_ptr<VirtualDir> VirtualDir::from_zip(const void *data, |
127 | size_t size) { |
128 | return ZipArchiveVirtualDir::from_zip(data, size); |
129 | } |
130 | std::unique_ptr<VirtualDir> VirtualDir::from_fs_dir( |
131 | const std::string &base_dir) { |
132 | return FilesystemVirtualDir::create(base_dir); |
133 | } |
134 | |
135 | } // namespace io |
136 | } // namespace taichi |
137 | |