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 "tools/objdump/extract_source.h"
16 
17 #include <gtest/gtest.h>
18 
19 #include <string>
20 
21 #include "source/opt/build_module.h"
22 #include "source/opt/ir_context.h"
23 #include "spirv-tools/libspirv.hpp"
24 #include "tools/util/cli_consumer.h"
25 
26 namespace {
27 
28 constexpr auto kDefaultEnvironment = SPV_ENV_UNIVERSAL_1_6;
29 
ExtractSource(const std::string & spv_source)30 std::pair<bool, std::unordered_map<std::string, std::string>> ExtractSource(
31     const std::string& spv_source) {
32   std::unique_ptr<spvtools::opt::IRContext> ctx = spvtools::BuildModule(
33       kDefaultEnvironment, spvtools::utils::CLIMessageConsumer, spv_source,
34       spvtools::SpirvTools::kDefaultAssembleOption |
35           SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
36   std::vector<uint32_t> binary;
37   ctx->module()->ToBinary(&binary, /* skip_nop = */ false);
38   std::unordered_map<std::string, std::string> output;
39   bool result = ExtractSourceFromModule(binary, &output);
40   return std::make_pair(result, std::move(output));
41 }
42 
43 }  // namespace
44 
TEST(ExtractSourceTest,no_debug)45 TEST(ExtractSourceTest, no_debug) {
46   std::string source = R"(
47            OpCapability Shader
48            OpCapability Linkage
49            OpMemoryModel Logical GLSL450
50    %void = OpTypeVoid
51       %2 = OpTypeFunction %void
52    %bool = OpTypeBool
53       %4 = OpUndef %bool
54       %5 = OpFunction %void None %2
55       %6 = OpLabel
56            OpReturn
57            OpFunctionEnd
58   )";
59 
60   auto[success, result] = ExtractSource(source);
61   ASSERT_TRUE(success);
62   ASSERT_TRUE(result.size() == 0);
63 }
64 
TEST(ExtractSourceTest,SimpleSource)65 TEST(ExtractSourceTest, SimpleSource) {
66   std::string source = R"(
67       OpCapability Shader
68       OpMemoryModel Logical GLSL450
69       OpEntryPoint GLCompute %1 "compute_1"
70       OpExecutionMode %1 LocalSize 1 1 1
71  %2 = OpString "compute.hlsl"
72       OpSource HLSL 660 %2 "[numthreads(1, 1, 1)] void compute_1(){ }"
73       OpName %1 "compute_1"
74  %3 = OpTypeVoid
75  %4 = OpTypeFunction %3
76  %1 = OpFunction %3 None %4
77  %5 = OpLabel
78       OpLine %2 1 41
79       OpReturn
80       OpFunctionEnd
81   )";
82 
83   auto[success, result] = ExtractSource(source);
84   ASSERT_TRUE(success);
85   ASSERT_TRUE(result.size() == 1);
86   ASSERT_TRUE(result["compute.hlsl"] ==
87               "[numthreads(1, 1, 1)] void compute_1(){ }");
88 }
89 
TEST(ExtractSourceTest,SourceContinued)90 TEST(ExtractSourceTest, SourceContinued) {
91   std::string source = R"(
92       OpCapability Shader
93       OpMemoryModel Logical GLSL450
94       OpEntryPoint GLCompute %1 "compute_1"
95       OpExecutionMode %1 LocalSize 1 1 1
96  %2 = OpString "compute.hlsl"
97       OpSource HLSL 660 %2 "[numthreads(1, 1, 1)] "
98       OpSourceContinued "void compute_1(){ }"
99       OpName %1 "compute_1"
100  %3 = OpTypeVoid
101  %4 = OpTypeFunction %3
102  %1 = OpFunction %3 None %4
103  %5 = OpLabel
104       OpLine %2 1 41
105       OpReturn
106       OpFunctionEnd
107   )";
108 
109   auto[success, result] = ExtractSource(source);
110   ASSERT_TRUE(success);
111   ASSERT_TRUE(result.size() == 1);
112   ASSERT_TRUE(result["compute.hlsl"] ==
113               "[numthreads(1, 1, 1)] void compute_1(){ }");
114 }
115 
TEST(ExtractSourceTest,OnlyFilename)116 TEST(ExtractSourceTest, OnlyFilename) {
117   std::string source = R"(
118       OpCapability Shader
119       OpMemoryModel Logical GLSL450
120       OpEntryPoint GLCompute %1 "compute_1"
121       OpExecutionMode %1 LocalSize 1 1 1
122  %2 = OpString "compute.hlsl"
123       OpSource HLSL 660 %2
124       OpName %1 "compute_1"
125  %3 = OpTypeVoid
126  %4 = OpTypeFunction %3
127  %1 = OpFunction %3 None %4
128  %5 = OpLabel
129       OpLine %2 1 41
130       OpReturn
131       OpFunctionEnd
132   )";
133 
134   auto[success, result] = ExtractSource(source);
135   ASSERT_TRUE(success);
136   ASSERT_TRUE(result.size() == 1);
137   ASSERT_TRUE(result["compute.hlsl"] == "");
138 }
139 
TEST(ExtractSourceTest,MultipleFiles)140 TEST(ExtractSourceTest, MultipleFiles) {
141   std::string source = R"(
142       OpCapability Shader
143       OpMemoryModel Logical GLSL450
144       OpEntryPoint GLCompute %1 "compute_1"
145       OpExecutionMode %1 LocalSize 1 1 1
146  %2 = OpString "compute1.hlsl"
147  %3 = OpString "compute2.hlsl"
148       OpSource HLSL 660 %2 "some instruction"
149       OpSource HLSL 660 %3 "some other instruction"
150       OpName %1 "compute_1"
151  %4 = OpTypeVoid
152  %5 = OpTypeFunction %4
153  %1 = OpFunction %4 None %5
154  %6 = OpLabel
155       OpLine %2 1 41
156       OpReturn
157       OpFunctionEnd
158   )";
159 
160   auto[success, result] = ExtractSource(source);
161   ASSERT_TRUE(success);
162   ASSERT_TRUE(result.size() == 2);
163   ASSERT_TRUE(result["compute1.hlsl"] == "some instruction");
164   ASSERT_TRUE(result["compute2.hlsl"] == "some other instruction");
165 }
166 
TEST(ExtractSourceTest,MultilineCode)167 TEST(ExtractSourceTest, MultilineCode) {
168   std::string source = R"(
169                OpCapability Shader
170                OpMemoryModel Logical GLSL450
171                OpEntryPoint GLCompute %1 "compute_1"
172                OpExecutionMode %1 LocalSize 1 1 1
173           %2 = OpString "compute.hlsl"
174                OpSource HLSL 660 %2 "[numthreads(1, 1, 1)]
175 void compute_1() {
176 }
177 "
178                OpName %1 "compute_1"
179           %3 = OpTypeVoid
180           %4 = OpTypeFunction %3
181           %1 = OpFunction %3 None %4
182           %5 = OpLabel
183                OpLine %2 3 1
184                OpReturn
185                OpFunctionEnd
186   )";
187 
188   auto[success, result] = ExtractSource(source);
189   ASSERT_TRUE(success);
190   ASSERT_TRUE(result.size() == 1);
191   ASSERT_TRUE(result["compute.hlsl"] ==
192               "[numthreads(1, 1, 1)]\nvoid compute_1() {\n}\n");
193 }
194 
TEST(ExtractSourceTest,EmptyFilename)195 TEST(ExtractSourceTest, EmptyFilename) {
196   std::string source = R"(
197                OpCapability Shader
198                OpMemoryModel Logical GLSL450
199                OpEntryPoint GLCompute %1 "compute_1"
200                OpExecutionMode %1 LocalSize 1 1 1
201           %2 = OpString ""
202                OpSource HLSL 660 %2 "void compute(){}"
203                OpName %1 "compute_1"
204           %3 = OpTypeVoid
205           %4 = OpTypeFunction %3
206           %1 = OpFunction %3 None %4
207           %5 = OpLabel
208                OpLine %2 3 1
209                OpReturn
210                OpFunctionEnd
211   )";
212 
213   auto[success, result] = ExtractSource(source);
214   ASSERT_TRUE(success);
215   ASSERT_TRUE(result.size() == 1);
216   ASSERT_TRUE(result["unnamed-0.hlsl"] == "void compute(){}");
217 }
218 
TEST(ExtractSourceTest,EscapeEscaped)219 TEST(ExtractSourceTest, EscapeEscaped) {
220   std::string source = R"(
221                OpCapability Shader
222                OpMemoryModel Logical GLSL450
223                OpEntryPoint GLCompute %1 "compute"
224                OpExecutionMode %1 LocalSize 1 1 1
225           %2 = OpString "compute.hlsl"
226                OpSource HLSL 660 %2 "// check \" escape removed"
227                OpName %1 "compute"
228           %3 = OpTypeVoid
229           %4 = OpTypeFunction %3
230           %1 = OpFunction %3 None %4
231           %5 = OpLabel
232                OpLine %2 6 1
233                OpReturn
234                OpFunctionEnd
235   )";
236 
237   auto[success, result] = ExtractSource(source);
238   ASSERT_TRUE(success);
239   ASSERT_TRUE(result.size() == 1);
240   ASSERT_TRUE(result["compute.hlsl"] == "// check \" escape removed");
241 }
242 
TEST(ExtractSourceTest,OpSourceWithNoSource)243 TEST(ExtractSourceTest, OpSourceWithNoSource) {
244   std::string source = R"(
245                OpCapability Shader
246                OpMemoryModel Logical GLSL450
247                OpEntryPoint GLCompute %1 "compute"
248                OpExecutionMode %1 LocalSize 1 1 1
249           %2 = OpString "compute.hlsl"
250                OpSource HLSL 660 %2
251                OpName %1 "compute"
252           %3 = OpTypeVoid
253           %4 = OpTypeFunction %3
254           %1 = OpFunction %3 None %4
255           %5 = OpLabel
256                OpLine %2 6 1
257                OpReturn
258                OpFunctionEnd
259   )";
260 
261   auto[success, result] = ExtractSource(source);
262   ASSERT_TRUE(success);
263   ASSERT_TRUE(result.size() == 1);
264   ASSERT_TRUE(result["compute.hlsl"] == "");
265 }
266