[Test] Allow specifying blocks by offset.

Author: nate-chandler

docs/SIL.rst

@@ -3950,6 +3950,8 @@ The following types of test arguments are supported:
             @function[uint] <-- function at index ``uint``
             @function[name] <-- function named ``name``
 - block: @block <-- the block containing the test_specification instruction
+         @block[+uint] <-- the block ``uint`` blocks after the containing block
+         @block[-uint] <-- the block ``uint`` blocks before the containing block
          @block[uint] <-- the block at index ``uint``
          @{function}.{block} <-- the indicated block in the indicated function
          Example: @function[foo].block[2]

lib/SILOptimizer/UtilityPasses/ParseTestSpecification.cpp

@@ -118,6 +118,23 @@ SILInstruction *getInstructionOffsetFrom(SILInstruction *base, long offset) {
   llvm_unreachable("never found instruction in its own function!?");
 }
 
+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);
+    }
+    return &*iterator;
+  }
+  auto iterator = base->getIterator();
+  for (auto counter = 0; counter < offset; ++counter) {
+    iterator = std::next(iterator);
+  }
+  return &*iterator;
+}
+
 SILBasicBlock *getBlock(SILFunction *function, unsigned long index) {
   auto iterator = function->begin();
   for (unsigned long i = 0; i < index; ++i) {
@@ -217,9 +234,28 @@ class ParseArgumentSpecification {
     return retval;
   }
 
-  Optional<TaggedUnion<unsigned long long, StringRef>> parseSubscript() {
+  Optional<TaggedUnion<unsigned long long, long long, StringRef>>
+  parseSubscript() {
     if (!consumePrefix("["))
       return llvm::None;
+    if (peekPrefix("+") || peekPrefix("-")) {
+      bool positive = false;
+      if (consumePrefix("+")) {
+        positive = true;
+      } else {
+        bool consumed = consumePrefix("-");
+        assert(consumed && "peeked a +/- but can't consume one!?");
+        (void)consumed;
+        positive = false;
+      }
+      unsigned long long index = 0;
+      bool consumed =
+          llvm::consumeUnsignedInteger(specification, /*radix=*/10, index);
+      assert(!consumed && "didn't find uint after sign!?");
+      demandPrefix("]");
+      long long offset = positive ? index : -index;
+      return {offset};
+    }
     unsigned long long index = 0;
     if (!llvm::consumeUnsignedInteger(specification, /*radix=*/10, index)) {
       demandPrefix("]");
@@ -308,6 +344,13 @@ class ParseArgumentSpecification {
       return getInstructionAtIndex(0, *within);
     }
     if (auto subscript = parseSubscript()) {
+      // If this is a bare @instruction[...] reference, it can refer either to
+      // an instruction counting from the beginning of the function or else to
+      // an instruction offset from the context of the test_specification.
+      if (!within && subscript->isa<long long>()) {
+        auto offset = subscript->get<long long>();
+        return getInstructionOffsetFrom(context, offset);
+      }
       auto index = subscript->get<unsigned long long>();
       return getInstructionAtIndex(index,
                                    within.getValueOr(context->getFunction()));
@@ -340,11 +383,15 @@ class ParseArgumentSpecification {
       return block;
     }
     if (auto subscript = parseSubscript()) {
-      if (!within) {
-        within = context->getFunction();
+      // If this is a bare @block[...] reference, it can refer either to a block
+      // counting from the beginning of the function or else to a block offset
+      // from the block containing the context of the test_specification.
+      if (!within && subscript->isa<long long>()) {
+        return getBlockOffsetFrom(context->getParent(),
+                                  subscript->get<long long>());
       }
       auto index = subscript->get<unsigned long long>();
-      auto *block = getBlock(within, index);
+      auto *block = getBlock(within ? within : context->getFunction(), index);
       return block;
     }
     llvm_unreachable("bad suffix after 'block'!?");

test/SILOptimizer/unit_test.sil

@@ -104,8 +104,18 @@ bb1:
 // CHECK: instruction: return
 // CHECK-LABEL: end running test 3 of {{[^,]+}} on test_contextual_arg_parsing: test-specification-parsing
 // CHECK-LABEL: begin running test 4 of {{[^,]+}} on test_contextual_arg_parsing: test-specification-parsing
-// CHECK: function: test_contextual_arg_parsing
+// CHECK: block
+// CHECK: bb0
+// CHECK: block
+// CHECK: bb1
+// CHECK: block
+// CHECK: bb2
+// CHECK: block
+// CHECK: bb3
 // CHECK-LABEL: end running test 4 of {{[^,]+}} on test_contextual_arg_parsing: test-specification-parsing
+// CHECK-LABEL: begin running test 5 of {{[^,]+}} on test_contextual_arg_parsing: test-specification-parsing
+// CHECK: function: test_contextual_arg_parsing
+// CHECK-LABEL: end running test 5 of {{[^,]+}} on test_contextual_arg_parsing: test-specification-parsing
 sil [ossa] @test_contextual_arg_parsing : $() -> () {
 entry:
   br one
@@ -115,6 +125,7 @@ one:
 two:
   test_specification "test-specification-parsing I @instruction"
   test_specification "test-specification-parsing IIII @instruction[-1] @instruction[+0] @instruction[+1] @instruction[+2]"
+  test_specification "test-specification-parsing BBBB @block[-2] @block[-1] @block[+0] @block[+1]"
   %retval = tuple ()
   br exit
 exit: