Keep record-flattening refactorings tractable in TLAPS

src/backend/smtlib.ml

@@ -456,6 +456,15 @@ let preprocess ~solver sq =
   let sq = sq
     |> debug "Original Obligation:"
     |> Encode.Rewrite.elim_flex
+    (* Decompose equalities of two record constructors with equal domains into
+       the field-wise conjunction, e.g.
+         [a |-> e1, b |-> e2] = [a |-> f1, b |-> f2]  -->  e1 = f1 /\ e2 = f2
+       Must run before Type.Synthesize: the constructors are then gone before
+       Axiomatize, so no record/function-extensionality axioms (FunExt,
+       RecDomDef, RecAppDef, ...) are emitted -- those are what make
+       wide-record equalities expensive for Z3.  Sound with no typing
+       assumption; see Encode.Rewrite.simpl_receq. *)
+    |> Encode.Rewrite.simpl_receq
     |> Type.Synthesize.main ~typelvl
     |> Encode.Rewrite.elim_notmem
     |> Encode.Rewrite.elim_compare

src/encode.mli

@@ -16,6 +16,7 @@ module Rewrite : sig
   val elim_tuples : sequent -> sequent
   val elim_records : sequent -> sequent
   val sort_recfields : sequent -> sequent
+  val simpl_receq : sequent -> sequent
 
   val simplify_range : sequent -> sequent
   val simplify_sets : ?limit:int -> ?rwlvl:int -> disable_arithmetic:bool -> sequent -> sequent

src/encode/n_rewrite.ml

