Make GenNavAST generate instances of a "Polyplate" class.

This provides transformM, which we can build applyX (etc.) on top of:
it takes a set of functions, and applies them wherever they could
match in a data structure, without automatically recursing. This is
done using a four(!)-argument typeclass, drawing inspiration from
Biplate.

The resulting 25,000-line set of instances takes a little while to
compile...

pass/Traversal.hs

@@ -28,6 +28,7 @@ module Traversal (
 import Data.Generics
 
 import GenericUtils
+import NavAST
 import Pass
 
 -- | A transformation for a single 'Data' type with explicit descent.

pregen/GenNavAST.hs

@@ -19,7 +19,9 @@ with this program.  If not, see <http://www.gnu.org/licenses/>.
 -- | Utilities for metaprogramming.
 module GenNavAST where
 
+import Data.Char
 import Data.Generics
+import Data.List
 import qualified Data.Set as Set
 
 import PregenUtils
@@ -28,7 +30,10 @@ import Utils
 header :: [String]
 header
     = [ "{-# OPTIONS_GHC -Werror -fwarn-overlapping-patterns -fwarn-unused-matches -fwarn-unused-binds -fwarn-incomplete-patterns #-}"
-      , "-- | Instances that allow the AST to be navigated efficiently."
+      , "-- | Type class and instances for transformations upon the AST."
+      , "--"
+      , "-- This was inspired by Neil Mitchell's Biplate class."
+      , "--"
       , "-- NOTE: This file is auto-generated by the GenNavAST program, "
       , "-- and should not be edited directly."
       , ""
@@ -37,38 +42,114 @@ header
       , "import qualified AST"
       , "import qualified Metadata"
       , ""
-      , "data Navigation = Hit | Through | Miss"
+      , "class Monad m => Polyplate m o o0 t where"
-      , ""
+      , "  transformM :: o -> o0 -> Bool -> t -> m t"
-      , "class Navigable f t where"
-      , "  navigate :: f -> t -> Navigation"
       , ""
       ]
 
-instancesFrom :: Data t => t -> [String]
+-- | Instances for a particular data type (i.e. where that data type is the
+-- last argument to 'Polyplate').
+instancesFrom :: forall t. Data t => t -> [String]
 instancesFrom w
-    = concat [inst c | DataBox c <- justBoxes $ astTypeMap]
+    = baseInst ++
+      concat [otherInst c | DataBox c <- justBoxes $ astTypeMap]
   where
     wName = show $ typeOf w
     wKey = typeKey w
+    wDType = dataTypeOf w
+    wCtrs = if isAlgType wDType then dataTypeConstrs wDType else []
+
+    -- The module prefix of this type, so we can use it in constructor names.
+    modPrefix
+        = if '.' `elem` (takeWhile (\c -> isAlphaNum c || c == '.') wName)
+            then takeWhile (/= '.') wName ++ "."
+            else ""
+
+    ctrArgs ctr
+        = gmapQ DataBox (fromConstr ctr :: t)
+    ctrArgTypes ctr
+        = [show $ typeOf w | DataBox w <- ctrArgs ctr]
+
+    -- | An instance that describes what to do when we have no transformations
+    -- left to apply.
+    baseInst :: [String]
+    baseInst
+        = [ "instance (" ++ concat (intersperse ", " context) ++ ") =>"
+          , "         Polyplate m () o0 (" ++ wName ++ ") where"
+          ] ++
+          (if isAlgType wDType
+            -- An algebraic type: apply to each child if we're following.
+            then ["  transformM () _ False v = return v"] ++
+                 (concatMap constrCase wCtrs)
+            -- A primitive type: just return it.
+            else ["  transformM () _ _ v = return v"]) ++
+          [""]
+
+    -- | Class context for 'baseInst'.
+    -- We need an instance of Polyplate for each of the types contained within
+    -- this type, so we can recurse into them.
+    context :: [String]
+    context
+        = ["Monad m"] ++
+          ["Polyplate m o0 o0 (" ++ argType ++ ")"
+           | argType <- nub $ sort $ concatMap ctrArgTypes wCtrs]
+
+    -- | A 'transformM' case for a particular constructor of this (algebraic)
+    -- data type: pull the value apart, apply 'transformM' to each part of it,
+    -- then stick it back together.
+    constrCase :: Constr -> [String]
+    constrCase ctr
+        = [ "  transformM () " ++ (if argNums == [] then "_" else "ops") ++
+            " True (" ++ ctrInput ++ ")"
+          , "    = do"
+          ] ++
+          [ "         r" ++ show i ++ " <- transformM ops ops False a" ++ show i
+           | i <- argNums] ++
+          [ "         return (" ++ ctrResult ++ ")"
+          ]
+      where
+        (isTuple, argNums)
+          -- FIXME: Should work for 3+-tuples too
+          | ctrS == "(,)" = (True, [0 .. 1])
+          | otherwise     = (False, [0 .. ((length $ ctrArgs ctr) - 1)])
+        ctrS = show ctr
+        ctrName = modPrefix ++ ctrS
+        makeCtr vs
+            = if isTuple
+                then "(" ++ (concat $ intersperse ", " vs) ++ ")"
+                else ctrName ++ concatMap (" " ++) vs
+        ctrInput = makeCtr ["a" ++ show i | i <- argNums]
+        ctrResult = makeCtr ["r" ++ show i | i <- argNums]
 
     containedKeys = Set.fromList [typeKey c
                                   | DataBox c <- justBoxes $ findTypesIn w]
 
-    inst c
+    -- | An instance that describes how to apply -- or not apply -- a
-        = [ "instance Navigable (" ++ wName ++ ") (" ++ cName ++ ") where"
+    -- transformation.
-          , "  navigate _ _ = " ++ result
+    otherInst c
+        = [ "instance (Monad m, Polyplate m r o0 (" ++ wName ++ ")) =>"
+          , "         Polyplate m ((" ++ cName ++ ") -> m (" ++ cName ++ "), r)"
+          , "                   o0 (" ++ wName ++ ") where"
+          , impl
           , ""
           ]
       where
         cName = show $ typeOf c
         cKey = typeKey c
-        result
+        impl
-          | wKey == cKey                    = "Hit"
+          -- This type matches the transformation: apply it.
-          | cKey `Set.member` containedKeys = "Through"
+          | wKey == cKey
-          | otherwise                       = "Miss"
+            = "  transformM (f, _) _ _ v = f v"
+          -- This type might contain the type that the transformation acts
+          -- upon: set the flag to say we need to recurse into it.
+          | cKey `Set.member` containedKeys
+            = "  transformM (_, rest) ops _ v = transformM rest ops True v"
+          -- This type can't contain the transformed type; just move on to the
+          -- next transformation.
+          | otherwise
+            = "  transformM (_, rest) ops b v = transformM rest ops b v"
 
 main :: IO ()
 main = putStr $ unlines $ header ++
                           concat [instancesFrom w
                                   | DataBox w <- justBoxes $ astTypeMap]
-