feat: #614 - Add forEachStructureChild

lib/ts-morph.d.ts

@@ -10802,9 +10802,19 @@ interface TypeParameterDeclarationSpecificStructure extends KindedStructure<Stru
     default?: string | WriterFunction;
 }
 
-export declare type OptionalKind<TStructure extends {
-        kind?: StructureKind;
-    }> = Pick<TStructure, Exclude<keyof TStructure, "kind">> & Partial<Pick<TStructure, "kind">>;
+/**
+ * Iterates over the elements in the provided array.
+ * @param structures - Array of structures to iterate over.
+ * @param callback - Callback to do on each element in the array. Returning a truthy value will return that value in the main function call.
+ */
+export declare function forEachStructureChild<TStructure>(structures: Structures[], callback: (child: Structures) => TStructure | void): TStructure | undefined;
+/**
+ * Iterates over the children of the provided array.
+ * @remarks If the children do not have a `kind` property, it will be automatically added.
+ * @param structure - Structure to iterate over.
+ * @param callback - Callback to do on each child of the provided structure. Returning a truthy value will return that value in the main function call.
+ */
+export declare function forEachStructureChild<TStructure>(structure: Structures, callback: (child: Structures) => TStructure | void): TStructure | undefined;
 
 export interface Structure {
     /**
@@ -10863,5 +10873,8 @@ export declare enum StructureKind {
     VariableStatement = 38
 }
 
+export declare type OptionalKind<TStructure extends {
+        kind?: StructureKind;
+    }> = Pick<TStructure, Exclude<keyof TStructure, "kind">> & Partial<Pick<TStructure, "kind">>;
 export { ts, SyntaxKind, CompilerOptions, EmitHint, ScriptKind, NewLineKind, LanguageVariant, ScriptTarget, TypeFlags, ObjectFlags, SymbolFlags, TypeFormatFlags, DiagnosticCategory, EditorSettings, ModuleResolutionKind };
 export * from "./code-block-writer";

package.json

@@ -60,7 +60,8 @@
     "include": [
       "src/**/*.ts",
       "!src/tests/**/*.ts",
-      "!src/utils/TypeGuards.ts"
+      "!src/utils/TypeGuards.ts",
+      "!src/structures/utils/forEachStructureChild.ts"
     ],
     "reporter": [
       "html"

scripts/generation/createForEachStructureChild.ts

@@ -0,0 +1,224 @@
+/**
+ * Code Manipulation - Create `forEachStructureKind`
+ * -------------------------------------------------
+ * This modifies the forEachStructureKind.ts file that
+ * is used for iterating over a structure's structures.
+ * -------------------------------------------------
+ */
+import { PropertySignature, SourceFile, Type, FunctionDeclarationStructure, StructureKind } from "ts-morph";
+import { TsMorphInspector, Structure } from "../inspectors";
+
+// very messy first pass... needs cleanup
+
+// todo: Shouldn't have functions for syntax kinds that will be no-ops (ex. `forDecoratableNode`)... actually, this will add the kind for them?
+
+interface StructureInfo {
+    name: string;
+    syntaxKind: string | undefined;
+    baseStructures: StructureInfo[];
+    members: MemberInfo[];
+}
+
+interface MemberInfo {
+    name: string;
+    syntaxKind: string | undefined;
+    allStructure: boolean;
+}
+
+export function createForEachStructureChild(inspector: TsMorphInspector) {
+    const forEachStructureChildFile = inspector.getProject().getSourceFileOrThrow("forEachStructureChild.ts");
+    clearPreviouslyGeneratedFunctions(forEachStructureChildFile);
+    const structureInfos = getStructureInfos(inspector);
+    addNewFunctions(forEachStructureChildFile, structureInfos);
+    updateForEachStructureChild(forEachStructureChildFile, structureInfos);
+    forEachStructureChildFile.fixMissingImports();
+}
+
+function clearPreviouslyGeneratedFunctions(sourceFile: SourceFile) {
+    // remove functions with an @generated jsdoc tag
+    for (const func of sourceFile.getFunctions()) {
+        if (func.getJsDocs().some(d => d.getTags().some(t => t.getTagName() === "generated")))
+            func.remove();
+    }
+}
+
+function updateForEachStructureChild(sourceFile: SourceFile, structureInfos: StructureInfo[]) {
+    const func = sourceFile.getFunctionOrThrow("forEachStructureChild");
+    func.removeStatements([0, func.getStatementsWithComments().length - 1]);
+
+    func.addStatements(writer => {
+        writer.writeLine("// automatically generated: run `yarn run code-generate` to update the code in here");
+        writer.write("if (structure instanceof Array)").block(() => {
+            writer.write("for (const item of structure)").block(() => {
+                writer.writeLine("const result = callback(item);");
+                writer.writeLine("if (result)");
+                writer.indent().write("return result;").newLine();
+            });
+            writer.writeLine("return undefined;");
+        });
+        writer.blankLine();
+        writer.write("switch (structure.kind)").block(() => {
+            for (const structureInfo of structureInfos.filter(s => s.syntaxKind != null)) {
+                writer.writeLine(`case StructureKind.${structureInfo.syntaxKind!}:`);
+                writer.indent().write(`return ${getInfoFunctionName(structureInfo)}(structure, callback);`).newLine();
+            }
+            writer.writeLine("default:");
+            writer.indent().write("return undefined;").newLine();
+        });
+    });
+}
+
+function addNewFunctions(sourceFile: SourceFile, structureInfos: StructureInfo[]) {
+    const functions: FunctionDeclarationStructure[] = [];
+    for (const info of structureInfos) {
+        functions.push({
+            kind: StructureKind.Function,
+            docs: ["@generated"],
+            name: getInfoFunctionName(info),
+            typeParameters: ["TStructure"],
+            parameters: [{ name: "structure", type: info.name }, { name: "callback", type: "(structure: Structures) => TStructure | void" }],
+            returnType: "TStructure | undefined",
+            statements: writer => {
+                let isFirst = true;
+                writer.write("return ");
+
+                for (const baseStructure of info.baseStructures)
+                    addExpression(getInfoFunctionName(baseStructure) + "(structure, callback)");
+
+                for (const member of info.members) {
+                    if (member.syntaxKind != null) {
+                        if (member.allStructure)
+                            addExpression(`forAll(structure.${member.name}, callback, StructureKind.${member.syntaxKind})`);
+                        else
+                            addExpression(`forAllIfStructure(structure.${member.name}, callback, StructureKind.${member.syntaxKind})`);
+                    }
+                    else
+                        addExpression(`forAllUnknownKindIfStructure(structure.${member.name}, callback)`);
+                }
+
+                writer.write(";");
+
+                function addExpression(text: string) {
+                    if (isFirst)
+                        isFirst = false;
+                    else
+                        writer.newLine().indent().write("|| ");
+
+                    writer.write(text);
+                }
+            }
+        });
+    }
+
+    const insertIndex = sourceFile.getFunctionOrThrow("forEachStructureChild").getChildIndex() + 1;
+    sourceFile.insertFunctions(insertIndex, functions);
+}
+
+function getStructureInfos(inspector: TsMorphInspector) {
+    const infos = new Map<Structure, StructureInfo>();
+    const structures = inspector.getStructures();
+
+    for (const structure of structures) {
+        if (structure.getName().endsWith("SpecificStructure") || structure.getName() === "KindedStructure" || structure.getName() === "Structure")
+            continue;
+
+        const type = structure.getType();
+        const kindSymbol = type.getProperty("kind");
+        if (kindSymbol == null)
+            continue;
+
+        handleStructure(structure);
+    }
+
+    return Array.from(infos.values()).filter(shouldIncludeInfo);
+
+    function handleStructure(structure: Structure): StructureInfo {
+        if (infos.has(structure))
+            return infos.get(structure)!;
+
+        const structureKindProperty = getTypeKindProperty(structure.getType());
+        const info: StructureInfo = {
+            name: structure.getName(),
+            baseStructures: [],
+            members: [],
+            syntaxKind: structureKindProperty == null ? undefined : getSyntaxKind(structureKindProperty.getTypeAtLocation(structure.node))
+        };
+        infos.set(structure, info);
+
+        for (const property of getProperties())
+            fillMemberInfo(property);
+
+        for (const baseStructure of structure.getBaseStructures()) {
+            if (!baseStructure.getName().endsWith("SpecificStructure") && baseStructure.getName() !== "Structure") {
+                const baseStructureInfo = handleStructure(baseStructure);
+                if (shouldIncludeInfo(baseStructureInfo))
+                    info.baseStructures.push(baseStructureInfo);
+            }
+        }
+
+        return info;
+
+        function getProperties() {
+            const result = structure.getProperties();
+            const specificStructureName = structure.getName().replace(/Structure$/, "SpecificStructure");
+            const specificStructure = structure.getBaseStructures().find(s => s.getName() === specificStructureName);
+            if (specificStructure != null)
+                result.push(...specificStructure.getProperties());
+            return result;
+        }
+
+        function fillMemberInfo(property: PropertySignature) {
+            const propertyType = property.getType();
+            const arrayTypes = getArrayTypes(propertyType.getNonNullableType());
+            if (getTypeKindProperty(propertyType.getNonNullableType()) != null)
+                throw new Error("Unhandled situation where there was a structure not as an array.");
+
+            const kinds: string[] = [];
+            let allStructure = true;
+            for (const arrayType of arrayTypes) {
+                const arrayElementType = arrayType.getArrayElementType();
+                for (const unionElementType of getTypeOrUnionElementTypes(arrayElementType)) {
+                    const kindProperty = getTypeKindProperty(unionElementType);
+                    if (kindProperty != null)
+                        kinds.push(getSyntaxKind(kindProperty.getTypeAtLocation(property.getTypeNodeOrThrow())));
+                    else
+                        allStructure = false;
+                }
+            }
+
+            if (kinds.length > 0) {
+                info.members.push({
+                    name: property.getName(),
+                    syntaxKind: kinds.length > 1 ? undefined : kinds[0],
+                    allStructure
+                });
+            }
+
+            function getArrayTypes(type: Type) {
+                if (type.isArray())
+                    return [type];
+                return type.getUnionTypes().filter(t => t.isArray());
+            }
+
+            function getTypeOrUnionElementTypes(type: Type) {
+                return type.isUnion() ? type.getUnionTypes() : [type];
+            }
+        }
+
+        function getTypeKindProperty(type: Type) {
+            return type.getProperty("kind");
+        }
+
+        function getSyntaxKind(type: Type) {
+            return type.getNonNullableType().getText().replace(/^.*\.([^\.]+)$/, "$1");
+        }
+    }
+
+    function shouldIncludeInfo(info: StructureInfo) {
+        return info.baseStructures.length > 0 || info.members.length > 0;
+    }
+}
+
+function getInfoFunctionName(info: StructureInfo) {
+    return `for${info.name.replace(/Structure$/, "")}`;
+}

scripts/generation/main.ts

@@ -5,6 +5,7 @@ import { createCodeBlockWriterFile } from "./createCodeBlockWriterFile";
 import { createKindToNodeMappings } from "./createKindToNodeMappings";
 import { createStructurePrinterFactory } from "./createStructurePrinterFactory";
 import { createDeclarationFile } from "./createDeclarationFile";
+import { createForEachStructureChild } from "./createForEachStructureChild";
 import { InspectorFactory } from "../inspectors";
 
 const args = process.argv.slice(2);
@@ -40,6 +41,10 @@ const tsInspector = factory.getTsInspector();
         console.log("Creating declaration file...");
         await createDeclarationFile();
     }
