{-# LANGUAGE OverloadedStrings, GeneralizedNewtypeDeriving, LambdaCase #-}
{-# LANGUAGE TupleSections #-}
-- | A generic Python code generator which is polymorphic in the type
-- of the operations.  Concretely, we use this to handle both
-- sequential and PyOpenCL Python code.
module Futhark.CodeGen.Backends.GenericPython
  ( compileProg
  , Constructor (..)
  , emptyConstructor

  , compileName
  , compileVar
  , compileDim
  , compileExp
  , compileCode
  , compilePrimValue
  , compilePrimType
  , compilePrimTypeExt
  , compilePrimToNp
  , compilePrimToExtNp

  , Operations (..)
  , defaultOperations

  , unpackDim

  , CompilerM (..)
  , OpCompiler
  , WriteScalar
  , ReadScalar
  , Allocate
  , Copy
  , StaticArray
  , EntryOutput
  , EntryInput

  , CompilerEnv(..)
  , CompilerState(..)
  , stm
  , atInit
  , collect'
  , collect
  , simpleCall

  , copyMemoryDefaultSpace
  ) where

import Control.Monad.Identity
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.RWS
import Data.Maybe
import qualified Data.Map as M

import Futhark.Representation.Primitive hiding (Bool)
import Futhark.MonadFreshNames
import Futhark.Representation.AST.Syntax (Space(..))
import qualified Futhark.CodeGen.ImpCode as Imp
import Futhark.CodeGen.Backends.GenericPython.AST
import Futhark.CodeGen.Backends.GenericPython.Options
import Futhark.CodeGen.Backends.GenericPython.Definitions
import Futhark.Util (zEncodeString)
import Futhark.Representation.AST.Attributes (isBuiltInFunction)

-- | A substitute expression compiler, tried before the main
-- compilation function.
type OpCompiler op s = op -> CompilerM op s ()

-- | Write a scalar to the given memory block with the given index and
-- in the given memory space.
type WriteScalar op s = PyExp -> PyExp -> PrimType -> Imp.SpaceId -> PyExp
                        -> CompilerM op s ()

-- | Read a scalar from the given memory block with the given index and
-- in the given memory space.
type ReadScalar op s = PyExp -> PyExp -> PrimType -> Imp.SpaceId
                       -> CompilerM op s PyExp

-- | Allocate a memory block of the given size in the given memory
-- space, saving a reference in the given variable name.
type Allocate op s = PyExp -> PyExp -> Imp.SpaceId
                     -> CompilerM op s ()

-- | Copy from one memory block to another.
type Copy op s = PyExp -> PyExp -> Imp.Space ->
                 PyExp -> PyExp -> Imp.Space ->
                 PyExp -> PrimType ->
                 CompilerM op s ()

-- | Create a static array of values - initialised at load time.
type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> Imp.ArrayContents -> CompilerM op s ()

-- | Construct the Python array being returned from an entry point.
type EntryOutput op s = VName -> Imp.SpaceId ->
                        PrimType -> Imp.Signedness ->
                        [Imp.DimSize] ->
                        CompilerM op s PyExp

-- | Unpack the array being passed to an entry point.
type EntryInput op s = PyExp -> Imp.SpaceId ->
                       PrimType -> Imp.Signedness ->
                       [Imp.DimSize] ->
                       PyExp ->
                       CompilerM op s ()


data Operations op s = Operations { Operations op s -> WriteScalar op s
opsWriteScalar :: WriteScalar op s
                                  , Operations op s -> ReadScalar op s
opsReadScalar :: ReadScalar op s
                                  , Operations op s -> Allocate op s
opsAllocate :: Allocate op s
                                  , Operations op s -> Copy op s
opsCopy :: Copy op s
                                  , Operations op s -> StaticArray op s
opsStaticArray :: StaticArray op s
                                  , Operations op s -> OpCompiler op s
opsCompiler :: OpCompiler op s
                                  , Operations op s -> EntryOutput op s
opsEntryOutput :: EntryOutput op s
                                  , Operations op s -> EntryInput op s
opsEntryInput :: EntryInput op s
                                  }

-- | A set of operations that fail for every operation involving
-- non-default memory spaces.  Uses plain pointers and @malloc@ for
-- memory management.
defaultOperations :: Operations op s
defaultOperations :: Operations op s
defaultOperations = Operations :: forall op s.
WriteScalar op s
-> ReadScalar op s
-> Allocate op s
-> Copy op s
-> StaticArray op s
-> OpCompiler op s
-> EntryOutput op s
-> EntryInput op s
-> Operations op s
Operations { opsWriteScalar :: WriteScalar op s
opsWriteScalar = WriteScalar op s
forall p p p p p a. p -> p -> p -> p -> p -> a
defWriteScalar
                               , opsReadScalar :: ReadScalar op s
opsReadScalar = ReadScalar op s
forall p p p p a. p -> p -> p -> p -> a
defReadScalar
                               , opsAllocate :: Allocate op s
opsAllocate  = Allocate op s
forall p p p a. p -> p -> p -> a
defAllocate
                               , opsCopy :: Copy op s
opsCopy = Copy op s
forall p p p p p p p p a. p -> p -> p -> p -> p -> p -> p -> p -> a
defCopy
                               , opsStaticArray :: StaticArray op s
opsStaticArray = StaticArray op s
forall p p p p a. p -> p -> p -> p -> a
defStaticArray
                               , opsCompiler :: OpCompiler op s
opsCompiler = OpCompiler op s
forall p a. p -> a
defCompiler
                               , opsEntryOutput :: EntryOutput op s
opsEntryOutput = EntryOutput op s
forall p p p p a. p -> p -> p -> p -> a
defEntryOutput
                               , opsEntryInput :: EntryInput op s
opsEntryInput = EntryInput op s
forall p p p p a. p -> p -> p -> p -> a
defEntryInput
                               }
  where defWriteScalar :: p -> p -> p -> p -> p -> a
defWriteScalar p
_ p
_ p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot write to non-default memory space because I am dumb"
        defReadScalar :: p -> p -> p -> p -> a
defReadScalar p
_ p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot read from non-default memory space"
        defAllocate :: p -> p -> p -> a
defAllocate p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot allocate in non-default memory space"
        defCopy :: p -> p -> p -> p -> p -> p -> p -> p -> a
defCopy p
_ p
_ p
_ p
_ p
_ p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot copy to or from non-default memory space"
        defStaticArray :: p -> p -> p -> p -> a
defStaticArray p
_ p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot create static array in non-default memory space"
        defCompiler :: p -> a
defCompiler p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"The default compiler cannot compile extended operations"
        defEntryOutput :: p -> p -> p -> p -> a
defEntryOutput p
_ p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot return array not in default memory space"
        defEntryInput :: p -> p -> p -> p -> a
defEntryInput p
_ p
_ p
_ p
_ =
          [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot accept array not in default memory space"

data CompilerEnv op s = CompilerEnv
  { CompilerEnv op s -> Operations op s
envOperations :: Operations op s
  , CompilerEnv op s -> Map VName PyExp
envVarExp :: M.Map VName PyExp
  }

envOpCompiler :: CompilerEnv op s -> OpCompiler op s
envOpCompiler :: CompilerEnv op s -> OpCompiler op s
envOpCompiler = Operations op s -> OpCompiler op s
forall op s. Operations op s -> OpCompiler op s
opsCompiler (Operations op s -> OpCompiler op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> OpCompiler op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envReadScalar :: CompilerEnv op s -> ReadScalar op s
envReadScalar :: CompilerEnv op s -> ReadScalar op s
envReadScalar = Operations op s -> ReadScalar op s
forall op s. Operations op s -> ReadScalar op s
opsReadScalar (Operations op s -> ReadScalar op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> ReadScalar op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envWriteScalar :: CompilerEnv op s -> WriteScalar op s
envWriteScalar :: CompilerEnv op s -> WriteScalar op s
envWriteScalar = Operations op s -> WriteScalar op s
forall op s. Operations op s -> WriteScalar op s
opsWriteScalar (Operations op s -> WriteScalar op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> WriteScalar op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envAllocate :: CompilerEnv op s -> Allocate op s
envAllocate :: CompilerEnv op s -> Allocate op s
envAllocate = Operations op s -> Allocate op s
forall op s. Operations op s -> Allocate op s
opsAllocate (Operations op s -> Allocate op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> Allocate op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envCopy :: CompilerEnv op s -> Copy op s
envCopy :: CompilerEnv op s -> Copy op s
envCopy = Operations op s -> Copy op s
forall op s. Operations op s -> Copy op s
opsCopy (Operations op s -> Copy op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> Copy op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envStaticArray :: CompilerEnv op s -> StaticArray op s
envStaticArray :: CompilerEnv op s -> StaticArray op s
envStaticArray = Operations op s -> StaticArray op s
forall op s. Operations op s -> StaticArray op s
opsStaticArray (Operations op s -> StaticArray op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> StaticArray op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envEntryOutput :: CompilerEnv op s -> EntryOutput op s
envEntryOutput :: CompilerEnv op s -> EntryOutput op s
envEntryOutput = Operations op s -> EntryOutput op s
forall op s. Operations op s -> EntryOutput op s
opsEntryOutput (Operations op s -> EntryOutput op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> EntryOutput op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

envEntryInput :: CompilerEnv op s -> EntryInput op s
envEntryInput :: CompilerEnv op s -> EntryInput op s
envEntryInput = Operations op s -> EntryInput op s
forall op s. Operations op s -> EntryInput op s
opsEntryInput (Operations op s -> EntryInput op s)
-> (CompilerEnv op s -> Operations op s)
-> CompilerEnv op s
-> EntryInput op s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Operations op s
forall op s. CompilerEnv op s -> Operations op s
envOperations

newCompilerEnv :: Operations op s -> CompilerEnv op s
newCompilerEnv :: Operations op s -> CompilerEnv op s
newCompilerEnv Operations op s
ops = CompilerEnv :: forall op s. Operations op s -> Map VName PyExp -> CompilerEnv op s
CompilerEnv { envOperations :: Operations op s
envOperations = Operations op s
ops
                                 , envVarExp :: Map VName PyExp
envVarExp = Map VName PyExp
forall a. Monoid a => a
mempty }

data CompilerState s = CompilerState {
    CompilerState s -> VNameSource
compNameSrc :: VNameSource
  , CompilerState s -> [PyStmt]
compInit :: [PyStmt]
  , CompilerState s -> s
compUserState :: s
}

newCompilerState :: VNameSource -> s -> CompilerState s
newCompilerState :: VNameSource -> s -> CompilerState s
newCompilerState VNameSource
src s
s = CompilerState :: forall s. VNameSource -> [PyStmt] -> s -> CompilerState s
CompilerState { compNameSrc :: VNameSource
compNameSrc = VNameSource
src
                                       , compInit :: [PyStmt]
compInit = []
                                       , compUserState :: s
compUserState = s
s }

newtype CompilerM op s a = CompilerM (RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a)
  deriving (a -> CompilerM op s b -> CompilerM op s a
(a -> b) -> CompilerM op s a -> CompilerM op s b
(forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b)
-> (forall a b. a -> CompilerM op s b -> CompilerM op s a)
-> Functor (CompilerM op s)
forall a b. a -> CompilerM op s b -> CompilerM op s a
forall a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a b. a -> CompilerM op s b -> CompilerM op s a
forall op s a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> CompilerM op s b -> CompilerM op s a
$c<$ :: forall op s a b. a -> CompilerM op s b -> CompilerM op s a
fmap :: (a -> b) -> CompilerM op s a -> CompilerM op s b
$cfmap :: forall op s a b. (a -> b) -> CompilerM op s a -> CompilerM op s b
Functor, Functor (CompilerM op s)
a -> CompilerM op s a
Functor (CompilerM op s)
-> (forall a. a -> CompilerM op s a)
-> (forall a b.
    CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b)
-> (forall a b c.
    (a -> b -> c)
    -> CompilerM op s a -> CompilerM op s b -> CompilerM op s c)
-> (forall a b.
    CompilerM op s a -> CompilerM op s b -> CompilerM op s b)
-> (forall a b.
    CompilerM op s a -> CompilerM op s b -> CompilerM op s a)
-> Applicative (CompilerM op s)
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
forall a. a -> CompilerM op s a
forall op s. Functor (CompilerM op s)
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a. a -> CompilerM op s a
forall a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall op s a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
forall op s a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: CompilerM op s a -> CompilerM op s b -> CompilerM op s a
$c<* :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s a
*> :: CompilerM op s a -> CompilerM op s b -> CompilerM op s b
$c*> :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
liftA2 :: (a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
$cliftA2 :: forall op s a b c.
(a -> b -> c)
-> CompilerM op s a -> CompilerM op s b -> CompilerM op s c
<*> :: CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
$c<*> :: forall op s a b.
CompilerM op s (a -> b) -> CompilerM op s a -> CompilerM op s b
pure :: a -> CompilerM op s a
$cpure :: forall op s a. a -> CompilerM op s a
$cp1Applicative :: forall op s. Functor (CompilerM op s)
Applicative, Applicative (CompilerM op s)
a -> CompilerM op s a
Applicative (CompilerM op s)
-> (forall a b.
    CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b)
-> (forall a b.
    CompilerM op s a -> CompilerM op s b -> CompilerM op s b)
-> (forall a. a -> CompilerM op s a)
-> Monad (CompilerM op s)
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a. a -> CompilerM op s a
forall op s. Applicative (CompilerM op s)
forall a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
forall op s a. a -> CompilerM op s a
forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
forall op s a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> CompilerM op s a
$creturn :: forall op s a. a -> CompilerM op s a
>> :: CompilerM op s a -> CompilerM op s b -> CompilerM op s b
$c>> :: forall op s a b.
CompilerM op s a -> CompilerM op s b -> CompilerM op s b
>>= :: CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
$c>>= :: forall op s a b.
CompilerM op s a -> (a -> CompilerM op s b) -> CompilerM op s b
$cp1Monad :: forall op s. Applicative (CompilerM op s)
Monad,
            MonadState (CompilerState s),
            MonadReader (CompilerEnv op s),
            MonadWriter [PyStmt])

instance MonadFreshNames (CompilerM op s) where
  getNameSource :: CompilerM op s VNameSource
getNameSource = (CompilerState s -> VNameSource) -> CompilerM op s VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> VNameSource
forall s. CompilerState s -> VNameSource
compNameSrc
  putNameSource :: VNameSource -> CompilerM op s ()
putNameSource VNameSource
src = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s -> CompilerState s
s { compNameSrc :: VNameSource
compNameSrc = VNameSource
src }

collect :: CompilerM op s () -> CompilerM op s [PyStmt]
collect :: CompilerM op s () -> CompilerM op s [PyStmt]
collect CompilerM op s ()
m = CompilerM op s ([PyStmt], [PyStmt] -> [PyStmt])
-> CompilerM op s [PyStmt]
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (CompilerM op s ([PyStmt], [PyStmt] -> [PyStmt])
 -> CompilerM op s [PyStmt])
-> CompilerM op s ([PyStmt], [PyStmt] -> [PyStmt])
-> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ do
  ((), [PyStmt]
w) <- CompilerM op s () -> CompilerM op s ((), [PyStmt])
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen CompilerM op s ()
m
  ([PyStmt], [PyStmt] -> [PyStmt])
-> CompilerM op s ([PyStmt], [PyStmt] -> [PyStmt])
forall (m :: * -> *) a. Monad m => a -> m a
return ([PyStmt]
w, [PyStmt] -> [PyStmt] -> [PyStmt]
forall a b. a -> b -> a
const [PyStmt]
forall a. Monoid a => a
mempty)

collect' :: CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' :: CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' CompilerM op s a
m = CompilerM op s ((a, [PyStmt]), [PyStmt] -> [PyStmt])
-> CompilerM op s (a, [PyStmt])
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (CompilerM op s ((a, [PyStmt]), [PyStmt] -> [PyStmt])
 -> CompilerM op s (a, [PyStmt]))
-> CompilerM op s ((a, [PyStmt]), [PyStmt] -> [PyStmt])
-> CompilerM op s (a, [PyStmt])
forall a b. (a -> b) -> a -> b
$ do
  (a
x, [PyStmt]
w) <- CompilerM op s a -> CompilerM op s (a, [PyStmt])
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen CompilerM op s a
m
  ((a, [PyStmt]), [PyStmt] -> [PyStmt])
-> CompilerM op s ((a, [PyStmt]), [PyStmt] -> [PyStmt])
forall (m :: * -> *) a. Monad m => a -> m a
return ((a
x, [PyStmt]
w), [PyStmt] -> [PyStmt] -> [PyStmt]
forall a b. a -> b -> a
const [PyStmt]
forall a. Monoid a => a
mempty)

atInit :: PyStmt -> CompilerM op s ()
atInit :: PyStmt -> CompilerM op s ()
atInit PyStmt
x = (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState s -> CompilerState s) -> CompilerM op s ())
-> (CompilerState s -> CompilerState s) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ \CompilerState s
s ->
  CompilerState s
s { compInit :: [PyStmt]
compInit = CompilerState s -> [PyStmt]
forall s. CompilerState s -> [PyStmt]
compInit CompilerState s
s [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt
x] }

stm :: PyStmt -> CompilerM op s ()
stm :: PyStmt -> CompilerM op s ()
stm PyStmt
x = [PyStmt] -> CompilerM op s ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [PyStmt
x]

futharkFun :: String -> String
futharkFun :: [Char] -> [Char]
futharkFun [Char]
s = [Char]
"futhark_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
zEncodeString [Char]
s

compileOutput :: [Imp.Param] -> [PyExp]
compileOutput :: [Param] -> [PyExp]
compileOutput = (Param -> PyExp) -> [Param] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> PyExp
Var ([Char] -> PyExp) -> (Param -> [Char]) -> Param -> PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
compileName (VName -> [Char]) -> (Param -> VName) -> Param -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName)

runCompilerM :: Operations op s
             -> VNameSource
             -> s
             -> CompilerM op s a
             -> a
runCompilerM :: Operations op s -> VNameSource -> s -> CompilerM op s a -> a
runCompilerM Operations op s
ops VNameSource
src s
userstate (CompilerM RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a
m) =
  (a, [PyStmt]) -> a
forall a b. (a, b) -> a
fst ((a, [PyStmt]) -> a) -> (a, [PyStmt]) -> a
forall a b. (a -> b) -> a -> b
$ RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a
-> CompilerEnv op s -> CompilerState s -> (a, [PyStmt])
forall r w s a. RWS r w s a -> r -> s -> (a, w)
evalRWS RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a
m (Operations op s -> CompilerEnv op s
forall op s. Operations op s -> CompilerEnv op s
newCompilerEnv Operations op s
ops) (VNameSource -> s -> CompilerState s
forall s. VNameSource -> s -> CompilerState s
newCompilerState VNameSource
src s
userstate)

standardOptions :: [Option]
standardOptions :: [Option]
standardOptions = [
  Option :: [Char] -> Maybe Char -> OptionArgument -> [PyStmt] -> Option
Option { optionLongName :: [Char]
optionLongName = [Char]
"write-runtime-to"
         , optionShortName :: Maybe Char
optionShortName = Char -> Maybe Char
forall a. a -> Maybe a
Just Char
't'
         , optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str"
         , optionAction :: [PyStmt]
optionAction =
           [
             PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"runtime_file")
             [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.close" []] []
           , PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"runtime_file") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
             [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"open" [[Char] -> PyExp
Var [Char]
"optarg", [Char] -> PyExp
String [Char]
"w"]
           ]
         },
  Option :: [Char] -> Maybe Char -> OptionArgument -> [PyStmt] -> Option
Option { optionLongName :: [Char]
optionLongName = [Char]
"runs"
         , optionShortName :: Maybe Char
optionShortName = Char -> Maybe Char
forall a. a -> Maybe a
Just Char
'r'
         , optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str"
         , optionAction :: [PyStmt]
optionAction =
           [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"num_runs") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"optarg"
           , PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"do_warmup_run") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ Bool -> PyExp
Bool Bool
True
           ]
         },
  Option :: [Char] -> Maybe Char -> OptionArgument -> [PyStmt] -> Option
Option { optionLongName :: [Char]
optionLongName = [Char]
"entry-point"
         , optionShortName :: Maybe Char
optionShortName = Char -> Maybe Char
forall a. a -> Maybe a
Just Char
'e'
         , optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"str"
         , optionAction :: [PyStmt]
optionAction =
           [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_point") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"optarg" ]
         },
  Option :: [Char] -> Maybe Char -> OptionArgument -> [PyStmt] -> Option
Option { optionLongName :: [Char]
optionLongName = [Char]
"binary-output"
         , optionShortName :: Maybe Char
optionShortName = Char -> Maybe Char
forall a. a -> Maybe a
Just Char
'b'
         , optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument
         , optionAction :: [PyStmt]
optionAction = [PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"binary_output") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ Bool -> PyExp
Bool Bool
True]
         },
  Option :: [Char] -> Maybe Char -> OptionArgument -> [PyStmt] -> Option
Option { optionLongName :: [Char]
optionLongName = [Char]
"tuning"
         , optionShortName :: Maybe Char
optionShortName = Maybe Char
forall a. Maybe a
Nothing
         , optionArgument :: OptionArgument
optionArgument = [Char] -> OptionArgument
RequiredArgument [Char]
"open"
         , optionAction :: [PyStmt]
optionAction = [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"read_tuning_file" [[Char] -> PyExp
Var [Char]
"sizes", [Char] -> PyExp
Var [Char]
"optarg"]]
         }
  ]


-- | The class generated by the code generator must have a
-- constructor, although it can be vacuous.
data Constructor = Constructor [String] [PyStmt]

-- | A constructor that takes no arguments and does nothing.
emptyConstructor :: Constructor
emptyConstructor :: Constructor
emptyConstructor = [[Char]] -> [PyStmt] -> Constructor
Constructor [[Char]
"self"] [PyStmt
Pass]

constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef
constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef
constructorToFunDef (Constructor [[Char]]
params [PyStmt]
body) [PyStmt]
at_init =
  [Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def [Char]
"__init__" [[Char]]
params ([PyStmt] -> PyFunDef) -> [PyStmt] -> PyFunDef
forall a b. (a -> b) -> a -> b
$ [PyStmt]
body [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. Semigroup a => a -> a -> a
<> [PyStmt]
at_init

compileProg :: MonadFreshNames m =>
               Maybe String
            -> Constructor
            -> [PyStmt]
            -> [PyStmt]
            -> Operations op s
            -> s
            -> [PyStmt]
            -> [Option]
            -> Imp.Definitions op
            -> m String
compileProg :: Maybe [Char]
-> Constructor
-> [PyStmt]
-> [PyStmt]
-> Operations op s
-> s
-> [PyStmt]
-> [Option]
-> Definitions op
-> m [Char]
compileProg Maybe [Char]
module_name Constructor
constructor [PyStmt]
imports [PyStmt]
defines Operations op s
ops s
userstate [PyStmt]
pre_timing [Option]
options Definitions op
prog = do
  VNameSource
src <- m VNameSource
forall (m :: * -> *). MonadFreshNames m => m VNameSource
getNameSource
  let prog' :: [PyStmt]
prog' = Operations op s
-> VNameSource -> s -> CompilerM op s [PyStmt] -> [PyStmt]
forall op s a.
Operations op s -> VNameSource -> s -> CompilerM op s a -> a
runCompilerM Operations op s
ops VNameSource
src s
userstate CompilerM op s [PyStmt]
forall s. CompilerM op s [PyStmt]
compileProg'
      maybe_shebang :: [Char]
maybe_shebang =
        case Maybe [Char]
module_name of Maybe [Char]
Nothing -> [Char]
"#!/usr/bin/env python\n"
                            Just [Char]
_  -> [Char]
""
  [Char] -> m [Char]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> m [Char]) -> [Char] -> m [Char]
forall a b. (a -> b) -> a -> b
$ [Char]
maybe_shebang [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++
    PyProg -> [Char]
forall a. Pretty a => a -> [Char]
pretty ([PyStmt] -> PyProg
PyProg ([PyStmt] -> PyProg) -> [PyStmt] -> PyProg
forall a b. (a -> b) -> a -> b
$ [PyStmt]
imports [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
            [ [Char] -> Maybe [Char] -> PyStmt
Import [Char]
"argparse" Maybe [Char]
forall a. Maybe a
Nothing
            , PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"sizes") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [(PyExp, PyExp)] -> PyExp
Dict []
            ] [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
            [PyStmt]
defines [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
            [[Char] -> PyStmt
Escape [Char]
pyUtility] [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
            [PyStmt]
prog')
  where Imp.Definitions Constants op
consts (Imp.Functions [(Name, Function op)]
funs) = Definitions op
prog
        compileProg' :: CompilerM op s [PyStmt]
compileProg' = Constants op -> CompilerM op s [PyStmt] -> CompilerM op s [PyStmt]
forall op s a. Constants op -> CompilerM op s a -> CompilerM op s a
withConstantSubsts Constants op
consts (CompilerM op s [PyStmt] -> CompilerM op s [PyStmt])
-> CompilerM op s [PyStmt] -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ do

          Constants op -> CompilerM op s ()
forall op s. Constants op -> CompilerM op s ()
compileConstants Constants op
consts

          [PyFunDef]
definitions <- ((Name, Function op) -> CompilerM op s PyFunDef)
-> [(Name, Function op)] -> CompilerM op s [PyFunDef]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, Function op) -> CompilerM op s PyFunDef
forall op s. (Name, Function op) -> CompilerM op s PyFunDef
compileFunc [(Name, Function op)]
funs
          [PyStmt]
at_inits <- (CompilerState s -> [PyStmt]) -> CompilerM op s [PyStmt]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CompilerState s -> [PyStmt]
forall s. CompilerState s -> [PyStmt]
compInit

          let constructor' :: PyFunDef
constructor' = Constructor -> [PyStmt] -> PyFunDef
constructorToFunDef Constructor
constructor [PyStmt]
at_inits

          case Maybe [Char]
module_name of
            Just [Char]
name -> do
              ([PyFunDef]
entry_points, [(PyExp, PyExp)]
entry_point_types) <-
                [(PyFunDef, (PyExp, PyExp))] -> ([PyFunDef], [(PyExp, PyExp)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PyFunDef, (PyExp, PyExp))] -> ([PyFunDef], [(PyExp, PyExp)]))
-> CompilerM op s [(PyFunDef, (PyExp, PyExp))]
-> CompilerM op s ([PyFunDef], [(PyExp, PyExp)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, Function op) -> CompilerM op s (PyFunDef, (PyExp, PyExp)))
-> [(Name, Function op)]
-> CompilerM op s [(PyFunDef, (PyExp, PyExp))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, Function op) -> CompilerM op s (PyFunDef, (PyExp, PyExp))
forall op s.
(Name, Function op) -> CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun (((Name, Function op) -> Bool)
-> [(Name, Function op)] -> [(Name, Function op)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Function op -> Bool
forall a. FunctionT a -> Bool
Imp.functionEntry (Function op -> Bool)
-> ((Name, Function op) -> Function op)
-> (Name, Function op)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Function op) -> Function op
forall a b. (a, b) -> b
snd) [(Name, Function op)]
funs)
              [PyStmt] -> CompilerM op s [PyStmt]
forall (m :: * -> *) a. Monad m => a -> m a
return [PyClassDef -> PyStmt
ClassDef (PyClassDef -> PyStmt) -> PyClassDef -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyStmt] -> PyClassDef
Class [Char]
name ([PyStmt] -> PyClassDef) -> [PyStmt] -> PyClassDef
forall a b. (a -> b) -> a -> b
$
                       PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_points") ([(PyExp, PyExp)] -> PyExp
Dict [(PyExp, PyExp)]
entry_point_types) PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
                       (PyFunDef -> PyStmt) -> [PyFunDef] -> [PyStmt]
forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef (PyFunDef
constructor' PyFunDef -> [PyFunDef] -> [PyFunDef]
forall a. a -> [a] -> [a]
: [PyFunDef]
definitions [PyFunDef] -> [PyFunDef] -> [PyFunDef]
forall a. [a] -> [a] -> [a]
++ [PyFunDef]
entry_points)]
            Maybe [Char]
Nothing -> do
              let classinst :: PyStmt
classinst = PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"self") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"internal" []
              ([PyFunDef]
entry_point_defs, [[Char]]
entry_point_names, [PyExp]
entry_points) <-
                [(PyFunDef, [Char], PyExp)] -> ([PyFunDef], [[Char]], [PyExp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(PyFunDef, [Char], PyExp)] -> ([PyFunDef], [[Char]], [PyExp]))
-> CompilerM op s [(PyFunDef, [Char], PyExp)]
-> CompilerM op s ([PyFunDef], [[Char]], [PyExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, Function op) -> CompilerM op s (PyFunDef, [Char], PyExp))
-> [(Name, Function op)]
-> CompilerM op s [(PyFunDef, [Char], PyExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([PyStmt]
-> (Name, Function op) -> CompilerM op s (PyFunDef, [Char], PyExp)
forall op s.
[PyStmt]
-> (Name, Function op) -> CompilerM op s (PyFunDef, [Char], PyExp)
callEntryFun [PyStmt]
pre_timing)
                (((Name, Function op) -> Bool)
-> [(Name, Function op)] -> [(Name, Function op)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Function op -> Bool
forall a. FunctionT a -> Bool
Imp.functionEntry (Function op -> Bool)
-> ((Name, Function op) -> Function op)
-> (Name, Function op)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Function op) -> Function op
forall a b. (a, b) -> b
snd) [(Name, Function op)]
funs)
              [PyStmt] -> CompilerM op s [PyStmt]
forall (m :: * -> *) a. Monad m => a -> m a
return ([PyStmt]
parse_options [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
                      PyClassDef -> PyStmt
ClassDef ([Char] -> [PyStmt] -> PyClassDef
Class [Char]
"internal" ([PyStmt] -> PyClassDef) -> [PyStmt] -> PyClassDef
forall a b. (a -> b) -> a -> b
$ (PyFunDef -> PyStmt) -> [PyFunDef] -> [PyStmt]
forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef ([PyFunDef] -> [PyStmt]) -> [PyFunDef] -> [PyStmt]
forall a b. (a -> b) -> a -> b
$
                                PyFunDef
constructor' PyFunDef -> [PyFunDef] -> [PyFunDef]
forall a. a -> [a] -> [a]
: [PyFunDef]
definitions) PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
                      PyStmt
classinst PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
                      (PyFunDef -> PyStmt) -> [PyFunDef] -> [PyStmt]
forall a b. (a -> b) -> [a] -> [b]
map PyFunDef -> PyStmt
FunDef [PyFunDef]
entry_point_defs [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
                      [[Char]] -> [PyExp] -> [PyStmt]
selectEntryPoint [[Char]]
entry_point_names [PyExp]
entry_points)

        parse_options :: [PyStmt]
parse_options =
          PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"runtime_file") PyExp
None PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
          PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"do_warmup_run") (Bool -> PyExp
Bool Bool
False) PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
          PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"num_runs") (Integer -> PyExp
Integer Integer
1) PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
          PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_point") ([Char] -> PyExp
String [Char]
"main") PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
          PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"binary_output") (Bool -> PyExp
Bool Bool
False) PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
          [Option] -> [PyStmt]
generateOptionParser ([Option]
standardOptions [Option] -> [Option] -> [Option]
forall a. [a] -> [a] -> [a]
++ [Option]
options)

        selectEntryPoint :: [[Char]] -> [PyExp] -> [PyStmt]
selectEntryPoint [[Char]]
entry_point_names [PyExp]
entry_points =
          [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_points") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
              [(PyExp, PyExp)] -> PyExp
Dict ([(PyExp, PyExp)] -> PyExp) -> [(PyExp, PyExp)] -> PyExp
forall a b. (a -> b) -> a -> b
$ [PyExp] -> [PyExp] -> [(PyExp, PyExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (([Char] -> PyExp) -> [[Char]] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map [Char] -> PyExp
String [[Char]]
entry_point_names) [PyExp]
entry_points,
            PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"entry_point_fun") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
              [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"entry_points.get" [[Char] -> PyExp
Var [Char]
"entry_point"],
            PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" ([Char] -> PyExp
Var [Char]
"entry_point_fun") PyExp
None)
              [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"sys.exit"
                  [PyExp -> [PyArg] -> PyExp
Call (PyExp -> [Char] -> PyExp
Field
                          ([Char] -> PyExp
String [Char]
"No entry point '{}'.  Select another with --entry point.  Options are:\n{}")
                          [Char]
"format")
                    [PyExp -> PyArg
Arg (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"entry_point",

                     PyExp -> PyArg
Arg (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ PyExp -> [PyArg] -> PyExp
Call (PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
String [Char]
"\n") [Char]
"join")
                     [PyExp -> PyArg
Arg (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"entry_points.keys" []]]]]
              [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"entry_point_fun" []]
          ]

withConstantSubsts :: Imp.Constants op -> CompilerM op s a -> CompilerM op s a
withConstantSubsts :: Constants op -> CompilerM op s a -> CompilerM op s a
withConstantSubsts (Imp.Constants [Param]
ps Code op
_) =
  (CompilerEnv op s -> CompilerEnv op s)
-> CompilerM op s a -> CompilerM op s a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((CompilerEnv op s -> CompilerEnv op s)
 -> CompilerM op s a -> CompilerM op s a)
-> (CompilerEnv op s -> CompilerEnv op s)
-> CompilerM op s a
-> CompilerM op s a
forall a b. (a -> b) -> a -> b
$ \CompilerEnv op s
env -> CompilerEnv op s
env { envVarExp :: Map VName PyExp
envVarExp = (Param -> Map VName PyExp) -> [Param] -> Map VName PyExp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Param -> Map VName PyExp
constExp [Param]
ps }
  where constExp :: Param -> Map VName PyExp
constExp Param
p =
          VName -> PyExp -> Map VName PyExp
forall k a. k -> a -> Map k a
M.singleton (Param -> VName
Imp.paramName Param
p) (PyExp -> Map VName PyExp) -> PyExp -> Map VName PyExp
forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx -> PyExp
Index ([Char] -> PyExp
Var [Char]
"self.constants") (PyIdx -> PyExp) -> PyIdx -> PyExp
forall a b. (a -> b) -> a -> b
$
          PyExp -> PyIdx
IdxExp (PyExp -> PyIdx) -> PyExp -> PyIdx
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
String ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ Param -> VName
Imp.paramName Param
p

compileConstants :: Imp.Constants op -> CompilerM op s ()
compileConstants :: Constants op -> CompilerM op s ()
compileConstants (Imp.Constants [Param]
_ Code op
init_consts) = do
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
atInit (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"self.constants") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [(PyExp, PyExp)] -> PyExp
Dict []
  (PyStmt -> CompilerM op s ()) -> [PyStmt] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
atInit ([PyStmt] -> CompilerM op s ())
-> CompilerM op s [PyStmt] -> CompilerM op s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
init_consts)

compileFunc :: (Name, Imp.Function op) -> CompilerM op s PyFunDef
compileFunc :: (Name, Function op) -> CompilerM op s PyFunDef
compileFunc (Name
fname, Imp.Function Bool
_ [Param]
outputs [Param]
inputs Code op
body [ExternalValue]
_ [ExternalValue]
_) = do
  [PyStmt]
body' <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
body
  let inputs' :: [[Char]]
inputs' = (Param -> [Char]) -> [Param] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> [Char]
compileName (VName -> [Char]) -> (Param -> VName) -> Param -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName) [Param]
inputs
  let ret :: PyStmt
ret = PyExp -> PyStmt
Return (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
tupleOrSingle ([PyExp] -> PyExp) -> [PyExp] -> PyExp
forall a b. (a -> b) -> a -> b
$ [Param] -> [PyExp]
compileOutput [Param]
outputs
  PyFunDef -> CompilerM op s PyFunDef
forall (m :: * -> *) a. Monad m => a -> m a
return (PyFunDef -> CompilerM op s PyFunDef)
-> PyFunDef -> CompilerM op s PyFunDef
forall a b. (a -> b) -> a -> b
$ [Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def ([Char] -> [Char]
futharkFun ([Char] -> [Char]) -> (Name -> [Char]) -> Name -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Char]
nameToString (Name -> [Char]) -> Name -> [Char]
forall a b. (a -> b) -> a -> b
$ Name
fname) ([Char]
"self" [Char] -> [[Char]] -> [[Char]]
forall a. a -> [a] -> [a]
: [[Char]]
inputs') ([PyStmt] -> PyFunDef) -> [PyStmt] -> PyFunDef
forall a b. (a -> b) -> a -> b
$
    [PyStmt]
body'[PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++[PyStmt
ret]

tupleOrSingle :: [PyExp] -> PyExp
tupleOrSingle :: [PyExp] -> PyExp
tupleOrSingle [PyExp
e] = PyExp
e
tupleOrSingle [PyExp]
es = [PyExp] -> PyExp
Tuple [PyExp]
es

-- | A 'Call' where the function is a variable and every argument is a
-- simple 'Arg'.
simpleCall :: String -> [PyExp] -> PyExp
simpleCall :: [Char] -> [PyExp] -> PyExp
simpleCall [Char]
fname = PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
fname) ([PyArg] -> PyExp) -> ([PyExp] -> [PyArg]) -> [PyExp] -> PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PyExp -> PyArg) -> [PyExp] -> [PyArg]
forall a b. (a -> b) -> [a] -> [b]
map PyExp -> PyArg
Arg

compileName :: VName -> String
compileName :: VName -> [Char]
compileName = [Char] -> [Char]
zEncodeString ([Char] -> [Char]) -> (VName -> [Char]) -> VName -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty

compileDim :: Imp.DimSize -> PyExp
compileDim :: DimSize -> PyExp
compileDim (Imp.Constant PrimValue
v) = PrimValue -> PyExp
compilePrimValue PrimValue
v
compileDim (Imp.Var VName
v) = [Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
compileName VName
v

unpackDim :: PyExp -> Imp.DimSize -> Int32 -> CompilerM op s ()
unpackDim :: PyExp -> DimSize -> Int32 -> CompilerM op s ()
unpackDim PyExp
arr_name (Imp.Constant PrimValue
c) Int32
i = do
  let shape_name :: PyExp
shape_name = PyExp -> [Char] -> PyExp
Field PyExp
arr_name [Char]
"shape"
  let constant_c :: PyExp
constant_c = PrimValue -> PyExp
compilePrimValue PrimValue
c
  let constant_i :: PyExp
constant_i = Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger Int32
i
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assert ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" PyExp
constant_c (PyExp -> PyIdx -> PyExp
Index PyExp
shape_name (PyIdx -> PyExp) -> PyIdx -> PyExp
forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx
IdxExp PyExp
constant_i)) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
    [Char] -> PyExp
String [Char]
"constant dimension wrong"
unpackDim PyExp
arr_name (Imp.Var VName
var) Int32
i = do
  let shape_name :: PyExp
shape_name = PyExp -> [Char] -> PyExp
Field PyExp
arr_name [Char]
"shape"
      src :: PyExp
src = PyExp -> PyIdx -> PyExp
Index PyExp
shape_name (PyIdx -> PyExp) -> PyIdx -> PyExp
forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx
IdxExp (PyExp -> PyIdx) -> PyExp -> PyIdx
forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger Int32
i
  PyExp
var' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
var
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
var' (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int32" [PyExp
src]

entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp
entryPointOutput :: ExternalValue -> CompilerM op s PyExp
entryPointOutput (Imp.OpaqueValue [Char]
desc [ValueDesc]
vs) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"opaque" ([PyExp] -> PyExp) -> ([PyExp] -> [PyExp]) -> [PyExp] -> PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char] -> PyExp
String ([Char] -> [Char]
forall a. Pretty a => a -> [Char]
pretty [Char]
desc)PyExp -> [PyExp] -> [PyExp]
forall a. a -> [a] -> [a]
:) ([PyExp] -> PyExp)
-> CompilerM op s [PyExp] -> CompilerM op s PyExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  (ValueDesc -> CompilerM op s PyExp)
-> [ValueDesc] -> CompilerM op s [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ExternalValue -> CompilerM op s PyExp
forall op s. ExternalValue -> CompilerM op s PyExp
entryPointOutput (ExternalValue -> CompilerM op s PyExp)
-> (ValueDesc -> ExternalValue)
-> ValueDesc
-> CompilerM op s PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueDesc -> ExternalValue
Imp.TransparentValue) [ValueDesc]
vs
entryPointOutput (Imp.TransparentValue (Imp.ScalarValue PrimType
bt Signedness
ept VName
name)) = do
  PyExp
name' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
  PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
tf [PyExp
name']
  where tf :: [Char]
tf = PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
bt Signedness
ept
entryPointOutput (Imp.TransparentValue (Imp.ArrayValue VName
mem (Imp.Space [Char]
sid) PrimType
bt Signedness
ept [DimSize]
dims)) = do
  EntryOutput op s
pack_output <- (CompilerEnv op s -> EntryOutput op s)
-> CompilerM op s (EntryOutput op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> EntryOutput op s
forall op s. CompilerEnv op s -> EntryOutput op s
envEntryOutput
  EntryOutput op s
pack_output VName
mem [Char]
sid PrimType
bt Signedness
ept [DimSize]
dims
entryPointOutput (Imp.TransparentValue (Imp.ArrayValue VName
mem Space
_ PrimType
bt Signedness
ept [DimSize]
dims)) = do
  PyExp
mem' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
mem
  let cast :: PyExp
cast = PyExp -> [Char] -> PyExp
Cast PyExp
mem' (PrimType -> Signedness -> [Char]
compilePrimTypeExt PrimType
bt Signedness
ept)
  PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"createArray" [PyExp
cast, [PyExp] -> PyExp
Tuple ([PyExp] -> PyExp) -> [PyExp] -> PyExp
forall a b. (a -> b) -> a -> b
$ (DimSize -> PyExp) -> [DimSize] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map DimSize -> PyExp
compileDim [DimSize]
dims]

badInput :: Int -> PyExp -> String -> PyStmt
badInput :: Int -> PyExp -> [Char] -> PyStmt
badInput Int
i PyExp
e [Char]
t =
  PyExp -> PyStmt
Raise (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"TypeError"
  [PyExp -> [PyArg] -> PyExp
Call (PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
String [Char]
err_msg) [Char]
"format")
   [PyExp -> PyArg
Arg ([Char] -> PyExp
String [Char]
t), PyExp -> PyArg
Arg (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"type" [PyExp
e], PyExp -> PyArg
Arg PyExp
e]]
  where err_msg :: [Char]
err_msg = [[Char]] -> [Char]
unlines [ [Char]
"Argument #" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" has invalid value"
                          , [Char]
"Futhark type: {}"
                          , [Char]
"Argument has Python type {} and value: {}"]


entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput :: (Int, ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput (Int
i, Imp.OpaqueValue [Char]
desc [ValueDesc]
vs, PyExp
e) = do
  let type_is_ok :: PyExp
type_is_ok = [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"and" ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"isinstance" [PyExp
e, [Char] -> PyExp
Var [Char]
"opaque"])
                               ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"desc") ([Char] -> PyExp
String [Char]
desc))
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp -> PyExp
UnOp [Char]
"not" PyExp
type_is_ok) [Int -> PyExp -> [Char] -> PyStmt
badInput Int
i PyExp
e [Char]
desc] []
  ((Int, ExternalValue, PyExp) -> CompilerM op s ())
-> [(Int, ExternalValue, PyExp)] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int, ExternalValue, PyExp) -> CompilerM op s ()
forall op s. (Int, ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput ([(Int, ExternalValue, PyExp)] -> CompilerM op s ())
-> [(Int, ExternalValue, PyExp)] -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [Int]
-> [ExternalValue] -> [PyExp] -> [(Int, ExternalValue, PyExp)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Int -> [Int]
forall a. a -> [a]
repeat Int
i) ((ValueDesc -> ExternalValue) -> [ValueDesc] -> [ExternalValue]
forall a b. (a -> b) -> [a] -> [b]
map ValueDesc -> ExternalValue
Imp.TransparentValue [ValueDesc]
vs) ([PyExp] -> [(Int, ExternalValue, PyExp)])
-> [PyExp] -> [(Int, ExternalValue, PyExp)]
forall a b. (a -> b) -> a -> b
$
    (Integer -> PyExp) -> [Integer] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map (PyExp -> PyIdx -> PyExp
Index (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"data") (PyIdx -> PyExp) -> (Integer -> PyIdx) -> Integer -> PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PyExp -> PyIdx
IdxExp (PyExp -> PyIdx) -> (Integer -> PyExp) -> Integer -> PyIdx
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> PyExp
Integer) [Integer
0..]

entryPointInput (Int
i, Imp.TransparentValue (Imp.ScalarValue PrimType
bt Signedness
s VName
name), PyExp
e) = do
  PyExp
vname' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
  let -- HACK: A Numpy int64 will signal an OverflowError if we pass
      -- it a number bigger than 2**63.  This does not happen if we
      -- pass e.g. int8 a number bigger than 2**7.  As a workaround,
      -- we first go through the corresponding ctypes type, which does
      -- not have this problem.
      ctobject :: [Char]
ctobject = PrimType -> [Char]
compilePrimType PrimType
bt
      ctcall :: PyExp
ctcall = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
ctobject [PyExp
e]
      npobject :: [Char]
npobject = PrimType -> [Char]
compilePrimToNp PrimType
bt
      npcall :: PyExp
npcall = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
npobject [PyExp
ctcall]
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [PyStmt] -> [PyExcept] -> PyStmt
Try [PyExp -> PyExp -> PyStmt
Assign PyExp
vname' PyExp
npcall]
    [PyExp -> [PyStmt] -> PyExcept
Catch ([PyExp] -> PyExp
Tuple [[Char] -> PyExp
Var [Char]
"TypeError", [Char] -> PyExp
Var [Char]
"AssertionError"])
     [Int -> PyExp -> [Char] -> PyStmt
badInput Int
i PyExp
e ([Char] -> PyStmt) -> [Char] -> PyStmt
forall a b. (a -> b) -> a -> b
$ Bool -> PrimType -> [Char]
prettySigned (Signedness
sSignedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
==Signedness
Imp.TypeUnsigned) PrimType
bt]]

entryPointInput (Int
i, Imp.TransparentValue (Imp.ArrayValue VName
mem (Imp.Space [Char]
sid) PrimType
bt Signedness
ept [DimSize]
dims), PyExp
e) = do
  EntryInput op s
unpack_input <- (CompilerEnv op s -> EntryInput op s)
-> CompilerM op s (EntryInput op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> EntryInput op s
forall op s. CompilerEnv op s -> EntryInput op s
envEntryInput
  PyExp
mem' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
mem
  [PyStmt]
unpack <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ EntryInput op s
unpack_input PyExp
mem' [Char]
sid PrimType
bt Signedness
ept [DimSize]
dims PyExp
e
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [PyStmt] -> [PyExcept] -> PyStmt
Try [PyStmt]
unpack
    [PyExp -> [PyStmt] -> PyExcept
Catch ([PyExp] -> PyExp
Tuple [[Char] -> PyExp
Var [Char]
"TypeError", [Char] -> PyExp
Var [Char]
"AssertionError"])
     [Int -> PyExp -> [Char] -> PyStmt
badInput Int
i PyExp
e ([Char] -> PyStmt) -> [Char] -> PyStmt
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Char] -> [[Char]]
forall a. Int -> a -> [a]
replicate ([DimSize] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) [Char]
"[]") [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++
     Bool -> PrimType -> [Char]
prettySigned (Signedness
eptSignedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
==Signedness
Imp.TypeUnsigned) PrimType
bt]]

entryPointInput (Int
i, Imp.TransparentValue (Imp.ArrayValue VName
mem Space
_ PrimType
t Signedness
s [DimSize]
dims), PyExp
e) = do
  let type_is_wrong :: PyExp
type_is_wrong =
        [Char] -> PyExp -> PyExp
UnOp [Char]
"not" (PyExp -> PyExp) -> PyExp -> PyExp
forall a b. (a -> b) -> a -> b
$
        [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"and"
        ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"in" ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"type" [PyExp
e]) ([PyExp] -> PyExp
List [[Char] -> PyExp
Var [Char]
"np.ndarray"]))
        ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"dtype") ([Char] -> PyExp
Var (PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
t Signedness
s)))
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If PyExp
type_is_wrong
    [Int -> PyExp -> [Char] -> PyStmt
badInput Int
i PyExp
e ([Char] -> PyStmt) -> [Char] -> PyStmt
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Char] -> [[Char]]
forall a. Int -> a -> [a]
replicate ([DimSize] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) [Char]
"[]") [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++
     Bool -> PrimType -> [Char]
prettySigned (Signedness
sSignedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
==Signedness
Imp.TypeUnsigned) PrimType
t]
    []

  (DimSize -> Int32 -> CompilerM op s ())
