7ac1d37 | plum

git clone https://git.pyrossh.dev/plum

A statically typed, imperative programming language inspired by rust, python

7ac1d37 — Peter John 2026-07-20T11:33:47+05:30

feat(plum-checker): general enum variant tags, field types, and construction checks

Files changed (2) hide show

plum-checker/src/lib.rs +66 -13
plum-checker/tests/checker_tests.rs +101 -0

plum-checker/src/lib.rs CHANGED Viewed

@@ -44,8 +44,17 @@ pub fn lookup(env: &TypeEnv, name: &str) -> Result<PlumType, String> {
 pub type ClassEnv = BTreeMap<String, Vec<(String, PlumType)>>;
 /// `(receiver type, method name) -> TFun` for every `name<Receiver>(...)` method.
 pub type MethodEnv = BTreeMap<(String, String), PlumType>;
+/// Info about one `enum` variant: which enum it belongs to, its 0-based runtime tag
+/// (numbering is shared across all of that enum's variants), and its payload field
+/// types (empty for a payload-free variant like `Red` or `None`).
+#[derive(Debug, Clone, PartialEq)]
+pub struct EnumVariantInfo {
+    pub enum_name: String,
+    pub tag: i32,
+    pub field_types: Vec<PlumType>,
+}
-/// Enum variant name -> owning enum name, e.g. `"True" -> "Bool"`.
+/// Enum variant name -> its info, e.g. `"True" -> { enum_name: "Bool", tag: 1, field_types: [] }`.
-pub type EnumVariants = BTreeMap<String, String>;
+pub type EnumVariants = BTreeMap<String, EnumVariantInfo>;
 /// Shared, read-only lookup tables built once from the whole source, threaded through
 /// every check/infer call alongside the (mutable, scope-local) `TypeEnv`.
@@ -66,8 +75,8 @@ pub fn build_global_tables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEn
     let mut enum_variants: EnumVariants = BTreeMap::new();
     // `Bool`'s variants are built in (see `infer_expr`'s TypeName handling) rather
     // than requiring every source file to redeclare `enum Bool = | True | False`.
-    enum_variants.insert("True".to_string(), "Bool".to_string());
+    enum_variants.insert("True".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 1, field_types: vec![] });
-    enum_variants.insert("False".to_string(), "Bool".to_string());
+    enum_variants.insert("False".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 0, field_types: vec![] });
     // First pass: register class fields and enum variants so later passes can
     // resolve `self.field`, `ClassName(...)`, and bare enum-tag patterns.
