Small changes suggested in review.

lib/prelude.dx

@@ -1592,6 +1592,3 @@ def throw (_:Unit) : {Except} a =
 
 def assert (b:Bool) : {Except} Unit =
   if not b then throw ()
-
--- TODO: remove this when we have modules that exercise the data-caching path
-justSomeDataToTestCaching = [toList [1,2,3], toList [], toList [1,2]]

src/Dex/Foreign/Context.hs

@@ -33,7 +33,7 @@ import PPrint
 
 import Dex.Foreign.Util
 
-data Context = Context EvalConfig Bindings
+data Context = Context EvalConfig TopEnv
 
 foreign import ccall "_internal_dexSetError" internalSetErrorPtr :: CString -> Int64 -> IO ()
 setError :: String -> IO ()
@@ -45,7 +45,7 @@ dexCreateContext = do
   let evalConfig = EvalConfig LLVM Nothing Nothing
   maybePreludeEnv <- evalPrelude evalConfig preludeSource
   case maybePreludeEnv of
-    Right preludeEnv -> toStablePtr $ Context evalConfig (topBindings preludeEnv)
+    Right preludeEnv -> toStablePtr $ Context evalConfig preludeEnv
     Left  err        -> nullPtr <$ setError ("Failed to initialize standard library: " ++ pprint err)
   where
     evalPrelude :: EvalConfig -> String -> IO (Either Err TopEnv)
@@ -66,19 +66,19 @@ dexEval :: Ptr Context -> CString -> IO (Ptr Context)
 dexEval ctxPtr sourcePtr = do
   Context evalConfig env <- fromStablePtr ctxPtr
   source <- peekCString sourcePtr
-  (results, finalEnv) <- runStateT (evalSource evalConfig source) (bindingsToTopEnv env)
+  (results, finalEnv) <- runStateT (evalSource evalConfig source) env
   let anyError = asum $ fmap (\case (_, Result _ (Left err)) -> Just err; _ -> Nothing) results
   case anyError of
-    Nothing  -> toStablePtr $ Context evalConfig (topBindings finalEnv)
+    Nothing  -> toStablePtr $ Context evalConfig finalEnv
     Just err -> setError (pprint err) $> nullPtr
 
 dexInsert :: Ptr Context -> CString -> Ptr Atom -> IO (Ptr Context)
 dexInsert ctxPtr namePtr atomPtr = do
   Context evalConfig env <- fromStablePtr ctxPtr
   name <- GlobalName . fromString <$> peekCString namePtr
   atom <- fromStablePtr atomPtr
-  let env' = env <> name @> (getType atom, LetBound PlainLet (Atom atom))
-  toStablePtr $ Context evalConfig env'
+  let newBinding = name @> (getType atom, LetBound PlainLet (Atom atom))
+  toStablePtr $ Context evalConfig $ env <> TopEnv newBinding mempty
 
 dexEvalExpr :: Ptr Context -> CString -> IO (Ptr Atom)
 dexEvalExpr ctxPtr sourcePtr = do
@@ -89,7 +89,7 @@ dexEvalExpr ctxPtr sourcePtr = do
       let (v, m) = exprAsModule expr
       let block = SourceBlock 0 0 LogNothing source (RunModule m) Nothing
       (resultEnv, Result [] maybeErr) <-
-          evalSourceBlock evalConfig (bindingsToTopEnv env) block
+          evalSourceBlock evalConfig env block
       case maybeErr of
         Right () -> do
           let (_, LetBound _ (Atom atom)) = topBindings resultEnv ! v
@@ -101,7 +101,7 @@ dexLookup :: Ptr Context -> CString -> IO (Ptr Atom)
 dexLookup ctxPtr namePtr = do
   Context _ env <- fromStablePtr ctxPtr
   name <- peekCString namePtr
-  case envLookup env (GlobalName $ fromString name) of
+  case envLookup (topBindings env) (GlobalName $ fromString name) of
     Just (_, LetBound _ (Atom atom)) -> toStablePtr atom
     Just _                           -> setError "Looking up an expression" $> nullPtr
     Nothing                          -> setError "Unbound name" $> nullPtr

src/Dex/Foreign/JIT.hs

@@ -37,6 +37,7 @@ import LLVMExec
 import JIT
 import Syntax  hiding (sizeOf)
 import Export
+import TopLevel
 
 import Dex.Foreign.Util
 import Dex.Foreign.Context
@@ -82,7 +83,8 @@ dexCompile jitPtr ctxPtr funcAtomPtr = do
   ForeignJIT{..} <- fromStablePtr jitPtr
   Context _ env <- fromStablePtr ctxPtr
   funcAtom <- fromStablePtr funcAtomPtr
-  let (impMod, nativeSignature) = prepareFunctionForExport env "userFunc" funcAtom
+  let (impMod, nativeSignature) = prepareFunctionForExport
+                                    (topBindings env) "userFunc" funcAtom
   nativeModule <- execLogger Nothing $ \logger -> do
     llvmAST <- impToLLVM logger impMod
     LLVM.JIT.compileModule jit llvmAST

src/lib/Serialize.hs

