chore: Add tokens and source files to the serialized AST

Author: robertoraggi

packages/cxx-gen-ast/src/gen_ast_decoder_cc.ts

@@ -48,14 +48,14 @@ export function gen_ast_decoder_cc({
     const className = makeClassName(base);
     emit();
     emit(
-      `  auto ASTDecoder::decode${className}(const void* ptr, io::${className} type) -> ${base}* {`,
+      `  auto ASTDecoder::decode${className}(const void* ptr, io::${className} type) -> ${base}* {`
     );
     emit(`  switch (type) {`);
     nodes.forEach(({ name }) => {
       const className = makeClassName(name);
       emit(`  case io::${baseClassName}_${className}:`);
       emit(
-        `    return decode${className}(reinterpret_cast<const io::${className}*>(ptr));`,
+        `    return decode${className}(reinterpret_cast<const io::${className}*>(ptr));`
       );
     });
     emit(`  default:`);
@@ -70,7 +70,7 @@ export function gen_ast_decoder_cc({
       const className = makeClassName(name);
       emit();
       emit(
-        `  auto ASTDecoder::decode${className}(const io::${className}* node) -> ${name}* {`,
+        `  auto ASTDecoder::decode${className}(const io::${className}* node) -> ${name}* {`
       );
       emit(`  if (!node) return nullptr;`);
       emit();
@@ -135,6 +135,8 @@ export function gen_ast_decoder_cc({
         } else if (m.kind == "attribute" && m.type === "TokenKind") {
           emit(`  ast->${m.name} = static_cast<TokenKind>(`);
           emit(`    node->${snakeName}());`);
+        } else if (m.kind == "token") {
+          emit(`  ast->${m.name} = SourceLocation(node->${snakeName}());`);
         }
       });
       emit(`  return ast;`);

packages/cxx-gen-ast/src/gen_ast_encoder_cc.ts

@@ -172,9 +172,8 @@ export function gen_ast_encoder_cc({
             emit(`    static_cast<std::uint32_t>(ast->${m.name}));`);
           });
         } else if (m.kind == "token") {
-          emit(`  auto ${m.name} = encodeSourceLocation(ast->${m.name});`);
           finalizers.push(() => {
-            emit(`  builder.add_${fieldName}(${m.name}.o);`);
+            emit(`  builder.add_${fieldName}(ast->${m.name}.index());`);
           });
         }
       });
