Go Encode/Decode fixes (#1865)

This PR does a couple things.

1. Adds some cffi roundtripping tests for go
2. Uses those tests to make sure that classes encode/decode properly
3. Uses those tests to make sure that map[string]primitive and
[]primitive work if they are return values or input parameters
4. Remove passing the typeMap recursively since its statically available
inside the baml code (you will notice the type signature is the only
change)
5. Added `BamlTypeName()` to the `BamlSerializer` class which allows us
to get a baml type name from classes/unions/enums.
6. Use that to support classes as values in maps

Still needs a test with a Union for map/list

Author: tberman

engine/language_client_cffi/Makefile.toml

@@ -24,7 +24,7 @@ dependencies = ["cbindgen"]
 
 [tasks.build-local]
 script = '''
-python build.py --targets aarch64-apple-darwin --debug
+cargo build
 '''
 dependencies = ["cbindgen"]
 

engine/language_client_cffi/src/cffi/cffi_generated.rs

@@ -1493,7 +1493,8 @@ impl<'a> flatbuffers::Follow<'a> for CFFIValueList<'a> {
 }
 
 impl<'a> CFFIValueList<'a> {
-  pub const VT_VALUES: flatbuffers::VOffsetT = 4;
+  pub const VT_FIELD_TYPE: flatbuffers::VOffsetT = 4;
+  pub const VT_VALUES: flatbuffers::VOffsetT = 6;
 
   #[inline]
   pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
@@ -1506,10 +1507,18 @@ impl<'a> CFFIValueList<'a> {
   ) -> flatbuffers::WIPOffset<CFFIValueList<'bldr>> {
     let mut builder = CFFIValueListBuilder::new(_fbb);
     if let Some(x) = args.values { builder.add_values(x); }
+    if let Some(x) = args.field_type { builder.add_field_type(x); }
     builder.finish()
   }
 
 
+  #[inline]
+  pub fn field_type(&self) -> Option<CFFIFieldTypeHolder<'a>> {
+    // Safety:
+    // Created from valid Table for this object
+    // which contains a valid value in this slot
+    unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder>>(CFFIValueList::VT_FIELD_TYPE, None)}
+  }
   #[inline]
   pub fn values(&self) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIValueHolder<'a>>>> {
     // Safety:
@@ -1526,18 +1535,21 @@ impl flatbuffers::Verifiable for CFFIValueList<'_> {
   ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
     use self::flatbuffers::Verifiable;
     v.visit_table(pos)?
+     .visit_field::<flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder>>("field_type", Self::VT_FIELD_TYPE, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<CFFIValueHolder>>>>("values", Self::VT_VALUES, false)?
      .finish();
     Ok(())
   }
 }
 pub struct CFFIValueListArgs<'a> {
+    pub field_type: Option<flatbuffers::WIPOffset<CFFIFieldTypeHolder<'a>>>,
     pub values: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIValueHolder<'a>>>>>,
 }
 impl<'a> Default for CFFIValueListArgs<'a> {
   #[inline]
   fn default() -> Self {
     CFFIValueListArgs {
+      field_type: None,
       values: None,
     }
   }
@@ -1548,6 +1560,10 @@ pub struct CFFIValueListBuilder<'a: 'b, 'b> {
   start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
 }
 impl<'a: 'b, 'b> CFFIValueListBuilder<'a, 'b> {
+  #[inline]
+  pub fn add_field_type(&mut self, field_type: flatbuffers::WIPOffset<CFFIFieldTypeHolder<'b >>) {
+    self.fbb_.push_slot_always::<flatbuffers::WIPOffset<CFFIFieldTypeHolder>>(CFFIValueList::VT_FIELD_TYPE, field_type);
+  }
   #[inline]
   pub fn add_values(&mut self, values: flatbuffers::WIPOffset<flatbuffers::Vector<'b , flatbuffers::ForwardsUOffset<CFFIValueHolder<'b >>>>) {
     self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(CFFIValueList::VT_VALUES, values);
@@ -1570,6 +1586,7 @@ impl<'a: 'b, 'b> CFFIValueListBuilder<'a, 'b> {
 impl core::fmt::Debug for CFFIValueList<'_> {
   fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
     let mut ds = f.debug_struct("CFFIValueList");
+      ds.field("field_type", &self.field_type());
       ds.field("values", &self.values());
       ds.finish()
   }