+    if (checkHasArg("create-for-each-structure-child")) {
+        console.log("Creating for each structure child...");
+        createForEachStructureChild(inspector);
+    }
 
     if (args.length > 0)
         console.error(`Unknown args: ${args}`);

scripts/inspectors/tsMorph/Structure.ts

@@ -1,10 +1,10 @@
 import { InterfaceDeclaration, TypeGuards } from "ts-morph";
-import { Memoize, ArrayUtils } from "../../../src/utils";
+import { Memoize } from "../../../src/utils";
 import { WrapperFactory } from "../WrapperFactory";
 import { isOverloadStructure } from "../../config";
 
 export class Structure {
-    constructor(private readonly wrapperFactory: WrapperFactory, private readonly node: InterfaceDeclaration) {
+    constructor(private readonly wrapperFactory: WrapperFactory, readonly node: InterfaceDeclaration) {
     }
 
     getName() {
@@ -15,6 +15,14 @@ export class Structure {
         return this.node.getSourceFile().getFilePath();
     }
 
+    getSymbol() {
+        return this.node.getSymbolOrThrow();
+    }
+
+    getType() {
+        return this.node.getType();
+    }
+
     getStartLineNumber() {
         return this.node.getStartLineNumber();
     }
@@ -23,6 +31,10 @@ export class Structure {
         return isOverloadStructure(this.getName());
     }
 
+    getProperties() {
+        return this.node.getProperties();
+    }
+
     @Memoize
     getBaseStructures() {
         return this.node.getBaseDeclarations().map(d => {

src/structures/index.ts

@@ -11,6 +11,7 @@ export * from "./jsx";
 export * from "./module";
 export * from "./statement";
 export * from "./type";
-export * from "./types";
+export * from "./utils";
 export * from "./Structure";
 export * from "./StructureKind";
+export * from "./types";