fix(typecheck): infer recursive functions' return type by fixpoint (#88)

compiler/lib/typecheck.ml

@@ -652,6 +652,109 @@ and check_pattern (gamma : env) (scrutinee_ty : ty) (pat : pattern)
 let valid_invariants =
   ["jones"; "alexander"; "homfly"; "kauffman"; "writhe"; "linking"]
 
+(** Does [e] contain a call to [f]?  Used to detect self-recursion in a
+    function definition, so that only recursive definitions pay the cost of the
+    return-type search below. *)
+let rec expr_calls (f : string) (e : expr) : bool =
+  let go = expr_calls f in
+  match e with
+  | Call (g, args) -> g = f || List.exists go args
+  | Match (scrut, arms) ->
+    go scrut || List.exists (fun a -> go a.arm_body) arms
+  | Let (_, e1, e2)      -> go e1 || go e2
+  | Pipeline (e1, e2)    -> go e1 || go e2
+  | BinOp (_, e1, e2)    -> go e1 || go e2
+  | Cap (e1, e2) | Cup (e1, e2) | Pair (e1, e2)
+  | EchoAdd (e1, e2) | EchoEq (e1, e2) -> go e1 || go e2
+  | UnaryOp (_, e1) | Close e1 | Mirror e1 | Reverse e1 | Simplify e1
+  | Twist e1 | EchoClose e1 | Lower e1 | Residue e1 | Fst e1 | Snd e1 -> go e1
+  | BraidLit _ | Identity | BoolLit _ | IntLit _ | FloatLit _
+  | StringLit _ | Var _ | Crossing _ -> false
+
+(** Candidate return types for a recursive function, in the order they are
+    tried.  Seeds drawn from the body come first (a match's non-recursive arms
+    are the usual source), then the scalar types, then the original [TWord 0]
+    placeholder so the previous behaviour remains reachable. *)
+let recursive_return_candidates (gamma : env) (f : string) (body : expr) : ty list =
+  (* Types of the arms that do NOT mention [f] — those can be inferred before
+     anything is known about the function's return type.  For
+       def length(w) = match w with
+         | identity  => 0                  <- typeable now: Num
+         | s1 . rest => 1 + length(rest)   <- not yet
+         | _         => 0                  <- typeable now: Num
+     this yields [TNum], which is the fixpoint. *)
+  let seeds_from_body =
+    match body with
+    | Match (_, arms) ->
+      List.filter_map (fun a ->
+        if expr_calls f a.arm_body then None
+        else (try Some (infer_expr gamma [] a.arm_body) with _ -> None)
+      ) arms
+    | _ -> []
+  in
+  let dedup l =
+    List.fold_left (fun acc x -> if List.mem x acc then acc else acc @ [x]) [] l
+  in
+  dedup (seeds_from_body @ [TNum; TBool; TStr; TWord 0])
+
+(** Infer the return type of a function definition and bind it in [gamma].
+    [T-Def-Fun].
+
+    This is the single implementation of function-definition typing.  It used
+    to exist twice — once here and once inlined in [check_program]'s pass 1b —
+    and only the copy in [check_program] was ever reached for whole programs.
+    A fix applied to one had no effect on the other, so they are now one
+    function called from both. *)
+let bind_function_def (gamma : env) (def : definition) : env =
+  let param_tys = List.map (fun _p -> TWord 0) def.def_params in
+  (* Bind the params, and the function itself so recursive calls resolve.
+     [seed] is the return type assumed for those recursive calls. *)
+  let with_seed (seed : ty) : env =
+    let fsig = { fsig_params = param_tys; fsig_return = seed } in
+    let g = env_bind_fun gamma def.def_name fsig in
+    List.fold_left2 (fun g pname pty -> env_bind_val g pname pty)
+      g def.def_params param_tys
+  in
+  let body_ty =
+    if not (expr_calls def.def_name def.def_body) then
+      (* Non-recursive: the seed is never consulted, so one pass suffices.
+         Unchanged from the original behaviour. *)
+      infer_expr (with_seed (TWord 0)) [] def.def_body
+    else begin
+      (* Recursive.  The return type occurs in its own derivation, so it must
+         be a FIXPOINT: assume a return type, check the body under that
+         assumption, and accept only if the body then has exactly the assumed
+         type.  Anything weaker is a guess that merely failed to crash.
+
+         Previously the seed was hard-coded to [TWord 0] and marked
+         "placeholder", with nothing checking the result against it.  So
+
+             def length(w) = match w with
+               | identity  => 0
+               | s1 . rest => 1 + length(rest)
+               | _         => 0
+
+         failed with "Cannot add Num and Word[0]" — the recursive call was
+         typed at the placeholder rather than at Num, which made every
+         recursive function returning a scalar unwritable.
+
+         The candidate list is finite and ordered, so this terminates.  If no
+         candidate is a fixpoint we fall back to the original single pass,
+         which re-raises its original error: a definition that did not
+         typecheck before still does not, with the same message. *)
+      let rec first_fixpoint = function
+        | [] -> infer_expr (with_seed (TWord 0)) [] def.def_body
+        | seed :: rest ->
+          (match infer_expr (with_seed seed) [] def.def_body with
+           | t when t = seed -> t            (* consistent: a real fixpoint *)
+           | _               -> first_fixpoint rest
+           | exception _     -> first_fixpoint rest)
+      in
+      first_fixpoint (recursive_return_candidates gamma def.def_name def.def_body)
+    end
+  in
+  env_bind_fun gamma def.def_name { fsig_params = param_tys; fsig_return = body_ty }
+
 (** Type-check a single statement, returning the (possibly extended) environment.
  *  Implements [T-Def-Val], [T-Def-Fun], [T-Assert], [T-Compute], [T-Weave].
  *)
