muddle/core/container.ts

import { UnknownProviderError } from "./errors.ts";
import { getLatestMetadata, getMetadata } from "./metadata.ts";

export type Type<T = unknown> = new (...args: never[]) => T;
export type AbstractType<T = unknown> = abstract new (...args: never[]) => T;
export type Token<T = unknown> = Type<T> | AbstractType<T> | symbol | string;

export type ProviderScope = "singleton" | "request" | "transient";

export interface ProviderDefinition<T = unknown> {
	/** Token the provider is registered under, defaults to `useClass`. */
	provide?: Token<T>;
	/** Static value. */
	useValue?: T;
	/** Class instantiated by the container. */
	useClass?: AbstractType<T>;
	/** Factory function, dependencies are resolved through `inject`. */
	useFactory?: (...dependencies: unknown[]) => T;
	/** Tokens handed to `useFactory` (or used for `useClass` overrides). */
	inject?: Token[];
	/** Lifetime of the provider. */
	scope?: ProviderScope;
}

interface FactoryEntry {
	token: Token;
	scope: ProviderScope;
	owner: Container;
	produce: (scope?: Container) => unknown;
	/** Asynchronous variant, used by {@link Container.resolveAsync}. */
	produceAsync?: (scope?: Container) => Promise<unknown>;
}

export const PARAM_TYPES_METADATA = "design:paramtypes";
export const INJECT_TOKENS_METADATA = "muddle:inject";
export const OPTIONAL_PARAMS_METADATA = "muddle:optional";

interface InjectEntry {
	index: number;
	token: Token;
}

/** Reads the `@Inject()` overrides of a constructor. */
export function getInjectTokens(target: object): Map<number, Token> {
	const entries = getLatestMetadata(INJECT_TOKENS_METADATA, target) as
		| InjectEntry[]
		| undefined;
	const map = new Map<number, Token>();
	for (const entry of entries ?? []) {
		map.set(entry.index, entry.token);
	}
	return map;
}

/** Reads the parameters marked with `@Optional()`. */
export function getOptionalParams(target: object): Set<number> {
	const entries = getLatestMetadata(OPTIONAL_PARAMS_METADATA, target) as
		| number[]
		| undefined;
	return new Set(entries ?? []);
}

export class Container {
	private readonly factories = new Map<Token, FactoryEntry>();
	readonly instances = new Map<Token, unknown>();

	/** Called when a token cannot be found in this container or its parents. */
	fallbackResolver?: (token: Token, scope?: Container) => unknown;

	constructor(readonly parent?: Container) {}

	/** Creates a child container that shares factories but isolates instances. */
	createScope(): Container {
		const scope = new Container(this);
		scope.fallbackResolver = (token, inner) =>
			this.lookupFallback(token, inner);
		return scope;
	}

	/** Registers a factory entry. */
	set(token: Token, entry: Omit<FactoryEntry, "token" | "owner">): void {
		this.factories.set(token, { ...entry, token, owner: this });
	}

	/** `true` when the token is known to this container or one of its parents. */
	has(token: Token): boolean {
		return this.findEntry(token) !== undefined;
	}

	/** Every token known to this container. */
	tokens(): Token[] {
		return [...this.factories.keys()];
	}

	/**
	 * Resolves a token.
	 * Providers built with an asynchronous factory return a promise, use
	 * {@link resolveAsync} to await them.
	 */
	resolve<T>(token: Token<T>, scope?: Container): T {
		const entry = this.findEntry(token);
		if (!entry) {
			const fallback = this.lookupFallback(token, scope);
			if (fallback !== undefined) return fallback as T;
			throw new UnknownProviderError(token);
		}

		if (entry.scope === "transient") {
			return entry.produce(scope) as T;
		}

		const cache = entry.scope === "request" ? (scope ?? this) : entry.owner;
		const cached = cache.instances.get(token);
		if (cached !== undefined) return cached as T;

		const value = entry.produce(scope);
		cache.instances.set(token, value);
		return value as T;
	}

	/** Resolves a token, awaits async providers. */
	async resolveAsync<T>(token: Token<T>, scope?: Container): Promise<T> {
		const entry = this.findEntry(token);
		if (!entry) {
			const fallback = this.lookupFallback(token, scope);
			if (fallback !== undefined) return (await fallback) as T;
			throw new UnknownProviderError(token);
		}

		if (entry.scope === "transient") {
			const value = entry.produceAsync
				? await entry.produceAsync(scope)
				: await entry.produce(scope);
			return value as T;
		}

		const cache = entry.scope === "request" ? (scope ?? this) : entry.owner;
		const cached = cache.instances.get(token);
		if (cached !== undefined) return (await cached) as T;

		const value = entry.produceAsync
			? await entry.produceAsync(scope)
			: await entry.produce(scope);
		cache.instances.set(token, value);
		return value as T;
	}

