Implement encoding/decoding of original scope variables

Author: szuend

src/codec.ts

@@ -37,6 +37,7 @@ export type Item =
   | typeof EmptyItem
   | OriginalScopeStartItem
   | OriginalScopeEndItem
+  | OriginalScopeVariablesItem
   | GeneratedRangeStartItem
   | GeneratedRangeEndItem;
 
@@ -55,6 +56,11 @@ interface OriginalScopeEndItem {
   column: number;
 }
 
+interface OriginalScopeVariablesItem {
+  tag: Tag.ORIGINAL_SCOPE_VARIABLES;
+  variableIdxs: number[];
+}
+
 interface GeneratedRangeStartItem {
   tag: Tag.GENERATED_RANGE_START;
   flags: number;

src/decode/decode.ts

@@ -28,6 +28,7 @@ const DEFAULT_SCOPE_STATE = {
   line: 0,
   name: 0,
   kind: 0,
+  variable: 0,
 };
 
 const DEFAULT_RANGE_STATE = {
@@ -95,6 +96,22 @@ class Decoder {
           this.#countToScope.set(this.#scopeCounter++, scope);
           break;
         }
+        case Tag.ORIGINAL_SCOPE_VARIABLES: {
+          const scope = this.#scopeStack.at(-1);
+          if (!scope) {
+            throw new Error(
+              "Encountered ORIGINAL_SCOPE_VARIABLES without surrounding ORIGINAL_SCOPE_START",
+            );
+          }
+
+          for (const variableIdx of item.variableIdxs) {
+            this.#scopeState.variable += variableIdx;
+            scope.variables.push(this.#names[this.#scopeState.variable]);
+
+            // TODO: Potentially throw if we decode an illegal index.
+          }
+          break;
+        }
         case Tag.ORIGINAL_SCOPE_END: {
           this.#scopeState.line += item.line;
 
@@ -226,6 +243,16 @@ class Decoder {
           yield item;
           break;
         }
+        case Tag.ORIGINAL_SCOPE_VARIABLES: {
+          const variableIdxs: number[] = [];
+
+          while (iter.hasNext() && iter.peek() !== ",") {
+            variableIdxs.push(iter.nextSignedVLQ());
+          }
+
+          yield { tag, variableIdxs };
+          break;
+        }
         case Tag.ORIGINAL_SCOPE_END: {
           yield {
             tag,

src/encode/encoder.ts

@@ -15,6 +15,7 @@ const DEFAULT_SCOPE_STATE = {
   column: 0,
   name: 0,
   kind: 0,
+  variable: 0,
 };
 
 const DEFAULT_RANGE_STATE = {
@@ -64,6 +65,7 @@ export class Encoder {
     }
 
     this.#encodeOriginalScopeStart(scope);
+    this.#encodeOriginalScopeVariables(scope);
     scope.children.forEach((child) => this.#encodeOriginalScope(child));
     this.#encodeOriginalScopeEnd(scope);
   }
@@ -104,6 +106,20 @@ export class Encoder {
     this.#scopeToCount.set(scope, this.#scopeCounter++);
   }
 
+  #encodeOriginalScopeVariables(scope: OriginalScope) {
+    if (scope.variables.length === 0) return;
+
+    this.#encodeTag(EncodedTag.ORIGINAL_SCOPE_VARIABLES);
+
+    for (const variable of scope.variables) {
+      const idx = this.#resolveNamesIdx(variable);
+      this.#encodeSigned(idx - this.#scopeState.variable);
+      this.#scopeState.variable = idx;
+    }
+
+    this.#finishItem();
+  }
+
   #encodeOriginalScopeEnd(scope: OriginalScope) {
     const { line, column } = scope.end;
     this.#verifyPositionWithScopeState(line, column);

src/roundtrip.test.ts

@@ -132,4 +132,12 @@ describe("round trip", () => {
 
     assertCodec(builder.build());
   });
+
+  it("handles OriginalScope variables", () => {
+    builder.startScope(0, 0, { variables: ["foo", "bar"] }).startScope(10, 0, {
+      variables: ["local1", "local2"],
+    }).endScope(20, 0).endScope(30, 0);
+
+    assertCodec(builder.build());
+  });
 });