[Test] Added arg convenience.

docs/SIL.rst

@@ -3959,6 +3959,10 @@ The following types of test arguments are supported:
          @trace[uint] <-- the ``debug_value [trace]`` at index ``uint``
          @{function}.{trace} <-- the indicated trace in the indicated function
          Example: @function[bar].trace
+- argument: @argument <-_ the first argument of the current block
+            @argument[uint] <-- the argument at index ``uint`` of the current block
+            @{block}.{argument} <-- the indicated argument in the indicated block
+            @{function}.{argument} <-- the indicated argument in the entry block of the indicated function
 - instruction: @instruction <-- the instruction after* the test_specification instruction
                @instruction[+uint] <-- the instruction ``uint`` instructions after* the test_specification instruction
                @instruction[-uint] <-- the instruction ``uint`` instructions before* the test_specification instruction

include/swift/SILOptimizer/Utils/ParseTestSpecification.h

@@ -18,6 +18,7 @@
 #define SWIFT_SIL_PARSETESTSPECIFICATION
 
 #include "swift/Basic/TaggedUnion.h"
+#include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILValue.h"
 #include "llvm/ADT/StringRef.h"
@@ -31,6 +32,7 @@ using llvm::StringRef;
 namespace swift {
 
 class SILFunction;
+class SILArgument;
 
 namespace test {
 
@@ -42,6 +44,7 @@ struct Argument {
     Value,
     Operand,
     Instruction,
+    BlockArgument,
     Block,
     Function,
   };
@@ -51,6 +54,7 @@ struct Argument {
                             SILValue,           // ValueArgument
                             Operand *,          // OperandArgument
                             SILInstruction *,   // InstructionArgument
+                            SILArgument *,      // BlockArgumentArgument
                             SILBasicBlock *,    // BlockArgument
                             SILFunction *       // FunctionArgument
                             >;
@@ -89,6 +93,11 @@ struct InstructionArgument
   InstructionArgument(SILInstruction *stored) : Super(stored) {}
 };
 
+struct BlockArgumentArgument
+    : ConcreteArgument<SILArgument *, Argument::Kind::BlockArgument> {
+  BlockArgumentArgument(SILArgument *stored) : Super(stored) {}
+};
+
 struct BlockArgument
     : ConcreteArgument<SILBasicBlock *, Argument::Kind::Block> {
   BlockArgument(SILBasicBlock *stored) : Super(stored) {}
@@ -146,6 +155,9 @@ struct Arguments {
       auto *instruction = cast<InstructionArgument>(argument).getValue();
       auto *svi = cast<SingleValueInstruction>(instruction);
       return svi;
+    } else if (isa<BlockArgumentArgument>(argument)) {
+      auto *arg = cast<BlockArgumentArgument>(argument).getValue();
+      return arg;
     }
     return cast<ValueArgument>(argument).getValue();
   }
@@ -155,6 +167,9 @@ struct Arguments {
   SILInstruction *takeInstruction() {
     return cast<InstructionArgument>(takeArgument()).getValue();
   }
+  SILArgument *takeBlockArgument() {
+    return cast<BlockArgumentArgument>(takeArgument()).getValue();
+  }
   SILBasicBlock *takeBlock() {
     return cast<BlockArgument>(takeArgument()).getValue();
   }

lib/SILOptimizer/UtilityPasses/ParseTestSpecification.cpp

@@ -13,6 +13,7 @@
 #include "swift/SILOptimizer/Utils/ParseTestSpecification.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/SIL/SILSuccessor.h"
 #include "swift/SILOptimizer/Utils/InstructionDeleter.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
@@ -84,55 +85,75 @@ SILInstruction *getInstruction(SILBasicBlock *block, unsigned long index) {
   llvm_unreachable("bad index!?");
 }
 
+SILBasicBlock *getNextBlock(SILBasicBlock *block) {
+  auto next = std::next(block->getIterator());
+  if (next == block->getFunction()->end())
+    return nullptr;
+  return &*next;
+}
+
+SILBasicBlock *getPreviousBlock(SILBasicBlock *block) {
+  auto iterator = block->getIterator();
+  if (iterator == block->getFunction()->begin())
+    return nullptr;
+  return &*std::prev(iterator);
+}
+
+SILInstruction *getNextInstructionIgnoringBlocks(SILInstruction *instruction) {
+  if (auto *nextInstruction = instruction->getNextInstruction())
+    return nextInstruction;
+  if (auto *nextBlock = getNextBlock(instruction->getParent()))
+    return &nextBlock->front();
+  return nullptr;
+}
+
+SILInstruction *
+getPreviousInstructionIgnoringBlocks(SILInstruction *instruction) {
+  if (auto *previousInstruction = instruction->getPreviousInstruction())
+    return previousInstruction;
+  if (auto *previousBlock = getPreviousBlock(instruction->getParent()))
+    return &previousBlock->back();
+  return nullptr;
+}
+
 SILInstruction *getInstructionOffsetFrom(SILInstruction *base, long offset) {
   if (offset == 0) {
     return base;
   }
+  auto *instruction = base;
   if (offset > 0) {
-    Optional<unsigned long> baseIndex;
-    unsigned long index = 0;
-    for (auto &block : *base->getFunction()) {
-      for (auto &other : block) {
-        if (baseIndex && index == (*baseIndex + offset)) {
-          return &other;
-        }
-        if (&other == base) {
-          baseIndex = index;
-        }
-        ++index;
-      }
-    }
-    llvm_unreachable("positive offset outside of function!?");
-  }
-  SmallVector<SILInstruction *, 64> instructions;
-  unsigned long index = 0;
-  for (auto &block : *base->getFunction()) {
-    for (auto &other : block) {
-      instructions.push_back(&other);
-      if (&other == base) {
-        return instructions[index + offset];
-      }
-      ++index;
+    for (auto index = 0; index < offset; ++index) {
+      instruction = getNextInstructionIgnoringBlocks(instruction);
+      assert(instruction && "too large an offset!?");
     }
+    return instruction;
   }
-  llvm_unreachable("never found instruction in its own function!?");
+  // offset < 0
+  for (auto index = 0; index > offset; --index) {
+    instruction = getPreviousInstructionIgnoringBlocks(instruction);
+    assert(instruction && "too negative an offset!?");
+  }
+  return instruction;
 }
 
 SILBasicBlock *getBlockOffsetFrom(SILBasicBlock *base, long offset) {
   if (offset == 0)
     return base;
-  if (offset < 0) {
-    auto iterator = base->getIterator();
-    for (auto counter = 0; counter > offset; --counter) {
-      iterator = std::prev(iterator);
+  if (offset > 0) {
+    auto *block = base;
+    for (auto counter = 0; counter < offset; ++counter) {
+      block = getNextBlock(block);
+      assert(block && "too large an offset!?");
     }
-    return &*iterator;
+    return block;
   }
-  auto iterator = base->getIterator();
-  for (auto counter = 0; counter < offset; ++counter) {
-    iterator = std::next(iterator);
+  // offset < 0
+  auto *block = base;
+  for (auto counter = 0; counter > offset; --counter) {
+    block = getPreviousBlock(block);
+    assert(block && "too negative an offset!?");
   }
-  return &*iterator;
+  return block;
 }
 
 SILBasicBlock *getBlock(SILFunction *function, unsigned long index) {
@@ -321,6 +342,31 @@ class ParseArgumentSpecification {
     return OperandArgument{operand};
   }
 
+  SILArgument *parseBlockArgumentComponent(SILBasicBlock *block) {
+    if (!consumePrefix("argument"))
+      return nullptr;
+    // If this is a bare @argument reference, it refers to the first argument
+    // of the block containing the test_specification.
+    if (!block) {
+      block = context->getParent();
+    }
+    if (empty()) {
+      return block->getArgument(0);
+    }
+    if (auto subscript = parseSubscript()) {
+      auto index = subscript->get<unsigned long long>();
+      return block->getArgument(index);
+    }
+    llvm_unreachable("bad suffix after 'argument'!?");
+  }
+
+  Optional<Argument> parseBlockArgumentReference(SILBasicBlock *block) {
+    auto *argument = parseBlockArgumentComponent(block);
+    if (!argument)
+      return llvm::None;
+    return BlockArgumentArgument{argument};
+  }
+
   using InstructionContext = TaggedUnion<SILFunction *, SILBasicBlock *>;
 
   SILInstruction *
@@ -403,8 +449,10 @@ class ParseArgumentSpecification {
       return llvm::None;
     if (!consumePrefix("."))
       return BlockArgument{block};
-    if (auto arg = parseInstructionReference({block}))
+    if (auto arg = parseBlockArgumentReference(block))
       return *arg;
+    if (auto inst = parseInstructionReference({block}))
+      return *inst;
     llvm_unreachable("unhandled suffix after 'block'!?");
   }
 
@@ -444,6 +492,8 @@ class ParseArgumentSpecification {
       return llvm::None;
     if (!consumePrefix("."))
       return FunctionArgument{function};
+    if (auto arg = parseBlockArgumentReference(function->getEntryBlock()))
+      return *arg;
     if (auto arg = parseInstructionReference({function}))
       return *arg;
     if (auto arg = parseTraceReference(function))
@@ -461,6 +511,8 @@ class ParseArgumentSpecification {
     if (auto arg =
             parseOperandReference(getInstruction(context->getFunction(), 0)))
       return *arg;
+    if (auto arg = parseBlockArgumentReference(nullptr))
+      return *arg;
     if (auto arg = parseInstructionReference(llvm::None))
       return *arg;
     if (auto arg = parseBlockReference(nullptr))

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -108,6 +108,7 @@ class UnitTest {
 
 // Arguments:
 // - string: list of characters, each of which specifies subsequent arguments
+//           - A: (block) argument
 //           - F: function
 //           - B: block
 //           - I: instruction
@@ -129,6 +130,12 @@ struct TestSpecificationTest : UnitTest {
     auto expectedFields = arguments.takeString();
     for (auto expectedField : expectedFields) {
       switch (expectedField) {
+      case 'A': {
+        auto *argument = arguments.takeBlockArgument();
+        llvm::errs() << "argument:\n";
+        argument->dump();
+        break;
+      }
       case 'F': {
         auto *function = arguments.takeFunction();
         llvm::errs() << "function: " << function->getName() << "\n";

test/SILOptimizer/unit_test.sil

@@ -133,6 +133,23 @@ exit:
   return %retval : $()
 }
 
+struct X {}
+
+// CHECK-LABEL: begin running test 1 of {{[^,]+}} on test_arg_arg_parsing: test-specification-parsing
+// CHECK: argument:
+// CHECK: %0 = argument of bb0 : $X
+// CHECK: argument:
+// CHECK: %0 = argument of bb0 : $X
+// CHECK: argument:
+// CHECK: %0 = argument of bb0 : $X
+// CHECK-LABEL: end running test 1 of {{[^,]+}} on test_arg_arg_parsing: test-specification-parsing
+sil [ossa] @test_arg_arg_parsing : $(X) -> () {
+entry(%instance : $X):
+  test_specification "test-specification-parsing AAA @argument @block.argument @function.argument[0]"
+  %retval = tuple ()
+  return %retval : $()
+}
+
 class C {}
 
 sil [ossa] @getC : $@convention(thin) () -> @owned C