plum

#treesitter#compiler#wasm

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

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


4fda634Peter John 2026-08-09T18:54:49+05:30
feat(plum-checker): validate enum discriminant values, enable field access on them
plum-checker/src/lib.rs CHANGED
@@ -53,31 +53,40 @@ pub struct EnumVariantInfo {
53
53
  pub enum_name: String,
54
54
  pub tag: i32,
55
55
  pub field_types: Vec<PlumType>,
56
+ pub values: Vec<ast::Expr>,
56
57
  }
57
58
  /// Enum variant name -> its info, e.g. `"True" -> { enum_name: "Bool", tag: 1, field_types: [] }`.
58
59
  pub type EnumVariants = BTreeMap<String, EnumVariantInfo>;
59
60
 
61
+ /// Field names and types for every discriminant enum's shared params (`enum Foo(n: Int) = ...`),
62
+ /// keyed by the ENUM's name (not a variant name) — e.g. `"Step" -> [("n", TInt)]`. Field
63
+ /// access on a value of this type must load/store at offset `(field_idx + 1) * 8`, NOT
64
+ /// `field_idx * 8` like a class — slot 0 is always the variant's tag.
65
+ pub type EnumParams = BTreeMap<String, Vec<(String, PlumType)>>;
66
+
60
67
  /// Shared, read-only lookup tables built once from the whole source, threaded through
61
68
  /// every check/infer call alongside the (mutable, scope-local) `TypeEnv`.
62
69
  pub struct CheckCtx<'a> {
63
70
  pub classes: &'a ClassEnv,
64
71
  pub methods: &'a MethodEnv,
65
72
  pub enum_variants: &'a EnumVariants,
73
+ pub enum_params: &'a EnumParams,
66
74
  }
67
75
 
68
76
  /// Builds the global lookup tables (function/const signatures, class fields, method
69
77
  /// signatures, enum variants) from a whole source. Shared by `check_source` and by
70
78
  /// `plum-wasm-codegen`, which needs the same tables to resolve `self`, field access,
71
79
  /// and method dispatch during code generation.
72
- pub fn build_global_tables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEnv, EnumVariants) {
80
+ pub fn build_global_tables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEnv, EnumVariants, EnumParams) {
73
81
  let mut global_env: TypeEnv = TypeEnv::new();
74
82
  let mut classes: ClassEnv = BTreeMap::new();
75
83
  let mut methods: MethodEnv = BTreeMap::new();
76
84
  let mut enum_variants: EnumVariants = BTreeMap::new();
85
+ let mut enum_params: EnumParams = BTreeMap::new();
77
86
  // `Bool`'s variants are built in (see `infer_expr`'s TypeName handling) rather
78
87
  // than requiring every source file to redeclare `enum Bool = | True | False`.
79
- enum_variants.insert("True".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 1, field_types: vec![] });
88
+ enum_variants.insert("True".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 1, field_types: vec![], values: vec![] });
80
- enum_variants.insert("False".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 0, field_types: vec![] });
89
+ enum_variants.insert("False".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 0, field_types: vec![], values: vec![] });
81
90
 
82
91
  // First pass: register class fields and enum variants so later passes can
83
92
  // resolve `self.field`, `ClassName(...)`, and bare enum-tag patterns.
@@ -90,14 +99,28 @@ pub fn build_global_tables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEn
90
99
  classes.insert(c.name.clone(), fields);
91
100
  }
