perf: memoize env lookups across top-level forms

.gitignore

@@ -5,6 +5,7 @@
 *.output.actual
 
 .stack-work/
+.stack-work-prof/
 
 /CarpHask-exe.prof
 

CarpHask.cabal

@@ -45,6 +45,7 @@ library
                        Interfaces,
                        Managed,
                        Map,
+                       Memo,
                        Memory,
                        Meta,
                        Obj,

src/Concretize.hs

@@ -36,6 +36,7 @@ import Data.List (foldl')
 import Data.Maybe (fromMaybe)
 import Debug.Trace
 import Env (EnvironmentError, empty, envIsExternal, getBinder, insert, insertX, search, searchTypeDef)
+import qualified Env
 import Forms
 import Info
 import InitialTypes
@@ -892,11 +893,17 @@ depsOfPolymorphicFunction typeEnv env visitedDefinitions functionName functionTy
             instancePath = SymPath ps (nm ++ polymorphicSuffix polyType functionType)
          in -- Already being concretized up the stack: emitted there. Re-listing
             -- it here would contribute a bodyless stub overwriting the real def.
-            if instancePath `elem` visitedDefinitions
+            -- Already defined by an earlier form: it and its own dependencies
+            -- are in the environment, so concretizing it again is wasted work.
+            if instancePath `elem` visitedDefinitions || isDefinedInstance instancePath
               then []
               else case concretizeDefinition False typeEnv env visitedDefinitions single functionType of
                 Left err -> error (show err)
                 Right (ok, deps) -> ok : deps
+      isDefinedInstance path = case Env.findBinderMaybe env path of
+        Just (Binder _ (XObj (Lst (XObj MetaStub _ _ : _)) _ _)) -> False
+        Just _ -> True
+        Nothing -> False
    in case allImplementations typeEnv env functionName functionType of
         [] ->
           (trace $ "[Warning] No '" ++ functionName ++ "' function found with type " ++ show functionType ++ ".")

src/Env.hs

@@ -44,6 +44,10 @@ module Env
     -------------------------
     -- finds
     findPoly,
+    findPolyIn,
+    polyEnvs,
+    cachedPolyEnvs,
+    cachedLookupEverywhere,
     findAllByMeta,
     findChildren,
     findImportedEnvs,
@@ -67,11 +71,14 @@ where
 
 import Data.Either (fromRight, rights)
 import Data.List (foldl', unfoldr)
-import Data.Maybe (fromMaybe, mapMaybe)
+import Data.Maybe (fromMaybe, listToMaybe, mapMaybe)
 import qualified Map
+import Memo
 import qualified Meta
 import Obj
 import qualified Set
+import System.IO.Unsafe (unsafePerformIO)
+import System.Mem.StableName (StableName)
 import Types
 
 --------------------------------------------------------------------------------
@@ -472,7 +479,42 @@ findPoly env name ty =
     Left _ -> foldl' go (Left (BindingNotFound name (prj env))) (findChildren env)
   where
     go x e = x <> (findPoly e name ty)
-    unify = areUnifiable ty . fromMaybe Universe . xobjTy . binderXObj
+    unify = polyMatches ty
+
+-- | Does the type of a binder unify with the type a polymorphic lookup asks for?
+polyMatches :: Ty -> Binder -> Bool
+polyMatches ty = areUnifiable ty . fromMaybe Universe . xobjTy . binderXObj
+
+-- | All environments that 'findPoly' visits, in the order it visits them.
+--
+-- Finding the children of an environment means scanning all of its binders,
+-- so for repeated lookups it is much cheaper to build this list once and
+-- search it with 'findPolyIn'.
+polyEnvs :: Environment e => e -> [e]
+polyEnvs e = e : concatMap polyEnvs (findChildren e)
+
+-- | 'polyEnvs' for a value environment, remembered for recent environments
+-- (see "Memo"), since many lookups in a row usually share an environment.
+cachedPolyEnvs :: Env -> [Env]
+cachedPolyEnvs = memoOnIdentity polyEnvsCache polyEnvs
+
+polyEnvsCache :: MemoCache (StableName Env) [Env]
+polyEnvsCache = unsafePerformIO newMemoCache
+{-# NOINLINE polyEnvsCache #-}
+
+-- | 'lookupEverywhere' for a value environment, with the expensive part, finding
+-- all module environments, remembered for recent environments (see "Memo").
+cachedLookupEverywhere :: Env -> String -> [(Env, Binder)]
+cachedLookupEverywhere = lookupExhuastive (memoOnIdentity lookupChildrenCache lookupChildren)
+
+lookupChildrenCache :: MemoCache (StableName Env) [Env]
+lookupChildrenCache = unsafePerformIO newMemoCache
+{-# NOINLINE lookupChildrenCache #-}
+
+-- | Like 'findPoly', but searches a list of environments built by 'polyEnvs'.
+findPolyIn :: Environment e => [e] -> String -> Ty -> Maybe (e, Binder)
+findPolyIn envs name ty =
+  listToMaybe [(e, b) | e <- envs, Right b <- [getBinder e name], polyMatches ty b]
 
 -- | Find all environments that are *direct* children of an environment (one
 -- level down).

src/Interfaces.hs

@@ -7,6 +7,7 @@ module Interfaces
     registerInInterface,
     retroactivelyRegisterInInterface,
     interfaceImplementedForTy,
+    interfaceImplementationTypes,
     removeInterfaceFromImplements,
     InterfaceError (..),
   )
@@ -159,9 +160,14 @@ retroactivelyRegisterInInterface ctx interface =
 -- | e.g. Is "delete" implemented for `(Fn [String] ())` ?
 interfaceImplementedForTy :: TypeEnv -> Env -> String -> Ty -> Bool
 interfaceImplementedForTy typeEnv globalEnv interfaceName matchingTy =
+  any (areUnifiable matchingTy) (interfaceImplementationTypes typeEnv globalEnv interfaceName)
+
+-- | The types of all registered implementations of an interface, or none if
+-- there is no such interface.
+interfaceImplementationTypes :: TypeEnv -> Env -> String -> [Ty]
+interfaceImplementationTypes typeEnv globalEnv interfaceName =
   case Env.getBinder typeEnv interfaceName of
     Right (Binder _ (XObj (Lst (XObj (Interface _ paths) _ _ : _)) _ _)) ->
       let lookupType' path = forceTy . binderXObj <$> (Env.searchBinder globalEnv path)
-          matches = filter (areUnifiable matchingTy) (rights (map lookupType' paths))
-       in not . null $ matches
-    _ -> False
+       in rights (map lookupType' paths)
+    _ -> []

src/Managed.hs

@@ -1,7 +1,10 @@
 module Managed where
 
 import Interfaces
+import Memo
 import Obj
+import System.IO.Unsafe (unsafePerformIO)
+import System.Mem.StableName (StableName)
 import Types
 
 -- | Should this type be handled by the memory management system.
@@ -11,21 +14,48 @@ import Types
 --
 -- TODO: When blit and delete are both implemented, issue a warning.
 isManaged :: TypeEnv -> Env -> Ty -> Bool
-isManaged typeEnv globalEnv structTy@StructTy {} =
-  not (isBlittable typeEnv globalEnv structTy)
-    && interfaceImplementedForTy typeEnv globalEnv "delete" (FuncTy [structTy] UnitTy StaticLifetimeTy)
-isManaged typeEnv globalEnv funcTy@FuncTy {} =
-  not (isBlittable typeEnv globalEnv funcTy)
-    && interfaceImplementedForTy typeEnv globalEnv "delete" (FuncTy [funcTy] UnitTy StaticLifetimeTy)
-isManaged _ _ StringTy =
-  True
-isManaged _ _ PatternTy =
-  True
-isManaged _ _ _ =
-  False
+isManaged typeEnv globalEnv = isManagedIn (findManagedTypes typeEnv globalEnv)
 
 -- | Returns true if this type implements the "blit" interface and is thus
 -- freely copyable.
 isBlittable :: TypeEnv -> Env -> Ty -> Bool
-isBlittable typeEnv globalEnv t =
-  interfaceImplementedForTy typeEnv globalEnv "blit" (FuncTy [t] t StaticLifetimeTy)
+isBlittable typeEnv globalEnv = isBlittableIn (findManagedTypes typeEnv globalEnv)
+
+-- | The types of the 'blit' and 'delete' implementations, looked up once so
+-- that many checks against the same environments can share the work.
+data ManagedTypes = ManagedTypes
+  { blitImplementationTypes :: [Ty],
+    deleteImplementationTypes :: [Ty]
+  }
+
+-- | Remembered for the most recent environments (see "Memo"), since the many
+-- checks during one top-level form share them.
+findManagedTypes :: TypeEnv -> Env -> ManagedTypes
+findManagedTypes = memoOnIdentity2 managedTypesCache $ \typeEnv globalEnv ->
+  ManagedTypes
+    (interfaceImplementationTypes typeEnv globalEnv "blit")
+    (interfaceImplementationTypes typeEnv globalEnv "delete")
+
+managedTypesCache :: MemoCache (StableName TypeEnv, StableName Env) ManagedTypes
+managedTypesCache = unsafePerformIO newMemoCache
+{-# NOINLINE managedTypesCache #-}
+
+-- | Like 'isManaged', but checks against precomputed 'ManagedTypes'.
+isManagedIn :: ManagedTypes -> Ty -> Bool
+isManagedIn managed structTy@StructTy {} =
+  not (isBlittableIn managed structTy)
+    && any (areUnifiable (FuncTy [structTy] UnitTy StaticLifetimeTy)) (deleteImplementationTypes managed)
+isManagedIn managed funcTy@FuncTy {} =
+  not (isBlittableIn managed funcTy)
+    && any (areUnifiable (FuncTy [funcTy] UnitTy StaticLifetimeTy)) (deleteImplementationTypes managed)
+isManagedIn _ StringTy =
+  True
+isManagedIn _ PatternTy =
+  True
+isManagedIn _ _ =
+  False
+
+-- | Like 'isBlittable', but checks against precomputed 'ManagedTypes'.
+isBlittableIn :: ManagedTypes -> Ty -> Bool
+isBlittableIn managed t =
+  any (areUnifiable (FuncTy [t] t StaticLifetimeTy)) (blitImplementationTypes managed)

src/Memo.hs

@@ -0,0 +1,63 @@
+-- | Memoization on the identity of immutable values.
+--
+-- The compiler repeatedly derives the same data from the same environment:
+-- an environment only changes between top-level forms, but lookups against
+-- it happen many thousands of times per form. Threading caches through all
+-- of those call sites is impractical, so these helpers remember the results
+-- for the most recently seen arguments instead.
+--
+-- This is safe because the arguments are immutable: the same heap object
+-- always gives the same result. A different object, even an equal one, is
+-- simply a cache miss and the result is computed as usual.
+module Memo
+  ( MemoCache,
+    newMemoCache,
+    memoOnIdentity,
+    memoOnIdentity2,
+  )
+where
+
+import Control.Exception (evaluate)
+import Data.IORef
+import System.IO.Unsafe (unsafePerformIO)
+import System.Mem.StableName
+
+-- | Remembers the results for the most recently seen keys, newest first.
+newtype MemoCache k v = MemoCache (IORef [(k, v)])
+
+-- | How many results a cache remembers. Lookups often alternate between a
+-- few environments (e.g. a local one and the global one), so one isn't enough.
+cacheSize :: Int
+cacheSize = 8
+
+-- | Create a cache. Bind each one to a top-level NOINLINE value.
+newMemoCache :: IO (MemoCache k v)
+newMemoCache = MemoCache <$> newIORef []
+
+lookupOrCompute :: Eq k => MemoCache k v -> k -> v -> IO v
+lookupOrCompute (MemoCache ref) key result = do
+  entries <- readIORef ref
+  case lookup key entries of
+    Just cached -> pure cached
+    Nothing -> do
+      writeIORef ref (take cacheSize ((key, result) : entries))
+      pure result
+
+-- | @memoOnIdentity cache f x@ is @f x@, reusing an earlier result when @x@
+-- is the same object as a recently seen argument.
+memoOnIdentity :: MemoCache (StableName a) v -> (a -> v) -> a -> v
+memoOnIdentity cache f x = unsafePerformIO $ do
+  x' <- evaluate x -- The stable names of a thunk and of its value differ.
+  name <- makeStableName x'
+  lookupOrCompute cache name (f x')
+{-# NOINLINE memoOnIdentity #-}
+
+-- | Like 'memoOnIdentity', for a function of two arguments.
+memoOnIdentity2 :: MemoCache (StableName a, StableName b) v -> (a -> b -> v) -> a -> b -> v
+memoOnIdentity2 cache f x y = unsafePerformIO $ do
+  x' <- evaluate x
+  y' <- evaluate y
+  nameX <- makeStableName x'
+  nameY <- makeStableName y'
+  lookupOrCompute cache (nameX, nameY) (f x' y')
+{-# NOINLINE memoOnIdentity2 #-}