-> [DimSize] -> [Int32] -> CompilerM op s ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (PyExp -> DimSize -> Int32 -> CompilerM op s ()
forall op s. PyExp -> DimSize -> Int32 -> CompilerM op s ()
unpackDim PyExp
e) [DimSize]
dims [Int32
0..]
  PyExp
dest <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
mem
  let unwrap_call :: PyExp
unwrap_call = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"unwrapArray" [PyExp
e]

  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
dest PyExp
unwrap_call

extValueDescName :: Imp.ExternalValue -> String
extValueDescName :: ExternalValue -> [Char]
extValueDescName (Imp.TransparentValue ValueDesc
v) = [Char] -> [Char]
extName ([Char] -> [Char]) -> [Char] -> [Char]
forall a b. (a -> b) -> a -> b
$ ValueDesc -> [Char]
valueDescName ValueDesc
v
extValueDescName (Imp.OpaqueValue [Char]
desc []) = [Char] -> [Char]
extName ([Char] -> [Char]) -> [Char] -> [Char]
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char]
zEncodeString [Char]
desc
extValueDescName (Imp.OpaqueValue [Char]
desc (ValueDesc
v:[ValueDesc]
_)) =
  [Char] -> [Char]
extName ([Char] -> [Char]) -> [Char] -> [Char]
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char]
zEncodeString [Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Pretty a => a -> [Char]
pretty (VName -> Int
baseTag (ValueDesc -> VName
valueDescVName ValueDesc
v))

extName :: String -> String
extName :: [Char] -> [Char]
extName = ([Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
"_ext")

valueDescName :: Imp.ValueDesc -> String
valueDescName :: ValueDesc -> [Char]
valueDescName = VName -> [Char]
compileName (VName -> [Char]) -> (ValueDesc -> VName) -> ValueDesc -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueDesc -> VName
valueDescVName

valueDescVName :: Imp.ValueDesc -> VName
valueDescVName :: ValueDesc -> VName
valueDescVName (Imp.ScalarValue PrimType
_ Signedness
_ VName
vname) = VName
vname
valueDescVName (Imp.ArrayValue VName
vname Space
_ PrimType
_ Signedness
_ [DimSize]
_) = VName
vname

-- Key into the FUTHARK_PRIMTYPES dict.
readTypeEnum :: PrimType -> Imp.Signedness -> String
readTypeEnum :: PrimType -> Signedness -> [Char]
readTypeEnum (IntType IntType
Int8)  Signedness
Imp.TypeUnsigned = [Char]
"u8"
readTypeEnum (IntType IntType
Int16) Signedness
Imp.TypeUnsigned = [Char]
"u16"
readTypeEnum (IntType IntType
Int32) Signedness
Imp.TypeUnsigned = [Char]
"u32"
readTypeEnum (IntType IntType
Int64) Signedness
Imp.TypeUnsigned = [Char]
"u64"
readTypeEnum (IntType IntType
Int8)  Signedness
Imp.TypeDirect   = [Char]
"i8"
readTypeEnum (IntType IntType
Int16) Signedness
Imp.TypeDirect   = [Char]
"i16"
readTypeEnum (IntType IntType
Int32) Signedness
Imp.TypeDirect   = [Char]
"i32"
readTypeEnum (IntType IntType
Int64) Signedness
Imp.TypeDirect   = [Char]
"i64"
readTypeEnum (FloatType FloatType
Float32) Signedness
_ = [Char]
"f32"
readTypeEnum (FloatType FloatType
Float64) Signedness
_ = [Char]
"f64"
readTypeEnum PrimType
Imp.Bool Signedness
_ = [Char]
"bool"
readTypeEnum PrimType
Cert Signedness
_ = [Char] -> [Char]
forall a. HasCallStack => [Char] -> a
error [Char]
"readTypeEnum: cert"

readInput :: Imp.ExternalValue -> PyStmt
readInput :: ExternalValue -> PyStmt
readInput (Imp.OpaqueValue [Char]
desc [ValueDesc]
_) =
  PyExp -> PyStmt
Raise (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"Exception"
  [[Char] -> PyExp
String ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot read argument of type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"."]

readInput decl :: ExternalValue
decl@(Imp.TransparentValue (Imp.ScalarValue PrimType
bt Signedness
ept VName
_)) =
  let type_name :: [Char]
type_name = PrimType -> Signedness -> [Char]
readTypeEnum PrimType
bt Signedness
ept
  in PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ ExternalValue -> [Char]
extValueDescName ExternalValue
decl) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"read_value" [[Char] -> PyExp
String [Char]
type_name]

readInput decl :: ExternalValue
decl@(Imp.TransparentValue (Imp.ArrayValue VName
_ Space
_ PrimType
bt Signedness
ept [DimSize]
dims)) =
  let type_name :: [Char]
type_name = PrimType -> Signedness -> [Char]
readTypeEnum PrimType
bt Signedness
ept
  in PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ ExternalValue -> [Char]
extValueDescName ExternalValue
decl) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"read_value"
     [[Char] -> PyExp
String ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Char] -> [[Char]]
forall a. Int -> a -> [a]
replicate ([DimSize] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) [Char]
"[]") [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
type_name]

