How REX Made Sage AST Checkpoints Round-Trip Through JSON

Posted on (Updated on )
REX now has an opt-in Sage AST JSON checkpoint tool. At selected OpenMP-facing compiler boundaries, it serializes the active SgSourceFile subtree to deterministic JSON, reads that JSON back, reconstructs a fresh Sage AST, replaces the original file inside the live SgProject, and continues the normal compiler path. The goal is not a pretty AST dump. It is a strict text-file interchange contract for a real compiler AST. That meant preserving source positions, token streams, types, scopes, symbol lookup behavior, external declaration identities, preprocessing attachments, and OpenMP/OpenACC side tables. The implementation deliberately fails hard when required state is missing instead of relying on later passes to repair the AST. Evaluation used direct AST JSON contract tests, OpenMP construction checkpoints, OpenMP lowering checkpoints, Rodinia-derived GPU offloading cases, Fortran lowering checkpoints, sanitizer runs, and focused Memcheck coverage.

Compilers are full of useful internal boundaries that are hard to make real.

In REX, one of those boundaries sits around the OpenMP path. A source file goes through Clang/Flang frontend work, Sage AST construction, OpenMP pragma parsing, OpenMP AST construction, lowering, unparsing, and sometimes generated helper files for offloading. Each step has a reasonable conceptual input and output. But the actual object moving through the compiler is not a small value. It is a pointer-rich Sage tree inside an SgProject, tied to types, scopes, symbol tables, file locations, token streams, comments, directives, hidden declarations, and transformation side tables.

That makes a simple request surprisingly hard:

1
2
pause here, write the AST to a text file, read it back, and continue as if
nothing happened

The new AST JSON checkpoint tool is REX’s first serious answer to that request. It is intentionally narrow in scope and strict in behavior. It serializes one complete SgSourceFile subtree at supported checkpoints, reconstructs a fresh Sage AST from the JSON, replaces the original source file in the live project, and then lets the rest of the compiler continue.

If the reconstructed AST cannot do the same job as the original AST, the tool is not successful. A file was written, but the compiler boundary was not real.

A diagram showing a normal REX compiler path with optional AST JSON checkpoints that serialize, deserialize, replace the live source file, and continue through OpenMP construction and lowering.

Figure 1. The checkpoint is inserted into the real compiler path. The reconstructed AST is not just inspected; it replaces the original SgSourceFile and must survive the remaining passes.

The Problem

Sage ASTs are not plain syntax trees.

A plain syntax tree can often be described as nested nodes: a function contains a body, a body contains statements, statements contain expressions, and so on. Sage has that structural tree, but REX depends on much more than parent-child shape.

For example, a usable REX AST also carries:

  • source positions and Sg_File_Info identity,
  • language and backend state on SgSourceFile,
  • type objects shared by declarations and expressions,
  • scope links and parent links,
  • symbol tables and duplicate-name lookup behavior,
  • declaration links such as first-nondefining and defining declarations,
  • token streams used by token-preserving output paths,
  • comments and preprocessing directives attached to nodes,
  • generated and external declarations used by runtime and lowering code,
  • OpenMP/OpenACC clause side tables that are not always ordinary tree edges.

That last category matters a lot for this work. OpenMP code in REX is not just ordinary C or C++ with a few comments. The parser and later construction passes attach structured information for clauses, array sections, mappers, uses_allocators, target-data behavior, and lowering state. If a checkpoint round trip loses one of those relationships, a later OpenMP pass may still see a tree, but it is no longer the same compiler input.

This is why a DOT graph or an inspection dump is not enough. DOT can help a human see a shape. It does not normally promise that every semantic relationship needed by the compiler can be reconstructed, installed back into the project, and passed through lowering.

The requirement here was stronger:

1
2
3
AST A -> JSON -> AST B

AST B must be a drop-in replacement for AST A at the chosen checkpoint.

That is the difference between an AST viewer and an interchange format.

The Motivation

The immediate motivation was to make important OpenMP compiler boundaries testable as file-based boundaries.

REX already has a working in-process compiler path. That is good for production use, but it makes architectural experiments hard. If every stage assumes the previous stage’s exact C++ object identity, then the only practical way to change the pipeline is to rewrite large parts at once. That is risky and hard to review.

A real checkpoint gives a better engineering shape:

1
2
3
4
5
run the normal compiler to a boundary
export the full AST contract
import that contract into a fresh AST
continue the normal compiler
compare the result

Once that works, a future stage can be swapped in gradually. The first stage does not need to implement a whole compiler. It only needs to consume and produce the same checkpoint contract for one boundary. REX can still run the rest of the compiler and compare behavior against the established path.

That only helps if the checkpoint contract is honest. If JSON export silently drops state and import relies on later compiler passes to patch the AST back together, the boundary becomes misleading. It would prove that the existing compiler can heal one specific broken import, not that the interchange format is complete.

