Merge pull request #25 from Virtual-Repetitions/better-join-tables

adding enum support (0.29.2)

Author: scoop1120

Cargo.lock

@@ -1,5 +1,5 @@
 [workspace.package]
-version = "0.29.0"
+version = "0.29.2"
 edition = "2024"
 
 # See https://github.com/mozilla/application-services/blob/main/Cargo.toml for the reasons why we use this structure

Cargo.toml

@@ -81,6 +81,9 @@ fn matches(expr: &AstExpr<Typed>, operation_name: &str, operation_kind: Operatio
             }
         },
         AstExpr::RelationalOp(_) => panic!("RelationalOp not supported in interceptor expression"),
+        AstExpr::EnumLiteral(_, _, _, _) => {
+            panic!("EnumLiteral not supported in interceptor expression")
+        }
         AstExpr::StringLiteral(value, _) => matches_str(value, operation_name, operation_kind),
         AstExpr::BooleanLiteral(value, _) => *value,
         AstExpr::NumberLiteral(_, _) => {

crates/builder/src/builder/interceptor_weaver.rs

@@ -9,11 +9,13 @@
 
 use std::collections::HashMap;
 
-use codemap_diagnostic::Diagnostic;
+use codemap_diagnostic::{Diagnostic, Level, SpanLabel, SpanStyle};
 use core_model::mapped_arena::MappedArena;
 use core_model_builder::typechecker::{Typed, annotation::AnnotationSpec};
 
-use crate::ast::ast_types::{AstExpr, FieldSelection, LogicalOp, RelationalOp, Untyped};
+use crate::ast::ast_types::{
+    AstExpr, FieldSelection, FieldSelectionElement, LogicalOp, RelationalOp, Untyped,
+};
 
 use super::{Scope, Type, TypecheckFrom};
 
@@ -27,6 +29,9 @@ impl TypecheckFrom<AstExpr<Untyped>> for AstExpr<Typed> {
             AstExpr::RelationalOp(relation) => {
                 AstExpr::RelationalOp(RelationalOp::shallow(relation))
             }
+            AstExpr::EnumLiteral(enum_name, value, span, _) => {
+                AstExpr::EnumLiteral(enum_name.clone(), value.clone(), *span, Type::Defer)
+            }
             AstExpr::StringLiteral(v, s) => AstExpr::StringLiteral(v.clone(), *s),
             AstExpr::BooleanLiteral(v, s) => AstExpr::BooleanLiteral(*v, *s),
             AstExpr::NumberLiteral(v, s) => AstExpr::NumberLiteral(v.clone(), *s),
@@ -49,11 +54,39 @@ impl TypecheckFrom<AstExpr<Untyped>> for AstExpr<Typed> {
         errors: &mut Vec<Diagnostic>,
     ) -> bool {
         match self {
-            AstExpr::FieldSelection(select) => select.pass(type_env, annotation_env, scope, errors),
+            AstExpr::FieldSelection(select) => {
+                if let Some((enum_name, value, span)) = enum_literal_parts(select) {
+                    if let Some(typ) =
+                        resolve_enum_literal(&enum_name, &value, span, type_env, scope, errors)
+                    {
+                        *self = AstExpr::EnumLiteral(enum_name, value, span, typ);
+                        true
+                    } else {
+                        select.pass(type_env, annotation_env, scope, errors)
+                    }
+                } else {
+                    select.pass(type_env, annotation_env, scope, errors)
+                }
+            }
             AstExpr::LogicalOp(logic) => logic.pass(type_env, annotation_env, scope, errors),
             AstExpr::RelationalOp(relation) => {
                 relation.pass(type_env, annotation_env, scope, errors)
             }
+            AstExpr::EnumLiteral(enum_name, value, span, typ) => {
+                if typ.is_incomplete() {
+                    if let Some(resolved) =
+                        resolve_enum_literal(enum_name, value, *span, type_env, scope, errors)
+                    {
+                        *typ = resolved;
+                        true
+                    } else {
+                        *typ = Type::Error;
+                        false
+                    }
+                } else {
+                    false
+                }
+            }
             AstExpr::StringList(_, _)
             | AstExpr::StringLiteral(_, _)
             | AstExpr::BooleanLiteral(_, _)
@@ -63,3 +96,76 @@ impl TypecheckFrom<AstExpr<Untyped>> for AstExpr<Typed> {
         }
     }
 }
+
+fn enum_literal_parts(
+    selection: &FieldSelection<Typed>,
+) -> Option<(String, String, codemap::Span)> {
+    match selection {
+        FieldSelection::Select(prefix, elem, span, _) => match (prefix.as_ref(), elem) {
+            (
+                FieldSelection::Single(prefix_elem, _),
+                FieldSelectionElement::Identifier(value, _, _),
+            ) => {
+                if let FieldSelectionElement::Identifier(enum_name, _, _) = prefix_elem {
+                    Some((enum_name.clone(), value.clone(), *span))
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+fn resolve_enum_literal(
+    enum_name: &str,
+    value: &str,
+    span: codemap::Span,
+    type_env: &MappedArena<Type>,
+    scope: &Scope,
+    errors: &mut Vec<Diagnostic>,
+) -> Option<Type> {
+    if scope.get_type(enum_name).is_some() {
+        return None;
+    }
+
+    let is_context = type_env
+        .get_by_key(enum_name)
+        .and_then(|t| match t {
+            Type::Composite(c) if c.kind == crate::ast::ast_types::AstModelKind::Context => Some(c),
+            _ => None,
+        })
+        .is_some();
+
+    if is_context {
+        return None;
+    }
+
+    let enum_type = type_env.get_by_key(enum_name).and_then(|t| match t {
+        Type::Enum(e) => Some(e.clone()),
+        _ => None,
+    });
+
+    #[allow(clippy::question_mark)]
+    let enum_type = match enum_type {
+        Some(enum_type) => enum_type,
+        None => return None,
+    };
+
+    if enum_type.fields.iter().any(|f| f.name == value) {
+        Some(Type::Enum(enum_type))
+    } else {
+        errors.push(Diagnostic {
+            level: Level::Error,
+            message: format!("Unknown variant '{value}' for enum '{enum_name}'"),
+            code: Some("C000".to_string()),
+            spans: vec![SpanLabel {
+                span,
+                style: SpanStyle::Primary,
+                label: Some("unknown enum variant".to_string()),
+            }],
+        });
+        Some(Type::Error)
+    }
+}

crates/builder/src/typechecker/expression.rs

@@ -84,7 +84,7 @@ impl TypecheckFrom<AstField<Untyped>> for AstField<Typed> {
                 }
             };
 
-            match *expr {
+            match expr {
                 AstExpr::StringLiteral(_, _) => assert_type(&[
                     "String",
                     "Decimal",
@@ -97,6 +97,22 @@ impl TypecheckFrom<AstField<Untyped>> for AstField<Typed> {
                 ]),
                 AstExpr::BooleanLiteral(_, _) => assert_type(&["Boolean"]),
                 AstExpr::NumberLiteral(_, _) => assert_type(&["Int", "Float"]),
+                AstExpr::EnumLiteral(enum_name, _, _, _) => {
+                    if enum_name != &type_name {
+                        errors.push(Diagnostic {
+                            level: Level::Error,
+                            message:
+                                "Literal specified for default value is not a valid type for field."
+                                    .to_string(),
+                            code: Some("C000".to_string()),
+                            spans: vec![SpanLabel {
+                                span: expr.span(),
+                                style: SpanStyle::Primary,
+                                label: Some(format!("should be of type {type_name}")),
+                            }],
+                        });
+                    }
+                }
                 AstExpr::FieldSelection(_) => {
                     // no type-checking here, since we don't have enough information.
                     // For example `user: User = AuthContext.id` should check that `AuthContext.id`

crates/builder/src/typechecker/field.rs

@@ -49,6 +49,7 @@ impl TypecheckFrom<AstFieldDefault<Untyped>> for AstFieldDefault<Typed> {
             AstExpr::BooleanLiteral(_, _)
             | AstExpr::StringLiteral(_, _)
             | AstExpr::NumberLiteral(_, _)
+            | AstExpr::EnumLiteral(_, _, _, _)
             | AstExpr::FieldSelection(_) => expr.pass(type_env, annotation_env, scope, errors),
 
             _ => {

crates/builder/src/typechecker/field_default_value.rs

@@ -494,6 +494,7 @@ pub mod test_support {
 #[cfg(test)]
 mod tests {
     use super::test_support::{build, parse_sorted};
+    use core_model_builder::ast::ast_types::{AstAnnotationParams, AstExpr};
     use multiplatform_test::multiplatform_test;
 
     // Due to a change in insta version 1.12, test names (hence the snapshot names) get derived
@@ -561,6 +562,57 @@ mod tests {
         assert_typechecking!(src, "with_auth_context_use_in_field_annotation");
     }
 
+    #[multiplatform_test]
+    fn enum_literal_in_access_predicate() {
+        let src = r#"
+        @postgres
+        module TagModule {
+            enum TagSharingType {
+                PublicPlaybook
+                PrivatePlaybook
+            }
+
+            type Tag {
+                sharing: TagSharingType
+                @access(query=self.sharing == TagSharingType.PublicPlaybook) title: String
+            }
+        }
+        "#;
+
+        let built = build(src).unwrap();
+        let module = built
+            .modules
+            .iter()
+            .find(|(_, module)| module.0.name == "TagModule")
+            .map(|(_, module)| module)
+            .unwrap();
+        let tag_type = module.0.types.iter().find(|t| t.name == "Tag").unwrap();
+        let title_field = tag_type.fields.iter().find(|f| f.name == "title").unwrap();
+        let access = title_field.annotations.annotations.get("access").unwrap();
+
+        let expr = match &access.params {
+            AstAnnotationParams::Map(params, _) => params
+                .iter()
+                .find(|(k, _)| k.as_str() == "query")
+                .map(|(_, v)| v)
+                .unwrap(),
+            _ => panic!("Expected map access expression"),
+        };
+
+        let AstExpr::RelationalOp(rel) = expr else {
+            panic!("Expected relational op in access expression");
+        };
+
+        let (_, right) = rel.sides();
+        match right {
+            AstExpr::EnumLiteral(enum_name, value, _, _) => {
+                assert_eq!(enum_name, "TagSharingType");
+                assert_eq!(value, "PublicPlaybook");
+            }
+            _ => panic!("Expected enum literal on right-hand side"),
+        }
+    }
+
     #[multiplatform_test]
     fn with_array_in_operator() {
         let src = r#"

crates/builder/src/typechecker/mod.rs

@@ -316,6 +316,15 @@ pub enum AstExpr<T: NodeTypedness> {
     FieldSelection(FieldSelection<T>),
     LogicalOp(LogicalOp<T>),
     RelationalOp(RelationalOp<T>),
+    EnumLiteral(
+        String,
+        String,
+        #[serde(skip_serializing)]
+        #[serde(skip_deserializing)]
+        #[serde(default = "default_span")]
+        Span,
+        T::Expr,
+    ),
     StringLiteral(
         String,
         #[serde(skip_serializing)]
@@ -362,6 +371,7 @@ impl<T: NodeTypedness> AstExpr<T> {
     pub fn span(&self) -> Span {
         match &self {
             AstExpr::FieldSelection(s) => *s.span(),
+            AstExpr::EnumLiteral(_, _, s, _) => *s,
             AstExpr::StringLiteral(_, s) => *s,
             AstExpr::LogicalOp(l) => match l {
                 LogicalOp::Not(_, s, _) => *s,

crates/core-subsystem/core-model-builder/src/ast/ast_types.rs

@@ -19,6 +19,7 @@ impl AstExpr<Typed> {
             AstExpr::FieldSelection(select) => select.typ().clone(),
             AstExpr::LogicalOp(logic) => logic.typ().clone(),
             AstExpr::RelationalOp(relation) => relation.typ().clone(),
+            AstExpr::EnumLiteral(_, _, _, typ) => typ.clone(),
             AstExpr::StringLiteral(_, _) => {
                 Type::Primitive(PrimitiveType::Plain(primitive_type::STRING_TYPE))
             }

crates/core-subsystem/core-model-builder/src/typechecker/expression.rs

@@ -135,6 +135,7 @@ pub enum CommonAccessPrimitiveExpression {
     StringLiteral(String),              // for example, "ADMIN"
     BooleanLiteral(bool),               // for example, true
     NumberLiteral(String),              // for example, integer (-13, 0, 300, 10.5, etc.)
+    EnumLiteral { enum_name: String, value: String },
     NullLiteral,
 }