printValue :: [(Imp.ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
printValue :: [(ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
printValue = ([[PyStmt]] -> [PyStmt])
-> CompilerM op s [[PyStmt]] -> CompilerM op s [PyStmt]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[PyStmt]] -> [PyStmt]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (CompilerM op s [[PyStmt]] -> CompilerM op s [PyStmt])
-> ([(ExternalValue, PyExp)] -> CompilerM op s [[PyStmt]])
-> [(ExternalValue, PyExp)]
-> CompilerM op s [PyStmt]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ExternalValue, PyExp) -> CompilerM op s [PyStmt])
-> [(ExternalValue, PyExp)] -> CompilerM op s [[PyStmt]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((ExternalValue -> PyExp -> CompilerM op s [PyStmt])
-> (ExternalValue, PyExp) -> CompilerM op s [PyStmt]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ExternalValue -> PyExp -> CompilerM op s [PyStmt]
forall (m :: * -> *).
Monad m =>
ExternalValue -> PyExp -> m [PyStmt]
printValue')
  -- We copy non-host arrays to the host before printing.  This is
  -- done in a hacky way - we assume the value has a .get()-method
  -- that returns an equivalent Numpy array.  This works for PyOpenCL,
  -- but we will probably need yet another plugin mechanism here in
  -- the future.
  where printValue' :: ExternalValue -> PyExp -> m [PyStmt]
printValue' (Imp.OpaqueValue [Char]
desc [ValueDesc]
_) PyExp
_ =
          [PyStmt] -> m [PyStmt]
forall (m :: * -> *) a. Monad m => a -> m a
return [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"sys.stdout.write"
                  [[Char] -> PyExp
String ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ [Char]
"#<opaque " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
desc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
">"]]
        printValue' (Imp.TransparentValue (Imp.ArrayValue VName
mem (Space [Char]
_) PrimType
bt Signedness
ept [DimSize]
shape)) PyExp
e =
          ExternalValue -> PyExp -> m [PyStmt]
