1 files changed, 16 insertions, 14 deletions

diff --git a/haddock-api/src/Haddock/ModuleTree.hs b/haddock-api/src/Haddock/ModuleTree.hs
index e6cf8201..a0be820a 100644
--- a/haddock-api/src/Haddock/ModuleTree.hs
+++ b/haddock-api/src/Haddock/ModuleTree.hs
@@ -20,39 +20,41 @@ import DynFlags      ( DynFlags )
 import Packages      ( lookupPackage )
 import PackageConfig ( sourcePackageIdString )
 
+import qualified Control.Applicative as A
 
-data ModuleTree = Node String Bool (Maybe String) (Maybe String) (Maybe (MDoc Name)) [ModuleTree]
+
+data ModuleTree = Node String (Maybe Module) (Maybe String) (Maybe String) (Maybe (MDoc Name)) [ModuleTree]
 
 
 mkModuleTree :: DynFlags -> Bool -> [(Module, Maybe (MDoc Name))] -> [ModuleTree]
 mkModuleTree dflags showPkgs mods =
-  foldr fn [] [ (splitModule mdl, modPkg mdl, modSrcPkg mdl, short) | (mdl, short) <- mods ]
+  foldr fn [] [ (mdl, splitModule mdl, modPkg mdl, modSrcPkg mdl, short) | (mdl, short) <- mods ]
   where
     modPkg mod_ | showPkgs = Just (unitIdString (moduleUnitId mod_))
                 | otherwise = Nothing
     modSrcPkg mod_ | showPkgs = fmap sourcePackageIdString
                                      (lookupPackage dflags (moduleUnitId mod_))
                    | otherwise = Nothing
-    fn (mod_,pkg,srcPkg,short) = addToTrees mod_ pkg srcPkg short
+    fn (m,mod_,pkg,srcPkg,short) = addToTrees mod_ m pkg srcPkg short
 
 
-addToTrees :: [String] -> Maybe String -> Maybe String -> Maybe (MDoc Name) -> [ModuleTree] -> [ModuleTree]
-addToTrees [] _ _ _ ts = ts
-addToTrees ss pkg srcPkg short [] = mkSubTree ss pkg srcPkg short
-addToTrees (s1:ss) pkg srcPkg short (t@(Node s2 leaf node_pkg node_srcPkg node_short subs) : ts)
-  | s1 >  s2  = t : addToTrees (s1:ss) pkg srcPkg short ts
-  | s1 == s2  = Node s2 (leaf || null ss) this_pkg this_srcPkg this_short (addToTrees ss pkg srcPkg short subs) : ts
-  | otherwise = mkSubTree (s1:ss) pkg srcPkg short ++ t : ts
+addToTrees :: [String] -> Module -> Maybe String -> Maybe String -> Maybe (MDoc Name) -> [ModuleTree] -> [ModuleTree]
+addToTrees [] _ _ _ _ ts = ts
+addToTrees ss m pkg srcPkg short [] = mkSubTree ss m pkg srcPkg short
+addToTrees (s1:ss) m pkg srcPkg short (t@(Node s2 leaf node_pkg node_srcPkg node_short subs) : ts)
+  | s1 >  s2  = t : addToTrees (s1:ss) m pkg srcPkg short ts
+  | s1 == s2  = Node s2 (leaf A.<|> (if null ss then Just m else Nothing)) this_pkg this_srcPkg this_short (addToTrees ss m pkg srcPkg short subs) : ts
+  | otherwise = mkSubTree (s1:ss) m pkg srcPkg short ++ t : ts
  where
   this_pkg = if null ss then pkg else node_pkg
   this_srcPkg = if null ss then srcPkg else node_srcPkg
   this_short = if null ss then short else node_short
 
 
-mkSubTree :: [String] -> Maybe String -> Maybe String -> Maybe (MDoc Name) -> [ModuleTree]
-mkSubTree []     _   _      _     = []
-mkSubTree [s]    pkg srcPkg short = [Node s True pkg srcPkg short []]
-mkSubTree (s:ss) pkg srcPkg short = [Node s (null ss) Nothing Nothing Nothing (mkSubTree ss pkg srcPkg short)]
+mkSubTree :: [String] -> Module -> Maybe String -> Maybe String -> Maybe (MDoc Name) -> [ModuleTree]
+mkSubTree []     _ _   _      _     = []
+mkSubTree [s]    m pkg srcPkg short = [Node s (Just m) pkg srcPkg short []]
+mkSubTree (s:s':ss) m pkg srcPkg short = [Node s Nothing Nothing Nothing Nothing (mkSubTree (s':ss) m pkg srcPkg short)]
 
 
 splitModule :: Module -> [String]