3254688 | plum

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

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

3254688 — Peter John 2026-07-20T13:27:12+05:30

fix(plum-wasm-codegen): propagate a value through tail-position match/if without explicit return

Files changed (2) hide show

plum-wasm-codegen/src/lib.rs +202 -89
plum-wasm-codegen/tests/codegen_tests.rs +132 -0

plum-wasm-codegen/src/lib.rs CHANGED Viewed

@@ -255,6 +255,10 @@ fn plum_type_to_valtype(t: &PlumType) -> ValType {
     }
 }
+fn block_type_for(result_vt: Option<ValType>) -> BlockType {
+    result_vt.map(BlockType::Result).unwrap_or(BlockType::Empty)
+}
 fn ret_type_to_wasm(ret: Option<&ast::ReturnType>) -> Option<ValType> {
     ret.and_then(|r| ast_type_to_wasm(&r.name))
 }
@@ -617,14 +621,14 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         bump_global: ctx.bump_global,
     };
-    let has_return_value = f.returns.as_ref().map(|r| r.name != "Unit").unwrap_or(false);
+    let result_vt = ret_type_to_wasm(f.returns.as_ref());
     match &f.body {
         ast::FnBody::Expr(e) => {
             compile_expr(e, &mut body, &local_ctx, state)?;
         }
         ast::FnBody::Block(block) => {
-            compile_block_as_fn_body(block, &mut body, &local_ctx, state, has_return_value)?;
+            compile_block_as_fn_body(block, &mut body, &local_ctx, state, result_vt)?;
         }
     }
@@ -639,70 +643,119 @@ fn compile_block(block: &ast::Block, body: &mut Vec<u8>, ctx: &LocalCtx, state:
     Ok(())
 }