@@ -201,8 +200,9 @@ export function gen_ast_encoder_cc({
 #include <cxx/literals.h>
 #include <cxx/names.h>
 #include <cxx/translation_unit.h>
-#include <format>
+#include <cxx/preprocessor.h>
 
+#include <format>
 #include <algorithm>
 
 namespace cxx {
@@ -214,17 +214,32 @@ auto ASTEncoder::operator()(TranslationUnit* unit) -> std::span<const std::uint8
   Table<StringLiteral> stringLiterals;
   Table<IntegerLiteral> integerLiterals;
   Table<FloatLiteral> floatLiterals;
-  SourceFiles sourceFiles;
-  SourceLines sourceLines;
 
   std::swap(unit_, unit);
   std::swap(identifiers_, identifiers);
   std::swap(charLiterals_, charLiterals);
   std::swap(stringLiterals_, stringLiterals);
   std::swap(integerLiterals_, integerLiterals);
   std::swap(floatLiterals_, floatLiterals);
-  std::swap(sourceFiles_, sourceFiles);
-  std::swap(sourceLines_, sourceLines);
+
+  std::vector<flatbuffers::Offset<io::Source>> sources;
+  for (const auto& source : unit_->preprocessor()->sources()) {
+    auto file_name = fbb_.CreateString(source.fileName);
+    std::vector<int> lineOffsets(source.lineOffsets.begin(),
+                                 source.lineOffsets.end());
+    auto line_offsets = fbb_.CreateVector(lineOffsets);
+    sources.push_back(io::CreateSource(fbb_, file_name, line_offsets));
+  }
+
+  auto source_list = fbb_.CreateVector(sources);
+
+  std::vector<std::uint64_t> tokens;
+  for (std::uint32_t i = 0; i < unit_->tokenCount(); ++i) {
+    const auto& token = unit_->tokenAt(SourceLocation(i));
+    tokens.push_back(token.raw());
+  }
+
+  auto token_list = fbb_.CreateVector(tokens);
 
   auto [unitOffset, unitType] = acceptUnit(unit_->ast());
 
@@ -234,15 +249,15 @@ auto ASTEncoder::operator()(TranslationUnit* unit) -> std::span<const std::uint8
   builder.add_unit(unitOffset);
   builder.add_unit_type(static_cast<io::Unit>(unitType));
   builder.add_file_name(file_name);
+  builder.add_source_list(source_list);
+  builder.add_token_list(token_list);
 
   std::swap(unit_, unit);
   std::swap(identifiers_, identifiers);
   std::swap(charLiterals_, charLiterals);
   std::swap(stringLiterals_, stringLiterals);
   std::swap(integerLiterals_, integerLiterals);
   std::swap(floatLiterals_, floatLiterals);
-  std::swap(sourceFiles_, sourceFiles);
-  std::swap(sourceLines_, sourceLines);
 
   fbb_.Finish(builder.Finish(), io::SerializedUnitIdentifier());
 
@@ -258,47 +273,6 @@ auto ASTEncoder::accept(AST* ast) -> flatbuffers::Offset<> {
   return offset;
 }
 
-auto ASTEncoder::encodeSourceLocation(const SourceLocation& loc)
-    -> flatbuffers::Offset<> {
-  if (!loc) {
-    return {};
-  }
-
-  const auto start = unit_->tokenStartPosition(loc);
-
-  flatbuffers::Offset<io::SourceLine> sourceLineOffset;
-
-  auto key = std::tuple(start.fileName, start.line);
-
-  if (sourceLines_.contains(key)) {
-    sourceLineOffset = sourceLines_.at(key).o;
-  } else {
-    flatbuffers::Offset<flatbuffers::String> fileNameOffset;
-
-    if (sourceFiles_.contains(start.fileName)) {
-      fileNameOffset = sourceFiles_.at(start.fileName);
-    } else {
-      fileNameOffset = fbb_.CreateString(start.fileName);
-      sourceFiles_.emplace(start.fileName, fileNameOffset.o);
-    }
-
-    io::SourceLineBuilder sourceLineBuilder{fbb_};
-    sourceLineBuilder.add_file_name(fileNameOffset);
-    sourceLineBuilder.add_line(start.line);
-    sourceLineOffset = sourceLineBuilder.Finish();
-    sourceLines_.emplace(std::move(key), sourceLineOffset.o);
-  }
-
-  io::SourceLocationBuilder sourceLocationBuilder{fbb_};
-  sourceLocationBuilder.add_source_line(sourceLineOffset);
-  sourceLocationBuilder.add_column(start.column);
-
-  auto offset = sourceLocationBuilder.Finish();
-
-  return offset.Union();
-}
-
-
 ${code.join("\n")}
 
 } // namespace cxx

packages/cxx-gen-ast/src/gen_ast_encoder_h.ts

@@ -43,22 +43,12 @@ export function gen_ast_encoder_h({
   emit(`  using Table = std::unordered_map<const T*,`);
   emit(`    flatbuffers::Offset<flatbuffers::String>>;`);
   emit();
-  emit(`  using SourceFiles = std::unordered_map<`);
-  emit(`    std::string_view,`);
-  emit(`    flatbuffers::Offset<flatbuffers::String>>;`);
-  emit();
-  emit(`  using SourceLines = std::map<`);
-  emit(`    std::tuple<std::string_view, std::uint32_t>,`);
-  emit(`    flatbuffers::Offset<flatbuffers::String>>;`);
-  emit();
   emit(`  TranslationUnit* unit_ = nullptr;`);
   emit(`  Table<Identifier> identifiers_;`);
   emit(`  Table<CharLiteral> charLiterals_;`);
   emit(`  Table<StringLiteral> stringLiterals_;`);
   emit(`  Table<IntegerLiteral> integerLiterals_;`);
   emit(`  Table<FloatLiteral> floatLiterals_;`);
-  emit(`  SourceFiles sourceFiles_;`);
-  emit(`  SourceLines sourceLines_;`);
   emit(`  flatbuffers::FlatBufferBuilder fbb_;`);
   emit(`  flatbuffers::Offset<> offset_;`);
   emit(`  std::uint32_t type_ = 0;`);
@@ -70,9 +60,6 @@ export function gen_ast_encoder_h({
   emit(`    -> std::span<const std::uint8_t>;`);
 
   emit(`private:`);
-  emit(`  auto encodeSourceLocation(const SourceLocation& loc)`);
-  emit(`    -> flatbuffers::Offset<>;`);
-  emit();
   emit(`  auto accept(AST* ast) -> flatbuffers::Offset<>;`);
   by_base.forEach((_nodes, base) => {
     if (base === "AST") return;

packages/cxx-gen-ast/src/gen_ast_fbs.ts

@@ -95,10 +95,10 @@ export function gen_ast_fbs({ ast, output }: { ast: AST; output: string }) {
             case "node-list":
               break;
             case "token":
-              emit(`  ${fieldName}: SourceLocation;`);
+              emit(`  ${fieldName}: uint32;`);
               break;
             case "token-list":
-              emit(`  ${fieldName}: [SourceLocation];`);
+              throw new Error("unexpected token-list");
               break;
             case "attribute": {
               break;
@@ -115,14 +115,9 @@ export function gen_ast_fbs({ ast, output }: { ast: AST; output: string }) {
   const out = `${cpy_header}
 namespace cxx.io;
 
-table SourceLine {
+table Source {
   file_name: string;
-  line: uint32;
-}
-
-table SourceLocation {
-  source_line: SourceLine;
-  column: uint32;
+  line_offsets: [int];
 }
 
 ${code.join("\n")}
@@ -131,6 +126,8 @@ table SerializedUnit {
   version: uint32;
   unit: Unit;
   file_name: string;
+  token_list: [uint64];
+  source_list: [Source];
 }
 
 root_type SerializedUnit;