-- | An interpreter operating on type-checked source Futhark terms.
-- Relatively slow.
module Language.Futhark.Interpreter
  ( Ctx (..),
    Env,
    InterpreterError,
    prettyInterpreterError,
    initialCtx,
    interpretExp,
    interpretDec,
    interpretImport,
    interpretFunction,
    ctxWithImports,
    ExtOp (..),
    BreakReason (..),
    StackFrame (..),
    typeCheckerEnv,

    -- * Values
    Value,
    fromTuple,
    isEmptyArray,
    prettyEmptyArray,
    prettyValue,
    valueText,
  )
where

import Control.Monad
import Control.Monad.Free.Church
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Trans.Maybe
import Data.Array
import Data.Bifunctor (first, second)
import Data.List
  ( find,
    foldl',
    genericLength,
    genericTake,
    isPrefixOf,
    transpose,
  )
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Maybe
import Data.Monoid hiding (Sum)
import Data.Text qualified as T
import Futhark.Data qualified as V
import Futhark.Util (chunk, maybeHead, splitFromEnd)
import Futhark.Util.Loc
import Futhark.Util.Pretty hiding (apply)
import Language.Futhark hiding (Shape, matchDims)
import Language.Futhark qualified as F
import Language.Futhark.Interpreter.Values hiding (Value)
import Language.Futhark.Interpreter.Values qualified
import Language.Futhark.Primitive (floatValue, intValue)
import Language.Futhark.Primitive qualified as P
import Language.Futhark.Semantic qualified as T
import Prelude hiding (break, mod)

data StackFrame = StackFrame
  { StackFrame -> Loc
stackFrameLoc :: Loc,
    StackFrame -> Ctx
stackFrameCtx :: Ctx
  }

instance Located StackFrame where
  locOf :: StackFrame -> Loc
locOf = StackFrame -> Loc
stackFrameLoc

-- | What is the reason for this break point?
data BreakReason
  = -- | An explicit breakpoint in the program.
    BreakPoint
  | -- | A
    BreakNaN

data ExtOp a
  = ExtOpTrace T.Text (Doc ()) a
  | ExtOpBreak Loc BreakReason (NE.NonEmpty StackFrame) a
  | ExtOpError InterpreterError

instance Functor ExtOp where
  fmap :: forall a b. (a -> b) -> ExtOp a -> ExtOp b
fmap a -> b
f (ExtOpTrace Text
w Doc ()
s a
x) = forall a. Text -> Doc () -> a -> ExtOp a
ExtOpTrace Text
w Doc ()
s forall a b. (a -> b) -> a -> b
$ a -> b
f a
x
  fmap a -> b
f (ExtOpBreak Loc
w BreakReason
why NonEmpty StackFrame
backtrace a
x) = forall a. Loc -> BreakReason -> NonEmpty StackFrame -> a -> ExtOp a
ExtOpBreak Loc
w BreakReason
why NonEmpty StackFrame
backtrace forall a b. (a -> b) -> a -> b
$ a -> b
f a
x
  fmap a -> b
_ (ExtOpError InterpreterError
err) = forall a. InterpreterError -> ExtOp a
ExtOpError InterpreterError
err

type Stack = [StackFrame]

type Sizes = M.Map VName Int64

-- | The monad in which evaluation takes place.
newtype EvalM a
  = EvalM
      ( ReaderT
          (Stack, M.Map ImportName Env)
          (StateT Sizes (F ExtOp))
          a
      )
  deriving
    ( Applicative EvalM
forall a. a -> EvalM a
forall a b. EvalM a -> EvalM b -> EvalM b
forall a b. EvalM a -> (a -> EvalM b) -> EvalM 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 :: forall a. a -> EvalM a
$creturn :: forall a. a -> EvalM a
>> :: forall a b. EvalM a -> EvalM b -> EvalM b
$c>> :: forall a b. EvalM a -> EvalM b -> EvalM b
>>= :: forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
$c>>= :: forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
Monad,
      Functor EvalM
forall a. a -> EvalM a
forall a b. EvalM a -> EvalM b -> EvalM a
forall a b. EvalM a -> EvalM b -> EvalM b
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall a b c. (a -> b -> c) -> EvalM a -> EvalM b -> EvalM 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
<* :: forall a b. EvalM a -> EvalM b -> EvalM a
$c<* :: forall a b. EvalM a -> EvalM b -> EvalM a
*> :: forall a b. EvalM a -> EvalM b -> EvalM b
$c*> :: forall a b. EvalM a -> EvalM b -> EvalM b
liftA2 :: forall a b c. (a -> b -> c) -> EvalM a -> EvalM b -> EvalM c
$cliftA2 :: forall a b c. (a -> b -> c) -> EvalM a -> EvalM b -> EvalM c
<*> :: forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
$c<*> :: forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
pure :: forall a. a -> EvalM a
$cpure :: forall a. a -> EvalM a
Applicative,
      forall a b. a -> EvalM b -> EvalM a
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> EvalM b -> EvalM a
$c<$ :: forall a b. a -> EvalM b -> EvalM a
fmap :: forall a b. (a -> b) -> EvalM a -> EvalM b
$cfmap :: forall a b. (a -> b) -> EvalM a -> EvalM b
Functor,
      MonadFree ExtOp,
      MonadReader (Stack, M.Map ImportName Env),
      MonadState Sizes
    )

runEvalM :: M.Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM :: forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM Map ImportName Env
imports (EvalM ReaderT
  ([StackFrame], Map ImportName Env) (StateT Sizes (F ExtOp)) a
m) = forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT
  ([StackFrame], Map ImportName Env) (StateT Sizes (F ExtOp)) a
m (forall a. Monoid a => a
mempty, Map ImportName Env
imports)) forall a. Monoid a => a
mempty

stacking :: SrcLoc -> Env -> EvalM a -> EvalM a
stacking :: forall a. SrcLoc -> Env -> EvalM a -> EvalM a
stacking SrcLoc
loc Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \([StackFrame]
ss, Map ImportName Env
imports) ->
  if SrcLoc -> Bool
isNoLoc SrcLoc
loc
    then ([StackFrame]
ss, Map ImportName Env
imports)
    else
      let s :: StackFrame
s = Loc -> Ctx -> StackFrame
StackFrame (forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Env -> Map ImportName Env -> Ctx
Ctx Env
env Map ImportName Env
imports)
       in (StackFrame
s forall a. a -> [a] -> [a]
: [StackFrame]
ss, Map ImportName Env
imports)
  where
    isNoLoc :: SrcLoc -> Bool
    isNoLoc :: SrcLoc -> Bool
isNoLoc = (forall a. Eq a => a -> a -> Bool
== Loc
NoLoc) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Located a => a -> Loc
locOf

stacktrace :: EvalM [Loc]
stacktrace :: EvalM [Loc]
stacktrace = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map StackFrame -> Loc
stackFrameLoc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst

lookupImport :: ImportName -> EvalM (Maybe Env)
lookupImport :: ImportName -> EvalM (Maybe Env)
lookupImport ImportName
f = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ImportName
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd

putExtSize :: VName -> Int64 -> EvalM ()
putExtSize :: VName -> Int64 -> EvalM ()
putExtSize VName
v Int64
x = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v Int64
x

getSizes :: EvalM Sizes
getSizes :: EvalM Sizes
getSizes = forall s (m :: * -> *). MonadState s m => m s
get

-- | Disregard any existential sizes computed during this action.
-- This is used so that existentials computed during one iteration of
-- a loop or a function call are not remembered the next time around.
localExts :: EvalM a -> EvalM a
localExts :: forall a. EvalM a -> EvalM a
localExts EvalM a
m = do
  Sizes
s <- forall s (m :: * -> *). MonadState s m => m s
get
  a
x <- EvalM a
m
  forall s (m :: * -> *). MonadState s m => s -> m ()
put Sizes
s
  forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

extSizeEnv :: EvalM Env
extSizeEnv :: EvalM Env
extSizeEnv = Sizes -> Env
i64Env forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM Sizes
getSizes

valueStructType :: ValueType -> StructType
valueStructType :: ValueType -> StructType
valueStructType = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Int64 -> Size
ConstSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral)

resolveTypeParams :: [VName] -> StructType -> StructType -> Env
resolveTypeParams :: [VName] -> StructType -> StructType -> Env
resolveTypeParams [VName]
names = StructType -> StructType -> Env
match
  where
    match :: StructType -> StructType -> Env
match (Scalar (TypeVar ()
_ Uniqueness
_ QualName VName
tn [TypeArg Size]
_)) StructType
t
      | forall vn. QualName vn -> vn
qualLeaf QualName VName
tn forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
names =
          Map VName StructType -> Env
typeEnv forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton (forall vn. QualName vn -> vn
qualLeaf QualName VName
tn) StructType
t
    match (Scalar (Record Map Name StructType
poly_fields)) (Scalar (Record Map Name StructType
fields)) =
      forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
        forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$
          forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith StructType -> StructType -> Env
match Map Name StructType
poly_fields Map Name StructType
fields
    match (Scalar (Sum Map Name [StructType]
poly_fields)) (Scalar (Sum Map Name [StructType]
fields)) =
      forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
        forall a b. (a -> b) -> [a] -> [b]
map forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
          forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$
            forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith StructType -> StructType -> Env
match) Map Name [StructType]
poly_fields Map Name [StructType]
fields
    match
      (Scalar (Arrow ()
_ PName
_ Diet
_ StructType
poly_t1 (RetType [VName]
_ StructType
poly_t2)))
      (Scalar (Arrow ()
_ PName
_ Diet
_ StructType
t1 (RetType [VName]
_ StructType
t2))) =
        StructType -> StructType -> Env
match StructType
poly_t1 StructType
t1 forall a. Semigroup a => a -> a -> a
<> StructType -> StructType -> Env
match StructType
poly_t2 StructType
t2
    match StructType
poly_t StructType
t
      | Size
d1 : [Size]
_ <- forall dim. Shape dim -> [dim]
shapeDims (forall dim as. TypeBase dim as -> Shape dim
arrayShape StructType
poly_t),
        Size
d2 : [Size]
_ <- forall dim. Shape dim -> [dim]
shapeDims (forall dim as. TypeBase dim as -> Shape dim
arrayShape StructType
t) =
          Size -> Size -> Env
matchDims Size
d1 Size
d2 forall a. Semigroup a => a -> a -> a
<> StructType -> StructType -> Env
match (forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 StructType
poly_t) (forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 StructType
t)
    match StructType
_ StructType
_ = forall a. Monoid a => a
mempty

    matchDims :: Size -> Size -> Env
matchDims (NamedSize (QualName [VName]
_ VName
d1)) (ConstSize Int64
d2)
      | VName
d1 forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
names =
          Sizes -> Env
i64Env forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton VName
d1 forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
d2
    matchDims Size
_ Size
_ = forall a. Monoid a => a
mempty

resolveExistentials :: [VName] -> StructType -> ValueShape -> M.Map VName Int64
resolveExistentials :: [VName] -> StructType -> ValueShape -> Sizes
resolveExistentials [VName]
names = forall {as} {a}. TypeBase Size as -> Shape a -> Map VName a
match
  where
    match :: TypeBase Size as -> Shape a -> Map VName a
match (Scalar (Record Map Name (TypeBase Size as)
poly_fields)) (ShapeRecord Map Name (Shape a)
fields) =
      forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
        forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$
          forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase Size as -> Shape a -> Map VName a
match Map Name (TypeBase Size as)
poly_fields Map Name (Shape a)
fields
    match (Scalar (Sum Map Name [TypeBase Size as]
poly_fields)) (ShapeSum Map Name [Shape a]
fields) =
      forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
        forall a b. (a -> b) -> [a] -> [b]
map forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
          forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$
            forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase Size as -> Shape a -> Map VName a
match) Map Name [TypeBase Size as]
poly_fields Map Name [Shape a]
fields
    match TypeBase Size as
poly_t (ShapeDim a
d2 Shape a
rowshape)
      | Size
d1 : [Size]
_ <- forall dim. Shape dim -> [dim]
shapeDims (forall dim as. TypeBase dim as -> Shape dim
arrayShape TypeBase Size as
poly_t) =
          forall {a}. Size -> a -> Map VName a
matchDims Size
d1 a
d2 forall a. Semigroup a => a -> a -> a
<> TypeBase Size as -> Shape a -> Map VName a
match (forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 TypeBase Size as
poly_t) Shape a
rowshape
    match TypeBase Size as
_ Shape a
_ = forall a. Monoid a => a
mempty

    matchDims :: Size -> a -> Map VName a
matchDims (NamedSize (QualName [VName]
_ VName
d1)) a
d2
      | VName
d1 forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
names = forall k a. k -> a -> Map k a
M.singleton VName
d1 a
d2
    matchDims Size
_ a
_ = forall a. Monoid a => a
mempty

checkShape :: Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape :: Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape (ShapeDim Maybe Int64
Nothing Shape (Maybe Int64)
shape1) (ShapeDim Int64
d2 ValueShape
shape2) =
  forall d. d -> Shape d -> Shape d
ShapeDim Int64
d2 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape Shape (Maybe Int64)
shape1 ValueShape
shape2
checkShape (ShapeDim (Just Int64
d1) Shape (Maybe Int64)
shape1) (ShapeDim Int64
d2 ValueShape
shape2) = do
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Int64
d1 forall a. Eq a => a -> a -> Bool
== Int64
d2
  forall d. d -> Shape d -> Shape d
ShapeDim Int64
d2 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape Shape (Maybe Int64)
shape1 ValueShape
shape2
checkShape (ShapeDim Maybe Int64
d1 Shape (Maybe Int64)
shape1) ValueShape
ShapeLeaf =
  -- This case is for handling polymorphism, when a function doesn't
  -- know that the array it produced actually has more dimensions.
  forall d. d -> Shape d -> Shape d
ShapeDim (forall a. a -> Maybe a -> a
fromMaybe Int64
0 Maybe Int64
d1) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape Shape (Maybe Int64)
shape1 forall d. Shape d
ShapeLeaf
checkShape (ShapeRecord Map Name (Shape (Maybe Int64))
shapes1) (ShapeRecord Map Name ValueShape
shapes2) =
  forall d. Map Name (Shape d) -> Shape d
ShapeRecord forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape Map Name (Shape (Maybe Int64))
shapes1 Map Name ValueShape
shapes2)
checkShape (ShapeRecord Map Name (Shape (Maybe Int64))
shapes1) ValueShape
ShapeLeaf =
  forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe Int64
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall d. Map Name (Shape d) -> Shape d
ShapeRecord Map Name (Shape (Maybe Int64))
shapes1
checkShape (ShapeSum Map Name [Shape (Maybe Int64)]
shapes1) (ShapeSum Map Name [ValueShape]
shapes2) =
  forall d. Map Name [Shape d] -> Shape d
ShapeSum forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence (forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape) Map Name [Shape (Maybe Int64)]
shapes1 Map Name [ValueShape]
shapes2)
checkShape (ShapeSum Map Name [Shape (Maybe Int64)]
shapes1) ValueShape
ShapeLeaf =
  forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe Int64
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall d. Map Name [Shape d] -> Shape d
ShapeSum Map Name [Shape (Maybe Int64)]
shapes1
checkShape Shape (Maybe Int64)
_ ValueShape
shape2 =
  forall a. a -> Maybe a
Just ValueShape
shape2

type Value = Language.Futhark.Interpreter.Values.Value EvalM

asInteger :: Value -> Integer
asInteger :: Value -> Integer
asInteger (ValuePrim (SignedValue IntValue
v)) = forall int. Integral int => IntValue -> int
P.valueIntegral IntValue
v
asInteger (ValuePrim (UnsignedValue IntValue
v)) =
  forall a. Integral a => a -> Integer
toInteger (forall int. Integral int => IntValue -> int
P.valueIntegral (IntValue -> IntType -> IntValue
P.doZExt IntValue
v IntType
Int64) :: Word64)
asInteger Value
v = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpectedly not an integer: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
v

asInt :: Value -> Int
asInt :: Value -> Int
asInt = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Integer
asInteger

asSigned :: Value -> IntValue
asSigned :: Value -> IntValue
asSigned (ValuePrim (SignedValue IntValue
v)) = IntValue
v
asSigned Value
v = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected not a signed integer: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
v

asInt64 :: Value -> Int64
asInt64 :: Value -> Int64
asInt64 = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Integer
asInteger

asBool :: Value -> Bool
asBool :: Value -> Bool
asBool (ValuePrim (BoolValue Bool
x)) = Bool
x
asBool Value
v = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpectedly not a boolean: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
v

lookupInEnv ::
  (Env -> M.Map VName x) ->
  QualName VName ->
  Env ->
  Maybe x
lookupInEnv :: forall x. (Env -> Map VName x) -> QualName VName -> Env -> Maybe x
lookupInEnv Env -> Map VName x
onEnv QualName VName
qv Env
env = Env -> [VName] -> Maybe x
f Env
env forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> [vn]
qualQuals QualName VName
qv
  where
    f :: Env -> [VName] -> Maybe x
f Env
m (VName
q : [VName]
qs) =
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm Env
m of
        Just (TermModule (Module Env
mod)) -> Env -> [VName] -> Maybe x
f Env
mod [VName]
qs
        Maybe TermBinding
_ -> forall a. Maybe a
Nothing
    f Env
m [] = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
qv) forall a b. (a -> b) -> a -> b
$ Env -> Map VName x
onEnv Env
m

lookupVar :: QualName VName -> Env -> Maybe TermBinding
lookupVar :: QualName VName -> Env -> Maybe TermBinding
lookupVar = forall x. (Env -> Map VName x) -> QualName VName -> Env -> Maybe x
lookupInEnv Env -> Map VName TermBinding
envTerm

lookupType :: QualName VName -> Env -> Maybe T.TypeBinding
lookupType :: QualName VName -> Env -> Maybe TypeBinding
lookupType = forall x. (Env -> Map VName x) -> QualName VName -> Env -> Maybe x
lookupInEnv Env -> Map VName TypeBinding
envType

-- | A TermValue with a 'Nothing' type annotation is an intrinsic.
data TermBinding
  = TermValue (Maybe T.BoundV) Value
  | -- | A polymorphic value that must be instantiated.
    TermPoly (Maybe T.BoundV) (StructType -> EvalM Value)
  | TermModule Module

data Module
  = Module Env
  | ModuleFun (Module -> EvalM Module)

-- | The actual type- and value environment.
data Env = Env
  { Env -> Map VName TermBinding
envTerm :: M.Map VName TermBinding,
    Env -> Map VName TypeBinding
envType :: M.Map VName T.TypeBinding
  }

instance Monoid Env where
  mempty :: Env
mempty = Map VName TermBinding -> Map VName TypeBinding -> Env
Env forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty

instance Semigroup Env where
  Env Map VName TermBinding
vm1 Map VName TypeBinding
tm1 <> :: Env -> Env -> Env
<> Env Map VName TermBinding
vm2 Map VName TypeBinding
tm2 = Map VName TermBinding -> Map VName TypeBinding -> Env
Env (Map VName TermBinding
vm1 forall a. Semigroup a => a -> a -> a
<> Map VName TermBinding
vm2) (Map VName TypeBinding
tm1 forall a. Semigroup a => a -> a -> a
<> Map VName TypeBinding
tm2)

-- | An error occurred during interpretation due to an error in the
-- user program.  Actual interpreter errors will be signaled with an
-- IO exception ('error').
newtype InterpreterError = InterpreterError T.Text

-- | Prettyprint the error for human consumption.
prettyInterpreterError :: InterpreterError -> Doc AnsiStyle
prettyInterpreterError :: InterpreterError -> Doc AnsiStyle
prettyInterpreterError (InterpreterError Text
e) = forall a ann. Pretty a => a -> Doc ann
pretty Text
e

