TypeSchema: Refactor type system for discriminated unions and type-safe resolution

CLAUDE.md

@@ -123,6 +123,7 @@ FHIR Package → TypeSchema Generator → TypeSchema Format → Code Generators
 - In code generators (writer-generator): use `curlyBlock` and `squareBlock` helpers for writing structured output instead of manual indent/deindent or string concatenation
 - Use `Record` instead of `Map` unless there is a significant reason for `Map` (e.g. non-string keys, iteration order guarantees, frequent deletion)
 - Prefer single-line guard clauses without braces: `if (!x) throw new Error("...");` instead of wrapping in `{ }`
+- Do not check `kind` of `Identifier`/`TypeIdentifier`/`TypeSchema` by manually comparing the `kind` field. Use dedicated predicates (`isPrimitiveIdentifier`, `isSpecializationTypeSchema`, etc.)
 
 ### Testing Strategy
 - Uses Bun's built-in test runner

docs/guides/typeschema-index.md

@@ -48,13 +48,13 @@ Why TypeSchemaIndex Matters
 Query schemas by type category:
 
 ```typescript
-collectComplexTypes(): SpecializationTypeSchema[]
+collectComplexTypes(): ComplexTypeTypeSchema[]
   Returns all complex types (datatypes, backbone elements)
 
-collectResources(): SpecializationTypeSchema[]
+collectResources(): ResourceTypeSchema[]
   Returns all FHIR resources (Patient, Observation, etc.)
 
-collectLogicalModels(): SpecializationTypeSchema[]
+collectLogicalModels(): LogicalTypeSchema[]
   Returns all logical models
 
 collectProfiles(): ProfileTypeSchema[]

examples/typescript-r4/fhir-types/type-tree.yaml

@@ -64,9 +64,7 @@ hl7.fhir.r4.core:
     http://hl7.org/fhir/StructureDefinition/Patient: {}
     http://hl7.org/fhir/StructureDefinition/Resource: {}
   value-set: {}
-  nested:
-    http://hl7.org/fhir/StructureDefinition/Observation#component: {}
-    http://hl7.org/fhir/StructureDefinition/Observation#referenceRange: {}
+  nested: {}
   binding: {}
   profile:
     http://hl7.org/fhir/StructureDefinition/patient-birthPlace: {}

examples/typescript-us-core/fhir-types/type-tree.yaml

@@ -85,10 +85,7 @@ hl7.fhir.r4.core:
     http://hl7.org/fhir/StructureDefinition/Patient: {}
     http://hl7.org/fhir/StructureDefinition/Resource: {}
   value-set: {}
-  nested:
-    http://hl7.org/fhir/StructureDefinition/Patient#communication: {}
-    http://hl7.org/fhir/StructureDefinition/Observation#component: {}
-    http://hl7.org/fhir/StructureDefinition/Observation#referenceRange: {}
+  nested: {}
   binding: {}
   profile:
     http://hl7.org/fhir/StructureDefinition/vitalsigns: {}

src/api/mustache/generator/ViewModelFactory.ts

@@ -15,9 +15,9 @@ import type { IsPrefixed } from "@root/utils/types";
 import {
     type ChoiceFieldInstance,
     type Field,
-    isComplexTypeIdentifier,
+    isComplexTypeTypeSchema,
     isNotChoiceDeclarationField,
-    isResourceIdentifier,
+    isResourceTypeSchema,
     type NestedTypeSchema,
     type RegularField,
     type TypeIdentifier,
@@ -100,7 +100,7 @@ export class ViewModelFactory {
         cache: ViewModelCache,
         nestedIn?: TypeSchema,
     ): TypeViewModel {
-        const type = this.tsIndex.resolve(typeRef);
+        const type = this.tsIndex.resolveType(typeRef);
         if (!type) {
             throw new Error(`ComplexType ${typeRef.name} not found`);
         }
@@ -113,7 +113,7 @@ export class ViewModelFactory {
     }
 
     private _createForResource(typeRef: TypeIdentifier, cache: ViewModelCache, nestedIn?: TypeSchema): TypeViewModel {
-        const type = this.tsIndex.resolve(typeRef);
+        const type = this.tsIndex.resolveType(typeRef);
         if (!type) {
             throw new Error(`Resource ${typeRef.name} not found`);
         }
@@ -126,14 +126,14 @@ export class ViewModelFactory {
     }
 
     private _createChildrenFor(typeRef: TypeIdentifier, cache: ViewModelCache, nestedIn?: TypeSchema): TypeViewModel[] {
-        const schema = this.tsIndex.resolve(typeRef);
-        if (!schema || !("typeFamily" in schema)) return [];
-        if (isComplexTypeIdentifier(typeRef)) {
+        const schema = this.tsIndex.resolveType(typeRef);
+        if (!schema) return [];
+        if (isComplexTypeTypeSchema(schema)) {
             return (schema.typeFamily?.complexTypes ?? [])
                 .filter(this.filterPred)
                 .map((childRef: TypeIdentifier) => this._createFor(childRef, cache, nestedIn));
         }
-        if (isResourceIdentifier(typeRef)) {
+        if (isResourceTypeSchema(schema)) {
             return (schema.typeFamily?.resources ?? [])
                 .filter(this.filterPred)
                 .map((childRef: TypeIdentifier) => this._createFor(childRef, cache, nestedIn));
@@ -146,7 +146,7 @@ export class ViewModelFactory {
         let parentRef: TypeIdentifier | undefined = "base" in base ? base.base : undefined;
         while (parentRef) {
             parents.push(this._createFor(parentRef, cache, undefined));
-            const parent = this.tsIndex.resolve(parentRef);
+            const parent = this.tsIndex.resolveType(parentRef);
             parentRef = parent && "base" in parent ? parent.base : undefined;
         }
         return parents;

src/api/writer-generator/csharp/csharp.ts

@@ -8,6 +8,7 @@ import type { Field, RegularField, TypeIdentifier } from "@typeschema/types";
 import {
     type ChoiceFieldInstance,
     isChoiceDeclarationField,
+    type NestedTypeSchema,
     type SpecializationTypeSchema,
 } from "@typeschema/types.ts";
 import type { TypeSchemaIndex } from "@typeschema/utils.ts";
@@ -53,12 +54,12 @@ const getFieldModifiers = (field: Field) => {
     return field.required ? ["required"] : [];
 };
 
-const formatClassName = (schema: SpecializationTypeSchema) => {
+const formatClassName = (schema: SpecializationTypeSchema | NestedTypeSchema) => {
     const name = prefixReservedTypeName(getResourceName(schema.identifier));
     return uppercaseFirstLetter(name);
 };
 
-const formatBaseClass = (schema: SpecializationTypeSchema) => {
+const formatBaseClass = (schema: SpecializationTypeSchema | NestedTypeSchema) => {
     return schema.base ? `: ${schema.base.name}` : "";
 };
 
@@ -138,7 +139,7 @@ export class CSharp extends Writer<CSharpGeneratorOptions> {
         this.generateHelperFile();
     }
 
-    private generateType(schema: SpecializationTypeSchema, packageName: string): void {
+    private generateType(schema: SpecializationTypeSchema | NestedTypeSchema, packageName: string): void {
         const className = formatClassName(schema);
         const baseClass = formatBaseClass(schema);
 
@@ -151,7 +152,7 @@ export class CSharp extends Writer<CSharpGeneratorOptions> {
         this.line();
     }
 
-    private generateFields(schema: SpecializationTypeSchema, packageName: string): void {
+    private generateFields(schema: SpecializationTypeSchema | NestedTypeSchema, packageName: string): void {
         if (!schema.fields) return;
 
         const sortedFields = Object.entries(schema.fields).sort(([a], [b]) => a.localeCompare(b));
@@ -161,7 +162,7 @@ export class CSharp extends Writer<CSharpGeneratorOptions> {
         }
     }
 
-    private generateNestedTypes(schema: SpecializationTypeSchema, packageName: string): void {
+    private generateNestedTypes(schema: SpecializationTypeSchema | NestedTypeSchema, packageName: string): void {
         if (!("nested" in schema) || !schema.nested) return;
 
         this.line();

src/api/writer-generator/python.ts

@@ -9,6 +9,7 @@ import {
     type EnumDefinition,
     type Field,
     isResourceTypeSchema,
+    type NestedTypeSchema,
     type SpecializationTypeSchema,
     type TypeIdentifier,
 } from "@typeschema/types.ts";
@@ -409,7 +410,7 @@ export class Python extends Writer<PythonGeneratorOptions> {
             this.pyImportFrom(`${this.opts.rootPackageName}.fhirpy_base_model`, "FhirpyBaseModel");
     }
 
-    private generateType(schema: SpecializationTypeSchema): void {
+    private generateType(schema: SpecializationTypeSchema | NestedTypeSchema): void {
         const className = deriveResourceName(schema.identifier);
         const superClasses = this.getSuperClasses(schema);
 
@@ -420,15 +421,15 @@ export class Python extends Writer<PythonGeneratorOptions> {
         this.line();
     }
 
-    private getSuperClasses(schema: SpecializationTypeSchema): string[] {
+    private getSuperClasses(schema: SpecializationTypeSchema | NestedTypeSchema): string[] {
         const bases: string[] = [];
         if (schema.base) bases.push(schema.base.name);
         bases.push(...this.injectSuperClasses(schema.identifier.url));
         if (schema.identifier.name in GENERIC_FIELD_REWRITES) bases.push("Generic[T]");
         return bases;
     }
 
-    private generateClassBody(schema: SpecializationTypeSchema): void {
+    private generateClassBody(schema: SpecializationTypeSchema | NestedTypeSchema): void {
         this.generateModelConfig();
 
         if (!schema.fields) {
@@ -473,7 +474,7 @@ export class Python extends Writer<PythonGeneratorOptions> {
         this.line(")");
     }
 
-    private generateFields(schema: SpecializationTypeSchema, schemaName: string): void {
+    private generateFields(schema: SpecializationTypeSchema | NestedTypeSchema, schemaName: string): void {
         const sortedFields = Object.entries(schema.fields ?? []).sort(([a], [b]) => a.localeCompare(b));
 
         for (const [fieldName, field] of sortedFields) {

src/api/writer-generator/typescript/profile-slices.ts

@@ -22,7 +22,7 @@ import type { TypeScript } from "./writer";
 /** Collect choice declaration field names from a base type schema */
 const collectChoiceBaseNames = (tsIndex: TypeSchemaIndex, typeId: TypeIdentifier): Set<string> => {
     const names = new Set<string>();
-    const schema = tsIndex.resolve(typeId);
+    const schema = tsIndex.resolveType(typeId);
     if (schema && "fields" in schema && schema.fields) {
         for (const [name, f] of Object.entries(schema.fields)) {
             if (isChoiceDeclarationField(f)) names.add(name);

src/api/writer-generator/typescript/profile.ts

@@ -168,7 +168,7 @@ const collectBaseRequiredParams = (
     coveredNames: string[],
 ) => {
     const covered = new Set(coveredNames);
-    const baseSchema = tsIndex.resolve(flatProfile.base);
+    const baseSchema = tsIndex.resolveType(flatProfile.base);
     if (!baseSchema || !("fields" in baseSchema) || !baseSchema.fields) return;
     for (const [name, field] of Object.entries(baseSchema.fields)) {
         if (covered.has(name)) continue;

src/api/writer-generator/typescript/writer.ts

@@ -11,6 +11,7 @@ import {
     isProfileTypeSchema,
     isResourceTypeSchema,
     isSpecializationTypeSchema,
+    type NestedTypeSchema,
     packageMeta,
     packageMetaToFhir,
     type SpecializationTypeSchema,
@@ -198,7 +199,7 @@ export class TypeScript extends Writer<TypeScriptOptions> {
         this.lineSM(`${extFieldName}?: ${typeExpr}`);
     }
 
-    generateType(tsIndex: TypeSchemaIndex, schema: SpecializationTypeSchema) {
+    generateType(tsIndex: TypeSchemaIndex, schema: SpecializationTypeSchema | NestedTypeSchema) {
         let name: string;
         // Generic types: Reference, Coding, CodeableConcept
         const genericTypes = ["Reference", "Coding", "CodeableConcept"];
@@ -212,7 +213,7 @@ export class TypeScript extends Writer<TypeScriptOptions> {
         const typeFamilyFields: { fieldName: string; familyTypeName: string }[] = [];
         for (const [fieldName, field] of Object.entries(schema.fields ?? {})) {
             if (isChoiceDeclarationField(field) || !field.type) continue;
-            const fieldTypeSchema = tsIndex.resolve(field.type);
+            const fieldTypeSchema = tsIndex.resolveType(field.type);
             if (
                 isSpecializationTypeSchema(fieldTypeSchema) &&
                 (fieldTypeSchema.typeFamily?.resources?.length ?? 0) > 0
@@ -285,7 +286,7 @@ export class TypeScript extends Writer<TypeScriptOptions> {
         });
     }
 
-    withPrimitiveTypeExtension(schema: TypeSchema): boolean {
+    withPrimitiveTypeExtension(schema: TypeSchema | NestedTypeSchema): boolean {
         if (!this.opts.primitiveTypeExtension) return false;
         if (!isSpecializationTypeSchema(schema)) return false;
         for (const field of Object.values(schema.fields ?? {})) {

src/typeschema/core/identifier.ts

@@ -8,11 +8,21 @@ import type { FHIRSchemaElement } from "@atomic-ehr/fhirschema";
 import type {
     BindingIdentifier,
     CanonicalUrl,
+    ComplexTypeIdentifier,
+    Identifier,
+    LogicalIdentifier,
     Name,
     PackageMeta,
+    PrimitiveIdentifier,
+    ProfileIdentifier,
+    ResourceIdentifier,
+    RichComplexTypeFHIRSchema,
     RichFHIRSchema,
+    RichLogicalFHIRSchema,
+    RichPrimitiveFHIRSchema,
+    RichProfileFHIRSchema,
+    RichResourceFHIRSchema,
     RichValueSet,
-    TypeIdentifier,
     ValueSetIdentifier,
 } from "@typeschema/types";
 import type { Register } from "../register";
@@ -27,23 +37,30 @@ function getVersionFromUrl(url: CanonicalUrl): string | undefined {
     return version;
 }
 
-function determineKind(fhirSchema: RichFHIRSchema): TypeIdentifier["kind"] {
-    if (fhirSchema.derivation === "constraint") return "profile";
-    if (fhirSchema.kind === "primitive-type") return "primitive-type";
-    if (fhirSchema.kind === "complex-type") return "complex-type";
-    if (fhirSchema.kind === "resource") return "resource";
-    if (fhirSchema.kind === "logical") return "logical";
-    return "resource";
-}
+const identifierBase = (fhirSchema: RichFHIRSchema) => ({
+    package: fhirSchema.package_meta.name,
+    version: fhirSchema.package_meta.version,
+    name: fhirSchema.name,
+    url: fhirSchema.url,
+});
 
-export function mkIdentifier(fhirSchema: RichFHIRSchema): TypeIdentifier {
-    return {
-        kind: determineKind(fhirSchema),
-        package: fhirSchema.package_meta.name,
-        version: fhirSchema.package_meta.version,
-        name: fhirSchema.name,
-        url: fhirSchema.url,
-    };
+export function mkIdentifier(fhirSchema: RichProfileFHIRSchema): ProfileIdentifier;
+export function mkIdentifier(fhirSchema: RichPrimitiveFHIRSchema): PrimitiveIdentifier;
+export function mkIdentifier(fhirSchema: RichComplexTypeFHIRSchema): ComplexTypeIdentifier;
+export function mkIdentifier(fhirSchema: RichResourceFHIRSchema): ResourceIdentifier;
+export function mkIdentifier(fhirSchema: RichLogicalFHIRSchema): LogicalIdentifier;
+export function mkIdentifier(
+    fhirSchema: RichComplexTypeFHIRSchema | RichResourceFHIRSchema | RichLogicalFHIRSchema,
+): ComplexTypeIdentifier | ResourceIdentifier | LogicalIdentifier;
+export function mkIdentifier(fhirSchema: RichFHIRSchema): Identifier;
+export function mkIdentifier(fhirSchema: RichFHIRSchema): Identifier {
+    const fields = identifierBase(fhirSchema);
+    if (fhirSchema.derivation === "constraint") return { kind: "profile", ...fields };
+    if (fhirSchema.kind === "primitive-type") return { kind: "primitive-type", ...fields };
+    if (fhirSchema.kind === "complex-type") return { kind: "complex-type", ...fields };
+    if (fhirSchema.kind === "resource") return { kind: "resource", ...fields };
+    if (fhirSchema.kind === "logical") return { kind: "logical", ...fields };
+    return { kind: "resource", ...fields };
 }
 
 const getValueSetName = (url: CanonicalUrl): Name => {