HLL::Compiler.compile: balance backend per-compile state on early exit

HLL::Backend.start (MoarVM) pushes a frame list onto
%COMPILING<moar><frames> at every compile's 'start' stage. The matching
pop runs at the 'mast' stage in QASTCompilerMAST.to_mast. A compile
that stops before 'mast' (e.g. with `:target('ast')`), or one whose
stage call throws, leaves the push dangling. The next nested compile's
MAST migration then attaches its frames to the dangling level rather
than the outer compile's, leaving those MAST::Frames unreachable from
any compunit's @!frames. The outer's deserialization code dies later
with `Could not find frame <unit>`.

Snapshot the backend's per-compile state before the stages loop and
rebalance it afterwards via two optional backend hooks
(compile_snapshot, compile_cleanup). Wrap each stage call in CATCH so
the cleanup also runs when a stage throws, capture the exception, exit
the stages loop, then rethrow with its original payload.

The legacy Rakudo frontend never stops short of 'mast' and never
throws across the stages loop boundary, so on every existing
non-RakuAST compile the snapshot matches the post-loop depth, the
cleanup pop loop runs no iterations, and the caught variable stays
null.

Author: ugexe

src/HLL/Compiler.nqp

@@ -503,6 +503,17 @@ class HLL::Compiler does HLL::Backend::Default {
             nqp::unshift(@stages, 'start');
         }
 
