Begin reviving the Python bridge

The Python bindings have not been updated since the safer name rewrite
so they were completely broken. This is the first change that tries to
restore them to their former glory.

Author: apaszke

.github/workflows/python-ci.yaml

@@ -4,7 +4,7 @@ on:
   push:
     branches: [ main ]
   pull_request:
-    branches: [ main ]
+    branches: [ main, safe-names-dev ]
 
 jobs:
   build:

makefile

@@ -96,11 +96,11 @@ install: dexrt-llvm
 	$(STACK) install $(STACK_BIN_PATH) --flag dex:optimized $(STACK_FLAGS)
 
 build-prof: dexrt-llvm
-	$(STACK) build $(STACK_FLAGS) $(PROF) --flag dex:-foreign --flag dex:debug
+	$(STACK) build $(STACK_FLAGS) $(PROF) --flag dex:debug
 
 # For some reason stack fails to detect modifications to foreign library files
 build-ffis: dexrt-llvm
-	$(STACK) build $(STACK_FLAGS) --force-dirty
+	$(STACK) build $(STACK_FLAGS) --force-dirty --flag dex:foreign
 	$(eval STACK_INSTALL_DIR=$(shell $(STACK) path --local-install-root))
 	cp $(STACK_INSTALL_DIR)/lib/libDex.so python/dex/
 	cp $(STACK_INSTALL_DIR)/lib/libDex.so julia/deps/

python/dex/__init__.py

@@ -21,18 +21,21 @@ class Module:
 
   def __init__(self, source):
     self._as_parameter_ = api.eval(prelude, api.as_cstr(source))
-    if not self._as_parameter_:
-      api.raise_from_dex()
+    if not self._as_parameter_: api.raise_from_dex()
+
+  @classmethod
+  def _from_ptr(cls, ptr):
+    if not ptr: api.raise_from_dex()
+    self = super().__new__(cls)
+    self._as_parameter_ = ptr
+    return self
 
   def __del__(self):
-    if api.nofree:
-      return
+    if api.nofree: return
     api.destroyContext(self)
 
   def __getattr__(self, name):
     result = api.lookup(self, api.as_cstr(name))
-    if not result:
-      api.raise_from_dex()
     return Atom._from_ptr(result, self)
 
 
@@ -46,13 +49,11 @@ def __init__(self):
 prelude = Prelude()
 
 
-def eval(expr: str, module=prelude, _env=None):
-  if _env is None:
-    _env = module
-  result = api.evalExpr(_env, api.as_cstr(expr))
-  if not result:
-    api.raise_from_dex()
-  return Atom._from_ptr(result, module)
+def eval(expr: str):
+  # TODO: Query a free source name
+  _final_env = Module._from_ptr(api.eval(prelude, api.as_cstr("python_result = " + expr)))
+  result = api.lookup(_final_env, api.as_cstr("python_result"))
+  return Atom._from_ptr(result, _final_env)
 
 
 class Atom:
@@ -77,6 +78,7 @@ def __init__(self, value):
 
   @classmethod
   def _from_ptr(cls, ptr, module):
+    if not ptr: api.raise_from_dex()
     self = super().__new__(cls)
     self._as_parameter_ = ptr
     self.module = module
@@ -88,7 +90,7 @@ def __del__(self):
 
   def __repr__(self):
     # TODO: Free!
-    return api.from_cstr(api.print(self))
+    return api.from_cstr(api.print(self.module, self))
 
   def __int__(self):
     return int(self._as_scalar())
@@ -107,6 +109,7 @@ def _as_scalar(self):
     return value.value
 
   def __call__(self, *args):
+    raise NotImplementedError()
     # TODO: Make those calls more hygenic
     env = self.module
     for i, atom in enumerate(it.chain((self,), args)):
@@ -118,6 +121,7 @@ def __call__(self, *args):
     return eval(" ".join(f"python_arg{i}" for i in range(len(args) + 1)), module=self.module, _env=env)
 
   def compile(self):
+    raise NotImplementedError()
     func_ptr = api.compile(api.jit, self.module, self)
     if not func_ptr:
       api.raise_from_dex()

python/dex/api.py

@@ -71,14 +71,12 @@ def dex_func(name, *signature):
 createContext  = dex_func('dexCreateContext',  HsContextPtr)
 destroyContext = dex_func('dexDestroyContext', HsContextPtr, None)
 
