feat: support inject Injector (#1)

* feat: support inject Injector

* refactor: rename parameter

* refactor: extract resolution creation

* refactor: rm redundant assertion

src/__tests__/inject.spec.ts

@@ -1,6 +1,7 @@
 import {afterEach, describe, expect, it, vi} from "vitest";
 
-import {Container, Inject, inject, injectAll} from "..";
+import {Build, Container, Inject, inject, injectAll, Injector, Scope, Scoped} from "..";
+import {useInjectionContext} from "../injection-context";
 
 describe("inject", () => {
   const container = new Container();
@@ -66,4 +67,58 @@ describe("inject", () => {
     expect(wizard.wand1.owner).toBe(wizard);
     expect(wizard.wand2.owner).toBe(wizard);
   });
+
+  describe("Injector", () => {
+    it("should inject injector", () => {
+      class Wizard {
+        injector = inject(Injector);
+      }
+
+      class Wand {
+        name = "Elder Wand";
+      }
+
+      const wizard = container.resolve(Wizard);
+      expect(wizard.injector.inject(Wand)).toBeInstanceOf(Wand);
+      expect(wizard.injector.injectAll(Wand)).toEqual([new Wand()]);
+    });
+
+    it("should use current context", () => {
+      class Wizard {
+        injector = inject(Injector);
+
+        context = inject(Build(useInjectionContext));
+
+        constructor() {
+          this.injector.inject(
+            Build(() => {
+              expect(useInjectionContext()).toBe(this.context);
+            }),
+          );
+        }
+      }
+
+      container.resolve(Wizard);
+    });
+
+    it("should have context of the dependent", () => {
+      const container = new Container({
+        autoRegister: true,
+      });
+
+      class Wand {
+        owner = inject(Wizard);
+      }
+
+      @Scoped(Scope.Container)
+      class Wizard {
+        injector = inject.by(this, Injector);
+      }
+
+      const wizard = container.resolve(Wizard);
+      const wand = wizard.injector.inject(Wand);
+      expect(wand.owner).toBe(wizard);
+      expect(container.getCached(Wand)).toBe(wand);
+    });
+  });
 });

src/container.ts

@@ -1,11 +1,10 @@
 import {assert, expectNever} from "./errors";
-import {type ResolvedScope, useInjectionContext, withInjectionContext} from "./injection-context";
+import {createResolution, useInjectionContext, withInjectionContext} from "./injection-context";
 import {getMetadata} from "./metadata";
 import {isClassProvider, isFactoryProvider, isValueProvider, type Provider} from "./provider";
 import {type Registration, type RegistrationOptions, Registry} from "./registry";
 import {Scope} from "./scope";
 import {type Constructor, isConstructor, type Token, type TokenList} from "./token";
-import {KeyedStack} from "./utils/keyed-stack";
 
 export interface ContainerOptions {
   parent?: Container;
@@ -184,11 +183,7 @@ export class Container {
     if (!context || context.container !== this) {
       return withInjectionContext({
         container: this,
-        resolution: {
-          stack: new KeyedStack(),
-          instances: new Map(),
-          dependents: new Map(),
-        },
+        resolution: createResolution(),
       }, () => this.getScopedInstance(registration, instantiate));
     }
 
@@ -233,7 +228,7 @@ export class Container {
     }
   }
 
-  private resolveScope(scope = this.defaultScope): ResolvedScope {
+  private resolveScope(scope = this.defaultScope) {
     let resolvedScope = scope;
     if (resolvedScope == Scope.Inherited) {
       const context = useInjectionContext();

src/index.ts

@@ -2,7 +2,7 @@ export type {ContainerOptions} from "./container";
 export {Container} from "./container";
 export type {ClassDecorator, ClassFieldDecorator, ClassFieldInitializer} from "./decorators";
 export {AutoRegister, Inject, Injectable, InjectAll, Scoped} from "./decorators";
-export {inject, injectAll, injectBy} from "./inject";
+export {inject, injectAll, injectBy, Injector} from "./inject";
 export type {InstanceRef} from "./instance";
 export type {ClassProvider, FactoryProvider, Provider, ValueProvider} from "./provider";
 export type {Registration, RegistrationMap, RegistrationOptions, Registry} from "./registry";

src/inject.ts

@@ -1,5 +1,6 @@
-import {ensureInjectionContext} from "./injection-context";
-import type {Token, TokenList} from "./token";
+import {ensureInjectionContext, useInjectionContext, withInjectionContext} from "./injection-context";
+import {Build} from "./registry";
+import type {Token, TokenList, Type} from "./token";
 import {invariant} from "./utils/invariant";
 
 export function inject<Values extends unknown[]>(...tokens: TokenList<Values>): Values[number];
@@ -17,15 +18,17 @@ inject.by = injectBy;
 export function injectBy<Values extends unknown[]>(thisArg: any, ...tokens: TokenList<Values>): Values[number];
 export function injectBy<Value>(thisArg: any, ...tokens: Token<Value>[]): Value {
   const context = ensureInjectionContext(injectBy);
+
   const currentFrame = context.resolution.stack.peek();
   invariant(currentFrame);
-  const provider = currentFrame.provider;
-  context.resolution.dependents.set(provider, {current: thisArg});
+
+  const currentProvider = currentFrame.provider;
+  context.resolution.dependents.set(currentProvider, {current: thisArg});
   try {
     return inject(...tokens);
   }
   finally {
-    context.resolution.dependents.delete(provider);
+    context.resolution.dependents.delete(currentProvider);
   }
 }
 
@@ -34,3 +37,41 @@ export function injectAll<Value>(...tokens: Token<Value>[]): NonNullable<Value>[
   const context = ensureInjectionContext(injectAll);
   return context.container.resolveAll(...tokens);
 }
+
+export interface Injector {
+  inject<Values extends unknown[]>(...tokens: TokenList<Values>): Values[number];
+  injectAll<Values extends unknown[]>(...tokens: TokenList<Values>): NonNullable<Values[number]>[];
+}
+
+export const Injector: Type<Injector> = /*@__PURE__*/ Build(function Injector() {
+  const context = ensureInjectionContext(Injector);
+
+  const dependentFrame = context.resolution.stack.peek(1);
+  invariant(dependentFrame);
+
+  const dependentProvider = dependentFrame.provider;
+  const dependentRef = context.resolution.dependents.get(dependentProvider);
+
+  const withCurrentContext = <R>(fn: () => R) => {
+    if (useInjectionContext()) {
+      return fn();
+    }
+    return withInjectionContext(context, () => {
+      context.resolution.stack.push(dependentProvider, dependentFrame);
+      if (dependentRef)
+        context.resolution.dependents.set(dependentProvider, dependentRef);
+      try {
+        return fn();
+      }
+      finally {
+        context.resolution.dependents.delete(dependentProvider);
+        context.resolution.stack.pop();
+      }
+    });
+  };
+
+  return {
+    inject: (...tokens) => withCurrentContext(() => inject(...tokens)),
+    injectAll: (...tokens) => withCurrentContext(() => injectAll(...tokens)),
+  };
+});

src/injection-context.ts

@@ -4,7 +4,8 @@ import type {InstanceRef} from "./instance";
 import type {Provider} from "./provider";
 import type {Scope} from "./scope";
 import {createContext} from "./utils/context";
-import type {KeyedStack} from "./utils/keyed-stack";
+import {KeyedStack} from "./utils/keyed-stack";
+import {WeakValueMap} from "./utils/weak-value-map";
 
 export interface InjectionContext {
   container: Container;
@@ -13,16 +14,23 @@ export interface InjectionContext {
 
 export interface Resolution {
   stack: KeyedStack<Provider, ResolutionFrame>;
-  instances: Map<Provider, InstanceRef>;
-  dependents: Map<Provider, InstanceRef>;
+  instances: WeakValueMap<Provider, InstanceRef>;
+  dependents: WeakValueMap<Provider, InstanceRef>;
 }
 
 export interface ResolutionFrame {
-  scope: ResolvedScope;
+  scope: Exclude<Scope, typeof Scope.Inherited>;
   provider: Provider;
 }
 
-export type ResolvedScope = Exclude<Scope, typeof Scope.Inherited>;
+// @internal
+export function createResolution(): Resolution {
+  return {
+    stack: new KeyedStack(),
+    instances: new WeakValueMap(),
+    dependents: new WeakValueMap(),
+  };
+}
 
 // @internal
 export const [withInjectionContext, useInjectionContext] = createContext<InjectionContext>();

src/utils/keyed-stack.ts

@@ -8,8 +8,8 @@ export class KeyedStack<K, V> {
     return this.keys.has(key);
   }
 
-  peek() {
-    const entry = this.entries.at(-1);
+  peek(n = 0) {
+    const entry = this.entries.at(-(n + 1));
     return entry?.value;
   }
 

src/utils/weak-value-map.ts

@@ -0,0 +1,26 @@
+import {invariant} from "./invariant";
+
+// @internal
+export class WeakValueMap<K, V extends object> {
+  private map = new Map<K, WeakRef<V>>();
+
+  delete(key: K) {
+    this.map.delete(key);
+  }
+
+  get(key: K) {
+    const ref = this.map.get(key);
+    if (ref) {
+      const value = ref.deref();
+      if (value) {
+        return value;
+      }
+      this.map.delete(key);
+    }
+  }
+
+  set(key: K, value: V) {
+    invariant(!this.get(key));
+    this.map.set(key, new WeakRef(value));
+  }
+}