Pretty-print context on failures

src/HackTestCLI.hack

@@ -87,8 +87,8 @@ final class HackTestCLI extends CLIWithRequiredArguments {
     $cf = $this->classFilter;
     $mf = $this->methodFilter;
     $output = $this->verbose
-      ? new _Private\VerboseCLIOutput()
-      : new _Private\ConciseCLIOutput();
+      ? new _Private\VerboseCLIOutput($this->getTerminal())
+      : new _Private\ConciseCLIOutput($this->getTerminal());
     $stdout = $this->getStdout();
 
     await HackTestRunner::runAsync(

src/_Private/CLIOutputHandler.hack

@@ -9,14 +9,19 @@
 
 namespace Facebook\HackTest\_Private;
 
-use namespace HH\Lib\{Dict, Math, Str};
+use namespace HH\Lib\{C, Dict, Math, Str, Vec};
 use namespace HH\Lib\Experimental\IO;
 use namespace Facebook\HackTest;
 
 abstract class CLIOutputHandler {
   <<__LateInit>> private dict<HackTest\TestResult, int> $resultCounts;
   <<__LateInit>> private vec<HackTest\ErrorProgressEvent> $errors;
 
+  const int CONTEXT_LINES = 3;
+
+  public function __construct(private \Facebook\CLILib\ITerminal $terminal) {
+  }
+
   final public async function writeProgressAsync(
     <<__AcceptDisposable>> IO\WriteHandle $handle,
     \Facebook\HackTest\ProgressEvent $e,
@@ -110,4 +115,96 @@ abstract class CLIOutputHandler {
       $result_counts[HackTest\TestResult::ERROR] ?? 0,
     ));
   }
+
+  final protected function getPrettyContext(
+    \Throwable $ex,
+    string $file,
+  ): ?string {
+    $frame = $ex->getTrace()
+      |> Vec\filter(
+        $$,
+        $row ==>
+          (($row as KeyedContainer<_, _>)['file'] ?? null) as ?string === $file,
+      )
+      |> C\last($$);
+
+    if (!$frame is KeyedContainer<_, _>) {
+      return null;
+    }
+    $colors = $this->terminal->supportsColors();
+    $c_light = $colors ? "\e[2m" : '';
+    $c_bold = $colors ? "\e[1m" : '';
+    $c_red = $colors ? "\e[31m" : '';
+    $c_reset = $colors ? "\e[0m" : '';
+
+    $line = $frame['line'] as int;
+    $line_number_width = Str\length((string)$line) + 2;
+
+    $first_line = Math\maxva(1, $line - self::CONTEXT_LINES);
+    $all_lines = \file_get_contents($file)
+      |> Str\split($$, "\n");
+
+    $context_lines = Vec\slice(
+      $all_lines,
+      $first_line - 1,
+      ($line - $first_line),
+    )
+      |> Vec\map_with_key(
+        $$,
+        ($n, $content) ==> Str\format(
+          "%s| %s%s%s",
+          Str\pad_left((string)($n + $first_line), $line_number_width, ' '),
+          $c_light,
+          $content,
+          $c_reset,
+        ),
+      );
+
+    $blame_line = $all_lines[$line - 1];
+    $fun = $frame['function'] as string;
+    $fun_offset = Str\search($blame_line, $fun.'(');
+    if ($fun_offset is null && Str\contains($fun, '\\')) {
+      $fun = Str\split($fun, '\\') |> C\lastx($$);
+      $fun_offset = Str\search($blame_line, $fun.'(');
+    }
+    if (
+      $fun_offset is null && $frame['function'] === "HH\\invariant_violation"
+    ) {
+      $fun = 'invariant';
+      $fun_offset = Str\search($blame_line, 'invariant(');
+    }
+
+    if ($fun_offset is null) {
+      $context_lines[] = Str\format(
+        "%s%s>%s %s%s",
+        Str\pad_left((string)$line, $line_number_width),
+        $c_red,
+        $c_reset.$c_bold,
+        $blame_line,
+        $c_reset,
+      );
+    } else {
+      $context_lines[] = Str\format(
+        "%s%s>%s %s%s%s%s%s%s",
+        Str\pad_left((string)$line, $line_number_width),
+        $c_red,
+        $c_reset.$c_bold,
+        Str\slice($blame_line, 0, $fun_offset),
+        $c_red,
+        $fun,
+        $c_reset.$c_bold,
+        Str\slice($blame_line, $fun_offset + Str\length($fun)),
+        $c_reset,
+      );
+
+      $context_lines[] = Str\format(
+        "%s%s%s%s",
+        Str\repeat(' ', $line_number_width + $fun_offset + 2),
+        $c_red,
+        Str\repeat('^', Str\length($fun)),
+        $c_reset,
+      );
+    }
+    return $file.':'.$line."\n".Str\join($context_lines, "\n");
+  }
 }

src/_Private/ConciseCLIOutput.hack

@@ -17,6 +17,7 @@ use namespace Facebook\HackTest;
 use type Facebook\HackTest\TestResult;
 
 final class ConciseCLIOutput extends CLIOutputHandler {
+
   <<__Override>>
   protected async function writeProgressImplAsync(
     <<__AcceptDisposable>> IO\WriteHandle $handle,
@@ -63,14 +64,16 @@ final class ConciseCLIOutput extends CLIOutputHandler {
         $message = 'Skipped: '.$ex->getMessage();
       } else if ($event is HackTest\FileProgressEvent) {
         $file = $event->getPath();
-        $trace = $ex->getTraceAsString()
+
+        $context = $this->getPrettyContext($ex, $file) ??
+          $ex->getTraceAsString()
           |> Str\split($$, '#')
           |> Vec\filter($$, $line ==> Str\contains($line, $file))
           |> Vec\map($$, $line ==> Str\strip_prefix($line, '  '))
           |> Str\join($$, "\n");
 
-        if ($trace !== '') {
-          $message .= "\n\n".$trace;
+        if ($context !== '') {
+          $message .= "\n\n".$context;
         }
       }
       await $handle->writeAsync($header.$message);

src/_Private/VerboseCLIOutput.hack

@@ -102,6 +102,13 @@ final class VerboseCLIOutput extends CLIOutputHandler {
           $message .= "\n\nPrevious exception:\n\n";
         }
       }
+      if ($event is HackTest\FileProgressEvent) {
+        $file = $event->getPath();
+        $context = $this->getPrettyContext($ex, $file);
+        if ($context is nonnull) {
+          $message .= "\n\n".$context;
+        }
+      }
       await $handle->writeAsync($header.$message);
     }
   }