da1c377 | plum

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

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

da1c377 — Peter John 2026-07-23T20:06:30+05:30

feat(plum-wasm-codegen): compile for-v-in-variadic-param iteration

Files changed (3) hide show

examples/functions.plum +4 -1
plum-wasm-codegen/src/lib.rs +114 -3
plum-wasm-codegen/tests/codegen_tests.rs +34 -0

examples/functions.plum CHANGED Viewed

@@ -8,7 +8,10 @@ withDefault(a: Int, step: Int = 1) -> Int =
   a + step
 sumAll(nums: ...Int) -> Int =
+  total = 0
+  for v in nums
+    total = total + v
-  todo
+  total
 wrap(value: a) -> Bool =
   True

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

@@ -299,6 +299,12 @@ struct LocalCtx<'a> {
     nested_class_scratch_base: u32,
     /// `CasePattern::Class` identity (pointer address) -> scratch slot offset.
     nested_class_scratch: HashMap<usize, u32>,
+    /// First local index reserved for variadic-`for` scratch temporaries (2 `i32`
+    /// slots per `for` statement that iterates a `TVariadic`: count, loop index).
+    variadic_for_scratch_base: u32,
+    /// `For` stmt identity (pointer address) -> slot number (multiply by 2 and add
+    /// `variadic_for_scratch_base` for the count local; +1 more for the index local).
+    variadic_for_scratch: HashMap<usize, u32>,
     func_ids: &'a HashMap<String, u32>,
     func_sigs: &'a HashMap<String, FuncSig>,
     closures: &'a HashMap<usize, ClosureInfo>,
@@ -975,8 +981,12 @@ impl<'a, 'c> ClosureWalker<'a, 'c> {
             }
             ast::Stmt::For(f) => {
                 self.walk_expr(&f.iter, None);
+                let elem_ty = match plum_checker::infer_expr(&f.iter, &self.env, &self.cctx) {
+                    Ok(PlumType::TVariadic(elem)) => *elem,
+                    _ => PlumType::TInt,
+                };
                 for v in &f.vars {
-                    self.env.insert(v.clone(), TypeScheme::mono(PlumType::TInt));
+                    self.env.insert(v.clone(), TypeScheme::mono(elem_ty.clone()));
                     self.locals.insert(v.clone());
                 }
                 self.walk_block(&f.body);
@@ -1539,6 +1549,13 @@ struct Collector<'a> {
     /// / `all_subjects`, so this only covers depth >= 1.
     nested_class_scratch: HashMap<usize, u32>,
     next_nested_class_slot: u32,
+    /// `For` stmt identity (pointer address) -> a slot number; each slot reserves 2
+    /// consecutive `i32` scratch locals for variadic iteration (`for v in nums`):
+    /// [count, loop index]. Only `for` statements whose iterable is a `TVariadic`
+    /// use this — an ordinary range `for` reuses its own loop var as the counter
+    /// and needs no extra scratch locals.
+    variadic_for_scratch: HashMap<usize, u32>,
+    next_variadic_for_slot: u32,
 }
 impl<'a> Collector<'a> {
@@ -1636,8 +1653,18 @@ impl<'a> Collector<'a> {
             }
             ast::Stmt::For(f) => {
                 self.walk_expr(&f.iter);
+                let iter_ty = plum_checker::infer_expr(&f.iter, &self.env, &self.cctx).unwrap_or(PlumType::TInt);
+                if let PlumType::TVariadic(elem) = &iter_ty {
+                    let idx = self.next_variadic_for_slot;
+                    self.next_variadic_for_slot += 1;
+                    self.variadic_for_scratch.insert(f as *const ast::For as usize, idx);
-                for v in &f.vars {
+                    for v in &f.vars {
+                        self.bind(v, (**elem).clone());
+                    }
+                } else {
+                    for v in &f.vars {
-                    self.bind(v, PlumType::TInt);
+                        self.bind(v, PlumType::TInt);
+                    }
                 }
                 self.walk_block(&f.body);
             }
@@ -1776,6 +1803,8 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         next_closure_slot: 0,
         nested_class_scratch: HashMap::new(),
         next_nested_class_slot: 0,
+        variadic_for_scratch: HashMap::new(),
+        next_variadic_for_slot: 0,
     };
     if let ast::FnBody::Block(block) = &f.body {
         collector.walk_block(block);
@@ -1828,6 +1857,14 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         idx += 1;
     }
+    let variadic_for_scratch_base = idx;
+    let variadic_for_scratch_count = collector.variadic_for_scratch.values().copied().max().map(|m| m + 1).unwrap_or(0);
+    for _ in 0..variadic_for_scratch_count {
+        groups.push(ValType::I32); // count
+        groups.push(ValType::I32); // loop index
+        idx += 2;
+    }
     let closure_scratch_base = idx;
     let closure_scratch_count = collector.closure_scratch.values().copied().max().map(|m| m + 2).unwrap_or(0);
     for _ in 0..closure_scratch_count {
@@ -1854,6 +1891,8 @@ fn compile_fn_body(f: &ast::Fn, ctx: &CompileCtx, state: &mut ModuleState) -> Re
         closure_scratch: collector.closure_scratch,
         nested_class_scratch_base,
         nested_class_scratch: collector.nested_class_scratch,
+        variadic_for_scratch_base,
+        variadic_for_scratch: collector.variadic_for_scratch,
         func_ids: &ctx.func_ids,
         func_sigs: &ctx.func_sigs,
         closures: &ctx.closures,
@@ -2118,6 +2157,66 @@ fn compile_stmt(stmt: &ast::Stmt, body: &mut Vec<u8>, ctx: &LocalCtx, state: &mu
                     return Ok(());
                 }
             }
+            if let PlumType::TVariadic(elem_ty) = infer_local_type(&f.iter, ctx) {
+                if f.vars.len() != 1 {
+                    return Err("codegen: for-loop over a variadic param must bind exactly one variable".to_string());
+                }
+                let var_name = &f.vars[0];
+                let var_idx = ctx
+                    .locals
+                    .get(var_name)
+                    .copied()
+                    .ok_or_else(|| format!("undeclared loop var '{}'", var_name))?;
+                ctx.type_env.borrow_mut().insert(var_name.clone(), TypeScheme::mono((*elem_ty).clone()));
+                let scratch_key = f as *const ast::For as usize;
+                let slot = *ctx
+                    .variadic_for_scratch
+                    .get(&scratch_key)
+                    .ok_or_else(|| "internal codegen error: missing variadic-for scratch slot".to_string())?;
+                let count_local = ctx.variadic_for_scratch_base + slot * 2;
+                let index_local = count_local + 1;
+                let elem_vt = plum_type_to_valtype(&elem_ty);
+                // count_local = i32.wrap_i64(load_i64([nums + 0]))
+                compile_expr(&f.iter, body, ctx, state)?;
+                Instruction::I64Load(MemArg { offset: 0, align: 3, memory_index: 0 }).encode(body);
+                Instruction::I32WrapI64.encode(body);
+                Instruction::LocalSet(count_local).encode(body);
+                // index_local = 0
+                Instruction::I32Const(0).encode(body);
+                Instruction::LocalSet(index_local).encode(body);
+                Instruction::Block(BlockType::Empty).encode(body);
+                Instruction::Loop(BlockType::Empty).encode(body);
+                Instruction::LocalGet(index_local).encode(body);
+                Instruction::LocalGet(count_local).encode(body);
+                Instruction::I32GeS.encode(body);
+                Instruction::BrIf(1).encode(body);
+                // var = load_elem([nums + 8 + index * 8])
+                compile_expr(&f.iter, body, ctx, state)?;
+                Instruction::I32Const(8).encode(body);
+                Instruction::I32Add.encode(body);
+                Instruction::LocalGet(index_local).encode(body);
+                Instruction::I32Const(8).encode(body);
+                Instruction::I32Mul.encode(body);
+                Instruction::I32Add.encode(body);
+                emit_load(elem_vt, 0, body);
+                Instruction::LocalSet(var_idx).encode(body);
+                compile_block(&f.body, body, ctx, state)?;
+                Instruction::LocalGet(index_local).encode(body);
+                Instruction::I32Const(1).encode(body);
+                Instruction::I32Add.encode(body);
+                Instruction::LocalSet(index_local).encode(body);
+                Instruction::Br(0).encode(body);
+                Instruction::End.encode(body);
+                Instruction::End.encode(body);
+                return Ok(());
+            }
             compile_expr(&f.iter, body, ctx, state)?;
             Instruction::Drop.encode(body);
         }