92
101
  ast::Item::Enum(e) => {
102
+ let shared_field_types: Vec<PlumType> = e.params.iter()
103
+ .map(|p| plum_type_from_ast(&p.ty))
104
+ .collect();
93
- for (tag, v) in e.variants.iter().enumerate() {
105
+ if !e.params.is_empty() {
94
- let field_types = v.fields.iter()
106
+ let params = e.params.iter()
95
- .map(|f| plum_type_from_ast(&ast::Type { name: f.clone(), generics: vec![] }))
107
+ .map(|p| (p.name.clone(), plum_type_from_ast(&p.ty)))
96
108
  .collect();
109
+ enum_params.insert(e.name.clone(), params);
110
+ }
111
+ for (tag, v) in e.variants.iter().enumerate() {
112
+ let field_types = if !e.params.is_empty() {
113
+ shared_field_types.clone()
114
+ } else {
115
+ v.fields.iter()
116
+ .map(|f| plum_type_from_ast(&ast::Type { name: f.clone(), generics: vec![] }))
117
+ .collect()
118
+ };
97
119
  enum_variants.insert(v.name.clone(), EnumVariantInfo {
98
120
  enum_name: e.name.clone(),
99
121
  tag: tag as i32,
100
122
  field_types,
123
+ values: v.values.clone(),
101
124
  });
102
125
  }
103
126
  }
@@ -139,14 +162,14 @@ pub fn build_global_tables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEn
139
162
  }
140
163
  }
141
164
 
142
- (global_env, classes, methods, enum_variants)
165
+ (global_env, classes, methods, enum_variants, enum_params)
143
166
  }
144
167
 
145
168
  pub fn check_source(source: &ast::Source) -> CheckResult<()> {
146
169
  let source = monomorphize::monomorphize_source(source).map_err(|e| vec![CheckError { message: e }])?;
147
170
  let mut errors: Vec<CheckError> = Vec::new();
148
- let (global_env, classes, methods, enum_variants) = build_global_tables(&source);
171
+ let (global_env, classes, methods, enum_variants, enum_params) = build_global_tables(&source);
149
- let ctx = CheckCtx { classes: &classes, methods: &methods, enum_variants: &enum_variants };
172
+ let ctx = CheckCtx { classes: &classes, methods: &methods, enum_variants: &enum_variants, enum_params: &enum_params };
150
173
 
151
174
  // A name that is both a class and an enum variant is ambiguous: `Name(...)`
152
175
  // could mean either construction, and downstream code (both the checker's
@@ -160,6 +183,49 @@ pub fn check_source(source: &ast::Source) -> CheckResult<()> {
160
183
  }
161
184
  }
162
185
 
186
+ // A discriminant enum (`enum Foo(n: Int) = ...`) requires every variant to supply
187
+ // exactly one value per declared param, unified against that param's type. An
188
+ // ordinary (param-less) enum must NOT have variants with values — most likely
189
+ // caused by writing `Some(5)` where the generic-payload form `Some[Int]` was meant.
190
+ for item in &source.items {
191
+ if let ast::Item::Enum(e) = item {
192
+ for v in &e.variants {
193
+ if e.params.is_empty() {
194
+ if !v.values.is_empty() {
195
+ errors.push(CheckError {
196
+ message: format!("enum '{}' variant '{}': has discriminant values but '{}' declares no params", e.name, v.name, e.name),
197
+ });
198
+ }
199
+ continue;
200
+ }
201
+ if v.values.len() != e.params.len() {
202
+ errors.push(CheckError {
203
+ message: format!(
204
+ "enum '{}' variant '{}': expected {} discriminant value(s), got {}",
205
+ e.name, v.name, e.params.len(), v.values.len()
206
+ ),
207
+ });
208
+ continue;
209
+ }
210
+ for (value, param) in v.values.iter().zip(e.params.iter()) {
211
+ match infer_expr(value, &global_env, &ctx) {
212
+ Ok(actual) => {
213
+ let expected = plum_type_from_ast(&param.ty);
214
+ if let Err(msg) = unify(&expected, &actual) {
215
+ errors.push(CheckError {
216
+ message: format!("enum '{}' variant '{}': param '{}': {}", e.name, v.name, param.name, msg),
217
+ });
218
+ }
219
+ }
220
+ Err(msg) => errors.push(CheckError {
221
+ message: format!("enum '{}' variant '{}': param '{}': {}", e.name, v.name, param.name, msg),
222
+ }),
223
+ }
224
+ }
225
+ }
226
+ }
227
+ }
228
+
163
229
  for item in &source.items {
164
230
  if let ast::Item::Fn(f) = item {
165
231
  let mut local_errors = check_fn(f, &global_env, &ctx);
@@ -639,8 +705,14 @@ pub fn infer_expr(expr: &ast::Expr, env: &TypeEnv, ctx: &CheckCtx) -> Result<Plu
639
705
  .find(|(n, _)| n == field_name)
640
706
  .map(|(_, t)| t.clone())
641
707
  .ok_or_else(|| format!("no field '{}' on type '{}'", field_name, class_name)),
708
+ None => match ctx.enum_params.get(class_name) {
709
+ Some(params) => params.iter()
710
+ .find(|(n, _)| n == field_name)
711
+ .map(|(_, t)| t.clone())
712
+ .ok_or_else(|| format!("no field '{}' on type '{}'", field_name, class_name)),
642
- // Unmodeled type: allow, codegen will catch.
713
+ // Unmodeled type: allow, codegen will catch.
643
- None => Ok(PlumType::TVar("_".to_string())),
714
+ None => Ok(PlumType::TVar("_".to_string())),
715
+ },
644
716
  },
645
717
  _ => Err(format!("cannot access field '{}' on non-class type {}", field_name, obj_ty)),
646
718
  },
