1#include "taichi/common/virtual_dir.h"
2#include "taichi/common/zip.h"
3
4namespace taichi {
5namespace io {
6
7struct 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};
49struct 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
99inline 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
116std::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
126std::unique_ptr<VirtualDir> VirtualDir::from_zip(const void *data,
127 size_t size) {
128 return ZipArchiveVirtualDir::from_zip(data, size);
129}
130std::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