Merge pull request purescript#1160 from puffnfresh/perf/dictionary-lookup

paf31 · paf31 · commit fedeee22bfa5 · 2015-06-10T20:38:40.000-07:00
Improve performance of type-class dict lookup
diff --git a/src/Language/PureScript/CodeGen/Externs.hs b/src/Language/PureScript/CodeGen/Externs.hs
@@ -107,7 +107,7 @@ moduleToPs (Module _ moduleName ds (Just exts)) env = intercalate "\n" . execWri
 
     exportToPs (TypeInstanceRef ident) = do
       let TypeClassDictionaryInScope { tcdClassName = className, tcdInstanceTypes = tys, tcdDependencies = deps} =
-            fromMaybe (error $ "Type class instance has no dictionary in exportToPs") . find (\tcd -> tcdName tcd == Qualified (Just moduleName) ident && tcdType tcd == TCDRegular) $ M.elems $ typeClassDictionaries env
+            fromMaybe (error $ "Type class instance has no dictionary in exportToPs") . find (\tcd -> tcdName tcd == Qualified (Just moduleName) ident && tcdType tcd == TCDRegular) . maybe [] M.elems . M.lookup (Just moduleName) $ typeClassDictionaries env
       let constraintsText = case fromMaybe [] deps of
                               [] -> ""
                               cs -> "(" ++ intercalate ", " (map (\(pn, tys') -> show pn ++ " " ++ unwords (map prettyPrintTypeAtom tys')) cs) ++ ") => "
diff --git a/src/Language/PureScript/Environment.hs b/src/Language/PureScript/Environment.hs
@@ -49,7 +49,7 @@ data Environment = Environment {
   -- |
   -- Available type class dictionaries
   --
-  , typeClassDictionaries :: M.Map (Qualified Ident, Maybe ModuleName) TypeClassDictionaryInScope
+  , typeClassDictionaries :: M.Map (Maybe ModuleName) (M.Map (Qualified Ident) TypeClassDictionaryInScope)
   -- |
   -- Type classes
   --
diff --git a/src/Language/PureScript/TypeChecker.hs b/src/Language/PureScript/TypeChecker.hs
@@ -90,8 +90,8 @@ addTypeClass moduleName pn args implies ds =
 
 addTypeClassDictionaries :: [TypeClassDictionaryInScope] -> Check ()
 addTypeClassDictionaries entries =
-  let mentries = M.fromList [ ((canonicalizeDictionary entry, mn), entry) | entry@TypeClassDictionaryInScope{ tcdName = Qualified mn _ } <- entries ]
-  in modify $ \st -> st { checkEnv = (checkEnv st) { typeClassDictionaries = (typeClassDictionaries . checkEnv $ st) `M.union` mentries } }
+  let mentries = M.fromListWith (M.union) [ (mn, M.singleton (canonicalizeDictionary entry) entry) | entry@TypeClassDictionaryInScope{ tcdName = Qualified mn _ } <- entries ]
+  in modify $ \st -> st { checkEnv = (checkEnv st) { typeClassDictionaries = M.unionWith (M.union) (typeClassDictionaries . checkEnv $ st) mentries } }
 
 checkDuplicateTypeArguments :: [String] -> Check ()
 checkDuplicateTypeArguments args = for_ firstDup $ \dup ->
@@ -223,8 +223,7 @@ typeCheckAll mainModuleName moduleName exps = go
     guardWith (errorMessage (OrphanFixityDeclaration name)) $ M.member (moduleName, Op name) $ names env
     return $ d : ds
   go (d@(ImportDeclaration importedModule _ _) : rest) = do
-    tcds <- getTypeClassDictionaries
-    let instances = filter (\tcd -> let Qualified (Just mn) _ = tcdName tcd in importedModule == mn) tcds
+    instances <- lookupTypeClassDictionaries $ Just importedModule
     addTypeClassDictionaries [ tcd { tcdName = Qualified (Just moduleName) ident, tcdType = TCDAlias (canonicalizeDictionary tcd) }
                              | tcd <- instances
                              , tcdExported tcd
diff --git a/src/Language/PureScript/TypeChecker/Monad.hs b/src/Language/PureScript/TypeChecker/Monad.hs
@@ -69,8 +69,8 @@ withScopedTypeVars mn ks = bindTypes (M.fromList (map (\(name, k) -> (Qualified
 withTypeClassDictionaries :: (MonadState CheckState m) => [TypeClassDictionaryInScope] -> m a -> m a
 withTypeClassDictionaries entries action = do
   orig <- get
-  let mentries = M.fromList [ ((canonicalizeDictionary entry, mn), entry) | entry@TypeClassDictionaryInScope{ tcdName = Qualified mn _ }  <- entries ]
-  modify $ \st -> st { checkEnv = (checkEnv st) { typeClassDictionaries = (typeClassDictionaries . checkEnv $ st) `M.union` mentries } }
+  let mentries = M.fromListWith (M.union) [ (mn, M.singleton (canonicalizeDictionary entry) entry) | entry@TypeClassDictionaryInScope{ tcdName = Qualified mn _ }  <- entries ]
+  modify $ \st -> st { checkEnv = (checkEnv st) { typeClassDictionaries = M.unionWith (M.union) (typeClassDictionaries . checkEnv $ st) mentries } }
   a <- action
   modify $ \st -> st { checkEnv = (checkEnv st) { typeClassDictionaries = typeClassDictionaries . checkEnv $ orig } }
   return a
@@ -79,7 +79,13 @@ withTypeClassDictionaries entries action = do
 -- Get the currently available list of type class dictionaries
 --
 getTypeClassDictionaries :: (Functor m, MonadState CheckState m) => m [TypeClassDictionaryInScope]
-getTypeClassDictionaries = M.elems . typeClassDictionaries . checkEnv <$> get
+getTypeClassDictionaries = (>>= M.elems) . M.elems . typeClassDictionaries . checkEnv <$> get
+
+-- |
+-- Lookup type class dictionaries in a module.
+--
+lookupTypeClassDictionaries :: (Functor m, MonadState CheckState m) => Maybe ModuleName -> m [TypeClassDictionaryInScope]
+lookupTypeClassDictionaries mn = maybe [] M.elems . M.lookup mn . typeClassDictionaries . checkEnv <$> get
 
 -- |
 -- Temporarily bind a collection of names to local variables