@@ -14,6 +14,8 @@ module Serialize (pprintVal, cached, getDexString, cachedWithSnapshot,
 import Prelude hiding (pi, abs)
 import Control.Monad
 import qualified Data.ByteString as BS
+import Data.ByteString.Internal (memcpy)
+import Data.ByteString.Unsafe (unsafeUseAsCString)
 import System.Directory
 import System.FilePath
 import Control.Monad.Writer
@@ -131,49 +133,55 @@ class HasPtrs a where
 
 takeSnapshot :: HasPtrs a => a -> IO (WithSnapshot a)
 takeSnapshot x =
+  -- TODO: we're using `Writer []` (as we do elsewhere) which has bad
+  -- asymptotics. We should switch all of these uses to use snoc lists instead.
   liftM (WithSnapshot x) $ execWriterT $ flip traversePtrs x \ptrTy ptrVal -> do
     snapshot <- lift $ takePtrSnapshot ptrTy ptrVal
     tell [snapshot]
     return ptrVal
 
-restoreSnapshot :: HasPtrs a => WithSnapshot a -> IO a
-restoreSnapshot (WithSnapshot x snapshots) =
-  flip evalStateT snapshots $ flip traversePtrs x \_ _ -> do
-    (s:ss) <- get
-    put ss
-    lift $ restorePtrSnapshot s
-
 takePtrSnapshot :: PtrType -> RawPtr -> IO PtrSnapshot
 takePtrSnapshot _ ptrVal | ptrVal == nullPtr = return NullPtr
-takePtrSnapshot (_, ptrTy) ptrVal = case ptrTy of
+takePtrSnapshot (Heap CPU, ptrTy) ptrVal = case ptrTy of
   PtrType eltTy -> do
     childPtrs <- loadPtrPtrs ptrVal
     PtrArray <$> mapM (takePtrSnapshot eltTy) childPtrs
   _ -> ByteArray <$> loadPtrBytes ptrVal
-
-restorePtrSnapshot :: PtrSnapshot -> IO RawPtr
-restorePtrSnapshot snapshot = case snapshot of
-  PtrArray  children -> storePtrPtrs =<< mapM restorePtrSnapshot children
-  ByteArray bytes    -> storePtrBytes bytes
-  NullPtr            -> return nullPtr
+takePtrSnapshot (Heap GPU, _) _ = error "Snapshots of GPU memory not implemented"
+takePtrSnapshot (Stack   , _) _ = error "Can't take snapshots of stack memory"
 
 loadPtrBytes :: RawPtr -> IO BS.ByteString
 loadPtrBytes ptr = do
   numBytes <- fromIntegral <$> dexAllocSize ptr
   liftM BS.pack $ peekArray numBytes $ castPtr ptr
 
-storePtrBytes :: BS.ByteString -> IO RawPtr
-storePtrBytes xs = do
-  let xs' = BS.unpack xs
-  ptr <- dexMalloc $ fromIntegral $ length xs'
-  pokeArray (castPtr ptr) xs'
-  return ptr
-
 loadPtrPtrs :: RawPtr -> IO [RawPtr]
 loadPtrPtrs ptr = do
   numBytes <- fromIntegral <$> dexAllocSize ptr
   peekArray (numBytes `div` ptrSize) $ castPtr ptr
 
+restoreSnapshot :: HasPtrs a => WithSnapshot a -> IO a
+restoreSnapshot (WithSnapshot x snapshots) =
+  flip evalStateT snapshots $ flip traversePtrs x \_ _ -> do
+    (s:ss) <- get
+    put ss
+    lift $ restorePtrSnapshot s
+
+restorePtrSnapshot :: PtrSnapshot -> IO RawPtr
+restorePtrSnapshot snapshot = case snapshot of
+  PtrArray  children -> storePtrPtrs =<< mapM restorePtrSnapshot children
+  ByteArray bytes    -> storePtrBytes bytes
+  NullPtr            -> return nullPtr
+
+storePtrBytes :: BS.ByteString -> IO RawPtr
+storePtrBytes xs = do
+  let numBytes = BS.length xs
+  destPtr <- dexMalloc $ fromIntegral numBytes
+  -- this is safe because we don't modify srcPtr's memory or let it escape
+  unsafeUseAsCString xs \srcPtr ->
+    memcpy (castPtr destPtr) (castPtr srcPtr) numBytes
+  return destPtr
+
 storePtrPtrs :: [RawPtr] -> IO RawPtr
 storePtrPtrs ptrs = do
   ptr <- dexMalloc $ fromIntegral $ length ptrs * ptrSize

src/lib/TopLevel.hs

@@ -9,7 +9,7 @@
 {-# LANGUAGE DeriveGeneric #-}
 
 module TopLevel (evalSourceBlock, evalDecl, evalSource, evalFile,
-                 bindingsToTopEnv, initTopEnv, EvalConfig (..), TopEnv (..)) where
+                 initTopEnv, EvalConfig (..), TopEnv (..)) where
 
 import Control.Monad.State.Strict
 import Control.Monad.Reader
@@ -330,9 +330,6 @@ instance Monoid TopEnv where
 moduleStatus :: ModuleName -> ModuleImportStatus -> TopEnv
 moduleStatus name status = mempty { modulesImported = M.singleton name status}
 
-bindingsToTopEnv :: Bindings -> TopEnv
-bindingsToTopEnv bindings = mempty { topBindings = bindings }
-
 instance HasPtrs TopEnv where
   traversePtrs f (TopEnv bindings status) =
     TopEnv <$> traverse (traversePtrs f) bindings <*> pure status

tests/io-tests.dx

@@ -72,6 +72,3 @@ unsafeIO \().
 
 :p dex_test_mode ()
 > True
-
-:p justSomeDataToTestCaching
-> [(AsList 3 [1, 2, 3]), (AsList 0 []), (AsList 2 [1, 2])]