@@ -556,6 +556,57 @@ let sort_recfields sq =
   snd (sort_recfields_visitor#sequent cx sq)
 
 
+(* {3 Record Equality Decomposition} *)
+
+(* Decompose an equality of two record constructors with the same field set
+   into the conjunction of field-wise equalities:
+     [h1 |-> a1, ...] = [h1 |-> b1, ...]  -->  a1 = b1 /\ ...
+   Run in the SMT-LIB pipeline before type synthesis.  The soundness argument
+   and the guards -- equal field set, pairing by name, field values copied
+   verbatim (so [=] is never pushed through IF), and rewriting equalities only
+   -- are stated and checked by test/unit/h_records/recordeq_*_smt_test.tla and
+   test/soundness_tests/recordeq_swap_stest.tla. *)
+let same_fieldset fs1 fs2 =
+  List.length fs1 = List.length fs2
+  && List.for_all (fun (h, _) -> List.mem_assoc h fs2) fs1
+  && List.for_all (fun (h, _) -> List.mem_assoc h fs1) fs2
+
+let rec mk_conj = function
+  | [] -> Internal B.TRUE %% []
+  | [ e ] -> e
+  | e :: es -> Apply (Internal B.Conj %% [], [ e ; mk_conj es ]) %% []
+
+let rec decompose_receq lhs rhs =
+  match lhs.core, rhs.core with
+  | Record fs1, Record fs2 when same_fieldset fs1 fs2 ->
+      mk_conj
+        (List.map (fun (h, e1) -> decompose_receq e1 (List.assoc h fs2)) fs1)
+  | _ ->
+      Apply (Internal B.Eq %% [], [ lhs ; rhs ]) %% []
+
+let simpl_receq_visitor = object (self : 'self)
+  inherit [unit] Visit.map as super
+
+  method expr scx oe =
+    match oe.core with
+    | Apply ({ core = Internal B.Eq } as op, [ e ; f ])
+      when not (has oe Props.tpars_prop) ->
+        let e = self#expr scx e in
+        let f = self#expr scx f in
+        begin match e.core, f.core with
+        | Record _, Record _ ->
+            (decompose_receq e f).core @@ oe
+        | _ ->
+            Apply (op, [ e ; f ]) @@ oe
+        end
+    | _ -> super#expr scx oe
+end
+
+let simpl_receq sq =
+  let cx = ((), Deque.empty) in
+  snd (simpl_receq_visitor#sequent cx sq)
+
+
 (* {3 Range Simplification} *)
 
 let simplify_range_visitor = object (self : 'self)

src/encode/n_rewrite.mli

@@ -39,6 +39,12 @@ val elim_records : sequent -> sequent
 (** Sort fields of records and record sets *)
 val sort_recfields : sequent -> sequent
 
+(** Decompose equalities between record constructors with identical field
+    sets into the conjunction of their field-wise equalities.  Sound without
+    typing assumptions; avoids the wide record-extensionality axiom on the
+    SMT side.  Intended to run before type synthesis. *)
+val simpl_receq : sequent -> sequent
+
 (** Simplify propositions involving ranges *)
 val simplify_range : sequent -> sequent
 

src/expr/e_elab.ml

@@ -102,33 +102,68 @@ let except_normalize =
 
     method expr scx e = match e.core with
       | Except (f, xs) ->
-          let rec simplify f x =
-            match x with
-            | [[tr], bod] -> { e with core = Except (f, [[tr], self#expr scx bod]) }
-            | [tr :: trs, bod] ->
-                let g = match tr with
-                  | Except_dot x -> { f with core = Dot (f, x) }
-                  | Except_apply e -> { f with core = FcnApp (f, [self#expr scx e]) }
+          (* Fold [f EXCEPT p1 = b1, ...] into [f] left-to-right, materialising
+             each shared prefix once, to avoid the multiplicative blow-up of the
+             naive desugaring (test/regression_tests/
+             record_except_explosion_test.tla).  Soundness of each reduction --
+             fold into a record constructor only on an existing field (its
+             DOMAIN), never push a selection/update through IF, group only
+             *adjacent* same-prefix updates and only on provably equal keys
+             (so opaque/CONSTANT keys are never merged or reordered) -- is
+             pinned by the [let%test_module] below.  Precondition: [@]/[At] are
+             already resolved by [Elab.desugar]. *)
+          let access b hd =
+            match hd with
+            | Except_dot x -> { b with core = Dot (b, x) }
+            | Except_apply k -> { b with core = FcnApp (b, [k]) }
+          in
+          let head_eq a b =
+            match a, b with
+            | Except_dot x, Except_dot y -> x = y
+            | Except_apply e1, Except_apply e2 -> E_eq.expr e1 e2
+            | _, _ -> false
+          in
+          let sel r hd =
+            match r.core, hd with
+            | Record fs, (Except_dot h | Except_apply {core = String h})
+              when List.mem_assoc h fs -> List.assoc h fs
+            | _ -> access r hd
+          in
+          let set b hd v =
+            match b.core, hd with
+            | Record fs, (Except_dot h | Except_apply {core = String h})
+              when List.mem_assoc h fs ->
+                { b with core =
+                    Record (List.map (fun (k, x) -> if k = h then (k, v) else (k, x)) fs) }
+            | _ -> { e with core = Except (b, [[hd], v]) }
+          in
+          let rec apply base subs =
+            match subs with
+            | [] -> base
+            | ([], bod) :: rest ->
+                apply bod rest
+            | (tr, _) :: _ ->
+                let hd = List.hd tr in
+                let same t = match t with h :: _ -> head_eq h hd | [] -> false in
+                let rec span acc = function
+                  | (t, b) :: tl when same t -> span ((List.tl t, b) :: acc) tl
+                  | rest -> (List.rev acc, rest)
                 in
-                { e with core = Except (f, [[tr], simplify g [trs, self#expr scx bod]]) }
-            | x :: xs ->
-                let ex = simplify f [x] in
-                simplify ex xs
-(*
-                let exs = simplify ex xs in
-                begin match exs.core with
-                  | Except (f, xs) ->
-                      { ex with core = Except (f, xs) }
-                  | _ ->
-                      Errors.bug ~at:ex "Expr.Elab.desugar: simplify/except/1"
-                end
-*)
-            | _ ->
-                Errors.bug ~at:f "Expr.Elab.desugar: simplify/except/2"
+                let subtails, rest = span [] subs in
+                let v = apply (sel base hd) subtails in
+                apply (set base hd v) rest
           in
           let f = self#expr scx f in
           let xs = List.map (self#exspec scx) xs in
-          simplify f xs
+          apply f xs
+      (* No top-level "collapse a selection over a record constructor" case:
+         that rewrite is equality-preserving but NOT proof-stable -- it would
+         collapse a hand-cited fact like [m2b.acc = self] to [self = self] and
+         drop it, erasing a term a backend needs (regressed
+         examples/ByzPaxos/BPConProof.tla).  It is also unnecessary: the fold
+         above already collapses every selection EXCEPT desugaring itself
+         produces.  See "a field selection over a record constructor is left
+         intact" in the [let%test_module] below. *)
       | _ -> super#expr scx e
   end in
   fun scx e -> visitor#expr scx e
@@ -165,8 +200,13 @@ let normalize cx e =
   let nte = non_temporal e in
   (* moved to action frontend *)
   (* let e = if nte then action_normalize scx e else e in *)
-  let e = if nte then except_normalize scx e else e in
+  (* let_normalize before except_normalize: refinement mappings bind the
+     updated state to LET operators, so inlining them first exposes the record
+     constructors that except_normalize folds into.  With the opposite order
+     the bases stay opaque LET variables and get re-embedded (and then
+     multiplied out) per path component. *)
   let e = let_normalize scx e in
+  let e = if nte then except_normalize scx e else e in
   (* moved to action frontend *)
   (* let e = if nte then unchanged_normalize scx e else e in
   let e = prime_normalize cx e in
@@ -238,6 +278,10 @@ let%test_module _ = (module struct
       [%test_eq: string] (prn_exp target_case) (prn_exp (normalize Deque.empty test_case))
 
   let%test_unit "t2" =
+    (* Only *adjacent* same-prefix updates are grouped, so the two updates to
+       key [0] are not merged and left-to-right order is preserved.  Keys are
+       compared by provable equality only, so opaque/CONSTANT keys are never
+       assumed equal and hence never merged or reordered. *)
     let test_case = create_expression "[[f EXCEPT ![0] = 10, ![1] = 1] EXCEPT ![0] = 0]" in
     let target_case = create_expression "[[[f EXCEPT ![0] = 10] EXCEPT ![1] = 1] EXCEPT ![0] = 0]" in
       [%test_eq: string] (prn_exp target_case) (prn_exp (normalize Deque.empty test_case))
@@ -265,6 +309,107 @@ let%test_module _ = (module struct
       [[arr EXCEPT ![x] = [arr[x] EXCEPT ![y] = foo]][u] EXCEPT ![v] = bar]]" in
         [%test_eq: string] (prn_exp target_case) (prn_exp (normalize Deque.empty test_case))
 
+  let%test_unit "except does not add a record field" =
+    (* EXCEPT preserves the domain of its base function.  Since [b] is not in
+       the domain of this record, selecting [.b] cannot be reduced to the
+       replacement value. *)
+    let test_case =
+      create_expression "[[a |-> 1] EXCEPT !.b = 2].b"
+    in
+    let target_case =
+      create_expression "[[a |-> 1] EXCEPT !.b = 2].b"
+    in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  let%test_unit "projection does not distribute through a non-Boolean IF" =
+    (* TLA+ is untyped.  IF is guaranteed to select a branch only when its
+       condition is Boolean, so this projection cannot be distributed without
+       first establishing that condition. *)
+    let conditional =
+      create_expression "IF 0 THEN [a |-> 1] ELSE [a |-> 2]"
+    in
+    (* Parentheses are retained explicitly by the parser and would hide the
+       [If] node from this normalization rule, so construct the projection AST
+       directly. *)
+    let test_case = Dot (conditional, "a") @@ conditional in
+    let target_case = Dot (conditional, "a") @@ conditional in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  let%test_unit "matching EXCEPT over an opaque base needs field membership" =
+    (* Unlike a record literal, an opaque base has an unknown domain.
+       [r EXCEPT !.b = 2].b] equals 2 only when [b \in DOMAIN r], which cannot
+       be assumed here, so the selection must not be reduced to 2. *)
+    let test_case = create_expression "[r EXCEPT !.b = 2].b" in
+    let target_case = create_expression "[r EXCEPT !.b = 2].b" in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  let%test_unit "nonmatching EXCEPT over an opaque base is not the base selection" =
+    (* [[r EXCEPT !.a = 2].b] is not provably [r.b]: when [b \notin DOMAIN r]
+       both sides are unspecified but need not be the same unspecified value
+       (their function arguments differ).  With an opaque, unknown-domain base
+       this reduction is therefore unsound. *)
+    let test_case = create_expression "[r EXCEPT !.a = 2].b" in
+    let target_case = create_expression "[r EXCEPT !.a = 2].b" in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  let%test_unit "projection does not distribute through an unknown-condition IF" =
+    (* Even when the condition is not a manifest non-Boolean, its Boolean-ness
+       is unknown for an opaque [p]; distributing the projection over the
+       branches is unsound unless [p \in BOOLEAN] has been established. *)
+    let conditional =
+      create_expression "IF p THEN [a |-> 1] ELSE [a |-> 2]"
+    in
+    let test_case = Dot (conditional, "a") @@ conditional in
+    let target_case = Dot (conditional, "a") @@ conditional in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  (* Positive counterparts of the guard tests above.  The blow-up in
+     test/regression_tests/record_except_explosion_test.tla is defused by
+     *folding* EXCEPT updates into the record constructor they update, so a wide
+     base is materialised once instead of being re-embedded per path component
+     -- NOT by collapsing field selections.  We deliberately do not collapse
+     selections (see the note in [except_normalize]): that is equality-
+     preserving but not proof-stable. *)
+
+  let%test_unit "EXCEPT over a record constructor folds the update in place" =
+    (* [b \in DOMAIN [a |-> 1, b |-> 2]], so the update is folded into the
+       constructor, keeping the result linear in the record's width. *)
+    let test_case = create_expression "[[a |-> 1, b |-> 2] EXCEPT !.b = 3]" in
+    let target_case = create_expression "[a |-> 1, b |-> 3]" in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  let%test_unit "multiple EXCEPT updates fold into a single constructor" =
+    let test_case =
+      create_expression "[[a |-> 1, b |-> 2] EXCEPT !.a = 3, !.b = 4]"
+    in
+    let target_case = create_expression "[a |-> 3, b |-> 4]" in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
+  let%test_unit "a field selection over a record constructor is left intact" =
+    (* Proof-stability: even though [[a |-> 1, b |-> 2].b = 2] is provable, we
+       must not rewrite it, since the same collapse applied to a hand-cited fact
+       (e.g. [m2b.acc = self]) erases the term a backend needs.  Folding keeps
+       the obligation linear without touching the selection. *)
+    let test_case = create_expression "[a |-> 1, b |-> 2].b" in
+    let target_case = create_expression "[a |-> 1, b |-> 2].b" in
+    [%test_eq: string]
+      (prn_exp target_case)
+      (prn_exp (normalize Deque.empty test_case))
+
   (*
   let%test_unit "t7" [@tags "disabled"] = (* doesnt work because we need to anonimie the created expressions from the parser*)
     let test_string = "f[x]'" in

test/regression_tests/record_except_explosion_test.tla

@@ -0,0 +1,73 @@
+---- MODULE record_except_explosion_test ----
+
+\* Regression guard for the EXCEPT-normalization fold in Expr.Elab
+\* (src/expr/e_elab.ml): a `[ recordLiteral EXCEPT ... ]` is folded back into a
+\* record literal instead of re-embedding the wide base once per path
+\* component.
+\*
+\* `Chain(z)` is a chain of nested, multi-component EXCEPT updates over a wide
+\* (20-field) record literal.  It is normalized during obligation generation,
+\* BEFORE any backend runs.  Without the fold the normalized term grows
+\* geometrically in the number of chain layers; the goal is deliberately
+\* trivial (`fa` is written once in `Step` and never touched again, so
+\* `Chain(z).fa = z` holds by inspection), leaving term SIZE as the only thing
+\* at stake.
+\*
+\* Measured with `--noproving --verbose | wc -l` at the 4 layers below:
+\*   fold present (Expr.Elab):   ~5000 lines
+\*   fold removed:              >=39000 lines  (>=9,000,000 lines at 6 layers)
+\* The size is monotone in the layer count, so the (500, 25000) window in the
+\* command below sits safely between the two: a regression that drops the fold
+\* makes this test fail.  `head` caps the captured output so the failing case
+\* stays cheap.
+\*
+\* This is a term-SIZE guard only; the semantic correctness of the fold (that
+\* the reduced value is right, and that no reduction is unsound) is covered by
+\* the inline tests in src/expr/e_elab.ml.
+
+EXTENDS Naturals, TLAPS
+
+VARIABLES p, dir, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, q, r, s
+CONSTANT X, Y
+
+\* A wide (20-field) record literal.  `fp` is accessed as a function (indexed
+\* `[X]`) and `fdir` as a record (sub-field `.local`), so the EXCEPT paths
+\* through them are multi-component.
+Rec ==
+  [ fp |-> p, fdir |-> dir,
+    fa |-> a, fb |-> b, fc |-> c, fd |-> d, fe |-> e, ff |-> f, fg |-> g, fh |-> h,
+    fi |-> i, fj |-> j, fk |-> k, fl |-> l, fm |-> m, fn |-> n, fo |-> o, fq |-> q,
+    fr |-> r, fs |-> s ]
+
+\* One update layer: multi-component paths and IF-valued sub-updates over the
+\* wide literal.
+Step(z) ==
+  [ Rec EXCEPT
+      !.fp   = IF z THEN [Rec.fp   EXCEPT ![X].c1 = TRUE, ![Y].c1 = FALSE]
+                    ELSE  Rec.fp,
+      !.fdir = IF z THEN [Rec.fdir EXCEPT !.local = TRUE, !.pending = FALSE]
+                    ELSE  Rec.fdir,
+      !.fa = z, !.fb = z, !.fc = z, !.fd = z, !.fe = z ]
+
+\* Four update layers over a common base bound by LET; each layer is a nested
+\* EXCEPT with multi-component paths, so the duplication compounds
+\* multiplicatively without the fold.
+Chain(z) ==
+  LET s1 == Step(z)
+      s2 == [ s1 EXCEPT !.fp[X].c1 = FALSE, !.fdir.local   = FALSE, !.ff = z, !.fg = z ]
+      s3 == [ s2 EXCEPT !.fp[Y].c1 = TRUE,  !.fdir.pending = TRUE,  !.fh = z, !.fi = z ]
+      s4 == [ s3 EXCEPT !.fp[X].c1 = TRUE,  !.fdir.local   = TRUE,  !.fj = z, !.fk = z ]
+  IN  s4
+
+LEMMA Explode ==
+  ASSUME NEW z \in BOOLEAN
+  PROVE  Chain(z).fa = z
+BY DEF Chain, Step, Rec
+
+====
+\* Bound the size of the normalized obligation (see header).  With the
+\* record-literal EXCEPT fold it is ~5000 lines; without it >=39000 and growing
+\* geometrically, so the (500, 25000) window fails on regression.  `head` caps
+\* the captured output so a regression cannot grow it unboundedly here.
+command: L=$( ${TLAPM} --noproving --verbose --nofp ${FILE} 2>&1 | head -n 60000 | wc -l | tr -d ' ' ); echo "normalized-obligation lines (cap 60000): $L"; test "$L" -gt 500 && test "$L" -lt 25000
+result: 0