Improve decoding of constant composites (#885)

* Improve decoding of constant composites

Signed-off-by: Alexey Sotkin <[email protected]>

Author: AlexeySotkin

lib/SPIRV/libSPIRV/SPIRVStream.cpp

@@ -296,4 +296,29 @@ spv_ostream &operator<<(spv_ostream &O, const SPIRVNL &E) {
   return O;
 }
 
+// Read the next word from the stream and if OpCode matches the argument,
+// decode the whole instruction. Multiple such instructions are possible. If
+// OpCode doesn't match the argument, set position of the next character to be
+// extracted from the stream to the beginning of the non-matching instruction.
+// Returns vector of extracted instructions.
+// Used to decode SPIRVTypeStructContinuedINTEL,
+// SPIRVConstantCompositeContinuedINTEL and
+// SPIRVSpecConstantCompositeContinuedINTEL.
+std::vector<SPIRVEntry *>
+SPIRVDecoder::getContinuedInstructions(const spv::Op ContinuedOpCode) {
+  std::vector<SPIRVEntry *> ContinuedInst;
+  std::streampos Pos = IS.tellg(); // remember position
+  getWordCountAndOpCode();
+  while (OpCode == ContinuedOpCode) {
+    SPIRVEntry *Entry = getEntry();
+    assert(Entry && "Failed to decode entry! Invalid instruction!");
+    M.add(Entry);
+    ContinuedInst.push_back(Entry);
+    Pos = IS.tellg();
+    getWordCountAndOpCode();
+  }
+  IS.seekg(Pos); // restore position
+  return ContinuedInst;
+}
+
 } // namespace SPIRV

lib/SPIRV/libSPIRV/SPIRVStream.h

@@ -76,6 +76,8 @@ class SPIRVDecoder {
   void validate() const;
   void ignore(size_t N);
   void ignoreInstruction();
+  std::vector<SPIRVEntry *>
+  getContinuedInstructions(const spv::Op ContinuedOpCode);
 
   std::istream &IS;
   SPIRVModule &M;

lib/SPIRV/libSPIRV/SPIRVType.h

@@ -700,18 +700,8 @@ class SPIRVTypeStruct : public SPIRVType {
     Decoder >> Id >> MemberTypeIdVec;
     Module->add(this);
 
-    Decoder.getWordCountAndOpCode();
-    while (!I.eof()) {
-      SPIRVEntry *Entry = Decoder.getEntry();
-      if (Entry != nullptr)
-        Module->add(Entry);
-      if (Entry && Decoder.OpCode == ContinuedOpCode) {
-        auto ContinuedInst = static_cast<ContinuedInstType>(Entry);
-        addContinuedInstruction(ContinuedInst);
-        Decoder.getWordCountAndOpCode();
-      } else {
-        break;
-      }
+    for (SPIRVEntry *E : Decoder.getContinuedInstructions(ContinuedOpCode)) {
+      addContinuedInstruction(static_cast<ContinuedInstType>(E));
     }
   }
 

lib/SPIRV/libSPIRV/SPIRVValue.h

@@ -344,18 +344,8 @@ template <spv::Op OC> class SPIRVConstantCompositeBase : public SPIRVValue {
     SPIRVDecoder Decoder = getDecoder(I);
     Decoder >> Type >> Id >> Elements;
 
-    Decoder.getWordCountAndOpCode();
-    while (!I.eof()) {
-      SPIRVEntry *Entry = Decoder.getEntry();
-      if (Entry != nullptr)
-        Module->add(Entry);
-      if (Entry && Decoder.OpCode == ContinuedOpCode) {
-        auto ContinuedInst = static_cast<ContinuedInstType>(Entry);
-        addContinuedInstruction(ContinuedInst);
-        Decoder.getWordCountAndOpCode();
-      } else {
-        break;
-      }
+    for (SPIRVEntry *E : Decoder.getContinuedInstructions(ContinuedOpCode)) {
+      addContinuedInstruction(static_cast<ContinuedInstType>(E));
     }
   }