@@ -671,19 +774,7 @@ let check_statement (gamma : env) (stmt : statement) : env =
        * A more sophisticated implementation would use constraint-based
        * inference; here we use a simple forward analysis.
        *)
-      let param_tys = List.map (fun _p -> TWord 0) def.def_params in
-      (* Bind params and the function itself (for recursion) into the env *)
-      let ret_ty = TWord 0 in  (* placeholder *)
-      let fsig = { fsig_params = param_tys; fsig_return = ret_ty } in
-      let gamma' = env_bind_fun gamma def.def_name fsig in
-      let gamma' = List.fold_left2 (fun g pname pty ->
-        env_bind_val g pname pty
-      ) gamma' def.def_params param_tys in
-      (* Infer the body type *)
-      let body_ty = infer_expr gamma' [] def.def_body in
-      (* Re-register with inferred return type *)
-      let fsig' = { fsig_params = param_tys; fsig_return = body_ty } in
-      env_bind_fun gamma def.def_name fsig'
+      bind_function_def gamma def
     end
 
   (* [T-Weave] (section 3.10) *)
@@ -813,16 +904,7 @@ let check_program (prog : program) : check_result =
           let ty = infer_expr gamma [] def.def_body in
           env_bind_val gamma def.def_name ty
         end else begin
-          let param_tys = List.map (fun _p -> TWord 0) def.def_params in
-          let fsig_placeholder =
-            { fsig_params = param_tys; fsig_return = TWord 0 } in
-          let gamma' = env_bind_fun gamma def.def_name fsig_placeholder in
-          let gamma' = List.fold_left2 (fun g pname pty ->
-            env_bind_val g pname pty
-          ) gamma' def.def_params param_tys in
-          let body_ty = infer_expr gamma' [] def.def_body in
-          let fsig = { fsig_params = param_tys; fsig_return = body_ty } in
-          env_bind_fun gamma def.def_name fsig
+          bind_function_def gamma def
         end
       with Type_error msg ->
         add_error_at def.def_line

compiler/test/test_typecheck.ml

@@ -623,6 +623,81 @@ let test_functions () =
     let r = check_ok prog in
     r.result_ok);
 
