Fix union encoding/decoding (#1898)

So there are a handful of union issues and this fixes some of them.

First, type aliases that point at classes, enums or unions were not
encoding properly. That is fixed by this PR.
Second, unions were not decoding properly. That is also fixed.

Third, type definitions that are unions are not decoding properly, not
yet fixed but the first two felt meaningful enough to put up.
<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Fixes union encoding/decoding issues and adds recursive union test.
> 
>   - **Behavior**:
> - Fixes type alias encoding for classes, enums, and unions in
`cffi_generated.rs` and `ctypes.rs`.
> - Fixes union decoding in `types-unions.go.j2` by switching from
`ValueTypeIndex` to `VariantName`.
>   - **Schema**:
> - Adds `variant_name` field to `CFFIValueUnionVariant` in `cffi.fbs`.
>   - **Tests**:
> - Adds `recursive-union.baml` test file for recursive union testing.
>   - **Misc**:
> - Updates `GoTypeAlias` in `generate_types.rs` to include `is_union`
and `is_baml_serializable` flags.
> - Modifies `EncodeUnion` function in `encode.go` to include
`variantName`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 7b1f7bc. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

---------

Co-authored-by: hellovai <vbv@boundaryml.com>

Author: tberman

engine/language_client_cffi/src/cffi/cffi_generated.rs

@@ -2336,9 +2336,10 @@ impl<'a> flatbuffers::Follow<'a> for CFFIValueUnionVariant<'a> {
 
 impl<'a> CFFIValueUnionVariant<'a> {
   pub const VT_NAME: flatbuffers::VOffsetT = 4;
-  pub const VT_FIELD_TYPES: flatbuffers::VOffsetT = 6;
-  pub const VT_VALUE_TYPE_INDEX: flatbuffers::VOffsetT = 8;
-  pub const VT_VALUE: flatbuffers::VOffsetT = 10;
+  pub const VT_VARIANT_NAME: flatbuffers::VOffsetT = 6;
+  pub const VT_FIELD_TYPES: flatbuffers::VOffsetT = 8;
+  pub const VT_VALUE_TYPE_INDEX: flatbuffers::VOffsetT = 10;
+  pub const VT_VALUE: flatbuffers::VOffsetT = 12;
 
   #[inline]
   pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
@@ -2353,6 +2354,7 @@ impl<'a> CFFIValueUnionVariant<'a> {
     if let Some(x) = args.value { builder.add_value(x); }
     builder.add_value_type_index(args.value_type_index);
     if let Some(x) = args.field_types { builder.add_field_types(x); }
+    if let Some(x) = args.variant_name { builder.add_variant_name(x); }
     if let Some(x) = args.name { builder.add_name(x); }
     builder.finish()
   }
@@ -2366,6 +2368,13 @@ impl<'a> CFFIValueUnionVariant<'a> {
     unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<&str>>(CFFIValueUnionVariant::VT_NAME, None)}
   }
   #[inline]
+  pub fn variant_name(&self) -> Option<&'a str> {
+    // Safety:
+    // Created from valid Table for this object
+    // which contains a valid value in this slot
+    unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<&str>>(CFFIValueUnionVariant::VT_VARIANT_NAME, None)}
+  }
+  #[inline]
   pub fn field_types(&self) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder<'a>>>> {
     // Safety:
     // Created from valid Table for this object
@@ -2396,6 +2405,7 @@ impl flatbuffers::Verifiable for CFFIValueUnionVariant<'_> {
     use self::flatbuffers::Verifiable;
     v.visit_table(pos)?
      .visit_field::<flatbuffers::ForwardsUOffset<&str>>("name", Self::VT_NAME, false)?