-eval     = dex_func('dexEval',     HsContextPtr, ctypes.c_char_p,            HsContextPtr)
-insert   = dex_func('dexInsert',   HsContextPtr, ctypes.c_char_p, HsAtomPtr, HsContextPtr)
-evalExpr = dex_func('dexEvalExpr', HsContextPtr, ctypes.c_char_p,            HsAtomPtr)
-lookup   = dex_func('dexLookup',   HsContextPtr, ctypes.c_char_p,            HsAtomPtr)
-
-print     = dex_func('dexPrint',     HsAtomPtr,           ctypes.c_char_p)
-toCAtom   = dex_func('dexToCAtom',   HsAtomPtr, CAtomPtr, ctypes.c_int)
-fromCAtom = dex_func('dexFromCAtom', CAtomPtr,            HsAtomPtr)
+eval     = dex_func('dexEval',     HsContextPtr, ctypes.c_char_p, HsContextPtr)
+lookup   = dex_func('dexLookup',   HsContextPtr, ctypes.c_char_p, HsAtomPtr)
+
+print     = dex_func('dexPrint',     HsContextPtr, HsAtomPtr, ctypes.c_char_p)
+toCAtom   = dex_func('dexToCAtom',   HsAtomPtr,    CAtomPtr,  ctypes.c_int)
+fromCAtom = dex_func('dexFromCAtom', CAtomPtr,                HsAtomPtr)
 
 createJIT  = dex_func('dexCreateJIT',  HsJITPtr)
 destroyJIT = dex_func('dexDestroyJIT', HsJITPtr, None)

python/tests/api_test.py

@@ -27,6 +27,7 @@ def test_module_attrs(self):
     assert str(m.x) == "2.5"
     assert str(m.y) == "[2, 3, 4]"
 
