coupling rule: one-sided and two-sided

namasikanam · namasikanam · commit 69efeaec8e4d · 2025-10-22T21:38:58.000+02:00
diff --git a/src/ecHiTacticals.ml b/src/ecHiTacticals.ml
@@ -200,6 +200,7 @@ and process1_phl (_ : ttenv) (t : phltactic located) (tc : tcenv1) =
     | Psetmatch info            -> EcPhlCodeTx.process_set_match info
     | Pweakmem info             -> EcPhlCodeTx.process_weakmem info
     | Prnd (side, pos, info)    -> EcPhlRnd.process_rnd side pos info
+    | Pcoupling (side, f)       -> EcPhlRnd.process_coupling side f
     | Prndsem (red, side, pos)  -> EcPhlRnd.process_rndsem ~reduce:red side pos
     | Pconseq (opt, info)       -> EcPhlConseq.process_conseq_opt opt info
     | Pconseqauto cm            -> process_conseqauto cm
diff --git a/src/ecLexer.mll b/src/ecLexer.mll
@@ -145,6 +145,7 @@
     "cfold"       , CFOLD      ;        (* KW: tactic *)
     "rnd"         , RND        ;        (* KW: tactic *)
     "rndsem"      , RNDSEM     ;        (* KW: tactic *)
+    "coupling"    , COUPLING   ;        (* KW: tactic *)
     "pr_bounded"  , PRBOUNDED  ;        (* KW: tactic *)
     "bypr"        , BYPR       ;        (* KW: tactic *)
     "byphoare"    , BYPHOARE   ;        (* KW: tactic *)
diff --git a/src/ecParser.mly b/src/ecParser.mly
@@ -416,6 +416,7 @@
 %token CONSEQ
 %token CONST
 %token COQ
+%token COUPLING
 %token CHECK
 %token EDIT
 %token FIX
@@ -3000,9 +3001,13 @@ interleave_info:
 | RND s=side? info=rnd_info c=prefix(COLON, semrndpos)?
     { Prnd (s, c, info) }
 
+
 | RNDSEM red=boption(STAR) s=side? c=codepos1
     { Prndsem (red, s, c) }
 
