{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

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

import Control.Monad.Except
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,
    intercalate,
    isPrefixOf,
    transpose,
  )
import qualified Data.List.NonEmpty as NE
import qualified Data.Map as M
import Data.Maybe
import Data.Monoid hiding (Sum)
import Futhark.Util (chunk, maybeHead, splitFromEnd)
import Futhark.Util.Loc
import Futhark.Util.Pretty hiding (apply, bool)
import Language.Futhark hiding (Shape, Value, matchDims)
import qualified Language.Futhark as F
import Language.Futhark.Primitive (floatValue, intValue)
import qualified Language.Futhark.Primitive as P
import qualified Language.Futhark.Semantic 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 String String 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 FilePath
w FilePath
s a
x) = forall a. FilePath -> FilePath -> a -> ExtOp a
ExtOpTrace FilePath
w FilePath
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 FilePath 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 FilePath Env),
      MonadState Sizes
    )

runEvalM :: M.Map FilePath Env -> EvalM a -> F ExtOp a
runEvalM :: forall a. Map FilePath Env -> EvalM a -> F ExtOp a
runEvalM Map FilePath Env
imports (EvalM ReaderT ([StackFrame], Map FilePath 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 FilePath Env) (StateT Sizes (F ExtOp)) a
m (forall a. Monoid a => a
mempty, Map FilePath 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 FilePath Env
imports) ->
  if SrcLoc -> Bool
isNoLoc SrcLoc
loc
    then ([StackFrame]
ss, Map FilePath Env
imports)
    else
      let s :: StackFrame
s = Loc -> Ctx -> StackFrame
StackFrame (forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Env -> Map FilePath Env -> Ctx
Ctx Env
env Map FilePath Env
imports)
       in (StackFrame
s forall a. a -> [a] -> [a]
: [StackFrame]
ss, Map FilePath 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 :: FilePath -> EvalM (Maybe Env)
lookupImport :: FilePath -> EvalM (Maybe Env)
lookupImport FilePath
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 FilePath
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

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

prettyRecord :: Pretty a => M.Map Name a -> Doc
prettyRecord :: forall a. Pretty a => Map Name a -> Doc
prettyRecord Map Name a
m
  | Just [a]
vs <- forall a. Map Name a -> Maybe [a]
areTupleFields Map Name a
m =
      Doc -> Doc
parens forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commasep forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Doc
ppr [a]
vs
  | Bool
otherwise =
      Doc -> Doc
braces forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commasep forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall {a} {a}. (Pretty a, Pretty a) => (a, a) -> Doc
field forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map Name a
m
  where
    field :: (a, a) -> Doc
field (a
k, a
v) = forall a. Pretty a => a -> Doc
ppr a
k Doc -> Doc -> Doc
<+> Doc
equals Doc -> Doc -> Doc
<+> forall a. Pretty a => a -> Doc
ppr a
v

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

-- | A shape is a tree to accomodate the case of records.  It is
-- parameterised over the representation of dimensions.
data Shape d
  = ShapeDim d (Shape d)
  | ShapeLeaf
  | ShapeRecord (M.Map Name (Shape d))
  | ShapeSum (M.Map Name [Shape d])
  deriving (Shape d -> Shape d -> Bool
forall d. Eq d => Shape d -> Shape d -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Shape d -> Shape d -> Bool
$c/= :: forall d. Eq d => Shape d -> Shape d -> Bool
== :: Shape d -> Shape d -> Bool
$c== :: forall d. Eq d => Shape d -> Shape d -> Bool
Eq, Int -> Shape d -> ShowS
forall d. Show d => Int -> Shape d -> ShowS
forall d. Show d => [Shape d] -> ShowS
forall d. Show d => Shape d -> FilePath
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [Shape d] -> ShowS
$cshowList :: forall d. Show d => [Shape d] -> ShowS
show :: Shape d -> FilePath
$cshow :: forall d. Show d => Shape d -> FilePath
showsPrec :: Int -> Shape d -> ShowS
$cshowsPrec :: forall d. Show d => Int -> Shape d -> ShowS
Show, forall a b. a -> Shape b -> Shape a
forall a b. (a -> b) -> Shape a -> Shape 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 -> Shape b -> Shape a
$c<$ :: forall a b. a -> Shape b -> Shape a
fmap :: forall a b. (a -> b) -> Shape a -> Shape b
$cfmap :: forall a b. (a -> b) -> Shape a -> Shape b
Functor, forall a. Eq a => a -> Shape a -> Bool
forall a. Num a => Shape a -> a
forall a. Ord a => Shape a -> a
forall m. Monoid m => Shape m -> m
forall a. Shape a -> Bool
forall a. Shape a -> Int
forall a. Shape a -> [a]
forall a. (a -> a -> a) -> Shape a -> a
forall m a. Monoid m => (a -> m) -> Shape a -> m
forall b a. (b -> a -> b) -> b -> Shape a -> b
forall a b. (a -> b -> b) -> b -> Shape a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => Shape a -> a
$cproduct :: forall a. Num a => Shape a -> a
sum :: forall a. Num a => Shape a -> a
$csum :: forall a. Num a => Shape a -> a
minimum :: forall a. Ord a => Shape a -> a
$cminimum :: forall a. Ord a => Shape a -> a
maximum :: forall a. Ord a => Shape a -> a
$cmaximum :: forall a. Ord a => Shape a -> a
elem :: forall a. Eq a => a -> Shape a -> Bool
$celem :: forall a. Eq a => a -> Shape a -> Bool
length :: forall a. Shape a -> Int
$clength :: forall a. Shape a -> Int
null :: forall a. Shape a -> Bool
$cnull :: forall a. Shape a -> Bool
toList :: forall a. Shape a -> [a]
$ctoList :: forall a. Shape a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Shape a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Shape a -> a
foldr1 :: forall a. (a -> a -> a) -> Shape a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> Shape a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> Shape a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Shape a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Shape a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Shape a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Shape a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Shape a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Shape a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> Shape a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> Shape a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Shape a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Shape a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Shape a -> m
fold :: forall m. Monoid m => Shape m -> m
$cfold :: forall m. Monoid m => Shape m -> m
Foldable, Functor Shape
Foldable Shape
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Shape (m a) -> m (Shape a)
forall (f :: * -> *) a. Applicative f => Shape (f a) -> f (Shape a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Shape a -> m (Shape b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Shape a -> f (Shape b)
sequence :: forall (m :: * -> *) a. Monad m => Shape (m a) -> m (Shape a)
$csequence :: forall (m :: * -> *) a. Monad m => Shape (m a) -> m (Shape a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Shape a -> m (Shape b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Shape a -> m (Shape b)
sequenceA :: forall (f :: * -> *) a. Applicative f => Shape (f a) -> f (Shape a)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Shape (f a) -> f (Shape a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Shape a -> f (Shape b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Shape a -> f (Shape b)
Traversable)

-- | The shape of an array.
type ValueShape = Shape Int64

instance Pretty d => Pretty (Shape d) where
  ppr :: Shape d -> Doc
ppr Shape d
ShapeLeaf = forall a. Monoid a => a
mempty
  ppr (ShapeDim d
d Shape d
s) = Doc -> Doc
brackets (forall a. Pretty a => a -> Doc
ppr d
d) forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Doc
ppr Shape d
s
  ppr (ShapeRecord Map Name (Shape d)
m) = forall a. Pretty a => Map Name a -> Doc
prettyRecord Map Name (Shape d)
m
  ppr (ShapeSum Map Name [Shape d]
cs) =
    forall a. Monoid a => [a] -> a
mconcat (Doc -> [Doc] -> [Doc]
punctuate (FilePath -> Doc
text FilePath
" | ") [Doc]
cs')
    where
      ppConstr :: (a, [a]) -> Doc
ppConstr (a
name, [a]
fs) = [Doc] -> Doc
sep forall a b. (a -> b) -> a -> b
$ (FilePath -> Doc
text FilePath
"#" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Doc
ppr a
name) forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Doc
ppr [a]
fs
      cs' :: [Doc]
cs' = forall a b. (a -> b) -> [a] -> [b]
map forall {a} {a}. (Pretty a, Pretty a) => (a, [a]) -> Doc
ppConstr forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map Name [Shape d]
cs

emptyShape :: ValueShape -> Bool
emptyShape :: ValueShape -> Bool
emptyShape (ShapeDim Int64
d ValueShape
s) = Int64
d forall a. Eq a => a -> a -> Bool
== Int64
0 Bool -> Bool -> Bool
|| ValueShape -> Bool
emptyShape ValueShape
s
emptyShape ValueShape
_ = Bool
False

typeShape :: M.Map VName (Shape d) -> TypeBase d () -> Shape d
typeShape :: forall d. Map VName (Shape d) -> TypeBase d () -> Shape d
typeShape Map VName (Shape d)
shapes = TypeBase d () -> Shape d
go
  where
    go :: TypeBase d () -> Shape d
go (Array ()
_ Uniqueness
_ Shape d
shape ScalarTypeBase d ()
et) =
      forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall d. d -> Shape d -> Shape d
ShapeDim (TypeBase d () -> Shape d
go (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase d ()
et)) forall a b. (a -> b) -> a -> b
$ forall dim. Shape dim -> [dim]
shapeDims Shape d
shape
    go (Scalar (Record Map Name (TypeBase d ())
fs)) =
      forall d. Map Name (Shape d) -> Shape d
ShapeRecord forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase d () -> Shape d
go Map Name (TypeBase d ())
fs
    go (Scalar (Sum Map Name [TypeBase d ()]
cs)) =
      forall d. Map Name [Shape d] -> Shape d
ShapeSum forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map TypeBase d () -> Shape d
go) Map Name [TypeBase d ()]
cs
    go (Scalar (TypeVar ()
_ Uniqueness
_ (QualName [] VName
v) []))
      | Just Shape d
shape <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (Shape d)
shapes =
          Shape d
shape
    go TypeBase d ()
_ =
      forall d. Shape d
ShapeLeaf

structTypeShape :: M.Map VName ValueShape -> StructType -> Shape (Maybe Int64)
structTypeShape :: Map VName ValueShape -> StructType -> Shape (Maybe Int64)
structTypeShape Map VName ValueShape
shapes = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {a}. Num a => Size -> Maybe a
dim forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. Map VName (Shape d) -> TypeBase d () -> Shape d
typeShape Map VName (Shape Size)
shapes'
  where
    dim :: Size -> Maybe a
dim (ConstSize Int
d) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d
    dim Size
_ = forall a. Maybe a
Nothing
    shapes' :: Map VName (Shape Size)
shapes' = forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ Int -> Size
ConstSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) Map VName ValueShape
shapes

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
_ StructType
poly_t1 (RetType [VName]
_ StructType
poly_t2)))
      (Scalar (Arrow ()
_ PName
_ 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 Int
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 Int
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

-- | A fully evaluated Futhark value.
data Value
  = ValuePrim !PrimValue
  | ValueArray ValueShape !(Array Int Value)
  | -- Stores the full shape.
    ValueRecord (M.Map Name Value)
  | ValueFun (Value -> EvalM Value)
  | -- Stores the full shape.
    ValueSum ValueShape Name [Value]
  | -- The update function and the array.
    ValueAcc (Value -> Value -> EvalM Value) !(Array Int Value)

instance Eq Value where
  ValuePrim (SignedValue IntValue
x) == :: Value -> Value -> Bool
== ValuePrim (SignedValue IntValue
y) =
    PrimValue -> PrimValue -> Bool
P.doCmpEq (IntValue -> PrimValue
P.IntValue IntValue
x) (IntValue -> PrimValue
P.IntValue IntValue
y)
  ValuePrim (UnsignedValue IntValue
x) == ValuePrim (UnsignedValue IntValue
y) =
    PrimValue -> PrimValue -> Bool
P.doCmpEq (IntValue -> PrimValue
P.IntValue IntValue
x) (IntValue -> PrimValue
P.IntValue IntValue
y)
  ValuePrim (FloatValue FloatValue
x) == ValuePrim (FloatValue FloatValue
y) =
    PrimValue -> PrimValue -> Bool
P.doCmpEq (FloatValue -> PrimValue
P.FloatValue FloatValue
x) (FloatValue -> PrimValue
P.FloatValue FloatValue
y)
  ValuePrim (BoolValue Bool
x) == ValuePrim (BoolValue Bool
y) =
    PrimValue -> PrimValue -> Bool
P.doCmpEq (Bool -> PrimValue
P.BoolValue Bool
x) (Bool -> PrimValue
P.BoolValue Bool
y)
  ValueArray ValueShape
_ Array Int Value
x == ValueArray ValueShape
_ Array Int Value
y = Array Int Value
x forall a. Eq a => a -> a -> Bool
== Array Int Value
y
  ValueRecord Map Name Value
x == ValueRecord Map Name Value
y = Map Name Value
x forall a. Eq a => a -> a -> Bool
== Map Name Value
y
  ValueSum ValueShape
_ Name
n1 [Value]
vs1 == ValueSum ValueShape
_ Name
n2 [Value]
vs2 = Name
n1 forall a. Eq a => a -> a -> Bool
== Name
n2 Bool -> Bool -> Bool
&& [Value]
vs1 forall a. Eq a => a -> a -> Bool
== [Value]
vs2
  ValueAcc Value -> Value -> EvalM Value
_ Array Int Value
x == ValueAcc Value -> Value -> EvalM Value
_ Array Int Value
y = Array Int Value
x forall a. Eq a => a -> a -> Bool
== Array Int Value
y
  Value
_ == Value
_ = Bool
False

instance Pretty Value where
  ppr :: Value -> Doc
ppr = forall a. Pretty a => Int -> a -> Doc
pprPrec Int
0
  pprPrec :: Int -> Value -> Doc
pprPrec Int
_ (ValuePrim PrimValue
v) = forall a. Pretty a => a -> Doc
ppr PrimValue
v
  pprPrec Int
_ (ValueArray ValueShape
_ Array Int Value
a) =
    let elements :: [Value]
elements = forall i e. Array i e -> [e]
elems Array Int Value
a -- [Value]
        (Value
x : [Value]
_) = [Value]
elements
        separator :: Doc
separator = case Value
x of
          ValueArray ValueShape
_ Array Int Value
_ -> Doc
comma forall a. Semigroup a => a -> a -> a
<> Doc
line
          Value
_ -> Doc
comma forall a. Semigroup a => a -> a -> a
<> Doc
space
     in Doc -> Doc
brackets forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
cat forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
punctuate Doc
separator (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Doc
ppr [Value]
elements)
  pprPrec Int
_ (ValueRecord Map Name Value
m) = forall a. Pretty a => Map Name a -> Doc
prettyRecord Map Name Value
m
  pprPrec Int
_ ValueFun {} = FilePath -> Doc
text FilePath
"#<fun>"
  pprPrec Int
_ ValueAcc {} = FilePath -> Doc
text FilePath
"#<acc>"
  pprPrec Int
p (ValueSum ValueShape
_ Name
n [Value]
vs) =
    Bool -> Doc -> Doc
parensIf (Int
p forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ FilePath -> Doc
text FilePath
"#" forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
sep (forall a. Pretty a => a -> Doc
ppr Name
n forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (forall a. Pretty a => Int -> a -> Doc
pprPrec Int
1) [Value]
vs)

valueShape :: Value -> ValueShape
valueShape :: Value -> ValueShape
valueShape (ValueArray ValueShape
shape Array Int Value
_) = ValueShape
shape
valueShape (ValueRecord Map Name Value
fs) = forall d. Map Name (Shape d) -> Shape d
ShapeRecord forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map Value -> ValueShape
valueShape Map Name Value
fs
valueShape (ValueSum ValueShape
shape Name
_ [Value]
_) = ValueShape
shape
valueShape Value
_ = forall d. Shape d
ShapeLeaf

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

-- | Does the value correspond to an empty array?
isEmptyArray :: Value -> Bool
isEmptyArray :: Value -> Bool
isEmptyArray = ValueShape -> Bool
emptyShape forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ValueShape
valueShape

-- | String representation of an empty array with the provided element
-- type.  This is pretty ad-hoc - don't expect good results unless the
-- element type is a primitive.
prettyEmptyArray :: TypeBase () () -> Value -> String
prettyEmptyArray :: TypeBase () () -> Value -> FilePath
prettyEmptyArray TypeBase () ()
t Value
v =
  FilePath
"empty(" forall a. [a] -> [a] -> [a]
++ forall {a}. Pretty a => Shape a -> FilePath
dims (Value -> ValueShape
valueShape Value
v) forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty TypeBase () ()
t' forall a. [a] -> [a] -> [a]
++ FilePath
")"
  where
    t' :: TypeBase () ()
t' = forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray (forall dim as. TypeBase dim as -> Int
arrayRank TypeBase () ()
t) TypeBase () ()
t
    dims :: Shape a -> FilePath
dims (ShapeDim a
n Shape a
rowshape) =
      FilePath
"[" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty a
n forall a. [a] -> [a] -> [a]
++ FilePath
"]" forall a. [a] -> [a] -> [a]
++ Shape a -> FilePath
dims Shape a
rowshape
    dims Shape a
_ = FilePath
""

-- | Create an array value; failing if that would result in an
-- irregular array.
mkArray :: TypeBase Int64 () -> [Value] -> Maybe Value
mkArray :: ValueType -> [Value] -> Maybe Value
mkArray ValueType
t [] =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
toArray (forall d. Map VName (Shape d) -> TypeBase d () -> Shape d
typeShape forall a. Monoid a => a
mempty ValueType
t) []
mkArray ValueType
_ (Value
v : [Value]
vs) = do
  let v_shape :: ValueShape
v_shape = Value -> ValueShape
valueShape Value
v
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((forall a. Eq a => a -> a -> Bool
== ValueShape
v_shape) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ValueShape
valueShape) [Value]
vs
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
toArray' ValueShape
v_shape forall a b. (a -> b) -> a -> b
$ Value
v forall a. a -> [a] -> [a]
: [Value]
vs

arrayLength :: Integral int => Array Int Value -> int
arrayLength :: forall int. Integral int => Array Int Value -> int
arrayLength = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
+ Int
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i e. Array i e -> (i, i)
bounds

toTuple :: [Value] -> Value
toTuple :: [Value] -> Value
toTuple = Map Name Value -> Value
ValueRecord forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames

fromTuple :: Value -> Maybe [Value]
fromTuple :: Value -> Maybe [Value]
fromTuple (ValueRecord Map Name Value
m) = forall a. Map Name a -> Maybe [a]
areTupleFields Map Name Value
m
fromTuple Value
_ = forall a. Maybe a
Nothing

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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Unexpectedly not an integer: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Unexpected not a signed integer: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Unexpectedly not a boolean: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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,
    -- | A mapping from type parameters to the shapes of
    -- the value to which they were initially bound.
    Env -> Map VName ValueShape
envShapes :: M.Map VName ValueShape
  }

instance Monoid Env where
  mempty :: Env
mempty = Map VName TermBinding
-> Map VName TypeBinding -> Map VName ValueShape -> Env
Env forall a. Monoid a => a
mempty 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 Map VName ValueShape
sm1 <> :: Env -> Env -> Env
<> Env Map VName TermBinding
vm2 Map VName TypeBinding
tm2 Map VName ValueShape
sm2 =
    Map VName TermBinding
-> Map VName TypeBinding -> Map VName ValueShape -> 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) (Map VName ValueShape
sm1 forall a. Semigroup a => a -> a -> a
<> Map VName ValueShape
sm2)

-- | 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 String

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,
      envShapes :: Map VName ValueShape
envShapes = 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,
      envShapes :: Map VName ValueShape
envShapes = 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,
      envShapes :: Map VName ValueShape
envShapes = forall a. Monoid a => a
mempty
    }
  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 Int64 -> (Maybe BoundV, Value)
f
  where
    f :: Int64 -> (Maybe BoundV, Value)
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,
        PrimValue -> Value
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 -> FilePath
show (InterpreterError FilePath
s) = FilePath
s

bad :: SrcLoc -> Env -> String -> EvalM a
bad :: forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad SrcLoc
loc Env
env FilePath
s = forall a. SrcLoc -> Env -> EvalM a -> EvalM a
stacking SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$ do
  [FilePath]
ss <- forall a b. (a -> b) -> [a] -> [b]
map (forall a. Located a => a -> FilePath
locStr 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 a b. (a -> b) -> a -> b
$ forall a. InterpreterError -> ExtOp a
ExtOpError forall a b. (a -> b) -> a -> b
$ FilePath -> InterpreterError
InterpreterError forall a b. (a -> b) -> a -> b
$ FilePath
"Error at\n" forall a. [a] -> [a] -> [a]
++ Int -> [FilePath] -> FilePath
prettyStacktrace Int
0 [FilePath]
ss forall a. [a] -> [a] -> [a]
++ FilePath
s

trace :: String -> Value -> EvalM ()
trace :: FilePath -> Value -> EvalM ()
trace FilePath
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. FilePath -> FilePath -> a -> ExtOp a
ExtOpTrace FilePath
w (forall a. Pretty a => a -> FilePath
prettyOneLine 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 :: Loc -> EvalM ()
break :: Loc -> EvalM ()
break Loc
loc = 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' -> 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected array value, but found: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

toArray :: ValueShape -> [Value] -> Value
toArray :: ValueShape -> [Value] -> Value
toArray ValueShape
shape [Value]
vs = ValueShape -> Array Int Value -> Value
ValueArray ValueShape
shape (forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, forall (t :: * -> *) a. Foldable t => t a -> Int
length [Value]
vs forall a. Num a => a -> a -> a
- Int
1) [Value]
vs)

toArray' :: ValueShape -> [Value] -> Value
toArray' :: ValueShape -> [Value] -> Value
toArray' ValueShape
rowshape [Value]
vs = ValueShape -> Array Int Value -> Value
ValueArray ValueShape
shape (forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, forall (t :: * -> *) a. Foldable t => t a -> Int
length [Value]
vs forall a. Num a => a -> a -> a
- Int
1) [Value]
vs)
  where
    shape :: ValueShape
shape = forall d. d -> Shape d -> Shape d
ShapeDim (forall i a. Num i => [a] -> i
genericLength [Value]
vs) ValueShape
rowshape

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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot apply non-function: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"matchPat: missing case for " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Pat
p forall a. [a] -> [a] -> [a]
++ FilePath
" and " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 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
$ PrimValue -> Value
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
  ppr :: Indexing -> Doc
ppr (IndexingFix Int64
i) = forall a. Pretty a => a -> Doc
ppr Int64
i
  ppr (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. Pretty a => a -> Doc
ppr Maybe Int64
i
      forall a. Semigroup a => a -> a -> a
<> FilePath -> Doc
text FilePath
":"
      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. Pretty a => a -> Doc
ppr Maybe Int64
j
      forall a. Semigroup a => a -> a -> a
<> FilePath -> Doc
text FilePath
":"
      forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Doc
ppr Int64
s
  ppr (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. Pretty a => a -> Doc
ppr Maybe Int64
i
      forall a. Semigroup a => a -> a -> a
<> FilePath -> Doc
text FilePath
":"
      forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Doc
ppr 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 ((FilePath -> Doc
text FilePath
":" forall a. Semigroup a => a -> a -> a
<>) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> Doc
ppr) Maybe Int64
s
  ppr (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. Pretty a => a -> Doc
ppr Maybe Int64
i forall a. Semigroup a => a -> a -> a
<> FilePath -> Doc
text FilePath
":"

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. Integral int => Array Int Value -> 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. Integral int => Array Int Value -> int
arrayLength Array Int Value
arr
  ValueShape -> [Value] -> Value
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 forall i e. Ix i => Array i e -> i -> e
!)) [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
$ ValueShape -> Array Int Value -> Value
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. Integral int => Array Int Value -> 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. Integral int => Array Int Value -> 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. Integral int => Array Int Value -> 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 (ValueShape -> Array Int Value -> Value
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 -> FilePath -> EvalM a
bad SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$
          FilePath
"Index ["
            forall a. Semigroup a => a -> a -> a
<> forall a. [a] -> [[a]] -> [a]
intercalate FilePath
", " (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> FilePath
pretty [Indexing]
is)
            forall a. Semigroup a => a -> a -> a
<> FilePath
"] out of bounds for array of shape "
            forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty (Value -> ValueShape
valueShape Value
arr)
            forall a. Semigroup a => a -> a -> a
<> FilePath
"."
  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).
evalType :: Env -> StructType -> StructType
evalType :: Env -> StructType -> StructType
evalType 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
evalType 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
evalType Env
env) Map Name StructType
fs
evalType Env
env (Scalar (Arrow () PName
p 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
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
p (Env -> StructType -> StructType
evalType Env
env StructType
t1) (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (Env -> StructType -> StructType
evalType Env
env StructType
t2))
evalType 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
evalType Env
env StructType
et
      shape' :: Shape Size
shape' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Size -> Size
evalDim Shape Size
shape
   in forall as dim.
Monoid as =>
Uniqueness -> Shape dim -> TypeBase dim as -> TypeBase dim as
arrayOf Uniqueness
u Shape Size
shape' StructType
et'
  where
    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 =
          Int -> 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
evalType 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
evalType (Map VName TermBinding
-> Map VName TypeBinding -> Map VName ValueShape -> Env
Env forall a. Monoid a => a
mempty Map VName TypeBinding
types forall a. Monoid a => a
mempty 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 Int
k) SrcLoc
_) =
      (forall k a. k -> a -> Map k a
M.singleton k
p forall a b. (a -> b) -> a -> b
$ Int -> Size
ConstSize Int
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
evalType 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
evalType 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
evalType Env
env) Map Name [StructType]
cs

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) -> do
      Env
size_env <- EvalM Env
extSizeEnv
      StructType -> EvalM Value
v forall a b. (a -> b) -> a -> b
$ Env -> StructType -> StructType
evalType (Env
size_env forall a. Semigroup a => a -> a -> a
<> 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"\"" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty QualName VName
qv forall a. Semigroup a => a -> a -> a
<> FilePath
"\" is not bound to a value."

typeValueShape :: Env -> StructType -> EvalM ValueShape
typeValueShape :: Env -> StructType -> EvalM ValueShape
typeValueShape Env
env StructType
t = do
  Env
size_env <- EvalM Env
extSizeEnv
  let t' :: StructType
t' = Env -> StructType -> StructType
evalType (Env
size_env forall a. Semigroup a => a -> a -> a
<> 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. Map VName (Shape d) -> TypeBase d () -> Shape d
typeShape forall a. Monoid a => a
mempty StructType
t' of
    Maybe ValueShape
Nothing -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"typeValueShape: failed to fully evaluate type " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty StructType
t'
    Just ValueShape
shape -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ValueShape
shape
  where
    dim :: Size -> Maybe a
dim (ConstSize Int
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
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
_ StructType
pt (RetType [VName]
_ TypeBase Size as
rt))) =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
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 <- 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 a b. (a -> b) -> a -> b
$
    (Value -> EvalM Value) -> Value
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
evalType 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 (Value -> 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 ret' :: StructType
ret' = Env -> StructType -> StructType
evalType Env
env forall a b. (a -> b) -> a -> b
$ forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase Size ()
ret
      arrow :: (PName, TypeBase dim ()) -> TypeBase dim () -> TypeBase dim ()
arrow (PName
xp, 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
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
xp 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, TypeBase dim ()) -> TypeBase dim () -> TypeBase dim ()
arrow forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat -> (PName, StructType)
patternParam) StructType
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
      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 (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (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 StructType
ret')
    else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      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 StructType
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 a b. (a -> b) -> a -> b
$
    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
evalType 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
$
        Value -> 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
$ ValueShape -> [Value] -> Value
toArray' forall d. Shape d
ShapeLeaf forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Integer -> Value
toInt [Integer
start', Integer
start' forall a. Num a => a -> a -> a
+ Integer
step .. Integer
end_adj]
    else forall a. SrcLoc -> Env -> FilePath -> 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 -> FilePath
badRange Integer
start' Maybe Integer
maybe_second' Inclusiveness Integer
end'
  where
    toInt :: Integer -> Value
toInt =
      case Exp -> PatType
typeOf Exp
start of
        Scalar (Prim (Signed IntType
t')) ->
          PrimValue -> Value
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')) ->
          PrimValue -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Nonsensical range type: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show PatType
t

    badRange :: a -> Maybe a -> Inclusiveness a -> FilePath
badRange a
start' Maybe a
maybe_second' Inclusiveness a
end' =
      FilePath
"Range "
        forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty a
start'
        forall a. [a] -> [a] -> [a]
++ ( case Maybe a
maybe_second' of
               Maybe a
Nothing -> FilePath
""
               Just a
second' -> FilePath
".." forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty a
second'
           )
        forall a. [a] -> [a] -> [a]
++ ( case Inclusiveness a
end' of
               DownToExclusive a
x -> FilePath
"..>" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty a
x
               ToInclusive a
x -> FilePath
"..." forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty a
x
               UpToExclusive a
x -> FilePath
"..<" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty a
x
           )
        forall a. [a] -> [a] -> [a]
++ FilePath
" is invalid."
evalAppExp Env
env StructType
t (Coerce Exp
e TypeExp VName
te SrcLoc
loc) = do
  Value
v <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  case Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape
checkShape (Map VName ValueShape -> StructType -> Shape (Maybe Int64)
structTypeShape (Env -> Map VName ValueShape
envShapes Env
env) StructType
t) (Value -> 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 -> FilePath -> EvalM a
bad SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$
        FilePath
"Value `" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty Value
v forall a. Semigroup a => a -> a -> a
<> FilePath
"` of shape `"
          forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty (Value -> ValueShape
valueShape Value
v)
          forall a. [a] -> [a] -> [a]
++ FilePath
"` cannot match shape of type `"
            forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty TypeExp VName
te
            forall a. Semigroup a => a -> a -> a
<> FilePath
"` (`"
            forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty StructType
t
            forall a. Semigroup a => a -> a -> a
<> FilePath
"`)"
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
evalType Env
env forall a b. (a -> b) -> a -> b
$ Pat -> StructType
patternStructType Pat
p
      v_s :: ValueShape
v_s = Value -> 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 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 -> FilePath
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
op) forall a. Eq a => a -> a -> Bool
== FilePath
"&&" = 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
$ PrimValue -> Value
ValuePrim forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
False
    | VName -> FilePath
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
op) forall a. Eq a => a -> a -> Bool
== FilePath
"||" = 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
$ PrimValue -> Value
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 Exp
x (Info (Diet
_, Maybe VName
ext)) SrcLoc
loc) = do
  -- It is important that 'x' is evaluated first in order to bring any
  -- sizes into scope that may be used in the type of 'f'.
  Value
x' <- Env -> Exp -> Maybe VName -> EvalM Value
evalArg Env
env Exp
x Maybe VName
ext
  Value
f' <- Env -> Exp -> EvalM Value
eval Env
env Exp
f
  SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env Value
f' Value
x'
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 -> FilePath -> EvalM a
bad SrcLoc
loc Env
env FilePath
"Bad update"
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)
              (Value -> 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)

    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 -> Exp -> EvalM Value
eval
              ( 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,
                        PrimValue -> Value
ValuePrim (IntValue -> PrimValue
SignedValue IntValue
i)
                      )
                  )
                  forall a. Semigroup a => a -> a -> a
<> Env
env'
              )
              Exp
body

    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 -> Exp -> EvalM Value
eval Env
env' Exp
body
        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 -> Exp -> EvalM Value
eval Env
env'' Exp
body
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 => FilePath -> a
error FilePath
"Pat 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
$ PrimValue -> Value
ValuePrim PrimValue
v
eval Env
env (Hole (Info PatType
t) SrcLoc
loc) = forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad SrcLoc
loc Env
env forall a b. (a -> b) -> a -> b
$ FilePath
"Hole of type: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
prettyOneLine 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Local open of module function at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> FilePath
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
_) = [Value] -> Value
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
_) =
  Map Name Value -> Value
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
$
    ValueShape -> [Value] -> Value
toArray' forall d. Shape d
ShapeLeaf forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
map (PrimValue -> Value
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
$ ValueShape -> [Value] -> Value
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
$ ValueShape -> [Value] -> Value
toArray' (Value -> 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))) =
  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
  where
    t' :: StructType
t' = Env -> StructType -> StructType
evalType Env
env forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
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 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
$ PrimValue -> Value
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
$ PrimValue -> Value
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
$ PrimValue -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"eval: nonsensical type for integer literal: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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
$ PrimValue -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"eval: nonsensical type for float literal: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty PatType
t
eval Env
env (Negate Exp
e SrcLoc
_) = do
  Value
ev <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  PrimValue -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot negate " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
ev
eval Env
env (Not Exp
e SrcLoc
_) = do
  Value
ev <- Env -> Exp -> EvalM Value
eval Env
env Exp
e
  PrimValue -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot logically negate " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 -> FilePath -> EvalM a
bad SrcLoc
loc Env
env FilePath
"Bad update"
eval Env
env (RecordUpdate Exp
src [Name]
all_fs Exp
v Info PatType
_ SrcLoc
_) =
  Value -> [Name] -> Value -> Value
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 -> [Name] -> Value -> Value
update Value
_ [] Value
v' = Value
v'
    update (ValueRecord Map Name Value
src') (Name
f : [Name]
fs) Value
v'
      | Just Value
f_v <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name Value
src' =
          Map Name Value -> Value
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 -> [Name] -> Value -> Value
update Value
f_v [Name]
fs Value
v') Map Name Value
src'
    update Value
_ [Name]
_ Value
_ = forall a. HasCallStack => FilePath -> a
error FilePath
"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 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
$
    (Value -> EvalM Value) -> Value
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
$ (Value -> EvalM Value) -> Value
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
$ (Value -> EvalM Value) -> Value
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 :: * -> *}. Applicative f => Value -> Name -> f Value
walk) [Name]
ks
  where
    walk :: Value -> Name -> f Value
walk (ValueRecord Map Name Value
fs) Name
f
      | 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'
    walk Value
_ Name
_ = forall a. HasCallStack => FilePath -> a
error FilePath
"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 => FilePath -> a
error FilePath
"Value does not have expected field."
eval Env
env (Assert Exp
what Exp
e (Info FilePath
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 -> FilePath -> EvalM a
bad SrcLoc
loc Env
env FilePath
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
$ ValueShape -> Name -> [Value] -> Value
ValueSum ValueShape
shape Name
c [Value]
vs
eval Env
env (Attr (AttrAtom (AtomName Name
"break") SrcLoc
_) Exp
e SrcLoc
loc) = do
  Loc -> EvalM ()
break (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
  FilePath -> Value -> EvalM ()
trace (forall a. Located a => a -> FilePath
locStr (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
  FilePath -> Value -> EvalM ()
trace (Name -> FilePath
nameToString 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 Map VName ValueShape
_)) =
      Env -> Module
Module forall a b. (a -> b) -> a -> b
$ Map VName TermBinding
-> Map VName TypeBinding -> Map VName ValueShape -> 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) forall a. Monoid a => a
mempty
    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 (forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe forall a. [a] -> Maybe a
maybeHead Map VName [VName]
rev_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 Int
x) = Int -> Size
ConstSize Int
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ ShowS
quote (forall a. Pretty a => a -> FilePath
pretty QualName VName
qv) forall a. Semigroup a => a -> a -> a
<> FilePath
" is not bound to a module."

evalModExp :: Env -> ModExp -> EvalM Module
evalModExp :: Env -> ModExp -> EvalM Module
evalModExp Env
_ (ModImport FilePath
_ (Info FilePath
f) SrcLoc
_) = do
  Maybe Env
f' <- FilePath -> EvalM (Maybe Env)
lookupImport FilePath
f
  Map FilePath 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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$
        [FilePath] -> FilePath
unlines
          [ FilePath
"Unknown interpreter import: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show FilePath
f,
            FilePath
"Known: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show (forall k a. Map k a -> [k]
M.keys Map FilePath 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 Map VName ValueShape
_ <- 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 -> Map VName ValueShape -> 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)
        forall a. Monoid a => a
mempty
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 => FilePath -> a
error FilePath
"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 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 => FilePath -> a
error FilePath
"Expected Module"
evalDec Env
env (ImportDec FilePath
name Info FilePath
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.
FilePath -> f FilePath -> SrcLoc -> ModExpBase f vn
ModImport FilePath
name Info FilePath
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 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
evalType 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 FilePath Env
ctxImports :: M.Map FilePath 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 FilePath Env -> Ctx
Ctx
    ( Map VName TermBinding
-> Map VName TypeBinding -> Map VName ValueShape -> Env
Env
        ( forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
            (Name -> Int -> VName
VName (FilePath -> Name
nameFromString FilePath
"intrinsics") Int
0)
            (Module -> TermBinding
TermModule (Env -> Module
Module forall a b. (a -> b) -> a -> b
$ Map VName TermBinding
-> Map VName TypeBinding -> Map VName ValueShape -> Env
Env Map VName TermBinding
terms Map VName TypeBinding
types forall a. Monoid a => a
mempty))
            Map VName TermBinding
terms
        )
        Map VName TypeBinding
types
        forall a. Monoid a => a
mempty
    )
    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
. FilePath -> Maybe TermBinding
def forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> FilePath
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
. FilePath -> Maybe TypeBinding
tdef forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> FilePath
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
$ (Value -> EvalM Value) -> Value
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 a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
x ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
y -> Value -> Value -> EvalM Value
f Value
x Value
y
    fun2t :: (Value -> Value -> EvalM Value) -> TermBinding
fun2t Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y] -> Value -> Value -> EvalM Value
f Value
x Value
y
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected pair; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v
    fun3t :: (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t Value -> Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y, Value
z] -> Value -> Value -> Value -> EvalM Value
f Value
x Value
y Value
z
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected triple; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

    fun5t :: (Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun5t Value -> Value -> Value -> Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y, Value
z, Value
a, Value
b] -> Value -> Value -> Value -> Value -> Value -> EvalM Value
f Value
x Value
y Value
z Value
a Value
b
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected pentuple; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

    fun6t :: (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6t Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y, Value
z, Value
a, Value
b, Value
c] -> Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value
f Value
x Value
y Value
z Value
a Value
b Value
c
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected sextuple; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

    fun7t :: (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun7t Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y, Value
z, Value
a, Value
b, Value
c, 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
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected septuple; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

    fun8t :: (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun8t Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y, Value
z, Value
a, Value
b, Value
c, Value
d, 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
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected sextuple; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

    fun10t :: (Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> Value
 -> EvalM Value)
-> TermBinding
fun10t Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
fun =
      Maybe BoundV -> Value -> TermBinding
TermValue forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case Value -> Maybe [Value]
fromTuple Value
v of
            Just [Value
x, Value
y, Value
z, Value
a, Value
b, Value
c, Value
d, Value
e, Value
f, Value
g] -> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> Value
-> EvalM Value
fun Value
x Value
y Value
z Value
a Value
b Value
c Value
d Value
e Value
f Value
g
            Maybe [Value]
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Expected octuple; got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
v

    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
$ PrimValue -> Value
ValuePrim PrimValue
z
        (Value, Value)
_ ->
          forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            FilePath
"Cannot apply operator to arguments "
              forall a. Semigroup a => a -> a -> a
<> ShowS
quote (forall a. Pretty a => a -> FilePath
pretty Value
x)
              forall a. Semigroup a => a -> a -> a
<> FilePath
" and "
              forall a. Semigroup a => a -> a -> a
<> ShowS
quote (forall a. Pretty a => a -> FilePath
pretty Value
y)
              forall a. Semigroup a => a -> a -> a
<> FilePath
"."
      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
$ PrimValue -> Value
ValuePrim PrimValue
r
        Value
_ ->
          forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            FilePath
"Cannot apply function to argument "
              forall a. Semigroup a => a -> a -> a
<> ShowS
quote (forall a. Pretty a => a -> FilePath
pretty Value
x)
              forall a. Semigroup a => a -> a -> a
<> FilePath
"."
      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 Value -> Maybe [Value]
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
$ PrimValue -> Value
ValuePrim PrimValue
z
        Maybe [Value]
_ ->
          forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            FilePath
"Cannot apply operator to argument "
              forall a. Semigroup a => a -> a -> a
<> ShowS
quote (forall a. Pretty a => a -> FilePath
pretty Value
v)
              forall a. Semigroup a => a -> a -> a
<> FilePath
"."

    def :: FilePath -> Maybe TermBinding
def FilePath
"!" =
      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 FilePath
"+" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Add` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FAdd
    def FilePath
"-" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Sub` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FSub
    def FilePath
"*" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Mul` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FMul
    def FilePath
"**" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp IntType -> BinOp
P.Pow FloatType -> BinOp
P.FPow
    def FilePath
"/" =
      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 FilePath
"%" =
      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 FilePath
"//" =
      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 FilePath
"%%" =
      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 FilePath
"^" = 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 FilePath
"&" = 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 FilePath
"|" = 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 FilePath
">>" = 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 FilePath
"<<" = 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 FilePath
">>>" = 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 FilePath
"==" = 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
$ PrimValue -> Value
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 FilePath
"!=" = 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
$ PrimValue -> Value
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 FilePath
"&&" = 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
$ PrimValue -> Value
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 FilePath
"||" = 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
$ PrimValue -> Value
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 FilePath
"<" =
      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 FilePath
">" =
      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 FilePath
"<=" =
      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 FilePath
">=" =
      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 FilePath
s
      | Just BinOp
bop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((FilePath
s forall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> FilePath
pretty) [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 ((FilePath
s forall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> FilePath
pretty) [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 ((FilePath
s forall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> FilePath
pretty) [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 ((FilePath
s forall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> FilePath
pretty) [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 FilePath
s Map FilePath ([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 -> Maybe PrimValue
getV' (ValuePrim PrimValue
v) = forall a. a -> Maybe a
Just PrimValue
v
                    getV' Value
_ = forall a. Maybe a
Nothing
                case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Value -> Maybe PrimValue
getV' forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value -> Maybe [Value]
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
$ PrimValue -> Value
ValuePrim PrimValue
res
                  Maybe [PrimValue]
_ ->
                    forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot apply " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty FilePath
s forall a. [a] -> [a] -> [a]
++ FilePath
" to " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
x
      | FilePath
"sign_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
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
$ PrimValue -> Value
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue IntValue
x'
                Value
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot sign: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
x
      | FilePath
"unsign_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
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
$ PrimValue -> Value
ValuePrim forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue IntValue
x'
                Value
_ -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot unsign: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
x
    def FilePath
s
      | FilePath
"map_stream" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
s =
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ (Value -> Value -> EvalM Value) -> TermBinding
fun2t Value -> Value -> EvalM Value
stream
    def FilePath
s | FilePath
"reduce_stream" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
s =
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t forall a b. (a -> b) -> a -> b
$ \Value
_ Value
f Value
arg -> Value -> Value -> EvalM Value
stream Value
f Value
arg
    def FilePath
"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
$
        (Value -> EvalM Value) -> Value
ValueFun forall a b. (a -> b) -> a -> b
$ \Value
v ->
          case (Value -> Maybe [Value]
fromTuple Value
v, forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType StructType
t) of
            (Just [Value
f, Value
xs], ([StructType
_], StructType
ret_t))
              | Just ValueShape
rowshape <- StructType -> Maybe ValueShape
typeRowShape StructType
ret_t ->
                  ValueShape -> [Value] -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Bad pure type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty StructType
ret_t
            (Maybe [Value], ([StructType], StructType))
_ ->
              forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$
                FilePath
"Invalid arguments to map intrinsic:\n"
                  forall a. [a] -> [a] -> [a]
++ [FilePath] -> FilePath
unlines [forall a. Pretty a => a -> FilePath
pretty StructType
t, forall a. Pretty a => a -> FilePath
pretty Value
v]
      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
. Map VName ValueShape -> StructType -> Shape (Maybe Int64)
structTypeShape forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1
    def FilePath
s | FilePath
"reduce" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` FilePath
s = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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 FilePath
"scan" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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')
        ValueShape -> [Value] -> Value
toArray' (Value -> 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 FilePath
"scatter" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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
$
              ValueShape -> Array Int Value -> Value
ValueArray ValueShape
shape forall a b. (a -> b) -> a -> b
$
                forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Array Int Value -> (Int, Value) -> Array Int Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"scatter expects array, but got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
arr
      where
        update :: Array Int Value -> (Int, Value) -> Array Int Value
update Array Int Value
arr' (Int
i, Value
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. Integral int => Array Int Value -> int
arrayLength Array Int Value
arr'
            then Array Int Value
arr' forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Int
i, Value
v)]
            else Array Int Value
arr'
    def FilePath
"scatter_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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 Value -> Maybe [Value]
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"scatter_2d expects array, but got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 => FilePath -> a
error FilePath
"scatter_2d expects 2-dimensional indices"
    def FilePath
"scatter_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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 Value -> Maybe [Value]
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"scatter_3d expects array, but got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty 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 => FilePath -> a
error FilePath
"scatter_3d expects 3-dimensional indices"
    def FilePath
"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
fun6t 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 FilePath
"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
fun6t 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 Value -> Maybe [Value]
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 => FilePath -> a
error FilePath
"hist_2d: bad index value"
    def FilePath
"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
fun6t 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 Value -> Maybe [Value]
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 => FilePath -> a
error FilePath
"hist_2d: bad index value"
    def FilePath
"partition" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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]] -> Value
pack [[Value]]
parts =
              [Value] -> Value
toTuple
                [ ValueShape -> [Value] -> Value
toArray' ValueShape
rowshape forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Value]]
parts,
                  ValueShape -> [Value] -> Value
toArray' ValueShape
rowshape forall a b. (a -> b) -> a -> b
$
                    forall a b. (a -> b) -> [a] -> [b]
map (PrimValue -> Value
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]]
parts
                ]

        [[Value]] -> Value
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 FilePath
"scatter_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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 = (Value -> Value -> EvalM Value) -> Array Int Value -> Value
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
$ ValueShape -> Array Int Value -> Value
ValueArray ValueShape
dest_shape Array Int Value
dest_arr'
                Value