printValue' (ValueDesc -> ExternalValue
Imp.TransparentValue (VName -> Space -> PrimType -> Signedness -> [DimSize] -> ValueDesc
Imp.ArrayValue VName
mem Space
DefaultSpace PrimType
bt Signedness
ept [DimSize]
shape)) (PyExp -> m [PyStmt]) -> PyExp -> m [PyStmt]
forall a b. (a -> b) -> a -> b
$
          [Char] -> [PyExp] -> PyExp
simpleCall (PyExp -> [Char]
forall a. Pretty a => a -> [Char]
pretty PyExp
e [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
".get") []
        printValue' (Imp.TransparentValue ValueDesc
_) PyExp
e =
          [PyStmt] -> m [PyStmt]
forall (m :: * -> *) a. Monad m => a -> m a
return [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
"write_value")
                   [PyExp -> PyArg
Arg PyExp
e,
                    [Char] -> PyExp -> PyArg
ArgKeyword [Char]
"binary" ([Char] -> PyExp
Var [Char]
"binary_output")],
                  PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"sys.stdout.write" [[Char] -> PyExp
String [Char]
"\n"]]

prepareEntry :: (Name, Imp.Function op) -> CompilerM op s
                (String, [String], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
                 [(Imp.ExternalValue, PyExp)], [PyStmt])
