Fix subtype bug with aliases and add tests

engine/baml-lib/baml-core/src/ir/ir_helpers/mod.rs

@@ -290,17 +290,17 @@ impl IRHelper for IntermediateRepr {
 
                         if !map_type.is_subtype_of(&field_type) {
                             anyhow::bail!("Could not unify {:?} with {:?}", map_type, field_type);
-                        } else {
-                            let mapped_fields: BamlMap<String, BamlValueWithMeta<FieldType>> =
+                        }
+
+                        let mapped_fields: BamlMap<String, BamlValueWithMeta<FieldType>> =
                                     pairs
                                     .into_iter()
                                     .map(|(key, val)| {
                                         let sub_value = self.distribute_type(val, item_type.clone())?;
                                         Ok((key, sub_value))
                                     })
                                     .collect::<anyhow::Result<BamlMap<String,BamlValueWithMeta<FieldType>>>>()?;
-                            Ok(BamlValueWithMeta::Map(mapped_fields, field_type))
-                        }
+                        Ok(BamlValueWithMeta::Map(mapped_fields, field_type))
                     }
                     None => Ok(BamlValueWithMeta::Map(BamlMap::new(), field_type)),
                 }

engine/baml-lib/baml-types/src/field_type/mod.rs

@@ -184,6 +184,8 @@ impl FieldType {
         }
 
         match (self, other) {
+            (FieldType::Alias { resolution, .. }, _) => resolution.is_subtype_of(other),
+            (_, FieldType::Alias { resolution, .. }) => self.is_subtype_of(resolution),
             (FieldType::Primitive(TypeValue::Null), FieldType::Optional(_)) => true,
             (FieldType::Optional(self_item), FieldType::Optional(other_item)) => {
                 self_item.is_subtype_of(other_item)
@@ -245,7 +247,6 @@ impl FieldType {
                         .zip(other_items)
                         .all(|(self_item, other_item)| self_item.is_subtype_of(other_item))
             }
-            (FieldType::Alias { resolution, .. }, _) => resolution.is_subtype_of(other),
             (FieldType::Tuple(_), _) => false,
             (FieldType::Primitive(_), _) => false,
             (FieldType::Enum(_), _) => false,

integ-tests/python/tests/test_functions.py

@@ -249,6 +249,24 @@ async def test_single_literal_string_key_in_map(self):
         res = await b.InOutSingleLiteralStringMapKey({"key": "1"})
         assert res["key"] == "1"
 
+    @pytest.mark.asyncio
+    async def test_primitive_union_alias(self):
+        res = await b.PrimitiveAlias("test")
+        assert res == "test"
+
+    @pytest.mark.asyncio
+    async def test_map_alias(self):
+        res = await b.MapAlias({"A": ["B", "C"], "B": [], "C": []})
+        assert res == {"A": ["B", "C"], "B": [], "C": []}
+
+    @pytest.mark.asyncio
+    async def test_alias_union(self):
+        res = await b.NestedAlias("test")
+        assert res == "test"
+
+        res = await b.NestedAlias({"A": ["B", "C"], "B": [], "C": []})
+        assert res == {"A": ["B", "C"], "B": [], "C": []}
+
 
 class MyCustomClass(NamedArgsSingleClass):
     date: datetime.datetime

integ-tests/ruby/test_functions.rb

@@ -78,6 +78,18 @@
 
     res = b.InOutSingleLiteralStringMapKey(m: {"key" => "1"})
     assert_equal res['key'], "1"
+
+    res = b.PrimitiveAlias(p: "test")
+    assert_equal res, "test"
+
+    res = b.MapAlias(m: {"A" => ["B", "C"], "B" => [], "C" => []})
+    assert_equal res, {"A" => ["B", "C"], "B" => [], "C" => []}
+
+    res = b.NestedAlias(c: "test")
+    assert_equal res, "test"
+
+    res = b.NestedAlias(c: {"A" => ["B", "C"], "B" => [], "C" => []})
+    assert_equal res, {"A" => ["B", "C"], "B" => [], "C" => []}
   end
 
   it "accepts subclass of baml type" do

integ-tests/typescript/tests/integ-tests.test.ts

@@ -148,6 +148,24 @@ describe('Integ tests', () => {
       const res = await b.InOutSingleLiteralStringMapKey({ key: '1' })
       expect(res).toHaveProperty('key', '1')
     })
+
+    it('primitive union alias', async () => {
+      const res = await b.PrimitiveAlias('test')
+      expect(res).toEqual('test')
+    })
+
+    it('map alias', async () => {
+      const res = await b.MapAlias({ A: ['B', 'C'], B: [], C: [] })
+      expect(res).toEqual({ A: ['B', 'C'], B: [], C: [] })
+    })
+
+    it('alias union', async () => {
+      let res = await b.NestedAlias('test')
+      expect(res).toEqual('test')
+
+      res = await b.NestedAlias({ A: ['B', 'C'], B: [], C: [] })
+      expect(res).toEqual({ A: ['B', 'C'], B: [], C: [] })
+    })
   })
 
   it('should work for all outputs', async () => {