An if's else arm is checked once and brought to the join, an abandoned check undoes only what it wrote, and an arm refused at the join says why at its value

lib/check.ml

@@ -34,6 +34,12 @@ let fail = Loc.fail
    language from the one the author decided on. *)
 let literal_at_want = "check/literal-at-want"
 
+(* A refusal that is two types failing to meet, a literal's included: what
+   an arm checked at another arm's type says when the two simply differ. *)
+let is_mismatch (d : Loc.diag) =
+  String.equal d.Loc.kind "check/type-mismatch"
+  || String.equal d.Loc.kind literal_at_want
+
 (* [List.map]'s evaluation order is unspecified, and checking allocates frame
    slots as a side effect. Left-to-right is required, not a preference: a later
    let binding sees an earlier one, and slot numbering must be reproducible. *)
@@ -121,6 +127,35 @@ type gstruct = {
    program asks for the copy itself. *)
 let struct_apps : (string, string * Types.t list) Hashtbl.t = Hashtbl.create 16
 
+(* The undo journal a check that may be abandoned writes into: every table
+   write a body's check makes goes through [jreplace]/[jremove]/[jset], which
+   note how to take it back while a [snapshot_env] is open. So abandoning a
+   check costs what it wrote, not the size of the tables it could have. *)
+let journal : (unit -> unit) list ref = ref []
+let journal_open = ref 0
+
+let jot undo = if !journal_open > 0 then journal := undo :: !journal
+
+let jreplace tbl k v =
+  (if !journal_open > 0 then
+     let old = Hashtbl.find_opt tbl k in
+     jot (fun () ->
+         match old with
+         | Some o -> Hashtbl.replace tbl k o
+         | None -> Hashtbl.remove tbl k));
+  Hashtbl.replace tbl k v
+
+let jremove tbl k =
+  (if !journal_open > 0 then
+     match Hashtbl.find_opt tbl k with
+     | Some o -> jot (fun () -> Hashtbl.replace tbl k o)
+     | None -> ());
+  Hashtbl.remove tbl k
+
+let jset r v =
+  (if !journal_open > 0 then let old = !r in jot (fun () -> r := old));
+  r := v
+
 (* The length every length variable has inside a generic body's abstract
    pass. Large so that no constant index into such an array is refused as out
    of bounds there, and within i32 so that [(length a)] is an ordinary index.
@@ -364,25 +399,25 @@ let new_env () = {
   guard_next = false;
 }
 
-(* Everything checking may write into [env], taken so that a check which is
+(* A check that may be abandoned — a trial, a probe, a return type read and
-   abandoned — a trial, a probe, a return type read and thrown away, a
+   thrown away, a tolerated body — opens one of these: [undo] puts back
-   tolerated body — can be undone as one piece. Partial undo is the bug this
+   everything it wrote into [env], [keep] closes it and leaves the writes,
+   which an enclosing one can still undo. Partial undo is the bug this
    exists for: rewinding [lifted] and not the generic cache left a copy made
    during a trial with the lambda it lifted gone, which is a link error. So
    the whole record is named, closed with warning 9: a new field stops this
    compiling until it is decided here.
 
-   Taken on every trial, so it copies only what a body's check writes: the
+   The tables a body's check writes go through the journal ([jreplace]);
-   struct table and its locations (a generic struct's copy, a closure's
+   the lists and flags are held here, which costs nothing; the rest is
-   environment), the struct-copy tables, and the generic cache, whose copies'
+   written by the declaration passes alone, before any body is checked. *)
-   entries in [fns] are the only ones a body adds. The rest is written by the
+let snapshot_env env : (unit -> unit) * (unit -> unit) =
-   declaration passes alone, before any body is checked, and is marked so. *)
+  let[@warning "+9"] { lifted; instances; tyvars; subst; tvpreds; chain;
-let snapshot_env env : unit -> unit =
+                       deferred; lenvars; len_placeholder; schain; in_field;
-  let[@warning "+9"] { structs; locs; copies; insts; lifted; instances;
+                       recovering; recovered; poison; speculating;
-                       tyvars; subst; tvpreds; chain; deferred; lenvars;
+                       guard_next;
-                       len_placeholder; schain; in_field; recovering;
+                       (* Journaled at their writes. *)
-                       recovered; poison; speculating; guard_next;
+                       structs = _; locs = _; copies = _; insts = _; fns = _;
-                       fns = _ (* through [insts], below *);
                        (* Declaration passes only. *)
                        datas = _; unions = _; cases = _; aliases = _;
                        consts = _; enums = _; parents = _; externs = _;
@@ -391,29 +426,22 @@ let snapshot_env env : unit -> unit =
                        generics = _; gsigs = _; refused_generics = _;
                        gstructs = _; broken = _; glens = _; classes = _;
                        tracks = _; inferred = _; infer_failed = _ } = env in
-  let keep t =
+  incr journal_open;
-    let c = Hashtbl.copy t in
+  let mark = !journal in
-    fun () -> Hashtbl.reset t; Hashtbl.iter (Hashtbl.add t) c
+  let close () =
+    decr journal_open;
+    if !journal_open = 0 then journal := []
   in
-  let tables = [ keep structs; keep locs; keep copies; keep struct_apps ] in
+  let undo () =
-  let cache = Hashtbl.fold (fun g r acc -> (g, r, !r) :: acc) insts [] in
+    let rec back l =
-  fun () ->
+      if l != mark then
-    List.iter (fun undo -> undo ()) tables;
+        match l with
-    (* A copy made since goes from [fns] with its cache entry. *)
+        | u :: rest -> u (); back rest
-    Hashtbl.filter_map_inplace
+        | [] -> ()
-      (fun g r ->
+    in
-         match List.find_opt (fun (h, _, _) -> String.equal g h) cache with
+    back !journal;
-         | None ->
+    journal := mark;
-           List.iter (fun (_, _, sym) -> Hashtbl.remove env.fns sym) !r;
+    close ();
-           None
-         | Some (_, _, before) ->
-           List.iter
-             (fun ((_, _, sym) as e) ->
-                if not (List.memq e before) then Hashtbl.remove env.fns sym)
-             !r;
-           r := before;
-           Some r)
-      insts;
     env.lifted <- lifted; env.instances <- instances; env.tyvars <- tyvars;
     env.subst <- subst; env.tvpreds <- tvpreds; env.chain <- chain;
     env.deferred <- deferred; env.lenvars <- lenvars;
@@ -421,6 +449,8 @@ let snapshot_env env : unit -> unit =
     env.in_field <- in_field; env.recovering <- recovering;
     env.recovered <- recovered; env.poison <- poison;
     env.speculating <- speculating; env.guard_next <- guard_next
+  in
+  (undo, close)
 
 (* A refusal [collect] can go on past: kept while a whole-file check is
    collecting, in the order found, and raised otherwise. *)
@@ -1469,7 +1499,7 @@ let struct_app g args =
            args)
   in
   if not (Hashtbl.mem struct_apps key) then begin
-    Hashtbl.replace struct_apps key (g, args);
+    jreplace struct_apps key (g, args);
     Hashtbl.replace Types.display key
       (Printf.sprintf "(%s %s)" g
          (String.concat " " (List.map Types.to_string args)))
@@ -1710,8 +1740,8 @@ and struct_copy ?(at_definition = false) env loc name targs =
   let key = struct_app name targs in
   if Hashtbl.mem env.copies key then key
   else if Hashtbl.mem env.broken name then begin
-    Hashtbl.replace env.copies key (List.exists generic_arg targs);
+    jreplace env.copies key (List.exists generic_arg targs);
-    Hashtbl.replace env.structs key { Tast.sname = key; fields = [] };
+    jreplace env.structs key { Tast.sname = key; fields = [] };
     key
   end
   else begin
@@ -1745,9 +1775,9 @@ and struct_copy ?(at_definition = false) env loc name targs =
     let generic = List.exists generic_arg targs in
     (* In before its fields, so a field that names the same copy through a
        pointer — [(defstruct Node [next (Ptr (Node $t))])] — finds it. *)
-    Hashtbl.replace env.copies key generic;
+    jreplace env.copies key generic;
-    Hashtbl.replace env.structs key { Tast.sname = key; fields = [] };
+    jreplace env.structs key { Tast.sname = key; fields = [] };
-    Hashtbl.replace env.locs key g.gloc;
+    jreplace env.locs key g.gloc;
     let saved =
       (env.subst, env.tyvars, env.lenvars, env.tvpreds, env.len_placeholder,
        env.in_field, env.schain)
@@ -1775,13 +1805,13 @@ and struct_copy ?(at_definition = false) env loc name targs =
     with
     | fields ->
       restore ();
-      Hashtbl.replace env.structs key { Tast.sname = key; fields };
+      jreplace env.structs key { Tast.sname = key; fields };
       finite_from env key;
       key
     | exception e ->
       restore ();
-      Hashtbl.remove env.copies key;
+      jremove env.copies key;
-      Hashtbl.remove env.structs key;
+      jremove env.structs key;
       (* A field refused inside the template says nothing about which use
          asked for this copy; the note names it, one per level of copies. *)
       (match e with
@@ -3011,7 +3041,7 @@ let close_over ~fname (octx : ctx) (fctx : ctx) loc =
         (fun (n, ((b : binding), _)) -> { Tast.fname = n; fty = b.bty })
         caught
     in
-    Hashtbl.replace fctx.env.structs ename { Tast.sname = ename; fields };
+    jreplace fctx.env.structs ename { Tast.sname = ename; fields };
     let ety = Types.Named ename in
     let eslot = fresh_slot fctx (Types.Ptr (Types.Mut, ety)) in
     let binds =
@@ -7367,47 +7397,42 @@ and check_if_once ctx ~tail ?want loc c t e =
        is checked once; a refused one re-checks each level below the refusal
        once more, the square of its depth. *)
     (* The else arm on its own terms, when nothing is wanted: the two arms
-       meet at [arm_join], so the order they are written in decides nothing.
+       meet at [arm_join], so the order they are written in decides nothing,
-       Kept when it has the then arm's type, or when the then arm is the one
+       and both are brought to the join as checked — the arm is never checked
-       that moves; when the else arm moves it is checked again below at the
+       twice, which in a chain of ifs would be twice per level. *)
-       then arm's type, so its literals are typed there as before. *)
     let joined =
       if want <> None || free_join || t.Tast.ty = Types.Never
          || t.Tast.ty = Types.Bool || and_sentinel e || lone_literal e
       then None
       else
-        (* A trial not kept leaves nothing: [trial] puts the context and
-           the environment back when it raises. *)
-        let not_kept = "check/arm-not-kept" in
         let alone () = branch ctx (fun () -> in_tail (fun () -> check ctx e)) in
-        match
+        match trial ctx alone with
-          trial ctx (fun () ->
+        | Ok v ->
-              let v = alone () in
+          (match arm_join t.Tast.ty v.Tast.ty with
-              match arm_join t.Tast.ty v.Tast.ty with
+           | Some j -> Some (j, expect ctx v.Tast.loc ~want:(Some j) v)
-              | Some j when Types.equal v.Tast.ty j -> v
+           (* No join: the refusal [if] gives its else arm. *)
-              | _ -> raise (Loc.Error (Loc.diag ~kind:not_kept v.Tast.loc "")))
+           | None ->
-        with
+             Some (t.Tast.ty, expect ctx v.Tast.loc ~want:(Some t.Tast.ty) v))
-        | Ok v -> Some (v.Tast.ty, v)
+        | Error own ->
-        | Error d when String.equal d.Loc.kind not_kept -> None
+          (* Refused on its own terms. The then arm's type may be what it
-        | Error _ ->
+             needed — [nil], a bare struct — and it is checked at it below,
-          (* Refused on its own terms. If the then arm's type is what it
+             whose refusal is then the one said. Only when that refusal is a
-             needed — [nil], a bare struct — the arm is checked at it below;
+             mismatch and the arm's own is not — an unknown name, say — is
-             if it is refused there too, its own error is the real one, so it
+             the arm's own error the real one, so nothing is invented about a
-             is checked for real on its own terms and nothing is invented
+             type it was never going to have. *)
-             about a type it was never going to have. *)
           (match
              trial ctx (fun () ->
                  branch ctx (fun () ->
                      in_tail (fun () -> check ctx ~want:t.Tast.ty e)))
            with
-           | Ok _ -> None
+           | Error d
-           | Error _ ->
+             when is_mismatch d && not (is_mismatch own) ->
              let v = alone () in
              (match arm_join t.Tast.ty v.Tast.ty with
               | Some j -> Some (j, expect ctx v.Tast.loc ~want:(Some j) v)
               | None ->
-                fail loc "the branches of this if have different types: %s and %s"
+                Some (t.Tast.ty, expect ctx v.Tast.loc ~want:(Some t.Tast.ty) v))
-                  (Types.to_string t.Tast.ty) (Types.to_string v.Tast.ty)))
+           | _ -> None)
     in
     match joined with
     | Some (j, v) ->