So the design rule was strict:

1
missing required AST state is a serializer/deserializer bug

The fix must be in the JSON contract or reconstruction code, not in a later pass that happens to run after import.

The Analysis

The hard part was not printing JSON.

The hard part was deciding what “the AST” means at a compiler checkpoint. Several categories turned out to be essential.

First, the structural node graph has to be complete for the supported surface. Each serialized node records a stable id, Sage kind, variant tag, file location, node properties, and named edges. During deserialization, those edges are used to rebuild children, references, parent links, scopes, declarations, and source-file state.

Second, source positions are real compiler data. REX uses Sg_File_Info for more than line and column display. File ids, physical and raw filenames, compiler-generated flags, transformation flags, token flags, default-argument flags, implicit-cast markers, and source-position-unavailable state can affect later behavior. The JSON format records that as structured data instead of collapsing it to a string.

Third, types and symbols cannot be treated as incidental pointers. Expressions and declarations need their type shapes reconstructed. Scopes need symbol tables. Duplicate lookup groups need to return the same preferred symbol after import that they returned before export. That is especially important in a compiler that handles C++, templates, namespace reopening, overloaded names, and generated declarations.

Fourth, not every required declaration is structurally inside the source-file subtree. Some relationships point to runtime functions, modules, or external declaration peers. The solution is not to print raw pointer addresses. The JSON schema has explicit external records for those cases, including enough identity and parameter-scope information to reconstruct the relationship safely.

Fifth, token and preprocessing state matters. A compiler that supports token-preserving paths has to remember tokens. A compiler that keeps comments and preprocessing directives attached to AST nodes has to serialize those attachments. Losing them might not break a minimal AST traversal, but it breaks the source-to-source compiler contract.

Finally, OpenMP and OpenACC bring side-table state. Array sections, iterator clauses, mapper clauses, uses-allocators clauses, target-data policies, and lowering information all have to survive if the checkpoint is going to be useful near OpenMP construction and lowering.

A diagram showing the Sage AST JSON contract as structural nodes, source locations, types, symbols, token and preprocessing data, external records, and OpenMP side tables.

Figure 2. The JSON contract is larger than a tree dump. It has to preserve the compiler relationships that later passes actually query.

The Solution

The tool lives under:

1
src/frontend/SageIII/astJson/

It is enabled only when requested. Users can choose checkpoints through command line options:

1
2
-rex:ast-json-checkpoint=<checkpoint>
-rex:ast-json-dir=<directory>

or with environment variables:

1
2
REX_AST_JSON_CHECKPOINT=<checkpoint>
REX_AST_JSON_DIR=<directory>

The supported checkpoint names are:

1
2
3
4
pre-omp-construction
post-omp-construction
post-omp-lowering
all

At each reached checkpoint, the tool performs the same sequence:

  1. Serialize the active SgSourceFile to JSON.
  2. Parse that JSON and validate the format, schema, root kind, node table, and required fields.
  3. Reconstruct a fresh SgSourceFile.
  4. Replace the old source file in the owning SgProject.
  5. Serialize the replacement again and compare canonical semantic signatures.
  6. Continue the normal compiler path with the reconstructed source file.

The current top-level JSON shape is deliberately explicit:

1
2
3
4
5
6
7
8
9
{
  "format": "rex-sage-ast-json",
  "schema_version": 27,
  "root_id": 1,
  "root_kind": "SgSourceFile",
  "node_count": 0,
  "metadata": {},
  "nodes": []
}

The schema version is not decorative. The reader rejects unsupported versions. The root kind must be SgSourceFile. The node count must match the node table. The format marker must match. This is a compiler interchange file, so accepting “close enough” input would be a bug.

Each node record has stable internal ids:

1
2
3
4
5
6
7
8
9
{
  "id": 1,
  "kind": "SgSourceFile",
  "variant": 0,
  "flags": {},
  "location": {},
  "properties": {},
  "edges": []
}

Those ids are local to the JSON file. They are used for references and edges, not as process addresses. When a reconstructed node is serialized again for validation, preserved ids are reused so comparison detects real semantic drift instead of allocator-order noise.

The implementation is split into conventional C++ translation units:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
sageAstJsonApi.cpp
sageAstJsonCollect.cpp
sageAstJsonDeserializeFinalize.cpp
sageAstJsonDeserializeNodes.cpp
sageAstJsonDeserializeSymbols.cpp
sageAstJsonDeserializeTypes.cpp
sageAstJsonFormat.cpp
sageAstJsonProject.cpp
sageAstJsonSerializeExternal.cpp
sageAstJsonSerializeNodes.cpp
sageAstJsonSerializeSupport.cpp

That split matters for maintainability. The serializer/deserializer is too large to live as one monolithic file, and it should not be hidden behind .inc implementation fragments. Each part has a clear responsibility: JSON parsing and formatting, node collection, external records, node serialization, deserialization, symbol-table reconstruction, type reconstruction, project replacement, and final validation.