@@ -1704,7 +1721,8 @@ impl<'a> flatbuffers::Follow<'a> for CFFIValueMap<'a> {
 }
 
 impl<'a> CFFIValueMap<'a> {
-  pub const VT_ENTRIES: flatbuffers::VOffsetT = 4;
+  pub const VT_FIELD_TYPES: flatbuffers::VOffsetT = 4;
+  pub const VT_ENTRIES: flatbuffers::VOffsetT = 6;
 
   #[inline]
   pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
@@ -1717,10 +1735,18 @@ impl<'a> CFFIValueMap<'a> {
   ) -> flatbuffers::WIPOffset<CFFIValueMap<'bldr>> {
     let mut builder = CFFIValueMapBuilder::new(_fbb);
     if let Some(x) = args.entries { builder.add_entries(x); }
+    if let Some(x) = args.field_types { builder.add_field_types(x); }
     builder.finish()
   }
 
 
+  #[inline]
+  pub fn field_types(&self) -> Option<CFFIFieldTypeHolder<'a>> {
+    // Safety:
+    // Created from valid Table for this object
+    // which contains a valid value in this slot
+    unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder>>(CFFIValueMap::VT_FIELD_TYPES, None)}
+  }
   #[inline]
   pub fn entries(&self) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIMapEntry<'a>>>> {
     // Safety:
@@ -1737,18 +1763,21 @@ impl flatbuffers::Verifiable for CFFIValueMap<'_> {
   ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
     use self::flatbuffers::Verifiable;
     v.visit_table(pos)?
+     .visit_field::<flatbuffers::ForwardsUOffset<CFFIFieldTypeHolder>>("field_types", Self::VT_FIELD_TYPES, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<CFFIMapEntry>>>>("entries", Self::VT_ENTRIES, false)?
      .finish();
     Ok(())
   }
 }
 pub struct CFFIValueMapArgs<'a> {
+    pub field_types: Option<flatbuffers::WIPOffset<CFFIFieldTypeHolder<'a>>>,
     pub entries: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIMapEntry<'a>>>>>,
 }
 impl<'a> Default for CFFIValueMapArgs<'a> {
   #[inline]
   fn default() -> Self {
     CFFIValueMapArgs {
+      field_types: None,
       entries: None,
     }
   }
@@ -1759,6 +1788,10 @@ pub struct CFFIValueMapBuilder<'a: 'b, 'b> {
   start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
 }
 impl<'a: 'b, 'b> CFFIValueMapBuilder<'a, 'b> {
+  #[inline]
+  pub fn add_field_types(&mut self, field_types: flatbuffers::WIPOffset<CFFIFieldTypeHolder<'b >>) {
+    self.fbb_.push_slot_always::<flatbuffers::WIPOffset<CFFIFieldTypeHolder>>(CFFIValueMap::VT_FIELD_TYPES, field_types);
+  }
   #[inline]
   pub fn add_entries(&mut self, entries: flatbuffers::WIPOffset<flatbuffers::Vector<'b , flatbuffers::ForwardsUOffset<CFFIMapEntry<'b >>>>) {
     self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(CFFIValueMap::VT_ENTRIES, entries);
@@ -1781,6 +1814,7 @@ impl<'a: 'b, 'b> CFFIValueMapBuilder<'a, 'b> {
 impl core::fmt::Debug for CFFIValueMap<'_> {
   fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
     let mut ds = f.debug_struct("CFFIValueMap");
+      ds.field("field_types", &self.field_types());
       ds.field("entries", &self.entries());
       ds.finish()
   }

engine/language_client_cffi/src/ctypes.rs

@@ -277,9 +277,12 @@ where
 
             let values = builder.create_vector_from_iter(items.into_iter());
 
+            let field_type = field_type_to_cffi_value_holder(value.field_type(), &mut builder);
+
             let value_list = CFFIValueList::create(
                 &mut builder,
                 &CFFIValueListArgs {
+                    field_type: Some(field_type),
                     values: Some(values),
                 },
             );
@@ -303,9 +306,12 @@ where
 
             let entries = builder.create_vector_from_iter(items.into_iter());
 