@@ -8878,12 +8903,16 @@ and check_match ctx ?(tail = false) ?want loc scrutinee arms =
                let at w () = block ctx ?want:w a.Ast.aloc a.Ast.body in
                match trial ctx (at None) with
                | Ok b -> b
-               | Error _ ->
+               | Error own ->
-                 (* Refused on its own terms: the join is what it needed,
+                 (* Refused on its own terms: the join may be what it
-                    or, refused there too, its own error is the real one. *)
+                    needed, and its refusal at the join is then the one
+                    said — unless that refusal is a mismatch and the arm's
+                    own is not, when the arm's own error is the real one. *)
                  (match trial ctx (at !want) with
-                  | Ok _ -> at !want ()
+                  | Error d
-                  | Error _ -> at None ())
+                    when is_mismatch d && not (is_mismatch own) ->
+                    at None ()
+                  | _ -> at !want ())
              else block ctx ?want:!want a.Ast.aloc a.Ast.body
            in
            (if body.Tast.ty <> Types.Never then
@@ -8900,11 +8929,32 @@ and check_match ctx ?(tail = false) ?want loc scrutinee arms =
   let checked =
     match free, !want with
     | true, Some j ->
+      (* Said at the arm's value, its last form, as a refusal checked at the
+         join would have been — not at its pattern. *)
+      let rec value (x : Ast.expr) =
+        match x.Ast.e with
+        | Ast.Do (_ :: _ as xs) | Ast.Let (_, (_ :: _ as xs)) ->
+          value (List.hd (List.rev xs))
+        | _ -> x.Ast.loc
+      in
+      let value_loc i =
+        let (a : Ast.arm), _, _ = List.nth resolved i in
+        match List.rev a.Ast.body with
+        | x :: _ -> value x
+        | [] -> a.Ast.aloc
+      in
+      (* Each arm refused on its own, so every arm that cannot meet the join
+         is said, as each would be checked at it. *)
       List.map
         (fun (i, (arm : Tast.arm)) ->
            match arm.Tast.abody with
            | [ b ] when not (Types.equal b.Tast.ty j || b.Tast.ty = Types.Never) ->
-             (i, { arm with Tast.abody = [ expect ctx b.Tast.loc ~want:(Some j) b ] })
+             let at = value_loc i in
+             let b =
+               try expect ctx at ~want:(Some j) b
+               with Loc.Error d -> refuse_or_poison ctx.env at d
+             in
+             (i, { arm with Tast.abody = [ b ] })
            | _ -> (i, arm))
         checked
     | _ -> checked
@@ -13346,7 +13396,7 @@ and instantiate env loc gname vars subst cparams cret =
   let cache =
     match Hashtbl.find_opt env.insts gname with
     | Some r -> r
-    | None -> let r = ref [] in Hashtbl.replace env.insts gname r; r
+    | None -> let r = ref [] in jreplace env.insts gname r; r
   in
   let same (ps, r, _) =
     List.length ps = List.length cparams
@@ -13392,8 +13442,8 @@ and instantiate env loc gname vars subst cparams cret =
     (* The entry goes in *before* the body is checked, which is what makes a
        recursive generic function terminate: the call to itself at the same
        types finds this and does not generate a second copy. *)
-    cache := (cparams, cret, sym) :: !cache;
+    jset cache ((cparams, cret, sym) :: !cache);
-    Hashtbl.replace env.fns sym (cparams, cret);
+    jreplace env.fns sym (cparams, cret);
     let saved_subst = env.subst and saved_vars = env.tyvars
     and saved_preds = env.tvpreds and saved_chain = env.chain in
     (* Inside the copy there are no variables left: [resolve_name] answers
@@ -13481,8 +13531,8 @@ and instantiate env loc gname vars subst cparams cret =
         (* A copy whose body did not check is not a copy. Both entries go back
            out, so a second call at the same types is the same refusal again
            rather than a cache hit on a function that does not exist. *)
-        cache := List.filter (fun (_, _, s) -> s <> sym) !cache;
+        jset cache (List.filter (fun (_, _, s) -> s <> sym) !cache);
-        Hashtbl.remove env.fns sym;
+        jremove env.fns sym;
         raise e
     in
     env.instances <- tfn :: env.instances;
@@ -13583,9 +13633,9 @@ and trial ctx f =
                        outer_what; caught; place_ok; envslot; parent = _;
                        in_frames; loops; tail; in_defer;
                        owner = _ } = ctx in
-  let undo = snapshot_env ctx.env in
+  let undo, keep = snapshot_env ctx.env in
   match speculate ctx.env f with
-  | r -> Ok r
+  | r -> keep (); Ok r
   | exception Loc.Error d ->
     undo ();
     ctx.slots <- slots; ctx.slot_tys <- slot_tys;
@@ -13596,6 +13646,7 @@ and trial ctx f =
     ctx.caught <- caught; ctx.place_ok <- place_ok; ctx.envslot <- envslot;
     ctx.loops <- loops; ctx.tail <- tail; ctx.in_defer <- in_defer;
     Error d
+  | exception e -> keep (); raise e
 
 (* Whether the trial's refusal is one worth reconsidering. A literal that did
    not fit is not, and neither is a refusal a program cannot make any use of
@@ -15147,9 +15198,9 @@ and check_generic env (fn : Ast.fn) =
      [env.tvpreds] whether the variable was declared to support it, and every
      instantiation asks the concrete type the same question again. *)
   env.tvpreds <- fn.Ast.fwhere;
-  Hashtbl.replace env.fns fn.Ast.name (params, ret);
+  jreplace env.fns fn.Ast.name (params, ret);
   let finish () =
-    Hashtbl.remove env.fns fn.Ast.name;
+    jremove env.fns fn.Ast.name;
     env.lifted <- saved_lifted;
     env.tyvars <- saved_vars;
     env.tvpreds <- saved_preds;
@@ -15186,7 +15237,7 @@ and read_return env (fn : Ast.fn) params =
   (* One check of the body against [ret], thrown away with everything it
      wrote into [env] ([snapshot_env]); pass two checks it again for real. *)
   let attempt ret =
-    let undo = snapshot_env env in
+    let undo, _ = snapshot_env env in
     let seen = !infer_seen in
     infer_seen := [];
     let restore () = undo (); infer_seen := seen in
@@ -16538,16 +16589,17 @@ let build_program ~keep_going ?tolerate ?previous (decls : Ast.decl list) :
          ([snapshot_env]): a copy it asked for would otherwise stay cached,
          and the next body asking for it would be handed the name of a copy
          the program does not have. *)
