[Test] Allowed referring to block arguments.

Author: nate-chandler

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

@@ -342,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 *
@@ -424,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'!?");
   }
 
@@ -465,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))
@@ -482,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