di-container.test.ts

import { createContainerBuilder } from "./di-container"

test.concurrent("no dependency", () => {
  const container = createContainerBuilder()
    .add("a", () => ({ a: "Hello" }))
    .add("b", () => ({ b: "World" }))
    .build()

  expect(container.get("a").a).toBe("Hello")
  expect(container.get("b").b).toBe("World")
})

test.concurrent("sync dependency", () => {
  const container = createContainerBuilder()
    .add("a", () => ({ name: "World" }))
    .add("b", ([{ name }]) => ({ msg: "Hello, " + name }), ["a"])
    .build()

  expect(container.get("b").msg).toBe("Hello, World")
})

test.concurrent("async dependency", async () => {
  const container = createContainerBuilder()
    .addAsync(
      "a",
      () =>
        new Promise<{ name: string }>((resolve) => {
          setTimeout(() => resolve({ name: "World" }), 0)
        })
    )
    .addAsync("b", ([{ name }]) => ({ msg: "Hello, " + name }), ["a"])
    .build()

  expect((await container.getAsync("b")).msg).toBe("Hello, World")
})

test.concurrent("lazy dependency", async () => {
  const container = createContainerBuilder()
    .addAsync(
      "a",
      () =>
        new Promise<{ name: string }>((resolve) => {
          setTimeout(() => resolve({ name: "World" }), 0)
        })
    )
    .addAsync(
      "b",
      async (_deps, [a]) => ({ msg: "Hello, " + (await a()).name }),
      [],
      ["a"]
    )
    .build()

  expect((await container.getAsync("b")).msg).toBe("Hello, World")
})

test("parallel resolve", async () => {
  const container = createContainerBuilder()
    .addAsync(
      "a",
      () =>
        new Promise<{ a: string }>((resolve) => {
          setTimeout(() => resolve({ a: "foo" }), 100)
        })
    )
    .addAsync(
      "b",
      () =>
        new Promise<{ b: string }>((resolve) => {
          setTimeout(() => resolve({ b: "bar" }), 100)
        })
    )
    .addAsync(
      "c",
      () =>
        new Promise<{ c: string }>((resolve) => {
          setTimeout(() => resolve({ c: "baz" }), 100)
        })
    )
    .addAsync("d", ([{ a }, { b }, { c }]) => ({ d: a + b + c }), [
      "a",
      "b",
      "c",
    ] as const)
    .build()

  // 2回目の getAsync は一瞬で終わる
  for (let i = 0; i < 2; i++) {
    expect((await container.getAsync("d")).d).toBe("foobarbaz")
  }
}, 250 /* finish before 100 * 3 ms */)

di-container.ts

export interface ContainerBuilder<SyncEntries, AsyncEntries> {
  add: <
    K extends EntryKey,
    V,
    Deps extends ReadonlyArray<keyof SyncEntries> = [],
    LazyDeps extends ReadonlyArray<keyof AsyncEntries | keyof SyncEntries> = []
  >(
    key: K,
    factory: (
      deps: DepsParameter<SyncEntries, Deps>,
      lazyDeps: LazyDepsParameter<AsyncEntries & SyncEntries, LazyDeps>
    ) => V,
    deps?: Deps,
    lazyDeps?: LazyDeps
  ) => ContainerBuilder<SyncEntries & { [_ in K]: V }, AsyncEntries>

  addAsync: <
    K extends EntryKey,
    V,
    Deps extends ReadonlyArray<keyof AsyncEntries | keyof SyncEntries> = [],
    LazyDeps extends ReadonlyArray<keyof AsyncEntries | keyof SyncEntries> = []
  >(
    key: K,
    factory: (
      deps: DepsParameter<AsyncEntries & SyncEntries, Deps>,
      lazyDeps: LazyDepsParameter<AsyncEntries & SyncEntries, LazyDeps>
    ) => Promise<V> | V,
    deps?: Deps,
    lazyDeps?: LazyDeps
  ) => ContainerBuilder<SyncEntries, AsyncEntries & { [_ in K]: V }>

  build: () => Container<SyncEntries, AsyncEntries>
}

export interface Container<SyncEntries, AsyncEntries> {
  get: <K extends keyof SyncEntries>(key: K) => SyncEntries[K]
  getAsync: <K extends keyof AsyncEntries | keyof SyncEntries>(
    key: K
  ) => Promise<
    K extends keyof SyncEntries
      ? SyncEntries[K]
      : K extends keyof AsyncEntries
      ? AsyncEntries[K]
      : never
  >
}

