From b6176ee615036ec1ea38ac06282bc6973f42dfc3 Mon Sep 17 00:00:00 2001
From: Neil Brown <neil@twistedsquare.com>
Date: Tue, 2 Dec 2008 14:10:43 +0000
Subject: [PATCH] Added functions and schemes related to the new query
 functions

---
 polyplate/Data/Generics/Polyplate.hs         | 32 ++++++++++++++++++--
 polyplate/Data/Generics/Polyplate/Schemes.hs | 26 ++++++++++++++--
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/polyplate/Data/Generics/Polyplate.hs b/polyplate/Data/Generics/Polyplate.hs
index a396529..3b62eb0 100644
--- a/polyplate/Data/Generics/Polyplate.hs
+++ b/polyplate/Data/Generics/Polyplate.hs
@@ -30,8 +30,11 @@ with this program.  If not, see <http://www.gnu.org/licenses/>.
 module Data.Generics.Polyplate (PolyplateM(..), Polyplate(..), PolyplateSpine(..),
   makeRecurseM, RecurseM, makeRecurse, Recurse,
   makeDescendM, DescendM, makeDescend, Descend,
+  makeRecurseQ, RecurseQ,
+  makeDescendQ, DescendQ,
   BaseOp, baseOp,
-  ExtOpM, extOpM, ExtOp, extOp, OneOpM, OneOp, TwoOpM, TwoOp) where
+  ExtOpM, extOpM, ExtOp, extOp, OneOpM, OneOp, TwoOpM, TwoOp,
+  ExtOpQ, extOpQ, OneOpQ, TwoOpQ) where
 
 import Control.Monad.Identity
 
@@ -155,6 +158,14 @@ type Descend opT = forall t. Polyplate t () opT => t -> t
 makeDescend :: opT -> Descend opT
 makeDescend ops = transform () ops
 
+type RecurseQ a opQ = forall t. PolyplateSpine t opQ () a => t -> Tree (Maybe a)
+type DescendQ a opQ = forall t. PolyplateSpine t () opQ a => t -> Tree (Maybe a)
+
+makeRecurseQ :: opQ -> RecurseQ a opQ
+makeRecurseQ ops = transformSpine ops ()
+
+makeDescendQ :: opQ -> DescendQ a opQ
+makeDescendQ ops = transformSpine () ops
 
 -- | The type of the empty set of operations
 type BaseOp = ()
@@ -173,6 +184,10 @@ type ExtOpM m opT t = (t -> m t, opT)
 --  You cannot mix monadic and non-monadic operations in the same list.
 type ExtOp opT t = (t -> t, opT)
 
+-- | The type that extends a query-ops set to be applied to the given type (t).
+-- Not to be mixed with modification operations.
+type ExtOpQ a opQ t = (t -> a, opQ)
+
 -- | The function that extends an ops set (opT) in the given monad (m) to be applied to
 -- the given type (t).  You cannot mix monadic and non-monadic operations in the
 -- same list.
@@ -185,13 +200,24 @@ extOpM ops f = (f, ops)
 extOp :: opT -> (t -> t) -> ExtOp opT t
 extOp ops f = (f, ops)
 
--- | A handy synonym for an ops set with only one item.
+-- | The function that extends a query-ops set to be applied to the given type (t).
+-- Not to be mixed with modification operations.
+extOpQ :: opQ -> (t -> a) -> ExtOpQ a opQ t
+extOpQ ops f = (f, ops)
+
+
+-- | A handy synonym for a monadic ops set with only one item.
 type OneOpM m t = ExtOpM m BaseOp t
 -- | A handy synonym for an ops set with only one item.
 type OneOp t = ExtOp BaseOp t
+-- | A handy synonym for a query ops set with only one item.
+type OneOpQ a t = ExtOpQ a BaseOp t
 
--- | A handy synonym for an ops set with only two items.
+-- | A handy synonym for a monadic ops set with only two items.
 type TwoOpM m s t = ExtOpM m (ExtOpM m BaseOp s) t
 -- | A handy synonym for an ops set with only two items.
 type TwoOp s t = ExtOp (ExtOp BaseOp s) t
+-- | A handy synonym for a monadic ops set with only two items.
+type TwoOpQ a s t = ExtOpQ a (ExtOpQ a BaseOp s) t
+
 
diff --git a/polyplate/Data/Generics/Polyplate/Schemes.hs b/polyplate/Data/Generics/Polyplate/Schemes.hs
index 8f1ca10..0c3bd16 100644
--- a/polyplate/Data/Generics/Polyplate/Schemes.hs
+++ b/polyplate/Data/Generics/Polyplate/Schemes.hs
@@ -19,6 +19,8 @@ with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 module Data.Generics.Polyplate.Schemes where
 
+import Data.Tree
+
 import Data.Generics.Polyplate
 
 -- | Given a list of operations and a modifier function, augments that modifier
@@ -60,8 +62,28 @@ makeCheckM ops f v
     descend = makeDescend ops
 -}
 
--- TODO also add a listify-like thing (maybe return a rose tree):
--- applyQuery :: (s -> a) -> t -> Tree a
+
+applyQuery :: PolyplateSpine t (OneOpQ a s) () a => (s -> a) -> t -> Tree (Maybe a)
+applyQuery qf = makeRecurseQ ops
+  where
+    ops = baseOp `extOpQ` qf
+
+applyQuery2 :: PolyplateSpine t (TwoOpQ a sA sB) () a => (sA -> a) -> (sB -> a) -> t -> Tree (Maybe a)
+applyQuery2 qfA qfB = makeRecurseQ ops
+  where
+    ops = baseOp `extOpQ` qfA `extOpQ` qfB 
+
+applyListify :: PolyplateSpine t (OneOpQ (Maybe s) s) () (Maybe s) => (s -> Bool) -> t -> Tree (Maybe s)
+applyListify qf = fmap joinMaybe . makeRecurseQ ops
+  where
+    qf' x = if qf x then Just x else Nothing
+
+    joinMaybe :: Maybe (Maybe a) -> Maybe a
+    joinMaybe Nothing = Nothing
+    joinMaybe (Just mx) = mx
+    
+    ops = baseOp `extOpQ` qf'
+
 
 -- | Given a monadic function that applies to a particular type (s), automatically
 -- applies that function to every instance of s in a larger structure of type t,