From 143eb793af5d4d677cf596e39b3639ca1993ac6a Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Wed, 10 Jun 2020 07:04:01 -0700
Subject: [PATCH] [sledge] Refactor: Add `let@`

Summary:
```
val ( let@ ) : ('a -> 'b) -> 'a -> 'b
(** [let@ x = e in b] is equivalent to [e @@ fun x -> b], that is,
    [e (fun x -> b)] *)
```

Reviewed By: jvillard

Differential Revision: D21721025

fbshipit-source-id: d8efdebbe
---
 sledge/src/control.ml        |  3 +--
 sledge/src/equality.ml       | 16 ++++++----------
 sledge/src/exp.ml            |  6 ++----
 sledge/src/global.ml         |  3 +--
 sledge/src/import/import.mli |  4 ++++
 sledge/src/import/import0.ml |  2 ++
 sledge/src/import/list.ml    |  3 +--
 sledge/src/import/map.ml     |  6 ++----
 sledge/src/import/set.ml     |  3 +--
 sledge/src/llair.ml          |  9 +++------
 sledge/src/sh.ml             |  3 +--
 sledge/src/solver.ml         |  3 +--
 sledge/src/term.ml           | 10 ++++------
 sledge/src/typ.ml            |  3 +--
 14 files changed, 30 insertions(+), 44 deletions(-)

diff --git a/sledge/src/control.ml b/sledge/src/control.ml
index e3124b198..bf3b5323f 100644
--- a/sledge/src/control.ml
+++ b/sledge/src/control.ml
@@ -87,8 +87,7 @@ module Make (Dom : Domain_intf.Dom) = struct
       | Empty -> ()
 
     let invariant s =
-      Invariant.invariant [%here] s [%sexp_of: t]
-      @@ fun () ->
+      let@ () = Invariant.invariant [%here] s [%sexp_of: t] in
       match s with
       | Return _ | Throw (_, Return _) | Empty -> ()
       | Throw _ -> fail "malformed stack: %a" print_abbrev s ()
diff --git a/sledge/src/equality.ml b/sledge/src/equality.ml
index 191d12b43..d23562d58 100644
--- a/sledge/src/equality.ml
+++ b/sledge/src/equality.ml
@@ -411,8 +411,7 @@ let congruent r a b = semi_congruent r (Term.map ~f:(Subst.norm r.rep) a) b
 (** Invariant *)
 
 let pre_invariant r =
-  Invariant.invariant [%here] r [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] r [%sexp_of: t] in
   Subst.iteri r.rep ~f:(fun ~key:trm ~data:_ ->
       (* no interpreted terms in carrier *)
       assert (non_interpreted trm || fail "non-interp %a" Term.pp trm ()) ;
@@ -425,8 +424,7 @@ let pre_invariant r =
                  subtrm Term.pp trm pp r () ) ) )
 
 let invariant r =