	private findEntry(token: Token): FactoryEntry | undefined {
		const own = this.factories.get(token);
		if (own) return own;
		return this.parent?.findEntry(token);
	}

	private lookupFallback(token: Token, scope?: Container): unknown {
		const fallback = this.fallbackResolver ?? this.parent?.fallbackResolver;
		if (!fallback) return undefined;
		return fallback(token, scope);
	}

	/** Registers a class, resolving its constructor dependencies automatically. */
	registerClass<T>(
		ctor: AbstractType<T>,
		options: { scope?: ProviderScope; inject?: Token[] } = {},
	): void {
		this.set(ctor as Token, {
			scope: options.scope ?? "singleton",
			produce: (scope) => this.instantiate(ctor, options.inject, scope),
			produceAsync: async (scope) =>
				this.instantiateAsync(ctor, options.inject, scope),
		});
	}

	/** Registers a plain value. */
	registerValue<T>(token: Token<T>, value: T): void {
		this.set(token as Token, { scope: "singleton", produce: () => value });
	}

	/** Registers a factory function. */
	registerFactory<T>(
		token: Token<T>,
		factory: (...dependencies: unknown[]) => T | Promise<T>,
		options: { inject?: Token[]; scope?: ProviderScope } = {},
	): void {
		this.set(token as Token, {
			scope: options.scope ?? "singleton",
			produce: (scope) => {
				const dependencies = (options.inject ?? []).map((dependency) =>
					this.resolve(dependency, scope),
				);
				return factory(...dependencies);
			},
			produceAsync: async (scope) => {
				const dependencies = await Promise.all(
					(options.inject ?? []).map((dependency) =>
						this.resolveAsync(dependency, scope),
					),
				);
				return factory(...dependencies);
			},
		});
	}

	/** Registers a provider definition. */
	registerProvider<T>(provider: ProviderDefinition<T>): void {
		if (provider.useValue !== undefined) {
			this.registerValue(
				(provider.provide ?? provider.useClass) as Token,
				provider.useValue,
			);
			return;
		}
		if (provider.useFactory) {
			const token = provider.provide ?? provider.useFactory;
			this.registerFactory(token as Token, provider.useFactory, {
				inject: provider.inject,
				scope: provider.scope,
			});
			return;
		}
		if (provider.useClass) {
			this.registerClass(provider.useClass, {
				scope: provider.scope,
				inject: provider.inject,
			});
			return;
		}
		throw new UnknownProviderError(provider.provide);
	}

	/** Instantiates a class, resolving constructor parameters by type. */
	instantiate<T>(
		ctor: AbstractType<T>,
		inject?: Token[],
		scope?: Container,
	): T {
		const paramTypes = (getMetadata(PARAM_TYPES_METADATA, ctor) ??
			[]) as unknown[];
		const overrides = getInjectTokens(ctor as object);
		const optional = getOptionalParams(ctor as object);
		const dependencies: unknown[] = [];

		const length = Math.max(
			paramTypes.length,
			overrides.size,
			inject?.length ?? 0,
		);
		for (let index = 0; index < length; index++) {
			const token =
				overrides.get(index) ??
				inject?.[index] ??
				(paramTypes[index] as Token | undefined);
			if (token === undefined || token === Object || token === Function) {
				if (optional.has(index)) {
					dependencies.push(undefined);
					continue;
				}
				throw new UnknownProviderError(
					`${ctor.name} parameter #${index} (use @Inject() to provide an explicit token)`,
				);
			}
			dependencies.push(this.resolve(token, scope));
		}

		return new (ctor as unknown as new (...args: unknown[]) => T)(
			...dependencies,
		);
	}

	/** Instantiates a class and awaits asynchronous dependencies. */
	async instantiateAsync<T>(
		ctor: AbstractType<T>,
		inject?: Token[],
		scope?: Container,
	): Promise<T> {
		const paramTypes = (getMetadata(PARAM_TYPES_METADATA, ctor) ??
			[]) as unknown[];
		const overrides = getInjectTokens(ctor as object);
		const optional = getOptionalParams(ctor as object);
		const dependencies: unknown[] = [];

		const length = Math.max(
			paramTypes.length,
			overrides.size,
			inject?.length ?? 0,
		);
		for (let index = 0; index < length; index++) {
			const token =
				overrides.get(index) ??
				inject?.[index] ??
				(paramTypes[index] as Token | undefined);
			if (token === undefined || token === Object || token === Function) {
				if (optional.has(index)) {
					dependencies.push(undefined);
					continue;
				}
				throw new UnknownProviderError(
					`${ctor.name} parameter #${index} (use @Inject() to provide an explicit token)`,
				);
			}
			dependencies.push(await this.resolveAsync(token, scope));
		}

		return new (ctor as unknown as new (...args: unknown[]) => T)(
			...dependencies,
		);
	}
}