+| COUPLING s=side? f=sform
+    { Pcoupling (s, f) }
+
 | INLINE s=side? u=inlineopt? o=occurences?
   { Pinline (`ByName(s, u, ([], o))) }
 
diff --git a/src/ecParsetree.ml b/src/ecParsetree.ml
@@ -748,6 +748,7 @@ type phltactic =
   | Pasgncase      of (oside * pcodepos)
   | Prnd           of oside * psemrndpos option * rnd_tac_info_f
   | Prndsem        of bool * oside * pcodepos1
+  | Pcoupling      of oside * pformula
   | Palias         of (oside * pcodepos * osymbol_r)
   | Pweakmem       of (oside * psymbol * fun_params)
   | Pset           of (oside * pcodepos * bool * psymbol * pexpr)
diff --git a/src/phl/ecPhlPrRw.ml b/src/phl/ecPhlPrRw.ml
@@ -339,4 +339,4 @@ let t_pr_rewrite (s, f) tc =
     let hyps = LDecl.push_active_ss mp hyps in
     {m;inv=EcProofTyping.process_form hyps f ty}
   in
-  t_pr_rewrite_low (s, omap to_env f) tc
+  t_pr_rewrite_low (s, omap to_env f) tc
diff --git a/src/phl/ecPhlRnd.ml b/src/phl/ecPhlRnd.ml
@@ -10,6 +10,7 @@ open EcSubst
 
 open EcMatching.Position
 open EcCoreGoal
+open EcCoreFol
 open EcLowGoal
 open EcLowPhlGoal
 
@@ -442,6 +443,89 @@ module Core = struct
       | `Right -> f_equivS (snd es.es_ml) mt (es_pr es) es.es_sl s (es_po es) in
     FApi.xmutate1 tc (`RndSem pos) [concl]
 
+  (* -------------------------------------------------------------------- *)
+  let t_equiv_coupling_r side g tc =
+    (* process the following pRHL goal, where g is a coupling of g1 and g2 *)
+    (*                   {phi} c1; x <$ g1 ~ c2; y <$ g2 {psi}             *)
+    let env = FApi.tc1_env tc in
+    let es = tc1_as_equivS tc in
+    let ml = fst es.es_ml in
+    let mr = fst es.es_mr in
+    let (lvL, muL), sl' = tc1_last_rnd tc es.es_sl in
+    let (lvR, muR), sr' = tc1_last_rnd tc es.es_sr in
+    let tyL = proj_distr_ty env (e_ty muL) in
+    let tyR = proj_distr_ty env (e_ty muR) in
+    let muL = (EcFol.ss_inv_of_expr ml muL).inv in
+    let muR = (EcFol.ss_inv_of_expr mr muR).inv in
+    let g = g.inv in
+
+    let goal =
+      match side with
+      | None ->
+        (* Goal: {phi} c1 ~ c2 {iscoupling g muL muR /\ (forall a b, (a, b) \in supp(g) => psi[x -> a, y -> b])} *)
+        (* Generate two free variables a and b and the pair (a, b) *)
+        let a_id = EcIdent.create "a" in
+        let b_id = EcIdent.create "b" in
+        let a = f_local a_id tyL in
+        let b = f_local b_id tyR in
+        let ab = f_tuple [a; b] in
+
+        (* Generate the coupling distribution type: (tyL * tyR) distr *)
+        let coupling_ty = ttuple [tyL; tyR] in
+        let g_app = f_app_simpl g [] (tdistr coupling_ty) in
+
+        let iscoupling_op = EcPath.extend EcCoreLib.p_top ["Distr"; "iscoupling"] in
+        let iscoupling_ty = tfun (tdistr tyL) (tfun (tdistr tyR) (tfun (tdistr coupling_ty) tbool)) in
+        let iscoupling_pred = f_app (f_op iscoupling_op [tyL; tyR] iscoupling_ty) [muL; muR; g_app] tbool in
+
+        (* Substitute in the postcondition *)
+        let post = (es_po es) in
+        let post_subst = subst_form_lv_left env lvL {ml=ml;mr=mr;inv=a} post in
+        let post_subst = subst_form_lv_right env lvR {ml=ml;mr=mr;inv=b} post_subst in
+        
+        let goal = map_ts_inv1 (f_imp (f_in_supp ab g_app)) post_subst in
+        let goal = map_ts_inv1 (f_forall_simpl [(a_id, GTty tyL); (b_id, GTty tyR)]) goal in
+        map_ts_inv1 (f_and iscoupling_pred) goal
+      | Some side ->
+        (* Goal (left):  {phi} c1 ~ c2 {dmap d1 g = d2 /\
+             forall a b, a \in supp(d1) -> b = g(a) -> psi[x -> a, y -> b]} *)
+        (* Goal (right): {phi} c1 ~ c2 {d1 = dmap d2 g /\
+             forall a b, b \in supp(d2) -> a = g(b) -> psi[x -> a, y -> b]} *)
+        let dmap_op = EcPath.extend EcCoreLib.p_top ["Distr"; "dmap"] in
+        let dmap_eq =
+          match side with
+          | `Left ->
+             let dmap_ty = tfun (tdistr tyL) (tfun (tfun tyL tyR) (tdistr tyR)) in
+             let dmap_pred = f_app (f_op dmap_op [tyL; tyR] dmap_ty) [muL ; g] (tdistr tyR) in
+             f_eq dmap_pred muR
+          | `Right -> 
+             let dmap_ty = tfun (tdistr tyR) (tfun (tfun tyR tyL) (tdistr tyL)) in
+             let dmap_pred = f_app (f_op dmap_op [tyR; tyL] dmap_ty) [muR ; g] (tdistr tyL) in
+             f_eq muL dmap_pred in
+
+        let a_id = EcIdent.create "a" in
+        let b_id = EcIdent.create "b" in
+        let a = f_local a_id tyL in
+        let b = f_local b_id tyR in
+        let post = es_po es in
+        let post_subst = subst_form_lv_left env lvL {ml=ml;mr=mr;inv=a} post in
+        let post_subst = subst_form_lv_right env lvR {ml=ml;mr=mr;inv=b} post_subst in
+
+        let goal = 
+          match side with
+          | `Left  -> map_ts_inv1
+                        (f_imp (f_in_supp a muL))
+                        (map_ts_inv1 (f_imp (f_eq (f_app_simpl g [a] tyR) b)) post_subst)
+          | `Right -> map_ts_inv1
+                        (f_imp (f_in_supp b muR))
+                        (map_ts_inv1 (f_imp (f_eq a (f_app_simpl g [b] tyL))) post_subst) in
+
+        let goal = map_ts_inv1 (f_forall_simpl [(a_id, GTty tyL); (b_id, GTty tyR)]) goal in
+        map_ts_inv1 (f_and dmap_eq) goal
+    in
+    let goal = f_equivS (snd es.es_ml) (snd es.es_mr) (es_pr es) sl' sr' goal in
+
+    FApi.xmutate1 tc `Rnd [goal]
 end (* Core *)
 
 (* -------------------------------------------------------------------- *)
@@ -666,32 +750,32 @@ let process_rnd side pos tac_info tc =
       t_bdhoare_rnd tac_info tc
 
   | _, _, _ when is_equivS concl ->
-    let process_form f ty1 ty2 =
-      TTC.tc1_process_prhl_form tc (tfun ty1 ty2) f in
-
-    let bij_info =
-      match tac_info with
-      | PNoRndParams -> None, None
-      | PSingleRndParam f -> Some (process_form f), None
-      | PTwoRndParams (f, finv) -> Some (process_form f), Some (process_form finv)
-      | _ -> tc_error !!tc "invalid arguments"
-    in
+         let process_form f ty1 ty2 =
+           TTC.tc1_process_prhl_form tc (tfun ty1 ty2) f in
+
+         let bij_info =
+           match tac_info with
+           | PNoRndParams -> None, None
+           | PSingleRndParam f -> Some (process_form f), None
+           | PTwoRndParams (f, finv) -> Some (process_form f), Some (process_form finv)
+           | _ -> tc_error !!tc "invalid arguments"
+         in
 
-    let pos = pos |> Option.map (function
-      | Single (b, p) ->
-          let p =
-            if Option.is_some side then
-              EcProofTyping.tc1_process_codepos1 tc (side, p)
-            else EcTyping.trans_codepos1 (FApi.tc1_env tc) p
-          in Single (b, p)
-      | Double ((b1, p1), (b2, p2)) ->
-          let p1 = EcProofTyping.tc1_process_codepos1 tc (Some `Left , p1) in
-          let p2 = EcProofTyping.tc1_process_codepos1 tc (Some `Right, p2) in
-          Double ((b1, p1), (b2, p2))
-    )
-    in
-    
-    t_equiv_rnd side ?pos bij_info tc
+         let pos = pos |> Option.map (function
+           | Single (b, p) ->
+               let p =
+                 if Option.is_some side then
+                   EcProofTyping.tc1_process_codepos1 tc (side, p)
+                 else EcTyping.trans_codepos1 (FApi.tc1_env tc) p
+               in Single (b, p)
+           | Double ((b1, p1), (b2, p2)) ->
+               let p1 = EcProofTyping.tc1_process_codepos1 tc (Some `Left , p1) in
+               let p2 = EcProofTyping.tc1_process_codepos1 tc (Some `Right, p2) in
+               Double ((b1, p1), (b2, p2))
+         )
+         in
+         
+         t_equiv_rnd side ?pos bij_info tc
 
   | _ -> tc_error !!tc "invalid arguments"
 
@@ -713,3 +797,24 @@ let process_rndsem ~reduce side pos tc =
   | Some side when is_equivS concl ->
      t_equiv_rndsem reduce side pos tc
   | _ -> tc_error !!tc "invalid arguments"
+
+let process_coupling side g tc =
+  let concl = FApi.tc1_goal tc in
+  
+  if not (is_equivS concl) then
+    tc_error !!tc "coupling can only be used on pRHL goals"
+  else
+    let env = FApi.tc1_env tc in
+    let es = tc1_as_equivS tc in
+    let (_, muL), _ = tc1_last_rnd tc es.es_sl in
+    let (_, muR), _ = tc1_last_rnd tc es.es_sr in
+    let tyL = proj_distr_ty env (e_ty muL) in
+    let tyR = proj_distr_ty env (e_ty muR) in
+
+    let coupling_ty =
+      match side with
+      | None -> tdistr (ttuple [tyL; tyR])
+      | Some `Left -> tfun tyL tyR
+      | Some `Right -> tfun tyR tyL in
+    let g_form = TTC.tc1_process_prhl_form tc coupling_ty g in
+    Core.t_equiv_coupling_r side g_form tc
diff --git a/src/phl/ecPhlRnd.mli b/src/phl/ecPhlRnd.mli
@@ -24,5 +24,8 @@ val t_equiv_rnd   : ?pos:semrndpos -> oside -> (mkbij_t option) pair -> backward
 (* -------------------------------------------------------------------- *)
 val process_rnd : oside -> psemrndpos option -> rnd_infos_t -> backward
 
+(* -------------------------------------------------------------------- *)
+val process_coupling : oside -> pformula -> backward
+
 (* -------------------------------------------------------------------- *)
 val process_rndsem : reduce:bool -> oside -> pcodepos1 -> backward
diff --git a/tests/coupling-rnd.ec b/tests/coupling-rnd.ec
@@ -0,0 +1,121 @@
+require import AllCore List.
+require import Distr DBool DList.
+(*---*) import Biased.
+require import StdBigop.
+(*---*) import Bigreal Bigreal.BRA.
+
+type t.
+
+op test : t -> bool.
+op D : t distr.
+
+axiom D_ll : is_lossless D.
+
+module B = {
+  proc f() : bool = {
+    var x: bool;
+    var garbage: int;
+    (* put some garbage to avoid the random sampling being the only instruction in a procedure *)
+    garbage <- 0; 
+    x <$ dbiased (mu D test);
+    return x;
+  }
+  proc g() : t = {
+    var l: t;
+    var garbage: int;
+    garbage <- 0;
+    l <$ D;
+    return l;
+  }
+}.
+
+op c : (bool * t) distr = dmap D (fun x => (test x, x)).
+      
+(* prove the lemma by directly providing a coupling *)
+lemma B_coupling:
+  equiv [ B.f ~ B.g : true ==> res{1} = test res{2} ].
+proof.
+proc.
+coupling c.
+wp; skip => /=; split; last first.
++ move => a b.
+  by rewrite /c => /supp_dmap [x] />.
+rewrite /iscoupling; split; last first.
++ by rewrite /c dmap_comp /(\o) /= dmap_id.
+rewrite /c dmap_comp /(\o) /=.
+apply eq_distr => b.
+rewrite dbiased1E.
+rewrite clamp_id.
++ exact mu_bounded.
+rewrite dmap1E /pred1 /(\o) /=.
+case b => _.
++ by apply mu_eq => x /#.
+rewrite (mu_eq _ _ (predC test)).
++ by smt().
+by rewrite mu_not D_ll.
+qed.
+
+(* prove the lemma by bypr, which means we compute the probability at eash side separately. *)
+lemma B_coupling_bypr:
+  equiv [ B.f ~ B.g : true ==> res{1} = test res{2} ].
+proof.
+bypr res{1} (test res{2}).
++ smt().
++ move => &1 &2 b.
+  have ->: Pr[B.f() @ &1 : res = b] = mu1 (dbiased (mu D test)) b.
+  + byphoare => //.
+    proc; rnd; wp; skip => />.
+  have -> //: Pr[B.g() @ &2 : test res = b] = mu1 (dbiased (mu D test)) b.
+  + byphoare => //.
+    proc; rnd; sp; skip => />.
+    rewrite dbiased1E.
+    rewrite clamp_id.
+    + exact mu_bounded.
+    case b => _.
+    + apply mu_eq => x /#.
+    rewrite (mu_eq _ _ (predC test)).
+    + by smt().
+    by rewrite mu_not D_ll.  
+qed.
+
+op f (x : t) : bool = test x.
+
+lemma B_compress:
+  equiv [ B.g ~ B.f : true ==> test res{1} = res{2} ].
+proof.
+proc.
+coupling{1} f.
+wp; skip => @/predT /=.
++ split.
+  + apply eq_distr => b.
+    rewrite dbiased1E.
+    rewrite clamp_id.
+    + exact mu_bounded.
+    rewrite dmap1E /pred1 /(\o) /=.
+    case b => _.
+    + by apply mu_eq => x /#.
+    rewrite (mu_eq _ _ (predC test)).
+    + by smt().
+    by rewrite mu_not D_ll.
+  + by move => @/f * //=.
+qed.
+
+lemma B_compress_reverse:
+  equiv [ B.f ~ B.g : true ==> test res{2} = res{1} ].
+proof.
+proc.
+coupling{2} f.
+wp; skip => @/predT /=.
++ split.
+  + apply eq_distr => b.
+    rewrite dbiased1E.
+    rewrite clamp_id.
+    + exact mu_bounded.
+    rewrite dmap1E /pred1 /(\o) /=.
+    case b => _.
+    + by apply mu_eq => x /#.
+    rewrite (mu_eq _ (fun (x : t) => f x = false) (predC test)).
+    + by smt().
+    by rewrite mu_not D_ll.
+  + by move => @/f //.
+qed.
diff --git a/theories/distributions/Distr.ec b/theories/distributions/Distr.ec
@@ -1437,6 +1437,26 @@ move=> a1 a2 _ /=; apply/iscpl_dlet; first by apply: cpl_d=> /#.
 by move=> b1 b2 _ /=; apply/iscpl_dunit.
 qed.
 
+(* -------------------------------------------------------------------- *)
+lemma iscpl_dmap1 ['a 'b] (d1 : 'a distr) (d2 : 'b distr) (f : 'a -> 'b) :
+  dmap d1 f = d2 => iscoupling<:'a, 'b> d1 d2 (dmap d1 (fun x => (x, f x))).
+proof.
+rewrite /iscoupling => ?.
+split.
++ by rewrite dmap_comp /(\o) /= dmap_id //.
++ by rewrite dmap_comp /(\o) /=.
+qed.
+
+(* -------------------------------------------------------------------- *)
+lemma iscpl_dmap2 ['a 'b] (d1 : 'a distr) (d2 : 'b distr) (f : 'b -> 'a) :
+  dmap d2 f = d1 => iscoupling<:'a, 'b> d1 d2 (dmap d2 (fun x => (f x, x))).
+proof.
+rewrite /iscoupling => ?.
+split.
++ by rewrite dmap_comp /(\o) /=.
++ by rewrite dmap_comp /(\o) /= dmap_id //.
+qed.
+
 (* -------------------------------------------------------------------- *)
 abbrev [-printing] dapply (F: 'a -> 'b) : 'a distr -> 'b distr =
   fun d => dmap d F.