-  Invariant.invariant [%here] r [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] r [%sexp_of: t] in
   pre_invariant r ;
   assert (
     (not r.sat)
@@ -449,13 +447,12 @@ let false_ = {true_ with sat= false}
 
 (** [lookup r a] is [b'] if [a ~ b = b'] for some equation [b = b'] in rep *)
 let lookup r a =
-  [%Trace.call fun {pf} -> pf "%a" Term.pp a]
+  ([%Trace.call fun {pf} -> pf "%a" Term.pp a]
   ;
-  ( with_return
-  @@ fun {return} ->
+  let@ {return} = with_return in
   Subst.iteri r.rep ~f:(fun ~key:b ~data:b' ->
       if semi_congruent r a b then return b' ) ;
-  a )
+  a)
   |>
   [%Trace.retn fun {pf} -> pf "%a" Term.pp]
 
@@ -507,8 +504,7 @@ let merge us a b r =
 
 (** find an unproved equation between congruent terms *)
 let find_missing r =
-  with_return
-  @@ fun {return} ->
+  let@ {return} = with_return in
   Subst.iteri r.rep ~f:(fun ~key:a ~data:a' ->
       let a_subnorm = Term.map ~f:(Subst.norm r.rep) a in
       Subst.iteri r.rep ~f:(fun ~key:b ~data:b' ->
diff --git a/sledge/src/exp.ml b/sledge/src/exp.ml
index 384d37fe0..6c8d42273 100644
--- a/sledge/src/exp.ml
+++ b/sledge/src/exp.ml
@@ -172,8 +172,7 @@ and pp_record fs elts =
 let valid_idx idx elts = 0 <= idx && idx < IArray.length elts
 
 let rec invariant exp =
-  Invariant.invariant [%here] exp [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] exp [%sexp_of: t] in
   match exp.desc with
   | Reg {typ} | Nondet {typ} -> assert (Typ.is_sized typ)
   | Integer {data; typ} -> (
@@ -314,8 +313,7 @@ module Reg = struct
     [@@warning "-9"]
 
   let invariant x =
-    Invariant.invariant [%here] x [%sexp_of: t]
-    @@ fun () ->
+    let@ () = Invariant.invariant [%here] x [%sexp_of: t] in
     match x.desc with Reg _ -> invariant x | _ -> assert false
 
   let name r =
diff --git a/sledge/src/global.ml b/sledge/src/global.ml
index c88e0d092..dce9085e5 100644
--- a/sledge/src/global.ml
+++ b/sledge/src/global.ml
@@ -25,8 +25,7 @@ let pp_defn fs {reg; init; loc} =
     (Option.map ~f:fst init)
 
 let invariant g =
-  Invariant.invariant [%here] g [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] g [%sexp_of: t] in
   let {reg} = g in
   assert (Typ.is_sized (Reg.typ reg)) ;
   assert (Var.is_global (Reg.var reg))
diff --git a/sledge/src/import/import.mli b/sledge/src/import/import.mli
index 42b346f24..312a3f1ab 100644
--- a/sledge/src/import/import.mli
+++ b/sledge/src/import/import.mli
@@ -11,6 +11,10 @@ include module type of Import0
 
 (** Function combinators *)
 
+val ( let@ ) : ('a -> 'b) -> 'a -> 'b
+(** [let@ x = e in b] is equivalent to [e @@ fun x -> b], that is,
+    [e (fun x -> b)] *)
+
 val ( >> ) : ('a -> 'b) -> ('b -> 'c) -> 'a -> 'c
 (** Composition of functions: [(f >> g) x] is exactly equivalent to
     [g (f (x))]. Left associative. *)
diff --git a/sledge/src/import/import0.ml b/sledge/src/import/import0.ml
index 07bf4018a..2ba9c4f9f 100644
--- a/sledge/src/import/import0.ml
+++ b/sledge/src/import/import0.ml
@@ -39,6 +39,8 @@ let ( <$ ) f x =
   f x ;
   x
 
+let ( let@ ) x f = x @@ f
+
 (** Pretty-printer for argument type. *)
 type 'a pp = Format.formatter -> 'a -> unit
 
diff --git a/sledge/src/import/list.ml b/sledge/src/import/list.ml
index ccf95310d..d5b7b1e34 100644
--- a/sledge/src/import/list.ml
+++ b/sledge/src/import/list.ml
@@ -42,8 +42,7 @@ let find_map_remove xs ~f =
   find_map_remove_ [] xs
 
 let fold_option xs ~init ~f =
-  with_return
-  @@ fun {return} ->
+  let@ {return} = with_return in
   Some
     (fold xs ~init ~f:(fun acc elt ->
          match f acc elt with Some res -> res | None -> return None ))
diff --git a/sledge/src/import/map.ml b/sledge/src/import/map.ml
index 0764bb7ea..7910ff2ad 100644
--- a/sledge/src/import/map.ml
+++ b/sledge/src/import/map.ml
@@ -52,16 +52,14 @@ end) : S with type key = Key.t = struct
     Container.fold_until ~fold ~init ~f ~finish m
 
   let root_key_exn m =
-    with_return
-    @@ fun {return} ->
+    let@ {return} = with_return in
     binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data:_ ->
         return key )
     |> ignore ;
     raise (Not_found_s (Atom __LOC__))
 
   let choose_exn m =
-    with_return
-    @@ fun {return} ->
+    let@ {return} = with_return in
     binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data ->
         return (key, data) )
     |> ignore ;
diff --git a/sledge/src/import/set.ml b/sledge/src/import/set.ml
index 7ed66ed68..ce224bb57 100644
--- a/sledge/src/import/set.ml
+++ b/sledge/src/import/set.ml
@@ -48,8 +48,7 @@ end) : S with type elt = Elt.t = struct
         | l2, None, r2 -> disjoint l1 l2 && disjoint r1 r2 )
 
   let choose_exn s =
-    with_return
-    @@ fun {return} ->
+    let@ {return} = with_return in
     binary_search_segmented s `Last_on_left ~segment_of:return |> ignore ;
     raise (Not_found_s (Atom __LOC__))
 
diff --git a/sledge/src/llair.ml b/sledge/src/llair.ml
index c14a19b9e..cd8e5d271 100644
--- a/sledge/src/llair.ml
+++ b/sledge/src/llair.ml
@@ -300,8 +300,7 @@ module Term = struct
   let pp = pp_term
 
   let invariant ?parent term =
-    Invariant.invariant [%here] term [%sexp_of: t]
-    @@ fun () ->
+    let@ () = Invariant.invariant [%here] term [%sexp_of: t] in
     match term with
     | Switch _ | Iswitch _ -> assert true
     | Call {typ; actuals; areturn; _} -> (
@@ -454,9 +453,8 @@ module Func = struct
             cfg )
 
   let invariant func =
-    Invariant.invariant [%here] func [%sexp_of: t]
-    @@ fun () ->
     assert (func == func.entry.parent) ;
+    let@ () = Invariant.invariant [%here] func [%sexp_of: t] in
     match Reg.typ func.name.reg with
     | Pointer {elt= Function {return; _}; _} ->
         assert (
@@ -579,8 +577,7 @@ let set_derived_metadata functions =
   functions
 
 let invariant pgm =
-  Invariant.invariant [%here] pgm [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] pgm [%sexp_of: t] in
   assert (
     not
       (IArray.contains_dup pgm.globals ~compare:(fun g1 g2 ->
diff --git a/sledge/src/sh.ml b/sledge/src/sh.ml
index 58466ec9b..5d5b10fb8 100644
--- a/sledge/src/sh.ml
+++ b/sledge/src/sh.ml
@@ -293,8 +293,7 @@ let invariant_pure = function
 let invariant_seg _ = ()
 
 let rec invariant q =
-  Invariant.invariant [%here] q [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] q [%sexp_of: t] in
   let {us; xs; cong; pure; heap; djns} = q in
   try
     assert (
diff --git a/sledge/src/solver.ml b/sledge/src/solver.ml
index 8dc59591d..a6a7af5be 100644
--- a/sledge/src/solver.ml
+++ b/sledge/src/solver.ml
@@ -73,8 +73,7 @@ end = struct
       sub
 
   let invariant g =
-    Invariant.invariant [%here] g [%sexp_of: t]
-    @@ fun () ->
+    let@ () = Invariant.invariant [%here] g [%sexp_of: t] in
     try
       let {us; com; min; xs; sub; zs; pgs= _} = g in
       assert (Var.Set.equal us com.us) ;
diff --git a/sledge/src/term.ml b/sledge/src/term.ml
index 7a0976a57..5f351f5f8 100644
--- a/sledge/src/term.ml
+++ b/sledge/src/term.ml
@@ -299,8 +299,7 @@ let rec assert_aggregate = function
   | _ -> assert false
 
 let invariant e =
-  Invariant.invariant [%here] e [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] e [%sexp_of: t] in
   match e with
   | And _ -> assert_conjunction e |> Fn.id
   | Or _ -> assert_disjunction e |> Fn.id
@@ -343,8 +342,8 @@ module Var = struct
   end
 
   let invariant x =
-    Invariant.invariant [%here] x [%sexp_of: t]
-    @@ fun () -> match x with Var _ -> invariant x | _ -> assert false
+    let@ () = Invariant.invariant [%here] x [%sexp_of: t] in
+    match x with Var _ -> invariant x | _ -> assert false
 
   let id = function Var v -> v.id | x -> violates invariant x
   let name = function Var v -> v.name | x -> violates invariant x
@@ -372,8 +371,7 @@ module Var = struct
     let t_of_sexp = Map.t_of_sexp T.t_of_sexp
 
     let invariant s =
-      Invariant.invariant [%here] s [%sexp_of: t]
-      @@ fun () ->
+      let@ () = Invariant.invariant [%here] s [%sexp_of: t] in
       let domain, range =
         Map.fold s ~init:(Set.empty, Set.empty)
           ~f:(fun ~key ~data (domain, range) ->
diff --git a/sledge/src/typ.ml b/sledge/src/typ.ml
index 1840b8b12..c2e66bc52 100644
--- a/sledge/src/typ.ml
+++ b/sledge/src/typ.ml
@@ -65,8 +65,7 @@ let is_sized = function
   | Opaque _ -> (* optimistically assume linking will make it sized *) true
 
 let invariant t =
-  Invariant.invariant [%here] t [%sexp_of: t]
-  @@ fun () ->
+  let@ () = Invariant.invariant [%here] t [%sexp_of: t] in
   match t with
   | Function {return; args} ->
       assert (Option.for_all ~f:is_sized return) ;