Checkpoint files are written atomically. The tool writes a per-process temporary file first and then publishes it with rename, so concurrent compiler invocations sharing a checkpoint directory do not see partial JSON files or clobber each other’s active output.

When validation fails, diagnostic artifacts are written next to the checkpoint file: the original JSON, reconstructed JSON, original canonical signature, and reconstructed canonical signature. Then the compiler fails hard.

That failure mode is intentional. A bad checkpoint should be loud.

What Had To Be Fixed Around It

Building this tool exposed gaps that normal tests did not always catch.

Some were directly in the AST JSON surface: detached external function identity, external parameter scopes, source-file token lists, duplicate symbol lookup preference, and type fixup behavior for reconstructed source files.

Others were ordinary compiler bugs or under-tested behavior made visible by round-tripping and by stricter validation:

  • optional OpenMP array-section bounds,
  • uses_allocators optional edges and per-entry isolation,
  • declare mapper typedef and template unparsing,
  • anonymous member name qualification,
  • OpenMP expression lookup restricted to the current file,
  • deterministic outlining variable-symbol ordering,
  • deterministic OpenMP root lowering order,
  • symbol-table rebuild coverage,
  • Clang frontend reads that needed to be made Valgrind-defined at the Clang AST boundary.

Those fixes were not papered over inside the JSON tool. They received targeted rex_test2026_* coverage so the behavior is tested even when AST JSON is not enabled.

That is an important lesson from this work: a checkpoint serializer is a good way to discover missing compiler contracts. But once it discovers them, the fix still belongs at the real root cause.

Evaluation

The evaluation had three layers.

The first layer is direct contract testing for the tool itself. The astJsonTests cases exercise reconstruction details that are easy to miss, such as external function parameter scopes and detached file type fixups.

The second layer is checkpoint comparison. OpenMP frontend tests now run normal compiler output and checkpointed output side by side, then compare the result. The same idea is used around OpenMP lowering, including mapper lowering, deterministic lowering checks, Rodinia-derived offloading cases, and Fortran lowering checkpoints.

The third layer is instrumented validation. The branch was tested under regular CTest, sanitizer builds, focused Memcheck runs, and targeted Valgrind-defined read cases for the Clang frontend boundary.

The large validation runs from the hardening branch were:

ValidationResult
Full regular CTest in the AST JSON build35360/35360 passed
Broad OpenMP/OpenACC/OpenMP-lowering checkpoint suite1152/1152 passed
Full sanitizer CTest with ASan, LSan, and UBSan enabled35360/35360 passed
Focused post-fix Valgrind-defined read testspassed
Final PR pre-push regular suite in the active merge checkout5387/5387 passed

The exact CTest count depends on the build configuration, enabled languages, and instrumentation mode. The important point is not the absolute number. It is that the checkpoint path was validated both as a direct serializer/deserializer and as a real compiler replacement point inside OpenMP construction and lowering.

A ladder diagram showing direct AST JSON contract tests, OpenMP checkpoint comparisons, lowering checkpoint comparisons, sanitizer validation, and Memcheck validation.

Figure 3. The tool was evaluated as a compiler boundary, not just as a JSON writer.

Why This Matters

This work gives REX a new way to reason about compiler stages.

Before this, many boundaries were conceptual. We could say “this pass consumes the AST after OpenMP construction” or “this pass transforms the AST before lowering,” but the data never had to survive leaving the process. Now a boundary can be tested as a file:

1
2
3
4
export complete AST state
import complete AST state
continue the compiler
compare behavior

That does not make the whole compiler modular overnight. It does make one important kind of modularity measurable. A stage is no longer “compatible” because it seems to call the same C++ functions. It is compatible when it can produce a checkpoint file that REX can import and run through the remaining tests without changing behavior.

The text-file part is also useful. JSON is not the fastest possible format, and it is not meant to be the final answer for every compiler storage problem. But for this job it has the right first property: it is inspectable, diffable, schema-versioned, and language-neutral. When a reconstruction mismatch happens, engineers can look at the file and the canonical signature artifacts instead of debugging only through live pointers.

Conclusion

The Sage AST JSON checkpoint tool is not an AST pretty-printer.

It is a strict interchange path for a real REX SgSourceFile at selected compiler checkpoints. It writes the AST to text, reads it back, reconstructs a new Sage AST, installs that AST into the live project, and demands that the rest of the compiler behave the same.

That is a hard bar, but it is the right one. Anything weaker would risk creating a file format that looks useful while quietly depending on the original in-process AST model.

The useful outcome is not just the tool. The useful outcome is the discipline it forces: explicit schema, deterministic serialization, strict reconstruction, targeted regression tests, and no hidden repair pass after import.

That gives REX a real checkpoint boundary to build on.