prepareEntry :: (Name, Function op)
-> CompilerM
     op
     s
     ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
      [(ExternalValue, PyExp)], [PyStmt])
prepareEntry (Name
fname, Imp.Function Bool
_ [Param]
outputs [Param]
inputs Code op
_ [ExternalValue]
results [ExternalValue]
args) = do
  let output_paramNames :: [[Char]]
output_paramNames = (Param -> [Char]) -> [Param] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> [Char]
compileName (VName -> [Char]) -> (Param -> VName) -> Param -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName) [Param]
outputs
      funTuple :: PyExp
funTuple = [PyExp] -> PyExp
tupleOrSingle ([PyExp] -> PyExp) -> [PyExp] -> PyExp
forall a b. (a -> b) -> a -> b
$ ([Char] -> PyExp) -> [[Char]] -> [PyExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Char] -> PyExp
Var [[Char]]
output_paramNames

  ([Maybe [Char]]
argexps_mem_copies, [PyStmt]
prepare_run) <- CompilerM op s [Maybe [Char]]
-> CompilerM op s ([Maybe [Char]], [PyStmt])
forall op s a. CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' (CompilerM op s [Maybe [Char]]
 -> CompilerM op s ([Maybe [Char]], [PyStmt]))
-> CompilerM op s [Maybe [Char]]
-> CompilerM op s ([Maybe [Char]], [PyStmt])
forall a b. (a -> b) -> a -> b
$ [Param]
-> (Param -> CompilerM op s (Maybe [Char]))
-> CompilerM op s [Maybe [Char]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param]
inputs ((Param -> CompilerM op s (Maybe [Char]))
 -> CompilerM op s [Maybe [Char]])
-> (Param -> CompilerM op s (Maybe [Char]))
-> CompilerM op s [Maybe [Char]]
forall a b. (a -> b) -> a -> b
$ \case
    Imp.MemParam VName
name Space
space -> do
      -- A program might write to its input parameters, so create a new memory
      -- block and copy the source there.  This way the program can be run more
      -- than once.
      VName
