FFI: Add explicit reloading and refactor eval code

src/Cryptol/Backend/FFI.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE BlockArguments            #-}
 {-# LANGUAGE CPP                       #-}
 {-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE LambdaCase                #-}
 {-# LANGUAGE RecordWildCards           #-}
 {-# LANGUAGE ScopedTypeVariables       #-}
 {-# LANGUAGE TypeApplications          #-}
@@ -11,6 +12,7 @@ module Cryptol.Backend.FFI
 #ifdef FFI_ENABLED
   ( ForeignSrc
   , loadForeignSrc
+  , unloadForeignSrc
   , ForeignImpl
   , loadForeignImpl
   , FFIArg
@@ -23,6 +25,8 @@ module Cryptol.Backend.FFI
 
 #ifdef FFI_ENABLED
 
+import           Control.Concurrent.MVar
+import           Control.DeepSeq
 import           Control.Exception
 import           Control.Monad
 import           Data.Bifunctor
@@ -36,22 +40,35 @@ import           System.IO.Error
 import           System.Posix.DynamicLinker
 
 import           Cryptol.Backend.FFI.Error
+import           Cryptol.Utils.Panic
 
 -- | A source from which we can retrieve implementations of foreign functions.
---
--- This is implemented as a 'ForeignPtr' wrapper around the pointer returned by
--- 'dlopen', where the destructor calls 'dlclose' when the library is no longer
--- needed. We keep references to the 'ForeignPtr' in each foreign function that
--- is in the evaluation environment, so when the Cryptol module is unloaded, the
--- shared library will be closed too.
-newtype ForeignSrc = ForeignSrc (ForeignPtr ())
+data ForeignSrc = ForeignSrc
+  { -- | The 'ForeignPtr' wraps the pointer returned by 'dlopen', where the
+    -- finalizer calls 'dlclose' when the library is no longer needed. We keep
+    -- references to the 'ForeignPtr' in each foreign function that is in the
+    -- evaluation environment, so that the shared library will stay open as long
+    -- as there are references to it.
+    foreignSrcFPtr   :: ForeignPtr ()
+    -- | We support explicit unloading of the shared library so we keep track of
+    -- if it has already been unloaded, and if so the finalizer does nothing.
+    -- This is updated atomically when the library is unloaded.
+  , foreignSrcLoaded :: MVar Bool }
+
+instance Show ForeignSrc where
+  show = show . foreignSrcFPtr
+
+instance NFData ForeignSrc where
+  rnf ForeignSrc {..} = foreignSrcFPtr `seq` foreignSrcLoaded `deepseq` ()
 
 -- | Load a 'ForeignSrc' for the given __Cryptol__ file path. The file path of
 -- the shared library that we try to load is the same as the Cryptol file path
 -- except with a platform specific extension.
 loadForeignSrc :: FilePath -> IO (Either FFILoadError ForeignSrc)
-loadForeignSrc = loadForeignLib >=> traverse \ptr ->
-  ForeignSrc <$> newForeignPtr ptr (unloadForeignLib ptr)
+loadForeignSrc = loadForeignLib >=> traverse \ptr -> do
+  foreignSrcLoaded <- newMVar True
+  foreignSrcFPtr <- newForeignPtr ptr (unloadForeignSrc' foreignSrcLoaded ptr)
+  pure ForeignSrc {..}
 
 loadForeignLib :: FilePath -> IO (Either FFILoadError (Ptr ()))
 #ifdef darwin_HOST_OS
@@ -70,22 +87,42 @@ loadForeignLib path =
         -- module loading time
         open ext = undl <$> dlopen (path -<.> ext) [RTLD_NOW]
 
+-- | Explicitly unload a 'ForeignSrc' immediately instead of waiting for the
+-- garbage collector to do it. This can be useful if you want to immediately
+-- load the same library again to pick up new changes.
+--
+-- The 'ForeignSrc' __must not__ be used in any way after this is called,
+-- including calling 'ForeignImpl's loaded from it.
+unloadForeignSrc :: ForeignSrc -> IO ()
+unloadForeignSrc ForeignSrc {..} = withForeignPtr foreignSrcFPtr $
+  unloadForeignSrc' foreignSrcLoaded
+
+unloadForeignSrc' :: MVar Bool -> Ptr () -> IO ()
+unloadForeignSrc' loaded lib = modifyMVar_ loaded \l -> do
+  when l $ unloadForeignLib lib
+  pure False
+
 unloadForeignLib :: Ptr () -> IO ()
 unloadForeignLib = dlclose . DLHandle
 
+withForeignSrc :: ForeignSrc -> (Ptr () -> IO a) -> IO a
+withForeignSrc ForeignSrc {..} f = withMVar foreignSrcLoaded \case
+  True  -> withForeignPtr foreignSrcFPtr f
+  False -> panic "[FFI] withForeignSrc" ["Use of foreign library after unload"]
+
 -- | An implementation of a foreign function.
 data ForeignImpl = ForeignImpl
   { foreignImplFun :: FunPtr ()
     -- | We don't need this to call the function but we want to keep the library
-    -- around as long as we still have a function from it so that the destructor
-    -- isn't called too early.
-  , foreignImplLib :: ForeignPtr ()
+    -- around as long as we still have a function from it so that it isn't
+    -- unloaded too early.
+  , foreignImplSrc :: ForeignSrc
   }
 
 -- | Load a 'ForeignImpl' with the given name from the given 'ForeignSrc'.
 loadForeignImpl :: ForeignSrc -> String -> IO (Either FFILoadError ForeignImpl)
-loadForeignImpl (ForeignSrc foreignImplLib) name =
-  withForeignPtr foreignImplLib \lib ->
+loadForeignImpl foreignImplSrc name =
+  withForeignSrc foreignImplSrc \lib ->
     tryLoad (CantLoadFFIImpl name) do
       foreignImplFun <- loadForeignFunPtr lib name
       pure ForeignImpl {..}
@@ -159,7 +196,7 @@ data SomeFFIArg = forall a. FFIArg a => SomeFFIArg a
 -- | Call a 'ForeignImpl' with the given arguments. The type parameter decides
 -- how the return value should be converted into a Haskell value.
 callForeignImpl :: forall a. FFIRet a => ForeignImpl -> [SomeFFIArg] -> IO a
-callForeignImpl ForeignImpl {..} xs = withForeignPtr foreignImplLib \_ ->
+callForeignImpl ForeignImpl {..} xs = withForeignSrc foreignImplSrc \_ ->
   callFFI foreignImplFun (ffiRet @a) $ map toArg xs
   where toArg (SomeFFIArg x) = ffiArg x
 

src/Cryptol/Eval/FFI.hs

@@ -1,11 +1,12 @@
 {-# LANGUAGE BlockArguments      #-}
 {-# LANGUAGE CPP                 #-}
-{-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE LambdaCase          #-}
+{-# LANGUAGE NamedFieldPuns      #-}
 {-# LANGUAGE PatternSynonyms     #-}
 {-# LANGUAGE RankNTypes          #-}
 {-# LANGUAGE RecordWildCards     #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections       #-}
 {-# LANGUAGE TypeApplications    #-}
 {-# LANGUAGE ViewPatterns        #-}
 
@@ -21,11 +22,11 @@ import           Cryptol.TypeCheck.AST
 
 #ifdef FFI_ENABLED
 
-import           Control.Monad.Except
-import           Control.Monad.Writer.Strict
-import           Data.Foldable
+import           Data.Either
 import           Data.IORef
+import           Data.Maybe
 import           Data.Proxy
+import           Data.Traversable
 import           Data.Word
 import           Foreign
 import           Foreign.C.Types
@@ -52,42 +53,34 @@ import           Cryptol.Utils.RecordMap
 -- | Find all the foreign declarations in the module and add them to the
 -- environment. This is a separate pass from the main evaluation functions in
 -- "Cryptol.Eval" since it only works for the Concrete backend.
+--
+-- Note: 'Right' is only returned if we successfully loaded some foreign
+-- functions and the environment was modified. If there were no foreign
+-- declarations at all then @Left []@ is returned, so 'Left' does not
+-- necessarily indicate an error.
 evalForeignDecls :: ModulePath -> Module -> EvalEnv ->
-  Eval (Either [FFILoadError] EvalEnv)
-evalForeignDecls path m env = do
-  -- We only want to try loading the 'ForeignSrc' when we encounter the first
-  -- foreign decl, so we store it in an 'IORef'
-  foreignSrc <- liftIO $ newIORef Nothing
-  let evalForeignDeclGroup e (Recursive ds)   = foldlM evalForeignDecl e ds
-      evalForeignDeclGroup e (NonRecursive d) = evalForeignDecl e d
-      evalForeignDecl e d = case dDefinition d of
-        DForeign ffiType -> do
-          fsrc <- liftIO (readIORef foreignSrc) >>= \case
-            Nothing -> case path of
-              InFile p -> do
-                -- If any error happens when loading the 'ForeignSrc', stop
-                -- processing the module.
-                fsrc <- liftEither =<<
-                  liftIO (canonicalizePath p >>= loadForeignSrc)
-                liftIO $ writeIORef foreignSrc $ Just fsrc
-                pure fsrc
-              -- We don't handle in-memory modules for now
-              InMem _ _ -> evalPanic "evalForeignDecls"
-                ["Can't find foreign source of in-memory module"]
-            Just fsrc -> pure fsrc
-          liftIO (loadForeignImpl fsrc $ unpackIdent $ nameIdent $ dName d)
-            >>= \case
-              -- If there is an error loading the 'ForeignImpl', record it and
-              -- keep going, because we want to check all the functions.
-              Left err -> tell [err] >> pure e
-              Right impl -> pure $ bindVarDirect (dName d)
-                (foreignPrimPoly (dName d) ffiType impl) e
-        _ -> pure e
-      report (Left err)         = Left [err]
-      report (Right (env', [])) = Right env'
-      report (Right (_, errs))  = Left errs
-  fmap report $ runExceptT $ runWriterT $
-    foldlM evalForeignDeclGroup env $ mDecls m
+  Eval (Either [FFILoadError] (ForeignSrc, EvalEnv))
+evalForeignDecls path m env = io
+  case mapMaybe getForeign $ mDecls m of
+    [] -> pure $ Left []
+    foreigns ->
+      case path of
+        InFile p -> canonicalizePath p >>= loadForeignSrc >>=
+          \case
+            Right fsrc -> collect <$> for foreigns \(name, ffiType) ->
+              fmap ((name,) . foreignPrimPoly name ffiType) <$>
+                loadForeignImpl fsrc (unpackIdent $ nameIdent name)
+              where collect (partitionEithers -> (errs, primMap))
+                      | null errs = Right
+                        (fsrc, foldr (uncurry bindVarDirect) env primMap)
+                      | otherwise = Left errs
+            Left err -> pure $ Left [err]
+        -- We don't handle in-memory modules for now
+        InMem _ _ -> evalPanic "evalForeignDecls"
+          ["Can't find foreign source of in-memory module"]
+  where getForeign (NonRecursive Decl { dName, dDefinition = DForeign ffiType })
+          = Just (dName, ffiType)
+        getForeign _ = Nothing
 
 -- | Generate a 'Prim' value representing the given foreign function, containing
 -- all the code necessary to marshal arguments and return values and do the

src/Cryptol/ModuleSystem.hs

@@ -64,17 +64,19 @@ findModule n env = runModuleM env (Base.findModule n)
 
 -- | Load the module contained in the given file.
 loadModuleByPath :: FilePath -> ModuleCmd (ModulePath,T.Module)
-loadModuleByPath path minp =
-  runModuleM minp{ minpModuleEnv = resetModuleEnv (minpModuleEnv minp) } $ do
+loadModuleByPath path minp = do
+  moduleEnv' <- resetModuleEnv $ minpModuleEnv minp
+  runModuleM minp{ minpModuleEnv = moduleEnv' } $ do
     unloadModule ((InFile path ==) . lmFilePath)
     m <- Base.loadModuleByPath path
     setFocusedModule (T.mName m)
     return (InFile path,m)
 
 -- | Load the given parsed module.
 loadModuleByName :: P.ModName -> ModuleCmd (ModulePath,T.Module)
-loadModuleByName n minp =
-  runModuleM minp{ minpModuleEnv = resetModuleEnv (minpModuleEnv minp) } $ do
+loadModuleByName n minp = do
+  moduleEnv' <- resetModuleEnv $ minpModuleEnv minp
+  runModuleM minp{ minpModuleEnv = moduleEnv' } $ do
     unloadModule ((n ==) . lmName)
     (path,m') <- Base.loadModuleFrom False (FromModule n)
     setFocusedModule (T.mName m')

src/Cryptol/ModuleSystem/Base.hs

@@ -19,6 +19,7 @@ module Cryptol.ModuleSystem.Base where
 
 import qualified Control.Exception as X
 import Control.Monad (unless,when)
+import Data.Functor.Compose
 import Data.Maybe (fromMaybe)
 import Data.Monoid ((<>))
 import Data.Text.Encoding (decodeUtf8')
@@ -239,11 +240,17 @@ doLoadModule quiet isrc path fp pm0 =
      let ?evalPrim = \i -> Right <$> Map.lookup i tbl
      callStacks <- getCallStacks
      let ?callStacks = callStacks
-     unless (T.isParametrizedModule tcm) $
-       modifyEvalEnvM (evalForeignDecls path tcm) >>= \case
-         Left errs -> ffiLoadErrors (T.mName tcm) errs
-         Right ()  -> modifyEvalEnv (E.moduleEnv Concrete tcm)
-     loadedModule path fp nameEnv tcm
+     foreignSrc <-
+       if T.isParametrizedModule tcm
+         then pure Nothing
+         else (getCompose
+                <$> modifyEvalEnvM (fmap Compose . evalForeignDecls path tcm)
+                >>= \case
+                  Left []                -> pure Nothing
+                  Left errs              -> ffiLoadErrors (T.mName tcm) errs
+                  Right (foreignSrc, ()) -> pure (Just foreignSrc))
+              <* modifyEvalEnv (E.moduleEnv Concrete tcm)
+     loadedModule path fp nameEnv foreignSrc tcm
 
      return tcm
   where

src/Cryptol/ModuleSystem/Env.hs

@@ -18,6 +18,7 @@ module Cryptol.ModuleSystem.Env where
 import Paths_cryptol (getDataDir)
 #endif
 
+import Cryptol.Backend.FFI (ForeignSrc, unloadForeignSrc)
 import Cryptol.Eval (EvalEnv)
 import Cryptol.ModuleSystem.Fingerprint
 import Cryptol.ModuleSystem.Interface
@@ -41,6 +42,7 @@ import System.Directory (getAppUserDataDirectory, getCurrentDirectory)
 import System.Environment(getExecutablePath)
 import System.FilePath ((</>), normalise, joinPath, splitPath, takeDirectory)
 import qualified Data.List as List
+import Data.Foldable
 
 import GHC.Generics (Generic)
 import Control.DeepSeq
@@ -99,14 +101,19 @@ data CoreLint = NoCoreLint        -- ^ Don't run core lint
               | CoreLint          -- ^ Run core lint
   deriving (Generic, NFData)
 
-resetModuleEnv :: ModuleEnv -> ModuleEnv
-resetModuleEnv env = env
-  { meLoadedModules = mempty
-  , meNameSeeds     = T.nameSeeds
-  , meEvalEnv       = mempty
-  , meFocusedModule = Nothing
-  , meDynEnv        = mempty
-  }
+resetModuleEnv :: ModuleEnv -> IO ModuleEnv
+resetModuleEnv env = do
+  for_ (getLoadedModules $ meLoadedModules env) $ \lm ->
+    case lmForeignSrc lm of
+      Just fsrc -> unloadForeignSrc fsrc
+      _         -> pure ()
+  pure env
+    { meLoadedModules = mempty
+    , meNameSeeds     = T.nameSeeds
+    , meEvalEnv       = mempty
+    , meFocusedModule = Nothing
+    , meDynEnv        = mempty
+    }
 
 initialModuleEnv :: IO ModuleEnv
 initialModuleEnv = do
@@ -342,6 +349,9 @@ data LoadedModule = LoadedModule
     -- ^ The actual type-checked module
 
   , lmFingerprint       :: Fingerprint
+
+  , lmForeignSrc        :: Maybe ForeignSrc
+    -- ^ The dynamically loaded source for any foreign functions in the module
   } deriving (Show, Generic, NFData)
 
 -- | Has this module been loaded already.
@@ -362,9 +372,9 @@ lookupModule mn me = search lmLoadedModules `mplus` search lmLoadedParamModules
 -- | Add a freshly loaded module.  If it was previously loaded, then
 -- the new version is ignored.
 addLoadedModule ::
-  ModulePath -> String -> Fingerprint -> R.NamingEnv -> T.Module ->
-  LoadedModules -> LoadedModules
-addLoadedModule path ident fp nameEnv tm lm
+  ModulePath -> String -> Fingerprint -> R.NamingEnv -> Maybe ForeignSrc ->
+  T.Module -> LoadedModules -> LoadedModules
+addLoadedModule path ident fp nameEnv fsrc tm lm
   | isLoaded (T.mName tm) lm  = lm
   | T.isParametrizedModule tm = lm { lmLoadedParamModules = loaded :
                                                 lmLoadedParamModules lm }
@@ -379,6 +389,7 @@ addLoadedModule path ident fp nameEnv tm lm
     , lmInterface       = T.genIface tm
     , lmModule          = tm
     , lmFingerprint     = fp
+    , lmForeignSrc      = fsrc
     }
 
 -- | Remove a previously loaded module.

src/Cryptol/ModuleSystem/Monad.hs

@@ -14,6 +14,7 @@ module Cryptol.ModuleSystem.Monad where
 
 import           Cryptol.Eval (EvalEnv,EvalOpts(..))
 
+import           Cryptol.Backend.FFI (ForeignSrc)
 import           Cryptol.Backend.FFI.Error
 import qualified Cryptol.Backend.Monad           as E
 
@@ -525,14 +526,15 @@ unloadModule rm = ModuleT $ do
   set $! env { meLoadedModules = removeLoadedModule rm (meLoadedModules env) }
 
 loadedModule ::
-  ModulePath -> Fingerprint -> NamingEnv -> T.Module -> ModuleM ()
-loadedModule path fp nameEnv m = ModuleT $ do
+  ModulePath -> Fingerprint -> NamingEnv -> Maybe ForeignSrc -> T.Module ->
+  ModuleM ()
+loadedModule path fp nameEnv fsrc m = ModuleT $ do
   env <- get
   ident <- case path of
              InFile p  -> unModuleT $ io (canonicalizePath p)
              InMem l _ -> pure l
 
-  set $! env { meLoadedModules = addLoadedModule path ident fp nameEnv m
+  set $! env { meLoadedModules = addLoadedModule path ident fp nameEnv fsrc m
                                                         (meLoadedModules env) }
 
 modifyEvalEnvM :: Traversable t =>