TSL: Make ArrayNode Generic

types/three/src/nodes/core/ArrayNode.d.ts

@@ -1,25 +1,33 @@
 import Node from "./Node.js";
 import TempNode from "./TempNode.js";
 
-declare class ArrayNode extends TempNode {
+export interface ArrayNodeInterface<TNodeType> {
     count: number;
-    values: Node[];
+    values: Node<TNodeType>[] | null;
     readonly isArrayNode: true;
-
-    constructor(nodeType: string, count: number, values: Node[]);
 }
 
+declare const ArrayNode: {
+    new<TNodeType>(
+        nodeType: TNodeType | null,
+        count: number,
+        values?: Node<TNodeType>[] | null,
+    ): ArrayNode<TNodeType>;
+};
+
+type ArrayNode<TNodeType> = TempNode<TNodeType> & ArrayNodeInterface<TNodeType>;
+
 export default ArrayNode;
 
 interface ArrayFunction {
-    (values: Node[]): ArrayNode;
-    (nodeType: string, count: number): ArrayNode;
+    <TNodeType>(values: Node<TNodeType>[]): ArrayNode<TNodeType>;
+    <TNodeType extends string>(nodeType: TNodeType, count: number): ArrayNode<TNodeType>;
 }
 
 export const array: ArrayFunction;
 
 declare module "./Node.js" {
-    interface NodeElements {
-        toArray: (count: number) => ArrayNode;
+    interface NodeExtensions<TNodeType> {
+        toArray: (count: number) => ArrayNode<TNodeType>;
     }
 }

types/three/src/nodes/tsl/TSLCore.d.ts

@@ -5,11 +5,13 @@ import { Matrix4 } from "../../math/Matrix4.js";
 import { Vector2 } from "../../math/Vector2.js";
 import { Vector3 } from "../../math/Vector3.js";
 import { Vector4 } from "../../math/Vector4.js";
+import ArrayNode from "../core/ArrayNode.js";
 import ConstNode from "../core/ConstNode.js";
 import Node, { MatType, NumOrBoolType } from "../core/Node.js";
 import NodeBuilder from "../core/NodeBuilder.js";
 import StackNode from "../core/StackNode.js";
 import VarNode from "../core/VarNode.js";
+import ArrayElementNode from "../utils/ArrayElementNode.js";
 import ConvertNode from "../utils/ConvertNode.js";
 import JoinNode from "../utils/JoinNode.js";
 
@@ -2398,14 +2400,18 @@ declare module "../core/Node.js" {
     }
 }
 
-export const element: (node: Node, indexNode: Node) => Node;
+export const element: <TNodeType>(node: ArrayNode<TNodeType>, indexNode: Node | number) => ArrayElementNode<TNodeType>;
 export const convert: (node: Node, types: string) => Node;
 export const split: (node: Node, channels?: string) => Node;
 
+declare module "../core/ArrayNode.js" {
+    interface ArrayNodeInterface<TNodeType> {
+        element: (indexNode: Node | number) => ArrayElementNode<TNodeType>;
+    }
+}
+
 declare module "../core/Node.js" {
     interface NodeElements {
-        element: (indexNode: Node) => Node;
-
         convert: (types: string) => Node;
     }
 }

types/three/src/nodes/utils/ArrayElementNode.d.ts

@@ -1,14 +1,15 @@
+import ArrayNode from "../core/ArrayNode.js";
 import Node from "../core/Node.js";
 
-interface ArrayElementNodeInterface {
-    node: Node;
+interface ArrayElementNodeInterface<TNodeType> {
+    node: ArrayNode<TNodeType>;
     indexNode: Node;
 }
 
 declare const ArrayElementNode: {
-    new<TNodeType>(node: Node<TNodeType>, indexNode: Node): ArrayElementNode<TNodeType>;
+    new<TNodeType>(node: ArrayNode<TNodeType>, indexNode: Node): ArrayElementNode<TNodeType>;
 };
 
-type ArrayElementNode<TNodeType> = Node<TNodeType> & ArrayElementNodeInterface;
+type ArrayElementNode<TNodeType> = Node<TNodeType> & ArrayElementNodeInterface<TNodeType>;
 
 export default ArrayElementNode;