{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# 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,
    CompilerMode,
    Constructor (..),
    emptyConstructor,
    compileName,
    compileVar,
    compileDim,
    compileExp,
    compilePrimExp,
    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.RWS
import qualified Data.Map as M
import Data.Maybe
import Futhark.CodeGen.Backends.GenericPython.AST
import Futhark.CodeGen.Backends.GenericPython.Definitions
import Futhark.CodeGen.Backends.GenericPython.Options
import qualified Futhark.CodeGen.ImpCode as Imp
import Futhark.Compiler.CLI (CompilerMode (..))
import Futhark.IR.Primitive hiding (Bool)
import Futhark.IR.Prop (isBuiltInFunction, subExpVars)
import Futhark.IR.Syntax (Space (..))
import Futhark.MonadFreshNames
import Futhark.Util (zEncodeString)

-- | 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]
"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"]]
      },
    Option :: [Char] -> Maybe Char -> OptionArgument -> [PyStmt] -> Option
Option
      { optionLongName :: [Char]
optionLongName = [Char]
"log",
        optionShortName :: Maybe Char
optionShortName = Char -> Maybe Char
forall a. a -> Maybe a
Just Char
'L',
        optionArgument :: OptionArgument
optionArgument = OptionArgument
NoArgument,
        optionAction :: [PyStmt]
optionAction = [PyStmt
Pass]
      }
  ]

executableOptions :: [Option]
executableOptions :: [Option]
executableOptions =
  [Option]
standardOptions
    [Option] -> [Option] -> [Option]
forall a. [a] -> [a] -> [a]
++ [ 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]
           }
       ]

-- | 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 =>
  CompilerMode ->
  String ->
  Constructor ->
  [PyStmt] ->
  [PyStmt] ->
  Operations op s ->
  s ->
  [PyStmt] ->
  [Option] ->
  Imp.Definitions op ->
  m String
compileProg :: CompilerMode
-> [Char]
-> Constructor
-> [PyStmt]
-> [PyStmt]
-> Operations op s
-> s
-> [PyStmt]
-> [Option]
-> Definitions op
-> m [Char]
compileProg CompilerMode
mode [Char]
class_name Constructor
constructor [PyStmt]
imports [PyStmt]
defines Operations op s
ops s
userstate [PyStmt]
sync [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 CompilerMode
mode of
          CompilerMode
ToLibrary -> [Char]
""
          CompilerMode
_ -> [Char]
"#!/usr/bin/env python3\n"
  [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]
pyValues,
                   [Char] -> PyStmt
Escape [Char]
pyFunctions,
                   [Char] -> PyStmt
Escape [Char]
pyPanic,
                   [Char] -> PyStmt
Escape [Char]
pyTuning,
                   [Char] -> PyStmt
Escape [Char]
pyUtility,
                   [Char] -> PyStmt
Escape [Char]
pyServer
                 ]
              [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 CompilerMode
mode of
        CompilerMode
ToLibrary -> 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
                ([PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (PyFunDef, (PyExp, PyExp))
forall op s.
[PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun [PyStmt]
sync ReturnTiming
DoNotReturnTiming)
                (((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]
class_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)
            ]
        CompilerMode
ToServer -> 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
                ([PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (PyFunDef, (PyExp, PyExp))
forall op s.
[PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun [PyStmt]
sync ReturnTiming
ReturnTiming)
                (((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] -> CompilerM op s [PyStmt])
-> [PyStmt] -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$
            [PyStmt]
parse_options_server
              [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [ PyClassDef -> PyStmt
ClassDef
                     ( [Char] -> [PyStmt] -> PyClassDef
Class [Char]
class_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)
                     ),
                   PyExp -> PyExp -> PyStmt
Assign
                     ([Char] -> PyExp
Var [Char]
"server")
                     ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"Server" [[Char] -> [PyExp] -> PyExp
simpleCall [Char]
class_name []]),
                   PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"server.run" []
                 ]
        CompilerMode
ToExecutable -> 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]
class_name []
          ([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]
sync)
                (((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] -> CompilerM op s [PyStmt])
-> [PyStmt] -> CompilerM op s [PyStmt]
forall a b. (a -> b) -> a -> b
$
            [PyStmt]
parse_options_executable
              [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ PyClassDef -> PyStmt
ClassDef
                ( [Char] -> [PyStmt] -> PyClassDef
Class [Char]
class_name ([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_executable :: [PyStmt]
parse_options_executable =
      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]
executableOptions [Option] -> [Option] -> [Option]
forall a. [a] -> [a] -> [a]
++ [Option]
options)

    parse_options_server :: [PyStmt]
parse_options_server =
      [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 -> CompilerM op s PyExp
compileDim :: DimSize -> CompilerM op s PyExp
compileDim (Imp.Constant PrimValue
v) = PyExp -> CompilerM op s PyExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PyExp -> CompilerM op s PyExp) -> PyExp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PyExp
compilePrimValue PrimValue
v
compileDim (Imp.Var VName
v) = VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar 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]
"Entry point arguments have invalid sizes."
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 -> [PyStmt] -> [PyStmt] -> PyStmt
If
      ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" PyExp
var' PyExp
None)
      [PyExp -> PyExp -> PyStmt
Assign PyExp
var' (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.int64" [PyExp
src]]
      [ PyExp -> PyExp -> PyStmt
Assert ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"==" PyExp
var' PyExp
src) (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
          [Char] -> PyExp
String [Char]
"Error: entry point arguments have invalid sizes."
      ]

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]
dims' <- (DimSize -> CompilerM op s PyExp)
-> [DimSize] -> CompilerM op s [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimSize -> CompilerM op s PyExp
forall op s. DimSize -> CompilerM op s PyExp
compileDim [DimSize]
dims
  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]
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: {}"
        ]

badInputType :: Int -> PyExp -> String -> PyExp -> PyExp -> PyStmt
badInputType :: Int -> PyExp -> [Char] -> PyExp -> PyExp -> PyStmt
badInputType Int
i PyExp
e [Char]
t PyExp
de PyExp
dg =
  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, PyExp -> PyArg
Arg PyExp
de, PyExp -> PyArg
Arg PyExp
dg]
      ]
  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: {}",
          [Char]