_ ->
                  forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"scatter_stream produced: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
acc'
          (Value, Value)
_ ->
            forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"scatter_stream expects array, but got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty (Value
dest, Value
vs)
    def FilePath
"hist_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun5t 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 = (Value -> Value -> EvalM Value) -> Array Int Value -> Value
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
$ ValueShape -> Array Int Value -> Value
ValueArray ValueShape
dest_shape Array Int Value
dest_arr'
                Value
_ ->
                  forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"hist_stream produced: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
acc'
          (Value, Value)
_ ->
            forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"hist_stream expects array, but got: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty (Value
dest, Value
vs)
    def FilePath
"acc_write" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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. Integral int => Array Int Value -> 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
$ (Value -> Value -> EvalM Value) -> Array Int Value -> Value
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 => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"acc_write invalid arguments: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty (Value
acc, Value
i, Value
v)
    --
    def FilePath
"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
fun6t 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] -> Maybe Value
mk1 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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] -> Maybe Value
mk2 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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 [Maybe Value] -> Maybe Value
mk1 [[Maybe Value] -> Maybe Value
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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            FilePath
"Index out of bounds: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty [(Int64
n1', Int64
s1', Int64
n2', Int64
s2')]
    --
    def FilePath
"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
fun5t 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 Value -> 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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                FilePath
"Index out of bounds: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty [(Int64
n1, Int64
s1', Int64
n2, Int64
s2')]
        ValueShape
s -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"flat_update_2d: invalid arg shape: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show ValueShape
s
    --
    def FilePath
"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
fun8t 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] -> Maybe Value
mk1 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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] -> Maybe Value
mk2 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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] -> Maybe Value
mk3 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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 [Maybe Value] -> Maybe Value
mk1 [[Maybe Value] -> Maybe Value
mk2 [[Maybe Value] -> Maybe Value
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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            FilePath
"Index out of bounds: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty [(Int64
n1', Int64
s1', Int64
n2', Int64
s2', Int64
n3', Int64
s3')]
    --
    def FilePath
"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
fun6t 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 Value -> 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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                FilePath
"Index out of bounds: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty [(Int64
n1, Int64
s1', Int64
n2, Int64
s2', Int64
n3, Int64
s3')]
        ValueShape
s -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"flat_update_3d: invalid arg shape: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show ValueShape
s
    --
    def FilePath
"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
fun10t 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] -> Maybe Value
mk1 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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] -> Maybe Value
mk2 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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] -> Maybe Value
mk3 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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] -> Maybe Value
mk4 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> [Value] -> Value
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 [Maybe Value] -> Maybe Value
mk1 [[Maybe Value] -> Maybe Value
mk2 [[Maybe Value] -> Maybe Value
mk3 [[Maybe Value] -> Maybe Value
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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
            FilePath
"Index out of bounds: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty [(Int64
n1', Int64
s1', Int64
n2', Int64
s2', Int64
n3', Int64
s3', Int64
n4', Int64
s4')]
    --
    def FilePath
"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
fun7t 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 Value -> 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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
                FilePath
"Index out of bounds: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty [(Int64
n1, Int64
s1', Int64
n2, Int64
s2', Int64
n3, Int64
s3', Int64
n4, Int64
s4')]
        ValueShape
s -> forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"flat_update_4d: invalid arg shape: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show ValueShape
s
    --
    def FilePath
"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) = Value -> 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], [Value]) -> [Value]
listPair ([Value]
xs, [Value]
ys) =
              [ValueShape -> [Value] -> Value
toArray' ValueShape
xs_shape [Value]
xs, ValueShape -> [Value] -> Value
toArray' ValueShape
ys_shape [Value]
ys]

        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Value] -> Value
toTuple forall a b. (a -> b) -> a -> b
$ ([Value], [Value]) -> [Value]
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}. Pretty b => Maybe [b] -> (b, b)
fromPair forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Maybe [Value]
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]
l = forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Not a pair: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Maybe [b]
l
    def FilePath
