From ff2043d0477ec7c87834a8f70f0bd8749f4d9d6e Mon Sep 17 00:00:00 2001
From: Samuel Chassot <samuel.chassot@epfl.ch>
Date: Fri, 11 Oct 2024 16:25:59 +0200
Subject: [PATCH] add flatmap postcondition

---
 frontends/library/stainless/lang/Set.scala | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/frontends/library/stainless/lang/Set.scala b/frontends/library/stainless/lang/Set.scala
index 8da2ba422..f9e491621 100644
--- a/frontends/library/stainless/lang/Set.scala
+++ b/frontends/library/stainless/lang/Set.scala
@@ -36,15 +36,10 @@ object Set {
       new Set(set.theSet.map(f))
     }
 
-    @extern @pure
-    def flatMap[B](f: A => Set[B]): Set[B] = {
-      new Set(set.theSet.flatMap(f(_).theSet))
-    }
-
     @extern @pure
     def mapPost1[B](f: A => B)(a: A): Unit = {
       ()
-   }.ensuring { _ =>
+    }.ensuring { _ =>
       !set.contains(a) || map[B](f).contains(f(a))
     }
    
@@ -52,10 +47,22 @@ object Set {
     def mapPost2[B](f: A => B)(b: B): A = {
       require(map[B](f).contains(b))
       (??? : A)
-   }.ensuring { (a:A) =>
+    }.ensuring { (a:A) =>
       b == f(a) && set.contains(a)
     }
 
+    @extern @pure
+    def flatMap[B](f: A => Set[B]): Set[B] = {
+      new Set(set.theSet.flatMap(f(_).theSet))
+    }
+
+    @extern @pure
+    def flatMapPost[B](f: A => B)(b: B) = {
+      (??? : A)
+    }.ensuring { _ =>
+      flatMap(f).contains(b) == set.exists(a => f(a).contains(b))
+    }
+
     @extern @pure
     def filter(p: A => Boolean): Set[A] = {
       new Set(set.theSet.filter(p))