"Expected array with elements of dtype: {}",
          [Char]
"The array given has elements of dtype: {}"
        ]

badInputDim :: Int -> PyExp -> String -> Int -> PyStmt
badInputDim :: Int -> PyExp -> [Char] -> Int -> PyStmt
badInputDim Int
i PyExp
e [Char]
typ Int
dimf =
  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 PyExp
eft, PyExp -> PyArg
Arg PyExp
aft]
      ]
  where
    eft :: PyExp
eft = [Char] -> PyExp
String ([[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Char] -> [[Char]]
forall a. Int -> a -> [a]
replicate Int
dimf [Char]
"[]") [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
typ)
    aft :: PyExp
aft = [Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"+" ([Char] -> PyExp -> PyExp -> PyExp
BinOp [Char]
"*" ([Char] -> PyExp
String [Char]
"[]") (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"ndim")) ([Char] -> PyExp
String [Char]
typ)
    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]
"Dimensionality mismatch",
          [Char]
"Expected Futhark type: {}",
          [Char]
"Bad Python value passed",
          [Char]
"Actual Futhark type: {}"
        ]

declEntryPointInputSizes :: [Imp.ExternalValue] -> CompilerM op s ()
declEntryPointInputSizes :: [ExternalValue] -> CompilerM op s ()
declEntryPointInputSizes = (VName -> CompilerM op s ()) -> [VName] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ VName -> CompilerM op s ()
forall op s. VName -> CompilerM op s ()
onSize ([VName] -> CompilerM op s ())
-> ([ExternalValue] -> [VName])
-> [ExternalValue]
-> CompilerM op s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ExternalValue -> [VName]) -> [ExternalValue] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ExternalValue -> [VName]
sizes
  where
    sizes :: ExternalValue -> [VName]
sizes (Imp.TransparentValue ValueDesc
v) = ValueDesc -> [VName]
valueSizes ValueDesc
v
    sizes (Imp.OpaqueValue [Char]
_ [ValueDesc]
vs) = (ValueDesc -> [VName]) -> [ValueDesc] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ValueDesc -> [VName]
valueSizes [ValueDesc]
vs
    valueSizes :: ValueDesc -> [VName]
