1 // Copyright (c) 2023 Google LLC.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include <filesystem>
16 #include <iostream>
17 
18 #include "extract_source.h"
19 #include "source/opt/log.h"
20 #include "tools/io.h"
21 #include "tools/util/cli_consumer.h"
22 #include "tools/util/flags.h"
23 
24 namespace {
25 
26 constexpr auto kHelpTextFmt =
27     R"(%s - Dumps information from a SPIR-V binary.
28 
29 Usage: %s [options] <filename>
30 
31 one of the following switches must be given:
32   --source        Extract source files obtained from debug symbols, output to stdout.
33   --entrypoint    Extracts the entrypoint name of the module, output to stdout.
34   --compiler-cmd  Extracts the command line used to compile this module, output to stdout.
35 
36 
37 General options:
38   -h, --help      Print this help.
39   --version       Display assembler version information.
40   -f,--force      Allow output file overwrite.
41 
42 Source dump options:
43   --list          Do not extract source code, only print filenames to stdout.
44   --outdir        Where shall the exrtacted HLSL/HLSL files be written to?
45                   File written to stdout if '-' is given. Default is `-`.
46 )";
47 
48 // Removes trailing '/' from `input`.
49 // A behavior difference has been observed between libc++ implementations.
50 // Fixing path to prevent this edge case to be reached.
51 // (https://github.com/llvm/llvm-project/issues/60634)
fixPathForLLVM(std::string input)52 std::string fixPathForLLVM(std::string input) {
53   while (!input.empty() && input.back() == '/') input.resize(input.size() - 1);
54   return input;
55 }
56 
57 // Write each HLSL file described in `sources` in a file in `outdirPath`.
58 // Doesn't ovewrite existing files, unless `overwrite` is set to true. The
59 // created HLSL file's filename is the path's filename obtained from `sources`.
60 // Returns true if all files could be written. False otherwise.
OutputSourceFiles(const std::unordered_map<std::string,std::string> & sources,const std::string & outdirPath,bool overwrite)61 bool OutputSourceFiles(
62     const std::unordered_map<std::string, std::string>& sources,
63     const std::string& outdirPath, bool overwrite) {
64   std::filesystem::path outdir(fixPathForLLVM(outdirPath));
65   if (!std::filesystem::is_directory(outdir)) {
66     if (!std::filesystem::create_directories(outdir)) {
67       std::cerr << "error: could not create output directory " << outdir
68                 << std::endl;
69       return false;
70     }
71   }
72 
73   for (const auto & [ filepath, code ] : sources) {
74     if (code.empty()) {
75       std::cout << "Ignoring source for " << filepath
76                 << ": no code source in debug infos." << std::endl;
77       continue;
78     }
79 
80     std::filesystem::path old_path(filepath);
81     std::filesystem::path new_path = outdir / old_path.filename();
82 
83     if (!overwrite && std::filesystem::exists(new_path)) {
84       std::cerr << "file " << filepath
85                 << " already exists, aborting (use --overwrite to allow it)."
86                 << std::endl;
87       return false;
88     }
89 
90     std::cout << "Exporting " << new_path << std::endl;
91     if (!WriteFile<char>(new_path.string().c_str(), "w", code.c_str(),
92                          code.size())) {
93       return false;
94     }
95   }
96   return true;
97 }
98 
99 }  // namespace
100 
101 // clang-format off
102 FLAG_SHORT_bool(  h,            /* default_value= */ false, /* required= */ false);
103 FLAG_LONG_bool(   help,         /* default_value= */ false, /* required= */ false);
104 FLAG_LONG_bool(   version,      /* default_value= */ false, /* required= */ false);
105 FLAG_LONG_bool(   source,       /* default_value= */ false, /* required= */ false);
106 FLAG_LONG_bool(   entrypoint,   /* default_value= */ false, /* required= */ false);
107 FLAG_LONG_bool(   compiler_cmd, /* default_value= */ false, /* required= */ false);
108 FLAG_SHORT_bool(  f,            /* default_value= */ false, /* required= */ false);
109 FLAG_LONG_bool(   force,        /* default_value= */ false, /* required= */ false);
110 FLAG_LONG_string( outdir,       /* default_value= */ "-",   /* required= */ false);
111 FLAG_LONG_bool(   list,         /* default_value= */ false, /* required= */ false);
112 // clang-format on
113 
main(int,const char ** argv)114 int main(int, const char** argv) {
115   if (!flags::Parse(argv)) {
116     return 1;
117   }
118   if (flags::h.value() || flags::help.value()) {
119     printf(kHelpTextFmt, argv[0], argv[0]);
120     return 0;
121   }
122   if (flags::version.value()) {
123     printf("%s\n", spvSoftwareVersionDetailsString());
124     return 0;
125   }
126 
127   if (flags::positional_arguments.size() != 1) {
128     std::cerr << "Expected exactly one input file." << std::endl;
129     return 1;
130   }
131   if (flags::entrypoint.value() || flags::compiler_cmd.value()) {
132     std::cerr << "Unimplemented flags." << std::endl;
133     return 1;
134   }
135 
136   std::vector<uint32_t> binary;
137   if (!ReadBinaryFile(flags::positional_arguments[0].c_str(), &binary)) {
138     return 1;
139   }
140 
141   if (flags::source.value()) {
142     std::unordered_map<std::string, std::string> sourceCode;
143     if (!ExtractSourceFromModule(binary, &sourceCode)) {
144       return 1;
145     }
146 
147     if (flags::list.value()) {
148       for (const auto & [ filename, source ] : sourceCode) {
149         printf("%s\n", filename.c_str());
150       }
151       return 0;
152     }
153 
154     const bool outputToConsole = flags::outdir.value() == "-";
155 
156     if (outputToConsole) {
157       for (const auto & [ filename, source ] : sourceCode) {
158         std::cout << filename << ":" << std::endl
159                   << source << std::endl
160                   << std::endl;
161       }
162       return 0;
163     }
164 
165     const std::filesystem::path outdirPath(flags::outdir.value());
166     if (!OutputSourceFiles(sourceCode, outdirPath.string(),
167                            flags::force.value())) {
168       return 1;
169     }
170   }
171 
172   // FIXME: implement logic.
173   return 0;
174 }
175