+/// Compiles a case/branch body either as an ordinary statement block (`result_vt: None`)
+/// or, when in value position, via `compile_block_in_value_position` so its own tail
+/// statement propagates a value instead of being dropped.
+fn compile_case_body(
+    block: &ast::Block,
+    result_vt: Option<ValType>,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+) -> Result<(), String> {
+    match result_vt {
+        Some(vt) => compile_block_in_value_position(block, vt, body, ctx, state),
+        None => compile_block(block, body, ctx, state),
+    }
+}
+/// Compiles a block whose value must be produced when control reaches its end — every
+/// statement except the last compiles normally; the last is compiled via
+/// `compile_stmt_in_value_position`.
+fn compile_block_in_value_position(
+    block: &ast::Block,
+    result_vt: ValType,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+) -> Result<(), String> {
+    let (last, rest) = block.stmts.split_last().ok_or_else(|| {
+        "codegen: function has a control-flow path that doesn't produce a return value (empty branch)".to_string()
+    })?;
+    for stmt in rest {
+        compile_stmt(stmt, body, ctx, state)?;
+    }
+    compile_stmt_in_value_position(last, result_vt, body, ctx, state)
+}
+/// Compiles a single statement in value position: a bare expression is left on the stack
+/// (not dropped); `return`/`todo` compile normally (both are stack-polymorphic in wasm —
-/// True if control can never fall through past this statement — every reachable path
+/// control never falls through past them, so no value is needed on this path); `if`/`match`
-/// ends in a `return`. Used to decide whether a tail-position If/Match needs a
+/// recurse so every arm/branch resolves the same way. Any other statement kind can't
-/// trailing `unreachable` to satisfy wasm's per-block (not whole-function) validation
+/// produce a value, so this returns a clear error instead of ever emitting wasm that
-/// when the function declares a non-Unit return type.
+/// would fail validation.
+fn compile_stmt_in_value_position(
+    stmt: &ast::Stmt,
+    result_vt: ValType,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
-fn stmt_always_diverges(stmt: &ast::Stmt) -> bool {
+) -> Result<(), String> {
     match stmt {
+        ast::Stmt::Expr(e) => compile_expr(e, body, ctx, state),
-        ast::Stmt::Return(_) | ast::Stmt::Todo => true,
+        ast::Stmt::Return(_) | ast::Stmt::Todo => compile_stmt(stmt, body, ctx, state),
-        ast::Stmt::If(if_) => {
-            if_.else_.is_some()
-                && block_always_diverges(&if_.body)
-                && if_.else_ifs.iter().all(|ei| block_always_diverges(&ei.body))
-                && if_.else_.as_ref().is_some_and(block_always_diverges)
+        ast::Stmt::If(if_) => compile_if(if_, Some(result_vt), body, ctx, state),
-        }
-        ast::Stmt::Match(m) => !m.cases.is_empty() && m.cases.iter().all(|c| block_always_diverges(&c.body)),
+        ast::Stmt::Match(m) => compile_match(m, body, ctx, state, Some(result_vt)),
-        _ => false,
+        _ => Err(
+            "codegen: function has a control-flow path that doesn't produce a return value".to_string(),
+        ),
     }
 }
+/// Compiles an `if`/`else if`/`else` chain. `result_vt` is `None` for an ordinary statement
+/// (each branch is `BlockType::Empty`, nothing left on the stack) or `Some(vt)` when this
+/// `if` is in value position — every branch must then leave a `vt` value on the stack, which
+/// requires an `else` (a value can't be produced on a path that doesn't exist).
+fn compile_if(
+    if_: &ast::If,
+    result_vt: Option<ValType>,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+) -> Result<(), String> {
+    if result_vt.is_some() && if_.else_.is_none() {
+        return Err(
+            "codegen: function has a control-flow path that doesn't produce a return value (if without else)".to_string(),
+        );
+    }
+    let bt = block_type_for(result_vt);
+    compile_expr(&if_.condition, body, ctx, state)?;
+    Instruction::If(bt).encode(body);
+    compile_case_body(&if_.body, result_vt, body, ctx, state)?;
+    if !if_.else_ifs.is_empty() || if_.else_.is_some() {
+        Instruction::Else.encode(body);
+        for ei in &if_.else_ifs {
+            compile_expr(&ei.condition, body, ctx, state)?;
+            Instruction::If(bt).encode(body);
+            compile_case_body(&ei.body, result_vt, body, ctx, state)?;
+            Instruction::Else.encode(body);
+        }
+        if let Some(else_block) = &if_.else_ {
-fn block_always_diverges(block: &ast::Block) -> bool {
+            compile_case_body(else_block, result_vt, body, ctx, state)?;
+        }
-    block.stmts.last().map(stmt_always_diverges).unwrap_or(false)
+        for _ in &if_.else_ifs {
+            Instruction::End.encode(body);
+        }
+    }
+    Instruction::End.encode(body);
+    Ok(())
 }
-/// Compiles a block that is the body of a function. If the function returns a value
+/// Compiles a block that is the body of a function. If the function returns a value,
-/// and the last statement is an expression, that expression's value is left on the
+/// its tail statement is compiled in value position (see `compile_stmt_in_value_position`)
+/// so a bare expression, or an `if`/`match` whose arms resolve to one, propagates that
-/// stack instead of being dropped.
+/// value instead of being dropped.
 fn compile_block_as_fn_body(
     block: &ast::Block,
     body: &mut Vec<u8>,
     ctx: &LocalCtx,
     state: &mut ModuleState,
-    has_return_value: bool,
+    result_vt: Option<ValType>,
 ) -> Result<(), String> {
-    let stmts = &block.stmts;
-    if has_return_value {
-        if let Some((last, rest)) = stmts.split_last() {
-            for stmt in rest {
-                compile_stmt(stmt, body, ctx, state)?;
-            }
-            match last {
+    match result_vt {
-                ast::Stmt::Expr(e) => {
-                    compile_expr(e, body, ctx, state)?;
+        Some(vt) => compile_block_in_value_position(block, vt, body, ctx, state),
-                    // do NOT drop — this is the return value
-                }
-                _ if stmt_always_diverges(last) => {
-                    compile_stmt(last, body, ctx, state)?;
+        None => compile_block(block, body, ctx, state),
-                    // Every branch of this If/Match already `return`s; if it somehow
-                    // falls through anyway (a bug, or a non-exhaustive match), trap
-                    // rather than continue with the required result value missing.
-                    // This also satisfies wasm validation, which requires a value at
-                    // the function's `end` regardless of whether every branch inside
-                    // an (BlockType::Empty) if/else chain already returned.
-                    Instruction::Unreachable.encode(body);
-                }
-                _ => {
-                    compile_stmt(last, body, ctx, state)?;
-                }
-            }
-            return Ok(());
-        }
     }
-    for stmt in stmts {
-        compile_stmt(stmt, body, ctx, state)?;
-    }
-    Ok(())
 }
 fn compile_stmt(stmt: &ast::Stmt, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mut ModuleState) -> Result<(), String> {
@@ -728,25 +781,7 @@ fn compile_stmt(stmt: &ast::Stmt, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mu
             Instruction::Return.encode(body);
         }
         ast::Stmt::If(if_) => {
-            compile_expr(&if_.condition, body, ctx, state)?;
-            Instruction::If(BlockType::Empty).encode(body);
-            compile_block(&if_.body, body, ctx, state)?;
-            if !if_.else_ifs.is_empty() || if_.else_.is_some() {
-                Instruction::Else.encode(body);
-                for ei in &if_.else_ifs {
-                    compile_expr(&ei.condition, body, ctx, state)?;
-                    Instruction::If(BlockType::Empty).encode(body);
-                    compile_block(&ei.body, body, ctx, state)?;
-                    Instruction::Else.encode(body);
-                }
-                if let Some(else_block) = &if_.else_ {
-                    compile_block(else_block, body, ctx, state)?;
+            compile_if(if_, None, body, ctx, state)?;
-                }
-                for _ in &if_.else_ifs {
-                    Instruction::End.encode(body);
-                }
-            }
-            Instruction::End.encode(body);
         }
         ast::Stmt::While(w) => {
             Instruction::Block(BlockType::Empty).encode(body);
@@ -805,7 +840,7 @@ fn compile_stmt(stmt: &ast::Stmt, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mu
             Instruction::Br(0).encode(body);
         }
         ast::Stmt::Match(m) => {
-            compile_match(m, body, ctx, state)?;
+            compile_match(m, body, ctx, state, None)?;
         }
         ast::Stmt::Assert(e) => {
             compile_expr(e, body, ctx, state)?;
@@ -841,7 +876,60 @@ fn expr_has_result(expr: &ast::Expr, ctx: &LocalCtx) -> bool {
     }
 }
+/// True if `cases` consists solely of enum-tag patterns (bare variant names or
+/// constructor patterns, no wildcard/binding/int/etc.) that between them cover every
+/// variant of a single enum type. When that holds, a match compiled in value position
+/// can never actually fall through past the last arm at runtime — even though the
+/// patterns don't include an explicit wildcard/binding catch-all — so the "ran out of
+/// patterns" fallback in `compile_match_arms` is provably unreachable code, not a real
+/// gap. `compile_match` uses this to append a synthetic trap-and-never-fall-through
+/// wildcard arm (rather than let the arms recursion hit its non-exhaustive-match error)
+/// so previously-working exhaustive enum matches (e.g. `Some`/`None`, `True`/`False`)
+/// keep compiling even without a trailing wildcard, while a genuinely non-exhaustive
+/// match (an `Int` match, or an enum match missing a variant) still gets a clear error.
-fn compile_match(m: &ast::Match, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mut ModuleState) -> Result<(), String> {
+fn match_covers_every_enum_variant(cases: &[ast::Case], subject_vt: ValType, ctx: &LocalCtx) -> bool {
+    if subject_vt != ValType::I32 {
+        return false;
+    }
+    let mut enum_name: Option<&str> = None;
+    let mut tags_seen: std::collections::BTreeSet<i32> = std::collections::BTreeSet::new();
+    for case in cases {
+        let variant_name = match case.patterns.first() {
+            Some(ast::CasePattern::Name(n))
+                if n.chars().next().map(|c| c.is_uppercase()).unwrap_or(false) && ctx.enum_variants.contains_key(n) =>
+            {
+                n.as_str()
+            }
+            Some(ast::CasePattern::Class { name, .. }) => name.as_str(),
+            _ => return false,
+        };
+        let info = match ctx.enum_variants.get(variant_name) {
+            Some(info) => info,
+            None => return false,
+        };
+        match enum_name {
+            Some(en) if en != info.enum_name => return false,
+            Some(_) => {}
+            None => enum_name = Some(&info.enum_name),
+        }
+        tags_seen.insert(info.tag);
+    }
+    match enum_name {
+        Some(en) => {
+            let total_variants = ctx.enum_variants.values().filter(|v| v.enum_name == en).count();
+            !tags_seen.is_empty() && tags_seen.len() == total_variants
+        }
+        None => false,
+    }
+}
+fn compile_match(
+    m: &ast::Match,
+    body: &mut Vec<u8>,
+    ctx: &LocalCtx,
+    state: &mut ModuleState,
+    result_vt: Option<ValType>,
+) -> Result<(), String> {
     if m.subjects.len() != 1 {
         return Err("codegen: multi-subject match is not yet supported".to_string());
     }
@@ -859,32 +947,55 @@ fn compile_match(m: &ast::Match, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mut
     compile_expr(subject, body, ctx, state)?;
     Instruction::LocalSet(scratch_local).encode(body);
+    // A match in value position whose arms already cover every variant of the
+    // subject's enum type (by explicit tag/constructor patterns, no wildcard) is
+    // exhaustive at runtime even though `compile_match_arms` can't see that from the
+    // remaining-cases slice alone. Give it a synthetic trailing `_ => todo` arm so the
+    // "out of patterns" fallback compiles to an (unreachable, but valid) trap instead
+    // of a spurious non-exhaustive-match error.
+    if result_vt.is_some() && match_covers_every_enum_variant(&m.cases, subject_vt, ctx) {
+        let mut cases = m.cases.clone();
+        cases.push(ast::Case {
+            patterns: vec![ast::CasePattern::Wildcard],
+            body: ast::Block { stmts: vec![ast::Stmt::Todo] },
+        });
+        return compile_match_arms(&cases, subject_vt, scratch_local, result_vt, body, ctx, state);
+    }
-    compile_match_arms(&m.cases, subject_vt, scratch_local, body, ctx, state)
+    compile_match_arms(&m.cases, subject_vt, scratch_local, result_vt, body, ctx, state)
 }
 fn compile_match_arms(
     cases: &[ast::Case],
     subject_vt: ValType,
     scratch_local: u32,
+    result_vt: Option<ValType>,
     body: &mut Vec<u8>,
     ctx: &LocalCtx,
     state: &mut ModuleState,
 ) -> Result<(), String> {
     let (case, rest) = match cases.split_first() {
+        None => {
+            return match result_vt {
+                Some(_) => Err(
+                    "codegen: function has a control-flow path that doesn't produce a return value (non-exhaustive match)".to_string(),
+                ),
-        None => return Ok(()),
+                None => Ok(()),
+            };
+        }
         Some(pair) => pair,
     };
     let pat = case.patterns.first().ok_or_else(|| "codegen: match case has no pattern".to_string())?;
     match pat {
         ast::CasePattern::Wildcard => {
             // Any cases after a wildcard are unreachable, matching real match semantics.
-            compile_block(&case.body, body, ctx, state)
+            compile_case_body(&case.body, result_vt, body, ctx, state)
         }
         ast::CasePattern::Name(n) => {
             let is_variant = n.chars().next().map(|c| c.is_uppercase()).unwrap_or(false)
                 && ctx.enum_variants.contains_key(n);
             if is_variant {
-                compile_variant_eq_arm(n, subject_vt, scratch_local, case, rest, body, ctx, state)
+                compile_variant_eq_arm(n, subject_vt, scratch_local, result_vt, case, rest, body, ctx, state)
             } else {
                 let idx = ctx
                     .locals
@@ -894,7 +1005,7 @@ fn compile_match_arms(
                 Instruction::LocalGet(scratch_local).encode(body);
                 Instruction::LocalSet(idx).encode(body);
                 ctx.type_env.borrow_mut().insert(n.clone(), TypeScheme::mono(plum_type_from_valtype_hint(subject_vt)));
-                compile_block(&case.body, body, ctx, state)
+                compile_case_body(&case.body, result_vt, body, ctx, state)
                 // A binding arm always matches — any following cases are unreachable.
             }
         }
@@ -905,17 +1016,17 @@ fn compile_match_arms(
             Instruction::LocalGet(scratch_local).encode(body);
             Instruction::I64Const(*n).encode(body);
             Instruction::I64Eq.encode(body);
-            Instruction::If(BlockType::Empty).encode(body);
+            Instruction::If(block_type_for(result_vt)).encode(body);
-            compile_block(&case.body, body, ctx, state)?;
+            compile_case_body(&case.body, result_vt, body, ctx, state)?;
             Instruction::Else.encode(body);
-            compile_match_arms(rest, subject_vt, scratch_local, body, ctx, state)?;
+            compile_match_arms(rest, subject_vt, scratch_local, result_vt, body, ctx, state)?;
             Instruction::End.encode(body);
             Ok(())
         }
         ast::CasePattern::String(_) => Err("codegen: string match patterns are not yet supported".to_string()),
         ast::CasePattern::Float(_) => Err("codegen: float match patterns are not yet supported".to_string()),
         ast::CasePattern::Class { name, fields } => {
-            compile_variant_constructor_arm(name, fields, subject_vt, scratch_local, case, rest, body, ctx, state)
+            compile_variant_constructor_arm(name, fields, subject_vt, scratch_local, result_vt, case, rest, body, ctx, state)
         }
     }
 }
@@ -925,6 +1036,7 @@ fn compile_variant_eq_arm(
     name: &str,
     subject_vt: ValType,
     scratch_local: u32,
+    result_vt: Option<ValType>,
     case: &ast::Case,
     rest: &[ast::Case],
     body: &mut Vec<u8>,
@@ -942,10 +1054,10 @@ fn compile_variant_eq_arm(
     Instruction::LocalGet(scratch_local).encode(body);
     Instruction::I32Const(tag).encode(body);
     Instruction::I32Eq.encode(body);
-    Instruction::If(BlockType::Empty).encode(body);
+    Instruction::If(block_type_for(result_vt)).encode(body);
-    compile_block(&case.body, body, ctx, state)?;
+    compile_case_body(&case.body, result_vt, body, ctx, state)?;
     Instruction::Else.encode(body);
-    compile_match_arms(rest, subject_vt, scratch_local, body, ctx, state)?;
+    compile_match_arms(rest, subject_vt, scratch_local, result_vt, body, ctx, state)?;
     Instruction::End.encode(body);
     Ok(())
 }
@@ -956,6 +1068,7 @@ fn compile_variant_constructor_arm(
     fields: &[ast::CasePattern],
     subject_vt: ValType,
     scratch_local: u32,
+    result_vt: Option<ValType>,
     case: &ast::Case,
     rest: &[ast::Case],
     body: &mut Vec<u8>,
@@ -994,7 +1107,7 @@ fn compile_variant_constructor_arm(
     Instruction::Else.encode(body);
     Instruction::I32Const(0).encode(body);
     Instruction::End.encode(body);
-    Instruction::If(BlockType::Empty).encode(body);
+    Instruction::If(block_type_for(result_vt)).encode(body);
     for (i, (pat, field_ty)) in fields.iter().zip(field_types.iter()).enumerate() {
         let bind_name = match pat {
             ast::CasePattern::Name(n) => Some(n.as_str()),
@@ -1018,9 +1131,9 @@ fn compile_variant_constructor_arm(
             ctx.type_env.borrow_mut().insert(n.to_string(), TypeScheme::mono(field_ty.clone()));
         }
     }
-    compile_block(&case.body, body, ctx, state)?;
+    compile_case_body(&case.body, result_vt, body, ctx, state)?;
     Instruction::Else.encode(body);
-    compile_match_arms(rest, subject_vt, scratch_local, body, ctx, state)?;
+    compile_match_arms(rest, subject_vt, scratch_local, result_vt, body, ctx, state)?;
     Instruction::End.encode(body);
     Ok(())
 }

plum-wasm-codegen/tests/codegen_tests.rs CHANGED Viewed

@@ -552,3 +552,135 @@ main() -> Int =
     let bytes = compile_source(&source).expect("compile failed");
     assert_eq!(run_main(&bytes), 42);
 }
+#[test]
+fn tail_match_without_return_runs_correctly() {
+    let src = "\
+bindExample(n: Int) -> Int =
+  match n
+    x =>
+      x
+main() -> Int =
+  bindExample(5)
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), 5);
+}
+#[test]
+fn tail_if_without_return_runs_correctly() {
+    let src = "\
+abs(n: Int) -> Int =
+  if n < 0
+    -n
+  else
+    n
+main() -> Int =
+  abs(-7)
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), 7);
+}
+#[test]
+fn tail_if_nested_inside_match_arm_without_return_runs_correctly() {
+    let src = "\
+classify(n: Int) -> Int =
+  match n
+    0 =>
+      1
+    x =>
+      if x < 0
+        -1
+      else
+        2
+main() -> Int =
+  classify(-5)
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), -1);
+}
+#[test]
+fn tail_match_mixing_return_and_bare_expr_arms_runs_correctly() {
+    let src = "\
+describe(n: Int) -> Int =
+  match n
+    0 =>
+      return 100
+    x =>
+      x * 2
+main() -> Int =
+  describe(21)
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), 42);
+}
+#[test]
+fn tail_enum_match_without_return_runs_correctly() {
+    let src = "\
+enum Option =
+  | Some(Int)
+  | None
+unwrapOr(o: Option, default: Int) -> Int =
+  match o
+    Some(v) =>
+      v
+    None =>
+      default
+main() -> Int =
+  unwrapOr(Some(9), 0)
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), 9);
+}
+#[test]
+fn tail_if_without_else_is_a_clear_error() {
+    let src = "\
+bad(n: Int) -> Int =
+  if n < 0
+    return 1
+";
+    let source = parse(src);
+    let err = compile_source(&source).expect_err("if without else in value position must be a clear error, not invalid wasm");
+    assert!(err.contains("doesn't produce a return value"), "got: {}", err);
+}
+#[test]
+fn tail_match_non_exhaustive_is_a_clear_error() {
+    let src = "\
+bad(n: Int) -> Int =
+  match n
+    0 =>
+      1
+";
+    let source = parse(src);
+    let err = compile_source(&source).expect_err("non-exhaustive match in value position must be a clear error, not invalid wasm");
+    assert!(err.contains("doesn't produce a return value"), "got: {}", err);
+}
+#[test]
+fn tail_match_arm_ending_in_non_value_statement_is_a_clear_error() {
+    let src = "\
+bad(n: Int) -> Int =
+  match n
+    x =>
+      y = x
+";
+    let source = parse(src);
+    let err = compile_source(&source).expect_err("a match arm ending in a non-value statement must be a clear error, not invalid wasm");
+    assert!(err.contains("doesn't produce a return value"), "got: {}", err);
+}