+        # Backends may keep per-compile state on %*COMPILING (e.g. a frame
+        # list that the 'start' stage pushes for the 'mast' stage to pop).
+        # Snapshot that state so compile_cleanup can rebalance it if the
+        # stages loop exits before the consuming stage runs, whether via a
+        # 'target' adverb shorter than the final stage or an in-flight
+        # exception.
+        my $compile_snapshot := nqp::can($!backend, 'compile_snapshot')
+            ?? $!backend.compile_snapshot
+            !! NQPMu;
+
+        my $caught := nqp::null;
         for @stages {
             $stderr.print(nqp::sprintf("Stage %-11s: ", [$_])) if nqp::defined($stagestats) && $_ ne "parse";
             my int $timestamp := nqp::time();
@@ -511,9 +522,22 @@ class HLL::Compiler does HLL::Backend::Default {
                 self.execute_stage($_, $result, %adverbs)
             }
 
-            $result := %adverbs<profile-stage> eq $_
-                ?? $!backend.run_profiled(&run, %adverbs<profile-filename> || %adverbs<profile-compile>, %adverbs<profile-kind>)
-                !! run();
+            # Only the stage call can throw something the cleanup below
+            # needs to run for. Wrap just that call so stagestats output
+            # and other bookkeeping stay outside the handler scope.
+            # The bare CATCH also matches CONTROL-category exceptions.
+            # Rakudo's resumable warnings (CX::Warn and friends) are
+            # absorbed by HLL-level handlers installed deeper in the
+            # call chain before they reach the stage call site, so they
+            # never bubble into this CATCH at runtime. If that ever
+            # changes, this scope would have to filter the category.
+            try {
+                $result := %adverbs<profile-stage> eq $_
+                    ?? $!backend.run_profiled(&run, %adverbs<profile-filename> || %adverbs<profile-compile>, %adverbs<profile-kind>)
+                    !! run();
+                CATCH { $caught := $_ }
+            }
+            last if !nqp::isnull($caught);
 
             my num $diff := nqp::div_n(nqp::time() - $timestamp,1000000000e0);
             if nqp::defined($stagestats) {
@@ -532,6 +556,11 @@ class HLL::Compiler does HLL::Backend::Default {
             last if $_ eq $target;
         }
 
+        nqp::can($!backend, 'compile_cleanup')
+            && $!backend.compile_cleanup($compile_snapshot);
+
+        nqp::rethrow($caught) if !nqp::isnull($caught);
+
         if %adverbs<compunit_ok> {
             return $result
         }

src/vm/moar/HLL/Backend.nqp

@@ -761,6 +761,10 @@ class HLL::Backend::MoarVM {
         0
     }
 
+    # Push a fresh entry onto %COMPILING<moar><frames>. The matching pop
+    # lives in QASTCompilerMAST.to_mast at the 'mast' stage. The safety
+    # net for a stages loop that exits before 'mast' lives in
+    # compile_snapshot / compile_cleanup below.
     method start($source, *%adverbs) {
         if nqp::existskey(%*COMPILING, 'moar') {
             nqp::push(%*COMPILING<moar><frames>, nqp::list);
@@ -780,6 +784,24 @@ class HLL::Backend::MoarVM {
         $source
     }
 
+    # compile_snapshot records the depth of %COMPILING<moar><frames>
+    # before the stages loop runs. compile_cleanup pops anything still
+    # above that mark, restoring the depth so an early-exit 'target' or
+    # in-flight exception cannot leave a partial frame entry behind for
+    # the next nested compile. Pairs with `start` above (the push) and
+    # QASTCompilerMAST.to_mast (the pop).
+    method compile_snapshot() {
+        nqp::existskey(%*COMPILING, 'moar')
+            ?? nqp::elems(nqp::atkey(nqp::atkey(%*COMPILING, 'moar'), 'frames'))
+            !! 0
+    }
+
+    method compile_cleanup($snapshot) {
+        return 0 unless nqp::existskey(%*COMPILING, 'moar');
+        my $frames := nqp::atkey(nqp::atkey(%*COMPILING, 'moar'), 'frames');
+        nqp::pop($frames) while nqp::elems($frames) > $snapshot;
+    }
+
     method mast($qast, *%adverbs) {
         nqp::getcomp('QAST').to_mast($qast);
     }

src/vm/moar/QAST/QASTCompilerMAST.nqp

@@ -358,6 +358,10 @@ my class MASTCompilerInstance {
             $annotations := %moar<annotations> := MAST::Bytecode.new;
         }
 
+        # Pops the entry HLL::Backend::MoarVM.start pushed for this
+        # compile. HLL::Compiler.compile runs compile_cleanup on the
+        # backend as a safety net for stages loops that exit before
+        # reaching us.
         my @frames := nqp::pop(%moar<frames>);
         $!writer := MoarVM::BytecodeWriter.new(:$string-heap, :$callsites, :$annotations);
         $!mast_compunit := MAST::CompUnit.new(

t/moar/55-compile-frame-balance.t

@@ -0,0 +1,72 @@
+plan(4);
+
+my $comp := nqp::getcomp('nqp');
+my $backend := $comp.backend;
+
+# HLL::Backend on MoarVM pushes a per-compile frame list onto
+# %COMPILING<moar><frames> at the 'start' stage and pops it at the
+# 'mast' stage (QASTCompilerMAST.to_mast). A compile that stops short
+# of 'mast' (e.g. `:target('ast')`), or one whose stages unwind via an
+# exception, leaves the push dangling. HLL::Compiler.compile rebalances
+# via compile_snapshot / compile_cleanup hooks on the backend.
+
+ok(nqp::can($backend, 'compile_snapshot'),
+    'MoarVM backend exposes compile_snapshot');
+ok(nqp::can($backend, 'compile_cleanup'),
+    'MoarVM backend exposes compile_cleanup');
+
+# Drive an outer compile whose BEGIN body runs several :target('ast')
+# compiles. Inner and outer share %*COMPILING. If the inner compiles
+# leak their start() push, the depth grows. The rebalance must keep
+# pre- and post-depth equal.
+my $caught := '';
+try {
+    $comp.compile(q|BEGIN {
+        my $inner := nqp::getcomp("nqp");
+        my $before := nqp::elems(%*COMPILING<moar><frames>);
+        $inner.compile("1", :target("ast"), :compunit_ok(1));
+        $inner.compile("2", :target("ast"), :compunit_ok(1));
+        $inner.compile("3", :target("ast"), :compunit_ok(1));
+        $inner.compile("4", :target("ast"), :compunit_ok(1));
+        $inner.compile("5", :target("ast"), :compunit_ok(1));
+        my $after := nqp::elems(%*COMPILING<moar><frames>);
+        if $before != $after {
+            nqp::die(":target('ast') leaked " ~
+                ($after - $before) ~ " frame level(s) onto " ~
+                "%COMPILING<moar><frames>");
+        }
+    };|);
+    CATCH { $caught := ~$!; }
+}
+is($caught, '',
+    "multiple :target('ast') compiles inside an outer compile's BEGIN " ~
+    "leave %COMPILING<moar><frames> balanced");
+
+# Stage exception path: an inner compile that throws during 'parse'
+# must run compile_cleanup before the rethrow, so the outer caller sees
+# the original exception and a balanced frame stack.
+my $caught2 := '';
+try {
+    $comp.compile(q|BEGIN {
+        my $inner := nqp::getcomp("nqp");
+        my $before := nqp::elems(%*COMPILING<moar><frames>);
+        my $threw := 0;
+        try {
+            $inner.compile("\"unterminated", :target("parse"), :compunit_ok(1));
+            CATCH { $threw := 1 }
+        }
+        my $after := nqp::elems(%*COMPILING<moar><frames>);
+        unless $threw {
+            nqp::die("inner compile did not throw on invalid source");
+        }
+        if $before != $after {
+            nqp::die("exception during 'parse' stage leaked " ~
+                ($after - $before) ~ " frame level(s) onto " ~
+                "%COMPILING<moar><frames>");
+        }
+    };|);
+    CATCH { $caught2 := ~$!; }
+}
+is($caught2, '',
+    "an exception during a stage runs compile_cleanup and rethrows " ~
+    "the original exception");