feat: add discriminated union types to nitrogen

apps/example/src/getTests.ts

@@ -1326,6 +1326,16 @@ export function getTests(
         .didNotThrow()
         .equals('Hello world!')
     ),
+    createTest('bounceVehicle(truck) works', () =>
+      it(() => testObject.bounceVehicle({ kind: 'truck', payload: 1000 }))
+        .didNotThrow()
+        .equals({ kind: 'truck', payload: 1000 })
+    ),
+    createTest('bounceVehicle(boat) works', () =>
+      it(() => testObject.bounceVehicle({ kind: 'boat', lengthMeters: 12.5 }))
+        .didNotThrow()
+        .equals({ kind: 'boat', lengthMeters: 12.5 })
+    ),
 
     // More complex variants...
     ...('getVariantTuple' in testObject

packages/nitrogen/src/syntax/c++/CppDiscriminatedUnion.ts

@@ -0,0 +1,137 @@
+import type { FileWithReferencedTypes } from '../SourceFile.js'
+import { createFileMetadataString, isNotDuplicate } from '../helpers.js'
+import { indent } from '../../utils.js'
+import { includeHeader, includeNitroHeader } from './includeNitroHeader.js'
+import { NitroConfig } from '../../config/NitroConfig.js'
+import type { DiscriminatedUnionType } from '../types/DiscriminatedUnionType.js'
+
+/**
+ * Generates a C++ JSIConverter specialization for a discriminated union.
+ * The discriminant property (e.g. `kind`) is read first to decide which
+ * struct to deserialize into, then delegated to that struct's own JSIConverter.
+ */
+export function createCppDiscriminatedUnion(
+  union: DiscriminatedUnionType
+): FileWithReferencedTypes {
+  const { unionName, discriminantKey, variants } = union
+  const fullyQualifiedVariants = variants
+    .map((v) => v.type.getCode('c++', { fullyQualified: true }))
+    .filter(isNotDuplicate)
+  const cxxVariantType = `std::variant<${fullyQualifiedVariants.join(', ')}>`
+
+  // fromJSI: switch on the discriminant string value
+  const fromJsiCases = variants
+    .map(
+      (v) =>
+        `case hashString("${v.discriminantValue}"): return JSIConverter<${v.type.getCode('c++', { fullyQualified: true })}>::fromJSI(runtime, arg);`
+    )
+    .join('\n')
+
+  // toJSI: generic lambda dispatch — serialize struct then inject discriminant key back
+  const toJsiDiscriminants = variants
+    .map(
+      (v) =>
+        `if constexpr (std::is_same_v<T, ${v.type.getCode('c++', { fullyQualified: true })}>)\n  obj.setProperty(runtime, PropNameIDCache::get(runtime, "${discriminantKey}"), JSIConverter<std::string>::toJSI(runtime, "${v.discriminantValue}"));`
+    )
+    .join('\nelse ')
+  // canConvert: check discriminant is present and is a known value
+  const canConvertCases = variants
+    .map((v) => `case hashString("${v.discriminantValue}"):`)
+    .join('\n')
+
+  // Includes for each constituent struct
+  const includedTypes = variants.flatMap((v) =>
+    v.type.getRequiredImports('c++')
+  )
+  const forwardDeclarations = includedTypes
+    .map((i) => i.forwardDeclaration)
+    .filter((v) => v != null)
+    .filter(isNotDuplicate)
+  const extraIncludes = includedTypes
+    .map((i) => includeHeader(i))
+    .filter(isNotDuplicate)
+
+  const cxxNamespace = NitroConfig.current.getCxxNamespace('c++')
+
+  const code = `
+${createFileMetadataString(`${unionName}.hpp`)}
+
+#pragma once
+
+#include <variant>
+${includeNitroHeader('NitroHash.hpp')}
+${includeNitroHeader('JSIConverter.hpp')}
+${includeNitroHeader('NitroDefines.hpp')}
+${includeNitroHeader('JSIHelpers.hpp')}
+${includeNitroHeader('PropNameIDCache.hpp')}
+
+${forwardDeclarations.join('\n')}
+
+${extraIncludes.join('\n')}
+
+namespace margelo::nitro {
+
+  // C++ ${cxxNamespace}::${unionName} <> JS ${unionName} (discriminated union on "${discriminantKey}")
+  template <>
+  struct JSIConverter<${cxxVariantType}> final {
+    static inline ${cxxVariantType} fromJSI(jsi::Runtime& runtime, const jsi::Value& arg) {
+      jsi::Object obj = arg.asObject(runtime);
+      std::string discriminant = JSIConverter<std::string>::fromJSI(
+        runtime, obj.getProperty(runtime, PropNameIDCache::get(runtime, "${discriminantKey}"))
+      );
+      switch (hashString(discriminant.c_str(), discriminant.size())) {
+        ${indent(fromJsiCases, '        ')}
+        default: [[unlikely]]
+          throw std::invalid_argument(
+            "Cannot convert JS object to ${unionName}: unknown discriminant \\"" + discriminant + "\\" for key \\"${discriminantKey}\\"!"
+          );
+      }
+    }
+    static inline jsi::Value toJSI(jsi::Runtime& runtime, const ${cxxVariantType}& arg) {
+      return std::visit(
+        [&runtime](const auto& val) {
+          // Serialize the struct, then inject the discriminant key back
+          // so JS receives the full discriminated object (e.g. { kind: 'truck', payload: 1000 })
+          using T = std::decay_t<decltype(val)>;
+          jsi::Value result = JSIConverter<T>::toJSI(runtime, val);
+          jsi::Object obj = result.asObject(runtime);
+          ${indent(toJsiDiscriminants, '          ')}
+          return jsi::Value(runtime, obj);
+        },
+        arg
+      );
+    }
+    static inline bool canConvert(jsi::Runtime& runtime, const jsi::Value& value) {
+      if (!value.isObject()) {
+        return false;
+      }
+      jsi::Object obj = value.getObject(runtime);
+      if (!nitro::isPlainObject(runtime, obj)) {
+        return false;
+      }
+      jsi::Value discriminantVal = obj.getProperty(runtime, PropNameIDCache::get(runtime, "${discriminantKey}"));
+      if (!discriminantVal.isString()) {
+        return false;
+      }
+      std::string discriminant = JSIConverter<std::string>::fromJSI(runtime, discriminantVal);
+      switch (hashString(discriminant.c_str(), discriminant.size())) {
+        ${indent(canConvertCases, '        ')}
+          return true;
+        default:
+          return false;
+      }
+    }
+  };
+
+} // namespace margelo::nitro
+  `
+
+  return {
+    content: code,
+    name: `${unionName}.hpp`,
+    subdirectory: [],
+    language: 'c++',
+    referencedTypes: variants.map((v) => v.type),
+    platform: 'shared',
+  }
+}

packages/nitrogen/src/syntax/createType.ts

@@ -17,6 +17,10 @@ import { OptionalType } from './types/OptionalType.js'
 import { NamedWrappingType } from './types/NamedWrappingType.js'
 import { getInterfaceProperties } from './getInterfaceProperties.js'
 import { VariantType } from './types/VariantType.js'
+import {
+  DiscriminatedUnionType,
+  type DiscriminatedVariant,
+} from './types/DiscriminatedUnionType.js'
 import { MapType } from './types/MapType.js'
 import { TupleType } from './types/TupleType.js'
 import {
@@ -227,6 +231,49 @@ export function addKnownType(
   knownTypes[language].set(key, type)
 }
 
+/**
+ * Checks if a union of object types is a discriminated union.
+ * Returns the discriminant key if every member has a unique string literal
+ * for the same property, otherwise returns null.
+ */
+function findDiscriminantKey(types: TSMorphType[]): string | null {
+  // Collect all property names from the first member
+  const firstProps = types[0]?.getProperties() ?? []
+  for (const prop of firstProps) {
+    const key = prop.getName()
+    const literalValues: string[] = []
+    let valid = true
+    for (const t of types) {
+      const p = t.getProperty(key)
+      if (p == null) {
+        valid = false
+        break
+      }
+      const decl = p.getDeclarations()[0] ?? p.getValueDeclaration()
+      if (decl == null) {
+        valid = false
+        break
+      }
+      const propType = p.getTypeAtLocation(decl)
+      if (!propType.isStringLiteral()) {
+        valid = false
+        break
+      }
+      const lit = propType.getLiteralValue()
+      if (typeof lit !== 'string') {
+        valid = false
+        break
+      }
+      literalValues.push(lit)
+    }
+    // Valid discriminant: all members have the key, all values are unique string literals
+    if (valid && new Set(literalValues).size === types.length) {
+      return key
+    }
+  }
+  return null
+}
+
 /**
  * Create a new type (or return it from cache if it is already known)
  */
@@ -372,6 +419,45 @@ export function createType(
         const typename = symbol.getEscapedName()
         return new EnumType(typename, type)
       } else {
+        // It consists of different types - check if it's a discriminated union first
+        const allAreObjects = nonNullTypes.every(
+          (t) => t.isInterface() || t.isObject()
+        )
+        if (allAreObjects && nonNullTypes.length >= 2) {
+          const discriminantKey = findDiscriminantKey(nonNullTypes)
+          if (discriminantKey != null) {
+            const name =
+              type.getAliasSymbol()?.getName() ?? 'UnknownDiscriminatedUnion'
+            const discriminatedVariants: DiscriminatedVariant[] =
+              nonNullTypes.map((t) => {
+                const prop = t.getProperty(discriminantKey)!
+                const propDecl =
+                  prop.getDeclarations()[0] ?? prop.getValueDeclaration()!
+                const discriminantValue = prop
+                  .getTypeAtLocation(propDecl)
+                  .getLiteralValueOrThrow() as string
+                const structName = (t.getAliasSymbol() ??
+                  t.getSymbol())!.getName()
+                // Strip the discriminant property from the struct — the JSIConverter
+                // for the union handles dispatch; the struct itself doesn't need it.
+                const filteredProps = getInterfaceProperties(
+                  language,
+                  t,
+                  new Set([discriminantKey])
+                )
+                return {
+                  discriminantValue,
+                  type: new StructType(structName, filteredProps),
+                }
+              })
+            return new DiscriminatedUnionType(
+              name,
+              discriminantKey,
+              discriminatedVariants
+            )
+          }
+        }
+
         // It consists of different types - that means it's a variant!
         const unionConstituents = getUnionConstituents(type, typeNode)
         let variants = unionConstituents

packages/nitrogen/src/syntax/getInterfaceProperties.ts

@@ -5,36 +5,40 @@ import type { Language } from '../getPlatformSpecs.js'
 
 export function getInterfaceProperties(
   language: Language,
-  interfaceType: Type<ts.ObjectType>
+  interfaceType: Type<ts.ObjectType>,
+  skipKeys?: Set<string>
 ): NamedType[] {
   const symbol = interfaceType.getAliasSymbol() ?? interfaceType.getSymbol()
   if (symbol == null)
     throw new Error(
       `Interface "${interfaceType.getText()}" does not have a Symbol!`
     )
-  return interfaceType.getProperties().map((prop) => {
-    const propDeclaration = prop
-      .getDeclarations()
-      .find((declaration) => Node.isPropertySignature(declaration))
-    let propType = prop.getDeclaredType()
-    if (propType.isAny() || propType.isUnknown()) {
-      // the interface is aliased/merged - we need to look into the actual declaration
-      for (const declaration of symbol.getDeclarations()) {
-        const declared = prop.getTypeAtLocation(declaration)
-        if (!declared.isAny() && !declared.isUnknown()) {
-          propType = declared
-          break
+  return interfaceType
+    .getProperties()
+    .filter((prop) => !skipKeys?.has(prop.getName()))
+    .map((prop) => {
+      const propDeclaration = prop
+        .getDeclarations()
+        .find((declaration) => Node.isPropertySignature(declaration))
+      let propType = prop.getDeclaredType()
+      if (propType.isAny() || propType.isUnknown()) {
+        // the interface is aliased/merged - we need to look into the actual declaration
+        for (const declaration of symbol.getDeclarations()) {
+          const declared = prop.getTypeAtLocation(declaration)
+          if (!declared.isAny() && !declared.isUnknown()) {
+            propType = declared
+            break
+          }
         }
       }
-    }
 
-    const refType = createNamedType(
-      language,
-      prop.getName(),
-      propType,
-      prop.isOptional() || propType.isNullable(),
-      propDeclaration?.getTypeNode()
-    )
-    return refType
-  })
+      const refType = createNamedType(
+        language,
+        prop.getName(),
+        propType,
+        prop.isOptional() || propType.isNullable(),
+        propDeclaration?.getTypeNode()
+      )
+      return refType
+    })
 }

packages/nitrogen/src/syntax/getReferencedTypes.ts

@@ -1,4 +1,5 @@
 import { ArrayType } from './types/ArrayType.js'
+import { DiscriminatedUnionType } from './types/DiscriminatedUnionType.js'
 import { FunctionType } from './types/FunctionType.js'
 import { getTypeAs } from './types/getTypeAs.js'
 import { OptionalType } from './types/OptionalType.js'
@@ -51,6 +52,10 @@ export function getReferencedTypes(type: Type): Type[] {
       const variant = getTypeAs(type, VariantType)
       return [type, ...variant.variants.flatMap((t) => getReferencedTypes(t))]
 
+    case 'discriminated-union':
+      const du = getTypeAs(type, DiscriminatedUnionType)
+      return [type, ...du.variants.flatMap((v) => getReferencedTypes(v.type))]
+
     default:
       return [type]
   }

packages/nitrogen/src/syntax/helpers.ts

@@ -132,6 +132,9 @@ function getTypeLooselyness(type: Type): number {
     case 'variant':
       // Pretty loose
       return 2
+    case 'discriminated-union':
+      // Pretty loose — same as variant
+      return 2
     case 'result-wrapper':
       // Not loose at all
       return 0