add validation in the right place

docs/docs/snippets/supported-types.mdx

@@ -217,12 +217,12 @@ A mapping of strings to elements of another type.
 
 **Example**: `map<string, string>`
 
-<Info>
+{/* <Info>
   For TS users: `map<string, ValueType>` will generate a 
   `Record<string, ValueType>` type annotation, but using any other type for the
   key will generate a `Map`, e.g. `map<int, string>` in BAML will generate a
   `Map<number, string>` type annotation in TypeScript.
-</Info>
+</Info> */}
 
 ### ❌ Set
 

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

@@ -1,7 +1,5 @@
 use baml_types::{BamlMap, BamlMediaType, BamlValue, FieldType, TypeValue};
 use core::result::Result;
-use log::kv;
-use std::ops::Deref;
 
 use crate::ir::IntermediateRepr;
 
@@ -68,7 +66,7 @@ pub fn validate_arg(
                             .and_then(|v| v.as_str())
                             .unwrap_or_default()
                             .to_string();
-                        Some(BamlValue::Media(baml_types::BamlMedia::url(
+                        Ok(BamlValue::Media(baml_types::BamlMedia::url(
                             BamlMediaType::Image,
                             s.to_string(),
                             Some(media_type_str),
@@ -79,7 +77,7 @@ pub fn validate_arg(
                             .and_then(|v| v.as_str())
                             .unwrap_or_default()
                             .to_string();
-                        Some(BamlValue::Media(baml_types::BamlMedia::base64(
+                        Ok(BamlValue::Media(baml_types::BamlMedia::base64(
                             BamlMediaType::Image,
                             s.to_string(),
                             media_type_str,
@@ -106,7 +104,7 @@ pub fn validate_arg(
                             .and_then(|v| v.as_str())
                             .unwrap_or_default()
                             .to_string();
-                        Some(BamlValue::Media(baml_types::BamlMedia::url(
+                        Ok(BamlValue::Media(baml_types::BamlMedia::url(
                             BamlMediaType::Audio,
                             s.to_string(),
                             Some(media_type_str),
@@ -117,7 +115,7 @@ pub fn validate_arg(
                             .and_then(|v| v.as_str())
                             .unwrap_or_default()
                             .to_string();
-                        Some(BamlValue::Media(baml_types::BamlMedia::base64(
+                        Ok(BamlValue::Media(baml_types::BamlMedia::base64(
                             BamlMediaType::Audio,
                             s.to_string(),
                             media_type_str,

engine/baml-lib/baml-core/src/validate/validation_pipeline/validations.rs

@@ -1,10 +1,10 @@
 mod classes;
 mod clients;
-mod common;
 mod configurations;
 mod cycle;
 mod enums;
 mod functions;
+mod types;
 mod variants;
 
 use super::context::Context;

engine/baml-lib/baml-core/src/validate/validation_pipeline/validations/classes.rs

@@ -1,18 +1,18 @@
 use crate::validate::validation_pipeline::context::Context;
 
-use super::common::validate_type_exists;
+use super::types::validate_type;
 
 pub(super) fn validate(ctx: &mut Context<'_>) {
     for cls in ctx.db.walk_classes() {
         let _ast_class = cls.ast_class();
 
         for c in cls.static_fields() {
             let field = c.ast_field();
-            validate_type_exists(ctx, &field.field_type);
+            validate_type(ctx, &field.field_type);
         }
         for c in cls.dynamic_fields() {
             let field = c.ast_field();
-            validate_type_exists(ctx, &field.field_type);
+            validate_type(ctx, &field.field_type);
         }
     }
 }

engine/baml-lib/baml-core/src/validate/validation_pipeline/validations/functions.rs

@@ -3,13 +3,13 @@ use internal_baml_schema_ast::ast::{WithIdentifier, WithName, WithSpan};
 
 use crate::validate::validation_pipeline::context::Context;
 
-use super::common::validate_type_exists;
+use super::types::validate_type;
 
 pub(super) fn validate(ctx: &mut Context<'_>) {
     for func in ctx.db.walk_old_functions() {
         for args in func.walk_input_args().chain(func.walk_output_args()) {
             let arg = args.ast_arg();
-            validate_type_exists(ctx, &arg.1.field_type)
+            validate_type(ctx, &arg.1.field_type)
         }
 
         // Check if the function has multiple impls, if it does,
@@ -114,7 +114,7 @@ pub(super) fn validate(ctx: &mut Context<'_>) {
     for func in ctx.db.walk_new_functions() {
         for args in func.walk_input_args().chain(func.walk_output_args()) {
             let arg = args.ast_arg();
-            validate_type_exists(ctx, &arg.1.field_type)
+            validate_type(ctx, &arg.1.field_type)
         }
 
         // Ensure the client is correct.

engine/baml-lib/baml-core/src/validate/validation_pipeline/validations/types.rs

@@ -0,0 +1,66 @@
+use baml_types::TypeValue;
+use internal_baml_diagnostics::DatamodelError;
+use internal_baml_schema_ast::ast::{FieldArity, FieldType, Identifier, WithName, WithSpan};
+
+use crate::validate::validation_pipeline::context::Context;
+
+fn errors_with_names<'a>(ctx: &'a mut Context<'_>, idn: &Identifier) {
+    // Push the error with the appropriate message
+    ctx.push_error(DatamodelError::new_type_not_found_error(
+        idn.name(),
+        ctx.db.valid_type_names(),
+        idn.span().clone(),
+    ));
+}
+
+/// Called for each type in the baml_src tree, validates that it is well-formed.
+///
+/// Operates in two passes:
+///
+///   1. Verify that the type is resolveable (for REF types)
+///   2. Verify that the type is well-formed/allowed in the language
+pub(crate) fn validate_type(ctx: &mut Context<'_>, field_type: &FieldType) {
+    validate_type_exists(ctx, field_type);
+    validate_type_allowed(ctx, field_type);
+}
+
+fn validate_type_exists(ctx: &mut Context<'_>, field_type: &FieldType) {
+    field_type
+        .flat_idns()
+        .iter()
+        .for_each(|f| match ctx.db.find_type(f) {
+            Some(_) => {}
+            None => match f {
+                Identifier::Primitive(..) => {}
+                _ => errors_with_names(ctx, f),
+            },
+        });
+}
+
+fn validate_type_allowed(ctx: &mut Context<'_>, field_type: &FieldType) {
+    match field_type {
+        FieldType::Map(kv_types, _) => {
+            match &kv_types.0 {
+                FieldType::Identifier(
+                    FieldArity::Required,
+                    Identifier::Primitive(TypeValue::String, _),
+                ) => {}
+                key_type => {
+                    ctx.push_error(DatamodelError::new_validation_error(
+                        "Maps may only have strings as keys",
+                        key_type.span().clone(),
+                    ));
+                }
+            }
+            validate_type_allowed(ctx, &kv_types.1);
+            // TODO:assert key_type is string or int or null
+        }
+        FieldType::Identifier(_, _) => {}
+        FieldType::List(field_type, _, _) => validate_type_allowed(ctx, field_type),
+        FieldType::Tuple(_, field_types, _) | FieldType::Union(_, field_types, _) => {
+            for field_type in field_types {
+                validate_type_allowed(ctx, field_type);
+            }
+        }
+    }
+}

engine/baml-lib/baml/tests/validation_files/class/map_types.baml

@@ -7,52 +7,58 @@ class MapDummy {
 class MapFields {
   a1 map<string, string>
   a2 map<string, int>
-  a2 map<string, MapDummy>
+  a3 map<string, MapDummy>
 
   b1 map<int, string>
   b2 map<float, string>
   b3 map<MapDummy, string>
+  b4 map<string?, string>
+  b5 map<string | int, string>
+
+  c1 string | map<string, string>
+  c2 string | map<int, string>
+  c3 string | map<string?, string>
 }
 
-// error: Error validating: This line is not a valid field or attribute definition. A valid class property looks like: 'myProperty string[] @description("This is a description")'
-//   -->  class/map_types.baml:8
-//    | 
-//  7 | class MapFields {
-//  8 |   a1 map<string, string>
-//  9 |   a2 map<string, int>
-//    | 
-// error: Error validating: This line is not a valid field or attribute definition. A valid class property looks like: 'myProperty string[] @description("This is a description")'
-//   -->  class/map_types.baml:9
-//    | 
-//  8 |   a1 map<string, string>
-//  9 |   a2 map<string, int>
-// 10 |   a2 map<string, MapDummy>
-//    | 
-// error: Error validating: This line is not a valid field or attribute definition. A valid class property looks like: 'myProperty string[] @description("This is a description")'
-//   -->  class/map_types.baml:10
-//    | 
-//  9 |   a2 map<string, int>
-// 10 |   a2 map<string, MapDummy>
-// 11 | 
-//    | 
-// error: Error validating: This line is not a valid field or attribute definition. A valid class property looks like: 'myProperty string[] @description("This is a description")'
+// error: Error validating: Maps may only have strings as keys
 //   -->  class/map_types.baml:12
 //    | 
 // 11 | 
 // 12 |   b1 map<int, string>
-// 13 |   b2 map<float, string>
 //    | 
-// error: Error validating: This line is not a valid field or attribute definition. A valid class property looks like: 'myProperty string[] @description("This is a description")'
+// error: Error validating: Maps may only have strings as keys
 //   -->  class/map_types.baml:13
 //    | 
 // 12 |   b1 map<int, string>
 // 13 |   b2 map<float, string>
-// 14 |   b3 map<MapDummy, string>
 //    | 
-// error: Error validating: This line is not a valid field or attribute definition. A valid class property looks like: 'myProperty string[] @description("This is a description")'
+// error: Error validating: Maps may only have strings as keys
 //   -->  class/map_types.baml:14
 //    | 
 // 13 |   b2 map<float, string>
 // 14 |   b3 map<MapDummy, string>
-// 15 | }
+//    | 
+// error: Error validating: Maps may only have strings as keys
+//   -->  class/map_types.baml:15
+//    | 
+// 14 |   b3 map<MapDummy, string>
+// 15 |   b4 map<string?, string>
+//    | 
+// error: Error validating: Maps may only have strings as keys
+//   -->  class/map_types.baml:16
+//    | 
+// 15 |   b4 map<string?, string>
+// 16 |   b5 map<string | int, string>
+//    | 
+// error: Error validating: Maps may only have strings as keys
+//   -->  class/map_types.baml:19
+//    | 
+// 18 |   c1 string | map<string, string>
+// 19 |   c2 string | map<int, string>
+//    | 
+// error: Error validating: Maps may only have strings as keys
+//   -->  class/map_types.baml:20
+//    | 
+// 19 |   c2 string | map<int, string>
+// 20 |   c3 string | map<string?, string>
 //    |