@@ -80,8 +89,15 @@ pub fn build_global_tables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEn
                 classes.insert(c.name.clone(), fields);
             }
             ast::Item::Enum(e) => {
-                for v in &e.variants {
+                for (tag, v) in e.variants.iter().enumerate() {
+                    let field_types = v.fields.iter()
+                        .map(|f| plum_type_from_ast(&ast::Type { name: f.clone(), generics: vec![] }))
+                        .collect();
-                    enum_variants.insert(v.name.clone(), e.name.clone());
+                    enum_variants.insert(v.name.clone(), EnumVariantInfo {
+                        enum_name: e.name.clone(),
+                        tag: tag as i32,
+                        field_types,
+                    });
                 }
             }
             _ => {}
@@ -332,7 +348,7 @@ fn check_match(m: &ast::Match, env: &TypeEnv, declared_ret: &PlumType, fn_name:
 /// Checks a single case pattern against the type of the subject it matches, binding any
 /// new names it introduces into `env`. Constructor-payload sub-patterns (`Some(x)`) bind
-/// against an unconstrained type since enum variants don't carry per-field type info (v1.5).
+/// against that variant's declared field types (see `EnumVariantInfo::field_types`).
 fn check_pattern(pat: &ast::CasePattern, subject_ty: &PlumType, env: &mut TypeEnv, ctx: &CheckCtx) -> Result<(), String> {
     match pat {
         ast::CasePattern::Wildcard => Ok(()),
@@ -349,12 +365,27 @@ fn check_pattern(pat: &ast::CasePattern, subject_ty: &PlumType, env: &mut TypeEn
                 Ok(())
             }
         }
-        ast::CasePattern::Class { name: _, fields } => {
+        ast::CasePattern::Class { name, fields } => match ctx.enum_variants.get(name) {
-            for f in fields {
+            Some(info) => {
+                if fields.len() != info.field_types.len() {
+                    return Err(format!(
+                        "constructor pattern '{}' expects {} field(s), got {}",
+                        name, info.field_types.len(), fields.len()
+                    ));
+                }
+                for (f, fty) in fields.iter().zip(info.field_types.iter()) {
-                check_pattern(f, &PlumType::TVar("_".to_string()), env, ctx)?;
+                    check_pattern(f, fty, env, ctx)?;
+                }
+                Ok(())
             }
+            // Unmodeled/builtin variant: allow, codegen will catch.
+            None => {
+                for f in fields {
+                    check_pattern(f, &PlumType::TVar("_".to_string()), env, ctx)?;
+                }
-            Ok(())
+                Ok(())
-        }
+            }
+        },
     }
 }
@@ -367,7 +398,11 @@ pub fn infer_expr(expr: &ast::Expr, env: &TypeEnv, ctx: &CheckCtx) -> Result<Plu
         ast::Expr::Self_ => lookup(env, "self"),
         ast::Expr::TypeName(n) => match n.as_str() {
             "True" | "False" => Ok(PlumType::TBool),
+            _ => match ctx.enum_variants.get(n) {
+                Some(info) => Ok(PlumType::TNamed(info.enum_name.clone())),
+                // Unmodeled/builtin type name: allow, codegen will catch.
-            other => Ok(PlumType::TNamed(other.to_string())),
+                None => Ok(PlumType::TNamed(n.to_string())),
+            },
         },
         ast::Expr::Paren(inner) => infer_expr(inner, env, ctx),
         ast::Expr::Not(inner) => {
@@ -407,6 +442,24 @@ pub fn infer_expr(expr: &ast::Expr, env: &TypeEnv, ctx: &CheckCtx) -> Result<Plu
             Ok(tt)
         }
         ast::Expr::FnCall(call) => {
+            if let Some(info) = ctx.enum_variants.get(&call.name) {
+                if call.args.len() != info.field_types.len() {
+                    return Err(format!(
+                        "variant '{}': expected {} arg(s), got {}",
+                        call.name, info.field_types.len(), call.args.len()
+                    ));
+                }
+                for (i, (arg, expected)) in call.args.iter().zip(info.field_types.iter()).enumerate() {
+                    let arg_expr = match arg {
+                        ast::Arg::Positional(e) => e,
+                        ast::Arg::Keyword { value, .. } => value,
+                        ast::Arg::Pair { value, .. } => value,
+                    };
+                    let actual = infer_expr(arg_expr, env, ctx)?;
+                    unify(expected, &actual).map_err(|e| format!("variant '{}' arg {}: {}", call.name, i, e))?;
+                }
+                return Ok(PlumType::TNamed(info.enum_name.clone()));
+            }
             match lookup(env, &call.name) {
                 Ok(PlumType::TFun(param_types, ret)) => {
                     if call.args.len() != param_types.len() {

plum-checker/tests/checker_tests.rs CHANGED Viewed

@@ -191,3 +191,104 @@ fn match_int_pattern_against_str_subject_is_error() {
     let result = check_source(&source);
     assert!(result.is_err());
 }
+#[test]
+fn bare_enum_tag_unifies_with_owning_enum_type() {
+    // Regression: a bare non-Bool tag like `None` used to type as `TNamed("None")`
+    // (itself, not its enum), so comparing it against an `Option` value would wrongly
+    // fail with a type mismatch.
+    let src = "\
+enum Option =
+  | Some(Int)
+  | None
+isNone(o: Option) -> Bool =
+  o == None
+";
+    let source = parse(src);
+    let result = check_source(&source);
+    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
+}
+#[test]
+fn variant_construction_checks_arg_count_and_types() {
+    let src = "\
+enum Option =
+  | Some(Int)
+  | None
+makeSome(v: Int) -> Option =
+  Some(v)
+";
+    let source = parse(src);
+    let result = check_source(&source);
+    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
+}
+#[test]
+fn variant_construction_wrong_arg_type_is_error() {
+    let src = "\
+enum Option =
+  | Some(Int)
+  | None
+bad() -> Option =
+  Some(\"x\")
+";
+    let source = parse(src);
+    let result = check_source(&source);
+    assert!(result.is_err());
+}
+#[test]
+fn variant_construction_wrong_arg_count_is_error() {
+    let src = "\
+enum Shape =
+  | Rect(Float, Float)
+  | Circle(Float)
+bad() -> Shape =
+  Rect(1.0)
+";
+    let source = parse(src);
+    let result = check_source(&source);
+    assert!(result.is_err());
+}
+#[test]
+fn constructor_pattern_binds_fields_to_declared_types() {
+    let src = "\
+enum Shape =
+  | Rect(Float, Float)
+  | Circle(Float)
+area(s: Shape) -> Float =
+  match s
+    Rect(w, h) =>
+      w * h
+    Circle(r) =>
+      r * r
+";
+    let source = parse(src);
+    let result = check_source(&source);
+    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
+}
+#[test]
+fn constructor_pattern_wrong_field_count_is_error() {
+    let src = "\
+enum Shape =
+  | Rect(Float, Float)
+  | Circle(Float)
+bad(s: Shape) -> Float =
+  match s
+    Rect(w) =>
+      w
+    _ =>
+      0.0
+";
+    let source = parse(src);
+    let result = check_source(&source);
+    assert!(result.is_err());
+}