|
| 1 | +import assert from "node:assert"; |
| 2 | +import type { CodegenLogger } from "@root/utils/codegen-logger"; |
| 3 | +import { extractDependencies } from "./core/transformer"; |
| 4 | +import type { ResolutionTree } from "./register"; |
| 5 | +import { |
| 6 | + type CanonicalUrl, |
| 7 | + type Field, |
| 8 | + isBindingSchema, |
| 9 | + isChoiceDeclarationField, |
| 10 | + isChoiceInstanceField, |
| 11 | + isNestedIdentifier, |
| 12 | + isNotChoiceDeclarationField, |
| 13 | + isPrimitiveTypeSchema, |
| 14 | + isSpecializationTypeSchema, |
| 15 | + isValueSetTypeSchema, |
| 16 | + type NestedType, |
| 17 | + type RegularTypeSchema, |
| 18 | + type TypeSchema, |
| 19 | +} from "./types"; |
| 20 | +import { mkTypeSchemaIndex, type TypeSchemaIndex } from "./utils"; |
| 21 | + |
| 22 | +export type TreeShake = Record<string, Record<string, TreeShakeRule>>; |
| 23 | + |
| 24 | +export type TreeShakeRule = { ignoreFields?: string[]; selectFields?: string[] }; |
| 25 | + |
| 26 | +const mutableSelectFields = (schema: RegularTypeSchema, selectFields: string[]) => { |
| 27 | + const selectedFields: Record<string, Field> = {}; |
| 28 | + |
| 29 | + const selectPolimorphic: Record<string, { declaration?: string[]; instances?: string[] }> = {}; |
| 30 | + for (const fieldName of selectFields) { |
| 31 | + const field = schema.fields?.[fieldName]; |
| 32 | + if (!schema.fields || !field) throw new Error(`Field ${fieldName} not found`); |
| 33 | + |
| 34 | + if (isChoiceDeclarationField(field)) { |
| 35 | + if (!selectPolimorphic[fieldName]) selectPolimorphic[fieldName] = {}; |
| 36 | + selectPolimorphic[fieldName].declaration = field.choices; |
| 37 | + } else if (isChoiceInstanceField(field)) { |
| 38 | + const choiceName = field.choiceOf; |
| 39 | + if (!selectPolimorphic[choiceName]) selectPolimorphic[choiceName] = {}; |
| 40 | + selectPolimorphic[choiceName].instances = [...(selectPolimorphic[choiceName].instances ?? []), fieldName]; |
| 41 | + } else { |
| 42 | + selectedFields[fieldName] = field; |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + for (const [choiceName, { declaration, instances }] of Object.entries(selectPolimorphic)) { |
| 47 | + const choices = instances ?? declaration; |
| 48 | + assert(choices); |
| 49 | + for (const choiceInstanceName of choices) { |
| 50 | + const field = schema.fields?.[choiceInstanceName]; |
| 51 | + assert(field); |
| 52 | + selectedFields[choiceInstanceName] = field; |
| 53 | + } |
| 54 | + const decl = schema.fields?.[choiceName]; |
| 55 | + assert(decl); |
| 56 | + selectedFields[choiceName] = { ...decl, choices: choices }; |
| 57 | + } |
| 58 | + schema.fields = selectedFields; |
| 59 | +}; |
| 60 | + |
| 61 | +const mutableIgnoreFields = (schema: RegularTypeSchema, ignoreFields: string[]) => { |
| 62 | + for (const fieldName of ignoreFields) { |
| 63 | + const field = schema.fields?.[fieldName]; |
| 64 | + if (!schema.fields || !field) throw new Error(`Field ${fieldName} not found`); |
| 65 | + if (schema.fields) { |
| 66 | + if (isChoiceDeclarationField(field)) { |
| 67 | + for (const choiceName of field.choices) { |
| 68 | + delete schema.fields[choiceName]; |
| 69 | + } |
| 70 | + } |
| 71 | + |
| 72 | + if (isChoiceInstanceField(field)) { |
| 73 | + const choiceDeclaration = schema.fields[field.choiceOf]; |
| 74 | + assert(isChoiceDeclarationField(choiceDeclaration)); |
| 75 | + choiceDeclaration.choices = choiceDeclaration.choices.filter((c) => c !== fieldName); |
| 76 | + if (choiceDeclaration.choices.length === 0) { |
| 77 | + delete schema.fields[field.choiceOf]; |
| 78 | + } |
| 79 | + } |
| 80 | + |
| 81 | + delete schema.fields[fieldName]; |
| 82 | + } |
| 83 | + } |
| 84 | +}; |
| 85 | + |
| 86 | +export const treeShakeTypeSchema = (schema: TypeSchema, rule: TreeShakeRule, _logger?: CodegenLogger): TypeSchema => { |
| 87 | + schema = structuredClone(schema); |
| 88 | + if (isPrimitiveTypeSchema(schema) || isValueSetTypeSchema(schema) || isBindingSchema(schema)) return schema; |
| 89 | + |
| 90 | + if (rule.selectFields) { |
| 91 | + if (rule.ignoreFields) throw new Error("Cannot use both ignoreFields and selectFields in the same rule"); |
| 92 | + mutableSelectFields(schema, rule.selectFields); |
| 93 | + } |
| 94 | + |
| 95 | + if (rule.ignoreFields) { |
| 96 | + if (rule.selectFields) throw new Error("Cannot use both ignoreFields and selectFields in the same rule"); |
| 97 | + mutableIgnoreFields(schema, rule.ignoreFields); |
| 98 | + } |
| 99 | + |
| 100 | + if (schema.nested) { |
| 101 | + const usedTypes = new Set<CanonicalUrl>(); |
| 102 | + const collectUsedNestedTypes = (s: RegularTypeSchema | NestedType) => { |
| 103 | + Object.values(s.fields ?? {}) |
| 104 | + .filter(isNotChoiceDeclarationField) |
| 105 | + .filter((f) => isNestedIdentifier(f.type)) |
| 106 | + .forEach((f) => { |
| 107 | + const url = f.type.url; |
| 108 | + if (!usedTypes.has(url)) { |
| 109 | + usedTypes.add(url); |
| 110 | + const nestedTypeDef = schema.nested?.find((f) => f.identifier.url === url); |
| 111 | + assert(nestedTypeDef); |
| 112 | + collectUsedNestedTypes(nestedTypeDef); |
| 113 | + } |
| 114 | + }); |
| 115 | + }; |
| 116 | + collectUsedNestedTypes(schema); |
| 117 | + schema.nested = schema.nested.filter((n) => usedTypes.has(n.identifier.url)); |
| 118 | + } |
| 119 | + |
| 120 | + schema.dependencies = extractDependencies(schema.identifier, schema.base, schema.fields, schema.nested); |
| 121 | + return schema; |
| 122 | +}; |
| 123 | + |
| 124 | +export const treeShake = ( |
| 125 | + tsIndex: TypeSchemaIndex, |
| 126 | + treeShake: TreeShake, |
| 127 | + { resolutionTree, logger }: { resolutionTree?: ResolutionTree; logger?: CodegenLogger }, |
| 128 | +): TypeSchemaIndex => { |
| 129 | + const focusedSchemas: TypeSchema[] = []; |
| 130 | + for (const [pkgId, requires] of Object.entries(treeShake)) { |
| 131 | + for (const [url, rule] of Object.entries(requires)) { |
| 132 | + const schema = tsIndex.resolveByUrl(pkgId, url as CanonicalUrl); |
| 133 | + if (!schema) throw new Error(`Schema not found for ${pkgId} ${url}`); |
| 134 | + const shaked = treeShakeTypeSchema(schema, rule); |
| 135 | + focusedSchemas.push(shaked); |
| 136 | + } |
| 137 | + } |
| 138 | + const collectDeps = (schemas: TypeSchema[], acc: Record<string, TypeSchema>): TypeSchema[] => { |
| 139 | + if (schemas.length === 0) return Object.values(acc); |
| 140 | + for (const schema of schemas) { |
| 141 | + acc[JSON.stringify(schema.identifier)] = schema; |
| 142 | + } |
| 143 | + |
| 144 | + const newSchemas: TypeSchema[] = []; |
| 145 | + |
| 146 | + for (const schema of schemas) { |
| 147 | + if (isSpecializationTypeSchema(schema)) { |
| 148 | + if (!schema.dependencies) continue; |
| 149 | + schema.dependencies.forEach((dep) => { |
| 150 | + const depSchema = tsIndex.resolve(dep); |
| 151 | + if (!depSchema) |
| 152 | + throw new Error( |
| 153 | + `Dependent schema ${JSON.stringify(dep)} not found for ${JSON.stringify(schema.identifier)}`, |
| 154 | + ); |
| 155 | + const id = JSON.stringify(depSchema.identifier); |
| 156 | + if (!acc[id]) newSchemas.push(depSchema); |
| 157 | + }); |
| 158 | + if (schema.nested) { |
| 159 | + for (const nest of schema.nested) { |
| 160 | + if (isNestedIdentifier(nest.identifier)) continue; |
| 161 | + const id = JSON.stringify(nest.identifier); |
| 162 | + if (!acc[id]) newSchemas.push(nest); |
| 163 | + } |
| 164 | + } |
| 165 | + } |
| 166 | + } |
| 167 | + return collectDeps(newSchemas, acc); |
| 168 | + }; |
| 169 | + |
| 170 | + const shaked = collectDeps(focusedSchemas, {}); |
| 171 | + return mkTypeSchemaIndex(shaked, { resolutionTree, logger }); |
| 172 | +}; |
0 commit comments