valueSizes (Imp.ArrayValue VName
_ Space
_ PrimType
_ Signedness
_ [DimSize]
dims) = [DimSize] -> [VName]
subExpVars [DimSize]
dims
    valueSizes Imp.ScalarValue {} = []
    onSize :: VName -> CompilerM op s ()
onSize 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 ([Char] -> PyExp
Var (VName -> [Char]
compileName VName
v)) PyExp
None

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
s Signedness -> 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
ept Signedness -> 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]
"in" ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"type" [PyExp
e]) (PyExp -> PyExp) -> PyExp -> PyExp
forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List [[Char] -> PyExp
Var [Char]
"np.ndarray"]
  let dtype_is_wrong :: PyExp
dtype_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]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"dtype") (PyExp -> PyExp) -> PyExp -> PyExp
forall a b. (a -> b) -> a -> b
$ [Char] -> PyExp
Var ([Char] -> PyExp) -> [Char] -> PyExp
forall a b. (a -> b) -> a -> b
$ PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
t Signedness
s
  let dim_is_wrong :: PyExp
dim_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]
"==" (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"ndim") (PyExp -> PyExp) -> PyExp -> PyExp
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. Integral a => a -> Integer
toInteger (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [DimSize] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims
  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
s Signedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
== Signedness
Imp.TypeUnsigned) PrimType
t
      ]
      []
  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
dtype_is_wrong
      [ Int -> PyExp -> [Char] -> PyExp -> PyExp -> PyStmt
badInputType
          Int
i
          PyExp
e
          ([[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
s Signedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
== Signedness
Imp.TypeUnsigned) PrimType
t)
          ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"np.dtype" [[Char] -> PyExp
Var (PrimType -> Signedness -> [Char]
compilePrimToExtNp PrimType
t Signedness
s)])
          (PyExp -> [Char] -> PyExp
Field PyExp
e [Char]
"dtype")
      ]
      []
  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