+            let field_types = field_type_to_cffi_value_holder(value.field_type(), &mut builder);
+
             let value_map = CFFIValueMap::create(
                 &mut builder,
                 &CFFIValueMapArgs {
+                    field_types: Some(field_types),
                     entries: Some(entries),
                 },
             );

engine/language_client_cffi/types/cffi.fbs

@@ -44,6 +44,7 @@ table CFFIValueBool {
 }
 
 table CFFIValueList {
+  field_type: CFFIFieldTypeHolder;
   values: [CFFIValueHolder];
 }
 
@@ -54,6 +55,7 @@ table CFFIMapEntry {
 }
 
 table CFFIValueMap {
+  field_types: CFFIFieldTypeHolder;
   entries: [CFFIMapEntry];
 }
 

engine/language_client_codegen/src/go/generate_types.rs

@@ -56,7 +56,7 @@ pub(crate) fn cast_value(container_variable_name: &str, field_type: &GoType) ->
 fn render_value_coercion(container_variable_name: &str, field_type: &GoType) -> String {
     if field_type.is_class {
         return format!(
-            "*baml.Decode({}, typeMap).(*{})",
+            "*baml.Decode({}).(*{})",
             container_variable_name,
             filters::type_name_without_pointer(&field_type.name)
                 .ok()
@@ -65,20 +65,20 @@ fn render_value_coercion(container_variable_name: &str, field_type: &GoType) ->
     } else if field_type.is_slice {
         let inner_type = field_type.underlying_type.as_ref().unwrap();
         return format!(
-            r#"baml.DecodeList({container_variable_name}, typeMap, func(__holder *cffi.CFFIValueHolder, typeMap baml.TypeMap) {} {{
+            r#"baml.DecodeList({container_variable_name}, func(__holder *cffi.CFFIValueHolder) {} {{
     return {}
 }})"#,
             inner_type.name,
             render_value_coercion("__holder", inner_type),
         );
     } else if field_type.is_union {
         return format!(
-            "*baml.Decode({container_variable_name}, typeMap).(*{})",
+            "*baml.Decode({container_variable_name}).(*{})",
             field_type.name
         );
     } else {
         return format!(
-            "baml.Decode({container_variable_name}, typeMap).({})",
+            "baml.Decode({container_variable_name}).({})",
             field_type.name
         );
     }

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

