fix(rosetta): tablet compression handled incorrectly in multiple places (#3670)

Unfortunately, #3652 was half-baked and this PR serves to finish handling tablet compression. It introduces the following:

- `compress-cache` cli option so that you can use a compressed cache file
- update  to all usages of `languageTablet.save()` to explicitly pass the compress flag
- `languageTablet.compressedSource`, which is set when we load from a compressed source, so that we can remember to save as a compressed source.
- `loadAllDefaultTablets` now handles compressed default tablets.
- a far more encompassing set of unit tests for the above features
---

By submitting this pull request, I confirm that my contribution is made under the terms of the [Apache 2.0 license].

[Apache 2.0 license]: https://www.apache.org/licenses/LICENSE-2.0

Author: kaizencc

packages/jsii-rosetta/bin/jsii-rosetta.ts

@@ -231,7 +231,12 @@ function main() {
           .options('compress-tablet', {
             alias: 'z',
             type: 'boolean',
-            describe: 'Compress the resulting tablet file',
+            describe: 'Compress the implicit tablet file',
+            default: false,
+          })
+          .options('compress-cache', {
+            type: 'boolean',
+            describe: 'Compress the cache-to file',
             default: false,
           })
           .conflicts('loose', 'strict')
@@ -259,6 +264,7 @@ function main() {
           trimCache: args['trim-cache'],
           loose: args.loose,
           compressTablet: args['compress-tablet'],
+          compressCacheToFile: args['compress-cache'],
         };
 
         const result = args.infuse

packages/jsii-rosetta/lib/commands/extract.ts

@@ -75,18 +75,26 @@ export interface ExtractOptions {
   readonly allowDirtyTranslations?: boolean;
 
   /**
-   * Compress the resulting tablet file
+   * Compress the implicit tablet files
    *
    * @default false
    */
   readonly compressTablet?: boolean;
+
+  /**
+   * Compress the cacheToFile tablet.
+   *
+   * @default false
+   */
+  readonly compressCacheToFile?: boolean;
 }
 
 export async function extractAndInfuse(assemblyLocations: string[], options: ExtractOptions): Promise<ExtractResult> {
   const result = await extractSnippets(assemblyLocations, options);
   await infuse(assemblyLocations, {
     cacheFromFile: options.cacheFromFile,
     cacheToFile: options.cacheToFile,
+    compressCacheToFile: options.compressCacheToFile,
   });
   return result;
 }
@@ -157,7 +165,7 @@ export async function extractSnippets(
     logging.info('Nothing left to translate');
   }
 
-  // Save to individual tablet files, and optionally append to the output file
+  // Save to individual tablet files
   if (options.writeToImplicitTablets ?? true) {
     await Promise.all(
       Object.entries(snippetsPerAssembly).map(async ([location, snips]) => {
@@ -175,13 +183,14 @@ export async function extractSnippets(
     );
   }
 
+  // optionally append to the output file
   if (options.cacheToFile) {
     logging.info(`Adding translations to ${options.cacheToFile}`);
     const output = options.trimCache
       ? new LanguageTablet()
       : await LanguageTablet.fromOptionalFile(options.cacheToFile);
     output.addTablets(translator.tablet);
-    await output.save(options.cacheToFile);
+    await output.save(options.cacheToFile, options.compressCacheToFile);
   }
 
   return { diagnostics, tablet: translator.tablet };

packages/jsii-rosetta/lib/commands/infuse.ts

@@ -12,7 +12,12 @@ import {
 import { renderMetadataline, TypeScriptSnippet } from '../snippet';
 import { SnippetSelector, mean, meanLength, shortest, longest } from '../snippet-selectors';
 import { snippetKey } from '../tablets/key';
-import { LanguageTablet, TranslatedSnippet, DEFAULT_TABLET_NAME } from '../tablets/tablets';
+import {
+  LanguageTablet,
+  TranslatedSnippet,
+  DEFAULT_TABLET_NAME,
+  DEFAULT_TABLET_NAME_COMPRESSED,
+} from '../tablets/tablets';
 import { isDefined, mkDict, indexBy } from '../util';
 
 export interface InfuseResult {
@@ -36,6 +41,11 @@ export interface InfuseOptions {
    * In addition to the implicit tablets, also write all added examples to this additional output tablet
    */
   readonly cacheToFile?: string;
+
+  /**
+   * Compress the cacheToFile
+   */
+  readonly compressCacheToFile?: boolean;
 }
 
 export const DEFAULT_INFUSION_RESULTS_NAME = 'infusion-results.html';
@@ -87,6 +97,10 @@ export async function infuse(assemblyLocations: string[], options?: InfuseOption
         stream?.write(`<h1>${assembly.name}</h1>\n`);
 
         const implicitTablet = defaultTablets[directory];
+        const implicitTabletFile = path.join(
+          directory,
+          implicitTablet.compressedSource ? DEFAULT_TABLET_NAME_COMPRESSED : DEFAULT_TABLET_NAME,
+        );
         if (!implicitTablet) {
           throw new Error(`No tablet found for ${directory}`);
         }
@@ -110,7 +124,7 @@ export async function infuse(assemblyLocations: string[], options?: InfuseOption
           // eslint-disable-next-line no-await-in-loop
           await Promise.all([
             replaceAssembly(assembly, directory),
-            implicitTablet.save(path.join(directory, DEFAULT_TABLET_NAME)),
+            implicitTablet.save(implicitTabletFile, implicitTablet.compressedSource),
           ]);
         }
 
@@ -130,7 +144,7 @@ export async function infuse(assemblyLocations: string[], options?: InfuseOption
   // If we copied examples onto different types, we'll also have inserted new snippets
   // with different keys into the tablet. We must now write the updated tablet somewhere.
   if (options?.cacheToFile) {
-    await additionalOutputTablet.save(options.cacheToFile);
+    await additionalOutputTablet.save(options.cacheToFile, options.compressCacheToFile);
   }
 
   return {

packages/jsii-rosetta/lib/commands/trim-cache.ts

@@ -25,7 +25,8 @@ export async function trimCache(options: TrimCacheOptions): Promise<void> {
   const original = await LanguageTablet.fromFile(options.cacheFile);
   const updated = new LanguageTablet();
   updated.addSnippets(...snippets.map((snip) => original.tryGetSnippet(snippetKey(snip))).filter(isDefined));
-  await updated.save(options.cacheFile);
+  // if the original file was compressed, then compress the updated file too
+  await updated.save(options.cacheFile, original.compressedSource);
 
   // eslint-disable-next-line prettier/prettier
   logging.info(`${options.cacheFile}: ${updated.count} snippets remaining (${original.count} - ${updated.count} trimmed)`);

packages/jsii-rosetta/lib/jsii/assemblies.ts

@@ -24,7 +24,7 @@ import {
   INITIALIZER_METHOD_NAME,
 } from '../snippet';
 import { enforcesStrictMode } from '../strict';
-import { LanguageTablet, DEFAULT_TABLET_NAME } from '../tablets/tablets';
+import { LanguageTablet, DEFAULT_TABLET_NAME, DEFAULT_TABLET_NAME_COMPRESSED } from '../tablets/tablets';
 import { fmap, mkDict, sortBy } from '../util';
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires, @typescript-eslint/no-require-imports
@@ -92,12 +92,15 @@ export function loadAssemblies(
 export async function loadAllDefaultTablets(asms: readonly LoadedAssembly[]): Promise<Record<string, LanguageTablet>> {
   return mkDict(
     await Promise.all(
-      asms.map(
-        async (a) =>
-          [a.directory, await LanguageTablet.fromOptionalFile(path.join(a.directory, DEFAULT_TABLET_NAME))] as const,
-      ),
+      asms.map(async (a) => [a.directory, await LanguageTablet.fromOptionalFile(guessTabletLocation(a))] as const),
     ),
   );
+
+  function guessTabletLocation(a: LoadedAssembly): string {
+    const defaultTablet = path.join(a.directory, DEFAULT_TABLET_NAME);
+    const compDefaultTablet = path.join(a.directory, DEFAULT_TABLET_NAME_COMPRESSED);
+    return fs.existsSync(defaultTablet) ? defaultTablet : compDefaultTablet;
+  }
 }
 
 export type AssemblySnippetSource =

packages/jsii-rosetta/lib/tablets/tablets.ts

@@ -54,6 +54,12 @@ export class LanguageTablet {
     return ret;
   }
 
+  /**
+   * Whether or not the LanguageTablet was loaded with a compressed source.
+   * This gets used to determine if it should be compressed when saved.
+   */
+  public compressedSource = false;
+
   private readonly snippets: Record<string, TranslatedSnippet> = {};
 
   /**
@@ -140,6 +146,7 @@ export class LanguageTablet {
     if (data[0] === 0x1f && data[1] === 0x8b && data[2] === 0x08) {
       // This is a gz object, so we decompress it now...
       data = zlib.gunzipSync(data);
+      this.compressedSource = true;
     }
 
     const obj: TabletSchema = JSON.parse(data.toString('utf-8'));

packages/jsii-rosetta/test/commands/extract.test.ts

@@ -10,6 +10,7 @@ import {
   DEFAULT_TABLET_NAME,
   TranslatedSnippet,
   typeScriptSnippetFromVisibleSource,
+  DEFAULT_TABLET_NAME_COMPRESSED,
 } from '../../lib';
 import * as extract from '../../lib/commands/extract';
 import { loadAssemblies } from '../../lib/jsii/assemblies';
@@ -71,19 +72,37 @@ test('extract samples from test assembly', async () => {
   await tablet.load(cacheToFile);
 
   expect(tablet.snippetKeys.length).toEqual(1);
+  expect(tablet.compressedSource).toBeFalsy();
 });
 
-test('extract can save/load compressed tablets', async () => {
+test('extract can compress cached tablet file', async () => {
   const compressedCacheFile = path.join(assembly.moduleDirectory, 'test.tabl.gz');
   await extract.extractSnippets([assembly.moduleDirectory], {
     cacheToFile: compressedCacheFile,
+    compressCacheToFile: true,
     ...defaultExtractOptions,
   });
 
   const tablet = new LanguageTablet();
   await tablet.load(compressedCacheFile);
 
   expect(tablet.snippetKeys.length).toEqual(1);
+  expect(tablet.compressedSource).toBeTruthy();
+});
+
+test('extract can compress implicit tablet file', async () => {
+  await extract.extractSnippets([assembly.moduleDirectory], {
+    ...defaultExtractOptions,
+    compressTablet: true,
+  });
+
+  const compImplicitTablet = path.join(assembly.moduleDirectory, DEFAULT_TABLET_NAME_COMPRESSED);
+  expect(fs.existsSync(compImplicitTablet)).toBeTruthy();
+  expect(fs.existsSync(path.join(assembly.moduleDirectory, DEFAULT_TABLET_NAME))).toBeFalsy();
+
+  const tablet = new LanguageTablet();
+  await tablet.load(compImplicitTablet);
+  expect(tablet.snippetKeys.length).toEqual(1);
 });
 
 test('extract works from compressed test assembly', async () => {

packages/jsii-rosetta/test/commands/infuse.test.ts

@@ -2,7 +2,7 @@ import { loadAssemblyFromPath, SPEC_FILE_NAME, SPEC_FILE_NAME_COMPRESSED } from
 import * as fs from 'fs-extra';
 import * as path from 'path';
 
-import { LanguageTablet, DEFAULT_TABLET_NAME } from '../../lib';
+import { LanguageTablet, DEFAULT_TABLET_NAME, DEFAULT_TABLET_NAME_COMPRESSED } from '../../lib';
 import { extractSnippets } from '../../lib/commands/extract';
 import { infuse, DEFAULT_INFUSION_RESULTS_NAME } from '../../lib/commands/infuse';
 import { loadAssemblies } from '../../lib/jsii/assemblies';
@@ -162,3 +162,41 @@ test('preserves the assembly compression if present', async () => {
   expect(types).toBeDefined();
   expect(types!['my_assembly.ClassA'].docs?.example).toBeDefined();
 });
+
+test('can infuse with compressed default tablets', async () => {
+  // remove any tablets that may currently exist
+  const implicitTablet = path.join(assembly.moduleDirectory, DEFAULT_TABLET_NAME);
+  const compImplicitTablet = path.join(assembly.moduleDirectory, DEFAULT_TABLET_NAME_COMPRESSED);
+  await fs.remove(implicitTablet);
+  await fs.remove(compImplicitTablet);
+
+  // create a compressed implicit tablet file via extract
+  await extractSnippets([assembly.moduleDirectory], {
+    includeCompilerDiagnostics: false,
+    validateAssemblies: false,
+    compressTablet: true,
+  });
+
+  expect(fs.existsSync(compImplicitTablet)).toBeTruthy();
+  expect(fs.existsSync(implicitTablet)).toBeFalsy();
+
+  // infuse can use compressed implicit tablets
+  await infuse([assembly.moduleDirectory]);
+
+  const assemblies = loadAssemblies([assembly.moduleDirectory], false);
+  const types = assemblies[0].assembly.types;
+  expect(types).toBeDefined();
+  expect(types!['my_assembly.ClassA'].docs?.example).toBeDefined();
+});
+
+test('can compress cacheToFile', async () => {
+  const compressedTabletFile = path.join(assembly.moduleDirectory, 'tabl.json.gz');
+
+  await infuse([assembly.moduleDirectory], {
+    cacheToFile: compressedTabletFile,
+    compressCacheToFile: true,
+  });
+
+  const tablet = await LanguageTablet.fromFile(compressedTabletFile);
+  expect(tablet.compressedSource).toBeTruthy();
+});

packages/jsii-rosetta/test/commands/trim-cache.test.ts

@@ -1,6 +1,12 @@
 import * as path from 'path';
 
-import { TranslatedSnippet, typeScriptSnippetFromVisibleSource, LanguageTablet, DEFAULT_TABLET_NAME } from '../../lib';
+import {
+  TranslatedSnippet,
+  typeScriptSnippetFromVisibleSource,
+  LanguageTablet,
+  DEFAULT_TABLET_NAME,
+  DEFAULT_TABLET_NAME_COMPRESSED,
+} from '../../lib';
 import { extractSnippets } from '../../lib/commands/extract';
 import { trimCache } from '../../lib/commands/trim-cache';
 import { TestJsiiModule, DUMMY_JSII_CONFIG, testSnippetLocation } from '../testutil';
@@ -77,6 +83,26 @@ test('trim-cache leaves used snippets', async () => {
   expect(updated.count).toEqual(1);
 });
 
+test('trim-cache preserves tablet compression', async () => {
+  const compDefaultTablet = path.join(assembly.moduleDirectory, DEFAULT_TABLET_NAME_COMPRESSED);
+
+  // GIVEN
+  const tbl = new LanguageTablet();
+  tbl.addSnippets(bogusTranslatedSnippet());
+  // explicitly compress the tablet file
+  await tbl.save(compDefaultTablet, true);
+
+  // WHEN
+  await trimCache({
+    assemblyLocations: [assembly.moduleDirectory],
+    cacheFile: compDefaultTablet,
+  });
+
+  // THEN
+  const updated = await LanguageTablet.fromFile(compDefaultTablet);
+  expect(updated.compressedSource).toBeTruthy();
+});
+
 function bogusTranslatedSnippet() {
   return TranslatedSnippet.fromTypeScript(
     typeScriptSnippetFromVisibleSource('console.log("hello");', testSnippetLocation('x.ts'), true),