type DepsParameter<Entries, Deps extends ReadonlyArray<keyof Entries>> = {
  -readonly [I in keyof Deps]: Deps[I] extends keyof Entries
    ? Entries[Deps[I]]
    : undefined
}

type LazyDepsParameter<Entries, Deps extends ReadonlyArray<keyof Entries>> = {
  -readonly [I in keyof Deps]: Deps[I] extends keyof Entries
    ? () => Promise<Entries[Deps[I]]>
    : undefined
}

type EntryKey = string | symbol
type AnyFactory = (deps: any[], lazyDeps: Array<() => Promise<any>>) => any

interface Entry {
  factory: AnyFactory
  deps?: readonly EntryKey[]
  lazyDeps?: readonly EntryKey[]
  isAsync: boolean
}

export function createContainerBuilder(): ContainerBuilder<{}, {}> {
  return new ContainerBuilderImpl({}) as any
}

class ContainerBuilderImpl {
  constructor(private readonly entries: Record<EntryKey, Entry>) {}

  add(
    key: EntryKey,
    factory: AnyFactory,
    deps: readonly EntryKey[],
    lazyDeps?: readonly EntryKey[]
  ): ContainerBuilderImpl {
    this.checkKey(key)
    return new ContainerBuilderImpl({
      ...this.entries,
      [key]: { factory, deps, lazyDeps, isAsync: false },
    })
  }

  addAsync(
    key: EntryKey,
    factory: AnyFactory,
    deps?: readonly EntryKey[],
    lazyDeps?: readonly EntryKey[]
  ): ContainerBuilderImpl {
    this.checkKey(key)
    return new ContainerBuilderImpl({
      ...this.entries,
      [key]: { factory, deps, lazyDeps, isAsync: true },
    })
  }

  build(): ContainerImpl {
    return new ContainerImpl(this.entries)
  }

  private checkKey(key: EntryKey): void {
    if (this.entries[key as any] != null) {
      throw new Error(
        `キー ${String(key)} に対応するサービスはすでに登録されています。`
      )
    }
  }
}

class ContainerImpl {
  private readonly resolved = new Map<EntryKey, any>()
  private readonly promises = new Map<EntryKey, Promise<any>>()

  constructor(private readonly entries: Record<EntryKey, Entry>) {}

  get(key: any): any {
    if (this.resolved.has(key)) {
      return this.resolved.get(key)
    }

    const entry = this.entries[key]
    if (entry == null) {
      throw new Error(
        `キー ${String(key)} に対応するサービスが登録されていません。`
      )
    }
    if (entry.isAsync) {
      throw new Error(
        `キー ${String(
          key
        )} に対応するサービスは非同期にインスタンス化する必要があります。`
      )
    }

    // Resolve dependencies
    const deps =
      entry.deps == null ? [] : entry.deps.map((dep) => this.get(dep))
    const lazyDeps =
      entry.lazyDeps == null
        ? []
        : entry.lazyDeps.map((dep) => () => this.getAsync(dep))

    const value = entry.factory(deps, lazyDeps)
    this.resolved.set(key, value)
    return value
  }

  async getAsync(key: any): Promise<any> {
    if (this.resolved.has(key)) {
      return this.resolved.get(key)
    }

    let promise = this.promises.get(key)
    if (promise != null) {
      return await promise
    }

    const entry = this.entries[key]
    if (entry == null) {
      throw new Error(
        `キー ${String(key)} に対応するサービスが登録されていません。`
      )
    }

    if (!entry.isAsync) {
      return this.get(key)
    }

    promise = (async () => {
      // Resolve dependencies
      const deps =
        entry.deps == null
          ? []
          : await Promise.all(entry.deps.map((dep) => this.getAsync(dep)))
      const lazyDeps =
        entry.lazyDeps == null
          ? []
          : entry.lazyDeps.map((dep) => () => this.getAsync(dep))

      const value = await entry.factory(deps, lazyDeps)
      this.resolved.set(key, value)

      // Cleanup
      this.promises.delete(key)

      return value
    })()

    this.promises.set(key, promise)

    return await promise
  }
}