@@ -52,7 +52,7 @@ func {{ fn.name }}(ctx context.Context{% for (name, type) in fn.args -%}
         , {{name}} {{type}}
         {%- endfor %}) (*{{ fn.return_type }}, error) {
 	args := map[string]any{ {% for (name, type) in fn.args -%}"{{name}}": {{name}},{% endfor %} }
-	encoded, err := baml.EncodeRoot(args, typeMap)
+	encoded, err := baml.EncodeRoot(args)
 	if err != nil {
 		panic(err)
 	}
@@ -78,7 +78,7 @@ func (*stream) {{ fn.name }}(ctx context.Context{% for (name, type) in fn.args -
 		, {{name}} {{type}}
 		{%- endfor %}) <-chan {{ fn.return_type }} {
 	args := map[string]any{ {% for (name, type) in fn.args -%}"{{name}}": {{name}},{% endfor %} }
-	encoded, err := baml.EncodeRoot(args, typeMap)
+	encoded, err := baml.EncodeRoot(args)
 	if err != nil {
 		panic(err)
 	}

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

@@ -8,4 +8,4 @@ var file_map = map[string]string{
 
 func getBamlFiles() map[string]string {
   return file_map
-}
+}

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

@@ -74,15 +74,19 @@ func (e *{{ enum.name }}) UnmarshalJSON(data []byte) error {
     return nil
 }
 
-func (e *{{ enum.name }}) Decode(holder cffi.CFFIValueEnum, typeMap baml.TypeMap) {
+func (e *{{ enum.name }}) Decode(holder cffi.CFFIValueEnum) {
     if string(holder.Name()) != "{{ enum.name }}" {
         panic(fmt.Sprintf("expected {{ enum.name }}, got %s", string(holder.Name())))
     }
     value := string(holder.Value())
     *e = {{ enum.name }}(value)
 }
 
-func (e {{ enum.name }}) Encode(builder *flatbuffers.Builder, typeMap baml.TypeMap) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
+func (e {{ enum.name }}) Encode(builder *flatbuffers.Builder) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
     return baml.EncodeEnum(builder, "{{ enum.name }}", string(e), false)
 }
+
+func (e {{ enum.name }}) BamlTypeName() string {
+    return "{{ enum.name }}"
+}
 {% endfor %}

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

@@ -17,7 +17,7 @@ type {{ union.name }} struct {
     {% endfor %}
 }
 
-func (u *{{ union.name }}) Decode(holder *cffi.CFFIValueUnionVariant, typeMap baml.TypeMap) {
+func (u *{{ union.name }}) Decode(holder *cffi.CFFIValueUnionVariant) {
     valueHolder := holder.Value(nil)
 	switch holder.ValueTypeIndex() {
     {% for variant in union.variants %}
@@ -30,11 +30,11 @@ func (u *{{ union.name }}) Decode(holder *cffi.CFFIValueUnionVariant, typeMap ba
 }
 
 
-func (u {{ union.name }}) Encode(builder *flatbuffers.Builder, typeMap baml.TypeMap) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
+func (u {{ union.name }}) Encode(builder *flatbuffers.Builder) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
     switch u.variant {
     {% for variant in union.variants %}
     case "{{ variant.0 }}":
-        return baml.EncodeUnion(builder, typeMap, u.variant, u.variant_{{ variant.0 }})
+        return baml.EncodeUnion(builder, u.variant, u.variant_{{ variant.0 }})
     {% endfor %}
     case "":
         return cffi.CFFIValueUnionNONE, 0, fmt.Errorf("invalid union variant: [unset]")
@@ -43,6 +43,10 @@ func (u {{ union.name }}) Encode(builder *flatbuffers.Builder, typeMap baml.Type
     return cffi.CFFIValueUnionNONE, 0, fmt.Errorf("invalid union variant: %s", u.variant)
 }
 
+func (u {{ union.name }}) BamlTypeName() string {
+    return "{{ union.name }}"
+}
+
 {% for variant in union.variants %}
 func (u *{{ union.name }}) Set{{ variant.0|exported_name }}(v {{ variant.1.name }}) {
     u.variant = "{{ variant.0 }}"

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

@@ -27,7 +27,7 @@ type {{ class.name }} struct {
 {% endif %}
 }
 
-func (c *{{ class.name }}) Decode(holder cffi.CFFIValueClass, typeMap baml.TypeMap) {
+func (c *{{ class.name }}) Decode(holder cffi.CFFIValueClass) {
 	if string(holder.Name()) != "{{ class.name }}" {
 		panic(fmt.Sprintf("expected {{ class.name }}, got %s", string(holder.Name())))
 	}
@@ -55,18 +55,22 @@ func (c *{{ class.name }}) Decode(holder cffi.CFFIValueClass, typeMap baml.TypeM
 		if holder.DynamicFields(&field, i) {
 			key := string(field.Key())
 			valueHolder := field.Value(nil)
-			c.DynamicProperties[key] = baml.Decode(valueHolder, typeMap)
+			c.DynamicProperties[key] = baml.Decode(valueHolder)
 		}
 	}
 {% endif %}
 }
 
-func (c {{ class.name }}) Encode(builder *flatbuffers.Builder, typeMap baml.TypeMap) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
+func (c {{ class.name }}) Encode(builder *flatbuffers.Builder) (cffi.CFFIValueUnion, flatbuffers.UOffsetT, error) {
 	fields := map[string]any{}
 	{% for field in class.fields %}
 	fields["{{ field.name }}"] = c.{{ field.name|exported_name }}
 	{% endfor %}
-	return baml.EncodeClass(builder, typeMap, "{{ class.name }}", fields, {% if class.dynamic %}&c.DynamicProperties{% else %}nil{% endif %})
+	return baml.EncodeClass(builder, "{{ class.name }}", fields, {% if class.dynamic %}&c.DynamicProperties{% else %}nil{% endif %})
+}
+
+func (c {{ class.name }}) BamlTypeName() string {
+	return "{{ class.name }}"
 }
 
 {% endfor %}