plum-checker/src/monomorphize.rs CHANGED
@@ -181,7 +181,7 @@ pub fn specialize_enum(e: &ast::Enum, subst: &Substitution, mangled_name: &str)
181
181
 
182
182
  use std::collections::BTreeSet;
183
183
  use crate::types::{TypeEnv, TypeScheme};
184
- use crate::{ClassEnv, MethodEnv, EnumVariants, EnumVariantInfo, CheckCtx};
184
+ use crate::{ClassEnv, MethodEnv, EnumVariants, EnumVariantInfo, EnumParams, CheckCtx};
185
185
 
186
186
  enum PendingSpecialization<'a> {
187
187
  Class { base: &'a ast::Class, subst: Substitution, mangled: String },
@@ -221,6 +221,7 @@ struct Monomorphizer<'a> {
221
221
  classes: ClassEnv,
222
222
  methods: MethodEnv,
223
223
  enum_variants: EnumVariants,
224
+ enum_params: EnumParams,
224
225
  specialized: BTreeSet<String>,
225
226
  enqueued: BTreeSet<String>,
226
227
  worklist: Vec<PendingSpecialization<'a>>,
@@ -229,7 +230,7 @@ struct Monomorphizer<'a> {
229
230
 
230
231
  impl<'a> Monomorphizer<'a> {
231
232
  fn infer(&self, e: &ast::Expr, env: &TypeEnv) -> PlumType {
232
- let ctx = CheckCtx { classes: &self.classes, methods: &self.methods, enum_variants: &self.enum_variants };
233
+ let ctx = CheckCtx { classes: &self.classes, methods: &self.methods, enum_variants: &self.enum_variants, enum_params: &self.enum_params };
233
234
  crate::infer_expr(e, env, &ctx).unwrap_or(PlumType::TVar("_".to_string()))
234
235
  }
235
236
 
@@ -641,6 +642,7 @@ impl<'a> Monomorphizer<'a> {
641
642
  enum_name: mangled.clone(),
642
643
  tag: tag as i32,
643
644
  field_types,
645
+ values: v.values.clone(),
644
646
  });
645
647
  }
646
648
  self.enum_variant_mangling.insert(mangled.clone(), table);
