Allow more cases in Tree.is (#20973)

## Description

This makes 3 changes to Tree.is:
1. calling Tree.is(node, SchemaNotUsedInTreeSchema) is no longer an
error. This undocumented edge case could have been problematic, and is
inconsistent since it would not fire if node was not a TreeNode.
2. Tree.is now takes in ImplicitAllowedTypes, making cases like
Tree.is(x, [schema.number, schema.string]) valid. THis is more
performant and more concise then doing two separate checks and ORing
them together. This also allows checking a node against an
ImplicitAllowedTypes pulled from a FieldSchema which could be handy for
some generic code.
3. The implementation of Tree.schema and Tree.is have been rewritten to
not rely on flex-schema as much, and fast path non-TreeNode inputs. This
should make it more maintainable and more performant.


Interestingly `#2` above is the only case that couldn't be covered by
`instanceof` (assuming TypeScript 5.3): we could make `instanceof` do
all narrowing currently done with Tree.is, except for this new case. The
presence of this case thus seems to motivate keeping `Tree.is` if for no
reason other than it can support this additional pattern which
`instancof` cannot.

packages/dds/tree/api-report/tree.api.md

@@ -1913,7 +1913,7 @@ export abstract class TreeNode implements WithType {
 
 // @public
 export interface TreeNodeApi {
-    is<TSchema extends TreeNodeSchema>(value: unknown, schema: TSchema): value is NodeFromSchema<TSchema>;
+    is<TSchema extends ImplicitAllowedTypes>(value: unknown, schema: TSchema): value is TreeNodeFromImplicitAllowedTypes<TSchema>;
     key(node: TreeNode): string | number;
     on<K extends keyof TreeChangeEvents>(node: TreeNode, eventName: K, listener: TreeChangeEvents[K]): () => void;
     parent(node: TreeNode): TreeNode | undefined;

packages/dds/tree/src/simple-tree/index.ts

@@ -32,7 +32,7 @@ export {
 } from "./schemaTypes.js";
 export { SchemaFactory, type ScopedSchemaName } from "./schemaFactory.js";
 export { getFlexNode } from "./proxyBinding.js";
-export { treeNodeApi, TreeNodeApi, TreeChangeEvents } from "./treeApi.js";
+export { treeNodeApi, TreeNodeApi, TreeChangeEvents } from "./treeNodeApi.js";
 export { toFlexConfig } from "./toFlexSchema.js";
 export {
 	ObjectFromSchemaRecordUnsafe,

packages/dds/tree/src/simple-tree/treeApi.ts → packages/dds/tree/src/simple-tree/treeNodeApi.ts

@@ -8,27 +8,34 @@ import { assert, unreachableCase } from "@fluidframework/core-utils/internal";
 import { Multiplicity, rootFieldKey } from "../core/index.js";
 import {
 	FieldKinds,
-	LeafNodeSchema,
+	LazyItem,
 	TreeStatus,
+	isLazy,
 	isTreeValue,
-	valueSchemaAllows,
 } from "../feature-libraries/index.js";
-import { fail, extractFromOpaque } from "../util/index.js";
+import { fail, extractFromOpaque, isReadonlyArray } from "../util/index.js";
 
-import { getOrCreateNodeProxy } from "./proxies.js";
-import { getFlexNode, tryGetFlexNode } from "./proxyBinding.js";
+import { getOrCreateNodeProxy, isTreeNode } from "./proxies.js";
+import { getFlexNode } from "./proxyBinding.js";
 import { tryGetSimpleNodeSchema } from "./schemaCaching.js";
-import { schemaFromValue } from "./schemaFactory.js";
 import {
-	NodeFromSchema,
 	NodeKind,
 	type TreeLeafValue,
 	TreeNodeSchema,
 	type ImplicitFieldSchema,
 	FieldSchema,
+	ImplicitAllowedTypes,
+	TreeNodeFromImplicitAllowedTypes,
 } from "./schemaTypes.js";
-import { getFlexSchema } from "./toFlexSchema.js";
 import { TreeNode } from "./types.js";
+import {
+	booleanSchema,
+	handleSchema,
+	nullSchema,
+	numberSchema,
+	stringSchema,
+} from "./leafNodeSchema.js";
+import { isFluidHandle } from "@fluidframework/runtime-utils/internal";
 
 /**
  * Provides various functions for analyzing {@link TreeNode}s.
@@ -51,15 +58,15 @@ export interface TreeNodeApi {
 	 * Narrow the type of the given value if it satisfies the given schema.
 	 * @example
 	 * ```ts
-	 * if (node.is(myNode, point)) {
-	 *     const y = myNode.y; // `myNode` is now known to satisfy the `point` schema and therefore has a `y` coordinate.
+	 * if (node.is(myNode, Point)) {
+	 *     const y = myNode.y; // `myNode` is now known to satisfy the `Point` schema and therefore has a `y` coordinate.
 	 * }
 	 * ```
 	 */
-	is<TSchema extends TreeNodeSchema>(
+	is<TSchema extends ImplicitAllowedTypes>(
 		value: unknown,
 		schema: TSchema,
-	): value is NodeFromSchema<TSchema>;
+	): value is TreeNodeFromImplicitAllowedTypes<TSchema>;
 
 	/**
 	 * Return the node under which this node resides in the tree (or undefined if this is a root node of the tree).
@@ -166,30 +173,30 @@ export const treeNodeApi: TreeNodeApi = {
 	status: (node: TreeNode) => {
 		return getFlexNode(node, true).treeStatus();
 	},
-	is: <TSchema extends TreeNodeSchema>(
+	is: <TSchema extends ImplicitAllowedTypes>(
 		value: unknown,
 		schema: TSchema,
-	): value is NodeFromSchema<TSchema> => {
-		const flexSchema = getFlexSchema(schema);
-		if (isTreeValue(value)) {
-			return (
-				flexSchema instanceof LeafNodeSchema && valueSchemaAllows(flexSchema.info, value)
-			);
+	): value is TreeNodeFromImplicitAllowedTypes<TSchema> => {
+		const actualSchema = tryGetSchema(value);
+		if (actualSchema === undefined) {
+			return false;
+		}
+		if (isReadonlyArray<LazyItem<TreeNodeSchema>>(schema)) {
+			for (const singleSchema of schema) {
+				const testSchema = isLazy(singleSchema) ? singleSchema() : singleSchema;
+				if (testSchema === actualSchema) {
+					return true;
+				}
+			}
+			return false;
+		} else {
+			return (schema as TreeNodeSchema) === actualSchema;
 		}
-		return tryGetFlexNode(value)?.is(flexSchema) ?? false;
 	},
 	schema<T extends TreeNode | TreeLeafValue>(
 		node: T,
 	): TreeNodeSchema<string, NodeKind, unknown, T> {
-		if (isTreeValue(node)) {
-			return schemaFromValue(node) as TreeNodeSchema<string, NodeKind, unknown, T>;
-		}
-		return tryGetSimpleNodeSchema(getFlexNode(node).schema) as TreeNodeSchema<
-			string,
-			NodeKind,
-			unknown,
-			T
-		>;
+		return tryGetSchema(node) ?? fail("Not a tree node");
 	},
 	shortId(node: TreeNode): number | string | undefined {
 		const flexNode = getFlexNode(node);
@@ -208,6 +215,38 @@ export const treeNodeApi: TreeNodeApi = {
 	},
 };
 
+/**
+ * Returns a schema for a value if the value is a {@link TreeNode} or a {@link TreeLeafValue}.
+ * Returns undefined for other values.
+ */
+export function tryGetSchema<T>(
+	value: T,
+): undefined | TreeNodeSchema<string, NodeKind, unknown, T> {
+	type TOut = TreeNodeSchema<string, NodeKind, unknown, T>;
+	switch (typeof value) {
+		case "string":
+			return stringSchema as TOut;
+		case "number":
+			return numberSchema as TOut;
+		case "boolean":
+			return booleanSchema as TOut;
+		case "object": {
+			if (isTreeNode(value)) {
+				// This case could be optimized, for example by placing the simple schema in a symbol on tree nodes.
+				return tryGetSimpleNodeSchema(getFlexNode(value).schema) as TOut;
+			}
+			if (value === null) {
+				return nullSchema as TOut;
+			}
+			if (isFluidHandle(value)) {
+				return handleSchema as TOut;
+			}
+		}
+		default:
+			return undefined;
+	}
+}
+
 /**
  * Gets the stored key with which the provided node is associated in the parent.
  */