dim_is_wrong
      [Int -> PyExp -> [Char] -> Int -> PyStmt
badInputDim Int
i PyExp
e (Bool -> PrimType -> [Char]
prettySigned (Signedness
s Signedness -> Signedness -> Bool
forall a. Eq a => a -> a -> Bool
== Signedness
Imp.TypeUnsigned) PrimType
t) ([DimSize] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimSize]
dims)]
      []

  (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')
  where
    -- 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.
    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
$ do
    [ExternalValue] -> CompilerM op s ()
forall op s. [ExternalValue] -> CompilerM op s ()
declEntryPointInputSizes [ExternalValue]
args
    ((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)

      -- 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"]

      call :: [[Char]] -> [PyStmt]
call [[Char]]
argexps =
        [ PyExp -> [PyStmt] -> PyStmt
With
            PyExp
errstate
            [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)]
        ]

  ([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,
      [[Char]] -> [PyStmt]
call [[Char]]
argexps_lib,
      [[Char]] -> [PyStmt]
call [[Char]]
argexps_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]

data ReturnTiming = ReturnTiming | DoNotReturnTiming

compileEntryFun ::
  [PyStmt] ->
  ReturnTiming ->
  (Name, Imp.Function op) ->
  CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun :: [PyStmt]
-> ReturnTiming
-> (Name, Function op)
-> CompilerM op s (PyFunDef, (PyExp, PyExp))
compileEntryFun [PyStmt]
sync ReturnTiming
timing (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 ([PyStmt]
maybe_sync, PyStmt
ret) =
        case ReturnTiming
timing of
          ReturnTiming
DoNotReturnTiming ->
            ( [],
              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
            )
          ReturnTiming
ReturnTiming ->
            ( [PyStmt]
sync,
              PyExp -> PyStmt
Return (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
                [PyExp] -> PyExp
Tuple
                  [ [Char] -> PyExp
Var [Char]
"runtime",
                    [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

      do_run :: [PyStmt]
do_run =
        PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"time_start") ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" []) PyStmt -> [PyStmt] -> [PyStmt]
forall a. a -> [a] -> [a]
:
        [PyStmt]
body_lib [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
maybe_sync
          [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [ PyExp -> PyExp -> PyStmt
Assign ([Char] -> PyExp
Var [Char]
"runtime") (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$
                 [Char] -> PyExp -> PyExp -> PyExp
BinOp
                   [Char]
"-"
                   (PyExp -> PyExp
toMicroseconds ([Char] -> [PyExp] -> PyExp
simpleCall [Char]
"time.time" []))
                   (PyExp -> PyExp
toMicroseconds ([Char] -> PyExp
Var [Char]
"time_start"))
             ]

  (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]
do_run [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
prepareOut [PyStmt] -> [PyStmt] -> [PyStmt]
forall a. [a] -> [a] -> [a]
++ [PyStmt]
sync [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

      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 [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"],
        PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ [Char] -> [PyExp] -> PyExp
simpleCall [Char]
"runtime_file.flush" []
      ]

toMicroseconds :: PyExp -> PyExp
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"
    FSignum {} -> [Char]
"np.sign"

compileBinOpLike ::
  Monad m =>
  (v -> m PyExp) ->
  Imp.PrimExp v ->
  Imp.PrimExp v ->
  m (PyExp, PyExp, String -> m PyExp)
compileBinOpLike :: (v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike v -> m PyExp
f PrimExp v
x PrimExp v
y = do
  PyExp
x' <- (v -> m PyExp) -> PrimExp v -> m PyExp
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
x
  PyExp
y' <- (v -> m PyExp) -> PrimExp v -> m PyExp
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
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)
-> m (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

-- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you.
compilePrimExp :: Monad m => (v -> m PyExp) -> Imp.PrimExp v -> m PyExp
compilePrimExp :: (v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
_ (Imp.ValueExp PrimValue
v) = 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
$ PrimValue -> PyExp
compilePrimValue PrimValue
v
compilePrimExp v -> m PyExp
f (Imp.LeafExp v
v PrimType
_) = v -> m PyExp
f v
v
compilePrimExp v -> m PyExp
f (Imp.BinOpExp BinOp
op PrimExp v
x PrimExp v
y) = do
  (PyExp
x', PyExp
y', [Char] -> m PyExp
simple) <- (v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike v -> m PyExp
f PrimExp v
x PrimExp v
y
  case BinOp
op of
    Add {} -> [Char] -> m PyExp
simple [Char]
"+"
    Sub {} -> [Char] -> m PyExp
simple [Char]
"-"
    Mul {} -> [Char] -> m PyExp
simple [Char]
"*"
    FAdd {} -> [Char] -> m PyExp
simple [Char]
"+"
    FSub {} -> [Char] -> m PyExp
simple [Char]
"-"
    FMul {} -> [Char] -> m PyExp
simple [Char]
"*"
    FDiv {} -> [Char] -> m PyExp
simple [Char]
"/"
    FMod {} -> [Char] -> m PyExp
simple [Char]
"%"
    Xor {} -> [Char] -> m PyExp
simple [Char]
"^"
    And {} -> [Char] -> m PyExp
simple [Char]
"&"
    Or {} -> [Char] -> m PyExp
simple [Char]
"|"
    Shl {} -> [Char] -> m PyExp
simple [Char]
"<<"
    LogAnd {} -> [Char] -> m PyExp
simple [Char]
"and"
    LogOr {} -> [Char] -> m PyExp
simple [Char]
"or"
    BinOp
_ -> 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
simpleCall (BinOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty BinOp
op) [PyExp
x', PyExp
y']
compilePrimExp v -> m PyExp
f (Imp.ConvOpExp ConvOp
conv PrimExp v
x) = do
  PyExp
x' <- (v -> m PyExp) -> PrimExp v -> m PyExp
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
x
  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
simpleCall (ConvOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty ConvOp
conv) [PyExp
x']
compilePrimExp v -> m PyExp
f (Imp.CmpOpExp CmpOp
cmp PrimExp v
x PrimExp v
y) = do
  (PyExp
x', PyExp
y', [Char] -> m PyExp
simple) <- (v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp)
-> PrimExp v -> PrimExp v -> m (PyExp, PyExp, [Char] -> m PyExp)
compileBinOpLike v -> m PyExp
f PrimExp v
x PrimExp v
y
  case CmpOp
cmp of
    CmpEq {} -> [Char] -> m PyExp
simple [Char]
"=="
    FCmpLt {} -> [Char] -> m PyExp
simple [Char]
"<"
    FCmpLe {} -> [Char] -> m PyExp
simple [Char]
"<="
    CmpOp
CmpLlt -> [Char] -> m PyExp
simple [Char]
"<"
    CmpOp
CmpLle -> [Char] -> m PyExp
simple [Char]
"<="
    CmpOp
_ -> 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
simpleCall (CmpOp -> [Char]
forall a. Pretty a => a -> [Char]
pretty CmpOp
cmp) [PyExp
x', PyExp
y']
compilePrimExp v -> m PyExp
f (Imp.UnOpExp UnOp
op PrimExp v
exp1) =
  [Char] -> PyExp -> PyExp
UnOp (UnOp -> [Char]
compileUnOp UnOp
op) (PyExp -> PyExp) -> m PyExp -> m PyExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (v -> m PyExp) -> PrimExp v -> m PyExp
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f PrimExp v
exp1
compilePrimExp v -> m PyExp
f (Imp.FunExp [Char]
h [PrimExp v]
args PrimType
_) =
  [Char] -> [PyExp] -> PyExp
simpleCall ([Char] -> [Char]
futharkFun ([Char] -> [Char]
forall a. Pretty a => a -> [Char]
pretty [Char]
h)) ([PyExp] -> PyExp) -> m [PyExp] -> m PyExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimExp v -> m PyExp) -> [PrimExp v] -> m [PyExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((v -> m PyExp) -> PrimExp v -> m PyExp
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp v -> m PyExp
f) [PrimExp v]
args

compileExp :: Imp.Exp -> CompilerM op s PyExp
compileExp :: Exp -> CompilerM op s PyExp
compileExp = (ExpLeaf -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall (m :: * -> *) v.
Monad m =>
(v -> m PyExp) -> PrimExp v -> m PyExp
compilePrimExp ExpLeaf -> CompilerM op s PyExp
forall op s. ExpLeaf -> CompilerM op s PyExp
compileLeaf
  where
    compileLeaf :: ExpLeaf -> CompilerM op s PyExp
compileLeaf (Imp.ScalarVar VName
vname) =
      VName -> CompilerM op s PyExp
forall op s. VName -> CompilerM op s PyExp
compileVar VName
vname
    compileLeaf (Imp.SizeOf PrimType
t) =
      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]
    compileLeaf (Imp.Index VName
src (Imp.Count TExp Int64
iexp) PrimType
restype (Imp.Space [Char]
space) Volatility
_) =
      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 (TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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
    compileLeaf (Imp.Index VName
src (Imp.Count TExp Int64
iexp) PrimType
bt Space
_ Volatility
_) = do
      PyExp
iexp' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp (Exp -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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]

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 TExp Bool
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 -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Bool -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Bool
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 TExp Bool
cond Code op
body) = do
  PyExp
cond' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp (Exp -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Bool -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Bool
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 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 (Exp -> PrimType
forall v. PrimExp v -> PrimType
Imp.primExpType Exp
bound)) [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 (Exp -> PrimType
forall v. PrimExp v -> PrimType
Imp.primExpType Exp
bound)) [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
      onPart (Imp.ErrorInt64 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
loc SrcLoc -> [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 (Imp.TPrimExp 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 (Imp.TPrimExp 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 TExp Int64
destoffset) Space
DefaultSpace VName
src (Imp.Count TExp Int64
srcoffset) Space
DefaultSpace (Imp.Count TExp Int64
size)) = do
  PyExp
destoffset' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp (Exp -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
destoffset
  PyExp
srcoffset' <- Exp -> CompilerM op s PyExp
forall op s. Exp -> CompilerM op s PyExp
compileExp (Exp -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 TExp Int64
destoffset) Space
destspace VName
src (Imp.Count TExp Int64
srcoffset) Space
srcspace (Imp.Count TExp Int64
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 (TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 (TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 (TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 TExp Int64
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 (TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 TExp Int64
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 -> CompilerM op s PyExp) -> Exp -> CompilerM op s PyExp
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> Exp
forall t v. TPrimExp t v -> PrimExp v
Imp.untyped TExp Int64
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 ()