TypeSchema: followReferences tree-shake option to generate referenced types

CLAUDE.md

@@ -319,7 +319,7 @@ Detection uses `mkIsFamilyType(tsIndex)` which checks `schema.typeFamily.resourc
 ### Working with Tree Shaking
 - Configured via `builder.typeSchema({ treeShake: {...} })`
 - Specify which resources and fields to include
-- Automatically resolves dependencies
+- Automatically resolves dependencies; reference targets are not dependencies — opt in per rule with `followReferences: true` (non-transitive), or for all roots via `treeShakeDefaults: { followReferences: true }`
 - Reference format: `"hl7.fhir.r4.core#4.0.1"`
 
 ## Key File Locations

README.md

@@ -279,7 +279,7 @@ Tree shaking optimizes the generated output by including only the resources you
 })
 ```
 
-This feature automatically resolves and includes all dependencies (referenced types, base resources, nested types, and extension definitions used by profiles) while excluding unused resources, significantly reducing the size of generated code and improving compilation times.
+This feature automatically resolves and includes all dependencies (field types, base resources, nested types, and extension definitions used by profiles) while excluding unused resources, significantly reducing the size of generated code and improving compilation times. Reference targets are not dependencies: a field typed `Reference<"Patient">` does not pull the `Patient` type into the output, since the reference is rendered as a string literal.
 
 ##### Field-Level Tree Shaking
 
@@ -311,6 +311,33 @@ Beyond resource-level filtering, tree shaking supports fine-grained field select
 
 FHIR choice types (like `multipleBirth[x]` which can be boolean or integer) are handled intelligently. Selecting/ignoring the base field affects all variants, while targeting specific variants only affects those types.
 