@@ -3248,6 +3347,8 @@ fn compile_closure_body(
         next_closure_slot: 0,
         nested_class_scratch: HashMap::new(),
         next_nested_class_slot: 0,
+        variadic_for_scratch: HashMap::new(),
+        next_variadic_for_slot: 0,
     };
     collector.walk_block(&cl.body);
@@ -3300,6 +3401,14 @@ fn compile_closure_body(
         idx += 1;
     }
+    let variadic_for_scratch_base = idx;
+    let variadic_for_scratch_count = collector.variadic_for_scratch.values().copied().max().map(|m| m + 1).unwrap_or(0);
+    for _ in 0..variadic_for_scratch_count {
+        groups.push(ValType::I32); // count
+        groups.push(ValType::I32); // loop index
+        idx += 2;
+    }
     let closure_scratch_base = idx;
     let closure_scratch_count = collector.closure_scratch.values().copied().max().map(|m| m + 2).unwrap_or(0);
     for _ in 0..closure_scratch_count {
@@ -3334,6 +3443,8 @@ fn compile_closure_body(
         closure_scratch: collector.closure_scratch,
         nested_class_scratch_base,
         nested_class_scratch: collector.nested_class_scratch,
+        variadic_for_scratch_base,
+        variadic_for_scratch: collector.variadic_for_scratch,
         func_ids: &ctx.func_ids,
         func_sigs: &ctx.func_sigs,
         closures: &ctx.closures,

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

@@ -1443,3 +1443,37 @@ main() -> Int =
     let bytes = compile_source(&source).expect("compile failed");
     assert_eq!(run_main(&bytes), 60);
 }
+#[test]
+fn sum_all_variadic_int_runs_correctly() {
+    let src = "\
+sumAll(nums: ...Int) -> Int =
+  total = 0
+  for v in nums
+    total = total + v
+  total
+main() -> Int =
+  sumAll(1, 2, 3, 4)
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), 10);
+}
+#[test]
+fn sum_all_variadic_int_with_zero_args_runs_correctly() {
+    let src = "\
+sumAll(nums: ...Int) -> Int =
+  total = 0
+  for v in nums
+    total = total + v
+  total
+main() -> Int =
+  sumAll()
+";
+    let source = parse(src);
+    let bytes = compile_source(&source).expect("compile failed");
+    assert_eq!(run_main(&bytes), 0);
+}