"zip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2t forall a b. (a -> b) -> a -> b
$ \Value
xs Value
ys -> do
        let ShapeDim Int64
_ ValueShape
xs_rowshape = Value -> ValueShape
valueShape Value
xs
            ShapeDim Int64
_ ValueShape
ys_rowshape = Value -> ValueShape
valueShape Value
ys
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
          ValueShape -> [Value] -> Value
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 [Value] -> Value
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 FilePath
"concat" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2t 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
$ ValueShape -> [Value] -> Value
toArray' ValueShape
rowshape forall a b. (a -> b) -> a -> b
$ [Value]
xs' forall a. [a] -> [a] -> [a]
++ [Value]
ys'
    def FilePath
"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
$
          ValueShape -> [Value] -> Value
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 (ValueShape -> [Value] -> Value
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 FilePath
"rotate" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2t 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 ValueShape -> [Value] -> Value
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 ValueShape -> [Value] -> Value
toArray ValueShape
shape forall a b. (a -> b) -> a -> b
$ [Value]
aft forall a. [a] -> [a] -> [a]
++ [Value]
bef
    def FilePath
"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
$ ValueShape -> [Value] -> Value
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 FilePath
"unflatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t 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 -> FilePath -> EvalM a
bad forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$
              FilePath
"Cannot unflatten array of shape ["
                forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty Int64
xs_size
                forall a. Semigroup a => a -> a -> a
<> FilePath
"] to array of shape ["
                forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty (Value -> Int64
asInt64 Value
n)
                forall a. Semigroup a => a -> a -> a
<> FilePath
"]["
                forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> FilePath
pretty (Value -> Int64
asInt64 Value
m)
                forall a. Semigroup a => a -> a -> a
<> FilePath
"]"
          else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
toArray ValueShape
shape forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (ValueShape -> [Value] -> Value
toArray ValueShape
rowshape) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [[a]]
chunk (Value -> Int
asInt Value
m) [Value]
xs'
    def FilePath
"vjp2" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t forall a b. (a -> b) -> a -> b
$
        \Value
_ Value
_ Value
_ -> forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty FilePath
"Interpreter does not support autodiff."
    def FilePath
"jvp2" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3t forall a b. (a -> b) -> a -> b
$
        \Value
_ Value
_ Value
_ -> forall a. SrcLoc -> Env -> FilePath -> EvalM a
bad forall a. IsLocation a => a
noLoc forall a. Monoid a => a
mempty FilePath
"Interpreter does not support autodiff."
    def FilePath
"acc" = forall a. Maybe a
Nothing
    def FilePath
s | FilePath -> Name
nameFromString FilePath
s forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map Name PrimType
namesToPrimTypes = forall a. Maybe a
Nothing
    def FilePath
s = forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Missing intrinsic: " forall a. [a] -> [a] -> [a]
++ FilePath
s

    tdef :: FilePath -> Maybe TypeBinding
tdef FilePath
s = do
      PrimType
t <- FilePath -> Name
nameFromString FilePath
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
n = PrimValue -> Value
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. Integral int => Array Int Value -> 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 Value
n Value
arg
    stream Value
_ Value
arg = forall a. HasCallStack => FilePath -> a
error forall a b. (a -> b) -> a -> b
$ FilePath
"Cannot stream: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> FilePath
pretty Value
arg

interpretExp :: Ctx -> Exp -> F ExtOp Value
interpretExp :: Ctx -> Exp -> F ExtOp Value
interpretExp Ctx
ctx Exp
e = forall a. Map FilePath Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map FilePath 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 FilePath Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map FilePath 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 -> (FilePath, Prog) -> F ExtOp Ctx
interpretImport :: Ctx -> (FilePath, Prog) -> F ExtOp Ctx
interpretImport Ctx
ctx (FilePath
fp, Prog
prog) = do
  Env
env <- forall a. Map FilePath Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map FilePath 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 FilePath Env
ctxImports = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert FilePath
fp Env
env forall a b. (a -> b) -> a -> b
$ Ctx -> Map FilePath 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}

