feat: free memory automatically (#5225)

lib/binding_web/src/finalization_registry.ts

@@ -0,0 +1,8 @@
+export function newFinalizer<T>(handler: (value: T) => void): FinalizationRegistry<T> | undefined {
+  try {
+    return new FinalizationRegistry(handler);
+  } catch(e) {
+    console.error('Unsupported FinalizationRegistry:', e);
+    return;
+  }
+}

lib/binding_web/src/lookahead_iterator.ts

@@ -1,5 +1,10 @@
 import { C, Internal, assertInternal } from './constants';
 import { Language } from './language';
+import { newFinalizer } from './finalization_registry';
+
+const finalizer = newFinalizer((address: number) => {
+  C._ts_lookahead_iterator_delete(address);
+});
 
 export class LookaheadIterator implements Iterable<string> {
   /** @internal */
@@ -13,6 +18,7 @@ export class LookaheadIterator implements Iterable<string> {
     assertInternal(internal);
     this[0] = address;
     this.language = language;
+    finalizer?.register(this, address, this);
   }
 
   /** Get the current symbol of the lookahead iterator. */
@@ -27,6 +33,7 @@ export class LookaheadIterator implements Iterable<string> {
 
   /** Delete the lookahead iterator, freeing its resources. */
   delete(): void {
+    finalizer?.unregister(this);
     C._ts_lookahead_iterator_delete(this[0]);
     this[0] = 0;
   }

lib/binding_web/src/parser.ts

@@ -3,6 +3,7 @@ import { Language } from './language';
 import { marshalRange, unmarshalRange } from './marshal';
 import { checkModule, initializeBinding } from './bindings';
 import { Tree } from './tree';
+import { newFinalizer } from './finalization_registry';
 
 /**
  * Options for parsing
@@ -82,6 +83,11 @@ export let LANGUAGE_VERSION: number;
  */
 export let MIN_COMPATIBLE_VERSION: number;
 
+const finalizer = newFinalizer((addresses: number[]) => {
+  C._ts_parser_delete(addresses[0]);
+  C._free(addresses[1]);
+});
+
 /**
  * A stateful object that is used to produce a {@link Tree} based on some
  * source code.
@@ -117,6 +123,7 @@ export class Parser {
    */
   constructor() {
     this.initialize();
+    finalizer?.register(this, [this[0], this[1]], this);
   }
 
   /** @internal */
@@ -131,6 +138,7 @@ export class Parser {
 
   /** Delete the parser, freeing its resources. */
   delete() {
+    finalizer?.unregister(this);
     C._ts_parser_delete(this[0]);
     C._free(this[1]);
     this[0] = 0;

lib/binding_web/src/query.ts

@@ -3,6 +3,7 @@ import { Node } from './node';
 import { marshalNode, unmarshalCaptures } from './marshal';
 import { TRANSFER_BUFFER } from './parser';
 import { Language } from './language';
+import { newFinalizer } from './finalization_registry';
 
 const PREDICATE_STEP_TYPE_CAPTURE = 1;
 
@@ -506,6 +507,10 @@ function parsePattern(
   }
 }
 
+const finalizer = newFinalizer((address: number) => {
+  C._ts_query_delete(address);
+});
+
 export class Query {
   /** @internal */
   private [0] = 0; // Internal handle for Wasm
@@ -687,10 +692,12 @@ export class Query {
     this.assertedProperties = assertedProperties;
     this.refutedProperties = refutedProperties;
     this.exceededMatchLimit = false;
+    finalizer?.register(this, address, this);
   }
 
   /** Delete the query, freeing its resources. */
   delete(): void {
+    finalizer?.unregister(this);
     C._ts_query_delete(this[0]);
     this[0] = 0;
   }

lib/binding_web/src/tree.ts

@@ -5,6 +5,7 @@ import { TreeCursor } from './tree_cursor';
 import { marshalEdit, marshalPoint, unmarshalNode, unmarshalRange } from './marshal';
 import { TRANSFER_BUFFER } from './parser';
 import { Edit } from './edit';
+import { newFinalizer } from './finalization_registry';
 
 /** @internal */
 export function getText(tree: Tree, startIndex: number, endIndex: number, startPosition: Point): string {
@@ -28,6 +29,10 @@ export function getText(tree: Tree, startIndex: number, endIndex: number, startP
   return result ?? '';
 }
 
+const finalizer = newFinalizer((address: number) => {
+  C._ts_tree_delete(address);
+});
+
 /** A tree that represents the syntactic structure of a source code file. */
 export class Tree {
   /** @internal */
@@ -45,6 +50,7 @@ export class Tree {
     this[0] = address;
     this.language = language;
     this.textCallback = textCallback;
+    finalizer?.register(this, address, this);
   }
 
   /** Create a shallow copy of the syntax tree. This is very fast. */
@@ -55,6 +61,7 @@ export class Tree {
 
   /** Delete the syntax tree, freeing its resources. */
   delete(): void {
+    finalizer?.unregister(this);
     C._ts_tree_delete(this[0]);
     this[0] = 0;
   }

lib/binding_web/src/tree_cursor.ts

@@ -3,6 +3,11 @@ import { marshalNode, marshalPoint, marshalTreeCursor, unmarshalNode, unmarshalP
 import { Node } from './node';
 import { TRANSFER_BUFFER } from './parser';
 import { getText, Tree } from './tree';
+import { newFinalizer } from './finalization_registry';
+
+const finalizer = newFinalizer((address: number) => {
+  C._ts_tree_cursor_delete_wasm(address);
+});
 
 /** A stateful object for walking a syntax {@link Tree} efficiently. */
 export class TreeCursor {
@@ -30,6 +35,7 @@ export class TreeCursor {
     assertInternal(internal);
     this.tree = tree;
     unmarshalTreeCursor(this);
+    finalizer?.register(this, this.tree[0], this);
   }
 
   /** Creates a deep copy of the tree cursor. This allocates new memory. */
@@ -42,6 +48,7 @@ export class TreeCursor {
 
   /** Delete the tree cursor, freeing its resources. */
   delete(): void {
+    finalizer?.unregister(this);
     marshalTreeCursor(this);
     C._ts_tree_cursor_delete_wasm(this.tree[0]);
     this[0] = this[1] = this[2] = 0;

lib/binding_web/test/memory.test.ts

@@ -0,0 +1,74 @@
+import { describe, expect, it } from 'vitest';
+import { gc, event, Finalizer } from './memory';
+
+// hijack finalization registry before import web-tree-sitter
+globalThis.FinalizationRegistry = Finalizer;
+
+describe('Memory Management', () => {
+  describe('call .delete()', () => {
+    it('test free memory manually', async () => {
+      const timer = setInterval(() => {
+        gc();
+      }, 100);
+      let done = 0; 
+      event.on('gc', () => {
+        done++;
+      });
+      await (async () => {
+        const { JavaScript } = await (await import('./helper')).default;
+        const  { Parser, Query } = await import('../src');
+        const parser = new Parser();
+        parser.setLanguage(JavaScript);
+        const tree = parser.parse('1+1')!;
+        const copyTree = tree.copy();
+        const cursor = tree.walk();
+        const copyCursor = cursor.copy();
+        const lookaheadIterator = JavaScript.lookaheadIterator(cursor.currentNode.nextParseState)!;
+        const query = new Query(JavaScript, '(identifier) @element');
+        parser.delete();
+        tree.delete();
+        copyTree.delete();
+        cursor.delete();
+        copyCursor.delete();
+        lookaheadIterator.delete();
+        query.delete();
+      })();
+      // wait for gc
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      clearInterval(timer);
+      // expect no gc event fired
+      expect(done).toBe(0);
+    });
+  });
+
+  describe('do not call .delete()', () => {
+    it('test free memory automatically', async () => {
+      const timer = setInterval(() => {
+        gc();
+      }, 100);
+      let done = 0;
+      const promise = new Promise((resolve) => {
+        event.on('gc', () => {
+          if (++done === 7) {
+            resolve(undefined);
+            clearInterval(timer);
+          }
+          console.log('free memory times: ', done);
+        });
+      });
+      await (async () => {
+        const { JavaScript } = await (await import('./helper')).default;
+        const  { Parser, Query } = await import('../src');
+        const parser = new Parser(); // 1
+        parser.setLanguage(JavaScript);
+        const tree = parser.parse('1+1')!; // 2
+        tree.copy(); // 3
+        const cursor = tree.walk(); // 4
+        cursor.copy(); // 5
+        JavaScript.lookaheadIterator(cursor.currentNode.nextParseState)!; // 6
+        new Query(JavaScript, '(identifier) @element'); // 7
+      })();
+      await promise;
+    });
+  });
+});

lib/binding_web/test/memory.ts

@@ -0,0 +1,20 @@
+import { EventEmitter } from 'events';
+import { Session } from 'inspector';
+
+const session = new Session();
+session.connect();
+
+export function gc() {
+  session.post('HeapProfiler.collectGarbage');
+}
+
+export const event = new EventEmitter();
+
+export class Finalizer<T> extends FinalizationRegistry<T> {
+  constructor(handler: (value: T) => void) {
+    super((value) => {
+      handler(value);
+      event.emit('gc');
+    });
+  }
+}

lib/binding_web/test/memory_unsupported.test.ts

@@ -0,0 +1,25 @@
+import { describe, it } from 'vitest';
+
+describe('FinalizationRegistry is unsupported', () => {
+  it('test FinalizationRegistry is unsupported', async () => {
+    // @ts-expect-error: test FinalizationRegistry is not supported
+    globalThis.FinalizationRegistry = undefined;
+    const { JavaScript } = await (await import('./helper')).default;
+    const  { Parser, Query } = await import('../src');
+    const parser = new Parser();
+    parser.setLanguage(JavaScript);
+    const tree = parser.parse('1+1')!;
+    const copyTree = tree.copy();
+    const cursor = tree.walk();
+    const copyCursor = cursor.copy();
+    const lookaheadIterator = JavaScript.lookaheadIterator(cursor.currentNode.nextParseState)!;
+    const query = new Query(JavaScript, '(identifier) @element');
+    parser.delete();
+    tree.delete();
+    copyTree.delete();
+    cursor.delete();
+    copyCursor.delete();
+    lookaheadIterator.delete();
+    query.delete();
+  });
+});