add match() method for storage maps

unyt-org · May 25, 2024 · 0e3224c · 0e3224c
1 parent b1f5380
commit 0e3224c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/types/storage-map.ts b/types/storage-map.ts
@@ -3,7 +3,10 @@
 import { Compiler } from "../compiler/compiler.ts";
 import { DX_PTR } from "../runtime/constants.ts";
 import { Pointer } from "../runtime/pointers.ts";
-import { Storage } from "../storage/storage.ts";
+import { MatchOptions, Storage } from "../storage/storage.ts";
+import { Class } from "../utils/global_types.ts";
+import { MatchInput, MatchResult, match } from "../utils/match.ts";
+import { Type } from "./type.ts";
 
 
 /**
@@ -275,5 +278,8 @@ export class StorageMap<K,V> extends StorageWeakMap<K,V> {
 		await Promise.all(promises);
 	}
 
+	match<Options extends MatchOptions, T extends V & object>(valueType:Class<T>|Type<T>, matchInput: MatchInput<T>, options?: Options): Promise<MatchResult<T, Options>> {
+		return match(this as unknown as StorageMap<unknown, T>, valueType, matchInput, options)
+	}
 
 }
diff --git a/utils/match.ts b/utils/match.ts
@@ -2,6 +2,7 @@ import { StorageSet } from "../types/storage_set.ts";
 import { Type } from "../types/type.ts";
 import type { Class } from "./global_types.ts";
 import { MatchInput, MatchResult, MatchOptions, Storage } from "../storage/storage.ts";
+import { StorageMap } from "../types/storage-map.ts";
 
 export type { MatchInput, MatchOptions, MatchResult } from "../storage/storage.ts";
 
@@ -12,7 +13,7 @@ export type { MatchInput, MatchOptions, MatchResult } from "../storage/storage.t
  * @param limit 
  * @returns 
  */
-export async function match<T extends object, Options extends MatchOptions>(inputSet: StorageSet<T>, valueType:Class<T>|Type<T>, match: MatchInput<T>, options?: Options): Promise<MatchResult<T, Options>> {
+export async function match<T extends object, Options extends MatchOptions>(inputSet: StorageSet<T>|StorageMap<unknown,T>, valueType:Class<T>|Type<T>, match: MatchInput<T>, options?: Options): Promise<MatchResult<T, Options>> {
 	options ??= {} as Options;
 	const found = new Set<T>();
 	const matchOrEntries = (match instanceof Array ? match : [match]).map(m => Object.entries(m)) as [keyof T, T[keyof T]][][];