From 051a5db131960d9e492df369cf591ed583e03067 Mon Sep 17 00:00:00 2001
From: Etienne Kneuss <ekneuss@gmail.com>
Date: Fri, 9 Nov 2012 17:51:09 +0100
Subject: [PATCH] Generic collection mechanism

---
 src/main/scala/leon/purescala/TreeOps.scala | 23 ++++++++++++---------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/main/scala/leon/purescala/TreeOps.scala b/src/main/scala/leon/purescala/TreeOps.scala
index 54bc87f1d..dc2352d4a 100644
--- a/src/main/scala/leon/purescala/TreeOps.scala
+++ b/src/main/scala/leon/purescala/TreeOps.scala
@@ -1317,21 +1317,24 @@ object TreeOps {
       es.map(formulaSize).foldRight(0)(_ + _)+1
   }
 
-  def collectChooses(e: Expr): List[Choose] = {
-    def post(e: Expr, cs: List[Choose]) = {
-      val newCs = e match {
-        case c: Choose => c :: cs
-        case _ => cs
+  def collect[C](f: PartialFunction[Expr, C])(e: Expr): List[C] = {
+    def post(e: Expr, cs: List[C]) = {
+      if (f.isDefinedAt(e)) {
+        (e, f(e) :: cs)
+      } else {
+        (e, cs)
       }
-
-      (e, newCs)
     }
 
-    def combiner(cs: Seq[List[Choose]]) = {
-      cs.foldLeft(List[Choose]())(_ ::: _)
+    def combiner(cs: Seq[List[C]]) = {
+      cs.foldLeft(List[C]())(_ ::: _)
     }
 
-    genericTransform[List[Choose]]((_, _), post, combiner)(List())(e)._2
+    genericTransform[List[C]]((_, _), post, combiner)(List())(e)._2
+  }
+
+  def collectChooses(e: Expr): List[Choose] = {
+    collect({ case c: Choose => c })(e)
   }
 
   def valuateWithModel(model: Map[Identifier, Expr])(id: Identifier): Expr = {
-- 
GitLab