Fix after unit refactoring

1 files changed, 6 insertions, 7 deletions

diff --git a/haddock-api/src/Haddock/ModuleTree.hs b/haddock-api/src/Haddock/ModuleTree.hs
index 688e3e71..d0a39322 100644
--- a/haddock-api/src/Haddock/ModuleTree.hs
+++ b/haddock-api/src/Haddock/ModuleTree.hs
@@ -14,10 +14,9 @@ module Haddock.ModuleTree ( ModuleTree(..), mkModuleTree ) where
 
 import Haddock.Types ( MDoc )
 
-import GHC                 ( Name )
-import GHC.Unit.Module    ( Module, moduleNameString, moduleName, moduleUnit, unitString )
-import GHC.Driver.Session  ( DynFlags )
-import GHC.Unit.State ( lookupUnit, unitPackageIdString )
+import GHC             ( Name )
+import GHC.Unit.Module ( Module, moduleNameString, moduleName, moduleUnit, unitString )
+import GHC.Unit.State  ( UnitState, lookupUnit, unitPackageIdString )
 
 import qualified Control.Applicative as A
 
@@ -25,14 +24,14 @@ import qualified Control.Applicative as A
 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 =
+mkModuleTree :: UnitState -> Bool -> [(Module, Maybe (MDoc Name))] -> [ModuleTree]
+mkModuleTree state showPkgs mods =
   foldr fn [] [ (mdl, splitModule mdl, modPkg mdl, modSrcPkg mdl, short) | (mdl, short) <- mods ]
   where
     modPkg mod_ | showPkgs = Just (unitString (moduleUnit mod_))
                 | otherwise = Nothing
     modSrcPkg mod_ | showPkgs = fmap unitPackageIdString
-                                     (lookupUnit dflags (moduleUnit mod_))
+                                     (lookupUnit state (moduleUnit mod_))
                    | otherwise = Nothing
     fn (m,mod_,pkg,srcPkg,short) = addToTrees mod_ m pkg srcPkg short