{-# LANGUAGE TypeOperators #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} module Annotations.MultiRec.Annotated ( AnnFix, AnnFix1, mkAnnFix, AnyAnnFix, unannotate, children, flatten, filterAnnFix, debugFlatten, allAnnotations, AnnZipper, focusAnn, explore, findLeftmostDeepest ) where import Annotations.ExploreHints import Annotations.MultiRec.ShowFam import Annotations.MultiRec.Any import Control.Monad.Writer (Writer, execWriter, tell) import Generics.MultiRec hiding (show) import Generics.MultiRec.HFix import Annotations.MultiRec.Zipper import Annotations.MultiRec.ZipperFix import Data.Maybe -- | A fully annotated tree. type AnnFix x s = HFix (K x :*: PF s) -- | A functor with fully annotated children. type AnnFix1 x s = (PF s) (AnnFix x s) -- | Supply a tree with an annotation at the top level. mkAnnFix :: x -> AnnFix1 x s ix -> AnnFix x s ix mkAnnFix x = HIn . (K x :*:) -- | A fixpoint of a data family @s@ annotated with an @x@ at every recursive position, with existentially quantified top-level index. type AnyAnnFix x s = AnyF s (AnnFix x s) -- | Removes all annotations from a recursively annotated fixpoint. unannotate :: HFunctor s (PF s) => s ix -> AnnFix x s ix -> HFix (PF s) ix unannotate p = HIn . hmap unannotate p . snd' . hout -- fst' :: (f :*: g) r ix -> f r ix -- fst' (x :*: _) = x snd' :: (f :*: g) r ix -> g r ix snd' (_ :*: y) = y -- | Collects the direct children of a functor in a list. children :: HFunctor s f => s ix -> f r ix -> [AnyF s r] children p x = execWriter (hmapM collect p x) where collect :: s ix -> r ix -> Writer [AnyF s r] (r ix) collect w x = tell [AnyF w x] >> return x -- | Flatten an annotated tree to a list of subtrees coupled with their annotations. flatten :: forall s x ix. (HFunctor s (PF s), Fam s) => s ix -> AnnFix x s ix -> [(x, Any s)] flatten p tree@(HIn (K x :*: y)) = (x, Any p (hto p (unannotate p tree :: HFix (PF s) ix))) : concatMap (flatten $?) (children p y) -- | Yield all subtrees whose annotation matches the predicate. filterAnnFix :: (Fam s, HFunctor s (PF s)) => s ix -> (x -> Bool) -> AnnFix x s ix -> [(x, Any s)] filterAnnFix p f = filter (f . fst) . flatten p -- | Flatten an annotated tree and print all subtrees to stdout. debugFlatten :: (HFunctor s (PF s), ShowFam s, Show x, Fam s) => s ix -> AnnFix x s ix -> IO () debugFlatten p = putStr . unlines . map show . flatten p -- | Recursively yield all annotations in the tree in preorder. allAnnotations :: HFunctor phi (PF phi) => phi ix -> AnnFix x phi ix -> [x] allAnnotations p y = f (AnyF p y) where f (AnyF p' (HIn (K x :*: y))) = x : concatMap f (children p' y) -- | Extract the annotation of the current focus. focusAnn :: Loc phi f (HFix (K x :*: g)) ix -> x focusAnn = on (\_ (HIn (K x :*: _)) -> x) type AnnZipper phi x = FixZipper phi (K x :*: PF phi) -- | Explore an annotated tree. Starting with the root of the tree, at each -- position the annotation at that position is matched against the -- 'ExploreHints' predicates and all the selections where 'matchHere' was -- positive are collected. The 'exploreRight' and 'exploreDown' allow pruning -- of the tree, preventing entire parts from being visited. explore :: Zipper phi (PF phi) => phi ix -> (x -> ExploreHints) -> (AnnFix x phi) ix -> [AnnZipper phi x ix] explore p hints = explore' hints . enter p explore' :: Zipper phi (PF phi) => (x -> ExploreHints) -> AnnZipper phi x ix -> [AnnZipper phi x ix] explore' hints root = [ z | (dirOk, zs) <- dirs, dirOk (hints x), z <- zs ] where x = focusAnn root dirs = [ (matchHere, [root]) , (exploreDown, exploreMore (down root)) , (exploreRight, exploreMore (right root)) ] exploreMore = maybe [] (explore' hints) -- | Find the deepest node in an annotated tree that matches the predicate. -- Starting with the root, the predicate tells whether a node's annotation -- matches. If so, the search continues at the node's children and the node's -- siblings to the right are excluded from further exploration. If no child -- matches, the node itself is returned. findLeftmostDeepest :: (Zipper phi (PF phi)) => phi ix -> (x -> Bool) -> AnnFix x phi ix -> Maybe (AnnZipper phi x ix) findLeftmostDeepest p down = listToMaybe . reverse . explore p hints where hints x | down x = ExploreHints True True False | otherwise = ExploreHints False False True