valEnv :: M.Map VName (Maybe T.BoundV, Value) -> Env
valEnv :: Map VName (Maybe BoundV, Value) -> Env
valEnv Map VName (Maybe BoundV, Value)
m =
  Env
    { envTerm :: Map VName TermBinding
envTerm = forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Maybe BoundV -> Value -> TermBinding
TermValue) Map VName (Maybe BoundV, Value)
m,
      envType :: Map VName TypeBinding
envType = forall a. Monoid a => a
mempty
    }

modEnv :: M.Map VName Module -> Env
modEnv :: Map VName Module -> Env
modEnv Map VName Module
m =
  Env
    { envTerm :: Map VName TermBinding
envTerm = forall a b k. (a -> b) -> Map k a -> Map k b
M.map Module -> TermBinding
TermModule Map VName Module
m,
      envType :: Map VName TypeBinding
envType = forall a. Monoid a => a
mempty
    }

typeEnv :: M.Map VName StructType -> Env
typeEnv :: Map VName StructType -> Env
typeEnv Map VName StructType
m =
  Env
    { envTerm :: Map VName TermBinding
envTerm = forall a. Monoid a => a
mempty,
      envType :: Map VName TypeBinding
envType = forall a b k. (a -> b) -> Map k a -> Map k b
M.map StructType -> TypeBinding
tbind Map VName StructType
m
    }
  where
    tbind :: StructType -> TypeBinding
tbind = Liftedness -> [TypeParam] -> RetTypeBase Size () -> TypeBinding
T.TypeAbbr Liftedness
Unlifted [] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []

i64Env :: M.Map VName Int64 -> Env
i64Env :: Sizes -> Env
i64Env = Map VName (Maybe BoundV, Value) -> Env
valEnv forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b k. (a -> b) -> Map k a -> Map k b
M.map forall {m :: * -> *}. Int64 -> (Maybe BoundV, Value m)
f
  where
    f :: Int64 -> (Maybe BoundV, Value m)
f Int64
x =
      ( forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
        forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
x
      )

instance Show InterpreterError where
  show :: InterpreterError -> [Char]
show (InterpreterError Text
s) = Text -> [Char]
T.unpack Text
s

bad :: SrcLoc -> Env -> T.Text -> EvalM a
bad :: forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env Text
s = forall a. SrcLoc -> Env -> EvalM a -> EvalM a
stacking SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$ do
  [Text]
ss <- forall a b. (a -> b) -> [a] -> [b]
map (forall a. Located a => a -> Text
locText forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Located a => a -> SrcLoc
srclocOf) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM [Loc]
stacktrace
  forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. InterpreterError -> ExtOp a
ExtOpError forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> InterpreterError
InterpreterError forall a b. (a -> b) -> a -> b
$
    Text
"Error at\n" forall a. Semigroup a => a -> a -> a
<> Int -> [Text] -> Text
prettyStacktrace Int
0 [Text]
ss forall a. Semigroup a => a -> a -> a
<> Text
s

trace :: T.Text -> Value -> EvalM ()
trace :: Text -> Value -> EvalM ()
trace Text
w Value
v = do
  forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF forall a b. (a -> b) -> a -> b
$ forall a. Text -> Doc () -> a -> ExtOp a
ExtOpTrace Text
w (forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
v) ()

typeCheckerEnv :: Env -> T.Env
typeCheckerEnv :: Env -> Env
typeCheckerEnv Env
env =
  -- FIXME: some shadowing issues are probably not right here.
  let valMap :: TermBinding -> Maybe BoundV
valMap (TermValue (Just BoundV
t) Value
_) = forall a. a -> Maybe a
Just BoundV
t
      valMap TermBinding
_ = forall a. Maybe a
Nothing
      vtable :: Map VName BoundV
vtable = forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe TermBinding -> Maybe BoundV
valMap forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm Env
env
      nameMap :: VName -> Maybe ((Namespace, Name), QualName VName)
nameMap VName
k
        | VName
k forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName BoundV
vtable = forall a. a -> Maybe a
Just ((Namespace
T.Term, VName -> Name
baseName VName
k), forall v. v -> QualName v
qualName VName
k)
        | Bool
otherwise = forall a. Maybe a
Nothing
   in forall a. Monoid a => a
mempty
        { envNameMap :: NameMap
T.envNameMap = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe ((Namespace, Name), QualName VName)
nameMap forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm Env
env,
          envVtable :: Map VName BoundV
T.envVtable = Map VName BoundV
vtable
        }

break :: Env -> Loc -> EvalM ()
break :: Env -> Loc -> EvalM ()
break Env
env Loc
loc = do
  Map ImportName Env
imports <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> b
snd
  NonEmpty StackFrame
backtrace <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Loc -> Ctx -> StackFrame
StackFrame Loc
loc (Env -> Map ImportName Env -> Ctx
Ctx Env
env Map ImportName Env
imports) NE.:|) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)
  forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF forall a b. (a -> b) -> a -> b
$ forall a. Loc -> BreakReason -> NonEmpty StackFrame -> a -> ExtOp a
ExtOpBreak Loc
loc BreakReason
BreakPoint NonEmpty StackFrame
backtrace ()

fromArray :: Value -> (ValueShape, [Value])
fromArray :: Value -> (ValueShape, [Value])
fromArray (ValueArray ValueShape
shape Array Int Value
as) = (ValueShape
shape, forall i e. Array i e -> [e]
elems Array Int Value
as)
fromArray Value
v = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Expected array value, but found: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
v

apply :: SrcLoc -> Env -> Value -> Value -> EvalM Value
apply :: SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env (ValueFun Value -> EvalM Value
f) Value
v = forall a. SrcLoc -> Env -> EvalM a -> EvalM a
stacking SrcLoc
loc Env
env (Value -> EvalM Value
f Value
v)
apply SrcLoc
_ Env
_ Value
f Value
_ = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot apply non-function: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
f

apply2 :: SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 :: SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 SrcLoc
loc Env
env Value
f Value
x Value
y = forall a. SrcLoc -> Env -> EvalM a -> EvalM a
stacking SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$ do
  Value
f' <- SrcLoc -> Env -> Value -> Value -> EvalM Value
apply forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f Value
x
  SrcLoc -> Env -> Value -> Value -> EvalM Value
apply forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f' Value
y

matchPat :: Env -> Pat -> Value -> EvalM Env
matchPat :: Env -> Pat -> Value -> EvalM Env
matchPat Env
env Pat
p Value
v = do
  Maybe Env
m <- forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT forall a b. (a -> b) -> a -> b
$ Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env Pat
p Value
v
  case Maybe Env
m of
    Maybe Env
Nothing -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"matchPat: missing case for " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString Pat
p forall a. [a] -> [a] -> [a]
++ [Char]
" and " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
v
    Just Env
env' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env'

patternMatch :: Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch :: Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env (Id VName
v (Info PatType
t) SrcLoc
_) Value
val =
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      Map VName (Maybe BoundV, Value) -> Env
valEnv (forall k a. k -> a -> Map k a
M.singleton VName
v (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t, Value
val)) forall a. Semigroup a => a -> a -> a
<> Env
env
patternMatch Env
env Wildcard {} Value
_ =
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env
patternMatch Env
env (TuplePat [Pat]
ps SrcLoc
_) (ValueRecord Map Name Value
vs) =
  forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' (Pat
p, Value
v) -> Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env' Pat
p Value
v) Env
env forall a b. (a -> b) -> a -> b
$
    forall a b. [a] -> [b] -> [(a, b)]
zip [Pat]
ps (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
sortFields Map Name Value
vs)
patternMatch Env
env (RecordPat [(Name, Pat)]
ps SrcLoc
_) (ValueRecord Map Name Value
vs) =
  forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' (Pat
p, Value
v) -> Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env' Pat
p Value
v) Env
env forall a b. (a -> b) -> a -> b
$
    forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, Pat)]
ps) Map Name Value
vs
patternMatch Env
env (PatParens Pat
p SrcLoc
_) Value
v = Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env Pat
p Value
v
patternMatch Env
env (PatAscription Pat
p TypeExp Info VName
_ SrcLoc
_) Value
v =
  Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env Pat
p Value
v
patternMatch Env
env (PatLit PatLit
l Info PatType
t SrcLoc
_) Value
v = do
  Value
l' <- case PatLit
l of
    PatLitInt Integer
x -> forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ Env -> Exp -> EvalM Value
eval Env
env forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
Integer -> f PatType -> SrcLoc -> ExpBase f vn
IntLit Integer
x Info PatType
t forall a. Monoid a => a
mempty
    PatLitFloat Double
x -> forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ Env -> Exp -> EvalM Value
eval Env
env forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
Double -> f PatType -> SrcLoc -> ExpBase f vn
FloatLit Double
x Info PatType
t forall a. Monoid a => a
mempty
    PatLitPrim PrimValue
lv -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
lv
  if Value
v forall a. Eq a => a -> a -> Bool
== Value
l'
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env
    else forall (m :: * -> *) a. MonadPlus m => m a
mzero
patternMatch Env
env (PatConstr Name
n Info PatType
_ [Pat]
ps SrcLoc
_) (ValueSum ValueShape
_ Name
n' [Value]
vs)
  | Name
n forall a. Eq a => a -> a -> Bool
== Name
n' =
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' (Pat
p, Value
v) -> Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env' Pat
p Value
v) Env
env forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Pat]
ps [Value]
vs
patternMatch Env
_ Pat
_ Value
_ = forall (m :: * -> *) a. MonadPlus m => m a
mzero

data Indexing
  = IndexingFix Int64
  | IndexingSlice (Maybe Int64) (Maybe Int64) (Maybe Int64)

instance Pretty Indexing where
  pretty :: forall ann. Indexing -> Doc ann
pretty (IndexingFix Int64
i) = forall a ann. Pretty a => a -> Doc ann
pretty Int64
i
  pretty (IndexingSlice Maybe Int64
i Maybe Int64
j (Just Int64
s)) =
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall a ann. Pretty a => a -> Doc ann
pretty Maybe Int64
i
      forall a. Semigroup a => a -> a -> a
<> Doc ann
":"
      forall a. Semigroup a => a -> a -> a
<> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall a ann. Pretty a => a -> Doc ann
pretty Maybe Int64
j
      forall a. Semigroup a => a -> a -> a
<> Doc ann
":"
      forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Int64
s
  pretty (IndexingSlice Maybe Int64
i (Just Int64
j) Maybe Int64
s) =
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall a ann. Pretty a => a -> Doc ann
pretty Maybe Int64
i
      forall a. Semigroup a => a -> a -> a
<> Doc ann
":"
      forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Int64
j
      forall a. Semigroup a => a -> a -> a
<> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty ((Doc ann
":" <>) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a ann. Pretty a => a -> Doc ann
pretty) Maybe Int64
s
  pretty (IndexingSlice Maybe Int64
i Maybe Int64
Nothing Maybe Int64
Nothing) =
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall a ann. Pretty a => a -> Doc ann
pretty Maybe Int64
i forall a. Semigroup a => a -> a -> a
<> Doc ann
":"

indexesFor ::
  Maybe Int64 ->
  Maybe Int64 ->
  Maybe Int64 ->
  Int64 ->
  Maybe [Int]
indexesFor :: Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride Int64
n
  | (Int64
start', Int64
end', Int64
stride') <- (Int64, Int64, Int64)
slice,
    Int64
end' forall a. Eq a => a -> a -> Bool
== Int64
start' Bool -> Bool -> Bool
|| forall p. (Eq p, Num p) => p -> p
signum' (Int64
end' forall a. Num a => a -> a -> a
- Int64
start') forall a. Eq a => a -> a -> Bool
== forall p. (Eq p, Num p) => p -> p
signum' Int64
stride',
    Int64
stride' forall a. Eq a => a -> a -> Bool
/= Int64
0,
    [Int64]
is <- [Int64
start', Int64
start' forall a. Num a => a -> a -> a
+ Int64
stride' .. Int64
end' forall a. Num a => a -> a -> a
- forall a. Num a => a -> a
signum Int64
stride'],
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Int64 -> Bool
inBounds [Int64]
is =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int64]
is
  | Bool
otherwise =
      forall a. Maybe a
Nothing
  where
    inBounds :: Int64 -> Bool
inBounds Int64
i = Int64
i forall a. Ord a => a -> a -> Bool
>= Int64
0 Bool -> Bool -> Bool
&& Int64
i forall a. Ord a => a -> a -> Bool
< Int64
n

    slice :: (Int64, Int64, Int64)
slice =
      case (Maybe Int64
start, Maybe Int64
end, Maybe Int64
stride) of
        (Just Int64
start', Maybe Int64
_, Maybe Int64
_) ->
          let end' :: Int64
end' = forall a. a -> Maybe a -> a
fromMaybe Int64
n Maybe Int64
end
           in (Int64
start', Int64
end', forall a. a -> Maybe a -> a
fromMaybe Int64
1 Maybe Int64
stride)
        (Maybe Int64
Nothing, Just Int64
end', Maybe Int64
_) ->
          let start' :: Int64
start' = Int64
0
           in (Int64
start', Int64
end', forall a. a -> Maybe a -> a
fromMaybe Int64
1 Maybe Int64
stride)
        (Maybe Int64
Nothing, Maybe Int64
Nothing, Just Int64
stride') ->
          ( if Int64
stride' forall a. Ord a => a -> a -> Bool
> Int64
0 then Int64
0 else Int64
n forall a. Num a => a -> a -> a
- Int64
1,
            if Int64
stride' forall a. Ord a => a -> a -> Bool
> Int64
0 then Int64
n else -Int64
1,
            Int64
stride'
          )
        (Maybe Int64
Nothing, Maybe Int64
Nothing, Maybe Int64
Nothing) ->
          (Int64
0, Int64
n, Int64
1)

-- | 'signum', but with 0 as 1.
signum' :: (Eq p, Num p) => p -> p
signum' :: forall p. (Eq p, Num p) => p -> p
signum' p
0 = p
1
signum' p
x = forall a. Num a => a -> a
signum p
x

indexShape :: [Indexing] -> ValueShape -> ValueShape
indexShape :: [Indexing] -> ValueShape -> ValueShape
indexShape (IndexingFix {} : [Indexing]
is) (ShapeDim Int64
_ ValueShape
shape) =
  [Indexing] -> ValueShape -> ValueShape
indexShape [Indexing]
is ValueShape
shape
indexShape (IndexingSlice Maybe Int64
start Maybe Int64
end Maybe Int64
stride : [Indexing]
is) (ShapeDim Int64
d ValueShape
shape) =
  forall d. d -> Shape d -> Shape d
ShapeDim Int64
n forall a b. (a -> b) -> a -> b
$ [Indexing] -> ValueShape -> ValueShape
indexShape [Indexing]
is ValueShape
shape
  where
    n :: Int64
n = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int64
0 forall i a. Num i => [a] -> i
genericLength forall a b. (a -> b) -> a -> b
$ Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride Int64
d
indexShape [Indexing]
_ ValueShape
shape =
  ValueShape
shape

indexArray :: [Indexing] -> Value -> Maybe Value
indexArray :: [Indexing] -> Value -> Maybe Value
indexArray (IndexingFix Int64
i : [Indexing]
is) (ValueArray ValueShape
_ Array Int Value
arr)
  | Int64
i forall a. Ord a => a -> a -> Bool
>= Int64
0,
    Int64
i forall a. Ord a => a -> a -> Bool
< Int64
n =
      [Indexing] -> Value -> Maybe Value
indexArray [Indexing]
is forall a b. (a -> b) -> a -> b
$ Array Int Value
arr forall i e. Ix i => Array i e -> i -> e
! forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i
  | Bool
otherwise =
      forall a. Maybe a
Nothing
  where
    n :: Int64
n = forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr
indexArray (IndexingSlice Maybe Int64
start Maybe Int64
end Maybe Int64
stride : [Indexing]
is) (ValueArray (ShapeDim Int64
_ ValueShape
rowshape) Array Int Value
arr) = do
  [Int]
js <- Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride forall a b. (a -> b) -> a -> b
$ forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr
  forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ([Indexing] -> ValueShape -> ValueShape
indexShape [Indexing]
is ValueShape
rowshape) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Indexing] -> Value -> Maybe Value
indexArray [Indexing]
is forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array Int Value
arr !)) [Int]
js
indexArray [Indexing]
_ Value
v = forall a. a -> Maybe a
Just Value
v

writeArray :: [Indexing] -> Value -> Value -> Maybe Value
writeArray :: [Indexing] -> Value -> Value -> Maybe Value
writeArray [Indexing]
slice Value
x Value
y = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (\Value
_ Value
y' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
y') [Indexing]
slice Value
x Value
y

updateArray ::
  Monad m =>
  (Value -> Value -> m Value) ->
  [Indexing] ->
  Value ->
  Value ->
  m (Maybe Value)
updateArray :: forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray Value -> Value -> m Value
f (IndexingFix Int64
i : [Indexing]
is) (ValueArray ValueShape
shape Array Int Value
arr) Value
v
  | Int64
i forall a. Ord a => a -> a -> Bool
>= Int64
0,
    Int64
i forall a. Ord a => a -> a -> Bool
< Int64
n = do
      Maybe Value
v' <- forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray Value -> Value -> m Value
f [Indexing]
is (Array Int Value
arr forall i e. Ix i => Array i e -> i -> e
! Int
i') Value
v
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ do
        Value
v'' <- Maybe Value
v'
        forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
shape forall a b. (a -> b) -> a -> b
$ Array Int Value
arr forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Int
i', Value
v'')]
  | Bool
otherwise =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
  where
    n :: Int64
n = forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr
    i' :: Int
i' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i
updateArray Value -> Value -> m Value
f (IndexingSlice Maybe Int64
start Maybe Int64
end Maybe Int64
stride : [Indexing]
is) (ValueArray ValueShape
shape Array Int Value
arr) (ValueArray ValueShape
_ Array Int Value
v)
  | Just [Int]
arr_is <- Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride forall a b. (a -> b) -> a -> b
$ forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr,
    forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
arr_is forall a. Eq a => a -> a -> Bool
== forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
v = do
      let update :: Maybe (Array Int Value)
-> (Int, Value) -> m (Maybe (Array Int Value))
update (Just Array Int Value
arr') (Int
i, Value
v') = do
            Maybe Value
x <- forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray Value -> Value -> m Value
f [Indexing]
is (Array Int Value
arr forall i e. Ix i => Array i e -> i -> e
! Int
i) Value
v'
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ do
              Value
x' <- Maybe Value
x
              forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Array Int Value
arr' forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Int
i, Value
x')]
          update Maybe (Array Int Value)
Nothing (Int, Value)
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
shape)) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Maybe (Array Int Value)
-> (Int, Value) -> m (Maybe (Array Int Value))
update (forall a. a -> Maybe a
Just Array Int Value
arr) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
arr_is forall a b. (a -> b) -> a -> b
$ forall i e. Array i e -> [e]
elems Array Int Value
v
  | Bool
otherwise =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
updateArray Value -> Value -> m Value
f [Indexing]
_ Value
x Value
y = forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Value -> m Value
f Value
x Value
y

evalDimIndex :: Env -> DimIndex -> EvalM Indexing
evalDimIndex :: Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env (DimFix Exp
x) =
  Int64 -> Indexing
IndexingFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
x
evalDimIndex Env
env (DimSlice Maybe Exp
start Maybe Exp
end Maybe Exp
stride) =
  Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Indexing
IndexingSlice
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Int64
asInt64 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Exp -> EvalM Value
eval Env
env) Maybe Exp
start
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Int64
asInt64 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Exp -> EvalM Value
eval Env
env) Maybe Exp
end
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Int64
asInt64 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Exp -> EvalM Value
eval Env
env) Maybe Exp
stride

evalIndex :: SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value
evalIndex :: SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value
evalIndex SrcLoc
loc Env
env [Indexing]
is Value
arr = do
  let oob :: EvalM a
oob =
        forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$
          Text
"Index ["
            forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Text
prettyText [Indexing]
is)
            forall a. Semigroup a => a -> a -> a
