Improve Syb code (#621)

Specialize.hs and Ast.hs are modified to have their Syb code not recurse into
Name or Id in HsSyn types.

Specialize.hs is refactored to have fewer calls to Syb functions.

Syb.hs has some foldl calls replaced with foldl' calls.

There is still a lot of performance on the floor of Ast.hs. The RenamedSource
is traversed many times, and lookupBySpan is very inefficient. everywhereBut and
lookupBySpan dominate the runtime whenever --hyperlinked-source is passed.

3 files changed, 88 insertions, 53 deletions

diff --git a/haddock-api/src/Haddock/Backends/Hyperlinker/Ast.hs b/haddock-api/src/Haddock/Backends/Hyperlinker/Ast.hs
index b97f0ead..78beacf2 100644
--- a/haddock-api/src/Haddock/Backends/Hyperlinker/Ast.hs
+++ b/haddock-api/src/Haddock/Backends/Hyperlinker/Ast.hs
@@ -2,12 +2,12 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE TypeFamilies #-}
-
+{-# LANGUAGE TypeApplications #-}
 
 module Haddock.Backends.Hyperlinker.Ast (enrich) where
 
 
-import Haddock.Syb
+import qualified Haddock.Syb as Syb
 import Haddock.Backends.Hyperlinker.Types
 
 import qualified GHC
@@ -16,6 +16,9 @@ import Control.Applicative
 import Data.Data
 import Data.Maybe
 
+everythingInRenamedSource :: (Alternative f, Data x)
+  => (forall a. Data a => a -> f r) -> x -> f r
+everythingInRenamedSource f = Syb.everythingButType @GHC.Name (<|>) f
 
 -- | Add more detailed information to token stream using GHC API.
 enrich :: GHC.RenamedSource -> [Token] -> [RichToken]
@@ -53,7 +56,7 @@ enrichToken _ _ = Nothing
 -- | Obtain details map for variables ("normally" used identifiers).
 variables :: GHC.RenamedSource -> DetailsMap
 variables =
-    everything (<|>) (var `combine` rec)
+    everythingInRenamedSource (var `Syb.combine` rec)
   where
     var term = case cast term of
         (Just (GHC.L sspan (GHC.HsVar name))) ->
@@ -68,8 +71,7 @@ variables =
 
 -- | Obtain details map for types.
 types :: GHC.RenamedSource -> DetailsMap
-types =
-    everything (<|>) ty
+types = everythingInRenamedSource ty
   where
     ty term = case cast term of
         (Just (GHC.L sspan (GHC.HsTyVar _ name))) ->
@@ -81,9 +83,10 @@ types =
 -- That includes both identifiers bound by pattern matching or declared using
 -- ordinary assignment (in top-level declarations, let-expressions and where
 -- clauses).
+
 binds :: GHC.RenamedSource -> DetailsMap
-binds =
-    everything (<|>) (fun `combine` pat `combine` tvar)
+binds = everythingInRenamedSource
+      (fun `Syb.combine` pat `Syb.combine` tvar)
   where
     fun term = case cast term of
         (Just (GHC.FunBind (GHC.L sspan name) _ _ _ _ :: GHC.HsBind GHC.Name)) ->
@@ -93,7 +96,7 @@ binds =
         (Just (GHC.L sspan (GHC.VarPat name))) ->
             pure (sspan, RtkBind (GHC.unLoc name))
         (Just (GHC.L _ (GHC.ConPatIn (GHC.L sspan name) recs))) ->
-            [(sspan, RtkVar name)] ++ everything (<|>) rec recs
+            [(sspan, RtkVar name)] ++ everythingInRenamedSource rec recs
         (Just (GHC.L _ (GHC.AsPat (GHC.L sspan name) _))) ->
             pure (sspan, RtkBind name)
         _ -> empty
@@ -112,8 +115,8 @@ binds =
 decls :: GHC.RenamedSource -> DetailsMap
 decls (group, _, _, _) = concatMap ($ group)
     [ concat . map typ . concat . map GHC.group_tyclds . GHC.hs_tyclds
-    , everything (<|>) fun . GHC.hs_valds
-    , everything (<|>) (con `combine` ins)
+    , everythingInRenamedSource fun . GHC.hs_valds
+    , everythingInRenamedSource (con `Syb.combine` ins)
     ]
   where
     typ (GHC.L _ t) = case t of
@@ -127,7 +130,8 @@ decls (group, _, _, _) = concatMap ($ group)
         _ -> empty
     con term = case cast term of
         (Just cdcl) ->
-            map decl (GHC.getConNames cdcl) ++ everything (<|>) fld cdcl
+            map decl (GHC.getConNames cdcl)
+              ++ everythingInRenamedSource fld cdcl
         Nothing -> empty
     ins term = case cast term of
         (Just (GHC.DataFamInstD inst)) -> pure . tyref $ GHC.dfid_tycon inst
@@ -149,7 +153,7 @@ decls (group, _, _, _) = concatMap ($ group)
 -- import lists.
 imports :: GHC.RenamedSource -> DetailsMap
 imports src@(_, imps, _, _) =
-    everything (<|>) ie src ++ mapMaybe (imp . GHC.unLoc) imps
+    everythingInRenamedSource ie src ++ mapMaybe (imp . GHC.unLoc) imps
   where
     ie term = case cast term of
         (Just (GHC.IEVar v)) -> pure $ var $ GHC.ieLWrappedName v
diff --git a/haddock-api/src/Haddock/Interface/Specialize.hs b/haddock-api/src/Haddock/Interface/Specialize.hs
index 28bbf305..8c28cd5a 100644
--- a/haddock-api/src/Haddock/Interface/Specialize.hs
+++ b/haddock-api/src/Haddock/Interface/Specialize.hs
@@ -1,9 +1,9 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE Rank2Types #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE RecordWildCards #-}
 
-
 module Haddock.Interface.Specialize
     ( specializeInstHead
     ) where
@@ -27,73 +27,66 @@ import Data.Map (Map)
 import qualified Data.Map as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
-
-
--- | Instantiate all occurrences of given name with particular type.
-specialize :: (Eq name, Typeable name)
-           => Data a
-           => name -> HsType name -> a -> a
-specialize name details =
-    everywhere $ mkT step
-  where
-    step (HsTyVar _ (L _ name')) | name == name' = details
-    step typ = typ
-
+import Data.Foldable
 
 -- | Instantiate all occurrences of given names with corresponding types.
---
--- It is just a convenience function wrapping 'specialize' that supports more
--- that one specialization.
-specialize' :: (Eq name, Typeable name)
+specialize :: forall name a. (Ord name, DataId name, NamedThing name)
             => Data a
             => [(name, HsType name)] -> a -> a
-specialize' = flip $ foldr (uncurry specialize)
+specialize specs = go
+  where
+    go :: forall x. Data x => x -> x
+    go = everywhereButType @name $ mkT $ sugar . specialize_ty_var
+    specialize_ty_var (HsTyVar _ (L _ name'))
+      | Just t <- Map.lookup name' spec_map = t
+    specialize_ty_var typ = typ
+    -- This is a tricky recursive definition that is guaranteed to terminate
+    -- because a type binder cannot be instantiated with a type that depends
+    -- on that binder. i.e. @a -> Maybe a@ is invalid
+    spec_map = Map.fromList [ (n, go t) | (n, t) <- specs]
 
 
 -- | Instantiate given binders with corresponding types.
 --
 -- Again, it is just a convenience function around 'specialize'. Note that
 -- length of type list should be the same as the number of binders.
-specializeTyVarBndrs :: (Eq name, DataId name)
+specializeTyVarBndrs :: (Ord name, DataId name, NamedThing name)
                      => Data a
                      => LHsQTyVars name -> [HsType name]
                      -> a -> a
 specializeTyVarBndrs bndrs typs =
-    specialize' $ zip bndrs' typs
+    specialize $ zip bndrs' typs
   where
     bndrs' = map (bname . unLoc) . hsq_explicit $ bndrs
     bname (UserTyVar (L _ name)) = name
     bname (KindedTyVar (L _ name) _) = name
 
 
-specializePseudoFamilyDecl :: (Eq name, DataId name)
+specializePseudoFamilyDecl :: (Ord name, DataId name, NamedThing name)
                            => LHsQTyVars name -> [HsType name]
                            -> PseudoFamilyDecl name
                            -> PseudoFamilyDecl name
 specializePseudoFamilyDecl bndrs typs decl =
-    decl { pfdTyVars = map specializeTyVars (pfdTyVars decl) }
-  where
-    specializeTyVars = specializeTyVarBndrs bndrs typs
-
+  decl {pfdTyVars = map (specializeTyVarBndrs bndrs typs) (pfdTyVars decl)}
 
-specializeSig :: forall name . (Eq name, DataId name, SetName name)
+specializeSig :: forall name . (Ord name, DataId name, SetName name, NamedThing name)
               => LHsQTyVars name -> [HsType name]
               -> Sig name
               -> Sig name
 specializeSig bndrs typs (TypeSig lnames typ) =
-    TypeSig lnames (typ { hswc_body = (hswc_body typ) { hsib_body = noLoc typ'}})
+  TypeSig lnames (typ {hswc_body = (hswc_body typ) {hsib_body = noLoc typ'}})
   where
     true_type :: HsType name
     true_type = unLoc (hsSigWcType typ)
     typ' :: HsType name
-    typ' = rename fv . sugar $ specializeTyVarBndrs bndrs typs true_type
+    typ' = rename fv $ specializeTyVarBndrs bndrs typs true_type
     fv = foldr Set.union Set.empty . map freeVariables $ typs
 specializeSig _ _ sig = sig
 
 
 -- | Make all details of instance head (signatures, associated types)
 -- specialized to that particular instance type.
-specializeInstHead :: (Eq name, DataId name, SetName name)
+specializeInstHead :: (Ord name, DataId name, SetName name, NamedThing name)
                    => InstHead name -> InstHead name
 specializeInstHead ihd@InstHead { ihdInstType = clsi@ClassInst { .. }, .. } =
     ihd { ihdInstType = instType' }
@@ -115,12 +108,7 @@ specializeInstHead ihd = ihd
 -- and @(a, b, c)@.
 sugar :: forall name. (NamedThing name, DataId name)
       => HsType name -> HsType name
-sugar =
-    everywhere $ mkT step
-  where
-    step :: HsType name -> HsType name
-    step = sugarOperators . sugarTuples . sugarLists
-
+sugar = sugarOperators . sugarTuples . sugarLists
 
 sugarLists :: NamedThing name => HsType name -> HsType name
 sugarLists (HsAppTy (L _ (HsTyVar _ (L _ name))) ltyp)
diff --git a/haddock-api/src/Haddock/Syb.hs b/haddock-api/src/Haddock/Syb.hs
index 4847e486..7e34ae8c 100644
--- a/haddock-api/src/Haddock/Syb.hs
+++ b/haddock-api/src/Haddock/Syb.hs
@@ -1,8 +1,11 @@
 {-# LANGUAGE Rank2Types #-}
-
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 module Haddock.Syb
-    ( everything, everythingWithState, everywhere
+    ( everything, everythingButType, everythingWithState
+    , everywhere, everywhereButType
     , mkT
     , combine
     ) where
@@ -10,16 +13,41 @@ module Haddock.Syb
 
 import Data.Data
 import Control.Applicative
+import Data.Maybe
+import Data.Foldable
 
+-- | Returns true if a == t.
+-- requires AllowAmbiguousTypes
+isType :: forall a b. (Typeable a, Typeable b) => b -> Bool
+isType _ = isJust $ eqT @a @b
 
 -- | Perform a query on each level of a tree.
 --
 -- This is stolen directly from SYB package and copied here to not introduce
 -- additional dependencies.
-everything :: (r -> r -> r) -> (forall a. Data a => a -> r)
+everything :: (r -> r -> r)
+           -> (forall a. Data a => a -> r)
            -> (forall a. Data a => a -> r)
-everything k f x = foldl k (f x) (gmapQ (everything k f) x)
+everything k f x = foldl' k (f x) (gmapQ (everything k f) x)
+
+-- | Variation of "everything" with an added stop condition
+-- Just like 'everything', this is stolen from SYB package.
+everythingBut :: (r -> r -> r)
+              -> (forall a. Data a => a -> (r, Bool))
+              -> (forall a. Data a => a -> r)
+everythingBut k f x = let (v, stop) = f x
+                      in if stop
+                           then v
+                           else foldl' k v (gmapQ (everythingBut k f) x)
 
+-- | Variation of "everything" that does not recurse into children of type t
+-- requires AllowAmbiguousTypes
+everythingButType ::
+  forall t r. (Typeable t)
+  => (r -> r -> r)
+  -> (forall a. Data a => a -> r)
+  -> (forall a. Data a => a -> r)
+everythingButType k f = everythingBut k $ (,) <$> f <*> isType @t
 
 -- | Perform a query with state on each level of a tree.
 --
@@ -31,8 +59,7 @@ everythingWithState :: s -> (r -> r -> r)
                     -> (forall a. Data a => a -> r)
 everythingWithState s k f x =
     let (r, s') = f x s
-    in foldl k r (gmapQ (everythingWithState s' k f) x)
-
+    in foldl' k r (gmapQ (everythingWithState s' k f) x)
 
 -- | Apply transformation on each level of a tree.
 --
@@ -40,6 +67,22 @@ everythingWithState s k f x =
 everywhere :: (forall a. Data a => a -> a) -> (forall a. Data a => a -> a)
 everywhere f = f . gmapT (everywhere f)
 
+-- | Variation on everywhere with an extra stop condition
+-- Just like 'everything', this is stolen from SYB package.
+everywhereBut :: (forall a. Data a => a -> Bool)
+  -> (forall a. Data a => a -> a)
+  -> (forall a. Data a => a -> a)
+everywhereBut q f x
+    | q x       = x
+    | otherwise = f (gmapT (everywhereBut q f) x)
+
+-- | Variation of "everywhere" that does not recurse into children of type t
+-- requires AllowAmbiguousTypes
+everywhereButType :: forall t . (Typeable t)
+  => (forall a. Data a => a -> a)
+  -> (forall a. Data a => a -> a)
+everywhereButType = everywhereBut (isType @t)
+
 -- | Create generic transformation.
 --
 -- Another function stolen from SYB package.