Simplify hooks

Author: LeaVerou

src/core/classes/hooks.ts

@@ -0,0 +1,107 @@
+/**
+ * A class for managing hooks for deep extensibility.
+ * Inspired https://www.npmjs.com/package/blissful-hooks
+ */
+export class Hooks {
+	private _all: Record<keyof HookEnv, HookCallback[]> = {};
+
+	/**
+	 * Adds the given callback to the list of callbacks for the given hook and returns a function that
+	 * removes the hook again when called.
+	 *
+	 * The callback will be invoked when the hook it is registered for is run.
+	 * Hooks are usually directly run by a highlight function but you can also run hooks yourself.
+	 *
+	 * One callback function can be registered to multiple hooks.
+	 *
+	 * A callback function must not be registered for the same hook multiple times. Doing so will cause
+	 * undefined behavior. However, registering a callback again after removing it is fine.
+	 *
+	 * @param name The name of the hook.
+	 * @param callback The callback function which is given environment variables.
+	 */
+	add<Name extends keyof HookEnv> (
+		name: Name | Name[] | MultipleHooks<Name>,
+		callback?: Name extends MultipleHooks<Name> ? never : HookCallback<Name>
+	): () => void {
+		if (Array.isArray(name)) {
+			// One function, multiple hooks
+			for (const n of name) {
+				this.add(n, callback);
+			}
+		}
+		else if (typeof name === 'object') {
+			// Multiple hooks
+			const hooks = name;
+
+			for (const name in hooks) {
+				const callback = hooks[name];
+				if (callback) {
+					this.add(name as string, callback as HookCallback);
+				}
+			}
+		}
+		else {
+			const hooks = (this._all[name] ??= []);
+			hooks.push(callback as never);
+		}
+
+		return () => {
+			this.remove(name, callback);
+		};
+	}
+
+	remove<Name extends keyof HookEnv> (
+		name: Name | Name[] | MultipleHooks<Name>,
+		callback?: Name extends MultipleHooks<Name> ? never : HookCallback<Name>
+	): void {
+		if (Array.isArray(name)) {
+			// Multiple hook names, same callback
+			for (const n of name) {
+				this.remove(n, callback);
+			}
+		}
+		else if (typeof name === 'object') {
+			// Map of hook names to callbacks
+			for (const n in name) {
+				this.remove(n, callback);
+			}
+		}
+		else {
+			const index = this._all[name]?.indexOf(callback as never);
+			if (index > -1) {
+				this._all[name].splice(index, 1);
+			}
+		}
+	}
+
+	/**
+	 * Runs a hook invoking all registered callbacks with the given environment variables.
+	 *
+	 * Callbacks will be invoked synchronously and in the order in which they were registered.
+	 *
+	 * @param name The name of the hook.
+	 * @param env The environment variables of the hook passed to all callbacks registered.
+	 */
+	run<Name extends keyof HookEnv> (name: Name, env: HookEnv[Name]): void {
+		const callbacks = this._all[name];
+		const context = env?.this ?? env?.context ?? env;
+
+		if (!callbacks || !callbacks.length) {
+			return;
+		}
+
+		for (const callback of callbacks) {
+			callback.call(context, env);
+		}
+	}
+}
+
+export interface BaseHookEnv {
+	context?: object;
+}
+export interface HookEnv extends BaseHookEnv, Record<string, any> {}
+
+export type HookCallback<T extends keyof HookEnv = string> = (env: HookEnv[T]) => void;
+
+export type MultipleHooks<T extends keyof HookEnv> = { [K in T]?: HookCallback<K> };

src/core/hook-state.ts

@@ -1,14 +1,13 @@
 import { getLanguage, setLanguage } from '../shared/dom-util';
 import { rest, tokenize } from '../shared/symbols';
 import { htmlEncode } from '../shared/util';
-import { HookState } from './hook-state';
-import { Hooks } from './hooks';
+import { Hooks } from './classes/hooks';
 import { LinkedList } from './linked-list';
 import { Registry } from './registry';
 import { Token } from './classes/token';
 import type { KnownPlugins } from '../known-plugins';
 import type { Grammar, GrammarToken, GrammarTokens, RegExpLike } from '../types';
-import type { HookEnvMap } from './hooks';
+import type { HookEnv } from './classes/hooks';
 import type { LinkedListHeadNode, LinkedListMiddleNode, LinkedListTailNode } from './linked-list';
 import type { TokenStream } from './classes/token';
 
@@ -36,18 +35,16 @@ export class Prism {
 	highlightAll (options: HighlightAllOptions = {}) {
 		const { root, async, callback } = options;
 
-		const env: HookEnvMap['before-highlightall'] | HookEnvMap['before-all-elements-highlight'] =
+		const env: HookEnv =
 			{
 				callback,
 				root: root ?? document,
 				selector:
 					'code[class*="language-"], [class*="language-"] code, code[class*="lang-"], [class*="lang-"] code',
-				state: new HookState(),
 			};
 
 		this.hooks.run('before-highlightall', env);
 
-		assertEnv<'before-all-elements-highlight'>(env);
 		env.elements = [...env.root.querySelectorAll(env.selector)];
 
 		this.hooks.run('before-all-elements-highlight', env);
@@ -93,16 +90,14 @@ export class Prism {
 
 		const code = element.textContent as string;
 
-		const env: HookEnvMap['before-sanity-check'] = {
+		const env: HookEnv = {
 			element,
 			language,
 			grammar,
 			code,
-			state: new HookState(),
 		};
 
 		const insertHighlightedCode = (highlightedCode: string) => {
-			assertEnv<'before-insert'>(env);
 			env.highlightedCode = highlightedCode;
 			this.hooks.run('before-insert', env);
 
@@ -168,7 +163,7 @@ export class Prism {
 		const languageId = this.components.resolveAlias(language);
 		const grammar = options?.grammar ?? this.components.getLanguage(languageId);
 
-		const env: HookEnvMap['before-tokenize'] | HookEnvMap['after-tokenize'] = {
+		const env: HookEnv = {
 			code: text,
 			grammar,
 			language,
@@ -178,7 +173,6 @@ export class Prism {
 			throw new Error('The language "' + env.language + '" has no grammar.');
 		}
 
-		assertEnv<'after-tokenize'>(env);
 		env.tokens = this.tokenize(env.code, env.grammar);
 		this.hooks.run('after-tokenize', env);
 
@@ -437,10 +431,6 @@ export interface HighlightOptions {
 	grammar?: Grammar;
 }
 
-function assertEnv<T extends keyof HookEnvMap> (env: unknown): asserts env is HookEnvMap[T] {
-	/* noop */
-}
-
 function matchPattern (pattern: RegExp, pos: number, text: string, lookbehind: boolean) {
 	pattern.lastIndex = pos;
 	const match = pattern.exec(text);
@@ -475,7 +465,7 @@ function stringify (o: string | Token | TokenStream, language: string, hooks: Ho
 		return s;
 	}
 
-	const env: HookEnvMap['wrap'] = {
+	const env: HookEnv = {
 		type: o.type,
 		content: stringify(o.content, language, hooks),
 		tag: 'span',