+  @unittest.skip
   def test_function_call(self):
     m = dex.Module(dedent("""
     def addOne (x: Float) : Float = x + 1.0

python/tests/jax_test.py

@@ -14,6 +14,7 @@
 from dex.interop.jax import primitive
 
 
+@unittest.skip
 class JAXTest(unittest.TestCase):
 
   def test_impl_scalar(self):

python/tests/jit_test.py

@@ -29,6 +29,7 @@ def test(self):
     return check_atom(dex.eval(dex_source), reference, args_iter)
   return test
 
+@unittest.skip
 class JITTest(unittest.TestCase):
   test_sigmoid = expr_test(r"\x:Float. 1.0 / (1.0 + exp(-x))",
                            lambda x: np.float32(1.0) / (np.float32(1.0) + np.exp(-x)),
@@ -75,7 +76,7 @@ def myTranspose (n: Int) ?-> (m: Int) ?->
   def test_tuple_return(self):
     dex_func = dex.eval(r"\x: ((Fin 10) => Float). (x, 2. .* x, 3. .* x)")
     reference = lambda x: (x, 2 * x, 3 * x)
-    
+
     x = np.arange(10, dtype=np.float32)
 
     dex_output = dex_func.compile()(x)

src/Dex/Foreign/API.hs

@@ -9,8 +9,6 @@ module Dex.Foreign.API where
 import Foreign.Ptr
 import Foreign.C
 
-import Syntax
-
 import Dex.Foreign.Context
 import Dex.Foreign.Serialize
 import Dex.Foreign.JIT
@@ -25,20 +23,19 @@ import Dex.Foreign.JIT
 -- Context
 foreign export ccall "dexCreateContext"  dexCreateContext  :: IO (Ptr Context)
 foreign export ccall "dexDestroyContext" dexDestroyContext :: Ptr Context -> IO ()
-foreign export ccall "dexInsert"   dexInsert   :: Ptr Context -> CString   -> Ptr Atom -> IO (Ptr Context)
+foreign export ccall "dexInsert"   dexInsert   :: Ptr Context -> CString   -> Ptr AtomEx -> IO (Ptr Context)
 foreign export ccall "dexEval"     dexEval     :: Ptr Context -> CString   -> IO (Ptr Context)
-foreign export ccall "dexEvalExpr" dexEvalExpr :: Ptr Context -> CString   -> IO (Ptr Atom)
-foreign export ccall "dexLookup"   dexLookup   :: Ptr Context -> CString   -> IO (Ptr Atom)
+foreign export ccall "dexLookup"   dexLookup   :: Ptr Context -> CString   -> IO (Ptr AtomEx)
 
 -- Serialization
-foreign export ccall "dexPrint"     dexPrint      :: Ptr Atom                      -> IO CString
-foreign export ccall "dexToCAtom"   dexToCAtom    :: Ptr Atom    -> Ptr CAtom      -> IO CInt
-foreign export ccall "dexFromCAtom" dexFromCAtom  :: Ptr CAtom                     -> IO (Ptr Atom)
+foreign export ccall "dexPrint"     dexPrint      :: Ptr Context -> Ptr AtomEx     -> IO CString
+foreign export ccall "dexToCAtom"   dexToCAtom    :: Ptr AtomEx  -> Ptr CAtom      -> IO CInt
+foreign export ccall "dexFromCAtom" dexFromCAtom  :: Ptr CAtom                     -> IO (Ptr AtomEx)
 
 -- JIT
 foreign export ccall "dexCreateJIT"  dexCreateJIT  :: IO (Ptr JIT)
 foreign export ccall "dexDestroyJIT" dexDestroyJIT :: Ptr JIT -> IO ()
-foreign export ccall "dexCompile"    dexCompile    :: Ptr JIT -> Ptr Context -> Ptr Atom -> IO (Ptr NativeFunction)
+foreign export ccall "dexCompile"    dexCompile    :: Ptr JIT -> Ptr Context -> Ptr AtomEx -> IO (Ptr NativeFunction)
 foreign export ccall "dexUnload"     dexUnload     :: Ptr JIT -> Ptr NativeFunction -> IO ()
 foreign export ccall "dexGetFunctionSignature"  dexGetFunctionSignature  :: Ptr JIT -> Ptr NativeFunction -> IO (Ptr ExportedSignature)
 foreign export ccall "dexFreeFunctionSignature" dexFreeFunctionSignature :: Ptr ExportedSignature -> IO ()

src/Dex/Foreign/Context.hs

@@ -4,18 +4,21 @@
 -- license that can be found in the LICENSE file or at
 -- https://developers.google.com/open-source/licenses/bsd
 
+{-# LANGUAGE GADTs #-}
+
 module Dex.Foreign.Context (
-  Context (..),
+  Context (..), AtomEx (..),
   setError,
   dexCreateContext, dexDestroyContext,
   dexInsert, dexLookup,
-  dexEval, dexEvalExpr,
+  dexEval,
   ) where
 
 import Foreign.Ptr
 import Foreign.StablePtr
 import Foreign.C.String
 
+import Control.Monad.IO.Class
 import Data.String
 import Data.Int
 import Data.Functor
@@ -26,18 +29,18 @@ import Resources
 import Syntax  hiding (sizeOf)
 import Type
 import TopLevel
-import Env hiding (Tag)
+import Name
 import PPrint
 import Err
+import Parser
+import Builder
 
 import Dex.Foreign.Util
 
-import SaferNames.Bridge
-import qualified SaferNames.Syntax as S
-import qualified SaferNames.Parser as S
-
 
 data Context = Context EvalConfig TopStateEx
+data AtomEx where
+  AtomEx :: Atom n -> AtomEx
 
 foreign import ccall "_internal_dexSetError" internalSetErrorPtr :: CString -> Int64 -> IO ()
 setError :: String -> IO ()
@@ -46,70 +49,45 @@ setError msg = withCStringLen msg $ \(ptr, len) ->
 
 dexCreateContext :: IO (Ptr Context)
 dexCreateContext = do
-  let evalConfig = EvalConfig LLVM Nothing Nothing
-  maybePreludeEnv <- evalPrelude evalConfig preludeSource
-  case maybePreludeEnv of
-    Success preludeEnv -> toStablePtr $ Context evalConfig preludeEnv
-    Failure  err       -> nullPtr <$ setError ("Failed to initialize standard library: " ++ pprint err)
-  where
-    evalPrelude :: EvalConfig -> String -> IO (Except TopStateEx)
-    evalPrelude opts sourceText = do
-      (results, env) <- runInterblockM opts initTopState $
-                            map snd <$> evalSourceText sourceText
-      return $ env `unlessError` results
-      where
-        unlessError :: TopStateEx -> [Result] -> Except TopStateEx
-        result `unlessError` []                        = Success result
-        _      `unlessError` ((Result _ (Failure err)):_) = Failure err
-        result `unlessError` (_:t                       ) = result `unlessError` t
+  let evalConfig = EvalConfig LLVM Nothing Nothing Nothing
+  cachedEnv <- loadCache
+  runTopperM evalConfig cachedEnv (evalSourceBlockRepl preludeImportBlock) >>= \case
+    (Result [] (Success  ()), preludeEnv) -> toStablePtr $ Context evalConfig preludeEnv
+    (Result _  (Failure err), _         ) -> nullPtr <$
+      setError ("Failed to initialize standard library: " ++ pprint err)
 
 dexDestroyContext :: Ptr Context -> IO ()
 dexDestroyContext = freeStablePtr . castPtrToStablePtr . castPtr
 
 dexEval :: Ptr Context -> CString -> IO (Ptr Context)
 dexEval ctxPtr sourcePtr = do
-  Context evalConfig env <- fromStablePtr ctxPtr
+  Context evalConfig initEnv <- fromStablePtr ctxPtr
   source <- peekCString sourcePtr
-  (results, finalEnv) <- runInterblockM evalConfig env $ evalSourceText source
+  (results, finalEnv) <- runTopperM evalConfig initEnv $ evalSourceText source
   let anyError = asum $ fmap (\case (_, Result _ (Failure err)) -> Just err; _ -> Nothing) results
   case anyError of
     Nothing  -> toStablePtr $ Context evalConfig finalEnv
     Just err -> setError (pprint err) $> nullPtr
 
-dexInsert :: Ptr Context -> CString -> Ptr Atom -> IO (Ptr Context)
+dexInsert :: Ptr Context -> CString -> Ptr AtomEx -> IO (Ptr Context)
 dexInsert ctxPtr namePtr atomPtr = do
-  Context evalConfig (TopStateEx env) <- fromStablePtr ctxPtr
-  name <- fromString <$> peekCString namePtr
-  atom <- fromStablePtr atomPtr
-  let freshName = genFresh (Name GenName (fromString name) 0) (topBindings $ topStateD env)
-  let newBinding = AtomBinderInfo (getType atom) (LetBound PlainLet (Atom atom))
-  let evaluated = EvaluatedModule (freshName @> newBinding) mempty
-                                  (SourceMap (M.singleton name (SrcAtomName freshName)))
-  let envNew = extendTopStateD env evaluated
-  toStablePtr $ Context evalConfig $ envNew
+  Context evalConfig initEnv <- fromStablePtr ctxPtr
+  sourceName <- peekCString namePtr
+  AtomEx atom <- fromStablePtr atomPtr
+  (_, finalEnv) <- runTopperM evalConfig initEnv do
+    -- TODO: Check if atom is compatible with context! Use module name?
+    name <- emitTopLet (fromString sourceName) PlainLet $ Atom $ unsafeCoerceE atom
+    emitSourceMap $ SourceMap $ M.singleton sourceName [ModuleVar Main $ Just $ UAtomVar name]
+  toStablePtr $ Context evalConfig finalEnv
 
-dexEvalExpr :: Ptr Context -> CString -> IO (Ptr Atom)
-dexEvalExpr ctxPtr sourcePtr = do
-  Context evalConfig env <- fromStablePtr ctxPtr
-  source <- peekCString sourcePtr
-  case S.parseExpr source of
-    Success expr -> do
-      let (v, m) = S.exprAsModule expr
-      let block = S.SourceBlock 0 0 LogNothing source (S.RunModule m) Nothing
-      (Result [] maybeErr, newState) <- runInterblockM evalConfig env $ evalSourceBlock block
-      case maybeErr of
-        Success () -> do
-          let Success (AtomBinderInfo _ (LetBound _ (Atom atom))) =
-                lookupSourceName newState v
-          toStablePtr atom
-        Failure err -> setError (pprint err) $> nullPtr
-    Failure err -> setError (pprint err) $> nullPtr
-
-dexLookup :: Ptr Context -> CString -> IO (Ptr Atom)
+dexLookup :: Ptr Context -> CString -> IO (Ptr AtomEx)
 dexLookup ctxPtr namePtr = do
-  Context _ env <- fromStablePtr ctxPtr
+  Context evalConfig env <- fromStablePtr ctxPtr
   name <- peekCString namePtr
-  case lookupSourceName env (fromString name) of
-    Success (AtomBinderInfo _ (LetBound _ (Atom atom))) -> toStablePtr atom
-    Failure _ -> setError "Unbound name" $> nullPtr
-    Success _ -> setError "Looking up an expression" $> nullPtr
+  fst <$> runTopperM evalConfig env do
+    lookupSourceMap name >>= \case
+      Just (UAtomVar v) -> lookupAtomName v >>= \case
+        LetBound (DeclBinding _ _ (Atom atom)) -> liftIO $ toStablePtr $ AtomEx atom
+        _ -> liftIO $ setError "Looking up an unevaluated atom?" $> nullPtr
+      Just _  -> liftIO $ setError "Only Atom names can be looked up" $> nullPtr
+      Nothing -> liftIO $ setError "Unbound name" $> nullPtr