+##### Following Reference Targets
+
+To also generate types for a schema's reference targets — both the base resources (`reference.resource`) and the target profiles (`reference.profiles`) — enable `followReferences` on the rule (default: `false`):
+
+```typescript
+.typeSchema({
+    treeShake: {
+        "kbv.basis": {
+            "https://fhir.kbv.de/StructureDefinition/KBV_PR_Base_Condition_Diagnosis": {
+                followReferences: true
+            }
+        }
+    }
+})
+```
+
+With this rule, `Condition.subject` targeting `KBV_PR_Base_Patient` pulls the Patient profile (and the `Patient`/`Group` resources) into the generated output, including their regular dependencies. Following is **non-transitive**: the pulled-in types keep their own reference targets as plain string literals, so a single flag cannot drag in the whole package. If a followed target is also listed as a tree-shake root, the root's rule wins.
+
+To enable it for every tree-shake root at once, use `treeShakeDefaults` (a rule's own `followReferences` still wins over the default):
+
+```typescript
+.typeSchema({
+    treeShake: { ... },
+    treeShakeDefaults: { followReferences: true }
+})
+```
+
 #### Logical Model Promotion
 
 Some implementation guides expose logical models (logical-kind StructureDefinitions) that are intended to be used like resources in generated SDKs. The code generator supports promoting selected logical models to behave as resources during generation.

src/api/builder.ts

@@ -394,6 +394,10 @@ export class APIBuilder {
             assert(this.options.typeSchema.treeShake === undefined, "treeShake option is already set");
             this.options.typeSchema.treeShake = cfg.treeShake;
         }
+        if (cfg.treeShakeDefaults) {
+            assert(this.options.typeSchema.treeShakeDefaults === undefined, "treeShakeDefaults option is already set");
+            this.options.typeSchema.treeShakeDefaults = cfg.treeShakeDefaults;
+        }
         if (cfg.promoteLogical) {
             assert(this.options.typeSchema.promoteLogical === undefined, "promoteLogical option is already set");
             this.options.typeSchema.promoteLogical = cfg.promoteLogical;
@@ -480,7 +484,12 @@ export class APIBuilder {
             };
             const tsIndexOpts = { register, irReport, logger: tsLogger };
             let tsIndex = mkTypeSchemaIndex(typeSchemas, tsIndexOpts);
-            if (this.options.typeSchema?.treeShake) tsIndex = treeShake(tsIndex, this.options.typeSchema.treeShake);
+            if (this.options.typeSchema?.treeShake)
+                tsIndex = treeShake(
+                    tsIndex,
+                    this.options.typeSchema.treeShake,
+                    this.options.typeSchema.treeShakeDefaults,
+                );
             if (this.options.typeSchema?.promoteLogical)
                 tsIndex = promoteLogical(tsIndex, this.options.typeSchema.promoteLogical);
 

src/typeschema/ir/tree-shake.ts

@@ -6,6 +6,7 @@ import {
     concatIdentifiers,
     extractExtensionDeps,
     type Field,
+    type Identifier,
     isBindingSchema,
     isChoiceDeclarationField,
     isChoiceInstanceField,
@@ -23,7 +24,7 @@ import {
     type TypeSchema,
 } from "../types";
 import type { TypeSchemaIndex } from "../utils";
-import type { IrReport, TreeShakeConf, TreeShakeReport, TreeShakeRule } from "./types";
+import type { IrReport, TreeShakeConf, TreeShakeDefaults, TreeShakeReport, TreeShakeRule } from "./types";
 
 const ensureIrReport = (indexOrReport: TypeSchemaIndex | IrReport): IrReport => {
     if ("irReport" in indexOrReport && typeof indexOrReport.irReport === "function") {
@@ -244,16 +245,47 @@ export const treeShakeTypeSchema = (schema: TypeSchema, rule: TreeShakeRule, _lo
     return schema;
 };
 
-export const treeShake = (tsIndex: TypeSchemaIndex, treeShake: TreeShakeConf): TypeSchemaIndex => {
+/** Reference target identifiers (base resources and target profiles) of a schema's fields, including nested types. */
+const collectReferenceTargets = (schema: TypeSchema): Identifier[] => {
+    if (!isSpecializationTypeSchema(schema) && !isProfileTypeSchema(schema)) return [];
+    const fieldSets = [schema.fields, ...(schema.nested ?? []).map((n) => n.fields)];
+    return fieldSets.flatMap((fields) =>
+        Object.values(fields ?? {})
+            .filter(isNotChoiceDeclarationField)
+            .flatMap((f) => [...(f.reference?.resource ?? []), ...(f.reference?.profiles ?? [])])
+            .filter((id): id is Identifier => !isNestedIdentifier(id)),
+    );
+};
+
+export const treeShake = (
+    tsIndex: TypeSchemaIndex,
+    treeShake: TreeShakeConf,
+    defaults?: TreeShakeDefaults,
+): TypeSchemaIndex => {
     const focusedSchemas: TypeSchema[] = [];
+    const followedSchemas: TypeSchema[] = [];
     for (const [pkgId, requires] of Object.entries(treeShake)) {
         for (const [url, rule] of Object.entries(requires)) {
             const schema = tsIndex.resolveByUrl(pkgId, url as CanonicalUrl);
             if (!schema || isNestedTypeSchema(schema)) throw new Error(`Schema not found for ${pkgId} ${url}`);
             const shaked = treeShakeTypeSchema(schema, rule);
             focusedSchemas.push(shaked);
+            if (rule.followReferences ?? defaults?.followReferences) {
+                for (const refId of collectReferenceTargets(shaked)) {
+                    const refSchema = tsIndex.resolve(refId);
+                    if (!refSchema)
+                        throw new Error(`Reference target ${JSON.stringify(refId)} not found for ${pkgId} ${url}`);
+                    followedSchemas.push(refSchema);
+                }
+            }
         }
     }
+    // Explicit roots win over followed reference targets: a root's rule must not
+    // be overridden by the untouched schema pulled in via followReferences.
+    const rootIds = new Set(focusedSchemas.map((s) => JSON.stringify(s.identifier)));
+    for (const schema of followedSchemas) {
+        if (!rootIds.has(JSON.stringify(schema.identifier))) focusedSchemas.push(schema);
+    }
     const collectDeps = (schemas: TypeSchema[], acc: Record<string, TypeSchema>): TypeSchema[] => {
         if (schemas.length === 0) return Object.values(acc);
         for (const schema of schemas) {

src/typeschema/ir/types.ts

@@ -17,15 +17,27 @@ export type ResolveCollisionsConf = Record<string, CollisionResolution>;
 
 export type IrConf = {
     treeShake?: TreeShakeConf;
+    /** Rule defaults applied to every treeShake root; a rule's own value wins. */
+    treeShakeDefaults?: TreeShakeDefaults;
     promoteLogical?: LogicalPromotionConf;
     resolveCollisions?: ResolveCollisionsConf;
 };
 
+export type TreeShakeDefaults = Pick<TreeShakeRule, "followReferences">;
+
 export type LogicalPromotionConf = Record<PkgName, CanonicalUrl[]>;
 
 export type TreeShakeConf = Record<string, Record<string, TreeShakeRule>>;
 
-export type TreeShakeRule = { ignoreFields?: string[]; selectFields?: string[]; ignoreExtensions?: string[] };
+export type TreeShakeRule = {
+    ignoreFields?: string[];
+    selectFields?: string[];
+    ignoreExtensions?: string[];
+    /** Also generate types for the schema's reference targets (base resources
+     *  and target profiles). Non-transitive: followed types keep their own
+     *  reference targets as plain string literals. Default: false. */
+    followReferences?: boolean;
+};
 
 export type IrReport = {
     treeShake?: TreeShakeReport;

test/unit/typeschema/ir/tree-shake.test.ts

@@ -59,6 +59,66 @@ describe("treeShake specific TypeSchema", async () => {
             expect(shaked.entityTree()).toMatchSnapshot();
         });
     });
+
+    describe("followReferences", () => {
+        const patientUrl = "http://hl7.org/fhir/StructureDefinition/Patient" as CanonicalUrl;
+        const observationRule = (rule: object) =>
+            treeShake(tsIndex, {
+                "hl7.fhir.r4.core": { "http://hl7.org/fhir/StructureDefinition/Observation": rule },
+            });
+
+        it("drops reference targets by default", () => {
+            const shaked = observationRule({});
+            expect(shaked.resolveByUrl("hl7.fhir.r4.core", patientUrl)).toBeUndefined();
+        });
+
+        it("keeps reference targets when enabled", () => {
+            const shaked = observationRule({ followReferences: true });
+            const patient = shaked.resolveByUrl("hl7.fhir.r4.core", patientUrl);
+            expect(patient).toBeDefined();
+            expect(patient?.identifier.kind).toBe("resource");
+            // Nested-type references are followed too (Observation.component has none,
+            // but Observation.performer lives on the root; device covers another target).
+            const device = shaked.resolveByUrl(
+                "hl7.fhir.r4.core",
+                "http://hl7.org/fhir/StructureDefinition/Device" as CanonicalUrl,
+            );
+            expect(device).toBeDefined();
+        });
+
+        it("treeShakeDefaults enables following for every root", () => {
+            const shaked = treeShake(
+                tsIndex,
+                { "hl7.fhir.r4.core": { "http://hl7.org/fhir/StructureDefinition/Observation": {} } },
+                { followReferences: true },
+            );
+            expect(shaked.resolveByUrl("hl7.fhir.r4.core", patientUrl)).toBeDefined();
+        });
+
+        it("rule-level followReferences wins over the default", () => {
+            const shaked = treeShake(
+                tsIndex,
+                {
+                    "hl7.fhir.r4.core": {
+                        "http://hl7.org/fhir/StructureDefinition/Observation": { followReferences: false },
+                    },
+                },
+                { followReferences: true },
+            );
+            expect(shaked.resolveByUrl("hl7.fhir.r4.core", patientUrl)).toBeUndefined();
+        });
+
+        it("does not follow references transitively", () => {
+            // Account is referenced by Encounter (followed) but not by Observation
+            // itself — followed types keep their own reference targets as literals.
+            const shaked = observationRule({ followReferences: true });
+            const account = shaked.resolveByUrl(
+                "hl7.fhir.r4.core",
+                "http://hl7.org/fhir/StructureDefinition/Account" as CanonicalUrl,
+            );
+            expect(account).toBeUndefined();
+        });
+    });
 });
 
 describe("treeShake specific TypeSchema", async () => {