a3ab768 | plum | pyrossh

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

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

a3ab768 — Peter John 2026-10-25T16:25:50+05:30

fix(monomorphize): resolve a closure's own generic from cb's declared return

Files changed (3) hide show

plum-checker/src/monomorphize.rs +69 -1
plum-checker/tests/checker_tests.rs +27 -0
plum-std/Try.plum +8 -0

plum-checker/src/monomorphize.rs CHANGED Viewed

@@ -643,6 +643,14 @@ struct Monomorphizer<'a> {
     /// the top of every `rewriteFnBody` call; never needs saving/restoring since
     /// closures aren't recursed into by this pass (see `rewriteExpr`'s `Closure` arm).
     current_return_type: Option<ast::Type>,
+    /// True while `current_return_type` holds a closure argument's declared
+    /// return (see `rewriteMethodClosureArgs`), which may still contain the
+    /// method's own unresolved generic letters.
+    return_type_is_closure_expected: bool,
+    /// The enclosing real function's return type while
+    /// `return_type_is_closure_expected` is set (fallback for letters the
+    /// closure's expected return leaves unresolved).
+    closure_outer_return_type: Option<ast::Type>,
     /// (receiver name or `None` for a free function, function/method name) ->
     /// (the generic enum its declared return type names, the concrete type args
     /// it names them with) — for every function/method whose OWN declared return