<> Text
"] out of bounds for array of shape "
            forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
arr)
            forall a. Semigroup a => a -> a -> a
<> Text
"."
  forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall {a}. EvalM a
oob forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Maybe Value
indexArray [Indexing]
is Value
arr

-- | Expand type based on information that was not available at
-- type-checking time (the structure of abstract types).
expandType :: Env -> StructType -> StructType
expandType :: Env -> StructType -> StructType
expandType Env
_ (Scalar (Prim PrimType
pt)) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
pt
expandType Env
env (Scalar (Record Map Name StructType
fs)) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Env -> StructType -> StructType
expandType Env
env) Map Name StructType
fs
expandType Env
env (Scalar (Arrow () PName
p Diet
d StructType
t1 (RetType [VName]
dims StructType
t2))) =
  forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> Diet
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
p Diet
d (Env -> StructType -> StructType
expandType Env
env StructType
t1) (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (Env -> StructType -> StructType
expandType Env
env StructType
t2))
expandType Env
env t :: StructType
t@(Array ()
_ Uniqueness
u Shape Size
shape ScalarTypeBase Size ()
_) =
  let et :: StructType
et = forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray (forall dim. Shape dim -> Int
shapeRank Shape Size
shape) StructType
t
      et' :: StructType
et' = Env -> StructType -> StructType
expandType Env
env StructType
et
   in forall as dim.
Monoid as =>
Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOf Uniqueness
u Shape Size
shape StructType
et'
expandType Env
env t :: StructType
t@(Scalar (TypeVar () Uniqueness
_ QualName VName
tn [TypeArg Size]
args)) =
  case QualName VName -> Env -> Maybe TypeBinding
lookupType QualName VName
tn Env
env of
    Just (T.TypeAbbr Liftedness
_ [TypeParam]
ps (RetType [VName]
_ StructType
t')) ->
      let (Map VName Size
substs, Map VName TypeBinding
types) = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {k}.
Ord k =>
TypeParamBase k -> TypeArg Size -> (Map k Size, Map k TypeBinding)
matchPtoA [TypeParam]
ps [TypeArg Size]
args
          onDim :: Size -> Size
onDim (NamedSize QualName VName
v) = forall a. a -> Maybe a -> a
fromMaybe (QualName VName -> Size
NamedSize QualName VName
v) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName Size
substs
          onDim Size
d = Size
d
       in if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParam]
ps
            then forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim StructType
t'
            else Env -> StructType -> StructType
expandType (Map VName TermBinding -> Map VName TypeBinding -> Env
Env forall a. Monoid a => a
mempty Map VName TypeBinding
types forall a. Semigroup a => a -> a -> a
<> Env
env) forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim StructType
t'
    Maybe TypeBinding
Nothing -> StructType
t
  where
    matchPtoA :: TypeParamBase k -> TypeArg Size -> (Map k Size, Map k TypeBinding)
matchPtoA (TypeParamDim k
p SrcLoc
_) (TypeArgDim (NamedSize QualName VName
qv) SrcLoc
_) =
      (forall k a. k -> a -> Map k a
M.singleton k
p forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize QualName VName
qv, forall a. Monoid a => a
mempty)
    matchPtoA (TypeParamDim k
p SrcLoc
_) (TypeArgDim (ConstSize Int64
k) SrcLoc
_) =
      (forall k a. k -> a -> Map k a
M.singleton k
p forall a b. (a -> b) -> a -> b
$ Int64 -> Size
ConstSize Int64
k, forall a. Monoid a => a
mempty)
    matchPtoA (TypeParamType Liftedness
l k
p SrcLoc
_) (TypeArgType StructType
t' SrcLoc
_) =
      let t'' :: StructType
t'' = Env -> StructType -> StructType
expandType Env
env StructType
t'
       in (forall a. Monoid a => a
mempty, forall k a. k -> a -> Map k a
M.singleton k
p forall a b. (a -> b) -> a -> b
$ Liftedness -> [TypeParam] -> RetTypeBase Size () -> TypeBinding
T.TypeAbbr Liftedness
l [] forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] StructType
t'')
    matchPtoA TypeParamBase k
_ TypeArg Size
_ = forall a. Monoid a => a
mempty
expandType Env
env (Scalar (Sum Map Name [StructType]
cs)) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (Env -> StructType -> StructType
expandType Env
env) Map Name [StructType]
cs

-- | First expand type abbreviations, then evaluate all possible
-- sizes.
evalType :: Env -> StructType -> EvalM StructType
evalType :: Env -> StructType -> EvalM StructType
evalType Env
outer_env StructType
t = do
  Env
size_env <- EvalM Env
extSizeEnv
  let env :: Env
env = Env
size_env forall a. Semigroup a => a -> a -> a
<> Env
outer_env
      evalDim :: Size -> Size
evalDim (NamedSize QualName VName
qn)
        | Just (TermValue Maybe BoundV
_ (ValuePrim (SignedValue (Int64Value Int64
x)))) <-
            QualName VName -> Env -> Maybe TermBinding
lookupVar QualName VName
qn Env
env =
            Int64 -> Size
ConstSize forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
x
      evalDim Size
d = Size
d
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
evalDim forall a b. (a -> b) -> a -> b
$ Env -> StructType -> StructType
expandType Env
env StructType
t

evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value
evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv StructType
t =
  case QualName VName -> Env -> Maybe TermBinding
lookupVar QualName VName
qv Env
env of
    Just (TermPoly Maybe BoundV
_ StructType -> EvalM Value
v) -> StructType -> EvalM Value
v forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> StructType -> EvalM StructType
evalType Env
env StructType
t
    Just (TermValue Maybe BoundV
_ Value
v) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
    Maybe TermBinding
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"\"" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString QualName VName
qv forall a. Semigroup a => a -> a -> a
<> [Char]
"\" is not bound to a value."

typeValueShape :: Env -> StructType -> EvalM ValueShape
typeValueShape :: Env -> StructType -> EvalM ValueShape
typeValueShape Env
env StructType
t = do
  StructType
t' <- Env -> StructType -> EvalM StructType
evalType Env
env StructType
t
  case forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall {a}. Num a => Size -> Maybe a
dim forall a b. (a -> b) -> a -> b
$ forall d. TypeBase d () -> Shape d
typeShape StructType
t' of
    Maybe ValueShape
Nothing -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"typeValueShape: failed to fully evaluate type " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString StructType
t'
    Just ValueShape
shape -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ValueShape
shape
  where
    dim :: Size -> Maybe a
dim (ConstSize Int64
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
x
    dim Size
_ = forall a. Maybe a
Nothing

evalFunction :: Env -> [VName] -> [Pat] -> Exp -> StructType -> EvalM Value
-- We treat zero-parameter lambdas as simply an expression to
-- evaluate immediately.  Note that this is *not* the same as a lambda
-- that takes an empty tuple '()' as argument!  Zero-parameter lambdas
-- can never occur in a well-formed Futhark program, but they are
-- convenient in the interpreter.
evalFunction :: Env -> [VName] -> [Pat] -> Exp -> StructType -> EvalM Value
evalFunction Env
env [VName]
_ [] Exp
body StructType
rettype =
  -- Eta-expand the rest to make any sizes visible.
  forall {as}. [Value] -> Env -> TypeBase Size as -> EvalM Value
etaExpand [] Env
env StructType
rettype
  where
    etaExpand :: [Value] -> Env -> TypeBase Size as -> EvalM Value
etaExpand [Value]
vs Env
env' (Scalar (Arrow as
_ PName
_ Diet
_ StructType
pt (RetType [VName]
_ TypeBase Size as
rt))) =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v -> do
        Env
env'' <- Env -> Pat -> Value -> EvalM Env
matchPat Env
env' (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct StructType
pt) forall a. IsLocation a => a
noLoc) Value
v
        [Value] -> Env -> TypeBase Size as -> EvalM Value
etaExpand (Value
v forall a. a -> [a] -> [a]
: [Value]
vs) Env
env'' TypeBase Size as
rt
    etaExpand [Value]
vs Env
env' TypeBase Size as
_ = do
      Value
f <- forall a. EvalM a -> EvalM a
localExts forall a b. (a -> b) -> a -> b
$ Env -> Exp -> EvalM Value
eval Env
env' Exp
body
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> EvalM Value
apply forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty) Value
f forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse [Value]
vs
evalFunction Env
env [VName]
missing_sizes (Pat
p : [Pat]
ps) Exp
body StructType
rettype =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v -> do
    Env
env' <- Env -> Pat -> Value -> EvalM Env
matchPat Env
env Pat
p Value
v
    -- Fix up the last sizes, if any.
    let p_t :: StructType
p_t = Env -> StructType -> StructType
expandType Env
env forall a b. (a -> b) -> a -> b
$ Pat -> StructType
patternStructType Pat
p
        env'' :: Env
env''
          | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
missing_sizes =
              Env
env'
          | Bool
otherwise =
              Env
env' forall a. Semigroup a => a -> a -> a
<> Sizes -> Env
i64Env ([VName] -> StructType -> ValueShape -> Sizes
resolveExistentials [VName]
missing_sizes StructType
p_t (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v))
    Env -> [VName] -> [Pat] -> Exp -> StructType -> EvalM Value
evalFunction Env
env'' [VName]
missing_sizes [Pat]
ps Exp
body StructType
rettype

evalFunctionBinding ::
  Env ->
  [TypeParam] ->
  [Pat] ->
  StructRetType ->
  Exp ->
  EvalM TermBinding
evalFunctionBinding :: Env
-> [TypeParam]
-> [Pat]
-> RetTypeBase Size ()
-> Exp
-> EvalM TermBinding
evalFunctionBinding Env
env [TypeParam]
tparams [Pat]
ps RetTypeBase Size ()
ret Exp
fbody = do
  let arrow :: (PName, Diet, TypeBase dim ())
-> TypeBase dim () -> TypeBase dim ()
arrow (PName
xp, Diet
d, TypeBase dim ()
xt) TypeBase dim ()
yt = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as.
as
-> PName
-> Diet
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
xp Diet
d TypeBase dim ()
xt forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase dim ()
yt
      ftype :: StructType
ftype = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall {dim}.
(PName, Diet, TypeBase dim ())
-> TypeBase dim () -> TypeBase dim ()
arrow forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat -> (PName, Diet, StructType)
patternParam) (forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase Size ()
ret) [Pat]
ps
      retext :: [VName]
retext = case [Pat]
ps of
        [] -> forall dim as. RetTypeBase dim as -> [VName]
retDims RetTypeBase Size ()
ret
        [Pat]
_ -> []

  -- Distinguish polymorphic and non-polymorphic bindings here.
  if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParam]
tparams
    then
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe BoundV -> Value -> TermBinding
TermValue (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] StructType
ftype))
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall als.
Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value
returned Env
env (forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase Size ()
ret) [VName]
retext
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> [VName] -> [Pat] -> Exp -> StructType -> EvalM Value
evalFunction Env
env [] [Pat]
ps Exp
fbody (forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase Size ()
ret)
    else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe BoundV -> (StructType -> EvalM Value) -> TermBinding
TermPoly (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] StructType
ftype) forall a b. (a -> b) -> a -> b
$ \StructType
ftype' -> do
      let tparam_names :: [VName]
tparam_names = forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
tparams
          env' :: Env
env' = [VName] -> StructType -> StructType -> Env
resolveTypeParams [VName]
tparam_names StructType
ftype StructType
ftype' forall a. Semigroup a => a -> a -> a
<> Env
env

          -- In some cases (abstract lifted types) there may be
          -- missing sizes that were not fixed by the type
          -- instantiation.  These will have to be set by looking
          -- at the actual function arguments.
          missing_sizes :: [VName]
missing_sizes =
            forall a. (a -> Bool) -> [a] -> [a]
filter (forall k a. Ord k => k -> Map k a -> Bool
`M.notMember` Env -> Map VName TermBinding
envTerm Env
env') forall a b. (a -> b) -> a -> b
$
              forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName (forall a. (a -> Bool) -> [a] -> [a]
filter forall vn. TypeParamBase vn -> Bool
isSizeParam [TypeParam]
tparams)
      forall als.
Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value
returned Env
env (forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase Size ()
ret) [VName]
retext
        forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> [VName] -> [Pat] -> Exp -> StructType -> EvalM Value
evalFunction Env
env' [VName]
missing_sizes [Pat]
ps Exp
fbody (forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase Size ()
ret)

evalArg :: Env -> Exp -> Maybe VName -> EvalM Value
evalArg :: Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
e Maybe VName
ext = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  case Maybe VName
ext of
    Just VName
ext' -> VName -> Int64 -> EvalM ()
putExtSize VName
ext' forall a b. (a -> b) -> a -> b
$ Value -> Int64
asInt64 Value
v
    Maybe VName
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v

returned :: Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value
returned :: forall als.
Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value
returned Env
_ TypeBase Size als
_ [] Value
v = forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
returned Env
env TypeBase Size als
ret [VName]
retext Value
v = do
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry VName -> Int64 -> EvalM ()
putExtSize) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$
    [VName] -> StructType -> ValueShape -> Sizes
resolveExistentials [VName]
retext (Env -> StructType -> StructType
expandType Env
env forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase Size als
ret) forall a b. (a -> b) -> a -> b
$
      forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v

evalAppExp :: Env -> StructType -> AppExp -> EvalM Value
evalAppExp :: Env -> StructType -> AppExp -> EvalM Value
evalAppExp Env
env StructType
_ (Range Exp
start Maybe Exp
maybe_second Inclusiveness Exp
end SrcLoc
loc) = do
  Integer
start' <- Value -> Integer
asInteger forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
start
  Maybe Integer
maybe_second' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Integer
asInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Exp -> EvalM Value
eval Env
env) Maybe Exp
maybe_second
  Inclusiveness Integer
end' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Integer
asInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Exp -> EvalM Value
eval Env
env) Inclusiveness Exp
end

  let (Integer
end_adj, Integer
step, Bool
ok) =
        case (Inclusiveness Integer
end', Maybe Integer
maybe_second') of
          (DownToExclusive Integer
end'', Maybe Integer
Nothing) ->
            (Integer
end'' forall a. Num a => a -> a -> a
+ Integer
1, -Integer
1, Integer
start' forall a. Ord a => a -> a -> Bool
>= Integer
end'')
          (DownToExclusive Integer
end'', Just Integer
second') ->
            (Integer
end'' forall a. Num a => a -> a -> a
+ Integer
1, Integer
second' forall a. Num a => a -> a -> a
- Integer
start', Integer
start' forall a. Ord a => a -> a -> Bool
>= Integer
end'' Bool -> Bool -> Bool
&& Integer
second' forall a. Ord a => a -> a -> Bool
< Integer
start')
          (ToInclusive Integer
end'', Maybe Integer
Nothing) ->
            (Integer
end'', Integer
1, Integer
start' forall a. Ord a => a -> a -> Bool
<= Integer
end'')
          (ToInclusive Integer
end'', Just Integer
second')
            | Integer
second' forall a. Ord a => a -> a -> Bool
> Integer
start' ->
                (Integer
end'', Integer
second' forall a. Num a => a -> a -> a
- Integer
start', Integer
start' forall a. Ord a => a -> a -> Bool
<= Integer
end'')
            | Bool
otherwise ->
                (Integer
end'', Integer
second' forall a. Num a => a -> a -> a
- Integer
start', Integer
start' forall a. Ord a => a -> a -> Bool
>= Integer
end'' Bool -> Bool -> Bool
&& Integer
second' forall a. Eq a => a -> a -> Bool
/= Integer
start')
          (UpToExclusive Integer
x, Maybe Integer
Nothing) ->
            (Integer
x forall a. Num a => a -> a -> a
- Integer
1, Integer
1, Integer
start' forall a. Ord a => a -> a -> Bool
<= Integer
x)
          (UpToExclusive Integer
x, Just Integer
second') ->
            (Integer
x forall a. Num a => a -> a -> a
- Integer
1, Integer
second' forall a. Num a => a -> a -> a
- Integer
start', Integer
start' forall a. Ord a => a -> a -> Bool
<= Integer
x Bool -> Bool -> Bool
&& Integer
second' forall a. Ord a => a -> a -> Bool
> Integer
start')

  if Bool
ok
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' forall d. Shape d
ShapeLeaf forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall {m :: * -> *}. Integer -> Value m
toInt [Integer
start', Integer
start' forall a. Num a => a -> a -> a
+ Integer
step .. Integer
end_adj]
    else forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$ forall {a} {a} {a}.
(Pretty a, Pretty a, Pretty a) =>
a -> Maybe a -> Inclusiveness a -> Text
badRange Integer
start' Maybe Integer
maybe_second' Inclusiveness Integer
end'
  where
    toInt :: Integer -> Value m
toInt =
      case Exp -> PatType
typeOf Exp
start of
        Scalar (Prim (Signed IntType
t')) ->
          forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
SignedValue forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t'
        Scalar (Prim (Unsigned IntType
t')) ->
          forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
UnsignedValue forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t'
        PatType
t -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Nonsensical range type: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show PatType
t

    badRange :: a -> Maybe a -> Inclusiveness a -> Text
badRange a
start' Maybe a
maybe_second' Inclusiveness a
end' =
      Text
"Range "
        forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
start'
        forall a. Semigroup a => a -> a -> a
<> ( case Maybe a
maybe_second' of
               Maybe a
Nothing -> Text
""
               Just a
second' -> Text
".." forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
second'
           )
        forall a. Semigroup a => a -> a -> a
<> ( case Inclusiveness a
end' of
               DownToExclusive a
x -> Text
"..>" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
x
               ToInclusive a
x -> Text
"..." forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
x
               UpToExclusive a
x -> Text
"..<" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
x
           )
        forall a. Semigroup a => a -> a -> a
<> Text
" is invalid."
evalAppExp Env
env StructType
t (Coerce Exp
e TypeExp Info VName
te SrcLoc
loc) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  case Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape (StructType -> Shape (Maybe Int64)
structTypeShape StructType
t) (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v) of
    Just ValueShape
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
    Maybe ValueShape
Nothing ->
      forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
        Doc Any
"Value `"
          forall a. Semigroup a => a -> a -> a
<> forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
v
          forall a. Semigroup a => a -> a -> a
<> Doc Any
"` of shape `"
          forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v)
          forall a. Semigroup a => a -> a -> a
<> Doc Any
"` cannot match shape of type `"
          forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty TypeExp Info VName
te
          forall a. Semigroup a => a -> a -> a
<> Doc Any
"` (`"
          forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty StructType
t
          forall a. Semigroup a => a -> a -> a
<> Doc Any
"`)"
evalAppExp Env
env StructType
_ (LetPat [SizeBinder VName]
sizes Pat
p Exp
e Exp
body SrcLoc
_) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  Env
env' <- Env -> Pat -> Value -> EvalM Env
matchPat Env
env Pat
p Value
v
  let p_t :: StructType
p_t = Env -> StructType -> StructType
expandType Env
env forall a b. (a -> b) -> a -> b
$ Pat -> StructType
patternStructType Pat
p
      v_s :: ValueShape
v_s = forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v
      env'' :: Env
env'' = Env
env' forall a. Semigroup a => a -> a -> a
<> Sizes -> Env
i64Env ([VName] -> StructType -> ValueShape -> Sizes
resolveExistentials (forall a b. (a -> b) -> [a] -> [b]
map forall vn. SizeBinder vn -> vn
sizeName [SizeBinder VName]
sizes) StructType
p_t ValueShape
v_s)
  Env -> Exp -> EvalM Value
eval Env
env'' Exp
body
evalAppExp Env
env StructType
_ (LetFun VName
f ([TypeParam]
tparams, [Pat]
ps, Maybe (TypeExp Info VName)
_, Info RetTypeBase Size ()
ret, Exp
fbody) Exp
body SrcLoc
_) = do
  TermBinding
binding <- Env
-> [TypeParam]
-> [Pat]
-> RetTypeBase Size ()
-> Exp
-> EvalM TermBinding
evalFunctionBinding Env
env [TypeParam]
tparams [Pat]
ps RetTypeBase Size ()
ret Exp
fbody
  Env -> Exp -> EvalM Value
eval (Env
env {envTerm :: Map VName TermBinding
envTerm = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
f TermBinding
binding forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm Env
env}) Exp
body
evalAppExp
  Env
env
  StructType
_
  (BinOp (QualName VName
op, SrcLoc
_) Info PatType
op_t (Exp
x, Info (StructType
_, Maybe VName
xext)) (Exp
y, Info (StructType
_, Maybe VName
yext)) SrcLoc
loc)
    | VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
op) forall a. Eq a => a -> a -> Bool
== [Char]
"&&" = do
        Bool
x' <- Value -> Bool
asBool forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
x
        if Bool
x'
          then Env -> Exp -> EvalM Value
eval Env
env Exp
y
          else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
False
    | VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
op) forall a. Eq a => a -> a -> Bool
== [Char]
"||" = do
        Bool
x' <- Value -> Bool
asBool forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
x
        if Bool
x'
          then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True
          else Env -> Exp -> EvalM Value
eval Env
env Exp
y
    | Bool
otherwise = do
        Value
op' <- Env -> Exp -> EvalM Value
eval Env
env forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
op Info PatType
op_t SrcLoc
loc
        Value
x' <- Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
x Maybe VName
xext
        Value
y' <- Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
y Maybe VName
yext
        SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 SrcLoc
loc Env
env Value
op' Value
x' Value
y'
evalAppExp Env
env StructType
_ (If Exp
cond Exp
e1 Exp
e2 SrcLoc
_) = do
  Bool
cond' <- Value -> Bool
asBool forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
cond
  if Bool
cond' then Env -> Exp -> EvalM Value
eval Env
env Exp
e1 else Env -> Exp -> EvalM Value
eval Env
env Exp
e2
evalAppExp Env
env StructType
_ (Apply Exp
f NonEmpty (Info (Diet, Maybe VName), Exp)
args SrcLoc
loc) = do
  -- It is important that 'arguments' are evaluated in reverse order
  -- in order to bring any sizes into scope that may be used in the
  -- type of the functions.
  [Value]
args' <- forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a}. (Info (a, Maybe VName), Exp) -> EvalM Value
evalArg' (forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Diet, Maybe VName), Exp)
args)
  Value
f' <- Env -> Exp -> EvalM Value
eval Env
env Exp
f
  forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env) Value
f' [Value]
args'
  where
    evalArg' :: (Info (a, Maybe VName), Exp) -> EvalM Value
evalArg' (Info (a
_, Maybe VName
ext), Exp
x) = Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
x Maybe VName
ext
evalAppExp Env
env StructType
_ (Index Exp
e SliceBase Info VName
is SrcLoc
loc) = do
  [Indexing]
is' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env) SliceBase Info VName
is
  Value
arr <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value
evalIndex SrcLoc
loc Env
env [Indexing]
is' Value
arr
evalAppExp Env
env StructType
_ (LetWith IdentBase Info VName
dest IdentBase Info VName
src SliceBase Info VName
is Exp
v Exp
body SrcLoc
loc) = do
  let Ident VName
src_vn (Info PatType
src_t) SrcLoc
_ = IdentBase Info VName
src
  Value
dest' <-
    forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall {a}. EvalM a
oob forall (f :: * -> *) a. Applicative f => a -> f a
pure
      forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> Value -> Maybe Value
writeArray
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env) SliceBase Info VName
is
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env (forall v. v -> QualName v
qualName VName
src_vn) (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
src_t)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> Exp -> EvalM Value
eval Env
env Exp
v
  let t :: BoundV
t = [TypeParam] -> StructType -> BoundV
T.BoundV [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType IdentBase Info VName
dest
  Env -> Exp -> EvalM Value
eval (Map VName (Maybe BoundV, Value) -> Env
valEnv (forall k a. k -> a -> Map k a
M.singleton (forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
dest) (forall a. a -> Maybe a
Just BoundV
t, Value
dest')) forall a. Semigroup a => a -> a -> a
<> Env
env) Exp
body
  where
    oob :: EvalM a
oob = forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env Text
"Update out of bounds"
evalAppExp Env
env StructType
_ (DoLoop [VName]
sparams Pat
pat Exp
init_e LoopFormBase Info VName
form Exp
body SrcLoc
_) = do
  Value
init_v <- Env -> Exp -> EvalM Value
eval Env
env Exp
init_e
  case LoopFormBase Info VName
form of
    For IdentBase Info VName
iv Exp
bound -> do
      IntValue
bound' <- Value -> IntValue
asSigned forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
bound
      VName -> IntValue -> IntValue -> Value -> EvalM Value
forLoop (forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
iv) IntValue
bound' (IntValue -> IntValue
zero IntValue
bound') Value
init_v
    ForIn Pat
in_pat Exp
in_e -> do
      (ValueShape
_, [Value]
in_vs) <- Value -> (ValueShape, [Value])
fromArray forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
in_e
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Pat -> Value -> Value -> EvalM Value
forInLoop Pat
in_pat) Value
init_v [Value]
in_vs
    While Exp
cond ->
      Exp -> Value -> EvalM Value
whileLoop Exp
cond Value
init_v
  where
    withLoopParams :: Value -> EvalM Env
withLoopParams Value
v =
      let sparams' :: Sizes
sparams' =
            [VName] -> StructType -> ValueShape -> Sizes
resolveExistentials
              [VName]
sparams
              (Pat -> StructType
patternStructType Pat
pat)
              (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v)
       in Env -> Pat -> Value -> EvalM Env
matchPat (Sizes -> Env
i64Env Sizes
sparams' forall a. Semigroup a => a -> a -> a
<> Env
env) Pat
pat Value
v

    inc :: IntValue -> IntValue
inc = (IntValue -> IntValue -> IntValue
`P.doAdd` Int64 -> IntValue
Int64Value Int64
1)
    zero :: IntValue -> IntValue
zero = (IntValue -> IntValue -> IntValue
`P.doMul` Int64 -> IntValue
Int64Value Int64
0)

    evalBody :: Env -> EvalM Value
evalBody Env
env' = forall a. EvalM a -> EvalM a
localExts forall a b. (a -> b) -> a -> b
$ Env -> Exp -> EvalM Value
eval Env
env' Exp
body

    forLoopEnv :: VName -> IntValue -> Env
forLoopEnv VName
iv IntValue
i =
      Map VName (Maybe BoundV, Value) -> Env
valEnv
        ( forall k a. k -> a -> Map k a
M.singleton
            VName
iv
            ( forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64,
              forall (m :: * -> *). PrimValue -> Value m
ValuePrim (IntValue -> PrimValue
SignedValue IntValue
i)
            )
        )

    forLoop :: VName -> IntValue -> IntValue -> Value -> EvalM Value
forLoop VName
iv IntValue
bound IntValue
i Value
v
      | IntValue
i forall a. Ord a => a -> a -> Bool
>= IntValue
bound = forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
      | Bool
otherwise = do
          Env
env' <- Value -> EvalM Env
withLoopParams Value
v
          VName -> IntValue -> IntValue -> Value -> EvalM Value
forLoop VName
iv IntValue
bound (IntValue -> IntValue
inc IntValue
i) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> EvalM Value
evalBody (VName -> IntValue -> Env
forLoopEnv VName
iv IntValue
i forall a. Semigroup a => a -> a -> a
<> Env
env')

    whileLoop :: Exp -> Value -> EvalM Value
whileLoop Exp
cond Value
v = do
      Env
env' <- Value -> EvalM Env
withLoopParams Value
v
      Bool
continue <- Value -> Bool
asBool forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env' Exp
cond
      if Bool
continue
        then Exp -> Value -> EvalM Value
whileLoop Exp
cond forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> EvalM Value
evalBody Env
env'
        else forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v

    forInLoop :: Pat -> Value -> Value -> EvalM Value
forInLoop Pat
in_pat Value
v Value
in_v = do
      Env
env' <- Value -> EvalM Env
withLoopParams Value
v
      Env
env'' <- Env -> Pat -> Value -> EvalM Env
matchPat Env
env' Pat
in_pat Value
in_v
      Env -> EvalM Value
evalBody Env
env''
evalAppExp Env
env StructType
_ (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
_) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  Value -> [CaseBase Info VName] -> EvalM Value
match Value
v (forall a. NonEmpty a -> [a]
NE.toList NonEmpty (CaseBase Info VName)
cs)
  where
    match :: Value -> [CaseBase Info VName] -> EvalM Value
match Value
_ [] =
      forall a. HasCallStack => [Char] -> a
error [Char]
"Pattern match failure."
    match Value
v (CaseBase Info VName
c : [CaseBase Info VName]
cs') = do
      Maybe Value
c' <- Value -> Env -> CaseBase Info VName -> EvalM (Maybe Value)
evalCase Value
v Env
env CaseBase Info VName
c
      case Maybe Value
c' of
        Just Value
v' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v'
        Maybe Value
Nothing -> Value -> [CaseBase Info VName] -> EvalM Value
match Value
v [CaseBase Info VName]
cs'

eval :: Env -> Exp -> EvalM Value
eval :: Env -> Exp -> EvalM Value
eval Env
_ (Literal PrimValue
v SrcLoc
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
v
eval Env
env (Hole (Info PatType
t) SrcLoc
loc) =
  forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$ Text
"Hole of type: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyTextOneLine PatType
t
eval Env
env (Parens Exp
e SrcLoc
_) = Env -> Exp -> EvalM Value
eval Env
env Exp
e
eval Env
env (QualParens (QualName VName
qv, SrcLoc
_) Exp
e SrcLoc
loc) = do
  Module
m <- Env -> QualName VName -> EvalM Module
evalModuleVar Env
env QualName VName
qv
  case Module
m of
    ModuleFun {} -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Local open of module function at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr SrcLoc
loc
    Module Env
m' -> Env -> Exp -> EvalM Value
eval (Env
m' forall a. Semigroup a => a -> a -> a
<> Env
env) Exp
e
eval Env
env (TupLit [Exp]
vs SrcLoc
_) = forall (m :: * -> *). [Value m] -> Value m
toTuple forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> Exp -> EvalM Value
eval Env
env) [Exp]
vs
eval Env
env (RecordLit [FieldBase Info VName]
fields SrcLoc
_) =
  forall (m :: * -> *). Map Name (Value m) -> Value m
ValueRecord forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> EvalM (Name, Value)
evalField [FieldBase Info VName]
fields
  where
    evalField :: FieldBase Info VName -> EvalM (Name, Value)
evalField (RecordFieldExplicit Name
k Exp
e SrcLoc
_) = do
      Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name
k, Value
v)
    evalField (RecordFieldImplicit VName
k Info PatType
t SrcLoc
loc) = do
      Value
v <- Env -> Exp -> EvalM Value
eval Env
env forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
k) Info PatType
t SrcLoc
loc
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Name
baseName VName
k, Value
v)
eval Env
_ (StringLit [Word8]
vs SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' forall d. Shape d
ShapeLeaf forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
map (forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
UnsignedValue forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> IntValue
Int8Value forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) [Word8]
vs
eval Env
env (ArrayLit [] (Info PatType
t) SrcLoc
_) = do
  ValueShape
t' <- Env -> StructType -> EvalM ValueShape
typeValueShape Env
env forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
t' []
eval Env
env (ArrayLit (Exp
v : [Exp]
vs) Info PatType
_ SrcLoc
_) = do
  Value
v' <- Env -> Exp -> EvalM Value
eval Env
env Exp
v
  [Value]
vs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> Exp -> EvalM Value
eval Env
env) [Exp]
vs
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v') (Value
v' forall a. a -> [a] -> [a]
: [Value]
vs')
eval Env
env (AppExp AppExp
e (Info (AppRes PatType
t [VName]
retext))) = do
  StructType
t' <- Env -> StructType -> EvalM StructType
evalType Env
env forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  forall als.
Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value
returned Env
env StructType
t' [VName]
retext forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> StructType -> AppExp -> EvalM Value
evalAppExp Env
env StructType
t' AppExp
e
eval Env
env (Var QualName VName
qv (Info PatType
t) SrcLoc
_) = Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)
eval Env
env (Ascript Exp
e TypeExp Info VName
_ SrcLoc
_) = Env -> Exp -> EvalM Value
eval Env
env Exp
e
eval Env
_ (IntLit Integer
v (Info PatType
t) SrcLoc
_) =
  case PatType
t of
    Scalar (Prim (Signed IntType
it)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v
    Scalar (Prim (Unsigned IntType
it)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v
    Scalar (Prim (FloatType FloatType
ft)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v
    PatType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"eval: nonsensical type for integer literal: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString PatType
t
eval Env
_ (FloatLit Double
v (Info PatType
t) SrcLoc
_) =
  case PatType
t of
    Scalar (Prim (FloatType FloatType
ft)) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v
    PatType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"eval: nonsensical type for float literal: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString PatType
t
eval Env
env (Negate Exp
e SrcLoc
_) = do
  Value
ev <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Value
ev of
    ValuePrim (SignedValue (Int8Value Int8
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int8 -> IntValue
Int8Value (-Int8
v)
    ValuePrim (SignedValue (Int16Value Int16
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int16 -> IntValue
Int16Value (-Int16
v)
    ValuePrim (SignedValue (Int32Value Int32
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int32 -> IntValue
Int32Value (-Int32
v)
    ValuePrim (SignedValue (Int64Value Int64
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (-Int64
v)
    ValuePrim (UnsignedValue (Int8Value Int8
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue forall a b. (a -> b) -> a -> b
$ Int8 -> IntValue
Int8Value (-Int8
v)
    ValuePrim (UnsignedValue (Int16Value Int16
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue forall a b. (a -> b) -> a -> b
$ Int16 -> IntValue
Int16Value (-Int16
v)
    ValuePrim (UnsignedValue (Int32Value Int32
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue forall a b. (a -> b) -> a -> b
$ Int32 -> IntValue
Int32Value (-Int32
v)
    ValuePrim (UnsignedValue (Int64Value Int64
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (-Int64
v)
    ValuePrim (FloatValue (Float16Value Half
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue forall a b. (a -> b) -> a -> b
$ Half -> FloatValue
Float16Value (-Half
v)
    ValuePrim (FloatValue (Float32Value Float
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue forall a b. (a -> b) -> a -> b
$ Float -> FloatValue
Float32Value (-Float
v)
    ValuePrim (FloatValue (Float64Value Double
v)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue forall a b. (a -> b) -> a -> b
$ Double -> FloatValue
Float64Value (-Double
v)
    Value
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot negate " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
ev
eval Env
env (Not Exp
e SrcLoc
_) = do
  Value
ev <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Value
ev of
    ValuePrim (BoolValue Bool
b) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not Bool
b
    ValuePrim (SignedValue IntValue
iv) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ IntValue -> IntValue
P.doComplement IntValue
iv
    ValuePrim (UnsignedValue IntValue
iv) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue forall a b. (a -> b) -> a -> b
$ IntValue -> IntValue
P.doComplement IntValue
iv
    Value
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot logically negate " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
ev
eval Env
env (Update Exp
src SliceBase Info VName
is Exp
v SrcLoc
loc) =
  forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall {a}. EvalM a
oob forall (f :: * -> *) a. Applicative f => a -> f a
pure
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> Value -> Maybe Value
writeArray forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env) SliceBase Info VName
is forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> Exp -> EvalM Value
eval Env
env Exp
src forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> Exp -> EvalM Value
eval Env
env Exp
v
  where
    oob :: EvalM a
oob = forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env Text
"Bad update"
eval Env
env (RecordUpdate Exp
src [Name]
all_fs Exp
v Info PatType
_ SrcLoc
_) =
  forall {m :: * -> *}. Value m -> [Name] -> Value m -> Value m
update forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
src forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
all_fs forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> Exp -> EvalM Value
eval Env
env Exp
v
  where
    update :: Value m -> [Name] -> Value m -> Value m
update Value m
_ [] Value m
v' = Value m
v'
    update (ValueRecord Map Name (Value m)
src') (Name
f : [Name]
fs) Value m
v'
      | Just Value m
f_v <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (Value m)
src' =
          forall (m :: * -> *). Map Name (Value m) -> Value m
ValueRecord forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
f (Value m -> [Name] -> Value m -> Value m
update Value m
f_v [Name]
fs Value m
v') Map Name (Value m)
src'
    update Value m
_ [Name]
_ Value m
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"eval RecordUpdate: invalid value."
-- We treat zero-parameter lambdas as simply an expression to
-- evaluate immediately.  Note that this is *not* the same as a lambda
-- that takes an empty tuple '()' as argument!  Zero-parameter lambdas
-- can never occur in a well-formed Futhark program, but they are
-- convenient in the interpreter.
eval Env
env (Lambda [Pat]
ps Exp
body Maybe (TypeExp Info VName)
_ (Info (Aliasing
_, RetType [VName]
_ StructType
rt)) SrcLoc
_) =
  Env -> [VName] -> [Pat] -> Exp -> StructType -> EvalM Value
evalFunction Env
env [] [Pat]
ps Exp
body StructType
rt
eval Env
env (OpSection QualName VName
qv (Info PatType
t) SrcLoc
_) =
  Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
eval Env
env (OpSectionLeft QualName VName
qv Info PatType
_ Exp
e (Info (PName
_, StructType
_, Maybe VName
argext), Info (PName, StructType)
_) (Info (RetType [VName]
_ PatType
t), Info [VName]
_) SrcLoc
loc) = do
  Value
v <- Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
e Maybe VName
argext
  Value
f <- Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t)
  SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env Value
f Value
v
eval Env
env (OpSectionRight QualName VName
qv Info PatType
_ Exp
e (Info (PName, StructType)
_, Info (PName
_, StructType
_, Maybe VName
argext)) (Info (RetType [VName]
_ PatType
t)) SrcLoc
loc) = do
  Value
y <- Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
e Maybe VName
argext
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x -> do
      Value
f <- Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
      SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 SrcLoc
loc Env
env Value
f Value
x Value
y
eval Env
env (IndexSection SliceBase Info VName
is Info PatType
_ SrcLoc
loc) = do
  [Indexing]
is' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env) SliceBase Info VName
is
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value
evalIndex SrcLoc
loc Env
env [Indexing]
is'
eval Env
_ (ProjectSection [Name]
ks Info PatType
_ SrcLoc
_) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {f :: * -> *} {m :: * -> *}.
Applicative f =>
Value m -> Name -> f (Value m)
walk) [Name]
ks
  where
    walk :: Value m -> Name -> f (Value m)
walk (ValueRecord Map Name (Value m)
fs) Name
f
      | Just Value m
v' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (Value m)
fs = forall (f :: * -> *) a. Applicative f => a -> f a
pure Value m
v'
    walk Value m
_ Name
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"Value does not have expected field."
eval Env
env (Project Name
f Exp
e Info PatType
_ SrcLoc
_) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  case Value
v of
    ValueRecord Map Name Value
fs | Just Value
v' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name Value
fs -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v'
    Value
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Value does not have expected field."
eval Env
env (Assert Exp
what Exp
e (Info Text
s) SrcLoc
loc) = do
  Bool
cond <- Value -> Bool
asBool forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Exp -> EvalM Value
eval Env
env Exp
what
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
cond forall a b. (a -> b) -> a -> b
$ forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env Text
s
  Env -> Exp -> EvalM Value
eval Env
env Exp
e
eval Env
env (Constr Name
c [Exp]
es (Info PatType
t) SrcLoc
_) = do
  [Value]
vs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Env -> Exp -> EvalM Value
eval Env
env) [Exp]
es
  ValueShape
shape <- Env -> StructType -> EvalM ValueShape
typeValueShape Env
env forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> Name -> [Value m] -> Value m
ValueSum ValueShape
shape Name
c [Value]
vs
eval Env
env (Attr (AttrAtom (AtomName Name
"break") SrcLoc
_) Exp
e SrcLoc
loc) = do
  Env -> Loc -> EvalM ()
break Env
env (forall a. Located a => a -> Loc
locOf SrcLoc
loc)
  Env -> Exp -> EvalM Value
eval Env
env Exp
e
eval Env
env (Attr (AttrAtom (AtomName Name
"trace") SrcLoc
_) Exp
e SrcLoc
loc) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  Text -> Value -> EvalM ()
trace (forall a. Located a => a -> Text
locText (forall a. Located a => a -> Loc
locOf SrcLoc
loc)) Value
v
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
eval Env
env (Attr (AttrComp Name
"trace" [AttrAtom (AtomName Name
tag) SrcLoc
_] SrcLoc
_) Exp
e SrcLoc
_) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  Text -> Value -> EvalM ()
trace (Name -> Text
nameToText Name
tag) Value
v
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
eval Env
env (Attr AttrInfo VName
_ Exp
e SrcLoc
_) =
  Env -> Exp -> EvalM Value
eval Env
env Exp
e

evalCase ::
  Value ->
  Env ->
  CaseBase Info VName ->
  EvalM (Maybe Value)
evalCase :: Value -> Env -> CaseBase Info VName -> EvalM (Maybe Value)
evalCase Value
v Env
env (CasePat Pat
p Exp
cExp SrcLoc
_) = forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT forall a b. (a -> b) -> a -> b
$ do
  Env
env' <- Env -> Pat -> Value -> MaybeT EvalM Env
patternMatch Env
env Pat
p Value
v
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ Env -> Exp -> EvalM Value
eval Env
env' Exp
cExp

-- We hackily do multiple substitutions in modules, because otherwise
-- we would lose in cases where the parameter substitutions are [a->x,
-- b->x] when we reverse. (See issue #1250.)
reverseSubstitutions :: M.Map VName VName -> M.Map VName [VName]
reverseSubstitutions :: Map VName VName -> Map VName [VName]
reverseSubstitutions =
  forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith forall a. Semigroup a => a -> a -> a
(<>) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a b c. (a -> b -> c) -> b -> a -> c
flip (,))) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList

substituteInModule :: M.Map VName VName -> Module -> Module
substituteInModule :: Map VName VName -> Module -> Module
substituteInModule Map VName VName
substs = Module -> Module
onModule
  where
    rev_substs :: Map VName [VName]
rev_substs = Map VName VName -> Map VName [VName]
reverseSubstitutions Map VName VName
substs
    replace :: VName -> [VName]
replace VName
v = forall a. a -> Maybe a -> a
fromMaybe [VName
v] forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName [VName]
rev_substs
    replaceQ :: QualName VName -> QualName VName
replaceQ QualName VName
v = forall b a. b -> (a -> b) -> Maybe a -> b
maybe QualName VName
v forall v. v -> QualName v
qualName forall a b. (a -> b) -> a -> b
$ forall a. [a] -> Maybe a
maybeHead forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName [VName]
rev_substs
    replaceM :: (t -> a) -> Map VName t -> Map VName a
replaceM t -> a
f Map VName t
m = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ do
      (VName
k, t
v) <- forall k a. Map k a -> [(k, a)]
M.toList Map VName t
m
      VName
k' <- VName -> [VName]
replace VName
k
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
k', t -> a
f t
v)
    onModule :: Module -> Module
onModule (Module (Env Map VName TermBinding
terms Map VName TypeBinding
types)) =
      Env -> Module
Module forall a b. (a -> b) -> a -> b
$ Map VName TermBinding -> Map VName TypeBinding -> Env
Env (forall {t} {a}. (t -> a) -> Map VName t -> Map VName a
replaceM TermBinding -> TermBinding
onTerm Map VName TermBinding
terms) (forall {t} {a}. (t -> a) -> Map VName t -> Map VName a
replaceM TypeBinding -> TypeBinding
onType Map VName TypeBinding
types)
    onModule (ModuleFun Module -> EvalM Module
f) =
      (Module -> EvalM Module) -> Module
ModuleFun forall a b. (a -> b) -> a -> b
$ \Module
m -> Module -> Module
onModule forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Module -> EvalM Module
f (Map VName VName -> Module -> Module
substituteInModule Map VName VName
substs Module
m)
    onTerm :: TermBinding -> TermBinding
onTerm (TermValue Maybe BoundV
t Value
v) = Maybe BoundV -> Value -> TermBinding
TermValue Maybe BoundV
t Value
v
    onTerm (TermPoly Maybe BoundV
t StructType -> EvalM Value
v) = Maybe BoundV -> (StructType -> EvalM Value) -> TermBinding
TermPoly Maybe BoundV
t StructType -> EvalM Value
v
    onTerm (TermModule Module
m) = Module -> TermBinding
TermModule forall a b. (a -> b) -> a -> b
$ Module -> Module
onModule Module
m
    onType :: TypeBinding -> TypeBinding
onType (T.TypeAbbr Liftedness
l [TypeParam]
ps RetTypeBase Size ()
t) = Liftedness -> [TypeParam] -> RetTypeBase Size () -> TypeBinding
T.TypeAbbr Liftedness
l [TypeParam]
ps forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim RetTypeBase Size ()
t
    onDim :: Size -> Size
onDim (NamedSize QualName VName
v) = QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ QualName VName -> QualName VName
replaceQ QualName VName
v
    onDim (ConstSize Int64
x) = Int64 -> Size
ConstSize Int64
x
    onDim (AnySize Maybe VName
v) = Maybe VName -> Size
AnySize Maybe VName
v

evalModuleVar :: Env -> QualName VName -> EvalM Module
evalModuleVar :: Env -> QualName VName -> EvalM Module
evalModuleVar Env
env QualName VName
qv =
  case QualName VName -> Env -> Maybe TermBinding
lookupVar QualName VName
qv Env
env of
    Just (TermModule Module
m) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Module
m
    Maybe TermBinding
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString QualName VName
qv forall a. Semigroup a => a -> a -> a
<> [Char]
" is not bound to a module."

evalModExp :: Env -> ModExp -> EvalM Module
evalModExp :: Env -> ModExp -> EvalM Module
evalModExp Env
_ (ModImport [Char]
_ (Info ImportName
f) SrcLoc
_) = do
  Maybe Env
f' <- ImportName -> EvalM (Maybe Env)
lookupImport ImportName
f
  Map ImportName Env
known <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> b
snd
  case Maybe Env
f' of
    Maybe Env
Nothing ->
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
        [[Char]] -> [Char]
unlines
          [ [Char]
"Unknown interpreter import: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show ImportName
f,
            [Char]
"Known: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show (forall k a. Map k a -> [k]
M.keys Map ImportName Env
known)
          ]
    Just Env
m -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Env -> Module
Module Env
m
evalModExp Env
env (ModDecs [DecBase Info VName]
ds SrcLoc
_) = do
  Env Map VName TermBinding
terms Map VName TypeBinding
types <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Env -> DecBase Info VName -> EvalM Env
evalDec Env
env [DecBase Info VName]
ds
  -- Remove everything that was present in the original Env.
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    Env -> Module
Module forall a b. (a -> b) -> a -> b
$
      Map VName TermBinding -> Map VName TypeBinding -> Env
Env
        (Map VName TermBinding
terms forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Env -> Map VName TermBinding
envTerm Env
env)
        (Map VName TypeBinding
types forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Env -> Map VName TypeBinding
envType Env
env)
evalModExp Env
env (ModVar QualName VName
qv SrcLoc
_) =
  Env -> QualName VName -> EvalM Module
evalModuleVar Env
env QualName VName
qv
evalModExp Env
env (ModAscript ModExp
me SigExpBase Info VName
_ (Info Map VName VName
substs) SrcLoc
_) =
  Map VName VName -> Module -> Module
substituteInModule Map VName VName
substs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ModExp -> EvalM Module
evalModExp Env
env ModExp
me
evalModExp Env
env (ModParens ModExp
me SrcLoc
_) = Env -> ModExp -> EvalM Module
evalModExp Env
env ModExp
me
evalModExp Env
env (ModLambda ModParamBase Info VName
p Maybe (SigExpBase Info VName, Info (Map VName VName))
ret ModExp
e SrcLoc
loc) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    (Module -> EvalM Module) -> Module
ModuleFun forall a b. (a -> b) -> a -> b
$ \Module
am -> do
      let env' :: Env
env' = Env
env {envTerm :: Map VName TermBinding
envTerm = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall (f :: * -> *) vn. ModParamBase f vn -> vn
modParamName ModParamBase Info VName
p) (Module -> TermBinding
TermModule Module
am) forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm Env
env}
      Env -> ModExp -> EvalM Module
evalModExp Env
env' forall a b. (a -> b) -> a -> b
$ case Maybe (SigExpBase Info VName, Info (Map VName VName))
ret of
        Maybe (SigExpBase Info VName, Info (Map VName VName))
Nothing -> ModExp
e
        Just (SigExpBase Info VName
se, Info (Map VName VName)
rsubsts) -> forall (f :: * -> *) vn.
ModExpBase f vn
-> SigExpBase f vn
-> f (Map VName VName)
-> SrcLoc
-> ModExpBase f vn
ModAscript ModExp
e SigExpBase Info VName
se Info (Map VName VName)
rsubsts SrcLoc
loc
evalModExp Env
env (ModApply ModExp
f ModExp
e (Info Map VName VName
psubst) (Info Map VName VName
rsubst) SrcLoc
_) = do
  Module
f' <- Env -> ModExp -> EvalM Module
evalModExp Env
env ModExp
f
  case Module
f' of
    ModuleFun Module -> EvalM Module
f'' -> do
      Module
e' <- Env -> ModExp -> EvalM Module
evalModExp Env
env ModExp
e
      Map VName VName -> Module -> Module
substituteInModule Map VName VName
rsubst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Module -> EvalM Module
f'' (Map VName VName -> Module -> Module
substituteInModule Map VName VName
psubst Module
e')
    Module
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Expected ModuleFun."

evalDec :: Env -> Dec -> EvalM Env
evalDec :: Env -> DecBase Info VName -> EvalM Env
evalDec Env
env (ValDec (ValBind Maybe (Info EntryPoint)
_ VName
v Maybe (TypeExp Info VName)
_ (Info RetTypeBase Size ()
ret) [TypeParam]
tparams [Pat]
ps Exp
fbody Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_)) = do
  TermBinding
binding <- Env
-> [TypeParam]
-> [Pat]
-> RetTypeBase Size ()
-> Exp
-> EvalM TermBinding
evalFunctionBinding Env
env [TypeParam]
tparams [Pat]
ps RetTypeBase Size ()
ret Exp
fbody
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Env
env {envTerm :: Map VName TermBinding
envTerm = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v TermBinding
binding forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm Env
env}
evalDec Env
env (OpenDec ModExp
me SrcLoc
_) = do
  Module
me' <- Env -> ModExp -> EvalM Module
evalModExp Env
env ModExp
me
  case Module
me' of
    Module Env
me'' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Env
me'' forall a. Semigroup a => a -> a -> a
<> Env
env
    Module
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Expected Module"
evalDec Env
env (ImportDec [Char]
name Info ImportName
name' SrcLoc
loc) =
  Env -> DecBase Info VName -> EvalM Env
evalDec Env
env forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. DecBase f vn -> SrcLoc -> DecBase f vn
LocalDec (forall (f :: * -> *) vn. ModExpBase f vn -> SrcLoc -> DecBase f vn
OpenDec (forall (f :: * -> *) vn.
[Char] -> f ImportName -> SrcLoc -> ModExpBase f vn
ModImport [Char]
name Info ImportName
name' SrcLoc
loc) SrcLoc
loc) SrcLoc
loc
evalDec Env
env (LocalDec DecBase Info VName
d SrcLoc
_) = Env -> DecBase Info VName -> EvalM Env
evalDec Env
env DecBase Info VName
d
evalDec Env
env SigDec {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env
evalDec Env
env (TypeDec (TypeBind VName
v Liftedness
l [TypeParam]
ps TypeExp Info VName
_ (Info (RetType [VName]
dims StructType
t)) Maybe DocComment
_ SrcLoc
_)) = do
  let abbr :: TypeBinding
abbr = Liftedness -> [TypeParam] -> RetTypeBase Size () -> TypeBinding
T.TypeAbbr Liftedness
l [TypeParam]
ps forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims forall a b. (a -> b) -> a -> b
$ Env -> StructType -> StructType
expandType Env
env StructType
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env {envType :: Map VName TypeBinding
envType = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v TypeBinding
abbr forall a b. (a -> b) -> a -> b
$ Env -> Map VName TypeBinding
envType Env
env}
evalDec Env
env (ModDec (ModBind VName
v [ModParamBase Info VName]
ps Maybe (SigExpBase Info VName, Info (Map VName VName))
ret ModExp
body Maybe DocComment
_ SrcLoc
loc)) = do
  Module
mod <- Env -> ModExp -> EvalM Module
evalModExp Env
env forall a b. (a -> b) -> a -> b
$ [ModParamBase Info VName] -> ModExp
wrapInLambda [ModParamBase Info VName]
ps
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Map VName Module -> Env
modEnv (forall k a. k -> a -> Map k a
M.singleton VName
v Module
mod) forall a. Semigroup a => a -> a -> a
<> Env
env
  where
    wrapInLambda :: [ModParamBase Info VName] -> ModExp
wrapInLambda [] = case Maybe (SigExpBase Info VName, Info (Map VName VName))
ret of
      Just (SigExpBase Info VName
se, Info (Map VName VName)
substs) -> forall (f :: * -> *) vn.
ModExpBase f vn
-> SigExpBase f vn
-> f (Map VName VName)
-> SrcLoc
-> ModExpBase f vn
ModAscript ModExp
body SigExpBase Info VName
se Info (Map VName VName)
substs SrcLoc
loc
      Maybe (SigExpBase Info VName, Info (Map VName VName))
Nothing -> ModExp
body
    wrapInLambda [ModParamBase Info VName
p] = forall (f :: * -> *) vn.
ModParamBase f vn
-> Maybe (SigExpBase f vn, f (Map VName VName))
-> ModExpBase f vn
-> SrcLoc
-> ModExpBase f vn
ModLambda ModParamBase Info VName
p Maybe (SigExpBase Info VName, Info (Map VName VName))
ret ModExp
body SrcLoc
loc
    wrapInLambda (ModParamBase Info VName
p : [ModParamBase Info VName]
ps') = forall (f :: * -> *) vn.
ModParamBase f vn
-> Maybe (SigExpBase f vn, f (Map VName VName))
-> ModExpBase f vn
-> SrcLoc
-> ModExpBase f vn
ModLambda ModParamBase Info VName
p forall a. Maybe a
Nothing ([ModParamBase Info VName] -> ModExp
wrapInLambda [ModParamBase Info VName]
ps') SrcLoc
loc

-- | The interpreter context.  All evaluation takes place with respect
-- to a context, and it can be extended with more definitions, which
-- is how the REPL works.
data Ctx = Ctx
  { Ctx -> Env
ctxEnv :: Env,
    Ctx -> Map ImportName Env
ctxImports :: M.Map ImportName Env
  }

nanValue :: PrimValue -> Bool
nanValue :: PrimValue -> Bool
nanValue (FloatValue FloatValue
v) =
  case FloatValue
v of
    Float16Value Half
x -> forall a. RealFloat a => a -> Bool
isNaN Half
x
    Float32Value Float
x -> forall a. RealFloat a => a -> Bool
isNaN Float
x
    Float64Value Double
x -> forall a. RealFloat a => a -> Bool
isNaN Double
x
nanValue PrimValue
_ = Bool
False

breakOnNaN :: [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN :: [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN [PrimValue]
inputs PrimValue
result
  | Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any PrimValue -> Bool
nanValue [PrimValue]
inputs) Bool -> Bool -> Bool
&& PrimValue -> Bool
nanValue PrimValue
result = do
      [StackFrame]
backtrace <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
      case forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty [StackFrame]
backtrace of
        Maybe (NonEmpty StackFrame)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just NonEmpty StackFrame
backtrace' ->
          let loc :: Loc
loc = StackFrame -> Loc
stackFrameLoc forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> a
NE.head NonEmpty StackFrame
backtrace'
           in forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF forall a b. (a -> b) -> a -> b
$ forall a. Loc -> BreakReason -> NonEmpty StackFrame -> a -> ExtOp a
ExtOpBreak Loc
loc BreakReason
BreakNaN NonEmpty StackFrame
backtrace' ()
breakOnNaN [PrimValue]
_ PrimValue
_ =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | The initial environment contains definitions of the various intrinsic functions.
initialCtx :: Ctx
initialCtx :: Ctx
initialCtx =
  Env -> Map ImportName Env -> Ctx
Ctx
    ( Map VName TermBinding -> Map VName TypeBinding -> Env
Env
        ( forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
            (Name -> Int -> VName
VName ([Char] -> Name
nameFromString [Char]
"intrinsics") Int
0)
            (Module -> TermBinding
TermModule (Env -> Module
Module forall a b. (a -> b) -> a -> b
$ Map VName TermBinding -> Map VName TypeBinding -> Env
Env Map VName TermBinding
terms Map VName TypeBinding
types))
            Map VName TermBinding
terms
        )
        Map VName TypeBinding
types
    )
    forall a. Monoid a => a
mempty
  where
    terms :: Map VName TermBinding
terms = forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey (forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Maybe TermBinding
def forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
baseString) Map VName Intrinsic
intrinsics
    types :: Map VName TypeBinding
types = forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey (forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Maybe TypeBinding
tdef forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [Char]
baseString) Map VName Intrinsic
intrinsics

    sintOp :: (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp IntType -> BinOp
f =
      [ (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int8)),
        (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int16)),
        (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int32)),
        (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int64))
      ]
    uintOp :: (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp IntType -> BinOp
f =
      [ (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int8)),
        (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int16)),
        (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int32)),
        (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int64))
      ]
    intOp :: (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
intOp IntType -> BinOp
f = (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp IntType -> BinOp
f forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp IntType -> BinOp
f
    floatOp :: (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
floatOp FloatType -> BinOp
f =
      [ (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (FloatType -> BinOp
f FloatType
Float16)),
        (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (FloatType -> BinOp
f FloatType
Float32)),
        (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (FloatType -> BinOp
f FloatType
Float64))
      ]
    arithOp :: (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp IntType -> BinOp
f FloatType -> BinOp
g = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
intOp IntType -> BinOp
f forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
floatOp FloatType -> BinOp
g

    flipCmps :: [(a, b, a -> b -> c)] -> [(a, b, b -> a -> c)]
flipCmps = forall a b. (a -> b) -> [a] -> [b]
map (\(a
f, b
g, a -> b -> c
h) -> (a
f, b
g, forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> c
h))
    sintCmp :: (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
sintCmp IntType -> CmpOp
f =
      [ (PrimValue -> Maybe PrimValue
getS, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int8)),
        (PrimValue -> Maybe PrimValue
getS, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int16)),
        (PrimValue -> Maybe PrimValue
getS, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int32)),
        (PrimValue -> Maybe PrimValue
getS, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int64))
      ]
    uintCmp :: (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
uintCmp IntType -> CmpOp
f =
      [ (PrimValue -> Maybe PrimValue
getU, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int8)),
        (PrimValue -> Maybe PrimValue
getU, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int16)),
        (PrimValue -> Maybe PrimValue
getU, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int32)),
        (PrimValue -> Maybe PrimValue
getU, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (IntType -> CmpOp
f IntType
Int64))
      ]
    floatCmp :: (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
floatCmp FloatType -> CmpOp
f =
      [ (PrimValue -> Maybe PrimValue
getF, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (FloatType -> CmpOp
f FloatType
Float16)),
        (PrimValue -> Maybe PrimValue
getF, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (FloatType -> CmpOp
f FloatType
Float32)),
        (PrimValue -> Maybe PrimValue
getF, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp (FloatType -> CmpOp
f FloatType
Float64))
      ]
    boolCmp :: CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
boolCmp CmpOp
f = [(PrimValue -> Maybe PrimValue
getB, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp CmpOp
f)]

    getV :: PrimValue -> Maybe PrimValue
getV (SignedValue IntValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getV (UnsignedValue IntValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getV (FloatValue FloatValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
P.FloatValue FloatValue
x
    getV (BoolValue Bool
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
P.BoolValue Bool
x
    putV :: PrimValue -> PrimValue
putV (P.IntValue IntValue
x) = IntValue -> PrimValue
SignedValue IntValue
x
    putV (P.FloatValue FloatValue
x) = FloatValue -> PrimValue
FloatValue FloatValue
x
    putV (P.BoolValue Bool
x) = Bool -> PrimValue
BoolValue Bool
x
    putV PrimValue
P.UnitValue = Bool -> PrimValue
BoolValue Bool
True

    getS :: PrimValue -> Maybe PrimValue
getS (SignedValue IntValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getS PrimValue
_ = forall a. Maybe a
Nothing
    putS :: PrimValue -> Maybe PrimValue
putS (P.IntValue IntValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue IntValue
x
    putS PrimValue
_ = forall a. Maybe a
Nothing

    getU :: PrimValue -> Maybe PrimValue
getU (UnsignedValue IntValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getU PrimValue
_ = forall a. Maybe a
Nothing
    putU :: PrimValue -> Maybe PrimValue
putU (P.IntValue IntValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue IntValue
x
    putU PrimValue
_ = forall a. Maybe a
Nothing

    getF :: PrimValue -> Maybe PrimValue
getF (FloatValue FloatValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
P.FloatValue FloatValue
x
    getF PrimValue
_ = forall a. Maybe a
Nothing
    putF :: PrimValue -> Maybe PrimValue
putF (P.FloatValue FloatValue
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue FloatValue
x
    putF PrimValue
_ = forall a. Maybe a
Nothing

    getB :: PrimValue -> Maybe PrimValue
getB (BoolValue Bool
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
P.BoolValue Bool
x
    getB PrimValue
_ = forall a. Maybe a
Nothing
    putB :: PrimValue -> Maybe PrimValue
putB (P.BoolValue Bool
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
x
    putB PrimValue
_ = forall a. Maybe a
Nothing

    fun1 :: (Value -> EvalM Value) -> TermBinding
fun1 Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x -> Value -> EvalM Value
f Value
x

    fun2 :: (Value -> Value -> EvalM Value) -> TermBinding
fun2 Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y -> Value -> Value -> EvalM Value
f Value
x Value
y

    fun3 :: (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 Value -> Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
z -> Value -> Value -> Value -> EvalM Value
f Value
x Value
y Value
z

    fun5 :: (Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun5 Value -> Value -> Value -> Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
z ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
a ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
b -> Value -> Value -> Value -> Value -> Value -> EvalM Value
f Value
x Value
y Value
z Value
a Value
b

    fun6 :: (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
z ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
a ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
b ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
c -> Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value
f Value
x Value
y Value
z Value
a Value
b Value
c

    fun7 :: (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun7 Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
z ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
a ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
b ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
c ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
d -> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f Value
x Value
y Value
z Value
a Value
b Value
c Value
d

    fun8 :: (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun8 Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
z ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
a ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
b ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
c ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
d ->
                    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
e -> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f Value
x Value
y Value
z Value
a Value
b Value
c Value
d Value
e

    fun10 :: (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun10 Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
z ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
a ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
b ->
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
c ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
d ->
                    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
e ->
                      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
g ->
                        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
h -> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f Value
x Value
y Value
z Value
a Value
b Value
c Value
d Value
e Value
g Value
h

    bopDef :: [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
fs = (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$ \Value
x Value
y ->
      case (Value
x, Value
y) of
        (ValuePrim PrimValue
x', ValuePrim PrimValue
y')
          | Just PrimValue
z <- forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall {m :: * -> *} {t} {t} {a} {b}.
Monad m =>
(t -> m t, a -> m b, t -> t -> m a) -> (t, t) -> m b
`bopDef'` (PrimValue
x', PrimValue
y')) [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
fs -> do
              [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN [PrimValue
x', PrimValue
y'] PrimValue
z
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
z
        (Value, Value)
_ ->
          forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
            Doc Any
"Cannot apply operator to arguments"
              forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
x)
              forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc Any
"and"
              forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
y)
                forall a. Semigroup a => a -> a -> a
<> Doc Any
"."
      where
        bopDef' :: (t -> m t, a -> m b, t -> t -> m a) -> (t, t) -> m b
bopDef' (t -> m t
valf, a -> m b
retf, t -> t -> m a
op) (t
x, t
y) = do
          t
x' <- t -> m t
valf t
x
          t
y' <- t -> m t
valf t
y
          a -> m b
retf forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< t -> t -> m a
op t
x' t
y'

    unopDef :: [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
-> TermBinding
unopDef [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
fs = (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
x ->
      case Value
x of
        (ValuePrim PrimValue
x')
          | Just PrimValue
r <- forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall {m :: * -> *} {t} {t} {a} {b}.
Monad m =>
(t -> m t, a -> m b, t -> m a) -> t -> m b
`unopDef'` PrimValue
x') [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
fs -> do
              [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN [PrimValue
x'] PrimValue
r
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
r
        Value
_ ->
          forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
            Doc Any
"Cannot apply function to argument"
              forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
x)
                forall a. Semigroup a => a -> a -> a
<> Doc Any
"."
      where
        unopDef' :: (t -> m t, a -> m b, t -> m a) -> t -> m b
unopDef' (t -> m t
valf, a -> m b
retf, t -> m a
op) t
x = do
          t
x' <- t -> m t
valf t
x
          a -> m b
retf forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< t -> m a
op t
x'

    tbopDef :: (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
tbopDef PrimValue -> PrimValue -> Maybe PrimValue
f = (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
v ->
      case forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple Value
v of
        Just [ValuePrim PrimValue
x, ValuePrim PrimValue
y]
          | Just PrimValue
x' <- PrimValue -> Maybe PrimValue
getV PrimValue
x,
            Just PrimValue
y' <- PrimValue -> Maybe PrimValue
getV PrimValue
y,
            Just PrimValue
z <- PrimValue -> PrimValue
putV forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimValue -> PrimValue -> Maybe PrimValue
f PrimValue
x' PrimValue
y' -> do
              [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN [PrimValue
x, PrimValue
y] PrimValue
z
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
z
        Maybe [Value]
_ ->
          forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
            Doc Any
"Cannot apply operator to argument"
              forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
dquotes (forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
v) forall a. Semigroup a => a -> a -> a
<> Doc Any
"."

    def :: [Char] -> Maybe TermBinding
def [Char]
"!" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
-> TermBinding
unopDef
          [ (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int8),
            (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int16),
            (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int32),
            (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int64),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int8),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int16),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int32),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int64),
            (PrimValue -> Maybe PrimValue
getB, PrimValue -> Maybe PrimValue
putB, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp UnOp
P.Not)
          ]
    def [Char]
"+" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Add` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FAdd
    def [Char]
"-" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Sub` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FSub
    def [Char]
"*" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Mul` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FMul
    def [Char]
"**" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp IntType -> BinOp
P.Pow FloatType -> BinOp
P.FPow
    def [Char]
"/" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp (IntType -> Safety -> BinOp
`P.SDiv` Safety
P.Unsafe)
            forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp (IntType -> Safety -> BinOp
`P.UDiv` Safety
P.Unsafe)
            forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
floatOp FloatType -> BinOp
P.FDiv
    def [Char]
"%" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp (IntType -> Safety -> BinOp
`P.SMod` Safety
P.Unsafe)
            forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp (IntType -> Safety -> BinOp
`P.UMod` Safety
P.Unsafe)
            forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
floatOp FloatType -> BinOp
P.FMod
    def [Char]
"//" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp (IntType -> Safety -> BinOp
`P.SQuot` Safety
P.Unsafe)
            forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp (IntType -> Safety -> BinOp
`P.UDiv` Safety
P.Unsafe)
    def [Char]
"%%" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp (IntType -> Safety -> BinOp
`P.SRem` Safety
P.Unsafe)
            forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp (IntType -> Safety -> BinOp
`P.UMod` Safety
P.Unsafe)
    def [Char]
"^" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
intOp IntType -> BinOp
P.Xor
    def [Char]
"&" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
intOp IntType -> BinOp
P.And
    def [Char]
"|" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
intOp IntType -> BinOp
P.Or
    def [Char]
">>" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp IntType -> BinOp
P.AShr forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp IntType -> BinOp
P.LShr
    def [Char]
"<<" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
intOp IntType -> BinOp
P.Shl
    def [Char]
">>>" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
sintOp IntType -> BinOp
P.LShr forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue)]
uintOp IntType -> BinOp
P.LShr
    def [Char]
"==" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$
        \Value
xs Value
ys -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue forall a b. (a -> b) -> a -> b
$ Value
xs forall a. Eq a => a -> a -> Bool
== Value
ys
    def [Char]
"!=" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$
        \Value
xs Value
ys -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue forall a b. (a -> b) -> a -> b
$ Value
xs forall a. Eq a => a -> a -> Bool
/= Value
ys
    -- The short-circuiting is handled directly in 'eval'; these cases
    -- are only used when partially applying and such.
    def [Char]
"&&" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$ \Value
x Value
y ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue forall a b. (a -> b) -> a -> b
$ Value -> Bool
asBool Value
x Bool -> Bool -> Bool
&& Value -> Bool
asBool Value
y
    def [Char]
"||" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$ \Value
x Value
y ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue forall a b. (a -> b) -> a -> b
$ Value -> Bool
asBool Value
x Bool -> Bool -> Bool
|| Value -> Bool
asBool Value
y
    def [Char]
"<" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
sintCmp IntType -> CmpOp
P.CmpSlt
            forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
uintCmp IntType -> CmpOp
P.CmpUlt
            forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
floatCmp FloatType -> CmpOp
P.FCmpLt
            forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
boolCmp CmpOp
P.CmpLlt
    def [Char]
">" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          forall {a} {b} {a} {b} {c}.
[(a, b, a -> b -> c)] -> [(a, b, b -> a -> c)]
flipCmps forall a b. (a -> b) -> a -> b
$
            (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
sintCmp IntType -> CmpOp
P.CmpSlt
              forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
uintCmp IntType -> CmpOp
P.CmpUlt
              forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
floatCmp FloatType -> CmpOp
P.FCmpLt
              forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
boolCmp CmpOp
P.CmpLlt
    def [Char]
"<=" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
sintCmp IntType -> CmpOp
P.CmpSle
            forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
uintCmp IntType -> CmpOp
P.CmpUle
            forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
floatCmp FloatType -> CmpOp
P.FCmpLe
            forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
boolCmp CmpOp
P.CmpLle
    def [Char]
">=" =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
        forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a)]
-> TermBinding
bopDef forall a b. (a -> b) -> a -> b
$
          forall {a} {b} {a} {b} {c}.
[(a, b, a -> b -> c)] -> [(a, b, b -> a -> c)]
flipCmps forall a b. (a -> b) -> a -> b
$
            (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
sintCmp IntType -> CmpOp
P.CmpSle
              forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
uintCmp IntType -> CmpOp
P.CmpUle
              forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
floatCmp FloatType -> CmpOp
P.FCmpLe
              forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool)]
boolCmp CmpOp
P.CmpLle
    def [Char]
s
      | Just BinOp
bop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
s ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [BinOp]
P.allBinOps =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
tbopDef forall a b. (a -> b) -> a -> b
$ BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp BinOp
bop
      | Just CmpOp
unop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
s ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [CmpOp]
P.allCmpOps =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
tbopDef forall a b. (a -> b) -> a -> b
$ \PrimValue
x PrimValue
y -> Bool -> PrimValue
P.BoolValue forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp CmpOp
unop PrimValue
x PrimValue
y
      | Just ConvOp
cop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
s ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [ConvOp]
P.allConvOps =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
-> TermBinding
unopDef [(PrimValue -> Maybe PrimValue
getV, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV, ConvOp -> PrimValue -> Maybe PrimValue
P.doConvOp ConvOp
cop)]
      | Just UnOp
unop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
s ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [UnOp]
P.allUnOps =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
-> TermBinding
unopDef [(PrimValue -> Maybe PrimValue
getV, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp UnOp
unop)]
      | Just ([PrimType]
pts, PrimType
_, [PrimValue] -> Maybe PrimValue
f) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup [Char]
s Map [Char] ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
P.primFuns =
          case forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimType]
pts of
            Int
1 -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a)]
-> TermBinding
unopDef [(PrimValue -> Maybe PrimValue
getV, forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV, [PrimValue] -> Maybe PrimValue
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure)]
            Int
_ -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
              (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
x -> do
                let getV' :: Value m -> Maybe PrimValue
getV' (ValuePrim PrimValue
v) = forall a. a -> Maybe a
Just PrimValue
v
                    getV' Value m
_ = forall a. Maybe a
Nothing
                case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {m :: * -> *}. Value m -> Maybe PrimValue
getV' forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple Value
x of
                  Just [PrimValue]
vs
                    | Just PrimValue
res <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PrimValue -> PrimValue
putV forall b c a. (b -> c) -> (a -> b) -> a -> c
. [PrimValue] -> Maybe PrimValue
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PrimValue -> Maybe PrimValue
getV [PrimValue]
vs -> do
                        [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN [PrimValue]
vs PrimValue
res
                        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
res
                  Maybe [PrimValue]
_ ->
                    forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot apply " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString [Char]
s forall a. [a] -> [a] -> [a]
++ [Char]
" to " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
x
      | [Char]
"sign_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
s =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
            (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
x ->
              case Value
x of
                (ValuePrim (UnsignedValue IntValue
x')) ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue IntValue
x'
                Value
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot sign: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
x
      | [Char]
"unsign_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
s =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
            (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
x ->
              case Value
x of
                (ValuePrim (SignedValue IntValue
x')) ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue IntValue
x'
                Value
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot unsign: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
x
    def [Char]
s
      | [Char]
"map_stream" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
s =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ (Value -> Value -> EvalM Value) -> TermBinding
fun2 Value -> Value -> EvalM Value
stream
    def [Char]
s | [Char]
"reduce_stream" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
s =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
_ Value
f Value
arg -> Value -> Value -> EvalM Value
stream Value
f Value
arg
    def [Char]
"map" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      Maybe BoundV -> (StructType -> EvalM Value) -> TermBinding
TermPoly forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ \StructType
t -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
f -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
xs ->
          case forall dim as.
TypeBase dim as -> ([(Diet, TypeBase dim ())], TypeBase dim ())
unfoldFunType StructType
t of
            ([(Diet, StructType)
_, (Diet, StructType)
_], StructType
ret_t)
              | Just ValueShape
rowshape <- StructType -> Maybe ValueShape
typeRowShape StructType
ret_t ->
                  forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (SrcLoc -> Env -> Value -> Value -> EvalM Value
apply forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs)
              | Bool
otherwise ->
                  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Bad return type: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString StructType
ret_t
            ([(Diet, StructType)], StructType)
_ ->
              forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
                [Char]
"Invalid arguments to map intrinsic:\n"
                  forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unlines [forall a. Pretty a => a -> [Char]
prettyString StructType
t, forall a. Show a => a -> [Char]
show Value
f, forall a. Show a => a -> [Char]
show Value
xs]
      where
        typeRowShape :: StructType -> Maybe ValueShape
typeRowShape = forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA forall b c a. (b -> c) -> (a -> b) -> a -> c
. StructType -> Shape (Maybe Int64)
structTypeShape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1
    def [Char]
s | [Char]
"reduce" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
s = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
f Value
ne Value
xs ->
        forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f) Value
ne forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs
    def [Char]
"scan" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
f Value
ne Value
xs -> do
        let next :: ([Value], Value) -> Value -> EvalM ([Value], Value)
next ([Value]
out, Value
acc) Value
x = do
              Value
x' <- SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f Value
acc Value
x
              forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value
x' forall a. a -> [a] -> [a]
: [Value]
out, Value
x')
        forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' (forall (m :: * -> *). Value m -> ValueShape
valueShape Value
ne) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Value], Value) -> Value -> EvalM ([Value], Value)
next ([], Value
ne) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs)
    def [Char]
"scatter" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
is Value
vs ->
        case Value
arr of
          ValueArray ValueShape
shape Array Int Value
arr' ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
              forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
shape forall a b. (a -> b) -> a -> b
$
                forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {m :: * -> *}.
Array Int (Value m) -> (Int, Value m) -> Array Int (Value m)
update Array Int Value
arr' forall a b. (a -> b) -> a -> b
$
                  forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Value -> Int
asInt forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs)
          Value
_ ->
            forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"scatter expects array, but got: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
arr
      where
        update :: Array Int (Value m) -> (Int, Value m) -> Array Int (Value m)
update Array Int (Value m)
arr' (Int
i, Value m
v) =
          if Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int (Value m)
arr'
            then Array Int (Value m)
arr' forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Int
i, Value m
v)]
            else Array Int (Value m)
arr'
    def [Char]
"scatter_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
is Value
vs ->
        case Value
arr of
          ValueArray ValueShape
_ Array Int Value
_ ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
              forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Value -> (Maybe [Value], Value) -> Value
update Value
arr forall a b. (a -> b) -> a -> b
$
                forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs)
          Value
_ ->
            forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"scatter_2d expects array, but got: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
arr
      where
        update :: Value -> (Maybe [Value], Value) -> Value
        update :: Value -> (Maybe [Value], Value) -> Value
update Value
arr (Just idxs :: [Value]
idxs@[Value
_, Value
_], Value
v) =
          forall a. a -> Maybe a -> a
fromMaybe Value
arr forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Value -> Maybe Value
writeArray (forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64) [Value]
idxs) Value
arr Value
v
        update Value
_ (Maybe [Value], Value)
_ =
          forall a. HasCallStack => [Char] -> a
error [Char]
"scatter_2d expects 2-dimensional indices"
    def [Char]
"scatter_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
is Value
vs ->
        case Value
arr of
          ValueArray ValueShape
_ Array Int Value
_ ->
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
              forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Value -> (Maybe [Value], Value) -> Value
update Value
arr forall a b. (a -> b) -> a -> b
$
                forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs)
          Value
_ ->
            forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"scatter_3d expects array, but got: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
arr
      where
        update :: Value -> (Maybe [Value], Value) -> Value
        update :: Value -> (Maybe [Value], Value) -> Value
update Value
arr (Just idxs :: [Value]
idxs@[Value
_, Value
_, Value
_], Value
v) =
          forall a. a -> Maybe a -> a
fromMaybe Value
arr forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Value -> Maybe Value
writeArray (forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64) [Value]
idxs) Value
arr Value
v
        update Value
_ (Maybe [Value], Value)
_ =
          forall a. HasCallStack => [Char] -> a
error [Char]
"scatter_3d expects 3-dimensional indices"
    def [Char]
"hist_1d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 forall a b. (a -> b) -> a -> b
$ \Value
_ Value
arr Value
fun Value
_ Value
is Value
vs ->
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
        (Value -> Value -> (Int64, Value) -> EvalM Value
update Value
fun)
        Value
arr
        (forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Value -> Int64
asInt64 forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs))
      where
        op :: Value -> Value -> Value -> EvalM Value
op = SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty
        update :: Value -> Value -> (Int64, Value) -> EvalM Value
update Value
fun Value
arr (Int64
i, Value
v) =
          forall a. a -> Maybe a -> a
fromMaybe Value
arr forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (Value -> Value -> Value -> EvalM Value
op Value
fun) [Int64 -> Indexing
IndexingFix Int64
i] Value
arr Value
v
    def [Char]
"hist_2d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 forall a b. (a -> b) -> a -> b
$ \Value
_ Value
arr Value
fun Value
_ Value
is Value
vs ->
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
        (Value -> Value -> (Maybe [Value], Value) -> EvalM Value
update Value
fun)
        Value
arr
        (forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs))
      where
        op :: Value -> Value -> Value -> EvalM Value
op = SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty
        update :: Value -> Value -> (Maybe [Value], Value) -> EvalM Value
update Value
fun Value
arr (Just idxs :: [Value]
idxs@[Value
_, Value
_], Value
v) =
          forall a. a -> Maybe a -> a
fromMaybe Value
arr
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (Value -> Value -> Value -> EvalM Value
op Value
fun) (forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64) [Value]
idxs) Value
arr Value
v
        update Value
_ Value
_ (Maybe [Value], Value)
_ =
          forall a. HasCallStack => [Char] -> a
error [Char]
"hist_2d: bad index value"
    def [Char]
"hist_3d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 forall a b. (a -> b) -> a -> b
$ \Value
_ Value
arr Value
fun Value
_ Value
is Value
vs ->
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
        (Value -> Value -> (Maybe [Value], Value) -> EvalM Value
update Value
fun)
        Value
arr
        (forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs))
      where
        op :: Value -> Value -> Value -> EvalM Value
op = SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty
        update :: Value -> Value -> (Maybe [Value], Value) -> EvalM Value
update Value
fun Value
arr (Just idxs :: [Value]
idxs@[Value
_, Value
_, Value
_], Value
v) =
          forall a. a -> Maybe a -> a
fromMaybe Value
arr
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (Value -> Value -> Value -> EvalM Value
op Value
fun) (forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64) [Value]
idxs) Value
arr Value
v
        update Value
_ Value
_ (Maybe [Value], Value)
_ =
          forall a. HasCallStack => [Char] -> a
error [Char]
"hist_2d: bad index value"
    def [Char]
"partition" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
k Value
f Value
xs -> do
        let (ShapeDim Int64
_ ValueShape
rowshape, [Value]
xs') = Value -> (ValueShape, [Value])
fromArray Value
xs

            next :: [[Value]] -> Value -> EvalM [[Value]]
next [[Value]]
outs Value
x = do
              Int
i <- Value -> Int
asInt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc -> Env -> Value -> Value -> EvalM Value
apply forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f Value
x
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {t} {t}. (Eq t, Num t) => t -> t -> [[t]] -> [[t]]
insertAt Int
i Value
x [[Value]]
outs
            pack :: [[Value m]] -> Value m
pack [[Value m]]
parts =
              forall (m :: * -> *). [Value m] -> Value m
toTuple
                [ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Value m]]
parts,
                  forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape forall a b. (a -> b) -> a -> b
$
                    forall a b. (a -> b) -> [a] -> [b]
map (forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
SignedValue forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i a. Num i => [a] -> i
genericLength) [[Value m]]
parts
                ]

        forall {m :: * -> *}. [[Value m]] -> Value m
pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. [a] -> [a]
reverse
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM [[Value]] -> Value -> EvalM [[Value]]
next (forall a. Int -> a -> [a]
replicate (Value -> Int
asInt Value
k) []) [Value]
xs'
      where
        insertAt :: t -> t -> [[t]] -> [[t]]
insertAt t
0 t
x ([t]
l : [[t]]
ls) = (t
x forall a. a -> [a] -> [a]
: [t]
l) forall a. a -> [a] -> [a]
: [[t]]
ls
        insertAt t
i t
x ([t]
l : [[t]]
ls) = [t]
l forall a. a -> [a] -> [a]
: t -> t -> [[t]] -> [[t]]
insertAt (t
i forall a. Num a => a -> a -> a
- t
1) t
x [[t]]
ls
        insertAt t
_ t
_ [[t]]
ls = [[t]]
ls
    def [Char]
"scatter_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
dest Value
f Value
vs ->
        case (Value
dest, Value
vs) of
          ( ValueArray ValueShape
dest_shape Array Int Value
dest_arr,
            ValueArray ValueShape
_ Array Int Value
vs_arr
            ) -> do
              let acc :: Value
acc = forall (m :: * -> *).
(Value m -> Value m -> m (Value m))
-> Array Int (Value m) -> Value m
ValueAcc (\Value
_ Value
x -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
x) Array Int Value
dest_arr
              Value
acc' <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f) Value
acc Array Int Value
vs_arr
              case Value
acc' of
                ValueAcc Value -> Value -> EvalM Value
_ Array Int Value
dest_arr' ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
dest_shape Array Int Value
dest_arr'
                Value
_ ->
                  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"scatter_stream produced: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
acc'
          (Value, Value)
_ ->
            forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"scatter_stream expects array, but got: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString (forall a. Show a => a -> [Char]
show Value
vs, forall a. Show a => a -> [Char]
show Value
vs)
    def [Char]
"hist_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun5 forall a b. (a -> b) -> a -> b
$ \Value
dest Value
op Value
_ne Value
f Value
vs ->
        case (Value
dest, Value
vs) of
          ( ValueArray ValueShape
dest_shape Array Int Value
dest_arr,
            ValueArray ValueShape
_ Array Int Value
vs_arr
            ) -> do
              let acc :: Value
acc = forall (m :: * -> *).
(Value m -> Value m -> m (Value m))
-> Array Int (Value m) -> Value m
ValueAcc (SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
op) Array Int Value
dest_arr
              Value
acc' <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f) Value
acc Array Int Value
vs_arr
              case Value
acc' of
                ValueAcc Value -> Value -> EvalM Value
_ Array Int Value
dest_arr' ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
dest_shape Array Int Value
dest_arr'
                Value
_ ->
                  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"hist_stream produced: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
acc'
          (Value, Value)
_ ->
            forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"hist_stream expects array, but got: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString (forall a. Show a => a -> [Char]
show Value
dest, forall a. Show a => a -> [Char]
show Value
vs)
    def [Char]
"acc_write" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
acc Value
i Value
v ->
        case (Value
acc, Value
i) of
          ( ValueAcc Value -> Value -> EvalM Value
op Array Int Value
acc_arr,
            ValuePrim (SignedValue (Int64Value Int64
i'))
            ) ->
              if Int64
i' forall a. Ord a => a -> a -> Bool
>= Int64
0 Bool -> Bool -> Bool
&& Int64
i' forall a. Ord a => a -> a -> Bool
< forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
acc_arr
                then do
                  let x :: Value
x = Array Int Value
acc_arr forall i e. Ix i => Array i e -> i -> e
! forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i'
                  Value
res <- Value -> Value -> EvalM Value
op Value
x Value
v
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
(Value m -> Value m -> m (Value m))
-> Array Int (Value m) -> Value m
ValueAcc Value -> Value -> EvalM Value
op forall a b. (a -> b) -> a -> b
$ Array Int Value
acc_arr forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i', Value
res)]
                else forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
acc
          (Value, Value)
_ ->
            forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"acc_write invalid arguments: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> [Char]
prettyString (forall a. Show a => a -> [Char]
show Value
acc, forall a. Show a => a -> [Char]
show Value
i, forall a. Show a => a -> [Char]
show Value
v)
    --
    def [Char]
"flat_index_2d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
offset Value
n1 Value
s1 Value
n2 Value
s2 -> do
      let offset' :: Int64
offset' = Value -> Int64
asInt64 Value
offset
          n1' :: Int64
n1' = Value -> Int64
asInt64 Value
n1
          n2' :: Int64
n2' = Value -> Int64
asInt64 Value
n2
          s1' :: Int64
s1' = Value -> Int64
asInt64 Value
s1
          s2' :: Int64
s2' = Value -> Int64
asInt64 Value
s2
          shapeFromDims :: [a] -> Shape a
shapeFromDims = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall d. d -> Shape d -> Shape d
ShapeDim forall d. Shape d
ShapeLeaf
          mk1 :: [Maybe (Value m)] -> Maybe (Value m)
mk1 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (forall {a}. [a] -> Shape a
shapeFromDims [Int64
n1', Int64
n2'])) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          mk2 :: [Maybe (Value m)] -> Maybe (Value m)
mk2 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> Shape a
shapeFromDims [Int64
n2']) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          iota :: a -> [a]
iota a
x = [a
0 .. a
x forall a. Num a => a -> a -> a
- a
1]
          f :: Int64 -> Int64 -> Maybe Value
f Int64
i Int64
j =
            [Indexing] -> Value -> Maybe Value
indexArray [Int64 -> Indexing
IndexingFix forall a b. (a -> b) -> a -> b
$ Int64
offset' forall a. Num a => a -> a -> a
+ Int64
i forall a. Num a => a -> a -> a
* Int64
s1' forall a. Num a => a -> a -> a
+ Int64
j forall a. Num a => a -> a -> a
* Int64
s2'] Value
arr

      case forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk1 [forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk2 [Int64 -> Int64 -> Maybe Value
f Int64
i Int64
j | Int64
j <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n2'] | Int64
i <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n1'] of
        Just Value
arr' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
arr'
        Maybe Value
Nothing ->
          forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            Text
"Index out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [((Int64
n1', Int64
s1'), (Int64
n2', Int64
s2'))]
    --
    def [Char]
"flat_update_2d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun5 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
offset Value
s1 Value
s2 Value
v -> do
      let offset' :: Int64
offset' = Value -> Int64
asInt64 Value
offset
          s1' :: Int64
s1' = Value -> Int64
asInt64 Value
s1
          s2' :: Int64
s2' = Value -> Int64
asInt64 Value
s2
      case forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v of
        ShapeDim Int64
n1 (ShapeDim Int64
n2 ValueShape
_) -> do
          let iota :: a -> [a]
iota a
x = [a
0 .. a
x forall a. Num a => a -> a -> a
- a
1]
              f :: Value -> (Int64, Int64) -> Maybe Value
f Value
arr' (Int64
i, Int64
j) =
                [Indexing] -> Value -> Value -> Maybe Value
writeArray [Int64 -> Indexing
IndexingFix forall a b. (a -> b) -> a -> b
$ Int64
offset' forall a. Num a => a -> a -> a
+ Int64
i forall a. Num a => a -> a -> a
* Int64
s1' forall a. Num a => a -> a -> a
+ Int64
j forall a. Num a => a -> a -> a
* Int64
s2'] Value
arr'
                  forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> Maybe Value
indexArray [Int64 -> Indexing
IndexingFix Int64
i, Int64 -> Indexing
IndexingFix Int64
j] Value
v
          case forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Value -> (Int64, Int64) -> Maybe Value
f Value
arr [(Int64
i, Int64
j) | Int64
i <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n1, Int64
j <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n2] of
            Just Value
arr' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
arr'
            Maybe Value
Nothing ->
              forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                Text
"Index out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [((Int64
n1, Int64
s1'), (Int64
n2, Int64
s2'))]
        ValueShape
s -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"flat_update_2d: invalid arg shape: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show ValueShape
s
    --
    def [Char]
"flat_index_3d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun8 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
offset Value
n1 Value
s1 Value
n2 Value
s2 Value
n3 Value
s3 -> do
      let offset' :: Int64
offset' = Value -> Int64
asInt64 Value
offset
          n1' :: Int64
n1' = Value -> Int64
asInt64 Value
n1
          n2' :: Int64
n2' = Value -> Int64
asInt64 Value
n2
          n3' :: Int64
n3' = Value -> Int64
asInt64 Value
n3
          s1' :: Int64
s1' = Value -> Int64
asInt64 Value
s1
          s2' :: Int64
s2' = Value -> Int64
asInt64 Value
s2
          s3' :: Int64
s3' = Value -> Int64
asInt64 Value
s3
          shapeFromDims :: [a] -> Shape a
shapeFromDims = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall d. d -> Shape d -> Shape d
ShapeDim forall d. Shape d
ShapeLeaf
          mk1 :: [Maybe (Value m)] -> Maybe (Value m)
mk1 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (forall {a}. [a] -> Shape a
shapeFromDims [Int64
n1', Int64
n2', Int64
n3'])) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          mk2 :: [Maybe (Value m)] -> Maybe (Value m)
mk2 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> Shape a
shapeFromDims [Int64
n2', Int64
n3']) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          mk3 :: [Maybe (Value m)] -> Maybe (Value m)
mk3 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> Shape a
shapeFromDims [Int64
n3']) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          iota :: a -> [a]
iota a
x = [a
0 .. a
x forall a. Num a => a -> a -> a
- a
1]
          f :: Int64 -> Int64 -> Int64 -> Maybe Value
f Int64
i Int64
j Int64
l =
            [Indexing] -> Value -> Maybe Value
indexArray [Int64 -> Indexing
IndexingFix forall a b. (a -> b) -> a -> b
$ Int64
offset' forall a. Num a => a -> a -> a
+ Int64
i forall a. Num a => a -> a -> a
* Int64
s1' forall a. Num a => a -> a -> a
+ Int64
j forall a. Num a => a -> a -> a
* Int64
s2' forall a. Num a => a -> a -> a
+ Int64
l forall a. Num a => a -> a -> a
* Int64
s3'] Value
arr

      case forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk1 [forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk2 [forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk3 [Int64 -> Int64 -> Int64 -> Maybe Value
f Int64
i Int64
j Int64
l | Int64
l <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n3'] | Int64
j <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n2'] | Int64
i <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n1'] of
        Just Value
arr' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
arr'
        Maybe Value
Nothing ->
          forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            Text
"Index out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [((Int64
n1', Int64
s1'), (Int64
n2', Int64
s2'), (Int64
n3', Int64
s3'))]
    --
    def [Char]
"flat_update_3d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
offset Value
s1 Value
s2 Value
s3 Value
v -> do
      let offset' :: Int64
offset' = Value -> Int64
asInt64 Value
offset
          s1' :: Int64
s1' = Value -> Int64
asInt64 Value
s1
          s2' :: Int64
s2' = Value -> Int64
asInt64 Value
s2
          s3' :: Int64
s3' = Value -> Int64
asInt64 Value
s3
      case forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v of
        ShapeDim Int64
n1 (ShapeDim Int64
n2 (ShapeDim Int64
n3 ValueShape
_)) -> do
          let iota :: a -> [a]
iota a
x = [a
0 .. a
x forall a. Num a => a -> a -> a
- a
1]
              f :: Value -> (Int64, Int64, Int64) -> Maybe Value
f Value
arr' (Int64
i, Int64
j, Int64
l) =
                [Indexing] -> Value -> Value -> Maybe Value
writeArray [Int64 -> Indexing
IndexingFix forall a b. (a -> b) -> a -> b
$ Int64
offset' forall a. Num a => a -> a -> a
+ Int64
i forall a. Num a => a -> a -> a
* Int64
s1' forall a. Num a => a -> a -> a
+ Int64
j forall a. Num a => a -> a -> a
* Int64
s2' forall a. Num a => a -> a -> a
+ Int64
l forall a. Num a => a -> a -> a
* Int64
s3'] Value
arr'
                  forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> Maybe Value
indexArray [Int64 -> Indexing
IndexingFix Int64
i, Int64 -> Indexing
IndexingFix Int64
j, Int64 -> Indexing
IndexingFix Int64
l] Value
v
          case forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Value -> (Int64, Int64, Int64) -> Maybe Value
f Value
arr [(Int64
i, Int64
j, Int64
l) | Int64
i <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n1, Int64
j <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n2, Int64
l <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n3] of
            Just Value
arr' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
arr'
            Maybe Value
Nothing ->
              forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                Text
"Index out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [((Int64
n1, Int64
s1'), (Int64
n2, Int64
s2'), (Int64
n3, Int64
s3'))]
        ValueShape
s -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"flat_update_3d: invalid arg shape: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show ValueShape
s
    --
    def [Char]
"flat_index_4d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun10 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
offset Value
n1 Value
s1 Value
n2 Value
s2 Value
n3 Value
s3 Value
n4 Value
s4 -> do
      let offset' :: Int64
offset' = Value -> Int64
asInt64 Value
offset
          n1' :: Int64
n1' = Value -> Int64
asInt64 Value
n1
          n2' :: Int64
n2' = Value -> Int64
asInt64 Value
n2
          n3' :: Int64
n3' = Value -> Int64
asInt64 Value
n3
          n4' :: Int64
n4' = Value -> Int64
asInt64 Value
n4
          s1' :: Int64
s1' = Value -> Int64
asInt64 Value
s1
          s2' :: Int64
s2' = Value -> Int64
asInt64 Value
s2
          s3' :: Int64
s3' = Value -> Int64
asInt64 Value
s3
          s4' :: Int64
s4' = Value -> Int64
asInt64 Value
s4
          shapeFromDims :: [a] -> Shape a
shapeFromDims = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall d. d -> Shape d -> Shape d
ShapeDim forall d. Shape d
ShapeLeaf
          mk1 :: [Maybe (Value m)] -> Maybe (Value m)
mk1 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (forall {a}. [a] -> Shape a
shapeFromDims [Int64
n1', Int64
n2', Int64
n3', Int64
n4'])) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          mk2 :: [Maybe (Value m)] -> Maybe (Value m)
mk2 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> Shape a
shapeFromDims [Int64
n2', Int64
n3', Int64
n4']) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          mk3 :: [Maybe (Value m)] -> Maybe (Value m)
mk3 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> Shape a
shapeFromDims [Int64
n3', Int64
n4']) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          mk4 :: [Maybe (Value m)] -> Maybe (Value m)
mk4 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> Shape a
shapeFromDims [Int64
n4']) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
          iota :: a -> [a]
iota a
x = [a
0 .. a
x forall a. Num a => a -> a -> a
- a
1]
          f :: Int64 -> Int64 -> Int64 -> Int64 -> Maybe Value
f Int64
i Int64
j Int64
l Int64
m =
            [Indexing] -> Value -> Maybe Value
indexArray [Int64 -> Indexing
IndexingFix forall a b. (a -> b) -> a -> b
$ Int64
offset' forall a. Num a => a -> a -> a
+ Int64
i forall a. Num a => a -> a -> a
* Int64
s1' forall a. Num a => a -> a -> a
+ Int64
j forall a. Num a => a -> a -> a
* Int64
s2' forall a. Num a => a -> a -> a
+ Int64
l forall a. Num a => a -> a -> a
* Int64
s3' forall a. Num a => a -> a -> a
+ Int64
m forall a. Num a => a -> a -> a
* Int64
s4'] Value
arr

      case forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk1 [forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk2 [forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk3 [forall {m :: * -> *}. [Maybe (Value m)] -> Maybe (Value m)
mk4 [Int64 -> Int64 -> Int64 -> Int64 -> Maybe Value
f Int64
i Int64
j Int64
l Int64
m | Int64
m <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n4'] | Int64
l <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n3'] | Int64
j <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n2'] | Int64
i <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n1'] of
        Just Value
arr' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
arr'
        Maybe Value
Nothing ->
          forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            Text
"Index out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [(((Int64
n1', Int64
s1'), (Int64
n2', Int64
s2')), ((Int64
n3', Int64
s3'), (Int64
n4', Int64
s4')))]
    --
    def [Char]
"flat_update_4d" = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun7 forall a b. (a -> b) -> a -> b
$ \Value
arr Value
offset Value
s1 Value
s2 Value
s3 Value
s4 Value
v -> do
      let offset' :: Int64
offset' = Value -> Int64
asInt64 Value
offset
          s1' :: Int64
s1' = Value -> Int64
asInt64 Value
s1
          s2' :: Int64
s2' = Value -> Int64
asInt64 Value
s2
          s3' :: Int64
s3' = Value -> Int64
asInt64 Value
s3
          s4' :: Int64
s4' = Value -> Int64
asInt64 Value
s4
      case forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v of
        ShapeDim Int64
n1 (ShapeDim Int64
n2 (ShapeDim Int64
n3 (ShapeDim Int64
n4 ValueShape
_))) -> do
          let iota :: a -> [a]
iota a
x = [a
0 .. a
x forall a. Num a => a -> a -> a
- a
1]
              f :: Value -> (Int64, Int64, Int64, Int64) -> Maybe Value
f Value
arr' (Int64
i, Int64
j, Int64
l, Int64
m) =
                [Indexing] -> Value -> Value -> Maybe Value
writeArray [Int64 -> Indexing
IndexingFix forall a b. (a -> b) -> a -> b
$ Int64
offset' forall a. Num a => a -> a -> a
+ Int64
i forall a. Num a => a -> a -> a
* Int64
s1' forall a. Num a => a -> a -> a
+ Int64
j forall a. Num a => a -> a -> a
* Int64
s2' forall a. Num a => a -> a -> a
+ Int64
l forall a. Num a => a -> a -> a
* Int64
s3' forall a. Num a => a -> a -> a
+ Int64
m forall a. Num a => a -> a -> a
* Int64
s4'] Value
arr'
                  forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> Maybe Value
indexArray [Int64 -> Indexing
IndexingFix Int64
i, Int64 -> Indexing
IndexingFix Int64
j, Int64 -> Indexing
IndexingFix Int64
l, Int64 -> Indexing
IndexingFix Int64
m] Value
v
          case forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Value -> (Int64, Int64, Int64, Int64) -> Maybe Value
f Value
arr [(Int64
i, Int64
j, Int64
l, Int64
m) | Int64
i <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n1, Int64
j <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n2, Int64
l <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n3, Int64
m <- forall {a}. (Num a, Enum a) => a -> [a]
iota Int64
n4] of
            Just Value
arr' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
arr'
            Maybe Value
Nothing ->
              forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                Text
"Index out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [(((Int64
n1, Int64
s1'), (Int64
n2, Int64
s2')), ((Int64
n3, Int64
s3'), (Int64
n4, Int64
s4')))]
        ValueShape
s -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"flat_update_4d: invalid arg shape: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show ValueShape
s
    --
    def [Char]
"unzip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
x -> do
        let ShapeDim Int64
_ (ShapeRecord Map Name ValueShape
fs) = forall (m :: * -> *). Value m -> ValueShape
valueShape Value
x
            Just [ValueShape
xs_shape, ValueShape
ys_shape] = forall a. Map Name a -> Maybe [a]
areTupleFields Map Name ValueShape
fs
            listPair :: ([Value m], [Value m]) -> [Value m]
listPair ([Value m]
xs, [Value m]
ys) =
              [forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
xs_shape [Value m]
xs, forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
ys_shape [Value m]
ys]

        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). [Value m] -> Value m
toTuple forall a b. (a -> b) -> a -> b
$ forall {m :: * -> *}. ([Value m], [Value m]) -> [Value m]
listPair forall a b. (a -> b) -> a -> b
$ forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall {b}. Maybe [b] -> (b, b)
fromPair forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple) forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
x
      where
        fromPair :: Maybe [b] -> (b, b)
fromPair (Just [b
x, b
y]) = (b
x, b
y)
        fromPair Maybe [b]
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"Not a pair"
    def [Char]
"zip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$ \Value
xs Value
ys -> do
        let ShapeDim Int64
_ ValueShape
xs_rowshape = forall (m :: * -> *). Value m -> ValueShape
valueShape Value
xs
            ShapeDim Int64
_ ValueShape
ys_rowshape = forall (m :: * -> *). Value m -> ValueShape
valueShape Value
ys
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
          forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' (forall d. Map Name (Shape d) -> Shape d
ShapeRecord (forall a. [a] -> Map Name a
tupleFields [ValueShape
xs_rowshape, ValueShape
ys_rowshape])) forall a b. (a -> b) -> a -> b
$
            forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). [Value m] -> Value m
toTuple forall a b. (a -> b) -> a -> b
$
              forall a. [[a]] -> [[a]]
transpose [forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs, forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
ys]
    def [Char]
"concat" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$ \Value
xs Value
ys -> do
        let (ShapeDim Int64
_ ValueShape
rowshape, [Value]
xs') = Value -> (ValueShape, [Value])
fromArray Value
xs
            (ValueShape
_, [Value]
ys') = Value -> (ValueShape, [Value])
fromArray Value
ys
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape forall a b. (a -> b) -> a -> b
$ [Value]
xs' forall a. [a] -> [a] -> [a]
++ [Value]
ys'
    def [Char]
"transpose" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
xs -> do
        let (ShapeDim Int64
n (ShapeDim Int64
m ValueShape
shape), [Value]
xs') = Value -> (ValueShape, [Value])
fromArray Value
xs
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
          forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (forall d. d -> Shape d -> Shape d
ShapeDim Int64
m (forall d. d -> Shape d -> Shape d
ShapeDim Int64
n ValueShape
shape)) forall a b. (a -> b) -> a -> b
$
            forall a b. (a -> b) -> [a] -> [b]
map (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (forall d. d -> Shape d -> Shape d
ShapeDim Int64
n ValueShape
shape)) forall a b. (a -> b) -> a -> b
$
              -- Slight hack to work around empty dimensions.
              forall i a. Integral i => i -> [a] -> [a]
genericTake Int64
m forall a b. (a -> b) -> a -> b
$
                forall a. [[a]] -> [[a]]
transpose (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> (ValueShape, [Value])
fromArray) [Value]
xs') forall a. [a] -> [a] -> [a]
++ forall a. a -> [a]
repeat []
    def [Char]
"rotate" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 forall a b. (a -> b) -> a -> b
$ \Value
i Value
xs -> do
        let (ValueShape
shape, [Value]
xs') = Value -> (ValueShape, [Value])
fromArray Value
xs
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
          let idx :: Int
idx = if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Value]
xs' then Int
0 else forall a. Integral a => a -> a -> a
rem (Value -> Int
asInt Value
i) (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Value]
xs')
           in if Int
idx forall a. Ord a => a -> a -> Bool
> Int
0
                then
                  let ([Value]
bef, [Value]
aft) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
idx [Value]
xs'
                   in forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
shape forall a b. (a -> b) -> a -> b
$ [Value]
aft forall a. [a] -> [a] -> [a]
++ [Value]
bef
                else
                  let ([Value]
bef, [Value]
aft) = forall a. Int -> [a] -> ([a], [a])
splitFromEnd (-Int
idx) [Value]
xs'
                   in forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
shape forall a b. (a -> b) -> a -> b
$ [Value]
aft forall a. [a] -> [a] -> [a]
++ [Value]
bef
    def [Char]
"flatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> EvalM Value) -> TermBinding
fun1 forall a b. (a -> b) -> a -> b
$ \Value
xs -> do
        let (ShapeDim Int64
n (ShapeDim Int64
m ValueShape
shape), [Value]
xs') = Value -> (ValueShape, [Value])
fromArray Value
xs
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (forall d. d -> Shape d -> Shape d
ShapeDim (Int64
n forall a. Num a => a -> a -> a
* Int64
m) ValueShape
shape) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> (ValueShape, [Value])
fromArray) [Value]
xs'
    def [Char]
"unflatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$ \Value
n Value
m Value
xs -> do
        let (ShapeDim Int64
xs_size ValueShape
innershape, [Value]
xs') = Value -> (ValueShape, [Value])
fromArray Value
xs
            rowshape :: ValueShape
rowshape = forall d. d -> Shape d -> Shape d
ShapeDim (Value -> Int64
asInt64 Value
m) ValueShape
innershape
            shape :: ValueShape
shape = forall d. d -> Shape d -> Shape d
ShapeDim (Value -> Int64
asInt64 Value
n) ValueShape
rowshape
        if Value -> Int64
asInt64 Value
n forall a. Num a => a -> a -> a
* Value -> Int64
asInt64 Value
m forall a. Eq a => a -> a -> Bool
/= Int64
xs_size
          then
            forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
              Text
"Cannot unflatten array of shape ["
                forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText Int64
xs_size
                forall a. Semigroup a => a -> a -> a
<> Text
"] to array of shape ["
                forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText (Value -> Int64
asInt64 Value
n)
                forall a. Semigroup a => a -> a -> a
<> Text
"]["
                forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText (Value -> Int64
asInt64 Value
m)
                forall a. Semigroup a => a -> a -> a
<> Text
"]"
          else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
shape forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
rowshape) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [[a]]
chunk (Value -> Int
asInt Value
m) [Value]
xs'
    def [Char]
"vjp2" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$
        \Value
_ Value
_ Value
_ -> forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Text
"Interpreter does not support autodiff."
    def [Char]
"jvp2" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 forall a b. (a -> b) -> a -> b
$
        \Value
_ Value
_ Value
_ -> forall a. SrcLoc -> Env -> Text -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Text
"Interpreter does not support autodiff."
    def [Char]
"acc" = forall a. Maybe a
Nothing
    def [Char]
s | [Char] -> Name
nameFromString [Char]
s forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map Name PrimType
namesToPrimTypes = forall a. Maybe a
Nothing
    def [Char]
s = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Missing intrinsic: " forall a. [a] -> [a] -> [a]
++ [Char]
s

    tdef :: [Char] -> Maybe TypeBinding
tdef [Char]
s = do
      PrimType
t <- [Char] -> Name
nameFromString [Char]
s forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map Name PrimType
namesToPrimTypes
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Liftedness -> [TypeParam] -> RetTypeBase Size () -> TypeBinding
T.TypeAbbr Liftedness
Unlifted [] forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t

    stream :: Value -> Value -> EvalM Value
stream Value
f arg :: Value
arg@(ValueArray ValueShape
_ Array Int Value
xs) =
      let n :: Value m
n = forall (m :: * -> *). PrimValue -> Value m
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value forall a b. (a -> b) -> a -> b
$ forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
xs
       in SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value
apply2 forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty Value
f forall {m :: * -> *}. Value m
n Value
arg
    stream Value
_ Value
arg = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot stream: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Value
arg

interpretExp :: Ctx -> Exp -> F ExtOp Value
interpretExp :: Ctx -> Exp -> F ExtOp Value
interpretExp Ctx
ctx Exp
e = forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map ImportName Env
ctxImports Ctx
ctx) forall a b. (a -> b) -> a -> b
$ Env -> Exp -> EvalM Value
eval (Ctx -> Env
ctxEnv Ctx
ctx) Exp
e

interpretDec :: Ctx -> Dec -> F ExtOp Ctx
interpretDec :: Ctx -> DecBase Info VName -> F ExtOp Ctx
interpretDec Ctx
ctx DecBase Info VName
d = do
  Env
env <- forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map ImportName Env
ctxImports Ctx
ctx) forall a b. (a -> b) -> a -> b
$ do
    Env
env <- Env -> DecBase Info VName -> EvalM Env
evalDec (Ctx -> Env
ctxEnv Ctx
ctx) DecBase Info VName
d
    -- We need to extract any new existential sizes and add them as
    -- ordinary bindings to the context, or we will not be able to
    -- look up their values later.
    Env
sizes <- EvalM Env
extSizeEnv
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Env
env forall a. Semigroup a => a -> a -> a
<> Env
sizes
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx
ctx {ctxEnv :: Env
ctxEnv = Env
env}

interpretImport :: Ctx -> (ImportName, Prog) -> F ExtOp Ctx
interpretImport :: Ctx -> (ImportName, Prog) -> F ExtOp Ctx
interpretImport Ctx
ctx (ImportName
fp, Prog
prog) = do
  Env
env <- forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map ImportName Env
ctxImports Ctx
ctx) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Env -> DecBase Info VName -> EvalM Env
evalDec (Ctx -> Env
ctxEnv Ctx
ctx) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs Prog
prog
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx
ctx {ctxImports :: Map ImportName Env
ctxImports = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ImportName
fp Env
env forall a b. (a -> b) -> a -> b
$ Ctx -> Map ImportName Env
ctxImports Ctx
ctx}

-- | Produce a context, based on the one passed in, where all of
-- the provided imports have been @open@ened in order.
ctxWithImports :: [Env] -> Ctx -> Ctx
ctxWithImports :: [Env] -> Ctx -> Ctx
ctxWithImports [Env]
envs Ctx
ctx = Ctx
ctx {ctxEnv :: Env
ctxEnv = forall a. Monoid a => [a] -> a
mconcat (forall a. [a] -> [a]
reverse [Env]
envs) forall a. Semigroup a => a -> a -> a
<> Ctx -> Env
ctxEnv Ctx
ctx}

valueType :: V.Value -> ValueType
valueType :: Value -> ValueType
valueType Value
v =
  let V.ValueType [Int]
shape PrimType
pt = Value -> ValueType
V.valueType Value
v
   in forall as dim.
Monoid as =>
Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOf forall a. Monoid a => a
mempty (forall dim. [dim] -> Shape dim
F.Shape (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
shape)) (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> PrimType
toPrim PrimType
pt)))
  where
    toPrim :: PrimType -> PrimType
toPrim PrimType
V.I8 = IntType -> PrimType
Signed IntType
Int8
    toPrim PrimType
V.I16 = IntType -> PrimType
Signed IntType
Int16
    toPrim PrimType
V.I32 = IntType -> PrimType
Signed IntType
Int32
    toPrim PrimType
V.I64 = IntType -> PrimType
Signed IntType
Int64
    toPrim PrimType
V.U8 = IntType -> PrimType
Unsigned IntType
Int8
    toPrim PrimType
V.U16 = IntType -> PrimType
Unsigned IntType
Int16
    toPrim PrimType
V.U32 = IntType -> PrimType
Unsigned IntType
Int32
    toPrim PrimType
V.U64 = IntType -> PrimType
Unsigned IntType
Int64
    toPrim PrimType
V.Bool = PrimType
Bool
    toPrim PrimType
V.F16 = FloatType -> PrimType
FloatType FloatType
Float16
    toPrim PrimType
V.F32 = FloatType -> PrimType
FloatType FloatType
Float32
    toPrim PrimType
V.F64 = FloatType -> PrimType
FloatType FloatType
Float64

checkEntryArgs :: VName -> [V.Value] -> StructType -> Either T.Text ()
checkEntryArgs :: VName -> [Value] -> StructType -> Either Text ()
checkEntryArgs VName
entry [Value]
args StructType
entry_t
  | [StructType]
args_ts forall a. Eq a => a -> a -> Bool
== forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Diet, StructType)]
param_ts =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise =
      forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
        forall {ann}. Doc ann
expected
          forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc Any
"Got input of types"
          forall ann. Doc ann -> Doc ann -> Doc ann
</> forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall a. [Doc a] -> Doc a
stack (forall a b. (a -> b) -> [a] -> [b]
map forall a ann. Pretty a => a -> Doc ann
pretty [StructType]
args_ts))
  where
    ([(Diet, StructType)]
param_ts, StructType
_) = forall dim as.
TypeBase dim as -> ([(Diet, TypeBase dim ())], TypeBase dim ())
unfoldFunType StructType
entry_t
    args_ts :: [StructType]
args_ts = forall a b. (a -> b) -> [a] -> [b]
map (ValueType -> StructType
valueStructType forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ValueType
valueType) [Value]
args
    expected :: Doc ann
expected
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Diet, StructType)]
param_ts =
          Doc ann
"Entry point " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
dquotes (forall v a. IsName v => v -> Doc a
prettyName VName
entry) forall a. Semigroup a => a -> a -> a
<> Doc ann
" is not a function."
      | Bool
otherwise =
          Doc ann
"Entry point " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
dquotes (forall v a. IsName v => v -> Doc a
prettyName VName
entry) forall a. Semigroup a => a -> a -> a
<> Doc ann
" expects input of type(s)"
            forall ann. Doc ann -> Doc ann -> Doc ann
</> forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (forall a. [Doc a] -> Doc a
stack (forall a b. (a -> b) -> [a] -> [b]
map forall a ann. Pretty a => a -> Doc ann
pretty [(Diet, StructType)]
param_ts))

-- | Execute the named function on the given arguments; may fail
-- horribly if these are ill-typed.
interpretFunction :: Ctx -> VName -> [V.Value] -> Either T.Text (F ExtOp Value)
interpretFunction :: Ctx -> VName -> [Value] -> Either Text (F ExtOp Value)
interpretFunction Ctx
ctx VName
fname [Value]
vs = do
  StructType
ft <- case QualName VName -> Env -> Maybe TermBinding
lookupVar (forall v. v -> QualName v
qualName VName
fname) forall a b. (a -> b) -> a -> b
$ Ctx -> Env
ctxEnv Ctx
ctx of
    Just (TermValue (Just (T.BoundV [TypeParam]
_ StructType
t)) Value
_) ->
      forall {as}.
[ValueType] -> TypeBase Size as -> Either Text (TypeBase Size as)
updateType (forall a b. (a -> b) -> [a] -> [b]
map Value -> ValueType
valueType [Value]
vs) StructType
t
    Just (TermPoly (Just (T.BoundV [TypeParam]
_ StructType
t)) StructType -> EvalM Value
_) ->
      forall {as}.
[ValueType] -> TypeBase Size as -> Either Text (TypeBase Size as)
updateType (forall a b. (a -> b) -> [a] -> [b]
map Value -> ValueType
valueType [Value]
vs) StructType
t
    Maybe TermBinding
_ ->
      forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ Text
"Unknown function `" forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText (forall v. IsName v => v -> Name
toName VName
fname) forall a. Semigroup a => a -> a -> a
<> Text
"`."

  let vs' :: [Value m]
vs' = forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Value -> Value m
fromDataValue [Value]
vs

  VName -> [Value] -> StructType -> Either Text ()
checkEntryArgs VName
fname [Value]
vs StructType
ft

  forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$
    forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map ImportName Env
ctxImports Ctx
ctx) forall a b. (a -> b) -> a -> b
$ do
      Value
f <- Env -> QualName VName -> StructType -> EvalM Value
evalTermVar (Ctx -> Env
ctxEnv Ctx
ctx) (forall v. v -> QualName v
qualName VName
fname) StructType
ft
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> EvalM Value
apply forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty) Value
f forall {m :: * -> *}. [Value m]
vs'
  where
    updateType :: [ValueType] -> TypeBase Size as -> Either Text (TypeBase Size as)
updateType (ValueType
vt : [ValueType]
vts) (Scalar (Arrow as
als PName
pn Diet
d StructType
pt (RetType [VName]
dims TypeBase Size as
rt))) = do
      ValueType -> StructType -> Either Text ()
checkInput ValueType
vt StructType
pt
      forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as.
as
-> PName
-> Diet
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
als PName
pn Diet
d (ValueType -> StructType
valueStructType ValueType
vt) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ValueType] -> TypeBase Size as -> Either Text (TypeBase Size as)
updateType [ValueType]
vts TypeBase Size as
rt
    updateType [ValueType]
_ TypeBase Size as
t =
      forall a b. b -> Either a b
Right TypeBase Size as
t

    -- FIXME: we don't check array sizes.
    checkInput :: ValueType -> StructType -> Either T.Text ()
    checkInput :: ValueType -> StructType -> Either Text ()
checkInput (Scalar (Prim PrimType
vt)) (Scalar (Prim PrimType
pt))
      | PrimType
vt forall a. Eq a => a -> a -> Bool
/= PrimType
pt = forall {a} {a} {b}. (Pretty a, Pretty a) => a -> a -> Either Text b
badPrim PrimType
vt PrimType
pt
    checkInput (Array ()
_ Uniqueness
_ Shape Int64
_ (Prim PrimType
vt)) (Array ()
_ Uniqueness
_ Shape Size
_ (Prim PrimType
pt))
      | PrimType
vt forall a. Eq a => a -> a -> Bool
/= PrimType
pt = forall {a} {a} {b}. (Pretty a, Pretty a) => a -> a -> Either Text b
badPrim PrimType
vt PrimType
pt
    checkInput ValueType
_ StructType
_ =
      forall a b. b -> Either a b
Right ()

    badPrim :: a -> a -> Either Text b
badPrim a
vt a
pt =
      forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$
        Doc Any
"Invalid argument type."
          forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc Any
"Expected:"
          forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty a
pt)
          forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc Any
"Got:     "
          forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty a
vt)