checkEntryArgs :: VName -> [F.Value] -> StructType -> Either String ()
checkEntryArgs :: VName -> [Value] -> StructType -> Either FilePath ()
checkEntryArgs VName
entry [Value]
args StructType
entry_t
  | [StructType]
args_ts forall a. Eq a => a -> a -> Bool
== [StructType]
param_ts =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise =
      forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$
        forall a. Pretty a => a -> FilePath
pretty forall a b. (a -> b) -> a -> b
$
          Doc
expected
            Doc -> Doc -> Doc
</> Doc
"Got input of types"
            Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 ([Doc] -> Doc
stack (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Doc
ppr [StructType]
args_ts))
  where
    ([StructType]
param_ts, StructType
_) = forall dim as.
TypeBase dim as -> ([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
expected
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StructType]
param_ts =
          Doc
"Entry point " forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
pquote (forall v. IsName v => v -> Doc
pprName VName
entry) forall a. Semigroup a => a -> a -> a
<> Doc
" is not a function."
      | Bool
otherwise =
          Doc
"Entry point " forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
pquote (forall v. IsName v => v -> Doc
pprName VName
entry) forall a. Semigroup a => a -> a -> a
<> Doc
" expects input of type(s)"
            Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 ([Doc] -> Doc
stack (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Doc
ppr [StructType]
param_ts))

-- | Execute the named function on the given arguments; may fail
-- horribly if these are ill-typed.
interpretFunction :: Ctx -> VName -> [F.Value] -> Either String (F ExtOp Value)
interpretFunction :: Ctx -> VName -> [Value] -> Either FilePath (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 FilePath (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 FilePath (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
$ FilePath
"Unknown function `" forall a. Semigroup a => a -> a -> a
<> forall v. IsName v => v -> FilePath
prettyName VName
fname forall a. Semigroup a => a -> a -> a
<> FilePath
"`."

  [Value]
vs' <- case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Value -> Maybe Value
convertValue [Value]
vs of
    Just [Value]
vs' -> forall a b. b -> Either a b
Right [Value]
vs'
    Maybe [Value]
Nothing -> forall a b. a -> Either a b
Left FilePath
"Invalid input: irregular array."

  VName -> [Value] -> StructType -> Either FilePath ()
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 FilePath Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map FilePath 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 [Value]
vs'
  where
    updateType :: [ValueType]
-> TypeBase Size as -> Either FilePath (TypeBase Size as)
updateType (ValueType
vt : [ValueType]
vts) (Scalar (Arrow as
als PName
u StructType
pt (RetType [VName]
dims TypeBase Size as
rt))) = do
      ValueType -> StructType -> Either FilePath ()
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
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow as
als PName
u (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 FilePath (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 String ()
    checkInput :: ValueType -> StructType -> Either FilePath ()
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 FilePath 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 FilePath b
badPrim PrimType
vt PrimType
pt
    checkInput ValueType
_ StructType
_ =
      forall a b. b -> Either a b
Right ()

    badPrim :: a -> a -> Either FilePath 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. Pretty a => a -> FilePath
pretty forall a b. (a -> b) -> a -> b
$
        Doc
"Invalid argument type."
          Doc -> Doc -> Doc
</> Doc
"Expected:"
          Doc -> Doc -> Doc
<+> Doc -> Doc
align (forall a. Pretty a => a -> Doc
ppr a
pt)
          Doc -> Doc -> Doc
</> Doc
"Got:     "
          Doc -> Doc -> Doc
<+> Doc -> Doc
align (forall a. Pretty a => a -> Doc
ppr a
vt)

    convertValue :: Value -> Maybe Value
convertValue (F.PrimValue PrimValue
p) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
ValuePrim PrimValue
p
    convertValue (F.ArrayValue Array Int Value
arr ValueType
t) = ValueType -> [Value] -> Maybe Value
mkArray ValueType
t 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 Value -> Maybe Value
convertValue (forall i e. Array i e -> [e]
elems Array Int Value
arr)