From 431a5b7635e4e5bb0aab4cad2199004a528fdc95 Mon Sep 17 00:00:00 2001
From: stefan-aws <stefanze@amazon.com>
Date: Mon, 23 Oct 2023 17:53:24 +0100
Subject: [PATCH] assumes

---
 .../DiscreteLaplace/Implementation.dfy         | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/Distributions/DiscreteLaplace/Implementation.dfy b/src/Distributions/DiscreteLaplace/Implementation.dfy
index 490bb650..7ef0d332 100644
--- a/src/Distributions/DiscreteLaplace/Implementation.dfy
+++ b/src/Distributions/DiscreteLaplace/Implementation.dfy
@@ -26,7 +26,7 @@ module DiscreteLaplace.Implementation {
       ghost var b' := b;
       ghost var y' := y;
       while b && y == 0
-        //invariant Model.SampleTailRecursive(scale, b', y')(old(s)) == Model.SampleTailRecursive(scale, b, y)(s)
+        invariant Model.SampleTailRecursive(scale, b', y')(old(s)) == Model.SampleTailRecursive(scale, b, y)(s)
         decreases *
       {
         b' := b;
@@ -37,32 +37,34 @@ module DiscreteLaplace.Implementation {
         label L2:
         var d := BernoulliExpNegSample(Rationals.Rational(u, scale.numer));
         assert (d, s) == Monad.Extract(BernoulliExpNeg.Model.Sample(Rationals.Rational(u, scale.numer))(old@L2(s)));
-        //assert Model.SampleTailRecursive(scale, b', y')(old@L1(s)) == Model.SampleTailRecursive(scale, b, y)(s);
         if !d {
+          assume {:axiom} false;
           continue;
         }
         label L3:
         var v := 0;
         var a := true;
         while a
-          invariant Model.SampleTailRecursiveHelper(scale)(old@L3(s)) == Model.SampleTailRecursiveHelper(scale, v, a)(s)
+          //invariant Model.SampleTailRecursiveHelper(scale)(old@L3(s)) == Model.SampleTailRecursiveHelper(scale, v, a)(s)
           decreases *
         {
+          label L4:
           a := BernoulliExpNegSample(Rationals.Int(1));
+          assert (a, s) == Monad.Extract(BernoulliExpNeg.Model.Sample(Rationals.Int(1))(old@L4(s)));
           if a {
             v := v + 1;
           }
         }
-        assert (v, s) == Monad.Extract(Model.SampleTailRecursiveHelper(scale)(old@L3(s)));
+        //assert (v, s) == Monad.Extract(Model.SampleTailRecursiveHelper(scale)(old@L3(s)));
         var x := u + scale.numer * v;
         y := x / scale.denom;
-        label L4:
+        label L5:
         b := BernoulliSample(Rationals.Rational(1, 2));
-        assert (b, s) == Monad.Extract(Bernoulli.Model.Sample(1, 2)(old@L4(s)));
-        //assume {:axiom} false; // add equivalence proof later
+        assert (b, s) == Monad.Extract(Bernoulli.Model.Sample(1, 2)(old@L5(s)));
+        assume Model.SampleTailRecursive(scale, b', y')(old(s)) == Model.SampleTailRecursive(scale, b, y)(s);
       }
+      assert Model.SampleTailRecursive(scale, b', y')(old(s)) == Model.SampleTailRecursive(scale, b, y)(s);
       z := if b then -y else y;
-      assume {:axiom} false; // add equivalence proof later
     }
 
   }