+     .visit_field::<flatbuffers::ForwardsUOffset<&str>>("variant_name", Self::VT_VARIANT_NAME, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder>>>>("field_types", Self::VT_FIELD_TYPES, false)?
      .visit_field::<i32>("value_type_index", Self::VT_VALUE_TYPE_INDEX, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<CFFIValueHolder>>("value", Self::VT_VALUE, false)?
@@ -2405,6 +2415,7 @@ impl flatbuffers::Verifiable for CFFIValueUnionVariant<'_> {
 }
 pub struct CFFIValueUnionVariantArgs<'a> {
     pub name: Option<flatbuffers::WIPOffset<&'a str>>,
+    pub variant_name: Option<flatbuffers::WIPOffset<&'a str>>,
     pub field_types: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder<'a>>>>>,
     pub value_type_index: i32,
     pub value: Option<flatbuffers::WIPOffset<CFFIValueHolder<'a>>>,
@@ -2414,6 +2425,7 @@ impl<'a> Default for CFFIValueUnionVariantArgs<'a> {
   fn default() -> Self {
     CFFIValueUnionVariantArgs {
       name: None,
+      variant_name: None,
       field_types: None,
       value_type_index: 0,
       value: None,
@@ -2431,6 +2443,10 @@ impl<'a: 'b, 'b> CFFIValueUnionVariantBuilder<'a, 'b> {
     self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(CFFIValueUnionVariant::VT_NAME, name);
   }
   #[inline]
+  pub fn add_variant_name(&mut self, variant_name: flatbuffers::WIPOffset<&'b  str>) {
+    self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(CFFIValueUnionVariant::VT_VARIANT_NAME, variant_name);
+  }
+  #[inline]
   pub fn add_field_types(&mut self, field_types: flatbuffers::WIPOffset<flatbuffers::Vector<'b , flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder<'b >>>>) {
     self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(CFFIValueUnionVariant::VT_FIELD_TYPES, field_types);
   }
@@ -2461,6 +2477,7 @@ impl core::fmt::Debug for CFFIValueUnionVariant<'_> {
   fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
     let mut ds = f.debug_struct("CFFIValueUnionVariant");
       ds.field("name", &self.name());
+      ds.field("variant_name", &self.variant_name());
       ds.field("field_types", &self.field_types());
       ds.field("value_type_index", &self.value_type_index());
       ds.field("value", &self.value());

engine/language_client_cffi/src/ctypes.rs

@@ -489,14 +489,17 @@ where
             .iter()
             .position(|t| real_type == *t)
             .expect("Failed to find target_type in options");
+        let variant_name = options[value_type_index].to_union_name();
         let options = builder.create_vector_from_iter(options_vec.into_iter());
-        // TODO: get the name from the target_type
-        let name = builder.create_string(&target_type.to_union_name());
+
+        let name_offset = builder.create_string(&target_type.to_union_name());
+        let variant_name_offset = builder.create_string(&variant_name);
 
         let value_union_variant = CFFIValueUnionVariant::create(
             &mut builder,
             &CFFIValueUnionVariantArgs {
-                name: Some(name),
+                name: Some(name_offset),
+                variant_name: Some(variant_name_offset),
                 field_types: Some(options),
                 value_type_index: value_type_index as i32,
                 value: Some(value_holder),

engine/language_client_cffi/types/cffi.fbs

@@ -84,6 +84,7 @@ table CFFIValueTuple {
 // For the Rust variant `Union(Vec<CFFIFieldType>, Box<CFFIValue>)`
 table CFFIValueUnionVariant {
   name: string;
+  variant_name: string;
   field_types: [CFFIFieldTypeHolder];
   value_type_index: int;
   value: CFFIValueHolder;

engine/language_client_codegen/src/go/generate_types.rs

@@ -28,15 +28,6 @@ pub(crate) fn cast_value(container_variable_name: &str, field_type: &GoType) ->
                 .ok()
                 .unwrap()
         );
-    } else if field_type.is_slice {
-        let inner_type = field_type.underlying_type.as_ref().unwrap();
-        return format!(
-            r#"castSlice({container_variable_name}, func(item any) {} {{
-    return {}
-}})"#,
-            inner_type.name,
-            cast_value("item", inner_type),
-        );
     } else if field_type.is_union {
         return format!("*({container_variable_name}).(*{})", field_type.name);
     } else if field_type.is_pointer {
@@ -249,6 +240,8 @@ pub struct GoType {
 struct GoTypeAlias<'ir> {
     name: Cow<'ir, str>,
     target: String,
+    is_baml_serializable: bool,
+    is_union: bool,
 }
 
 #[derive(askama::Template)]
@@ -415,9 +408,12 @@ impl<'ir> From<ClassWalker<'ir>> for GoClass<'ir> {
 // TODO: Define AliasWalker to simplify type.
 impl<'ir> From<Walker<'ir, (&'ir String, &'ir FieldType)>> for GoTypeAlias<'ir> {
     fn from(walker: Walker<(&'ir String, &'ir FieldType)>) -> Self {
+        let type_ref = walker.item.1.to_type_ref_2(walker.ir, false);
         GoTypeAlias {
             name: Cow::Borrowed(walker.item.0),
-            target: walker.item.1.to_type_ref_2(walker.ir, false).name,
+            target: type_ref.name,
+            is_union: type_ref.is_union,
+            is_baml_serializable: type_ref.is_class || type_ref.is_enum || type_ref.is_union,
         }
     }
 }

engine/language_client_codegen/src/go/templates/client.go.j2

@@ -29,15 +29,6 @@ type stream struct {}
 
 var Stream = &stream{}
 
-func castSlice[T any](result any, castResult func(any) T) []T {
-	items := result.([]any)
-	casted := make([]T, len(items))
-	for i, item := range items {
-		casted[i] = castResult(item)
-	}
-	return casted
-}
-
 func castOptional[T any](result any, castResult func(any) T) *T {
 	if result == nil {
 		return nil

engine/language_client_codegen/src/go/templates/types-unions.go.j2

@@ -19,9 +19,9 @@ type {{ union.name }} struct {
 
 func (u *{{ union.name }}) Decode(holder *cffi.CFFIValueUnionVariant) {
     valueHolder := holder.Value(nil)
-	switch holder.ValueTypeIndex() {
+	switch string(holder.VariantName()) {
     {% for variant in union.variants %}
-	    case {{ loop.index0 }}:
+	    case "{{ variant.0 }}":
             u.variant = "{{ variant.0 }}"
             value := {{ crate::go::generate_types::render_value_coercion("valueHolder", variant.1) }}
             u.variant_{{ variant.0 }} = &value
@@ -34,7 +34,7 @@ func (u {{ union.name }}) Encode(builder *flatbuffers.Builder) (cffi.CFFIValueUn
     switch u.variant {
     {% for variant in union.variants %}
     case "{{ variant.0 }}":
-        return baml.EncodeUnion(builder, u.variant, u.variant_{{ variant.0 }})
+        return baml.EncodeUnion(builder, "{{ union.name }}", u.variant, u.variant_{{ variant.0 }})
     {% endfor %}
     case "":
         return cffi.CFFIValueUnionNONE, 0, fmt.Errorf("invalid union variant: [unset]")

engine/language_client_codegen/src/go/templates/types.go.j2

@@ -78,4 +78,23 @@ func (c {{ class.name }}) BamlTypeName() string {
 {#- Type Aliases -#}
 {% for alias in structural_recursive_alias_cycles %}
 type {{alias.name}} {{alias.target}}
+
+{% if alias.is_baml_serializable %}
+func (c {{alias.name}}) BamlTypeName() string {
+	return {{alias.target}}(c).BamlTypeName()
+}
+
+func (c {{alias.name}}) Encode(builder *flatbuffers.Builder) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
+	return {{alias.target}}(c).Encode(builder)
+}
+{% endif %}
+
+{% if alias.is_union %}
+func (u *{{alias.name}}) Decode(holder *cffi.CFFIValueUnionVariant) {
+	decodedUnion := {{alias.target}}{}
+	decodedUnion.Decode(holder)
+	*u = {{alias.name}}(decodedUnion)
+}
+{% endif %}
+
 {% endfor %}

engine/language_client_go/pkg/cffi/CFFIValueUnionVariant.go

@@ -63,21 +63,17 @@ func EncodeEnum(builder *flatbuffers.Builder, name string, value string, isDynam
 	return cffi.CFFIValueUnionCFFIValueEnum, cffi.CFFIValueEnumEnd(builder), nil
 }
 
-func EncodeUnion(builder *flatbuffers.Builder, variantName string, value any) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
-	nameOffset := builder.CreateString(variantName)
+func EncodeUnion(builder *flatbuffers.Builder, name string, variantName string, value any) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
+	nameOffset := builder.CreateString(name)
+	variantNameOffset := builder.CreateString(variantName)
 	valueHolderOffset, err := Encode(builder, value)
 	if err != nil {
 		return cffi.CFFIValueUnionNONE, 0, fmt.Errorf("encoding inner value for union variant '%s': %w", variantName, err)
 	}
 
-	// Note: The CFFIValueUnionVariant schema has field_types and value_type_index.
-	// The BamlUnionSerializer interface currently doesn't provide these.
-	// For now, we are omitting them, which might be incomplete depending on
-	// how the receiving end uses this structure. If they are needed, the
-	// interface and this function must be updated.
-
 	cffi.CFFIValueUnionVariantStart(builder)
 	cffi.CFFIValueUnionVariantAddName(builder, nameOffset)
+	cffi.CFFIValueUnionVariantAddVariantName(builder, variantNameOffset)
 	cffi.CFFIValueUnionVariantAddValue(builder, valueHolderOffset)
 	// cffi.CFFIValueUnionVariantAddFieldTypes(builder, ...) // Omitted
 	// cffi.CFFIValueUnionVariantAddValueTypeIndex(builder, ...) // Omitted
@@ -250,6 +246,7 @@ func encodeList(builder *flatbuffers.Builder, value reflect.Value) (flatbuffers.
 
 // encodeMap now accepts and passes TypeMap
 func encodeMap(builder *flatbuffers.Builder, mapValue reflect.Value) (flatbuffers.UOffsetT, error) {
+
 	mapLength := mapValue.Len()
 	entryOffsets := make([]flatbuffers.UOffsetT, 0, mapLength)
 	// Iterate map and build entries (order doesn't strictly matter for map, but FB requires building bottom-up)

engine/language_client_go/pkg/encode.go

@@ -0,0 +1,23 @@
+type RecursiveUnion = string | map<string, RecursiveUnion>
+
+function RecursiveUnionTest(input: RecursiveUnion) -> RecursiveUnion {
+  client GPT35
+  prompt #"
+    Return back the same value you were given: {{ input }}
+    
+    {{ ctx.output_format }}
+  "#
+}
+
+test RecursiveUnionTest {
+  functions [RecursiveUnionTest]
+  args {
+    input {
+      "key" "value"
+      "key2" {
+        "key" "value2"
+        "key2" "value3"
+      }
+    }
+  }
+}