Using the instruction folder to fold OpSpecConstantOp (#2598)

In order to try to reduce code duplication and to be able
to fold more cases, we want to use the instruction folder
when folding an OpSpecConstantOp with constant operands.

A couple other changes are need to make this work.  First
GetDefiningInstruction| in the constant manager is able
to handle |type_id| being logically equivalent to another
type, so we updated the interface, and removed the assert.

Some tests were also updated because we not generate
better code because constants are not duplicated as much
as before.

No need for new tests.  The functionality of the instruction folder is
already tested.  There are tests check that the instruction folder is
being used correctly for OpCompositeExtract and OpVectorShuffle in the
existing test cases.

Fixes #2585.
diff --git a/source/opt/constants.cpp b/source/opt/constants.cpp
index 8dce8ec..3c05f9e 100644
--- a/source/opt/constants.cpp
+++ b/source/opt/constants.cpp
@@ -185,8 +185,6 @@
 
 Instruction* ConstantManager::GetDefiningInstruction(
     const Constant* c, uint32_t type_id, Module::inst_iterator* pos) {
-  assert(type_id == 0 ||
-         context()->get_type_mgr()->GetType(type_id) == c->type());
   uint32_t decl_id = FindDeclaredConstant(c, type_id);
   if (decl_id == 0) {
     auto iter = context()->types_values_end();
diff --git a/source/opt/constants.h b/source/opt/constants.h
index de2dfc3..5ca8518 100644
--- a/source/opt/constants.h
+++ b/source/opt/constants.h
@@ -524,12 +524,7 @@
   // instruction at the end of the current module's types section.
   //
   // |type_id| is an optional argument for disambiguating equivalent types. If
-  // |type_id| is specified, it is used as the type of the constant when a new
-  // instruction is created. Otherwise the type of the constant is derived by
-  // getting an id from the type manager for |c|.
-  //
-  // When |type_id| is not zero, the type of |c| must be the type returned by
-  // type manager when given |type_id|.
+  // |type_id| is specified, the contant returned will have that type id.
   Instruction* GetDefiningInstruction(const Constant* c, uint32_t type_id = 0,
                                       Module::inst_iterator* pos = nullptr);
 
diff --git a/source/opt/fold_spec_constant_op_and_composite_pass.cpp b/source/opt/fold_spec_constant_op_and_composite_pass.cpp
index 663d112..d61daae 100644
--- a/source/opt/fold_spec_constant_op_and_composite_pass.cpp
+++ b/source/opt/fold_spec_constant_op_and_composite_pass.cpp
@@ -120,19 +120,14 @@
 
   switch (static_cast<SpvOp>(inst->GetSingleWordInOperand(0))) {
     case SpvOp::SpvOpCompositeExtract:
-      folded_inst = DoCompositeExtract(pos);
-      break;
     case SpvOp::SpvOpVectorShuffle:
-      folded_inst = DoVectorShuffle(pos);
-      break;
-
     case SpvOp::SpvOpCompositeInsert:
-      // Current Glslang does not generate code with OpSpecConstantOp
-      // CompositeInsert instruction, so this is not implmented so far.
-      // TODO(qining): Implement CompositeInsert case.
-      return false;
-
+      folded_inst = FoldWithInstructionFolder(pos);
+      break;
     default:
+      // TODO: This should use the instruction folder as well, but some folding
+      // rules are missing.
+
       // Component-wise operations.
       folded_inst = DoComponentWiseOperation(pos);
       break;
@@ -157,54 +152,65 @@
   return subtype;
 }
 
-Instruction* FoldSpecConstantOpAndCompositePass::DoCompositeExtract(
-    Module::inst_iterator* pos) {
-  Instruction* inst = &**pos;
-  assert(inst->NumInOperands() - 1 >= 2 &&
-         "OpSpecConstantOp CompositeExtract requires at least two non-type "
-         "non-opcode operands.");
-  assert(inst->GetInOperand(1).type == SPV_OPERAND_TYPE_ID &&
-         "The composite operand must have a SPV_OPERAND_TYPE_ID type");
-  assert(
-      inst->GetInOperand(2).type == SPV_OPERAND_TYPE_LITERAL_INTEGER &&
-      "The literal operand must have a SPV_OPERAND_TYPE_LITERAL_INTEGER type");
-
-  // Note that for OpSpecConstantOp, the second in-operand is the first id
-  // operand. The first in-operand is the spec opcode.
-  uint32_t source = inst->GetSingleWordInOperand(1);
-  uint32_t type = context()->get_def_use_mgr()->GetDef(source)->type_id();
-  const analysis::Constant* first_operand_const =
-      context()->get_constant_mgr()->FindDeclaredConstant(source);
-  if (!first_operand_const) return nullptr;
-
-  const analysis::Constant* current_const = first_operand_const;
-  for (uint32_t i = 2; i < inst->NumInOperands(); i++) {
-    uint32_t literal = inst->GetSingleWordInOperand(i);
-    type = GetTypeComponent(type, literal);
-  }
-  for (uint32_t i = 2; i < inst->NumInOperands(); i++) {
-    uint32_t literal = inst->GetSingleWordInOperand(i);
-    if (const analysis::CompositeConstant* composite_const =
-            current_const->AsCompositeConstant()) {
-      // Case 1: current constant is a non-null composite type constant.
-      assert(literal < composite_const->GetComponents().size() &&
-             "Literal index out of bound of the composite constant");
-      current_const = composite_const->GetComponents().at(literal);
-    } else if (current_const->AsNullConstant()) {
-      // Case 2: current constant is a constant created with OpConstantNull.
-      // Because components of a NullConstant are always NullConstants, we can
-      // return early with a NullConstant in the result type.
-      return context()->get_constant_mgr()->BuildInstructionAndAddToModule(
-          context()->get_constant_mgr()->GetConstant(
-              context()->get_constant_mgr()->GetType(inst), {}),
-          pos, type);
-    } else {
-      // Dereferencing a non-composite constant. Invalid case.
+Instruction* FoldSpecConstantOpAndCompositePass::FoldWithInstructionFolder(
+    Module::inst_iterator* inst_iter_ptr) {
+  // If one of operands to the instruction is not a
+  // constant, then we cannot fold this spec constant.
+  for (uint32_t i = 1; i < (*inst_iter_ptr)->NumInOperands(); i++) {
+    const Operand& operand = (*inst_iter_ptr)->GetInOperand(i);
+    if (operand.type != SPV_OPERAND_TYPE_ID &&
+        operand.type != SPV_OPERAND_TYPE_OPTIONAL_ID) {
+      continue;
+    }
+    uint32_t id = operand.words[0];
+    if (context()->get_constant_mgr()->FindDeclaredConstant(id) == nullptr) {
       return nullptr;
     }
   }
-  return context()->get_constant_mgr()->BuildInstructionAndAddToModule(
-      current_const, pos);
+
+  // All of the operands are constant.  Construct a regular version of the
+  // instruction and pass it to the instruction folder.
+  std::unique_ptr<Instruction> inst((*inst_iter_ptr)->Clone(context()));
+  inst->SetOpcode(
+      static_cast<SpvOp>((*inst_iter_ptr)->GetSingleWordInOperand(0)));
+  inst->RemoveOperand(2);
+
+  // We want the current instruction to be replaced by an |OpConstant*|
+  // instruction in the same position. We need to keep track of which constants
+  // the instruction folder creates, so we can move them into the correct place.
+  auto last_type_value_iter = (context()->types_values_end());
+  --last_type_value_iter;
+  Instruction* last_type_value = &*last_type_value_iter;
+
+  auto identity_map = [](uint32_t id) { return id; };
+  Instruction* new_const_inst =
+      context()->get_instruction_folder().FoldInstructionToConstant(
+          inst.get(), identity_map);
+  assert(new_const_inst != nullptr &&
+         "Failed to fold instruction that must be folded.");
+
+  // Get the instruction before |pos| to insert after.  |pos| cannot be the
+  // first instruction in the list because its type has to come first.
+  Instruction* insert_pos = (*inst_iter_ptr)->PreviousNode();
+  assert(insert_pos != nullptr &&
+         "pos is the first instruction in the types and values.");
+  bool need_to_clone = true;
+  for (Instruction* i = last_type_value->NextNode(); i != nullptr;
+       i = last_type_value->NextNode()) {
+    if (i == new_const_inst) {
+      need_to_clone = false;
+    }
+    i->InsertAfter(insert_pos);
+    insert_pos = insert_pos->NextNode();
+  }
+
+  if (need_to_clone) {
+    new_const_inst = new_const_inst->Clone(context());
+    new_const_inst->SetResultId(TakeNextId());
+    new_const_inst->InsertAfter(insert_pos);
+    get_def_use_mgr()->AnalyzeInstDefUse(new_const_inst);
+  }
+  return new_const_inst;
 }
 
 Instruction* FoldSpecConstantOpAndCompositePass::DoVectorShuffle(
diff --git a/source/opt/fold_spec_constant_op_and_composite_pass.h b/source/opt/fold_spec_constant_op_and_composite_pass.h
index 1627125..361d3ca 100644
--- a/source/opt/fold_spec_constant_op_and_composite_pass.h
+++ b/source/opt/fold_spec_constant_op_and_composite_pass.h
@@ -54,11 +54,9 @@
   // it.
   bool ProcessOpSpecConstantOp(Module::inst_iterator* pos);
 
-  // Try to fold the OpSpecConstantOp CompositeExtract instruction pointed by
-  // the given instruction iterator to a normal constant defining instruction.
-  // Returns the pointer to the new constant defining instruction if succeeded.
-  // Otherwise returns nullptr.
-  Instruction* DoCompositeExtract(Module::inst_iterator* inst_iter_ptr);
+  // Returns the result of folding the OpSpecConstantOp instruction
+  // |inst_iter_ptr| using the instruction folder.
+  Instruction* FoldWithInstructionFolder(Module::inst_iterator* inst_iter_ptr);
 
   // Try to fold the OpSpecConstantOp VectorShuffle instruction pointed by the
   // given instruction iterator to a normal constant defining instruction.
diff --git a/test/opt/fold_spec_const_op_composite_test.cpp b/test/opt/fold_spec_const_op_composite_test.cpp
index 12f0cd1..a9e7dd3 100644
--- a/test/opt/fold_spec_const_op_composite_test.cpp
+++ b/test/opt/fold_spec_const_op_composite_test.cpp
@@ -1135,14 +1135,14 @@
                 "%outer = OpConstantComposite %outer_struct %inner %signed_one",
                 "%extract_inner = OpSpecConstantOp %inner_struct CompositeExtract %outer 0",
                 "%extract_int = OpSpecConstantOp %int CompositeExtract %outer 1",
-                "%extract_inner_float = OpSpecConstantOp %int CompositeExtract %outer 0 2",
+                "%extract_inner_float = OpSpecConstantOp %float CompositeExtract %outer 0 2",
               },
               // expected
               {
                 "%float_1 = OpConstant %float 1",
                 "%inner = OpConstantComposite %inner_struct %bool_true %signed_null %float_1",
                 "%outer = OpConstantComposite %outer_struct %inner %signed_one",
-                "%extract_inner = OpConstantComposite %flat_struct %bool_true %signed_null %float_1",
+                "%extract_inner = OpConstantComposite %inner_struct %bool_true %signed_null %float_1",
                 "%extract_int = OpConstant %int 1",
                 "%extract_inner_float = OpConstant %float 1",
               },
@@ -1256,13 +1256,9 @@
               },
               // expected
               {
-                "%60 = OpConstantNull %int",
                 "%a = OpConstantComposite %v2int %signed_null %signed_null",
-                "%62 = OpConstantNull %int",
                 "%b = OpConstantComposite %v2int %signed_zero %signed_one",
-                "%64 = OpConstantNull %int",
                 "%c = OpConstantComposite %v2int %signed_three %signed_null",
-                "%66 = OpConstantNull %int",
                 "%d = OpConstantComposite %v2int %signed_null %signed_null",
               }
             },