-      let undo = snapshot_env env in
+      let undo, keep = snapshot_env env in
       (match f () with
-       | x -> x
+       | x -> keep (); x
        | exception ((Loc.Error d | Loc.Errors (d :: _)) as e) ->
          if ok env name d then begin
            undo ();
            tolerated := name :: !tolerated;
            None
          end
-         else raise e)
+         else (keep (); raise e)
+       | exception e -> keep (); raise e)
   in
   let decls, prelude_warnings =
     shadow_prelude (Parse.program (Prelude.forms ())) decls

test/test_flan.ml

@@ -7862,6 +7862,69 @@ let () =
      incr failures;
      Printf.printf "FAIL  an else arm's own error alone: %d errors\n" n
    end);
+  (* An arm that needs the other's type, and is refused at it too, says why
+     at that type — not that it has no type on its own. *)
+  rejects_check "a bare struct else arm with an unknown name in it"
+    ~needle:"unknown name q2"
+    "(defstruct P [x i32 y i32])\n\
+     (defn b3 [c bool p P] i32 (let [v (if c p {.x q2 .y 2})] (.y v)))\n\
+     (defn main [] i32 0)";
+  rejects_check "a bare struct match arm with an unknown name in it"
+    ~needle:"unknown name q2"
+    "(defstruct P [x i32 y i32])\n\
+     (defn a3 [o (Option i32) p P] i32 \
+        (let [v (match o (Some q) p None {.x q2 .y 2})] (.y v)))\n\
+     (defn main [] i32 0)";
+  (* A match arm that does not meet the others is refused at its value. *)
+  (match
+     checked
+       "(defn m4 [o (Option i32) x i8] i32 \
+          (let [v (match o (Some q) x None (do (println \"a\") \"lit\"))] 0))\n\
+        (defn main [] i32 0)"
+   with
+   | _ -> check "a match arm of another type is refused" false
+   | exception Loc.Error { Loc.dloc; dmsg; _ } ->
+     check "a match arm of another type is refused at its value"
+       (dloc.Loc.col = 87 && contains dmsg "expected i8, found string"));
+  (* Nested arms that meet at a wider type are checked once each, not once
+     per level above them. *)
+  (let nest kind depth =
+     let rec go k e =
+       if k = 0 then e
+       else
+         go (k - 1)
+           (match kind with
+            | `If -> Printf.sprintf "(if c a (+ (idg b) (i32 %s)))" e
+            | `Match ->
+              Printf.sprintf "(match o (Some q) a None (+ (idg b) (i32 %s)))" e)
+     in
+     go depth "(i32 b)"
+   in
+   let structs =
+     String.concat ""
+       (List.init 300 (Printf.sprintf "(defstruct S%d [a i32 b i64])\n"))
+   in
+   List.iter
+     (fun (kind, sg, what) ->
+        let src =
+          "(defn idg [x $t] $t x)\n" ^ structs
+          ^ "(defn f " ^ sg ^ " i64 (let [v " ^ nest kind 20 ^ "] v))\n\
+             (defn g " ^ sg ^ " _ " ^ nest kind 20 ^ ")"
+        in
+        let t0 = Unix.gettimeofday () in
+        match checked src with
+        | p ->
+          check (what ^ " twenty deep checks fast")
+            (Unix.gettimeofday () -. t0 < 3.0);
+          check (what ^ " twenty deep meets at i64")
+            (List.exists
+               (fun (f : Tast.fn) ->
+                  f.Tast.name = "g" && Types.equal f.Tast.ret (Types.Int Types.I64))
+               p.Tast.fns)
+        | exception Loc.Error { Loc.dmsg; _ } ->
+          check (what ^ " twenty deep checks: " ^ dmsg) false)
+     [ (`If, "[c bool a i64 b i32]", "an if");
+       (`Match, "[o (Option i32) a i64 b i32]", "a match") ]);
   rejects_check "an else arm's own error is the one reported"
     ~needle:"unknown function nope2"
     "(defn g [c bool a i32 b i64] i64 (let [v (if c a (+ b (nope2 1)))] v))\n\