diff --git a/Language/Haskell/GhcMod/Gap.hs b/Language/Haskell/GhcMod/Gap.hs
index 7d8706f..6f699f5 100644
--- a/Language/Haskell/GhcMod/Gap.hs
+++ b/Language/Haskell/GhcMod/Gap.hs
@@ -200,23 +200,24 @@ withContext action = gbracket setup teardown body
     setup = getContext
     teardown = setContext
     body _ = do
-        topImports >>= setContext
+        topImports >>= setCtx
         action
     topImports = do
         mss <- getModuleGraph
-#if __GLASGOW_HASKELL__ >= 706
-        let modName = IIModule . moduleName . ms_mod
-#elif __GLASGOW_HASKELL__ >= 704
-        let modName = IIModule . ms_mod
-#else
-        let modName = ms_mod
-#endif
         map modName <$> filterM isTop mss
-      where
-        isTop mos = lookupMod ||> returnFalse
-          where
-            lookupMod = lookupModule (ms_mod_name mos) Nothing >> return True
-            returnFalse = return False
+    isTop mos = lookupMod mos ||> returnFalse
+    lookupMod mos = lookupModule (ms_mod_name mos) Nothing >> return True
+    returnFalse = return False
+#if __GLASGOW_HASKELL__ >= 706
+    modName = IIModule . moduleName . ms_mod
+    setCtx = setContext
+#elif __GLASGOW_HASKELL__ >= 704
+    modName = IIModule . ms_mod
+    setCtx = setContext
+#else
+    modName = ms_mod
+    setCtx = flip setContext []
+#endif
 
 showSeverityCaption :: Severity -> String
 #if __GLASGOW_HASKELL__ >= 706