+  (* ---------------------------------------------------------------- *)
+  (* #88: recursive functions returning a SCALAR.                       *)
+  (* The return type occurs in its own derivation, so it must be a      *)
+  (* fixpoint. Previously the recursive call was typed at the hard-coded *)
+  (* [TWord 0] "placeholder" seed with nothing checking the result       *)
+  (* against it, so every one of these failed with                       *)
+  (*   "Cannot add Num and Word[0]"                                      *)
+  (* and recursive scalar functions were simply unwritable.              *)
+  (* ---------------------------------------------------------------- *)
+
+  test "#88 recursive fn returning Num (the examples/ `length` shape)" (fun () ->
+    (* def length(w) = match w with
+         | identity  => 0
+         | s1 . rest => 1 + length(rest)
+         | _         => 0 *)
+    let prog = [
+      def_fun "length" ["w"] (Match (Var "w", [
+        { arm_pattern = PatIdentity; arm_body = IntLit 0 };
+        { arm_pattern = PatCons ({ gpat_index = 1; gpat_exponent = 1 }, PatVar "rest");
+          arm_body = BinOp (Add, IntLit 1, Call ("length", [Var "rest"])) };
+        { arm_pattern = PatWildcard; arm_body = IntLit 0 };
+      ]))
+    ] in
+    let r = check_ok prog in
+    r.result_ok);
+
+  test "#88 recursive fn's return type is Num, not the seed" (fun () ->
+    (* Guard: it must not merely typecheck — the signature must be right,
+       otherwise callers would still see Word[0]. *)
+    let prog = [
+      def_fun "length" ["w"] (Match (Var "w", [
+        { arm_pattern = PatIdentity; arm_body = IntLit 0 };
+        { arm_pattern = PatCons ({ gpat_index = 1; gpat_exponent = 1 }, PatVar "rest");
+          arm_body = BinOp (Add, IntLit 1, Call ("length", [Var "rest"])) };
+        { arm_pattern = PatWildcard; arm_body = IntLit 0 };
+      ]))
+    ] in
+    let r = check_ok prog in
+    (match env_lookup r.result_env "length" with
+     | Some (EFun s) -> s.fsig_return = TNum
+     | _ -> false));
+
+  test "#88 recursive fn result is usable at Num by a caller" (fun () ->
+    let prog = [
+      def_fun "length" ["w"] (Match (Var "w", [
+        { arm_pattern = PatIdentity; arm_body = IntLit 0 };
+        { arm_pattern = PatCons ({ gpat_index = 1; gpat_exponent = 1 }, PatVar "rest");
+          arm_body = BinOp (Add, IntLit 1, Call ("length", [Var "rest"])) };
+        { arm_pattern = PatWildcard; arm_body = IntLit 0 };
+      ]));
+      (* Only well-typed if `length` really returns Num. *)
+      def_val "total" (BinOp (Add, Call ("length", [BraidLit [sigma 1]]), IntLit 1));
+    ] in
+    let r = check_ok prog in
+    r.result_ok);
+
+  test "#88 non-recursive definitions are unaffected" (fun () ->
+    (* Regression guard: the fixpoint search must only engage for genuinely
+       recursive bodies. A word-returning function still returns a word. *)
+    let prog = [def_fun "twice" ["x"] (BinOp (Compose, Var "x", Var "x"))] in
+    let r = check_ok prog in
+    r.result_ok &&
+    (match env_lookup r.result_env "twice" with
+     | Some (EFun s) -> s.fsig_return = TWord 0
+     | _ -> false));
+
+  test "#88 a genuinely ill-typed recursive fn still fails" (fun () ->
+    (* The fallback must not turn errors into silent successes: no candidate
+       return type makes `identity + 1` well-formed. *)
+    let prog = [
+      def_fun "bad" ["w"] (BinOp (Add, Identity, Call ("bad", [Var "w"])))
+    ] in
+    let r = check_ok prog in
+    not r.result_ok);
+
   (* [T-App]: function application *)
   test "T-App" (fun () ->
     let prog = [