@@ -772,6 +780,7 @@ impl<'a> Monomorphizer<'a> {
     /// the body's construction site has just been specialized to a mangled name.
     fn rewriteFnBody(&mut self, f: &mut ast::Fn, resolve_return: bool) -> Result<(), String> {
         self.current_return_type = f.returns.clone();
+        self.return_type_is_closure_expected = false;
         self.local_types_by_name.clear();
         let mut env = self.global_env.clone();
         if let Some(recv) = &f.type_param {
@@ -1329,6 +1338,23 @@ impl<'a> Monomorphizer<'a> {
                 },
                 None => Vec::new(),
             };
+            // The closure's own declared return, with the receiver's bindings
+            // substituted (`andThen`'s `cb: fn(T) -> Try[U, E]` on a
+            // `Try[Int, [Bad]]` receiver -> `Try[U, [Bad]]`). A bare `Ok(n * 2)`
+            // in the closure body can only pin `T`, never `E`; without this it
+            // fell back to the ENCLOSING function's return type (wrong scope)
+            // and stayed unmangled, so `resolveMethodOwnGenerics` couldn't
+            // recover `U`. Only a non-bare-letter return is used (a bare
+            // `fn(T) -> U` says nothing more than the outer fallback did).
+            let closure_expected_ret: Option<ast::Type> = match &method_and_bindings {
+                Some((method, class_bindings)) => match method.params.get(i).map(|p| &p.ty) {
+                    Some(ast::ParamType::Fn(_, Some(ret))) if !ret.generics.is_empty() => {
+                        Some(substituteType(ret, &Substitution(class_bindings.clone())))
+                    }
+                    _ => None,
+                },
+                None => None,
+            };
             let mut closure_env = env.clone();
             for (p, ty) in cl.params.iter().zip(closure_param_tys.iter()) {
                 closure_env.insert(p.clone(), TypeScheme::mono(ty.clone()));
@@ -1338,7 +1364,23 @@ impl<'a> Monomorphizer<'a> {
                     .entry(p.clone())
                     .or_insert_with(|| TypeScheme::mono(PlumType::TVar("_".to_string())));
             }
+            if let Some(expected) = closure_expected_ret {
+                let saved_ret = self.current_return_type.replace(expected);
+                let saved_flag = std::mem::replace(&mut self.return_type_is_closure_expected, true);
+                // Nested closures keep the outermost real function's return.
+                let saved_outer = if saved_flag {
+                    self.closure_outer_return_type.clone()
+                } else {
+                    std::mem::replace(&mut self.closure_outer_return_type, saved_ret.clone())
+                };
+                let result = self.rewriteBlock(&mut cl.body, &mut closure_env);
+                self.current_return_type = saved_ret;
+                self.return_type_is_closure_expected = saved_flag;
+                self.closure_outer_return_type = saved_outer;
+                result?;
+            } else {
-            self.rewriteBlock(&mut cl.body, &mut closure_env)?;
+                self.rewriteBlock(&mut cl.body, &mut closure_env)?;
+            }
         }
         Ok(())
     }
@@ -2036,6 +2078,30 @@ impl<'a> Monomorphizer<'a> {
             if let Some(rt) = self.current_return_type.clone() {
                 if rt.name == e.name && rt.generics.len() == params.len() {
                     for (p, gt) in params.iter().zip(rt.generics.iter()) {
+                        // A closure's expected return can still name the
+                        // method's own not-yet-inferred letter (`U` in
+                        // `Try[U, [Bad]]`) — that's what this site is about
+                        // to determine, not a binding to copy.
+                        // Fall back to the enclosing function's return (the
+                        // pre-closure-expected behavior), so `return
+                        // x.andThen(|n| Err(Bad))` in a `-> Try[Int, [Bad]]`
+                        // function still pins `T`.
+                        if self.return_type_is_closure_expected
+                            && gt.generics.is_empty()
+                            && isGenericParamName(&gt.name)
+                        {
+                            if bindings.contains_key(p) {
+                                continue;
+                            }
+                            if let Some(outer) = self.closure_outer_return_type.clone() {
+                                if outer.name == e.name && outer.generics.len() == params.len() {
+                                    let og = &outer.generics[params.iter().position(|q| q == p).unwrap()];
+                                    let resolved = self.resolveFieldType(og);
+                                    bindings.insert(p.clone(), crate::plumTypeFromAst(&resolved));
+                                }
+                            }
+                            continue;
+                        }
                         if !bindings.contains_key(p) {
                             // See the matching comment in
                             // `resolveClassInstantiation` — `gt` needs
@@ -2686,6 +2752,8 @@ pub fn monomorphizeSource(source: &ast::Source) -> Result<ast::Source, String> {
         worklist: Vec::new(),
         produced: Vec::new(),
         current_return_type: None,
+        return_type_is_closure_expected: false,
+        closure_outer_return_type: None,
         fn_return_generic_enum: BTreeMap::new(),
         local_types_by_name: BTreeMap::new(),
         rename_counter: 0,

plum-checker/tests/checker_tests.rs CHANGED Viewed

@@ -1680,3 +1680,30 @@ fun f() -> Try[Int, [Red]] =
     );
     assert!(bad.is_err());
 }
+#[test]
+fn andThenClosureReturningBareOkResolvesOwnGeneric() {
+    // `cb: fn(T) -> Try[U, E]`: a bare `Ok(n * 2)` only pins `T`; `E` must
+    // come from the receiver via the closure's declared return, or `U` is lost.
+    let src = "\
+enum Try[T, E] =
+  | Ok(T)
+  | Err(E)
+  fun andThen(self, cb: fn(T) -> Try[U, E]) -> Try[U, E] =
+    match self
+      Ok(v) =>
+        return cb(v)
+      Err(e) =>
+        return Err(e)
+fun five() -> Try[Int, [Bad]] =
+  return Ok(5)
+fun main() =
+  r := five().andThen(|n| Ok(n * 2))
+";
+    let source = parseClean(src);
+    let mono = plum_checker::monomorphize::monomorphizeSource(&source);
+    assert!(mono.is_ok(), "expected Ok, got {:?}", mono.err());
+}

plum-std/Try.plum CHANGED Viewed

@@ -104,6 +104,9 @@ fun mergeForTryTest(s: Str, i: Int) -> Try[Int, [BadNumStr, OutOfBounds]] =
   x := xs.get(i)?
   return Ok(n + x)
+fun andThenBareErrForTryTest() -> Try[Int, [Bad]] =
+  return okIntForTryTest().andThen(|n| Err(Bad))
 fun closureMapForTryTest(s: Str) -> Try[Int, [InvalidFirst]] =
   n := parseInt(s) ? |e| InvalidFirst
   return Ok(n)
@@ -121,6 +124,11 @@ test "andThen chains fallible steps"
   assert a.unwrap() == 1
   assert errIntForTryTest().andThen(|n| fifthForTryTest(n)).isErr()
+test "andThen with closure returning bare Ok"
+  assert okIntForTryTest().andThen(|n| Ok(n * 2)).unwrap() == 10
+  assert errIntForTryTest().andThen(|n| Ok(n * 2)).isErr()
+  assert andThenBareErrForTryTest().isErr()
 test "unwrapOr and unwrapOrElse fall back on Err"
   assert errIntForTryTest().unwrapOr(9) == 9
   assert okIntForTryTest().unwrapOr(9) == 5