@@ -736,7 +738,7 @@ impl<'a> Monomorphizer<'a> {
736
738
  /// result has no generic syntax left in it — `check_source`/`compile_source` run
737
739
  /// on it completely unmodified.
738
740
  pub fn monomorphize_source(source: &ast::Source) -> Result<ast::Source, String> {
739
- let (global_env, classes, methods, enum_variants) = crate::build_global_tables(source);
741
+ let (global_env, classes, methods, enum_variants, enum_params) = crate::build_global_tables(source);
740
742
 
741
743
  let mut m = Monomorphizer {
742
744
  classes_generic: BTreeMap::new(),
@@ -750,6 +752,7 @@ pub fn monomorphize_source(source: &ast::Source) -> Result<ast::Source, String>
750
752
  classes,
751
753
  methods,
752
754
  enum_variants,
755
+ enum_params,
753
756
  specialized: BTreeSet::new(),
754
757
  enqueued: BTreeSet::new(),
755
758
  worklist: Vec::new(),
plum-checker/tests/checker_tests.rs CHANGED
@@ -254,6 +254,89 @@ fun bad() -> Shape =
254
254
  assert!(result.is_err());
255
255
  }
256
256
 
257
+ #[test]
258
+ fn enum_discriminant_values_type_check_with_no_errors() {
259
+ let src = "\
260
+ enum Step(n: Int) =
261
+ | ReadMin(0)
262
+ | ReadMax(1)
263
+ ";
264
+ let source = parse(src);
265
+ let result = check_source(&source);
266
+ assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
267
+ }
268
+
269
+ #[test]
270
+ fn enum_discriminant_wrong_value_count_is_error() {
271
+ let src = "\
272
+ enum Step(n: Int) =
273
+ | ReadMin(0)
274
+ | ReadMax
275
+ ";
276
+ let source = parse(src);
277
+ let result = check_source(&source);
278
+ assert!(result.is_err());
279
+ }
280
+
281
+ #[test]
282
+ fn enum_discriminant_wrongly_typed_value_is_error() {
283
+ let src = "\
284
+ enum Step(n: Int) =
285
+ | ReadMin(0)
286
+ | ReadMax(\"x\")
287
+ ";
288
+ let source = parse(src);
289
+ let result = check_source(&source);
290
+ assert!(result.is_err());
291
+ }
292
+
293
+ #[test]
294
+ fn enum_discriminant_value_on_param_less_enum_is_error() {
295
+ let src = "\
296
+ enum Option =
297
+ | Some(5)
298
+ | None
299
+ ";
300
+ let source = parse(src);
301
+ let result = check_source(&source);
302
+ assert!(result.is_err());
303
+ }
304
+
305
+ #[test]
306
+ fn field_access_on_discriminant_enum_receiver_type_checks() {
307
+ let src = "\
308
+ enum Step(n: Int) =
309
+ | ReadMin(0)
310
+ | ReadMax(1)
311
+
312
+ fun toNumber(self) -> Int =
313
+ self.n
314
+ ";
315
+ let source = parse(src);
316
+ let result = check_source(&source);
317
+ assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
318
+ }
319
+
320
+ #[test]
321
+ fn field_access_on_ordinary_enum_receiver_is_unaffected_by_enum_params() {
322
+ // An ordinary enum (no discriminant `params`) isn't in `ctx.classes` OR
323
+ // `ctx.enum_params`, so it falls through to the same permissive "unmodeled
324
+ // type" escape hatch every other type not in `ctx.classes` gets (codegen,
325
+ // not the checker, is what would catch a genuinely bad field access here) —
326
+ // exactly as it did before discriminant enums existed.
327
+ let src = "\
328
+ enum Option =
329
+ | Some[Int]
330
+ | None
331
+
332
+ fun bad(self) -> Int =
333
+ self.n
334
+ ";
335
+ let source = parse(src);
336
+ let result = check_source(&source);
337
+ assert!(result.is_ok(), "expected Ok, got {:?}", result.err());
338
+ }
339
+
257
340
  #[test]
258
341
  fn constructor_pattern_binds_fields_to_declared_types() {
259
342
  let src = "\
plum-checker/tests/monomorphize_tests.rs CHANGED
@@ -42,9 +42,10 @@ fn fn_generic_params_detects_implicit_uppercase_letter_types_in_order() {
42
42
  fn enum_generic_params_detects_implicit_uppercase_letter_variant_fields() {
43
43
  let e = ast::Enum {
44
44
  name: "Option".to_string(),
45
+ params: vec![],
45
46
  variants: vec![
46
- ast::EnumVariant { name: "Some".to_string(), fields: vec!["T".to_string()] },
47
+ ast::EnumVariant { name: "Some".to_string(), fields: vec!["T".to_string()], values: vec![] },
47
- ast::EnumVariant { name: "None".to_string(), fields: vec![] },
48
+ ast::EnumVariant { name: "None".to_string(), fields: vec![], values: vec![] },
48
49
  ],
49
50
  };
50
51
  assert_eq!(enum_generic_params(&e), vec!["T".to_string()]);
@@ -114,9 +115,10 @@ fn specialize_fn_sets_new_receiver_for_a_method() {
114
115
  fn specialize_enum_substitutes_generic_variant_field_names() {
115
116
  let e = ast::Enum {
116
117
  name: "Option".to_string(),
118
+ params: vec![],
117
119
  variants: vec![
118
- ast::EnumVariant { name: "Some".to_string(), fields: vec!["T".to_string()] },
120
+ ast::EnumVariant { name: "Some".to_string(), fields: vec!["T".to_string()], values: vec![] },
119
- ast::EnumVariant { name: "None".to_string(), fields: vec![] },
121
+ ast::EnumVariant { name: "None".to_string(), fields: vec![], values: vec![] },
120
122
  ],
121
123
  };
122
124
  let mut bindings = std::collections::BTreeMap::new();
plum-wasm-codegen/src/lib.rs CHANGED
@@ -3,7 +3,7 @@ use std::cell::RefCell;
3
3
  use std::collections::{HashMap, HashSet};
4
4
  use plum_core::ast;
5
5
  use plum_checker::types::{PlumType, TypeEnv, TypeScheme};
6
- use plum_checker::{ClassEnv, MethodEnv, EnumVariants, EnumVariantInfo};
6
+ use plum_checker::{ClassEnv, MethodEnv, EnumVariants, EnumVariantInfo, EnumParams};
7
7
 
8
8
  /// Bump-allocated heap for class instances starts at the second 64KiB page so it can
9
9
  /// never collide with the (small, compile-time-sized) string literal data area below it.
@@ -247,6 +247,7 @@ pub struct CompileCtx<'a> {
247
247
  pub classes: ClassEnv,
248
248
  pub methods: MethodEnv,
249
249
  pub enum_variants: EnumVariants,
250
+ pub enum_params: EnumParams,
250
251
  pub global_env: TypeEnv,
251
252
  pub bump_global: u32,
252
253
  /// Closure literal (keyed by `&Expr::Closure` pointer identity) -> its `ClosureInfo`.
@@ -315,6 +316,7 @@ struct LocalCtx<'a> {
315
316
  classes: &'a ClassEnv,
316
317
  methods: &'a MethodEnv,
317
318
  enum_variants: &'a EnumVariants,
319
+ enum_params: &'a EnumParams,
318
320
  /// Tracks each binding's inferred type as compilation proceeds through
319
321
  /// statements in order, mirroring `plum-checker`'s own env evolution — needed
320
322
  /// to resolve `Attribute`/`ClassCall` targets and pick the right load/store width.
@@ -407,8 +409,8 @@ fn encode_leb128_u32(mut val: u32) -> Vec<u8> {
407
409
  bytes
408
410
  }
409
411
 
410
- fn check_ctx_of<'a>(ctx_classes: &'a ClassEnv, ctx_methods: &'a MethodEnv, ctx_enum_variants: &'a EnumVariants) -> plum_checker::CheckCtx<'a> {
412
+ fn check_ctx_of<'a>(ctx_classes: &'a ClassEnv, ctx_methods: &'a MethodEnv, ctx_enum_variants: &'a EnumVariants, ctx_enum_params: &'a EnumParams) -> plum_checker::CheckCtx<'a> {
411
- plum_checker::CheckCtx { classes: ctx_classes, methods: ctx_methods, enum_variants: ctx_enum_variants }
413
+ plum_checker::CheckCtx { classes: ctx_classes, methods: ctx_methods, enum_variants: ctx_enum_variants, enum_params: ctx_enum_params }
412
414
  }
413
415
 
414
416
  /// Infers an expression's type using the function's current (mutable, evolving) type
@@ -417,13 +419,13 @@ fn check_ctx_of<'a>(ctx_classes: &'a ClassEnv, ctx_methods: &'a MethodEnv, ctx_e
417
419
  /// codegen is being driven directly on unchecked input (as some tests do).
418
420
  fn infer_local_type(expr: &ast::Expr, ctx: &LocalCtx) -> PlumType {
419
421
  let env = ctx.type_env.borrow();
420
- let cctx = check_ctx_of(ctx.classes, ctx.methods, ctx.enum_variants);
422
+ let cctx = check_ctx_of(ctx.classes, ctx.methods, ctx.enum_variants, ctx.enum_params);
421
423
  plum_checker::infer_expr(expr, &env, &cctx).unwrap_or(PlumType::TInt)
422
424
  }
423
425
 
424
426
  pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
425
427
  let source = &plum_checker::monomorphize::monomorphize_source(source)?;
426
- let (global_env, classes, methods, enum_variants) = plum_checker::build_global_tables(source);
428
+ let (global_env, classes, methods, enum_variants, enum_params) = plum_checker::build_global_tables(source);
427
429
 
428
430
  let mut module = WasmModule::new();
429
431
  module.add_memory(2, None);
@@ -502,7 +504,7 @@ pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
502
504
  }
503
505
  let mut walker = ClosureWalker {
504
506
  env,
505
- cctx: check_ctx_of(&classes, &methods, &enum_variants),
507
+ cctx: check_ctx_of(&classes, &methods, &enum_variants, &enum_params),
506
508
  fn_decls: &fn_decls,
507
509
  found: Vec::new(),
508
510
  locals,
@@ -596,7 +598,7 @@ pub fn compile_source(source: &ast::Source) -> Result<Vec<u8>, String> {
596
598
  let int_to_string_func = register_int_to_string_helper(&mut module, bump_global);
597
599
 
598
600
  let ctx = CompileCtx {
599
- func_ids, func_sigs, classes, methods, enum_variants, global_env, bump_global,
601
+ func_ids, func_sigs, classes, methods, enum_variants, enum_params, global_env, bump_global,
600
602
  closures, closure_asts, closure_call_types, named_fn_values,
601
603
  string_concat_func, int_to_string_func,
602
604
  };
@@ -1793,7 +1795,7 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
1793
1795
 
1794
1796
  let mut collector = Collector {
1795
1797
  env: base_env.clone(),
1796
- cctx: check_ctx_of(&ctx.classes, &ctx.methods, &ctx.enum_variants),
1798
+ cctx: check_ctx_of(&ctx.classes, &ctx.methods, &ctx.enum_variants, &ctx.enum_params),
1797
1799
  named: Vec::new(),
1798
1800
  named_set: Default::default(),
1799
1801
  classcall_scratch: HashMap::new(),
@@ -1903,6 +1905,7 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
1903
1905
  classes: &ctx.classes,
1904
1906
  methods: &ctx.methods,
1905
1907
  enum_variants: &ctx.enum_variants,
1908
+ enum_params: &ctx.enum_params,
1906
1909
  type_env: RefCell::new(base_env),
1907
1910
  closure_local_sigs: RefCell::new(HashMap::new()),
1908
1911
  bump_global: ctx.bump_global,
@@ -3337,7 +3340,7 @@ fn compile_closure_body(
3337
3340
 
3338
3341
  let mut collector = Collector {
3339
3342
  env: base_env.clone(),
3340
- cctx: check_ctx_of(&ctx.classes, &ctx.methods, &ctx.enum_variants),
3343
+ cctx: check_ctx_of(&ctx.classes, &ctx.methods, &ctx.enum_variants, &ctx.enum_params),
3341
3344
  named: Vec::new(),
3342
3345
  named_set: Default::default(),
3343
3346
  classcall_scratch: HashMap::new(),
@@ -3455,6 +3458,7 @@ fn compile_closure_body(
3455
3458
  classes: &ctx.classes,
3456
3459
  methods: &ctx.methods,
3457
3460
  enum_variants: &ctx.enum_variants,
3461
+ enum_params: &ctx.enum_params,
3458
3462
  type_env: RefCell::new(base_env),
3459
3463
  closure_local_sigs: RefCell::new(HashMap::new()),
3460
3464
  bump_global: ctx.bump_global,