name' <- [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> CompilerM op s VName) -> [Char] -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString VName
name [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_copy"
      Copy op s
copy <- (CompilerEnv op s -> Copy op s) -> CompilerM op s (Copy op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> Copy op s
forall op s. CompilerEnv op s -> Copy op s
envCopy
      Allocate op s
allocate <- (CompilerEnv op s -> Allocate op s)
-> CompilerM op s (Allocate op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> Allocate op s
forall op s. CompilerEnv op s -> Allocate op s
envAllocate
      let size :: PyExp
size = [Char] -> PyExp
Var ([Char] -> [Char]
extName (VName -> [Char]
compileName VName
name) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
".nbytes") -- FIXME
          dest :: VName
dest = VName
name'
          src :: VName
src = VName
name
          offset :: PyExp
offset = Integer -> PyExp
Integer Integer
0
      case Space
space of
        Space [Char]
sid ->
          Allocate op s
allocate ([Char] -> PyExp
Var (VName -> [Char]
compileName VName
name')) PyExp
size [Char]
sid
        Space
_ ->
          PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var (VName -> [Char]
compileName VName
name'))
                       ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"allocateMem" [PyExp
size]) -- FIXME
      PyExp
dest' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
      PyExp
src' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
      Copy op s
copy PyExp
dest' PyExp
offset Space
space PyExp
src' PyExp
offset Space
space PyExp
size (IntType -> PrimType
IntType IntType
Int32) -- FIXME
      Maybe [Char] -> CompilerM op s (Maybe [Char])
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe [Char] -> CompilerM op s (Maybe [Char]))
-> Maybe [Char] -> CompilerM op s (Maybe [Char])
forall a b. (a -> b) -> a -> b
$ [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just ([Char] -> Maybe [Char]) -> [Char] -> Maybe [Char]
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
compileName VName
name'
    Param
_ -> Maybe [Char] -> CompilerM op s (Maybe [Char])
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe [Char]
forall a. Maybe a
Nothing

  [PyStmt]
prepareIn <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ ((Int, ExternalValue, PyExp) -> CompilerM op s ())
-> [(Int, ExternalValue, PyExp)] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int, ExternalValue, PyExp) -> CompilerM op s ()
forall op s. (Int, ExternalValue, PyExp) -> CompilerM op s ()
entryPointInput ([(Int, ExternalValue, PyExp)] -> CompilerM op s ())
-> [(Int, ExternalValue, PyExp)] -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [Int]
-> [ExternalValue] -> [PyExp] -> [(Int, ExternalValue, PyExp)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Int
0..] [ExternalValue]
args ([PyExp] -> [(Int, ExternalValue, PyExp)])
-> [PyExp] -> [(Int, ExternalValue, PyExp)]
forall a b. (a -> b) -> a -> b
$
               (ExternalValue -> PyExp) -> [ExternalValue] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> PyExp
Var ([Char] -> PyExp)
-> (ExternalValue -> [Char]) -> ExternalValue -> PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExternalValue -> [Char]
extValueDescName) [ExternalValue]
args
  ([PyExp]
res, [PyStmt]
prepareOut) <- CompilerM op s [PyExp] -> CompilerM op s ([PyExp], [PyStmt])
forall op s a. CompilerM op s a -> CompilerM op s (a, [PyStmt])
collect' (CompilerM op s [PyExp] -> CompilerM op s ([PyExp], [PyStmt]))
-> CompilerM op s [PyExp] -> CompilerM op s ([PyExp], [PyStmt])
forall a b. (a -> b) -> a -> b
$ (ExternalValue -> CompilerM op s PyExp)
-> [ExternalValue] -> CompilerM op s [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExternalValue -> CompilerM op s PyExp
forall op s. ExternalValue -> CompilerM op s PyExp
entryPointOutput [ExternalValue]
results

  let argexps_lib :: [[Char]]
argexps_lib = (Param -> [Char]) -> [Param] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> [Char]
compileName (VName -> [Char]) -> (Param -> VName) -> Param -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param -> VName
Imp.paramName) [Param]
inputs
      argexps_bin :: [[Char]]
argexps_bin = ([Char] -> Maybe [Char] -> [Char])
-> [[Char]] -> [Maybe [Char]] -> [[Char]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [Char] -> Maybe [Char] -> [Char]
forall a. a -> Maybe a -> a
fromMaybe [[Char]]
argexps_lib [Maybe [Char]]
argexps_mem_copies
      fname' :: [Char]
fname' = [Char]
"self." [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
futharkFun (Name -> [Char]
nameToString Name
fname)
      call_lib :: [PyStmt]
call_lib = [PyExp -> PyExp -> PyStmt
Assign PyExp
funTuple (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
fname' (([Char] -> PyExp) -> [[Char]] -> [PyExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Char] -> PyExp
Var [[Char]]
argexps_lib)]
      call_bin :: [PyStmt]
call_bin = [PyExp -> PyExp -> PyStmt
Assign PyExp
funTuple (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
fname' (([Char] -> PyExp) -> [[Char]] -> [PyExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Char] -> PyExp
Var [[Char]]
argexps_bin)]

  ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
 [(ExternalValue, PyExp)], [PyStmt])
-> CompilerM
     op
     s
     ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
      [(ExternalValue, PyExp)], [PyStmt])
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> [Char]
nameToString Name
fname, (ExternalValue -> [Char]) -> [ExternalValue] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ExternalValue -> [Char]
extValueDescName [ExternalValue]
args,
          [PyStmt]
prepareIn, [PyStmt]
call_lib, [PyStmt]
call_bin, [PyStmt]
prepareOut,
          [ExternalValue] -> [PyExp] -> [(ExternalValue, PyExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [ExternalValue]
results [PyExp]
res, [PyStmt]
prepare_run)

copyMemoryDefaultSpace :: PyExp -> PyExp -> PyExp -> PyExp -> PyExp ->
                          CompilerM op s ()
copyMemoryDefaultSpace :: PyExp -> PyExp -> PyExp -> PyExp -> PyExp -> CompilerM op s ()
copyMemoryDefaultSpace PyExp
destmem PyExp
destidx PyExp
srcmem PyExp
srcidx PyExp
nbytes = do
  let offset_call1 :: PyExp
offset_call1 = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"addressOffset"
                     [PyExp
destmem, PyExp
destidx, [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  let offset_call2 :: PyExp
offset_call2 = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"addressOffset"
                     [PyExp
srcmem, PyExp
srcidx, [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"ct.memmove" [PyExp
offset_call1, PyExp
offset_call2, PyExp
nbytes]

compileEntryFun :: (Name, Imp.Function op)
                -> CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun :: (Name, Function op) -> CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun (Name, Function op)
entry = do
  ([Char]
fname', [[Char]]
params, [PyStmt]
prepareIn, [PyStmt]
body_lib, [PyStmt]
_, [PyStmt]
prepareOut, [(ExternalValue, PyExp)]
res, [PyStmt]
_) <- (Name, Function op)
-> CompilerM
     op
     s
     ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
      [(ExternalValue, PyExp)], [PyStmt])
forall op s.
(Name, Function op)
-> CompilerM
     op
     s
     ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
      [(ExternalValue, PyExp)], [PyStmt])
prepareEntry (Name, Function op)
entry
  let ret :: PyStmt
ret = PyExp -> PyStmt
Return (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
tupleOrSingle ([PyExp] -> PyExp) -> [PyExp] -> PyExp
forall a b. (a -> b) -> a -> b
$ ((ExternalValue, PyExp) -> PyExp)
-> [(ExternalValue, PyExp)] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map (ExternalValue, PyExp) -> PyExp
forall a b. (a, b) -> b
snd [(ExternalValue, PyExp)]
res
      ([[Char]]
pts, [[Char]]
rts) = Function op -> ([[Char]], [[Char]])
forall op. Function op -> ([[Char]], [[Char]])
entryTypes (Function op -> ([[Char]], [[Char]]))
-> Function op -> ([[Char]], [[Char]])
forall a b. (a -> b) -> a -> b
$ (Name, Function op) -> Function op
forall a b. (a, b) -> b
snd (Name, Function op)
entry
  (PyFunDef, (PyExp, PyExp))
-> CompilerM op s (PyFunDef, (PyExp, PyExp))
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def [Char]
fname' ([Char]
"self" [Char] -> [[Char]] -> [[Char]]
forall a. a -> [a] -> [a]
: [[Char]]
params) ([PyStmt] -> PyFunDef) -> [PyStmt] -> PyFunDef
forall a b. (a -> b) -> a -> b
$
           [PyStmt]
prepareIn [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
body_lib [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
prepareOut [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt
ret],
          ([Char] -> PyExp
String [Char]
fname', [PyExp] -> PyExp
Tuple [[PyExp] -> PyExp
List (([Char] -> PyExp) -> [[Char]] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map [Char] -> PyExp
String [[Char]]
pts), [PyExp] -> PyExp
List (([Char] -> PyExp) -> [[Char]] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map [Char] -> PyExp
String [[Char]]
rts)]))

entryTypes :: Imp.Function op -> ([String], [String])
entryTypes :: Function op -> ([[Char]], [[Char]])
entryTypes Function op
func = ((ExternalValue -> [Char]) -> [ExternalValue] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ExternalValue -> [Char]
desc ([ExternalValue] -> [[Char]]) -> [ExternalValue] -> [[Char]]
forall a b. (a -> b) -> a -> b
$ Function op -> [ExternalValue]
forall a. FunctionT a -> [ExternalValue]
Imp.functionArgs Function op
func,
                   (ExternalValue -> [Char]) -> [ExternalValue] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ExternalValue -> [Char]
desc ([ExternalValue] -> [[Char]]) -> [ExternalValue] -> [[Char]]
forall a b. (a -> b) -> a -> b
$ Function op -> [ExternalValue]
forall a. FunctionT a -> [ExternalValue]
Imp.functionResult Function op
func)
  where desc :: ExternalValue -> [Char]
desc (Imp.OpaqueValue [Char]
d [ValueDesc]
_) = [Char]
d
        desc (Imp.TransparentValue (Imp.ScalarValue PrimType
pt Signedness
s VName
_)) = PrimType -> Signedness -> [Char]
readTypeEnum PrimType
pt Signedness
s
        desc (Imp.TransparentValue (Imp.ArrayValue VName
_ Space
_ PrimType
pt Signedness
s [DimSize]
dims)) =
          [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Char] -> [[Char]]
forall a. Int -> a -> [a]
replicate ([DimSize] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims) [Char]
"[]") [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PrimType -> Signedness -> [Char]
readTypeEnum PrimType
pt Signedness
s

callEntryFun :: [PyStmt] -> (Name, Imp.Function op)
             -> CompilerM op s (PyFunDef, String, PyExp)
callEntryFun :: [PyStmt]
-> (Name, Function op) -> CompilerM op s (PyFunDef, [Char], PyExp)
callEntryFun [PyStmt]
pre_timing entry :: (Name, Function op)
entry@(Name
fname, Imp.Function Bool
_ [Param]
_ [Param]
_ Code op
_ [ExternalValue]
_ [ExternalValue]
decl_args) = do
  ([Char]
_, [[Char]]
_, [PyStmt]
prepare_in, [PyStmt]
_, [PyStmt]
body_bin, [PyStmt]
_, [(ExternalValue, PyExp)]
res, [PyStmt]
prepare_run) <- (Name, Function op)
-> CompilerM
     op
     s
     ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
      [(ExternalValue, PyExp)], [PyStmt])
forall op s.
(Name, Function op)
-> CompilerM
     op
     s
     ([Char], [[Char]], [PyStmt], [PyStmt], [PyStmt], [PyStmt],
      [(ExternalValue, PyExp)], [PyStmt])
prepareEntry (Name, Function op)
entry

  let str_input :: [PyStmt]
str_input = (ExternalValue -> PyStmt) -> [ExternalValue] -> [PyStmt]
forall a b. (a -> b) -> [a] -> [b]
map ExternalValue -> PyStmt
readInput [ExternalValue]
decl_args
      end_of_input :: [PyStmt]
end_of_input = [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"end_of_input" [[Char] -> PyExp
String ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ Name -> [Char]
forall a. Pretty a => a -> [Char]
pretty Name
fname]]

      exitcall :: [PyStmt]
exitcall = [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"sys.exit" [PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
String [Char]
"Assertion.{} failed") [Char]
"format(e)"]]
      except' :: PyExcept
except' = PyExp -> [PyStmt] -> PyExcept
Catch ([Char] -> PyExp
Var [Char]
"AssertionError") [PyStmt]
exitcall
      do_run :: [PyStmt]
do_run = [PyStmt]
body_bin [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
pre_timing
      ([PyStmt]
do_run_with_timing, PyStmt
close_runtime_file) = [PyStmt] -> ([PyStmt], PyStmt)
addTiming [PyStmt]
do_run

      -- We ignore overflow errors and the like for executable entry
      -- points.  These are (somewhat) well-defined in Futhark.
      ignore :: [Char] -> PyArg
ignore [Char]
s = [Char] -> PyExp -> PyArg
ArgKeyword [Char]
s (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
String [Char]
"ignore"
      errstate :: PyExp
errstate = PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
"np.errstate") ([PyArg] -> PyExp) -> [PyArg] -> PyExp
forall a b. (a -> b) -> a -> b
$ ([Char] -> PyArg) -> [[Char]] -> [PyArg]
forall a b. (a -> b) -> [a] -> [b]
map [Char] -> PyArg
ignore [[Char]
"divide", [Char]
"over", [Char]
"under", [Char]
"invalid"]

      do_warmup_run :: PyStmt
do_warmup_run =
        PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"do_warmup_run") ([PyStmt]
prepare_run [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
do_run) []

      do_num_runs :: PyStmt
do_num_runs =
        [Char] -> PyExp -> [PyStmt] -> PyStmt
For [Char]
"i" ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"range" [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
"int" [[Char] -> PyExp
Var [Char]
"num_runs"]])
        ([PyStmt]
prepare_run [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
do_run_with_timing)

  [PyStmt]
str_output <- [(ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
forall op s. [(ExternalValue, PyExp)] -> CompilerM op s [PyStmt]
printValue [(ExternalValue, PyExp)]
res

  let fname' :: [Char]
fname' = [Char]
"entry_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
nameToString Name
fname

  (PyFunDef, [Char], PyExp)
-> CompilerM op s (PyFunDef, [Char], PyExp)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> [[Char]] -> [PyStmt] -> PyFunDef
Def [Char]
fname' [] ([PyStmt] -> PyFunDef) -> [PyStmt] -> PyFunDef
forall a b. (a -> b) -> a -> b
$
           [PyStmt]
str_input [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
end_of_input [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
prepare_in [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
           [[PyStmt] -> [PyExcept] -> PyStmt
Try [PyExp -> [PyStmt] -> PyStmt
With PyExp
errstate [PyStmt
do_warmup_run, PyStmt
do_num_runs]] [PyExcept
except']] [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
           [PyStmt
close_runtime_file] [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
           [PyStmt]
str_output,

          Name -> [Char]
nameToString Name
fname,

          [Char] -> PyExp
Var [Char]
fname')

addTiming :: [PyStmt] -> ([PyStmt], PyStmt)
addTiming :: [PyStmt] -> ([PyStmt], PyStmt)
addTiming [PyStmt]
statements =
  ([ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"time_start") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" [] ] [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
   [PyStmt]
statements [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++
   [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"time_end") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" []
   , PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"runtime_file") [PyStmt]
print_runtime [] ],

   PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If ([Char] -> PyExp
Var [Char]
"runtime_file") [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.close" []] [])
  where print_runtime :: [PyStmt]
print_runtime =
          [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.write"
           [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
"str"
            [[Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"-"
             (PyExp -> PyExp
toMicroseconds ([Char] -> PyExp
Var [Char]
"time_end"))
             (PyExp -> PyExp
toMicroseconds ([Char] -> PyExp
Var [Char]
"time_start"))]],
           PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.write" [[Char] -> PyExp
String [Char]
"\n"]]
        toMicroseconds :: PyExp -> PyExp
toMicroseconds PyExp
x =
          [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"int" [[Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"*" PyExp
x (PyExp -> PyExp) -> PyExp -> PyExp
forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer Integer
1000000]

compileUnOp :: Imp.UnOp -> String
compileUnOp :: UnOp -> [Char]
compileUnOp UnOp
op =
  case UnOp
op of
    UnOp
Not -> [Char]
"not"
    Complement{} -> [Char]
"~"
    Abs{} -> [Char]
"abs"
    FAbs{} -> [Char]
"abs"
    SSignum{} -> [Char]
"ssignum"
    USignum{} -> [Char]
"usignum"

compileBinOpLike :: Monad m =>
                    Imp.Exp -> Imp.Exp
                 -> CompilerM op s (PyExp, PyExp, String -> m PyExp)
compileBinOpLike :: Exp -> Exp -> CompilerM op s (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike Exp
x Exp
y = do
  PyExp
x' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
  PyExp
y' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
y
  let simple :: [Char] -> m PyExp
simple [Char]
s = PyExp -> m PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> m PyExp) -> PyExp -> m PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
s PyExp
x' PyExp
y'
  (PyExp, PyExp, [Char] -> m PyExp)
-> CompilerM op s (PyExp, PyExp, [Char] -> m PyExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp
x', PyExp
y', [Char] -> m PyExp
forall (m :: * -> *). Monad m => [Char] -> m PyExp
simple)

-- | The ctypes type corresponding to a 'PrimType'.
compilePrimType :: PrimType -> String
compilePrimType :: PrimType -> [Char]
compilePrimType PrimType
t =
  case PrimType
t of
    IntType IntType
Int8 -> [Char]
"ct.c_int8"
    IntType IntType
Int16 -> [Char]
"ct.c_int16"
    IntType IntType
Int32 -> [Char]
"ct.c_int32"
    IntType IntType
Int64 -> [Char]
"ct.c_int64"
    FloatType FloatType
Float32 -> [Char]
"ct.c_float"
    FloatType FloatType
Float64 -> [Char]
"ct.c_double"
    PrimType
Imp.Bool -> [Char]
"ct.c_bool"
    PrimType
Cert -> [Char]
"ct.c_bool"

-- | The ctypes type corresponding to a 'PrimType', taking sign into account.
compilePrimTypeExt :: PrimType -> Imp.Signedness -> String
compilePrimTypeExt :: PrimType -> Signedness -> [Char]
compilePrimTypeExt PrimType
t Signedness
ept =
  case (PrimType
t, Signedness
ept) of
    (IntType IntType
Int8, Signedness
Imp.TypeUnsigned) -> [Char]
"ct.c_uint8"
    (IntType IntType
Int16, Signedness
Imp.TypeUnsigned) -> [Char]
"ct.c_uint16"
    (IntType IntType
Int32, Signedness
Imp.TypeUnsigned) -> [Char]
"ct.c_uint32"
    (IntType IntType
Int64, Signedness
Imp.TypeUnsigned) -> [Char]
"ct.c_uint64"
    (IntType IntType
Int8, Signedness
_) -> [Char]
"ct.c_int8"
    (IntType IntType
Int16, Signedness
_) -> [Char]
"ct.c_int16"
    (IntType IntType
Int32, Signedness
_) -> [Char]
"ct.c_int32"
    (IntType IntType
Int64, Signedness
_) -> [Char]
"ct.c_int64"
    (FloatType FloatType
Float32, Signedness
_) -> [Char]
"ct.c_float"
    (FloatType FloatType
Float64, Signedness
_) -> [Char]
"ct.c_double"
    (PrimType
Imp.Bool, Signedness
_) -> [Char]
"ct.c_bool"
    (PrimType
Cert, Signedness
_) -> [Char]
"ct.c_byte"

-- | The Numpy type corresponding to a 'PrimType'.
compilePrimToNp :: Imp.PrimType -> String
compilePrimToNp :: PrimType -> [Char]
compilePrimToNp PrimType
bt =
  case PrimType
bt of
    IntType IntType
Int8 -> [Char]
"np.int8"
    IntType IntType
Int16 -> [Char]
"np.int16"
    IntType IntType
Int32 -> [Char]
"np.int32"
    IntType IntType
Int64 -> [Char]
"np.int64"
    FloatType FloatType
Float32 -> [Char]
"np.float32"
    FloatType FloatType
Float64 -> [Char]
"np.float64"
    PrimType
Imp.Bool -> [Char]
"np.byte"
    PrimType
Cert -> [Char]
"np.byte"

-- | The Numpy type corresponding to a 'PrimType', taking sign into account.
compilePrimToExtNp :: Imp.PrimType -> Imp.Signedness -> String
compilePrimToExtNp :: PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
bt Signedness
ept =
  case (PrimType
bt,Signedness
ept) of
    (IntType IntType
Int8, Signedness
Imp.TypeUnsigned) -> [Char]
"np.uint8"
    (IntType IntType
Int16, Signedness
Imp.TypeUnsigned) -> [Char]
"np.uint16"
    (IntType IntType
Int32, Signedness
Imp.TypeUnsigned) -> [Char]
"np.uint32"
    (IntType IntType
Int64, Signedness
Imp.TypeUnsigned) -> [Char]
"np.uint64"
    (IntType IntType
Int8, Signedness
_) -> [Char]
"np.int8"
    (IntType IntType
Int16, Signedness
_) -> [Char]
"np.int16"
    (IntType IntType
Int32, Signedness
_) -> [Char]
"np.int32"
    (IntType IntType
Int64, Signedness
_) -> [Char]
"np.int64"
    (FloatType FloatType
Float32, Signedness
_) -> [Char]
"np.float32"
    (FloatType FloatType
Float64, Signedness
_) -> [Char]
"np.float64"
    (PrimType
Imp.Bool, Signedness
_) -> [Char]
"np.bool_"
    (PrimType
Cert, Signedness
_) -> [Char]
"np.byte"

compilePrimValue :: Imp.PrimValue -> PyExp
compilePrimValue :: PrimValue -> PyExp
compilePrimValue (IntValue (Int8Value Int8
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int8" [Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int8 -> Integer
forall a. Integral a => a -> Integer
toInteger Int8
v]
compilePrimValue (IntValue (Int16Value Int16
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int16" [Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int16 -> Integer
forall a. Integral a => a -> Integer
toInteger Int16
v]
compilePrimValue (IntValue (Int32Value Int32
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int32" [Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger Int32
v]
compilePrimValue (IntValue (Int64Value Int64
v)) =
  [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int64" [Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int64 -> Integer
forall a. Integral a => a -> Integer
toInteger Int64
v]
compilePrimValue (FloatValue (Float32Value Float
v))
  | Float -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Float
v =
      if Float
v Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
> Float
0 then [Char] -> PyExp
Var [Char]
"np.inf" else [Char] -> PyExp
Var [Char]
"-np.inf"
  | Float -> Bool
forall a. RealFloat a => a -> Bool
isNaN Float
v =
      [Char] -> PyExp
Var [Char]
"np.nan"
  | Bool
otherwise = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.float32" [Double -> PyExp
Float (Double -> PyExp) -> Double -> PyExp
forall a b. (a -> b) -> a -> b
$ Rational -> Double
forall a. Fractional a => Rational -> a
fromRational (Rational -> Double) -> Rational -> Double
forall a b. (a -> b) -> a -> b
$ Float -> Rational
forall a. Real a => a -> Rational
toRational Float
v]
compilePrimValue (FloatValue (Float64Value Double
v))
  | Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
v =
      if Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then [Char] -> PyExp
Var [Char]
"np.inf" else [Char] -> PyExp
Var [Char]
"-np.inf"
  | Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
v =
      [Char] -> PyExp
Var [Char]
"np.nan"
  | Bool
otherwise = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.float64" [Double -> PyExp
Float (Double -> PyExp) -> Double -> PyExp
forall a b. (a -> b) -> a -> b
$ Rational -> Double
forall a. Fractional a => Rational -> a
fromRational (Rational -> Double) -> Rational -> Double
forall a b. (a -> b) -> a -> b
$ Double -> Rational
forall a. Real a => a -> Rational
toRational Double
v]
compilePrimValue (BoolValue Bool
v) = Bool -> PyExp
Bool Bool
v
compilePrimValue PrimValue
Checked = [Char] -> PyExp
Var [Char]
"True"

compileVar :: VName -> CompilerM op s PyExp
compileVar :: VName -> CompilerM op s PyExp
compileVar VName
v =
  (CompilerEnv op s -> PyExp) -> CompilerM op s PyExp
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((CompilerEnv op s -> PyExp) -> CompilerM op s PyExp)
-> (CompilerEnv op s -> PyExp) -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ PyExp -> Maybe PyExp -> PyExp
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
compileName VName
v) (Maybe PyExp -> PyExp)
-> (CompilerEnv op s -> Maybe PyExp) -> CompilerEnv op s -> PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Map VName PyExp -> Maybe PyExp
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName PyExp -> Maybe PyExp)
-> (CompilerEnv op s -> Map VName PyExp)
-> CompilerEnv op s
-> Maybe PyExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompilerEnv op s -> Map VName PyExp
forall op s. CompilerEnv op s -> Map VName PyExp
envVarExp

compileExp :: Imp.Exp -> CompilerM op s PyExp

compileExp :: Exp -> CompilerM op s PyExp
compileExp (Imp.ValueExp PrimValue
v) = PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PyExp
compilePrimValue PrimValue
v

compileExp (Imp.LeafExp (Imp.ScalarVar VName
vname) PrimType
_) =
  VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
vname

compileExp (Imp.LeafExp (Imp.SizeOf PrimType
t) PrimType
_) =
  PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimToNp (PrimType -> [Char]) -> PrimType -> [Char]
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32) [Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ PrimType -> Integer
forall a. Num a => PrimType -> a
primByteSize PrimType
t]

compileExp (Imp.LeafExp (Imp.Index VName
src (Imp.Count Exp
iexp) PrimType
restype (Imp.Space [Char]
space) Volatility
_) PrimType
_) =
  CompilerM op s (CompilerM op s PyExp) -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s PyExp) -> CompilerM op s PyExp)
-> CompilerM op s (CompilerM op s PyExp) -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ (CompilerEnv op s -> ReadScalar op s)
-> CompilerM op s (ReadScalar op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> ReadScalar op s
forall op s. CompilerEnv op s -> ReadScalar op s
envReadScalar
    CompilerM op s (ReadScalar op s)
-> CompilerM op s PyExp
-> CompilerM
     op s (PyExp -> PrimType -> [Char] -> CompilerM op s PyExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
src CompilerM
  op s (PyExp -> PrimType -> [Char] -> CompilerM op s PyExp)
-> CompilerM op s PyExp
-> CompilerM op s (PrimType -> [Char] -> CompilerM op s PyExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
iexp
    CompilerM op s (PrimType -> [Char] -> CompilerM op s PyExp)
-> CompilerM op s PrimType
-> CompilerM op s ([Char] -> CompilerM op s PyExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimType -> CompilerM op s PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
restype CompilerM op s ([Char] -> CompilerM op s PyExp)
-> CompilerM op s [Char] -> CompilerM op s (CompilerM op s PyExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> CompilerM op s [Char]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space

compileExp (Imp.LeafExp (Imp.Index VName
src (Imp.Count Exp
iexp) PrimType
bt Space
_ Volatility
_) PrimType
_) = do
  PyExp
iexp' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
iexp
  let bt' :: [Char]
bt' = PrimType -> [Char]
compilePrimType PrimType
bt
      nptype :: [Char]
nptype = PrimType -> [Char]
compilePrimToNp PrimType
bt
  PyExp
src' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
  PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"indexArray" [PyExp
src', PyExp
iexp', [Char] -> PyExp
Var [Char]
bt', [Char] -> PyExp
Var [Char]
nptype]

compileExp (Imp.BinOpExp BinOp
op Exp
x Exp
y) = do
  (PyExp
x', PyExp
y', [Char] -> CompilerM op s PyExp
simple) <- Exp
-> Exp
-> CompilerM op s (PyExp, PyExp, [Char] -> CompilerM op s PyExp)
forall (m :: * -> *) op s.
Monad m =>
Exp -> Exp -> CompilerM op s (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike Exp
x Exp
y
  case BinOp
op of
    Add{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"+"
    Sub{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"-"
    Mul{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"*"
    FAdd{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"+"
    FSub{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"-"
    FMul{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"*"
    FDiv{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"/"
    FMod{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"%"
    Xor{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"^"
    And{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"&"
    Or{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"|"
    Shl{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"<<"
    LogAnd{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"and"
    LogOr{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"or"
    BinOp
_ -> PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (BinOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty BinOp
op) [PyExp
x', PyExp
y']

compileExp (Imp.ConvOpExp ConvOp
conv Exp
x) = do
  PyExp
x' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
  PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (ConvOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty ConvOp
conv) [PyExp
x']

compileExp (Imp.CmpOpExp CmpOp
cmp Exp
x Exp
y) = do
  (PyExp
x', PyExp
y', [Char] -> CompilerM op s PyExp
simple) <- Exp
-> Exp
-> CompilerM op s (PyExp, PyExp, [Char] -> CompilerM op s PyExp)
forall (m :: * -> *) op s.
Monad m =>
Exp -> Exp -> CompilerM op s (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike Exp
x Exp
y
  case CmpOp
cmp of
    CmpEq{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"=="
    FCmpLt{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"<"
    FCmpLe{} -> [Char] -> CompilerM op s PyExp
simple [Char]
"<="
    CmpOp
CmpLlt -> [Char] -> CompilerM op s PyExp
simple [Char]
"<"
    CmpOp
CmpLle -> [Char] -> CompilerM op s PyExp
simple [Char]
"<="
    CmpOp
_ -> PyExp -> CompilerM op s PyExp
forall (m :: * -> *) a. Monad m => a -> m a
return (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (CmpOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty CmpOp
cmp) [PyExp
x', PyExp
y']

compileExp (Imp.UnOpExp UnOp
op Exp
exp1) =
  [Char] -> PyExp -> PyExp
UnOp (UnOp -> [Char]
compileUnOp UnOp
op) (PyExp -> PyExp) -> CompilerM op s PyExp -> CompilerM op s PyExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
exp1

compileExp (Imp.FunExp [Char]
h [Exp]
args PrimType
_) =
  [Char] -> [PyExp] -> PyExp
simpleCall ([Char] -> [Char]
futharkFun ([Char] -> [Char]
forall a. Pretty a => a -> [Char]
pretty [Char]
h)) ([PyExp] -> PyExp)
-> CompilerM op s [PyExp] -> CompilerM op s PyExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> CompilerM op s PyExp) -> [Exp] -> CompilerM op s [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp [Exp]
args

compileCode :: Imp.Code op -> CompilerM op s ()

compileCode :: Code op -> CompilerM op s ()
compileCode Imp.DebugPrint{} =
  () -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

compileCode (Imp.Op op
op) =
  CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s ()) -> CompilerM op s ())
-> CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ (CompilerEnv op s -> OpCompiler op s)
-> CompilerM op s (OpCompiler op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> OpCompiler op s
forall op s. CompilerEnv op s -> OpCompiler op s
envOpCompiler CompilerM op s (OpCompiler op s)
-> CompilerM op s op -> CompilerM op s (CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> op -> CompilerM op s op
forall (f :: * -> *) a. Applicative f => a -> f a
pure op
op

compileCode (Imp.If Exp
cond Code op
tb Code op
fb) = do
  PyExp
cond' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
cond
  [PyStmt]
tb' <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
tb
  [PyStmt]
fb' <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
fb
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> [PyStmt] -> PyStmt
If PyExp
cond' [PyStmt]
tb' [PyStmt]
fb'

compileCode (Code op
c1 Imp.:>>: Code op
c2) = do
  Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
c1
  Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
c2

compileCode (Imp.While Exp
cond Code op
body) = do
  PyExp
cond' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
cond
  [PyStmt]
body' <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
body
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> [PyStmt] -> PyStmt
While PyExp
cond' [PyStmt]
body'

compileCode (Imp.For VName
i IntType
it Exp
bound Code op
body) = do
  PyExp
bound' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
bound
  let i' :: [Char]
i' = VName -> [Char]
compileName VName
i
  [PyStmt]
body' <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
body
  [Char]
counter <- VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty (VName -> [Char]) -> CompilerM op s VName -> CompilerM op s [Char]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"counter"
  [Char]
one <- VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty (VName -> [Char]) -> CompilerM op s VName -> CompilerM op s [Char]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"one"
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
i') (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimToNp (IntType -> PrimType
IntType IntType
it)) [Integer -> PyExp
Integer Integer
0]
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
one) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall (PrimType -> [Char]
compilePrimToNp (IntType -> PrimType
IntType IntType
it)) [Integer -> PyExp
Integer Integer
1]
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp -> [PyStmt] -> PyStmt
For [Char]
counter ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"range" [PyExp
bound']) ([PyStmt] -> PyStmt) -> [PyStmt] -> PyStmt
forall a b. (a -> b) -> a -> b
$
    [PyStmt]
body' [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [[Char] -> PyExp -> PyExp -> PyStmt
AssignOp [Char]
"+" ([Char] -> PyExp
Var [Char]
i') ([Char] -> PyExp
Var [Char]
one)]

compileCode (Imp.SetScalar VName
name Exp
exp1) =
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ())
-> CompilerM op s PyStmt -> CompilerM op s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign (PyExp -> PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s (PyExp -> PyStmt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name CompilerM op s (PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s PyStmt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
exp1

compileCode Imp.DeclareMem{} = () -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
compileCode (Imp.DeclareScalar VName
v Volatility
_ PrimType
Cert) = do
  PyExp
v' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
v
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
v' (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var [Char]
"True"
compileCode Imp.DeclareScalar{} = () -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

compileCode (Imp.DeclareArray VName
name (Space [Char]
space) PrimType
t ArrayContents
vs) =
  CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s ()) -> CompilerM op s ())
-> CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ (CompilerEnv op s -> StaticArray op s)
-> CompilerM op s (StaticArray op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> StaticArray op s
forall op s. CompilerEnv op s -> StaticArray op s
envStaticArray CompilerM op s (StaticArray op s)
-> CompilerM op s VName
-> CompilerM
     op s ([Char] -> PrimType -> ArrayContents -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
  VName -> CompilerM op s VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
name CompilerM
  op s ([Char] -> PrimType -> ArrayContents -> CompilerM op s ())
-> CompilerM op s [Char]
-> CompilerM op s (PrimType -> ArrayContents -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> CompilerM op s [Char]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space CompilerM op s (PrimType -> ArrayContents -> CompilerM op s ())
-> CompilerM op s PrimType
-> CompilerM op s (ArrayContents -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimType -> CompilerM op s PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t CompilerM op s (ArrayContents -> CompilerM op s ())
-> CompilerM op s ArrayContents
-> CompilerM op s (CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ArrayContents -> CompilerM op s ArrayContents
forall (f :: * -> *) a. Applicative f => a -> f a
pure ArrayContents
vs

compileCode (Imp.DeclareArray VName
name Space
_ PrimType
t ArrayContents
vs) = do
  let arr_name :: [Char]
arr_name = VName -> [Char]
compileName VName
name [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_arr"
  -- It is important to store the Numpy array in a temporary variable
  -- to prevent it from going "out-of-scope" before calling
  -- unwrapArray (which internally uses the .ctype method); see
  -- https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.ctypes.html
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
atInit (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign (PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
Var [Char]
"self") [Char]
arr_name) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ case ArrayContents
vs of
    Imp.ArrayValues [PrimValue]
vs' ->
      PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
"np.array")
      [PyExp -> PyArg
Arg (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List ([PyExp] -> PyExp) -> [PyExp] -> PyExp
forall a b. (a -> b) -> a -> b
$ (PrimValue -> PyExp) -> [PrimValue] -> [PyExp]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> PyExp
compilePrimValue [PrimValue]
vs',
       [Char] -> PyExp -> PyArg
ArgKeyword [Char]
"dtype" (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ PrimType -> [Char]
compilePrimToNp PrimType
t]
    Imp.ArrayZeros Int
n ->
      PyExp -> [PyArg] -> PyExp
Call ([Char] -> PyExp
Var [Char]
"np.zeros")
      [PyExp -> PyArg
Arg (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ Integer -> PyExp
Integer (Integer -> PyExp) -> Integer -> PyExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n,
       [Char] -> PyExp -> PyArg
ArgKeyword [Char]
"dtype" (PyExp -> PyArg) -> PyExp -> PyArg
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ PrimType -> [Char]
compilePrimToNp PrimType
t]
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
atInit (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$
    PyExp -> PyExp -> PyStmt
Assign (PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
Var [Char]
"self") (VName -> [Char]
compileName VName
name)) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
    [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"unwrapArray" [PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
Var [Char]
"self") [Char]
arr_name]
  PyExp
name' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assign PyExp
name' (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ PyExp -> [Char] -> PyExp
Field ([Char] -> PyExp
Var [Char]
"self") (VName -> [Char]
compileName VName
name)

compileCode (Imp.Comment [Char]
s Code op
code) = do
  [PyStmt]
code' <- CompilerM op s () -> CompilerM op s [PyStmt]
forall op s. CompilerM op s () -> CompilerM op s [PyStmt]
collect (CompilerM op s () -> CompilerM op s [PyStmt])
-> CompilerM op s () -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$ Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
code
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyStmt] -> PyStmt
Comment [Char]
s [PyStmt]
code'

compileCode (Imp.Assert Exp
e (Imp.ErrorMsg [ErrorMsgPart Exp]
parts) (SrcLoc
loc,[SrcLoc]
locs)) = do
  PyExp
e' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
  let onPart :: ErrorMsgPart Exp -> CompilerM op s (a, PyExp)
onPart (Imp.ErrorString [Char]
s) = (a, PyExp) -> CompilerM op s (a, PyExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
"%s", [Char] -> PyExp
String [Char]
s)
      onPart (Imp.ErrorInt32 Exp
x) = (a
"%d",) (PyExp -> (a, PyExp))
-> CompilerM op s PyExp -> CompilerM op s (a, PyExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
x
  ([[Char]]
formatstrs, [PyExp]
formatargs) <- [([Char], PyExp)] -> ([[Char]], [PyExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([Char], PyExp)] -> ([[Char]], [PyExp]))
-> CompilerM op s [([Char], PyExp)]
-> CompilerM op s ([[Char]], [PyExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ErrorMsgPart Exp -> CompilerM op s ([Char], PyExp))
-> [ErrorMsgPart Exp] -> CompilerM op s [([Char], PyExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ErrorMsgPart Exp -> CompilerM op s ([Char], PyExp)
forall a op s.
IsString a =>
ErrorMsgPart Exp -> CompilerM op s (a, PyExp)
onPart [ErrorMsgPart Exp]
parts
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyExp -> PyStmt
Assert PyExp
e' ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"%"
                   ([Char] -> PyExp
String ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ [Char]
"Error: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Char]]
formatstrs [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"\n\nBacktrace:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
stacktrace)
                   ([PyExp] -> PyExp
Tuple [PyExp]
formatargs))
  where stacktrace :: [Char]
stacktrace = Int -> [[Char]] -> [Char]
prettyStacktrace Int
0 ([[Char]] -> [Char]) -> [[Char]] -> [Char]
forall a b. (a -> b) -> a -> b
$ (SrcLoc -> [Char]) -> [SrcLoc] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr ([SrcLoc] -> [[Char]]) -> [SrcLoc] -> [[Char]]
forall a b. (a -> b) -> a -> b
$ SrcLoc
locSrcLoc -> [SrcLoc] -> [SrcLoc]
forall a. a -> [a] -> [a]
:[SrcLoc]
locs

compileCode (Imp.Call [VName]
dests Name
fname [Arg]
args) = do
  [PyExp]
args' <- (Arg -> CompilerM op s PyExp) -> [Arg] -> CompilerM op s [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Arg -> CompilerM op s PyExp
forall op s. Arg -> CompilerM op s PyExp
compileArg [Arg]
args
  PyExp
dests' <- [PyExp] -> PyExp
tupleOrSingle ([PyExp] -> PyExp)
-> CompilerM op s [PyExp] -> CompilerM op s PyExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> CompilerM op s PyExp)
-> [VName] -> CompilerM op s [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar [VName]
dests
  let fname' :: [Char]
fname'
        | Name -> Bool
isBuiltInFunction Name
fname = [Char] -> [Char]
futharkFun (Name -> [Char]
forall a. Pretty a => a -> [Char]
pretty  Name
fname)
        | Bool
otherwise               = [Char]
"self." [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
futharkFun (Name -> [Char]
forall a. Pretty a => a -> [Char]
pretty  Name
fname)
      call' :: PyExp
call' = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
fname' [PyExp]
args'
  -- If the function returns nothing (is called only for side
  -- effects), take care not to assign to an empty tuple.
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ if [VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
dests
        then PyExp -> PyStmt
Exp PyExp
call'
        else PyExp -> PyExp -> PyStmt
Assign PyExp
dests' PyExp
call'
  where compileArg :: Arg -> CompilerM op s PyExp
compileArg (Imp.MemArg VName
m) = VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
m
        compileArg (Imp.ExpArg Exp
e) = Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e

compileCode (Imp.SetMem VName
dest VName
src Space
_) =
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ())
-> CompilerM op s PyStmt -> CompilerM op s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign (PyExp -> PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s (PyExp -> PyStmt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest CompilerM op s (PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s PyStmt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
src

compileCode (Imp.Allocate VName
name (Imp.Count Exp
e) (Imp.Space [Char]
space)) =
  CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s ()) -> CompilerM op s ())
-> CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ (CompilerEnv op s -> Allocate op s)
-> CompilerM op s (Allocate op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> Allocate op s
forall op s. CompilerEnv op s -> Allocate op s
envAllocate
    CompilerM op s (Allocate op s)
-> CompilerM op s PyExp
-> CompilerM op s (PyExp -> [Char] -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name
    CompilerM op s (PyExp -> [Char] -> CompilerM op s ())
-> CompilerM op s PyExp
-> CompilerM op s ([Char] -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
    CompilerM op s ([Char] -> CompilerM op s ())
-> CompilerM op s [Char] -> CompilerM op s (CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> CompilerM op s [Char]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space

compileCode (Imp.Allocate VName
name (Imp.Count Exp
e) Space
_) = do
  PyExp
e' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
e
  let allocate' :: PyExp
allocate' = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"allocateMem" [PyExp
e']
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ())
-> CompilerM op s PyStmt -> CompilerM op s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign (PyExp -> PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s (PyExp -> PyStmt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name CompilerM op s (PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s PyStmt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PyExp -> CompilerM op s PyExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure PyExp
allocate'

compileCode (Imp.Free VName
name Space
_) =
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ())
-> CompilerM op s PyStmt -> CompilerM op s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PyExp -> PyExp -> PyStmt
Assign (PyExp -> PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s (PyExp -> PyStmt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
name CompilerM op s (PyExp -> PyStmt)
-> CompilerM op s PyExp -> CompilerM op s PyStmt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PyExp -> CompilerM op s PyExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure PyExp
None

compileCode (Imp.Copy VName
dest (Imp.Count Exp
destoffset) Space
DefaultSpace VName
src (Imp.Count Exp
srcoffset) Space
DefaultSpace (Imp.Count Exp
size)) = do
  PyExp
destoffset' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
destoffset
  PyExp
srcoffset' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
srcoffset
  PyExp
dest' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
  PyExp
src' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
src
  PyExp
size' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
size
  let offset_call1 :: PyExp
offset_call1 = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"addressOffset" [PyExp
dest', PyExp
destoffset', [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  let offset_call2 :: PyExp
offset_call2 = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"addressOffset" [PyExp
src', PyExp
srcoffset', [Char] -> PyExp
Var [Char]
"ct.c_byte"]
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"ct.memmove" [PyExp
offset_call1, PyExp
offset_call2, PyExp
size']

compileCode (Imp.Copy VName
dest (Imp.Count Exp
destoffset) Space
destspace VName
src (Imp.Count Exp
srcoffset) Space
srcspace (Imp.Count Exp
size)) = do
  Copy op s
copy <- (CompilerEnv op s -> Copy op s) -> CompilerM op s (Copy op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> Copy op s
forall op s. CompilerEnv op s -> Copy op s
envCopy
  CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s ()) -> CompilerM op s ())
-> CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ Copy op s
copy
    Copy op s
-> CompilerM op s PyExp
-> CompilerM
     op
     s
     (PyExp
      -> Space
      -> PyExp
      -> PyExp
      -> Space
      -> PyExp
      -> PrimType
      -> CompilerM op s ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest CompilerM
  op
  s
  (PyExp
   -> Space
   -> PyExp
   -> PyExp
   -> Space
   -> PyExp
   -> PrimType
   -> CompilerM op s ())
-> CompilerM op s PyExp
-> CompilerM
     op
     s
     (Space
      -> PyExp
      -> PyExp
      -> Space
      -> PyExp
      -> PrimType
      -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
destoffset CompilerM
  op
  s
  (Space
   -> PyExp
   -> PyExp
   -> Space
   -> PyExp
   -> PrimType
   -> CompilerM op s ())
-> CompilerM op s Space
-> CompilerM
     op
     s
     (PyExp -> PyExp -> Space -> PyExp -> PrimType -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Space -> CompilerM op s Space
forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
destspace
    CompilerM
  op
  s
  (PyExp -> PyExp -> Space -> PyExp -> PrimType -> CompilerM op s ())
-> CompilerM op s PyExp
-> CompilerM
     op s (PyExp -> Space -> PyExp -> PrimType -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
src CompilerM
  op s (PyExp -> Space -> PyExp -> PrimType -> CompilerM op s ())
-> CompilerM op s PyExp
-> CompilerM op s (Space -> PyExp -> PrimType -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
srcoffset CompilerM op s (Space -> PyExp -> PrimType -> CompilerM op s ())
-> CompilerM op s Space
-> CompilerM op s (PyExp -> PrimType -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Space -> CompilerM op s Space
forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
srcspace
    CompilerM op s (PyExp -> PrimType -> CompilerM op s ())
-> CompilerM op s PyExp
-> CompilerM op s (PrimType -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
size CompilerM op s (PrimType -> CompilerM op s ())
-> CompilerM op s PrimType -> CompilerM op s (CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimType -> CompilerM op s PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType -> PrimType
IntType IntType
Int32) -- FIXME

compileCode (Imp.Write VName
dest (Imp.Count Exp
idx) PrimType
elemtype (Imp.Space [Char]
space) Volatility
_ Exp
elemexp) =
  CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (CompilerM op s (CompilerM op s ()) -> CompilerM op s ())
-> CompilerM op s (CompilerM op s ()) -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ (CompilerEnv op s -> WriteScalar op s)
-> CompilerM op s (WriteScalar op s)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks CompilerEnv op s -> WriteScalar op s
forall op s. CompilerEnv op s -> WriteScalar op s
envWriteScalar
    CompilerM op s (WriteScalar op s)
-> CompilerM op s PyExp
-> CompilerM
     op s (PyExp -> PrimType -> [Char] -> PyExp -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
    CompilerM
  op s (PyExp -> PrimType -> [Char] -> PyExp -> CompilerM op s ())
-> CompilerM op s PyExp
-> CompilerM
     op s (PrimType -> [Char] -> PyExp -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
idx
    CompilerM op s (PrimType -> [Char] -> PyExp -> CompilerM op s ())
-> CompilerM op s PrimType
-> CompilerM op s ([Char] -> PyExp -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimType -> CompilerM op s PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
elemtype
    CompilerM op s ([Char] -> PyExp -> CompilerM op s ())
-> CompilerM op s [Char]
-> CompilerM op s (PyExp -> CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> CompilerM op s [Char]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
space
    CompilerM op s (PyExp -> CompilerM op s ())
-> CompilerM op s PyExp -> CompilerM op s (CompilerM op s ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
elemexp

compileCode (Imp.Write VName
dest (Imp.Count Exp
idx) PrimType
elemtype Space
_ Volatility
_ Exp
elemexp) = do
  PyExp
idx' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
idx
  PyExp
elemexp' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp Exp
elemexp
  PyExp
dest' <- VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
dest
  let elemtype' :: [Char]
elemtype' = PrimType -> [Char]
compilePrimType PrimType
elemtype
      ctype :: PyExp
ctype = [Char] -> [PyExp] -> PyExp
simpleCall [Char]
elemtype' [PyExp
elemexp']
  PyStmt -> CompilerM op s ()
forall op s. PyStmt -> CompilerM op s ()
stm (PyStmt -> CompilerM op s ()) -> PyStmt -> CompilerM op s ()
forall a b. (a -> b) -> a -> b
$ PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"writeScalarArray" [PyExp
dest', PyExp
idx', PyExp
ctype]

compileCode Code op
Imp.Skip = () -> CompilerM op s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()