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

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

import Control.Monad
import Control.Monad.Free.Church
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Trans.Maybe
import Data.Array
import Data.Bifunctor
import Data.Bitraversable
import Data.List
  ( find,
    foldl',
    genericLength,
    genericTake,
    transpose,
  )
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Maybe
import Data.Monoid hiding (Sum)
import Data.Ord
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.Data qualified as V
import Futhark.Util (chunk)
import Futhark.Util.Loc
import Futhark.Util.Pretty hiding (apply)
import Language.Futhark hiding (Shape, matchDims)
import Language.Futhark qualified as F
import Language.Futhark.Interpreter.AD qualified as AD
import Language.Futhark.Interpreter.Values hiding (Value)
import Language.Futhark.Interpreter.Values qualified
import Language.Futhark.Primitive (floatValue, intValue)
import Language.Futhark.Primitive qualified as P
import Language.Futhark.Semantic qualified as T
import Language.Futhark.TypeChecker.Types (Subst (..), applySubst)
import Prelude hiding (break, mod)

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

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

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

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

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

type Stack = [StackFrame]

type Exts = M.Map VName Value

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

runEvalM :: M.Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM :: forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM Map ImportName Env
imports (EvalM ReaderT
  ([StackFrame], Map ImportName Env) (StateT Exts (F ExtOp)) a
m) = StateT Exts (F ExtOp) a -> Exts -> F ExtOp a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (ReaderT
  ([StackFrame], Map ImportName Env) (StateT Exts (F ExtOp)) a
-> ([StackFrame], Map ImportName Env) -> StateT Exts (F ExtOp) a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT
  ([StackFrame], Map ImportName Env) (StateT Exts (F ExtOp)) a
m ([StackFrame]
forall a. Monoid a => a
mempty, Map ImportName Env
imports)) Exts
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 = (([StackFrame], Map ImportName Env)
 -> ([StackFrame], Map ImportName Env))
-> EvalM a -> EvalM a
forall a.
(([StackFrame], Map ImportName Env)
 -> ([StackFrame], Map ImportName Env))
-> EvalM a -> EvalM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((([StackFrame], Map ImportName Env)
  -> ([StackFrame], Map ImportName Env))
 -> EvalM a -> EvalM a)
-> (([StackFrame], Map ImportName Env)
    -> ([StackFrame], Map ImportName Env))
-> EvalM a
-> EvalM a
forall a b. (a -> b) -> a -> b
$ \([StackFrame]
ss, Map ImportName Env
imports) ->
  if SrcLoc -> Bool
isNoLoc SrcLoc
loc
    then ([StackFrame]
ss, Map ImportName Env
imports)
    else
      let s :: StackFrame
s = Loc -> Ctx -> StackFrame
StackFrame (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc) (Env -> Map ImportName Env -> Ctx
Ctx Env
env Map ImportName Env
imports)
       in (StackFrame
s StackFrame -> [StackFrame] -> [StackFrame]
forall a. a -> [a] -> [a]
: [StackFrame]
ss, Map ImportName Env
imports)
  where
    isNoLoc :: SrcLoc -> Bool
    isNoLoc :: SrcLoc -> Bool
isNoLoc = (Loc -> Loc -> Bool
forall a. Eq a => a -> a -> Bool
== Loc
NoLoc) (Loc -> Bool) -> (SrcLoc -> Loc) -> SrcLoc -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf

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

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

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

getExts :: EvalM Exts
getExts :: EvalM Exts
getExts = EvalM Exts
forall s (m :: * -> *). MonadState s m => m s
get

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

extEnv :: EvalM Env
extEnv :: EvalM Env
extEnv = Map VName (Maybe BoundV, Value) -> Env
valEnv (Map VName (Maybe BoundV, Value) -> Env)
-> (Exts -> Map VName (Maybe BoundV, Value)) -> Exts -> Env
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> (Maybe BoundV, Value))
-> Exts -> Map VName (Maybe BoundV, Value)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Value -> (Maybe BoundV, Value)
forall {b} {a}. b -> (Maybe a, b)
f (Exts -> Env) -> EvalM Exts -> EvalM Env
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM Exts
getExts
  where
    f :: b -> (Maybe a, b)
f b
v =
      ( Maybe a
forall a. Maybe a
Nothing,
        b
v
      )

valueStructType :: ValueType -> StructType
valueStructType :: ValueType -> StructType
valueStructType = (Int64 -> ExpBase Info VName) -> ValueType -> StructType
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Int64 -> ExpBase Info VName) -> ValueType -> StructType)
-> (Int64 -> ExpBase Info VName) -> ValueType -> StructType
forall a b. (a -> b) -> a -> b
$ (Integer -> SrcLoc -> ExpBase Info VName)
-> SrcLoc -> Integer -> ExpBase Info VName
forall a b c. (a -> b -> c) -> b -> a -> c
flip Integer -> SrcLoc -> ExpBase Info VName
sizeFromInteger SrcLoc
forall a. Monoid a => a
mempty (Integer -> ExpBase Info VName)
-> (Int64 -> Integer) -> Int64 -> ExpBase Info VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Integer
forall a. Integral a => a -> Integer
toInteger

-- | An expression along with an environment in which to evaluate that
-- expression. Used to represent non-interpreted size expressions,
-- which may still be in reference to some environment.
data SizeClosure = SizeClosure Env Size
  deriving (Int -> SizeClosure -> ShowS
[SizeClosure] -> ShowS
SizeClosure -> [Char]
(Int -> SizeClosure -> ShowS)
-> (SizeClosure -> [Char])
-> ([SizeClosure] -> ShowS)
-> Show SizeClosure
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SizeClosure -> ShowS
showsPrec :: Int -> SizeClosure -> ShowS
$cshow :: SizeClosure -> [Char]
show :: SizeClosure -> [Char]
$cshowList :: [SizeClosure] -> ShowS
showList :: [SizeClosure] -> ShowS
Show)

instance Pretty SizeClosure where
  pretty :: forall ann. SizeClosure -> Doc ann
pretty (SizeClosure Env
_ ExpBase Info VName
e) = ExpBase Info VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ExpBase Info VName -> Doc ann
pretty ExpBase Info VName
e

instance Pretty (F.Shape SizeClosure) where
  pretty :: forall ann. Shape SizeClosure -> Doc ann
pretty = [Doc ann] -> Doc ann
forall a. Monoid a => [a] -> a
mconcat ([Doc ann] -> Doc ann)
-> (Shape SizeClosure -> [Doc ann]) -> Shape SizeClosure -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SizeClosure -> Doc ann) -> [SizeClosure] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
braces (Doc ann -> Doc ann)
-> (SizeClosure -> Doc ann) -> SizeClosure -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SizeClosure -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SizeClosure -> Doc ann
pretty) ([SizeClosure] -> [Doc ann])
-> (Shape SizeClosure -> [SizeClosure])
-> Shape SizeClosure
-> [Doc ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape SizeClosure -> [SizeClosure]
forall dim. Shape dim -> [dim]
shapeDims

-- | A type where the sizes are unevaluated expressions.
type EvalType = TypeBase SizeClosure NoUniqueness

structToEval :: Env -> StructType -> EvalType
structToEval :: Env -> StructType -> TypeBase SizeClosure NoUniqueness
structToEval Env
env = (ExpBase Info VName -> SizeClosure)
-> StructType -> TypeBase SizeClosure NoUniqueness
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Env -> ExpBase Info VName -> SizeClosure
SizeClosure Env
env)

evalToStruct :: EvalType -> StructType
evalToStruct :: TypeBase SizeClosure NoUniqueness -> StructType
evalToStruct = (SizeClosure -> ExpBase Info VName)
-> TypeBase SizeClosure NoUniqueness -> StructType
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (\(SizeClosure Env
_ ExpBase Info VName
e) -> ExpBase Info VName
e)

resolveTypeParams ::
  [VName] ->
  StructType ->
  EvalType ->
  ([(VName, ([VName], EvalType))], [(VName, SizeClosure)])
resolveTypeParams :: [VName]
-> StructType
-> TypeBase SizeClosure NoUniqueness
-> ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
    [(VName, SizeClosure)])
resolveTypeParams [VName]
names StructType
orig_t1 TypeBase SizeClosure NoUniqueness
orig_t2 =
  State
  ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
   [(VName, SizeClosure)])
  ()
-> ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
    [(VName, SizeClosure)])
-> ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
    [(VName, SizeClosure)])
forall s a. State s a -> s -> s
execState ([VName]
-> StructType
-> TypeBase SizeClosure NoUniqueness
-> State
     ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
      [(VName, SizeClosure)])
     ()
forall {m :: * -> *} {p :: * -> * -> *}.
(Bifunctor p,
 MonadState
   (p [(VName, ([VName], TypeBase SizeClosure NoUniqueness))]
      [(VName, SizeClosure)])
   m) =>
[VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
forall a. Monoid a => a
mempty StructType
orig_t1 TypeBase SizeClosure NoUniqueness
orig_t2) ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
 [(VName, SizeClosure)])
forall a. Monoid a => a
mempty
  where
    addType :: a -> b -> m ()
addType a
v b
t = (p [(a, b)] c -> p [(a, b)] c) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((p [(a, b)] c -> p [(a, b)] c) -> m ())
-> (p [(a, b)] c -> p [(a, b)] c) -> m ()
forall a b. (a -> b) -> a -> b
$ ([(a, b)] -> [(a, b)]) -> p [(a, b)] c -> p [(a, b)] c
forall a b c. (a -> b) -> p a c -> p b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (([(a, b)] -> [(a, b)]) -> p [(a, b)] c -> p [(a, b)] c)
-> ([(a, b)] -> [(a, b)]) -> p [(a, b)] c -> p [(a, b)] c
forall a b. (a -> b) -> a -> b
$ ((a, b) -> (a, b) -> Ordering) -> (a, b) -> [(a, b)] -> [(a, b)]
forall a. (a -> a -> Ordering) -> a -> [a] -> [a]
L.insertBy (((a, b) -> a) -> (a, b) -> (a, b) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (a, b) -> a
forall a b. (a, b) -> a
fst) (a
v, b
t)
    addDim :: a -> b -> m ()
addDim a
v b
e = (p a [(a, b)] -> p a [(a, b)]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((p a [(a, b)] -> p a [(a, b)]) -> m ())
-> (p a [(a, b)] -> p a [(a, b)]) -> m ()
forall a b. (a -> b) -> a -> b
$ ([(a, b)] -> [(a, b)]) -> p a [(a, b)] -> p a [(a, b)]
forall b c a. (b -> c) -> p a b -> p a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (([(a, b)] -> [(a, b)]) -> p a [(a, b)] -> p a [(a, b)])
-> ([(a, b)] -> [(a, b)]) -> p a [(a, b)] -> p a [(a, b)]
forall a b. (a -> b) -> a -> b
$ ((a, b) -> (a, b) -> Ordering) -> (a, b) -> [(a, b)] -> [(a, b)]
forall a. (a -> a -> Ordering) -> a -> [a] -> [a]
L.insertBy (((a, b) -> a) -> (a, b) -> (a, b) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (a, b) -> a
forall a b. (a, b) -> a
fst) (a
v, b
e)

    match :: [VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound (Scalar (TypeVar NoUniqueness
_ QualName VName
tn [TypeArg (ExpBase Info VName)]
_)) TypeBase SizeClosure NoUniqueness
t
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
tn VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
names = VName -> ([VName], TypeBase SizeClosure NoUniqueness) -> m ()
forall {p :: * -> * -> *} {a} {b} {c} {m :: * -> *}.
(MonadState (p [(a, b)] c) m, Bifunctor p, Ord a) =>
a -> b -> m ()
addType (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
tn) ([VName]
bound, TypeBase SizeClosure NoUniqueness
t)
    match [VName]
bound (Scalar (Record Map Name StructType
poly_fields)) (Scalar (Record Map Name (TypeBase SizeClosure NoUniqueness)
fields)) =
      [m ()] -> m ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([m ()] -> m ())
-> (Map Name (m ()) -> [m ()]) -> Map Name (m ()) -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name (m ()) -> [m ()]
forall k a. Map k a -> [a]
M.elems (Map Name (m ()) -> m ()) -> Map Name (m ()) -> m ()
forall a b. (a -> b) -> a -> b
$
        (StructType -> TypeBase SizeClosure NoUniqueness -> m ())
-> Map Name StructType
-> Map Name (TypeBase SizeClosure NoUniqueness)
-> Map Name (m ())
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ([VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound) Map Name StructType
poly_fields Map Name (TypeBase SizeClosure NoUniqueness)
fields
    match [VName]
bound (Scalar (Sum Map Name [StructType]
poly_fields)) (Scalar (Sum Map Name [TypeBase SizeClosure NoUniqueness]
fields)) =
      [m ()] -> m ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([m ()] -> m ())
-> (Map Name [m ()] -> [m ()]) -> Map Name [m ()] -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[m ()]] -> [m ()]
forall a. Monoid a => [a] -> a
mconcat ([[m ()]] -> [m ()])
-> (Map Name [m ()] -> [[m ()]]) -> Map Name [m ()] -> [m ()]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [m ()] -> [[m ()]]
forall k a. Map k a -> [a]
M.elems (Map Name [m ()] -> m ()) -> Map Name [m ()] -> m ()
forall a b. (a -> b) -> a -> b
$
        ([StructType] -> [TypeBase SizeClosure NoUniqueness] -> [m ()])
-> Map Name [StructType]
-> Map Name [TypeBase SizeClosure NoUniqueness]
-> Map Name [m ()]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((StructType -> TypeBase SizeClosure NoUniqueness -> m ())
-> [StructType] -> [TypeBase SizeClosure NoUniqueness] -> [m ()]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((StructType -> TypeBase SizeClosure NoUniqueness -> m ())
 -> [StructType] -> [TypeBase SizeClosure NoUniqueness] -> [m ()])
-> (StructType -> TypeBase SizeClosure NoUniqueness -> m ())
-> [StructType]
-> [TypeBase SizeClosure NoUniqueness]
-> [m ()]
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound) Map Name [StructType]
poly_fields Map Name [TypeBase SizeClosure NoUniqueness]
fields
    match
      [VName]
bound
      (Scalar (Arrow NoUniqueness
_ PName
p1 Diet
_ StructType
poly_t1 (RetType [VName]
dims1 TypeBase (ExpBase Info VName) Uniqueness
poly_t2)))
      (Scalar (Arrow NoUniqueness
_ PName
p2 Diet
_ TypeBase SizeClosure NoUniqueness
t1 (RetType [VName]
dims2 TypeBase SizeClosure Uniqueness
t2))) = do
        let bound' :: [VName]
bound' = (PName -> Maybe VName) -> [PName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe PName -> Maybe VName
paramName [PName
p1, PName
p2] [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims1 [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims2 [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
bound
        [VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound' StructType
poly_t1 TypeBase SizeClosure NoUniqueness
t1
        [VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound' (TypeBase (ExpBase Info VName) Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase (ExpBase Info VName) Uniqueness
poly_t2) (TypeBase SizeClosure Uniqueness
-> TypeBase SizeClosure NoUniqueness
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase SizeClosure Uniqueness
t2)
    match [VName]
bound StructType
poly_t TypeBase SizeClosure NoUniqueness
t
      | ExpBase Info VName
d1 : [ExpBase Info VName]
_ <- Shape (ExpBase Info VName) -> [ExpBase Info VName]
forall dim. Shape dim -> [dim]
shapeDims (StructType -> Shape (ExpBase Info VName)
forall dim as. TypeBase dim as -> Shape dim
arrayShape StructType
poly_t),
        SizeClosure
d2 : [SizeClosure]
_ <- Shape SizeClosure -> [SizeClosure]
forall dim. Shape dim -> [dim]
shapeDims (TypeBase SizeClosure NoUniqueness -> Shape SizeClosure
forall dim as. TypeBase dim as -> Shape dim
arrayShape TypeBase SizeClosure NoUniqueness
t) = do
          [VName] -> ExpBase Info VName -> SizeClosure -> m ()
forall {f :: * -> *} {p :: * -> * -> *} {a} {t :: * -> *}.
(MonadState (p a [(VName, SizeClosure)]) f, Bifunctor p,
 Foldable t) =>
t VName -> ExpBase Info VName -> SizeClosure -> f ()
matchDims [VName]
bound ExpBase Info VName
d1 SizeClosure
d2
          [VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound (Int -> StructType -> StructType
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 StructType
poly_t) (Int
-> TypeBase SizeClosure NoUniqueness
-> TypeBase SizeClosure NoUniqueness
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 TypeBase SizeClosure NoUniqueness
t)
    match [VName]
bound StructType
t1 TypeBase SizeClosure NoUniqueness
t2
      | Just StructType
t1' <- StructType -> Maybe StructType
forall d u. TypeBase d u -> Maybe (TypeBase d NoUniqueness)
isAccType StructType
t1,
        Just TypeBase SizeClosure NoUniqueness
t2' <- TypeBase SizeClosure NoUniqueness
-> Maybe (TypeBase SizeClosure NoUniqueness)
forall d u. TypeBase d u -> Maybe (TypeBase d NoUniqueness)
isAccType TypeBase SizeClosure NoUniqueness
t2 =
          [VName] -> StructType -> TypeBase SizeClosure NoUniqueness -> m ()
match [VName]
bound StructType
t1' TypeBase SizeClosure NoUniqueness
t2'
    match [VName]
_ StructType
_ TypeBase SizeClosure NoUniqueness
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
forall a. Monoid a => a
mempty

    matchDims :: t VName -> ExpBase Info VName -> SizeClosure -> f ()
matchDims t VName
bound ExpBase Info VName
e1 (SizeClosure Env
env ExpBase Info VName
e2)
      | ExpBase Info VName
e1 ExpBase Info VName -> ExpBase Info VName -> Bool
forall a. Eq a => a -> a -> Bool
== ExpBase Info VName
anySize Bool -> Bool -> Bool
|| ExpBase Info VName
e2 ExpBase Info VName -> ExpBase Info VName -> Bool
forall a. Eq a => a -> a -> Bool
== ExpBase Info VName
anySize = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
forall a. Monoid a => a
mempty
      | Bool
otherwise = t VName -> Env -> ExpBase Info VName -> ExpBase Info VName -> f ()
forall {p :: * -> * -> *} {a} {m :: * -> *} {t :: * -> *}.
(MonadState (p a [(VName, SizeClosure)]) m, Bifunctor p,
 Foldable t) =>
t VName -> Env -> ExpBase Info VName -> ExpBase Info VName -> m ()
matchExps t VName
bound Env
env ExpBase Info VName
e1 ExpBase Info VName
e2

    matchExps :: t VName -> Env -> ExpBase Info VName -> ExpBase Info VName -> m ()
matchExps t VName
bound Env
env (Var (QualName [VName]
_ VName
d1) Info StructType
_ SrcLoc
_) ExpBase Info VName
e
      | VName
d1 VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
names,
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any VName -> Bool
problematic (Set VName -> Bool) -> Set VName -> Bool
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e =
          VName -> SizeClosure -> m ()
forall {p :: * -> * -> *} {a} {a} {b} {m :: * -> *}.
(MonadState (p a [(a, b)]) m, Bifunctor p, Ord a) =>
a -> b -> m ()
addDim VName
d1 (Env -> ExpBase Info VName -> SizeClosure
SizeClosure Env
env ExpBase Info VName
e)
      where
        problematic :: VName -> Bool
problematic VName
v = VName
v VName -> t VName -> Bool
forall a. Eq a => a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
bound Bool -> Bool -> Bool
|| VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
names
    matchExps t VName
bound Env
env ExpBase Info VName
e1 ExpBase Info VName
e2
      | Just [(ExpBase Info VName, ExpBase Info VName)]
es <- ExpBase Info VName
-> ExpBase Info VName
-> Maybe [(ExpBase Info VName, ExpBase Info VName)]
similarExps ExpBase Info VName
e1 ExpBase Info VName
e2 =
          ((ExpBase Info VName, ExpBase Info VName) -> m ())
-> [(ExpBase Info VName, ExpBase Info VName)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((ExpBase Info VName -> ExpBase Info VName -> m ())
-> (ExpBase Info VName, ExpBase Info VName) -> m ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((ExpBase Info VName -> ExpBase Info VName -> m ())
 -> (ExpBase Info VName, ExpBase Info VName) -> m ())
-> (ExpBase Info VName -> ExpBase Info VName -> m ())
-> (ExpBase Info VName, ExpBase Info VName)
-> m ()
forall a b. (a -> b) -> a -> b
$ t VName -> Env -> ExpBase Info VName -> ExpBase Info VName -> m ()
matchExps t VName
bound Env
env) [(ExpBase Info VName, ExpBase Info VName)]
es
    matchExps t VName
_ Env
_ ExpBase Info VName
_ ExpBase Info VName
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
forall a. Monoid a => a
mempty

evalWithExts :: Env -> Exp -> EvalM Value
evalWithExts :: Env -> ExpBase Info VName -> EvalM Value
evalWithExts Env
env ExpBase Info VName
e = do
  Env
size_env <- EvalM Env
extEnv
  Env -> ExpBase Info VName -> EvalM Value
eval (Env
size_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env) ExpBase Info VName
e

evalResolved ::
  ([(VName, ([VName], EvalType))], [(VName, SizeClosure)]) ->
  EvalM Env
evalResolved :: ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
 [(VName, SizeClosure)])
-> EvalM Env
evalResolved ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))]
ts, [(VName, SizeClosure)]
ds) = do
  [(VName, StructType)]
ts' <- ((VName, ([VName], TypeBase SizeClosure NoUniqueness))
 -> EvalM (VName, StructType))
-> [(VName, ([VName], TypeBase SizeClosure NoUniqueness))]
-> EvalM [(VName, StructType)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((([VName], TypeBase SizeClosure NoUniqueness) -> EvalM StructType)
-> (VName, ([VName], TypeBase SizeClosure NoUniqueness))
-> EvalM (VName, StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (VName, a) -> f (VName, b)
traverse ((([VName], TypeBase SizeClosure NoUniqueness) -> EvalM StructType)
 -> (VName, ([VName], TypeBase SizeClosure NoUniqueness))
 -> EvalM (VName, StructType))
-> (([VName], TypeBase SizeClosure NoUniqueness)
    -> EvalM StructType)
-> (VName, ([VName], TypeBase SizeClosure NoUniqueness))
-> EvalM (VName, StructType)
forall a b. (a -> b) -> a -> b
$ \([VName]
bound, TypeBase SizeClosure NoUniqueness
t) -> (Either Int64 SizeClosure -> ExpBase Info VName)
-> TypeBase (Either Int64 SizeClosure) NoUniqueness -> StructType
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Either Int64 SizeClosure -> ExpBase Info VName
forall {a}.
Integral a =>
Either a SizeClosure -> ExpBase Info VName
onDim (TypeBase (Either Int64 SizeClosure) NoUniqueness -> StructType)
-> EvalM (TypeBase (Either Int64 SizeClosure) NoUniqueness)
-> EvalM StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> TypeBase SizeClosure NoUniqueness
-> EvalM (TypeBase (Either Int64 SizeClosure) NoUniqueness)
evalType ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
bound) TypeBase SizeClosure NoUniqueness
t) [(VName, ([VName], TypeBase SizeClosure NoUniqueness))]
ts
  [(VName, Int64)]
ds' <- ((VName, SizeClosure) -> EvalM (VName, Int64))
-> [(VName, SizeClosure)] -> EvalM [(VName, Int64)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SizeClosure -> EvalM Int64)
-> (VName, SizeClosure) -> EvalM (VName, Int64)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (VName, a) -> f (VName, b)
traverse ((SizeClosure -> EvalM Int64)
 -> (VName, SizeClosure) -> EvalM (VName, Int64))
-> (SizeClosure -> EvalM Int64)
-> (VName, SizeClosure)
-> EvalM (VName, Int64)
forall a b. (a -> b) -> a -> b
$ \(SizeClosure Env
env ExpBase Info VName
e) -> Value -> Int64
asInt64 (Value -> Int64) -> EvalM Value -> EvalM Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
evalWithExts Env
env ExpBase Info VName
e) [(VName, SizeClosure)]
ds
  Env -> EvalM Env
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env -> EvalM Env) -> Env -> EvalM Env
forall a b. (a -> b) -> a -> b
$ Map VName StructType -> Env
typeEnv ([(VName, StructType)] -> Map VName StructType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, StructType)]
ts') Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Map VName Int64 -> Env
i64Env ([(VName, Int64)] -> Map VName Int64
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, Int64)]
ds')
  where
    onDim :: Either a SizeClosure -> ExpBase Info VName
onDim (Left a
x) = Integer -> SrcLoc -> ExpBase Info VName
sizeFromInteger (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
x) SrcLoc
forall a. Monoid a => a
mempty
    onDim (Right (SizeClosure Env
_ ExpBase Info VName
e)) = ExpBase Info VName
e -- FIXME

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

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

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

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

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

asInt :: Value -> Int
asInt :: Value -> Int
asInt = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> (Value -> Integer) -> Value -> Int
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 (ValueAD Int
_ ADVariable
v)
  | P.IntValue IntValue
v' <- ADVariable -> PrimValue
AD.varPrimal ADVariable
v = IntValue
v'
asSigned Value
v = [Char] -> IntValue
forall a. HasCallStack => [Char] -> a
error ([Char] -> IntValue) -> [Char] -> IntValue
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpectedly not a signed integer: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v

asInt64 :: Value -> Int64
asInt64 :: Value -> Int64
asInt64 = Integer -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int64) -> (Value -> Integer) -> Value -> Int64
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 (ValueAD Int
_ ADVariable
v)
  | P.BoolValue Bool
v' <- ADVariable -> PrimValue
AD.varPrimal ADVariable
v = Bool
v'
asBool Value
v = [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error ([Char] -> Bool) -> [Char] -> Bool
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpectedly not a boolean: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v

lookupInEnv ::
  (Env -> M.Map VName x) ->
  QualName VName ->
  Env ->
  Maybe x
lookupInEnv :: forall x. (Env -> Map VName x) -> QualName VName -> Env -> Maybe x
lookupInEnv Env -> Map VName x
onEnv QualName VName
qv Env
env = Env -> [VName] -> Maybe x
f Env
env ([VName] -> Maybe x) -> [VName] -> Maybe x
forall a b. (a -> b) -> a -> b
$ QualName VName -> [VName]
forall vn. QualName vn -> [vn]
qualQuals QualName VName
qv
  where
    f :: Env -> [VName] -> Maybe x
f Env
m (VName
q : [VName]
qs) =
      case VName -> Map VName TermBinding -> Maybe TermBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q (Map VName TermBinding -> Maybe TermBinding)
-> Map VName TermBinding -> Maybe TermBinding
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
_ -> Maybe x
forall a. Maybe a
Nothing
    f Env
m [] = VName -> Map VName x -> Maybe x
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qv) (Map VName x -> Maybe x) -> Map VName x -> Maybe x
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 = (Env -> Map VName TermBinding)
-> QualName VName -> Env -> Maybe TermBinding
forall x. (Env -> Map VName x) -> QualName VName -> Env -> Maybe x
lookupInEnv Env -> Map VName TermBinding
envTerm

lookupType :: QualName VName -> Env -> Maybe TypeBinding
lookupType :: QualName VName -> Env -> Maybe TypeBinding
lookupType = (Env -> Map VName TypeBinding)
-> QualName VName -> Env -> Maybe TypeBinding
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 or
-- an existential.
data TermBinding
  = TermValue (Maybe T.BoundV) Value
  | -- | A polymorphic value that must be instantiated.  The
    --  'StructType' provided is un-evaluated, but parts of it can be
    --  evaluated using the provided 'Eval' function.
    TermPoly (Maybe T.BoundV) (EvalType -> EvalM Value)
  | TermModule Module

instance Show TermBinding where
  show :: TermBinding -> [Char]
show (TermValue Maybe BoundV
bv Value
v) = [[Char]] -> [Char]
unwords [[Char]
"TermValue", Maybe BoundV -> [Char]
forall a. Show a => a -> [Char]
show Maybe BoundV
bv, Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v]
  show (TermPoly Maybe BoundV
bv TypeBase SizeClosure NoUniqueness -> EvalM Value
_) = [[Char]] -> [Char]
unwords [[Char]
"TermPoly", Maybe BoundV -> [Char]
forall a. Show a => a -> [Char]
show Maybe BoundV
bv]
  show (TermModule Module
m) = [[Char]] -> [Char]
unwords [[Char]
"TermModule", Module -> [Char]
forall a. Show a => a -> [Char]
show Module
m]

data TypeBinding = TypeBinding Env [TypeParam] StructRetType
  deriving (Int -> TypeBinding -> ShowS
[TypeBinding] -> ShowS
TypeBinding -> [Char]
(Int -> TypeBinding -> ShowS)
-> (TypeBinding -> [Char])
-> ([TypeBinding] -> ShowS)
-> Show TypeBinding
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TypeBinding -> ShowS
showsPrec :: Int -> TypeBinding -> ShowS
$cshow :: TypeBinding -> [Char]
show :: TypeBinding -> [Char]
$cshowList :: [TypeBinding] -> ShowS
showList :: [TypeBinding] -> ShowS
Show)

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

instance Show Module where
  show :: Module -> [Char]
show (Module Env
env) = [Char]
"(" [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [[Char]] -> [Char]
unwords [[Char]
"Module", Env -> [Char]
forall a. Show a => a -> [Char]
show Env
env] [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
")"
  show (ModuleFun Module -> EvalM Module
_) = [Char]
"(ModuleFun _)"

-- | 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 TypeBinding
  }
  deriving (Int -> Env -> ShowS
[Env] -> ShowS
Env -> [Char]
(Int -> Env -> ShowS)
-> (Env -> [Char]) -> ([Env] -> ShowS) -> Show Env
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Env -> ShowS
showsPrec :: Int -> Env -> ShowS
$cshow :: Env -> [Char]
show :: Env -> [Char]
$cshowList :: [Env] -> ShowS
showList :: [Env] -> ShowS
Show)

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

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

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

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

valEnv :: M.Map VName (Maybe T.BoundV, Value) -> Env
valEnv :: Map VName (Maybe BoundV, Value) -> Env
valEnv Map VName (Maybe BoundV, Value)
m =
  Env
    { envTerm :: Map VName TermBinding
envTerm = ((Maybe BoundV, Value) -> TermBinding)
-> Map VName (Maybe BoundV, Value) -> Map VName TermBinding
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((Maybe BoundV -> Value -> TermBinding)
-> (Maybe BoundV, Value) -> TermBinding
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 = Map VName TypeBinding
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 = (Module -> TermBinding)
-> Map VName Module -> Map VName TermBinding
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 = Map VName TypeBinding
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 = Map VName TermBinding
forall a. Monoid a => a
mempty,
      envType :: Map VName TypeBinding
envType = (StructType -> TypeBinding)
-> Map VName StructType -> Map VName TypeBinding
forall a b k. (a -> b) -> Map k a -> Map k b
M.map StructType -> TypeBinding
tbind Map VName StructType
m
    }
  where
    tbind :: StructType -> TypeBinding
tbind = Env -> [TypeParam] -> StructRetType -> TypeBinding
TypeBinding Env
forall a. Monoid a => a
mempty [] (StructRetType -> TypeBinding)
-> (StructType -> StructRetType) -> StructType -> TypeBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> StructType -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType []

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

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

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

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

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

break :: Env -> Loc -> EvalM ()
break :: Env -> Loc -> EvalM ()
break Env
env Loc
loc = do
  Map ImportName Env
imports <- (([StackFrame], Map ImportName Env) -> Map ImportName Env)
-> EvalM (Map ImportName Env)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ([StackFrame], Map ImportName Env) -> Map ImportName Env
forall a b. (a, b) -> b
snd
  NonEmpty StackFrame
backtrace <- (([StackFrame], Map ImportName Env) -> NonEmpty StackFrame)
-> EvalM (NonEmpty StackFrame)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Loc -> Ctx -> StackFrame
StackFrame Loc
loc (Env -> Map ImportName Env -> Ctx
Ctx Env
env Map ImportName Env
imports) NE.:|) ([StackFrame] -> NonEmpty StackFrame)
-> (([StackFrame], Map ImportName Env) -> [StackFrame])
-> ([StackFrame], Map ImportName Env)
-> NonEmpty StackFrame
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([StackFrame], Map ImportName Env) -> [StackFrame]
forall a b. (a, b) -> a
fst)
  ExtOp () -> EvalM ()
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (ExtOp () -> EvalM ()) -> ExtOp () -> EvalM ()
forall a b. (a -> b) -> a -> b
$ Loc -> BreakReason -> NonEmpty StackFrame -> () -> ExtOp ()
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, Array Int Value -> [Value]
forall i e. Array i e -> [e]
elems Array Int Value
as)
fromArray Value
v = [Char] -> (ValueShape, [Value])
forall a. HasCallStack => [Char] -> a
error ([Char] -> (ValueShape, [Value]))
-> [Char] -> (ValueShape, [Value])
forall a b. (a -> b) -> a -> b
$ [Char]
"Expected array value, but found: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v

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

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

matchPat :: Env -> Pat (TypeBase Size u) -> Value -> EvalM Env
matchPat :: forall u.
Env -> Pat (TypeBase (ExpBase Info VName) u) -> Value -> EvalM Env
matchPat Env
env Pat (TypeBase (ExpBase Info VName) u)
p Value
v = do
  Maybe Env
m <- MaybeT EvalM Env -> EvalM (Maybe Env)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT EvalM Env -> EvalM (Maybe Env))
-> MaybeT EvalM Env -> EvalM (Maybe Env)
forall a b. (a -> b) -> a -> b
$ Env
-> Pat (TypeBase (ExpBase Info VName) u)
-> Value
-> MaybeT EvalM Env
forall u.
Env
-> Pat (TypeBase (ExpBase Info VName) u)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env Pat (TypeBase (ExpBase Info VName) u)
p Value
v
  case Maybe Env
m of
    Maybe Env
Nothing -> [Char] -> EvalM Env
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Env) -> [Char] -> EvalM Env
forall a b. (a -> b) -> a -> b
$ [Char]
"matchPat: missing case for " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> PatBase Info VName StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (TypeBase (ExpBase Info VName) u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct (TypeBase (ExpBase Info VName) u -> StructType)
-> Pat (TypeBase (ExpBase Info VName) u)
-> PatBase Info VName StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat (TypeBase (ExpBase Info VName) u)
p) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" and " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v
    Just Env
env' -> Env -> EvalM Env
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env'

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

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

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

indexesFor ::
  Maybe Int64 ->
  Maybe Int64 ->
  Maybe Int64 ->
  Int64 ->
  Maybe [Int]
indexesFor :: Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride Int64
n
  | (Int64
start', Int64
end', Int64
stride') <- (Int64, Int64, Int64)
slice,
    Int64
end' Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
start' Bool -> Bool -> Bool
|| Int64 -> Int64
forall p. (Eq p, Num p) => p -> p
signum' (Int64
end' Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
start') Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64 -> Int64
forall p. (Eq p, Num p) => p -> p
signum' Int64
stride',
    Int64
stride' Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Int64
0,
    [Int64]
is <- [Int64
start', Int64
start' Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
stride' .. Int64
end' Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64 -> Int64
forall a. Num a => a -> a
signum Int64
stride'],
    (Int64 -> Bool) -> [Int64] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Int64 -> Bool
inBounds [Int64]
is =
      [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just ([Int] -> Maybe [Int]) -> [Int] -> Maybe [Int]
forall a b. (a -> b) -> a -> b
$ (Int64 -> Int) -> [Int64] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int64]
is
  | Bool
otherwise =
      Maybe [Int]
forall a. Maybe a
Nothing
  where
    inBounds :: Int64 -> Bool
inBounds Int64
i = Int64
i Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Int64
0 Bool -> Bool -> Bool
&& Int64
i Int64 -> Int64 -> Bool
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' = Int64 -> Maybe Int64 -> Int64
forall a. a -> Maybe a -> a
fromMaybe Int64
n Maybe Int64
end
           in (Int64
start', Int64
end', Int64 -> Maybe Int64 -> Int64
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', Int64 -> Maybe Int64 -> Int64
forall a. a -> Maybe a -> a
fromMaybe Int64
1 Maybe Int64
stride)
        (Maybe Int64
Nothing, Maybe Int64
Nothing, Just Int64
stride') ->
          ( if Int64
stride' Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Int64
0 then Int64
0 else Int64
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1,
            if Int64
stride' Int64 -> Int64 -> Bool
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 = p -> p
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) =
  Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim Int64
n (ValueShape -> ValueShape) -> ValueShape -> ValueShape
forall a b. (a -> b) -> a -> b
$ [Indexing] -> ValueShape -> ValueShape
indexShape [Indexing]
is ValueShape
shape
  where
    n :: Int64
n = Int64 -> ([Int] -> Int64) -> Maybe [Int] -> Int64
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int64
0 [Int] -> Int64
forall i a. Num i => [a] -> i
genericLength (Maybe [Int] -> Int64) -> Maybe [Int] -> Int64
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 Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Int64
0,
    Int64
i Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
n =
      [Indexing] -> Value -> Maybe Value
indexArray [Indexing]
is (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Array Int Value
arr Array Int Value -> Int -> Value
forall i e. Ix i => Array i e -> i -> e
! Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i
  | Bool
otherwise =
      Maybe Value
forall a. Maybe a
Nothing
  where
    n :: Int64
n = Array Int Value -> Int64
forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr
indexArray (IndexingSlice Maybe Int64
start Maybe Int64
end Maybe Int64
stride : [Indexing]
is) (ValueArray (ShapeDim Int64
_ ValueShape
rowshape) Array Int Value
arr) = do
  [Int]
js <- Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride (Int64 -> Maybe [Int]) -> Int64 -> Maybe [Int]
forall a b. (a -> b) -> a -> b
$ Array Int Value -> Int64
forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr
  ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ([Indexing] -> ValueShape -> ValueShape
indexShape [Indexing]
is ValueShape
rowshape) ([Value] -> Value) -> Maybe [Value] -> Maybe Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Maybe Value) -> [Int] -> Maybe [Value]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Indexing] -> Value -> Maybe Value
indexArray [Indexing]
is (Value -> Maybe Value) -> (Int -> Value) -> Int -> Maybe Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array Int Value
arr !)) [Int]
js
indexArray [Indexing]
_ Value
v = Value -> Maybe Value
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 = Identity (Maybe Value) -> Maybe Value
forall a. Identity a -> a
runIdentity (Identity (Maybe Value) -> Maybe Value)
-> Identity (Maybe Value) -> Maybe Value
forall a b. (a -> b) -> a -> b
$ (Value -> Value -> Identity Value)
-> [Indexing] -> Value -> Value -> Identity (Maybe Value)
forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (\Value
_ Value
y' -> Value -> Identity Value
forall a. a -> Identity a
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 Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Int64
0,
    Int64
i Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
n = do
      Maybe Value
v' <- (Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
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 Array Int Value -> Int -> Value
forall i e. Ix i => Array i e -> i -> e
! Int
i') Value
v
      Maybe Value -> m (Maybe Value)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Value -> m (Maybe Value)) -> Maybe Value -> m (Maybe Value)
forall a b. (a -> b) -> a -> b
$ do
        Value
v'' <- Maybe Value
v'
        Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> Array Int Value -> Value
forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
shape (Array Int Value -> Value) -> Array Int Value -> Value
forall a b. (a -> b) -> a -> b
$ Array Int Value
arr Array Int Value -> [(Int, Value)] -> Array Int Value
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Int
i', Value
v'')]
  | Bool
otherwise =
      Maybe Value -> m (Maybe Value)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Value
forall a. Maybe a
Nothing
  where
    n :: Int64
n = Array Int Value -> Int64
forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr
    i' :: Int
i' = Int64 -> Int
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 (Int64 -> Maybe [Int]) -> Int64 -> Maybe [Int]
forall a b. (a -> b) -> a -> b
$ Array Int Value -> Int64
forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
arr,
    [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
arr_is Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Array Int Value -> Int
forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int Value
v = do
      let update :: Maybe (Array Int Value)
-> (Int, Value) -> m (Maybe (Array Int Value))
update (Just Array Int Value
arr') (Int
i, Value
v') = do
            Maybe Value
x <- (Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
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 Array Int Value -> Int -> Value
forall i e. Ix i => Array i e -> i -> e
! Int
i) Value
v'
            Maybe (Array Int Value) -> m (Maybe (Array Int Value))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Array Int Value) -> m (Maybe (Array Int Value)))
-> Maybe (Array Int Value) -> m (Maybe (Array Int Value))
forall a b. (a -> b) -> a -> b
$ do
              Value
x' <- Maybe Value
x
              Array Int Value -> Maybe (Array Int Value)
forall a. a -> Maybe a
Just (Array Int Value -> Maybe (Array Int Value))
-> Array Int Value -> Maybe (Array Int Value)
forall a b. (a -> b) -> a -> b
$ Array Int Value
arr' Array Int Value -> [(Int, Value)] -> Array Int Value
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)
_ = Maybe (Array Int Value) -> m (Maybe (Array Int Value))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Array Int Value)
forall a. Maybe a
Nothing
      (Maybe (Array Int Value) -> Maybe Value)
-> m (Maybe (Array Int Value)) -> m (Maybe Value)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Array Int Value -> Value)
-> Maybe (Array Int Value) -> Maybe Value
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ValueShape -> Array Int Value -> Value
forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
shape)) (m (Maybe (Array Int Value)) -> m (Maybe Value))
-> m (Maybe (Array Int Value)) -> m (Maybe Value)
forall a b. (a -> b) -> a -> b
$ (Maybe (Array Int Value)
 -> (Int, Value) -> m (Maybe (Array Int Value)))
-> Maybe (Array Int Value)
-> [(Int, Value)]
-> m (Maybe (Array Int Value))
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 (Array Int Value -> Maybe (Array Int Value)
forall a. a -> Maybe a
Just Array Int Value
arr) ([(Int, Value)] -> m (Maybe (Array Int Value)))
-> [(Int, Value)] -> m (Maybe (Array Int Value))
forall a b. (a -> b) -> a -> b
$ [Int] -> [Value] -> [(Int, Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
arr_is ([Value] -> [(Int, Value)]) -> [Value] -> [(Int, Value)]
forall a b. (a -> b) -> a -> b
$ Array Int Value -> [Value]
forall i e. Array i e -> [e]
elems Array Int Value
v
  | Bool
otherwise =
      Maybe Value -> m (Maybe Value)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Value
forall a. Maybe a
Nothing
updateArray Value -> Value -> m Value
f [Indexing]
_ Value
x Value
y = Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> m Value -> m (Maybe Value)
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 ExpBase Info VName
x) =
  Int64 -> Indexing
IndexingFix (Int64 -> Indexing) -> (Value -> Int64) -> Value -> Indexing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64 (Value -> Indexing) -> EvalM Value -> EvalM Indexing
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
x
evalDimIndex Env
env (DimSlice Maybe (ExpBase Info VName)
start Maybe (ExpBase Info VName)
end Maybe (ExpBase Info VName)
stride) =
  Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Indexing
IndexingSlice
    (Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Indexing)
-> EvalM (Maybe Int64)
-> EvalM (Maybe Int64 -> Maybe Int64 -> Indexing)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> EvalM Int64)
-> Maybe (ExpBase Info VName) -> EvalM (Maybe Int64)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ((Value -> Int64) -> EvalM Value -> EvalM Int64
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Int64
asInt64 (EvalM Value -> EvalM Int64)
-> (ExpBase Info VName -> EvalM Value)
-> ExpBase Info VName
-> EvalM Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ExpBase Info VName -> EvalM Value
eval Env
env) Maybe (ExpBase Info VName)
start
    EvalM (Maybe Int64 -> Maybe Int64 -> Indexing)
-> EvalM (Maybe Int64) -> EvalM (Maybe Int64 -> Indexing)
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (ExpBase Info VName -> EvalM Int64)
-> Maybe (ExpBase Info VName) -> EvalM (Maybe Int64)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ((Value -> Int64) -> EvalM Value -> EvalM Int64
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Int64
asInt64 (EvalM Value -> EvalM Int64)
-> (ExpBase Info VName -> EvalM Value)
-> ExpBase Info VName
-> EvalM Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ExpBase Info VName -> EvalM Value
eval Env
env) Maybe (ExpBase Info VName)
end
    EvalM (Maybe Int64 -> Indexing)
-> EvalM (Maybe Int64) -> EvalM Indexing
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (ExpBase Info VName -> EvalM Int64)
-> Maybe (ExpBase Info VName) -> EvalM (Maybe Int64)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ((Value -> Int64) -> EvalM Value -> EvalM Int64
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Int64
asInt64 (EvalM Value -> EvalM Int64)
-> (ExpBase Info VName -> EvalM Value)
-> ExpBase Info VName
-> EvalM Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ExpBase Info VName -> EvalM Value
eval Env
env) Maybe (ExpBase Info VName)
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 =
        SrcLoc -> Env -> Text -> EvalM a
forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env (Text -> EvalM a) -> Text -> EvalM a
forall a b. (a -> b) -> a -> b
$
          Text
"Index ["
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " ((Indexing -> Text) -> [Indexing] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Indexing -> Text
forall a. Pretty a => a -> Text
prettyText [Indexing]
is)
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"] out of bounds for array of shape "
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ValueShape -> Text
forall a. Pretty a => a -> Text
prettyText (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
arr)
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"."
  EvalM Value -> (Value -> EvalM Value) -> Maybe Value -> EvalM Value
forall b a. b -> (a -> b) -> Maybe a -> b
maybe EvalM Value
forall {a}. EvalM a
oob Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Value -> EvalM Value) -> Maybe Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Maybe Value
indexArray [Indexing]
is Value
arr

-- | Expand type based on information that was not available at
-- type-checking time (the structure of abstract types).
expandType :: (Pretty u) => Env -> TypeBase Size u -> TypeBase SizeClosure u
expandType :: forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
_ (Scalar (Prim PrimType
pt)) = ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u)
-> ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase SizeClosure u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
pt
expandType Env
env (Scalar (Record Map Name (TypeBase (ExpBase Info VName) u)
fs)) = ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u)
-> ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase SizeClosure u) -> ScalarTypeBase SizeClosure u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase SizeClosure u) -> ScalarTypeBase SizeClosure u)
-> Map Name (TypeBase SizeClosure u)
-> ScalarTypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ (TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u)
-> Map Name (TypeBase (ExpBase Info VName) u)
-> Map Name (TypeBase SizeClosure u)
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env) Map Name (TypeBase (ExpBase Info VName) u)
fs
expandType Env
env (Scalar (Arrow u
u PName
p Diet
d StructType
t1 (RetType [VName]
dims TypeBase (ExpBase Info VName) Uniqueness
t2))) =
  ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u)
-> ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ u
-> PName
-> Diet
-> TypeBase SizeClosure NoUniqueness
-> RetTypeBase SizeClosure Uniqueness
-> ScalarTypeBase SizeClosure u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
u PName
p Diet
d (Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env StructType
t1) ([VName]
-> TypeBase SizeClosure Uniqueness
-> RetTypeBase SizeClosure Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (Env
-> TypeBase (ExpBase Info VName) Uniqueness
-> TypeBase SizeClosure Uniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env TypeBase (ExpBase Info VName) Uniqueness
t2))
expandType Env
env t :: TypeBase (ExpBase Info VName) u
t@(Array u
u Shape (ExpBase Info VName)
shape ScalarTypeBase (ExpBase Info VName) NoUniqueness
_) =
  let et :: TypeBase (ExpBase Info VName) u
et = Int
-> TypeBase (ExpBase Info VName) u
-> TypeBase (ExpBase Info VName) u
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray (Shape (ExpBase Info VName) -> Int
forall dim. Shape dim -> Int
shapeRank Shape (ExpBase Info VName)
shape) TypeBase (ExpBase Info VName) u
t
      et' :: TypeBase SizeClosure u
et' = Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env TypeBase (ExpBase Info VName) u
et
      shape' :: Shape SizeClosure
shape' = (ExpBase Info VName -> SizeClosure)
-> Shape (ExpBase Info VName) -> Shape SizeClosure
forall a b. (a -> b) -> Shape a -> Shape b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Env -> ExpBase Info VName -> SizeClosure
SizeClosure Env
env) Shape (ExpBase Info VName)
shape
   in (NoUniqueness -> u)
-> TypeBase SizeClosure NoUniqueness -> TypeBase SizeClosure u
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (u -> NoUniqueness -> u
forall a b. a -> b -> a
const u
u) (Shape SizeClosure
-> TypeBase SizeClosure NoUniqueness
-> TypeBase SizeClosure NoUniqueness
forall dim.
Shape dim -> TypeBase dim NoUniqueness -> TypeBase dim NoUniqueness
arrayOf Shape SizeClosure
shape' (TypeBase SizeClosure NoUniqueness
 -> TypeBase SizeClosure NoUniqueness)
-> TypeBase SizeClosure NoUniqueness
-> TypeBase SizeClosure NoUniqueness
forall a b. (a -> b) -> a -> b
$ TypeBase SizeClosure u -> TypeBase SizeClosure NoUniqueness
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase SizeClosure u
et')
expandType Env
env (Scalar (TypeVar u
u QualName VName
tn [TypeArg (ExpBase Info VName)]
args)) =
  case QualName VName -> Env -> Maybe TypeBinding
lookupType QualName VName
tn Env
env of
    Just (TypeBinding Env
tn_env [TypeParam]
ps (RetType [VName]
ext StructType
t')) ->
      let (Map VName (Subst t)
substs, Map VName TypeBinding
types) = [(Map VName (Subst t), Map VName TypeBinding)]
-> (Map VName (Subst t), Map VName TypeBinding)
forall a. Monoid a => [a] -> a
mconcat ([(Map VName (Subst t), Map VName TypeBinding)]
 -> (Map VName (Subst t), Map VName TypeBinding))
-> [(Map VName (Subst t), Map VName TypeBinding)]
-> (Map VName (Subst t), Map VName TypeBinding)
forall a b. (a -> b) -> a -> b
$ (TypeParam
 -> TypeArg (ExpBase Info VName)
 -> (Map VName (Subst t), Map VName TypeBinding))
-> [TypeParam]
-> [TypeArg (ExpBase Info VName)]
-> [(Map VName (Subst t), Map VName TypeBinding)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeParam
-> TypeArg (ExpBase Info VName)
-> (Map VName (Subst t), Map VName TypeBinding)
forall {k} {t}.
Ord k =>
TypeParamBase k
-> TypeArg (ExpBase Info VName)
-> (Map k (Subst t), Map k TypeBinding)
matchPtoA [TypeParam]
ps [TypeArg (ExpBase Info VName)]
args
          onDim :: SizeClosure -> SizeClosure
onDim (SizeClosure Env
dim_env ExpBase Info VName
dim)
            | (VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
ext) (Set VName -> Bool) -> Set VName -> Bool
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> FV
freeInExp ExpBase Info VName
dim =
                -- The case can occur when a type with existential
                -- size has been hidden by a module ascription, e.g.
                -- tests/modules/sizeparams4.fut.
                Env -> ExpBase Info VName -> SizeClosure
SizeClosure Env
forall a. Monoid a => a
mempty ExpBase Info VName
anySize
            | Bool
otherwise =
                Env -> ExpBase Info VName -> SizeClosure
SizeClosure (Env
env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
dim_env) (ExpBase Info VName -> SizeClosure)
-> ExpBase Info VName -> SizeClosure
forall a b. (a -> b) -> a -> b
$
                  TypeSubs -> ExpBase Info VName -> ExpBase Info VName
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
forall {t}. Map VName (Subst t)
substs) ExpBase Info VName
dim
       in (SizeClosure -> SizeClosure)
-> (NoUniqueness -> u)
-> TypeBase SizeClosure NoUniqueness
-> TypeBase SizeClosure u
forall a b c d.
(a -> b) -> (c -> d) -> TypeBase a c -> TypeBase b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap SizeClosure -> SizeClosure
onDim (u -> NoUniqueness -> u
forall a b. a -> b -> a
const u
u) (TypeBase SizeClosure NoUniqueness -> TypeBase SizeClosure u)
-> TypeBase SizeClosure NoUniqueness -> TypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType (Map VName TermBinding -> Map VName TypeBinding -> Env
Env Map VName TermBinding
forall a. Monoid a => a
mempty Map VName TypeBinding
types Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
tn_env) StructType
t'
    Maybe TypeBinding
Nothing ->
      -- This case only happens for built-in abstract types,
      -- e.g. accumulators.
      ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (u
-> QualName VName
-> [TypeArg SizeClosure]
-> ScalarTypeBase SizeClosure u
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar u
u QualName VName
tn ([TypeArg SizeClosure] -> ScalarTypeBase SizeClosure u)
-> [TypeArg SizeClosure] -> ScalarTypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ (TypeArg (ExpBase Info VName) -> TypeArg SizeClosure)
-> [TypeArg (ExpBase Info VName)] -> [TypeArg SizeClosure]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg (ExpBase Info VName) -> TypeArg SizeClosure
expandArg [TypeArg (ExpBase Info VName)]
args)
  where
    matchPtoA :: TypeParamBase k
-> TypeArg (ExpBase Info VName)
-> (Map k (Subst t), Map k TypeBinding)
matchPtoA (TypeParamDim k
p SrcLoc
_) (TypeArgDim ExpBase Info VName
e) =
      (k -> Subst t -> Map k (Subst t)
forall k a. k -> a -> Map k a
M.singleton k
p (Subst t -> Map k (Subst t)) -> Subst t -> Map k (Subst t)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> Subst t
forall t. ExpBase Info VName -> Subst t
ExpSubst ExpBase Info VName
e, Map k TypeBinding
forall a. Monoid a => a
mempty)
    matchPtoA (TypeParamType Liftedness
_ k
p SrcLoc
_) (TypeArgType StructType
t') =
      let t'' :: StructType
t'' = TypeBase SizeClosure NoUniqueness -> StructType
evalToStruct (TypeBase SizeClosure NoUniqueness -> StructType)
-> TypeBase SizeClosure NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env StructType
t' -- FIXME, we are throwing away the closure here.
       in (Map k (Subst t)
forall a. Monoid a => a
mempty, k -> TypeBinding -> Map k TypeBinding
forall k a. k -> a -> Map k a
M.singleton k
p (Env -> [TypeParam] -> StructRetType -> TypeBinding
TypeBinding Env
forall a. Monoid a => a
mempty [] (StructRetType -> TypeBinding) -> StructRetType -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] StructType
t''))
    matchPtoA TypeParamBase k
_ TypeArg (ExpBase Info VName)
_ = (Map k (Subst t), Map k TypeBinding)
forall a. Monoid a => a
mempty
    expandArg :: TypeArg (ExpBase Info VName) -> TypeArg SizeClosure
expandArg (TypeArgDim ExpBase Info VName
s) = SizeClosure -> TypeArg SizeClosure
forall dim. dim -> TypeArg dim
TypeArgDim (SizeClosure -> TypeArg SizeClosure)
-> SizeClosure -> TypeArg SizeClosure
forall a b. (a -> b) -> a -> b
$ Env -> ExpBase Info VName -> SizeClosure
SizeClosure Env
env ExpBase Info VName
s
    expandArg (TypeArgType StructType
t) = TypeBase SizeClosure NoUniqueness -> TypeArg SizeClosure
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (TypeBase SizeClosure NoUniqueness -> TypeArg SizeClosure)
-> TypeBase SizeClosure NoUniqueness -> TypeArg SizeClosure
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env StructType
t
expandType Env
env (Scalar (Sum Map Name [TypeBase (ExpBase Info VName) u]
cs)) = ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u)
-> ScalarTypeBase SizeClosure u -> TypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase SizeClosure u] -> ScalarTypeBase SizeClosure u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase SizeClosure u] -> ScalarTypeBase SizeClosure u)
-> Map Name [TypeBase SizeClosure u]
-> ScalarTypeBase SizeClosure u
forall a b. (a -> b) -> a -> b
$ (([TypeBase (ExpBase Info VName) u] -> [TypeBase SizeClosure u])
-> Map Name [TypeBase (ExpBase Info VName) u]
-> Map Name [TypeBase SizeClosure u]
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([TypeBase (ExpBase Info VName) u] -> [TypeBase SizeClosure u])
 -> Map Name [TypeBase (ExpBase Info VName) u]
 -> Map Name [TypeBase SizeClosure u])
-> ((TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u)
    -> [TypeBase (ExpBase Info VName) u] -> [TypeBase SizeClosure u])
-> (TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u)
-> Map Name [TypeBase (ExpBase Info VName) u]
-> Map Name [TypeBase SizeClosure u]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u)
-> [TypeBase (ExpBase Info VName) u] -> [TypeBase SizeClosure u]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env) Map Name [TypeBase (ExpBase Info VName) u]
cs

-- | Evaluate all possible sizes, except those that contain free
-- variables in the set of names.
evalType :: S.Set VName -> EvalType -> EvalM (TypeBase (Either Int64 SizeClosure) NoUniqueness)
evalType :: Set VName
-> TypeBase SizeClosure NoUniqueness
-> EvalM (TypeBase (Either Int64 SizeClosure) NoUniqueness)
evalType Set VName
outer_bound TypeBase SizeClosure NoUniqueness
t = do
  let evalDim :: Set VName -> p -> SizeClosure -> EvalM (Either Int64 SizeClosure)
evalDim Set VName
bound p
_ (SizeClosure Env
env ExpBase Info VName
e)
        | Set VName -> ExpBase Info VName -> Bool
canBeEvaluated Set VName
bound ExpBase Info VName
e =
            Int64 -> Either Int64 SizeClosure
forall a b. a -> Either a b
Left (Int64 -> Either Int64 SizeClosure)
-> (Value -> Int64) -> Value -> Either Int64 SizeClosure
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64 (Value -> Either Int64 SizeClosure)
-> EvalM Value -> EvalM (Either Int64 SizeClosure)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
evalWithExts Env
env ExpBase Info VName
e
      evalDim Set VName
_ p
_ SizeClosure
e = Either Int64 SizeClosure -> EvalM (Either Int64 SizeClosure)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Int64 SizeClosure -> EvalM (Either Int64 SizeClosure))
-> Either Int64 SizeClosure -> EvalM (Either Int64 SizeClosure)
forall a b. (a -> b) -> a -> b
$ SizeClosure -> Either Int64 SizeClosure
forall a b. b -> Either a b
Right SizeClosure
e
  (Set VName
 -> DimPos -> SizeClosure -> EvalM (Either Int64 SizeClosure))
-> TypeBase SizeClosure NoUniqueness
-> EvalM (TypeBase (Either Int64 SizeClosure) NoUniqueness)
forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName
-> DimPos -> SizeClosure -> EvalM (Either Int64 SizeClosure)
forall {p}.
Set VName -> p -> SizeClosure -> EvalM (Either Int64 SizeClosure)
evalDim TypeBase SizeClosure NoUniqueness
t
  where
    canBeEvaluated :: Set VName -> ExpBase Info VName -> Bool
canBeEvaluated Set VName
bound ExpBase Info VName
e =
      let free :: Set VName
free = FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
       in Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) Set VName
free Bool -> Bool -> Bool
|| (VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
outer_bound) Set VName
free

-- | Evaluate all sizes, and it better work. This implies it must be a
-- size-dependent function type, or one that has existentials.
evalTypeFully :: EvalType -> EvalM ValueType
evalTypeFully :: TypeBase SizeClosure NoUniqueness -> EvalM ValueType
evalTypeFully TypeBase SizeClosure NoUniqueness
t = do
  let evalDim :: SizeClosure -> EvalM Int64
evalDim (SizeClosure Env
env ExpBase Info VName
e) = Value -> Int64
asInt64 (Value -> Int64) -> EvalM Value -> EvalM Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
evalWithExts Env
env ExpBase Info VName
e
  (SizeClosure -> EvalM Int64)
-> (NoUniqueness -> EvalM NoUniqueness)
-> TypeBase SizeClosure NoUniqueness
-> EvalM ValueType
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse SizeClosure -> EvalM Int64
evalDim NoUniqueness -> EvalM NoUniqueness
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase SizeClosure NoUniqueness
t

evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value
evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv StructType
t =
  case QualName VName -> Env -> Maybe TermBinding
lookupVar QualName VName
qv Env
env of
    Just (TermPoly Maybe BoundV
_ TypeBase SizeClosure NoUniqueness -> EvalM Value
v) -> TypeBase SizeClosure NoUniqueness -> EvalM Value
v (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> TypeBase SizeClosure NoUniqueness -> EvalM Value
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env StructType
t
    Just (TermValue Maybe BoundV
_ Value
v) -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
    Maybe TermBinding
x -> do
      [Text]
ss <- (Loc -> Text) -> [Loc] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (SrcLoc -> Text
forall a. Located a => a -> Text
locText (SrcLoc -> Text) -> (Loc -> SrcLoc) -> Loc -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Loc -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf) ([Loc] -> [Text]) -> EvalM [Loc] -> EvalM [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM [Loc]
stacktrace
      [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$
        QualName VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString QualName VName
qv
          [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
" is not bound to a value.\n"
          [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> [Char]
T.unpack (Int -> [Text] -> Text
prettyStacktrace Int
0 [Text]
ss)
          [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
"Bound to\n"
          [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Maybe TermBinding -> [Char]
forall a. Show a => a -> [Char]
show Maybe TermBinding
x

typeValueShape :: Env -> StructType -> EvalM ValueShape
typeValueShape :: Env -> StructType -> EvalM ValueShape
typeValueShape Env
env StructType
t = ValueType -> ValueShape
forall d u. TypeBase d u -> Shape d
typeShape (ValueType -> ValueShape) -> EvalM ValueType -> EvalM ValueShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase SizeClosure NoUniqueness -> EvalM ValueType
evalTypeFully (Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env StructType
t)

-- Sometimes type instantiation is not quite enough - then we connect
-- up the missing sizes here.  In particular used for eta-expanded
-- entry points.
linkMissingSizes :: [VName] -> Pat (TypeBase Size u) -> Value -> Env -> Env
linkMissingSizes :: forall u.
[VName]
-> Pat (TypeBase (ExpBase Info VName) u) -> Value -> Env -> Env
linkMissingSizes [] Pat (TypeBase (ExpBase Info VName) u)
_ Value
_ Env
env = Env
env
linkMissingSizes [VName]
missing_sizes Pat (TypeBase (ExpBase Info VName) u)
p Value
v Env
env =
  Env
env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Map VName Int64 -> Env
i64Env ([VName] -> StructType -> ValueShape -> Map VName Int64
resolveExistentials [VName]
missing_sizes StructType
p_t (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v))
  where
    p_t :: StructType
p_t = TypeBase SizeClosure NoUniqueness -> StructType
evalToStruct (TypeBase SizeClosure NoUniqueness -> StructType)
-> TypeBase SizeClosure NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env (StructType -> TypeBase SizeClosure NoUniqueness)
-> StructType -> TypeBase SizeClosure NoUniqueness
forall a b. (a -> b) -> a -> b
$ Pat (TypeBase (ExpBase Info VName) u) -> StructType
forall u. Pat (TypeBase (ExpBase Info VName) u) -> StructType
patternStructType Pat (TypeBase (ExpBase Info VName) u)
p

evalFunction :: Env -> [VName] -> [Pat ParamType] -> Exp -> ResType -> 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 ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Uniqueness
-> EvalM Value
evalFunction Env
env [VName]
missing_sizes [] ExpBase Info VName
body TypeBase (ExpBase Info VName) Uniqueness
rettype =
  -- Eta-expand the rest to make any sizes visible.
  [Value]
-> Env -> TypeBase (ExpBase Info VName) Uniqueness -> EvalM Value
etaExpand [] Env
env TypeBase (ExpBase Info VName) Uniqueness
rettype
  where
    etaExpand :: [Value]
-> Env -> TypeBase (ExpBase Info VName) Uniqueness -> EvalM Value
etaExpand [Value]
vs Env
env' (Scalar (Arrow Uniqueness
_ PName
_ Diet
_ StructType
p_t (RetType [VName]
_ TypeBase (ExpBase Info VName) Uniqueness
rt))) = do
      Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value)
-> ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value)
-> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> EvalM Value)
-> (Value -> EvalM Value) -> EvalM Value
forall a b. (a -> b) -> a -> b
$ \Value
v -> do
        let p :: PatBase Info vn StructType
p = Info StructType -> SrcLoc -> PatBase Info vn StructType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (StructType -> Info StructType
forall a. a -> Info a
Info StructType
p_t) SrcLoc
forall a. IsLocation a => a
noLoc
        Env
env'' <- [VName] -> PatBase Info VName StructType -> Value -> Env -> Env
forall u.
[VName]
-> Pat (TypeBase (ExpBase Info VName) u) -> Value -> Env -> Env
linkMissingSizes [VName]
missing_sizes PatBase Info VName StructType
forall {vn}. PatBase Info vn StructType
p Value
v (Env -> Env) -> EvalM Env -> EvalM Env
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> PatBase Info VName StructType -> Value -> EvalM Env
forall u.
Env -> Pat (TypeBase (ExpBase Info VName) u) -> Value -> EvalM Env
matchPat Env
env' PatBase Info VName StructType
forall {vn}. PatBase Info vn StructType
p Value
v
        [Value]
-> Env -> TypeBase (ExpBase Info VName) Uniqueness -> EvalM Value
etaExpand (Value
v Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
vs) Env
env'' TypeBase (ExpBase Info VName) Uniqueness
rt
    etaExpand [Value]
vs Env
env' TypeBase (ExpBase Info VName) Uniqueness
_ = do
      Value
f <- EvalM Value -> EvalM Value
forall a. EvalM a -> EvalM a
localExts (EvalM Value -> EvalM Value) -> EvalM Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ Env -> ExpBase Info VName -> EvalM Value
eval Env
env' ExpBase Info VName
body
      (Value -> Value -> EvalM Value) -> Value -> [Value] -> EvalM Value
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty) Value
f ([Value] -> EvalM Value) -> [Value] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
vs
evalFunction Env
env [VName]
missing_sizes (Pat ParamType
p : [Pat ParamType]
ps) ExpBase Info VName
body TypeBase (ExpBase Info VName) Uniqueness
rettype =
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value)
-> ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value)
-> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> EvalM Value)
-> (Value -> EvalM Value) -> EvalM Value
forall a b. (a -> b) -> a -> b
$ \Value
v -> do
    Env
env' <- [VName] -> Pat ParamType -> Value -> Env -> Env
forall u.
[VName]
-> Pat (TypeBase (ExpBase Info VName) u) -> Value -> Env -> Env
linkMissingSizes [VName]
missing_sizes Pat ParamType
p Value
v (Env -> Env) -> EvalM Env -> EvalM Env
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> Pat ParamType -> Value -> EvalM Env
forall u.
Env -> Pat (TypeBase (ExpBase Info VName) u) -> Value -> EvalM Env
matchPat Env
env Pat ParamType
p Value
v
    Env
-> [VName]
-> [Pat ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Uniqueness
-> EvalM Value
evalFunction Env
env' [VName]
missing_sizes [Pat ParamType]
ps ExpBase Info VName
body TypeBase (ExpBase Info VName) Uniqueness
rettype

evalFunctionBinding ::
  Env ->
  [TypeParam] ->
  [Pat ParamType] ->
  ResRetType ->
  Exp ->
  EvalM TermBinding
evalFunctionBinding :: Env
-> [TypeParam]
-> [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Uniqueness
-> ExpBase Info VName
-> EvalM TermBinding
evalFunctionBinding Env
env [TypeParam]
tparams [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Uniqueness
ret ExpBase Info VName
fbody = do
  let ftype :: StructType
ftype = [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Uniqueness -> StructType
funType [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Uniqueness
ret
      retext :: [VName]
retext = case [Pat ParamType]
ps of
        [] -> RetTypeBase (ExpBase Info VName) Uniqueness -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims RetTypeBase (ExpBase Info VName) Uniqueness
ret
        [Pat ParamType]
_ -> []

  -- Distinguish polymorphic and non-polymorphic bindings here.
  if [TypeParam] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParam]
tparams
    then
      (Value -> TermBinding) -> EvalM Value -> EvalM TermBinding
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe BoundV -> Value -> TermBinding
TermValue (BoundV -> Maybe BoundV
forall a. a -> Maybe a
Just (BoundV -> Maybe BoundV) -> BoundV -> Maybe BoundV
forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] StructType
ftype))
        (EvalM Value -> EvalM TermBinding)
-> (Value -> EvalM Value) -> Value -> EvalM TermBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env
-> TypeBase (ExpBase Info VName) Uniqueness
-> [VName]
-> Value
-> EvalM Value
forall u.
Env
-> TypeBase (ExpBase Info VName) u
-> [VName]
-> Value
-> EvalM Value
returned Env
env (RetTypeBase (ExpBase Info VName) Uniqueness
-> TypeBase (ExpBase Info VName) Uniqueness
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase (ExpBase Info VName) Uniqueness
ret) [VName]
retext
        (Value -> EvalM TermBinding) -> EvalM Value -> EvalM TermBinding
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env
-> [VName]
-> [Pat ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Uniqueness
-> EvalM Value
evalFunction Env
env [] [Pat ParamType]
ps ExpBase Info VName
fbody (RetTypeBase (ExpBase Info VName) Uniqueness
-> TypeBase (ExpBase Info VName) Uniqueness
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase (ExpBase Info VName) Uniqueness
ret)
    else TermBinding -> EvalM TermBinding
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TermBinding -> EvalM TermBinding)
-> ((TypeBase SizeClosure NoUniqueness -> EvalM Value)
    -> TermBinding)
-> (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> EvalM TermBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe BoundV
-> (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> TermBinding
TermPoly (BoundV -> Maybe BoundV
forall a. a -> Maybe a
Just (BoundV -> Maybe BoundV) -> BoundV -> Maybe BoundV
forall a b. (a -> b) -> a -> b
$ [TypeParam] -> StructType -> BoundV
T.BoundV [] StructType
ftype) ((TypeBase SizeClosure NoUniqueness -> EvalM Value)
 -> EvalM TermBinding)
-> (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> EvalM TermBinding
forall a b. (a -> b) -> a -> b
$ \TypeBase SizeClosure NoUniqueness
ftype' -> do
      let resolved :: ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
 [(VName, SizeClosure)])
resolved = [VName]
-> StructType
-> TypeBase SizeClosure NoUniqueness
-> ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
    [(VName, SizeClosure)])
resolveTypeParams ((TypeParam -> VName) -> [TypeParam] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> VName
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
tparams) StructType
ftype TypeBase SizeClosure NoUniqueness
ftype'
      Env
tparam_env <- ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
 [(VName, SizeClosure)])
-> EvalM Env
evalResolved ([(VName, ([VName], TypeBase SizeClosure NoUniqueness))],
 [(VName, SizeClosure)])
resolved
      let env' :: Env
env' = Env
tparam_env Env -> Env -> Env
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 =
            (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Map VName TermBinding -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.notMember` Env -> Map VName TermBinding
envTerm Env
env') ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$
              (TypeParam -> VName) -> [TypeParam] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> VName
forall vn. TypeParamBase vn -> vn
typeParamName ((TypeParam -> Bool) -> [TypeParam] -> [TypeParam]
forall a. (a -> Bool) -> [a] -> [a]
filter TypeParam -> Bool
forall vn. TypeParamBase vn -> Bool
isSizeParam [TypeParam]
tparams)
      Env
-> TypeBase (ExpBase Info VName) Uniqueness
-> [VName]
-> Value
-> EvalM Value
forall u.
Env
-> TypeBase (ExpBase Info VName) u
-> [VName]
-> Value
-> EvalM Value
returned Env
env (RetTypeBase (ExpBase Info VName) Uniqueness
-> TypeBase (ExpBase Info VName) Uniqueness
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase (ExpBase Info VName) Uniqueness
ret) [VName]
retext
        (Value -> EvalM Value) -> EvalM Value -> EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env
-> [VName]
-> [Pat ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Uniqueness
-> EvalM Value
evalFunction Env
env' [VName]
missing_sizes [Pat ParamType]
ps ExpBase Info VName
fbody (RetTypeBase (ExpBase Info VName) Uniqueness
-> TypeBase (ExpBase Info VName) Uniqueness
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType RetTypeBase (ExpBase Info VName) Uniqueness
ret)

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

returned :: Env -> TypeBase Size u -> [VName] -> Value -> EvalM Value
returned :: forall u.
Env
-> TypeBase (ExpBase Info VName) u
-> [VName]
-> Value
-> EvalM Value
returned Env
_ TypeBase (ExpBase Info VName) u
_ [] Value
v = Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
returned Env
env TypeBase (ExpBase Info VName) u
ret [VName]
retext Value
v = do
  ((VName, Int64) -> EvalM ()) -> [(VName, Int64)] -> EvalM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((VName -> Value -> EvalM ()) -> (VName, Value) -> EvalM ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry VName -> Value -> EvalM ()
putExtSize ((VName, Value) -> EvalM ())
-> ((VName, Int64) -> (VName, Value)) -> (VName, Int64) -> EvalM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int64 -> Value) -> (VName, Int64) -> (VName, Value)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> (Int64 -> PrimValue) -> Int64 -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
SignedValue (IntValue -> PrimValue)
-> (Int64 -> IntValue) -> Int64 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value))
    ([(VName, Int64)] -> EvalM ())
-> (Map VName Int64 -> [(VName, Int64)])
-> Map VName Int64
-> EvalM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Int64 -> [(VName, Int64)]
forall k a. Map k a -> [(k, a)]
M.toList
    (Map VName Int64 -> EvalM ()) -> Map VName Int64 -> EvalM ()
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> ValueShape -> Map VName Int64
resolveExistentials [VName]
retext (TypeBase SizeClosure NoUniqueness -> StructType
evalToStruct (TypeBase SizeClosure NoUniqueness -> StructType)
-> TypeBase SizeClosure NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env (StructType -> TypeBase SizeClosure NoUniqueness)
-> StructType -> TypeBase SizeClosure NoUniqueness
forall a b. (a -> b) -> a -> b
$ TypeBase (ExpBase Info VName) u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase (ExpBase Info VName) u
ret)
    (ValueShape -> Map VName Int64) -> ValueShape -> Map VName Int64
forall a b. (a -> b) -> a -> b
$ Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v

evalAppExp :: Env -> AppExp -> EvalM Value
evalAppExp :: Env -> AppExp -> EvalM Value
evalAppExp Env
env (Range ExpBase Info VName
start Maybe (ExpBase Info VName)
maybe_second Inclusiveness (ExpBase Info VName)
end SrcLoc
loc) = do
  Integer
start' <- Value -> Integer
asInteger (Value -> Integer) -> EvalM Value -> EvalM Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
start
  Maybe Integer
maybe_second' <- (ExpBase Info VName -> EvalM Integer)
-> Maybe (ExpBase Info VName) -> EvalM (Maybe Integer)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse ((Value -> Integer) -> EvalM Value -> EvalM Integer
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Integer
asInteger (EvalM Value -> EvalM Integer)
-> (ExpBase Info VName -> EvalM Value)
-> ExpBase Info VName
-> EvalM Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ExpBase Info VName -> EvalM Value
eval Env
env) Maybe (ExpBase Info VName)
maybe_second
  Inclusiveness Integer
end' <- (ExpBase Info VName -> EvalM Integer)
-> Inclusiveness (ExpBase Info VName)
-> EvalM (Inclusiveness Integer)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Inclusiveness a -> f (Inclusiveness b)
traverse ((Value -> Integer) -> EvalM Value -> EvalM Integer
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> Integer
asInteger (EvalM Value -> EvalM Integer)
-> (ExpBase Info VName -> EvalM Value)
-> ExpBase Info VName
-> EvalM Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ExpBase Info VName -> EvalM Value
eval Env
env) Inclusiveness (ExpBase Info VName)
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'' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1, -Integer
1, Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
end'')
          (DownToExclusive Integer
end'', Just Integer
second') ->
            (Integer
end'' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1, Integer
second' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
start', Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
end'' Bool -> Bool -> Bool
&& Integer
second' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
start')
          (ToInclusive Integer
end'', Maybe Integer
Nothing) ->
            (Integer
end'', Integer
1, Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
end'')
          (ToInclusive Integer
end'', Just Integer
second')
            | Integer
second' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
start' ->
                (Integer
end'', Integer
second' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
start', Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
end'')
            | Bool
otherwise ->
                (Integer
end'', Integer
second' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
start', Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
end'' Bool -> Bool -> Bool
&& Integer
second' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
start')
          (UpToExclusive Integer
x, Maybe Integer
Nothing) ->
            (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1, Integer
1, Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
x)
          (UpToExclusive Integer
x, Just Integer
second') ->
            (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1, Integer
second' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
start', Integer
start' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
x Bool -> Bool -> Bool
&& Integer
second' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
start')

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

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

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

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

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

    forLoop :: VName -> IntValue -> IntValue -> Value -> EvalM Value
forLoop VName
iv IntValue
bound IntValue
i Value
v
      | IntValue
i IntValue -> IntValue -> Bool
forall a. Ord a => a -> a -> Bool
>= IntValue
bound = Value -> EvalM Value
forall a. a -> EvalM a
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) (Value -> EvalM Value) -> EvalM Value -> EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> EvalM Value
evalBody (VName -> IntValue -> Env
forLoopEnv VName
iv IntValue
i Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env')

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

    forInLoop :: Pat (TypeBase (ExpBase Info VName) u)
-> Value -> Value -> EvalM Value
forInLoop Pat (TypeBase (ExpBase Info VName) u)
in_pat Value
v Value
in_v = do
      Env
env' <- Value -> EvalM Env
withLoopParams Value
v
      Env
env'' <- Env -> Pat (TypeBase (ExpBase Info VName) u) -> Value -> EvalM Env
forall u.
Env -> Pat (TypeBase (ExpBase Info VName) u) -> Value -> EvalM Env
matchPat Env
env' Pat (TypeBase (ExpBase Info VName) u)
in_pat Value
in_v
      Env -> EvalM Value
evalBody Env
env''
evalAppExp Env
env (Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs SrcLoc
_) = do
  Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  Value -> [CaseBase Info VName] -> EvalM Value
match Value
v (NonEmpty (CaseBase Info VName) -> [CaseBase Info VName]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (CaseBase Info VName)
cs)
  where
    match :: Value -> [CaseBase Info VName] -> EvalM Value
match Value
_ [] =
      [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error [Char]
"Pattern match failure."
    match Value
v (CaseBase Info VName
c : [CaseBase Info VName]
cs') = do
      Maybe Value
c' <- Value -> Env -> CaseBase Info VName -> EvalM (Maybe Value)
evalCase Value
v Env
env CaseBase Info VName
c
      case Maybe Value
c' of
        Just Value
v' -> Value -> EvalM Value
forall a. a -> EvalM a
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 -> ExpBase Info VName -> EvalM Value
eval Env
_ (Literal PrimValue
v SrcLoc
_) = Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
v
eval Env
env (Hole (Info StructType
t) SrcLoc
loc) =
  SrcLoc -> Env -> Text -> EvalM Value
forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env (Text -> EvalM Value) -> Text -> EvalM Value
forall a b. (a -> b) -> a -> b
$ Text
"Hole of type: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> StructType -> Text
forall a. Pretty a => a -> Text
prettyTextOneLine StructType
t
eval Env
env (Parens ExpBase Info VName
e SrcLoc
_) = Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
eval Env
env (QualParens (QualName VName
qv, SrcLoc
_) ExpBase Info VName
e SrcLoc
loc) = do
  Module
m <- Env -> QualName VName -> EvalM Module
evalModuleVar Env
env QualName VName
qv
  case Module
m of
    ModuleFun {} -> [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"Local open of module function at " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc
    Module Env
m' -> Env -> ExpBase Info VName -> EvalM Value
eval (Env
m' Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env) ExpBase Info VName
e
eval Env
env (TupLit [ExpBase Info VName]
vs SrcLoc
_) = [Value] -> Value
forall (m :: * -> *). [Value m] -> Value m
toTuple ([Value] -> Value) -> EvalM [Value] -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> EvalM Value)
-> [ExpBase Info VName] -> EvalM [Value]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> ExpBase Info VName -> EvalM Value
eval Env
env) [ExpBase Info VName]
vs
eval Env
env (RecordLit [FieldBase Info VName]
fields SrcLoc
_) =
  Map Name Value -> Value
forall (m :: * -> *). Map Name (Value m) -> Value m
ValueRecord (Map Name Value -> Value)
-> ([(Name, Value)] -> Map Name Value) -> [(Name, Value)] -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, Value)] -> Map Name Value
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, Value)] -> Value) -> EvalM [(Name, Value)] -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> EvalM (Name, Value))
-> [FieldBase Info VName] -> EvalM [(Name, Value)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM FieldBase Info VName -> EvalM (Name, Value)
evalField [FieldBase Info VName]
fields
  where
    evalField :: FieldBase Info VName -> EvalM (Name, Value)
evalField (RecordFieldExplicit (L Loc
_ Name
k) ExpBase Info VName
e SrcLoc
_) = do
      Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
      (Name, Value) -> EvalM (Name, Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name
k, Value
v)
    evalField (RecordFieldImplicit (L Loc
_ VName
k) Info StructType
t SrcLoc
loc) = do
      Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env (ExpBase Info VName -> EvalM Value)
-> ExpBase Info VName -> EvalM Value
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info StructType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
k) Info StructType
t SrcLoc
loc
      (Name, Value) -> EvalM (Name, Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Name
baseName VName
k, Value
v)
eval Env
_ (StringLit [Word8]
vs SrcLoc
_) =
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
    ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
forall d. Shape d
ShapeLeaf ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$
      (Word8 -> Value) -> [Word8] -> [Value]
forall a b. (a -> b) -> [a] -> [b]
map (PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> (Word8 -> PrimValue) -> Word8 -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
UnsignedValue (IntValue -> PrimValue)
-> (Word8 -> IntValue) -> Word8 -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> IntValue
Int8Value (Int8 -> IntValue) -> (Word8 -> Int8) -> Word8 -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral) [Word8]
vs
eval Env
env (ArrayLit [] (Info StructType
t) SrcLoc
_) = do
  ValueShape
t' <- Env -> StructType -> EvalM ValueShape
typeValueShape Env
env (StructType -> EvalM ValueShape) -> StructType -> EvalM ValueShape
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
t' []
eval Env
env (ArrayLit (ExpBase Info VName
v : [ExpBase Info VName]
vs) Info StructType
_ SrcLoc
_) = do
  Value
v' <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
v
  [Value]
vs' <- (ExpBase Info VName -> EvalM Value)
-> [ExpBase Info VName] -> EvalM [Value]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> ExpBase Info VName -> EvalM Value
eval Env
env) [ExpBase Info VName]
vs
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v') (Value
v' Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
vs')
eval Env
_ (ArrayVal [PrimValue]
vs PrimType
_ SrcLoc
_) =
  -- Probably will not ever be used.
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
forall d. Shape d
ShapeLeaf ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Value) -> [PrimValue] -> [Value]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim [PrimValue]
vs
eval Env
env (AppExp AppExp
e (Info (AppRes StructType
t [VName]
retext))) = do
  Value
v <- Env -> AppExp -> EvalM Value
evalAppExp Env
env AppExp
e
  Env -> StructType -> [VName] -> Value -> EvalM Value
forall u.
Env
-> TypeBase (ExpBase Info VName) u
-> [VName]
-> Value
-> EvalM Value
returned Env
env (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t) [VName]
retext Value
v
eval Env
env (Var QualName VName
qv (Info StructType
t) SrcLoc
_) = Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t)
eval Env
env (Ascript ExpBase Info VName
e TypeExp (ExpBase Info VName) VName
_ SrcLoc
_) = Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
eval Env
env (Coerce ExpBase Info VName
e TypeExp (ExpBase Info VName) VName
te (Info StructType
t) SrcLoc
loc) = do
  Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  ValueType
t' <- TypeBase SizeClosure NoUniqueness -> EvalM ValueType
evalTypeFully (TypeBase SizeClosure NoUniqueness -> EvalM ValueType)
-> TypeBase SizeClosure NoUniqueness -> EvalM ValueType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoUniqueness
forall u.
Pretty u =>
Env -> TypeBase (ExpBase Info VName) u -> TypeBase SizeClosure u
expandType Env
env StructType
t
  case ValueShape -> ValueShape -> Maybe ValueShape
checkShape (ValueType -> ValueShape
forall d u. TypeBase d u -> Shape d
typeShape ValueType
t') (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v) of
    Just ValueShape
_ -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
    Maybe ValueShape
Nothing ->
      SrcLoc -> Env -> Text -> EvalM Value
forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env (Text -> EvalM Value)
-> (Doc Any -> Text) -> Doc Any -> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc Any -> Text
forall a. Doc a -> Text
docText (Doc Any -> EvalM Value) -> Doc Any -> EvalM Value
forall a b. (a -> b) -> a -> b
$
        Doc Any
"Value `"
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Value -> Doc Any
forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
v
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Doc Any
"` of shape `"
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> ValueShape -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. ValueShape -> Doc ann
pretty (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v)
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Doc Any
"` cannot match shape of type `"
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> TypeExp (ExpBase Info VName) VName -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeExp (ExpBase Info VName) VName -> Doc ann
pretty TypeExp (ExpBase Info VName) VName
te
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Doc Any
"` (`"
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> ValueType -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. ValueType -> Doc ann
pretty ValueType
t'
          Doc Any -> Doc Any -> Doc Any
forall a. Semigroup a => a -> a -> a
<> Doc Any
"`)"
eval Env
_ (IntLit Integer
v (Info StructType
t) SrcLoc
_) =
  case StructType
t of
    Scalar (Prim (Signed IntType
it)) ->
      Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v
    Scalar (Prim (Unsigned IntType
it)) ->
      Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v
    Scalar (Prim (FloatType FloatType
ft)) ->
      Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Integer -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v
    StructType
_ -> [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"eval: nonsensical type for integer literal: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
eval Env
_ (FloatLit Double
v (Info StructType
t) SrcLoc
_) =
  case StructType
t of
    Scalar (Prim (FloatType FloatType
ft)) ->
      Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v
    StructType
_ -> [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"eval: nonsensical type for float literal: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
eval Env
env (Negate ExpBase Info VName
e SrcLoc
loc) = do
  Value
ev <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env Value
intrinsicsNeg Value
ev
eval Env
env (Not ExpBase Info VName
e SrcLoc
loc) =
  SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env Value
intrinsicsNot (Value -> EvalM Value) -> EvalM Value -> EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
eval Env
env (Update ExpBase Info VName
src SliceBase Info VName
is ExpBase Info VName
v SrcLoc
loc) =
  EvalM Value -> (Value -> EvalM Value) -> Maybe Value -> EvalM Value
forall b a. b -> (a -> b) -> Maybe a -> b
maybe EvalM Value
forall {a}. EvalM a
oob Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (Maybe Value -> EvalM Value) -> EvalM (Maybe Value) -> EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> Value -> Maybe Value
writeArray ([Indexing] -> Value -> Value -> Maybe Value)
-> EvalM [Indexing] -> EvalM (Value -> Value -> Maybe Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (DimIndex -> EvalM Indexing)
-> SliceBase Info VName -> EvalM [Indexing]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env) SliceBase Info VName
is EvalM (Value -> Value -> Maybe Value)
-> EvalM Value -> EvalM (Value -> Maybe Value)
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
src EvalM (Value -> Maybe Value) -> EvalM Value -> EvalM (Maybe Value)
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
v
  where
    oob :: EvalM a
oob = SrcLoc -> Env -> Text -> EvalM a
forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env Text
"Bad update"
eval Env
env (RecordUpdate ExpBase Info VName
src [Name]
all_fs ExpBase Info VName
v Info StructType
_ SrcLoc
_) =
  Value -> [Name] -> Value -> Value
forall {m :: * -> *}. Value m -> [Name] -> Value m -> Value m
update (Value -> [Name] -> Value -> Value)
-> EvalM Value -> EvalM ([Name] -> Value -> Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
src EvalM ([Name] -> Value -> Value)
-> EvalM [Name] -> EvalM (Value -> Value)
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Name] -> EvalM [Name]
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
all_fs EvalM (Value -> Value) -> EvalM Value -> EvalM Value
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
v
  where
    update :: Value m -> [Name] -> Value m -> Value m
update Value m
_ [] Value m
v' = Value m
v'
    update (ValueRecord Map Name (Value m)
src') (Name
f : [Name]
fs) Value m
v'
      | Just Value m
f_v <- Name -> Map Name (Value m) -> Maybe (Value m)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (Value m)
src' =
          Map Name (Value m) -> Value m
forall (m :: * -> *). Map Name (Value m) -> Value m
ValueRecord (Map Name (Value m) -> Value m) -> Map Name (Value m) -> Value m
forall a b. (a -> b) -> a -> b
$ Name -> Value m -> Map Name (Value m) -> Map Name (Value m)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
f (Value m -> [Name] -> Value m -> Value m
update Value m
f_v [Name]
fs Value m
v') Map Name (Value m)
src'
    update Value m
_ [Name]
_ Value m
_ = [Char] -> Value m
forall a. HasCallStack => [Char] -> a
error [Char]
"eval RecordUpdate: invalid value."
-- We treat zero-parameter lambdas as simply an expression to
-- evaluate immediately.  Note that this is *not* the same as a lambda
-- that takes an empty tuple '()' as argument!  Zero-parameter lambdas
-- can never occur in a well-formed Futhark program, but they are
-- convenient in the interpreter.
eval Env
env (Lambda [Pat ParamType]
ps ExpBase Info VName
body Maybe (TypeExp (ExpBase Info VName) VName)
_ (Info (RetType [VName]
_ TypeBase (ExpBase Info VName) Uniqueness
rt)) SrcLoc
_) =
  Env
-> [VName]
-> [Pat ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Uniqueness
-> EvalM Value
evalFunction Env
env [] [Pat ParamType]
ps ExpBase Info VName
body TypeBase (ExpBase Info VName) Uniqueness
rt
eval Env
env (OpSection QualName VName
qv (Info StructType
t) SrcLoc
_) =
  Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv (StructType -> EvalM Value) -> StructType -> EvalM Value
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t
eval Env
env (OpSectionLeft QualName VName
qv Info StructType
_ ExpBase Info VName
e (Info (PName
_, ParamType
_, Maybe VName
argext), Info (PName, ParamType)
_) (Info (RetType [VName]
_ TypeBase (ExpBase Info VName) Uniqueness
t), Info [VName]
_) SrcLoc
loc) = do
  Value
v <- Env -> ExpBase Info VName -> Maybe VName -> EvalM Value
evalArg Env
env ExpBase Info VName
e Maybe VName
argext
  Value
f <- Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv (TypeBase (ExpBase Info VName) Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase (ExpBase Info VName) Uniqueness
t)
  SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
loc Env
env Value
f Value
v
eval Env
env (OpSectionRight QualName VName
qv Info StructType
_ ExpBase Info VName
e (Info (PName, ParamType)
_, Info (PName
_, ParamType
_, Maybe VName
argext)) (Info (RetType [VName]
_ TypeBase (ExpBase Info VName) Uniqueness
t)) SrcLoc
loc) = do
  Value
y <- Env -> ExpBase Info VName -> Maybe VName -> EvalM Value
evalArg Env
env ExpBase Info VName
e Maybe VName
argext
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
    (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value) -> Value
forall a b. (a -> b) -> a -> b
$ \Value
x -> do
      Value
f <- Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv (StructType -> EvalM Value) -> StructType -> EvalM Value
forall a b. (a -> b) -> a -> b
$ TypeBase (ExpBase Info VName) Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase (ExpBase Info VName) Uniqueness
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 StructType
_ SrcLoc
loc) = do
  [Indexing]
is' <- (DimIndex -> EvalM Indexing)
-> SliceBase Info VName -> EvalM [Indexing]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> DimIndex -> EvalM Indexing
evalDimIndex Env
env) SliceBase Info VName
is
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value) -> Value
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 StructType
_ SrcLoc
_) =
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value) -> Value
forall a b. (a -> b) -> a -> b
$ (Value -> [Name] -> EvalM Value) -> [Name] -> Value -> EvalM Value
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Value -> Name -> EvalM Value) -> Value -> [Name] -> EvalM Value
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Value -> Name -> EvalM Value
forall {f :: * -> *} {m :: * -> *}.
Applicative f =>
Value m -> Name -> f (Value m)
walk) [Name]
ks
  where
    walk :: Value m -> Name -> f (Value m)
walk (ValueRecord Map Name (Value m)
fs) Name
f
      | Just Value m
v' <- Name -> Map Name (Value m) -> Maybe (Value m)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (Value m)
fs = Value m -> f (Value m)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value m
v'
    walk Value m
_ Name
_ = [Char] -> f (Value m)
forall a. HasCallStack => [Char] -> a
error [Char]
"Value does not have expected field."
eval Env
env (Project Name
f ExpBase Info VName
e Info StructType
_ SrcLoc
_) = do
  Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  case Value
v of
    ValueRecord Map Name Value
fs | Just Value
v' <- Name -> Map Name Value -> Maybe Value
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name Value
fs -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v'
    Value
_ -> [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error [Char]
"Value does not have expected field."
eval Env
env (Assert ExpBase Info VName
what ExpBase Info VName
e (Info Text
s) SrcLoc
loc) = do
  Bool
cond <- Value -> Bool
asBool (Value -> Bool) -> EvalM Value -> EvalM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
what
  Bool -> EvalM () -> EvalM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
cond (EvalM () -> EvalM ()) -> EvalM () -> EvalM ()
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Env -> Text -> EvalM ()
forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
loc Env
env Text
s
  Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
eval Env
env (Constr Name
c [ExpBase Info VName]
es (Info StructType
t) SrcLoc
_) = do
  [Value]
vs <- (ExpBase Info VName -> EvalM Value)
-> [ExpBase Info VName] -> EvalM [Value]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Env -> ExpBase Info VName -> EvalM Value
eval Env
env) [ExpBase Info VName]
es
  ValueShape
shape <- Env -> StructType -> EvalM ValueShape
typeValueShape Env
env (StructType -> EvalM ValueShape) -> StructType -> EvalM ValueShape
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> Name -> [Value] -> Value
forall (m :: * -> *). ValueShape -> Name -> [Value m] -> Value m
ValueSum ValueShape
shape Name
c [Value]
vs
eval Env
env (Attr (AttrAtom (AtomName Name
"break") SrcLoc
_) ExpBase Info VName
e SrcLoc
loc) = do
  Env -> Loc -> EvalM ()
break Env
env (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc)
  Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
eval Env
env (Attr (AttrAtom (AtomName Name
"trace") SrcLoc
_) ExpBase Info VName
e SrcLoc
loc) = do
  Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  Text -> Value -> EvalM ()
trace (Loc -> Text
forall a. Located a => a -> Text
locText (SrcLoc -> Loc
forall a. Located a => a -> Loc
locOf SrcLoc
loc)) Value
v
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
eval Env
env (Attr (AttrComp Name
"trace" [AttrAtom (AtomName Name
tag) SrcLoc
_] SrcLoc
_) ExpBase Info VName
e SrcLoc
_) = do
  Value
v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  Text -> Value -> EvalM ()
trace (Name -> Text
nameToText Name
tag) Value
v
  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
eval Env
env (Attr AttrInfo VName
_ ExpBase Info VName
e SrcLoc
_) =
  Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
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 PatBase Info VName StructType
p ExpBase Info VName
cExp SrcLoc
_) = MaybeT EvalM Value -> EvalM (Maybe Value)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT EvalM Value -> EvalM (Maybe Value))
-> MaybeT EvalM Value -> EvalM (Maybe Value)
forall a b. (a -> b) -> a -> b
$ do
  Env
env' <- Env -> PatBase Info VName StructType -> Value -> MaybeT EvalM Env
forall u.
Env
-> Pat (TypeBase (ExpBase Info VName) u)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env PatBase Info VName StructType
p Value
v
  EvalM Value -> MaybeT EvalM Value
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (EvalM Value -> MaybeT EvalM Value)
-> EvalM Value -> MaybeT EvalM Value
forall a b. (a -> b) -> a -> b
$ Env -> ExpBase Info VName -> EvalM Value
eval Env
env' ExpBase Info VName
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 =
  ([VName] -> [VName] -> [VName])
-> [(VName, [VName])] -> Map VName [VName]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
(<>) ([(VName, [VName])] -> Map VName [VName])
-> (Map VName VName -> [(VName, [VName])])
-> Map VName VName
-> Map VName [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, VName) -> (VName, [VName]))
-> [(VName, VName)] -> [(VName, [VName])]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> [VName]) -> (VName, VName) -> (VName, [VName])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second VName -> [VName]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((VName, VName) -> (VName, [VName]))
-> ((VName, VName) -> (VName, VName))
-> (VName, VName)
-> (VName, [VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> VName -> (VName, VName))
-> (VName, VName) -> (VName, VName)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((VName -> VName -> (VName, VName))
-> VName -> VName -> (VName, VName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,))) ([(VName, VName)] -> [(VName, [VName])])
-> (Map VName VName -> [(VName, VName)])
-> Map VName VName
-> [(VName, [VName])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName VName -> [(VName, VName)]
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 = [VName] -> Maybe [VName] -> [VName]
forall a. a -> Maybe a -> a
fromMaybe [VName
v] (Maybe [VName] -> [VName]) -> Maybe [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ VName -> Map VName [VName] -> Maybe [VName]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName [VName]
rev_substs
    replaceM :: (t -> a) -> Map VName t -> Map VName a
replaceM t -> a
f Map VName t
m = [(VName, a)] -> Map VName a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, a)] -> Map VName a) -> [(VName, a)] -> Map VName a
forall a b. (a -> b) -> a -> b
$ do
      (VName
k, t
v) <- Map VName t -> [(VName, t)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName t
m
      VName
k' <- VName -> [VName]
replace VName
k
      (VName, a) -> [(VName, a)]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
k', t -> a
f t
v)
    onEnv :: Env -> Env
onEnv (Env Map VName TermBinding
terms Map VName TypeBinding
types) =
      Map VName TermBinding -> Map VName TypeBinding -> Env
Env ((TermBinding -> TermBinding)
-> Map VName TermBinding -> Map VName TermBinding
forall {t} {a}. (t -> a) -> Map VName t -> Map VName a
replaceM TermBinding -> TermBinding
onTerm Map VName TermBinding
terms) ((TypeBinding -> TypeBinding)
-> Map VName TypeBinding -> Map VName TypeBinding
forall {t} {a}. (t -> a) -> Map VName t -> Map VName a
replaceM TypeBinding -> TypeBinding
onType Map VName TypeBinding
types)
    onModule :: Module -> Module
onModule (Module Env
env) =
      Env -> Module
Module (Env -> Module) -> Env -> Module
forall a b. (a -> b) -> a -> b
$ Env -> Env
onEnv Env
env
    onModule (ModuleFun Module -> EvalM Module
f) =
      (Module -> EvalM Module) -> Module
ModuleFun ((Module -> EvalM Module) -> Module)
-> (Module -> EvalM Module) -> Module
forall a b. (a -> b) -> a -> b
$ \Module
m -> Module -> Module
onModule (Module -> Module) -> EvalM Module -> EvalM Module
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Module -> EvalM Module
f (Map VName VName -> Module -> Module
substituteInModule Map VName VName
substs Module
m)
    onTerm :: TermBinding -> TermBinding
onTerm (TermValue Maybe BoundV
t Value
v) = Maybe BoundV -> Value -> TermBinding
TermValue Maybe BoundV
t Value
v
    onTerm (TermPoly Maybe BoundV
t TypeBase SizeClosure NoUniqueness -> EvalM Value
v) = Maybe BoundV
-> (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> TermBinding
TermPoly Maybe BoundV
t TypeBase SizeClosure NoUniqueness -> EvalM Value
v
    onTerm (TermModule Module
m) = Module -> TermBinding
TermModule (Module -> TermBinding) -> Module -> TermBinding
forall a b. (a -> b) -> a -> b
$ Module -> Module
onModule Module
m
    onType :: TypeBinding -> TypeBinding
onType (TypeBinding Env
env [TypeParam]
ps StructRetType
t) = Env -> [TypeParam] -> StructRetType -> TypeBinding
TypeBinding (Env -> Env
onEnv Env
env) [TypeParam]
ps StructRetType
t

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) -> Module -> EvalM Module
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Module
m
    Maybe TermBinding
_ -> [Char] -> EvalM Module
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Module) -> [Char] -> EvalM Module
forall a b. (a -> b) -> a -> b
$ QualName VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString QualName VName
qv [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
" is not bound to a module."

-- We also return a new Env here, because we want the definitions
-- inside any constructed modules to also be in scope at the top
-- level. This is because types may contain un-qualified references to
-- definitions in modules, and sometimes those definitions may not
-- actually *have* any qualified name!  See tests/modules/sizes7.fut.
-- This occurs solely because of evalType.
evalModExp :: Env -> ModExp -> EvalM (Env, Module)
evalModExp :: Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
_ (ModImport [Char]
_ (Info ImportName
f) SrcLoc
_) = do
  Maybe Env
f' <- ImportName -> EvalM (Maybe Env)
lookupImport ImportName
f
  Map ImportName Env
known <- (([StackFrame], Map ImportName Env) -> Map ImportName Env)
-> EvalM (Map ImportName Env)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ([StackFrame], Map ImportName Env) -> Map ImportName Env
forall a b. (a, b) -> b
snd
  case Maybe Env
f' of
    Maybe Env
Nothing ->
      [Char] -> EvalM (Env, Module)
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM (Env, Module)) -> [Char] -> EvalM (Env, Module)
forall a b. (a -> b) -> a -> b
$
        [[Char]] -> [Char]
unlines
          [ [Char]
"Unknown interpreter import: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ImportName -> [Char]
forall a. Show a => a -> [Char]
show ImportName
f,
            [Char]
"Known: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [ImportName] -> [Char]
forall a. Show a => a -> [Char]
show (Map ImportName Env -> [ImportName]
forall k a. Map k a -> [k]
M.keys Map ImportName Env
known)
          ]
    Just Env
m -> (Env, Module) -> EvalM (Env, Module)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env
forall a. Monoid a => a
mempty, Env -> Module
Module Env
m)
evalModExp Env
env (ModDecs [DecBase Info VName]
ds SrcLoc
_) = do
  Env Map VName TermBinding
terms Map VName TypeBinding
types <- (Env -> DecBase Info VName -> EvalM Env)
-> Env -> [DecBase Info VName] -> EvalM Env
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.
  let env' :: Env
env' =
        Map VName TermBinding -> Map VName TypeBinding -> Env
Env
          (Map VName TermBinding
terms Map VName TermBinding
-> Map VName TermBinding -> Map VName TermBinding
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 Map VName TypeBinding
-> Map VName TypeBinding -> Map VName TypeBinding
forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Env -> Map VName TypeBinding
envType Env
env)
  (Env, Module) -> EvalM (Env, Module)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env
env', Env -> Module
Module Env
env')
evalModExp Env
env (ModVar QualName VName
qv SrcLoc
_) =
  (Env
forall a. Monoid a => a
mempty,) (Module -> (Env, Module)) -> EvalM Module -> EvalM (Env, Module)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> QualName VName -> EvalM Module
evalModuleVar Env
env QualName VName
qv
evalModExp Env
env (ModAscript ModExp
me ModTypeExpBase Info VName
_ (Info Map VName VName
substs) SrcLoc
_) =
  (Env -> Env)
-> (Module -> Module) -> (Env, Module) -> (Env, Module)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Env -> Env
substituteInEnv (Map VName VName -> Module -> Module
substituteInModule Map VName VName
substs) ((Env, Module) -> (Env, Module))
-> EvalM (Env, Module) -> EvalM (Env, Module)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env ModExp
me
  where
    substituteInEnv :: Env -> Env
substituteInEnv Env
env' =
      let Module Env
env'' = Map VName VName -> Module -> Module
substituteInModule Map VName VName
substs (Env -> Module
Module Env
env') in Env
env''
evalModExp Env
env (ModParens ModExp
me SrcLoc
_) =
  Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env ModExp
me
evalModExp Env
env (ModLambda ModParamBase Info VName
p Maybe (ModTypeExpBase Info VName, Info (Map VName VName))
ret ModExp
e SrcLoc
loc) =
  (Env, Module) -> EvalM (Env, Module)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( Env
forall a. Monoid a => a
mempty,
      (Module -> EvalM Module) -> Module
ModuleFun ((Module -> EvalM Module) -> Module)
-> (Module -> EvalM Module) -> Module
forall a b. (a -> b) -> a -> b
$ \Module
am -> do
        let env' :: Env
env' = Env
env {envTerm = M.insert (modParamName p) (TermModule am) $ envTerm env}
        ((Env, Module) -> Module) -> EvalM (Env, Module) -> EvalM Module
forall a b. (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Env, Module) -> Module
forall a b. (a, b) -> b
snd (EvalM (Env, Module) -> EvalM Module)
-> (ModExp -> EvalM (Env, Module)) -> ModExp -> EvalM Module
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env' (ModExp -> EvalM Module) -> ModExp -> EvalM Module
forall a b. (a -> b) -> a -> b
$ case Maybe (ModTypeExpBase Info VName, Info (Map VName VName))
ret of
          Maybe (ModTypeExpBase Info VName, Info (Map VName VName))
Nothing -> ModExp
e
          Just (ModTypeExpBase Info VName
se, Info (Map VName VName)
rsubsts) -> ModExp
-> ModTypeExpBase Info VName
-> Info (Map VName VName)
-> SrcLoc
-> ModExp
forall (f :: * -> *) vn.
ModExpBase f vn
-> ModTypeExpBase f vn
-> f (Map VName VName)
-> SrcLoc
-> ModExpBase f vn
ModAscript ModExp
e ModTypeExpBase 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
  (Env
f_env, Module
f') <- Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env ModExp
f
  (Env
e_env, Module
e') <- Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env ModExp
e
  case Module
f' of
    ModuleFun Module -> EvalM Module
f'' -> do
      Module
res_mod <- Map VName VName -> Module -> Module
substituteInModule Map VName VName
rsubst (Module -> Module) -> EvalM Module -> EvalM Module
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')
      let res_env :: Env
res_env = case Module
res_mod of
            Module Env
x -> Env
x
            Module
_ -> Env
forall a. Monoid a => a
mempty
      (Env, Module) -> EvalM (Env, Module)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env
f_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
e_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
res_env, Module
res_mod)
    Module
_ -> [Char] -> EvalM (Env, Module)
forall a. HasCallStack => [Char] -> a
error [Char]
"Expected ModuleFun."

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

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

nanValue :: PrimValue -> Bool
nanValue :: PrimValue -> Bool
nanValue (FloatValue FloatValue
v) =
  case FloatValue
v of
    Float16Value Half
x -> Half -> Bool
forall a. RealFloat a => a -> Bool
isNaN Half
x
    Float32Value Float
x -> Float -> Bool
forall a. RealFloat a => a -> Bool
isNaN Float
x
    Float64Value Double
x -> Double -> Bool
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 ((PrimValue -> Bool) -> [PrimValue] -> Bool
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 <- (([StackFrame], Map ImportName Env) -> [StackFrame])
-> EvalM [StackFrame]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ([StackFrame], Map ImportName Env) -> [StackFrame]
forall a b. (a, b) -> a
fst
      case [StackFrame] -> Maybe (NonEmpty StackFrame)
forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty [StackFrame]
backtrace of
        Maybe (NonEmpty StackFrame)
Nothing -> () -> EvalM ()
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just NonEmpty StackFrame
backtrace' ->
          let loc :: Loc
loc = StackFrame -> Loc
stackFrameLoc (StackFrame -> Loc) -> StackFrame -> Loc
forall a b. (a -> b) -> a -> b
$ NonEmpty StackFrame -> StackFrame
forall a. NonEmpty a -> a
NE.head NonEmpty StackFrame
backtrace'
           in ExtOp () -> EvalM ()
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (ExtOp () -> EvalM ()) -> ExtOp () -> EvalM ()
forall a b. (a -> b) -> a -> b
$ Loc -> BreakReason -> NonEmpty StackFrame -> () -> ExtOp ()
forall a. Loc -> BreakReason -> NonEmpty StackFrame -> a -> ExtOp a
ExtOpBreak Loc
loc BreakReason
BreakNaN NonEmpty StackFrame
backtrace' ()
breakOnNaN [PrimValue]
_ PrimValue
_ =
  () -> EvalM ()
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

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

    sintOp :: (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp IntType -> BinOp
f =
      [ (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int8), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (IntType -> BinOp
f IntType
Int64))
      ]
    uintOp :: (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp IntType -> BinOp
f =
      [ (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (IntType -> BinOp
f IntType
Int8), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (IntType -> BinOp
f IntType
Int64))
      ]
    intOp :: (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
intOp IntType -> BinOp
f = (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp IntType -> BinOp
f [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp IntType -> BinOp
f
    floatOp :: (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
floatOp FloatType -> BinOp
f =
      [ (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp (FloatType -> BinOp
f FloatType
Float16), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ BinOp -> Op
AD.OpBin (FloatType -> BinOp
f FloatType
Float64))
      ]
    arithOp :: (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp IntType -> BinOp
f FloatType -> BinOp
g = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
intOp IntType -> BinOp
f [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
floatOp FloatType -> BinOp
g

    flipCmps :: [(a, b, a -> b -> c, a -> b -> c)]
-> [(a, b, b -> a -> c, b -> a -> c)]
flipCmps = ((a, b, a -> b -> c, a -> b -> c)
 -> (a, b, b -> a -> c, b -> a -> c))
-> [(a, b, a -> b -> c, a -> b -> c)]
-> [(a, b, b -> a -> c, b -> a -> c)]
forall a b. (a -> b) -> [a] -> [b]
map (\(a
f, b
g, a -> b -> c
h, a -> b -> c
o) -> (a
f, b
g, (a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> c
h, (a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> c
o))
    sintCmp :: (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
sintCmp IntType -> CmpOp
f =
      [ (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int8)),
        (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int16)),
        (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int32)),
        (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int64))
      ]
    uintCmp :: (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
uintCmp IntType -> CmpOp
f =
      [ (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int8)),
        (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int16)),
        (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int32)),
        (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (IntType -> CmpOp
f IntType
Int64))
      ]
    floatCmp :: (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
floatCmp FloatType -> CmpOp
f =
      [ (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (FloatType -> CmpOp
f FloatType
Float16)),
        (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (FloatType -> CmpOp
f FloatType
Float32)),
        (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
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), Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp (FloatType -> CmpOp
f FloatType
Float64))
      ]
    boolCmp :: CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
boolCmp CmpOp
f = [(PrimValue -> Maybe PrimValue
getB, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (Bool -> PrimValue) -> Bool -> Maybe PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue, CmpOp -> PrimValue -> PrimValue -> Maybe Bool
P.doCmpOp CmpOp
f, Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp (Op -> ADValue -> ADValue -> Maybe ADValue)
-> Op -> ADValue -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp CmpOp
f)]

    getV :: PrimValue -> Maybe PrimValue
getV (SignedValue IntValue
x) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getV (UnsignedValue IntValue
x) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getV (FloatValue FloatValue
x) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
P.FloatValue FloatValue
x
    getV (BoolValue Bool
x) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
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) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
P.IntValue IntValue
x
    getS PrimValue
_ = Maybe PrimValue
forall a. Maybe a
Nothing
    putS :: PrimValue -> Maybe PrimValue
putS (P.IntValue IntValue
x) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue IntValue
x
    putS PrimValue
_ = Maybe PrimValue
forall a. Maybe a
Nothing

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

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

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

    getAD :: Value m -> Maybe ADValue
getAD (ValuePrim PrimValue
v) = PrimValue -> ADValue
AD.Constant (PrimValue -> ADValue) -> Maybe PrimValue -> Maybe ADValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimValue -> Maybe PrimValue
getV PrimValue
v
    getAD (ValueAD Int
d ADVariable
v) = ADValue -> Maybe ADValue
forall a. a -> Maybe a
Just (ADValue -> Maybe ADValue) -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ Int -> ADVariable -> ADValue
AD.Variable Int
d ADVariable
v
    getAD Value m
_ = Maybe ADValue
forall a. Maybe a
Nothing
    putAD :: ADValue -> Value m
putAD (AD.Variable Int
d ADVariable
s) = Int -> ADVariable -> Value m
forall (m :: * -> *). Int -> ADVariable -> Value m
ValueAD Int
d ADVariable
s
    putAD (AD.Constant PrimValue
v) = PrimValue -> Value m
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value m) -> PrimValue -> Value m
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
putV PrimValue
v

    adToPrim :: ADValue -> PrimValue
adToPrim ADValue
v = PrimValue -> PrimValue
putV (PrimValue -> PrimValue) -> PrimValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ ADValue -> PrimValue
AD.primitive ADValue
v

    adBinOp :: Op -> ADValue -> ADValue -> Maybe ADValue
adBinOp Op
op ADValue
x ADValue
y = Op -> [ADValue] -> Maybe ADValue
AD.doOp Op
op [ADValue
x, ADValue
y]
    adUnOp :: Op -> ADValue -> Maybe ADValue
adUnOp Op
op ADValue
x = Op -> [ADValue] -> Maybe ADValue
AD.doOp Op
op [ADValue
x]

    fun1 :: (Value -> EvalM Value) -> TermBinding
fun1 Value -> EvalM Value
f =
      Maybe BoundV -> Value -> TermBinding
TermValue Maybe BoundV
forall a. Maybe a
Nothing (Value -> TermBinding) -> Value -> TermBinding
forall a b. (a -> b) -> a -> b
$ (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value) -> Value
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 Maybe BoundV
forall a. Maybe a
Nothing (Value -> TermBinding)
-> ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value)
-> TermBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x ->
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value)
-> ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value)
-> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> EvalM Value)
-> (Value -> EvalM Value) -> EvalM Value
forall a b. (a -> b) -> a -> b
$ \Value
y -> Value -> Value -> EvalM Value
f Value
x Value
y

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

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

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

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

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

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

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

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

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

    def :: T.Text -> Maybe TermBinding
    def :: Text -> Maybe TermBinding
def Text
"!" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> Maybe PrimValue, ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Maybe ADValue)]
-> TermBinding
unopDef
          [ (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int8, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int16, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int32, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int64, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int8, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int16, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int32, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> UnOp
P.Complement IntType
Int64, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
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 -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg PrimType
P.Bool, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg PrimType
P.Bool)
          ]
    def Text
"neg" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> Maybe PrimValue, ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Maybe ADValue)]
-> TermBinding
unopDef
          [ (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int8, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int8),
            (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int16, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int16),
            (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int32, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int32),
            (PrimValue -> Maybe PrimValue
getS, PrimValue -> Maybe PrimValue
putS, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int64, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int64),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int8, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int8),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int16, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int16),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int32, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int32),
            (PrimValue -> Maybe PrimValue
getU, PrimValue -> Maybe PrimValue
putU, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int64, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
P.IntType IntType
Int64),
            (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
P.FloatType FloatType
Float16, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
P.FloatType FloatType
Float16),
            (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
P.FloatType FloatType
Float32, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
P.FloatType FloatType
Float32),
            (PrimValue -> Maybe PrimValue
getF, PrimValue -> Maybe PrimValue
putF, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
P.FloatType FloatType
Float64, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg (PrimType -> UnOp) -> PrimType -> UnOp
forall a b. (a -> b) -> a -> b
$ FloatType -> PrimType
P.FloatType FloatType
Float64),
            (PrimValue -> Maybe PrimValue
getB, PrimValue -> Maybe PrimValue
putB, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp (UnOp -> PrimValue -> Maybe PrimValue)
-> UnOp -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg PrimType
P.Bool, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn (UnOp -> Op) -> UnOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> UnOp
P.Neg PrimType
P.Bool)
          ]
    def Text
"+" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Add` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FAdd
    def Text
"-" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Sub` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FSub
    def Text
"*" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp (IntType -> Overflow -> BinOp
`P.Mul` Overflow
P.OverflowWrap) FloatType -> BinOp
P.FMul
    def Text
"**" = (IntType -> BinOp) -> (FloatType -> BinOp) -> Maybe TermBinding
arithOp IntType -> BinOp
P.Pow FloatType -> BinOp
P.FPow
    def Text
"/" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp (IntType -> Safety -> BinOp
`P.SDiv` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp (IntType -> Safety -> BinOp
`P.UDiv` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
floatOp FloatType -> BinOp
P.FDiv
    def Text
"%" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp (IntType -> Safety -> BinOp
`P.SMod` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp (IntType -> Safety -> BinOp
`P.UMod` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
floatOp FloatType -> BinOp
P.FMod
    def Text
"//" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp (IntType -> Safety -> BinOp
`P.SQuot` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp (IntType -> Safety -> BinOp
`P.UDiv` Safety
P.Unsafe)
    def Text
"%%" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp (IntType -> Safety -> BinOp
`P.SRem` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp (IntType -> Safety -> BinOp
`P.UMod` Safety
P.Unsafe)
    def Text
"^" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
intOp IntType -> BinOp
P.Xor
    def Text
"&" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
intOp IntType -> BinOp
P.And
    def Text
"|" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
intOp IntType -> BinOp
P.Or
    def Text
">>" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp IntType -> BinOp
P.AShr [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp IntType -> BinOp
P.LShr
    def Text
"<<" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
intOp IntType -> BinOp
P.Shl
    def Text
">>>" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
sintOp IntType -> BinOp
P.LShr [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Maybe ADValue)]
uintOp IntType -> BinOp
P.LShr
    def Text
"==" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 ((Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$
        \Value
xs Value
ys -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> Bool -> PrimValue
forall a b. (a -> b) -> a -> b
$ Value
xs Value -> Value -> Bool
forall a. Eq a => a -> a -> Bool
== Value
ys
    def Text
"!=" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 ((Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$
        \Value
xs Value
ys -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> Bool -> PrimValue
forall a b. (a -> b) -> a -> b
$ Value
xs Value -> Value -> Bool
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 Text
"&&" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 ((Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x Value
y ->
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> Bool -> PrimValue
forall a b. (a -> b) -> a -> b
$ Value -> Bool
asBool Value
x Bool -> Bool -> Bool
&& Value -> Bool
asBool Value
y
    def Text
"||" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 ((Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x Value
y ->
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> Bool -> PrimValue
forall a b. (a -> b) -> a -> b
$ Value -> Bool
asBool Value
x Bool -> Bool -> Bool
|| Value -> Bool
asBool Value
y
    def Text
"<" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
sintCmp IntType -> CmpOp
P.CmpSlt
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
uintCmp IntType -> CmpOp
P.CmpUlt
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
floatCmp FloatType -> CmpOp
P.FCmpLt
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
boolCmp CmpOp
P.CmpLlt
    def Text
">" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall {a} {b} {a} {b} {c} {a} {b} {c}.
[(a, b, a -> b -> c, a -> b -> c)]
-> [(a, b, b -> a -> c, b -> a -> c)]
flipCmps ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Maybe ADValue)]
 -> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
      PrimValue -> PrimValue -> Maybe Bool,
      ADValue -> ADValue -> Maybe ADValue)])
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a b. (a -> b) -> a -> b
$
            (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
sintCmp IntType -> CmpOp
P.CmpSlt
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
uintCmp IntType -> CmpOp
P.CmpUlt
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
floatCmp FloatType -> CmpOp
P.FCmpLt
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
boolCmp CmpOp
P.CmpLlt
    def Text
"<=" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
sintCmp IntType -> CmpOp
P.CmpSle
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
uintCmp IntType -> CmpOp
P.CmpUle
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
floatCmp FloatType -> CmpOp
P.FCmpLe
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
boolCmp CmpOp
P.CmpLle
    def Text
">=" =
      TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
        [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Maybe ADValue)]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall {a} {b} {a} {b} {c} {a} {b} {c}.
[(a, b, a -> b -> c, a -> b -> c)]
-> [(a, b, b -> a -> c, b -> a -> c)]
flipCmps ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Maybe ADValue)]
 -> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
      PrimValue -> PrimValue -> Maybe Bool,
      ADValue -> ADValue -> Maybe ADValue)])
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a b. (a -> b) -> a -> b
$
            (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
sintCmp IntType -> CmpOp
P.CmpSle
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
uintCmp IntType -> CmpOp
P.CmpUle
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
floatCmp FloatType -> CmpOp
P.FCmpLe
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Maybe ADValue)]
boolCmp CmpOp
P.CmpLle
    def Text
s
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Text
s ==) (Text -> Bool) -> (BinOp -> Text) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> Text
forall a. Pretty a => a -> Text
prettyText) [BinOp]
P.allBinOps =
          TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ Op -> (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
tbopDef (BinOp -> Op
AD.OpBin BinOp
bop) ((PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding)
-> (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
forall a b. (a -> b) -> a -> b
$ BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
P.doBinOp BinOp
bop
      | Just CmpOp
unop <- (CmpOp -> Bool) -> [CmpOp] -> Maybe CmpOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Text
s ==) (Text -> Bool) -> (CmpOp -> Text) -> CmpOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmpOp -> Text
forall a. Pretty a => a -> Text
prettyText) [CmpOp]
P.allCmpOps =
          TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ Op -> (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
tbopDef (CmpOp -> Op
AD.OpCmp CmpOp
unop) ((PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding)
-> (PrimValue -> PrimValue -> Maybe PrimValue) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \PrimValue
x PrimValue
y -> Bool -> PrimValue
P.BoolValue (Bool -> PrimValue) -> Maybe Bool -> Maybe PrimValue
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 <- (ConvOp -> Bool) -> [ConvOp] -> Maybe ConvOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Text
s ==) (Text -> Bool) -> (ConvOp -> Text) -> ConvOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConvOp -> Text
forall a. Pretty a => a -> Text
prettyText) [ConvOp]
P.allConvOps =
          TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> Maybe PrimValue, ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Maybe ADValue)]
-> TermBinding
unopDef [(PrimValue -> Maybe PrimValue
getV, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (PrimValue -> PrimValue) -> PrimValue -> Maybe PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV, ConvOp -> PrimValue -> Maybe PrimValue
P.doConvOp ConvOp
cop, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ ConvOp -> Op
AD.OpConv ConvOp
cop)]
      | Just UnOp
unop <- (UnOp -> Bool) -> [UnOp] -> Maybe UnOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Text
s ==) (Text -> Bool) -> (UnOp -> Text) -> UnOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnOp -> Text
forall a. Pretty a => a -> Text
prettyText) [UnOp]
P.allUnOps =
          TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> Maybe PrimValue, ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Maybe ADValue)]
-> TermBinding
unopDef [(PrimValue -> Maybe PrimValue
getV, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (PrimValue -> PrimValue) -> PrimValue -> Maybe PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV, UnOp -> PrimValue -> Maybe PrimValue
P.doUnOp UnOp
unop, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ UnOp -> Op
AD.OpUn UnOp
unop)]
      | Just ([PrimType]
pts, PrimType
_, [PrimValue] -> Maybe PrimValue
f) <- Text
-> Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> Maybe ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
s Map Text ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
P.primFuns =
          case [PrimType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimType]
pts of
            Int
1 -> TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> Maybe PrimValue, ADValue -> Maybe ADValue)]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Maybe ADValue)]
-> TermBinding
unopDef [(PrimValue -> Maybe PrimValue
getV, PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue)
-> (PrimValue -> PrimValue) -> PrimValue -> Maybe PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV, [PrimValue] -> Maybe PrimValue
f ([PrimValue] -> Maybe PrimValue)
-> (PrimValue -> [PrimValue]) -> PrimValue -> Maybe PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> [PrimValue]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure, Op -> ADValue -> Maybe ADValue
adUnOp (Op -> ADValue -> Maybe ADValue) -> Op -> ADValue -> Maybe ADValue
forall a b. (a -> b) -> a -> b
$ Text -> Op
AD.OpFn Text
s)]
            Int
_ -> TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
              (Value -> EvalM Value) -> TermBinding
fun1 ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x -> do
                let getV' :: Value m -> Maybe PrimValue
getV' (ValuePrim PrimValue
v) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just PrimValue
v
                    getV' Value m
_ = Maybe PrimValue
forall a. Maybe a
Nothing
                case (Value -> Maybe PrimValue) -> [Value] -> Maybe [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Value -> Maybe PrimValue
forall {m :: * -> *}. Value m -> Maybe PrimValue
getV' ([Value] -> Maybe [PrimValue])
-> Maybe [Value] -> Maybe [PrimValue]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value -> Maybe [Value]
forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple Value
x of
                  Just [PrimValue]
vs
                    | Just PrimValue
res <- (PrimValue -> PrimValue) -> Maybe PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PrimValue -> PrimValue
putV (Maybe PrimValue -> Maybe PrimValue)
-> ([PrimValue] -> Maybe PrimValue)
-> [PrimValue]
-> Maybe PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [PrimValue] -> Maybe PrimValue
f ([PrimValue] -> Maybe PrimValue)
-> Maybe [PrimValue] -> Maybe PrimValue
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (PrimValue -> Maybe PrimValue) -> [PrimValue] -> Maybe [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM PrimValue -> Maybe PrimValue
getV [PrimValue]
vs -> do
                        [PrimValue] -> PrimValue -> EvalM ()
breakOnNaN [PrimValue]
vs PrimValue
res
                        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim PrimValue
res
                  Maybe [PrimValue]
_ ->
                    [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot apply " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Text
s [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" to " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
x
      | Text
"sign_" Text -> Text -> Bool
`T.isPrefixOf` Text
s =
          TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
            (Value -> EvalM Value) -> TermBinding
fun1 ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x ->
              case Value
x of
                (ValuePrim (UnsignedValue IntValue
x')) ->
                  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
SignedValue IntValue
x'
                Value
_ -> [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot sign: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
x
      | Text
"unsign_" Text -> Text -> Bool
`T.isPrefixOf` Text
s =
          TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
            (Value -> EvalM Value) -> TermBinding
fun1 ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x ->
              case Value
x of
                (ValuePrim (SignedValue IntValue
x')) ->
                  Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
UnsignedValue IntValue
x'
                Value
_ -> [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot unsign: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
x
    def Text
"map" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      Maybe BoundV
-> (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> TermBinding
TermPoly Maybe BoundV
forall a. Maybe a
Nothing ((TypeBase SizeClosure NoUniqueness -> EvalM Value) -> TermBinding)
-> (TypeBase SizeClosure NoUniqueness -> EvalM Value)
-> TermBinding
forall a b. (a -> b) -> a -> b
$ \TypeBase SizeClosure NoUniqueness
t -> do
        ValueType
t' <- TypeBase SizeClosure NoUniqueness -> EvalM ValueType
evalTypeFully TypeBase SizeClosure NoUniqueness
t
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value) -> Value
forall a b. (a -> b) -> a -> b
$ \Value
f -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value)
-> ((Value -> EvalM Value) -> Value)
-> (Value -> EvalM Value)
-> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> EvalM Value) -> Value
forall (m :: * -> *). (Value m -> m (Value m)) -> Value m
ValueFun ((Value -> EvalM Value) -> EvalM Value)
-> (Value -> EvalM Value) -> EvalM Value
forall a b. (a -> b) -> a -> b
$ \Value
xs ->
          case ValueType -> ([TypeBase Int64 Diet], ValueType)
forall dim as.
TypeBase dim as -> ([TypeBase dim Diet], TypeBase dim NoUniqueness)
unfoldFunType ValueType
t' of
            ([TypeBase Int64 Diet
_, TypeBase Int64 Diet
_], ValueType
ret_t)
              | ValueShape
rowshape <- ValueType -> ValueShape
forall d u. TypeBase d u -> Shape d
typeShape (ValueType -> ValueShape) -> ValueType -> ValueShape
forall a b. (a -> b) -> a -> b
$ Int -> ValueType -> ValueType
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 ValueType
ret_t ->
                  ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape ([Value] -> Value) -> EvalM [Value] -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> EvalM Value) -> [Value] -> EvalM [Value]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs)
            ([TypeBase Int64 Diet], ValueType)
_ ->
              [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$
                [Char]
"Invalid arguments to map intrinsic:\n"
                  [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unlines [TypeBase SizeClosure NoUniqueness -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase SizeClosure NoUniqueness
t, Value -> [Char]
forall a. Show a => a -> [Char]
show Value
f, Value -> [Char]
forall a. Show a => a -> [Char]
show Value
xs]
    def Text
s | Text
"reduce" Text -> Text -> Bool
`T.isPrefixOf` Text
s = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
f Value
ne Value
xs ->
        (Value -> Value -> EvalM Value) -> Value -> [Value] -> EvalM Value
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 SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f) Value
ne ([Value] -> EvalM Value) -> [Value] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs
    def Text
"scan" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
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 SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f Value
acc Value
x
              ([Value], Value) -> EvalM ([Value], Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value
x' Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
out, Value
x')
        ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
ne) ([Value] -> Value)
-> (([Value], Value) -> [Value]) -> ([Value], Value) -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Value] -> [Value]
forall a. [a] -> [a]
reverse ([Value] -> [Value])
-> (([Value], Value) -> [Value]) -> ([Value], Value) -> [Value]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Value], Value) -> [Value]
forall a b. (a, b) -> a
fst
          (([Value], Value) -> Value)
-> EvalM ([Value], Value) -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([Value], Value) -> Value -> EvalM ([Value], Value))
-> ([Value], Value) -> [Value] -> EvalM ([Value], Value)
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) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs)
    def Text
"scatter" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
arr Value
is Value
vs ->
        case Value
arr of
          ValueArray ValueShape
shape Array Int Value
arr' ->
            Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
              ValueShape -> Array Int Value -> Value
forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
ValueArray ValueShape
shape (Array Int Value -> Value) -> Array Int Value -> Value
forall a b. (a -> b) -> a -> b
$
                (Array Int Value -> (Int, Value) -> Array Int Value)
-> Array Int Value -> [(Int, Value)] -> Array Int Value
forall b a. (b -> a -> b) -> 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
forall {m :: * -> *}.
Array Int (Value m) -> (Int, Value m) -> Array Int (Value m)
update Array Int Value
arr' ([(Int, Value)] -> Array Int Value)
-> [(Int, Value)] -> Array Int Value
forall a b. (a -> b) -> a -> b
$
                  [Int] -> [Value] -> [(Int, Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Value -> Int) -> [Value] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Int
asInt ([Value] -> [Int]) -> [Value] -> [Int]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs)
          Value
_ ->
            [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"scatter expects array, but got: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
arr
      where
        update :: Array Int (Value m) -> (Int, Value m) -> Array Int (Value m)
update Array Int (Value m)
arr' (Int
i, Value m
v) =
          if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array Int (Value m) -> Int
forall int (m :: * -> *).
Integral int =>
Array Int (Value m) -> int
arrayLength Array Int (Value m)
arr'
            then Array Int (Value m)
arr' Array Int (Value m) -> [(Int, Value m)] -> Array Int (Value m)
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Int
i, Value m
v)]
            else Array Int (Value m)
arr'
    def Text
"scatter_2d" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
arr Value
is Value
vs ->
        case Value
arr of
          ValueArray ValueShape
_ Array Int Value
_ ->
            Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
              (Value -> (Maybe [Value], Value) -> Value)
-> Value -> [(Maybe [Value], Value)] -> Value
forall b a. (b -> a -> b) -> 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 ([(Maybe [Value], Value)] -> Value)
-> [(Maybe [Value], Value)] -> Value
forall a b. (a -> b) -> a -> b
$
                [Maybe [Value]] -> [Value] -> [(Maybe [Value], Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Value -> Maybe [Value]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Maybe [Value]
forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple ([Value] -> [Maybe [Value]]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs)
          Value
_ ->
            [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"scatter_2d expects array, but got: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
arr
      where
        update :: Value -> (Maybe [Value], Value) -> Value
        update :: Value -> (Maybe [Value], Value) -> Value
update Value
arr (Just idxs :: [Value]
idxs@[Value
_, Value
_], Value
v) =
          Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
arr (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Value -> Maybe Value
writeArray ((Value -> Indexing) -> [Value] -> [Indexing]
forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix (Int64 -> Indexing) -> (Value -> Int64) -> Value -> Indexing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64) [Value]
idxs) Value
arr Value
v
        update Value
_ (Maybe [Value], Value)
_ =
          [Char] -> Value
forall a. HasCallStack => [Char] -> a
error [Char]
"scatter_2d expects 2-dimensional indices"
    def Text
"scatter_3d" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
arr Value
is Value
vs ->
        case Value
arr of
          ValueArray ValueShape
_ Array Int Value
_ ->
            Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
              (Value -> (Maybe [Value], Value) -> Value)
-> Value -> [(Maybe [Value], Value)] -> Value
forall b a. (b -> a -> b) -> 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 ([(Maybe [Value], Value)] -> Value)
-> [(Maybe [Value], Value)] -> Value
forall a b. (a -> b) -> a -> b
$
                [Maybe [Value]] -> [Value] -> [(Maybe [Value], Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Value -> Maybe [Value]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Maybe [Value]
forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple ([Value] -> [Maybe [Value]]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
vs)
          Value
_ ->
            [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM Value) -> [Char] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Char]
"scatter_3d expects array, but got: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
arr
      where
        update :: Value -> (Maybe [Value], Value) -> Value
        update :: Value -> (Maybe [Value], Value) -> Value
update Value
arr (Just idxs :: [Value]
idxs@[Value
_, Value
_, Value
_], Value
v) =
          Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
arr (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Value -> Maybe Value
writeArray ((Value -> Indexing) -> [Value] -> [Indexing]
forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix (Int64 -> Indexing) -> (Value -> Int64) -> Value -> Indexing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Int64
asInt64) [Value]
idxs) Value
arr Value
v
        update Value
_ (Maybe [Value], Value)
_ =
          [Char] -> Value
forall a. HasCallStack => [Char] -> a
error [Char]
"scatter_3d expects 3-dimensional indices"
    def Text
"hist_1d" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> ((Value
     -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
    -> TermBinding)
-> (Value
    -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> Maybe TermBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 ((Value
  -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
 -> Maybe TermBinding)
-> (Value
    -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
_ Value
arr Value
fun Value
_ Value
is Value
vs ->
      (Value -> (Int64, Value) -> EvalM Value)
-> Value -> [(Int64, Value)] -> EvalM Value
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
        ([Int64] -> [Value] -> [(Int64, Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Value -> Int64) -> [Value] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Int64
asInt64 ([Value] -> [Int64]) -> [Value] -> [Int64]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
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 SrcLoc
forall a. Monoid a => a
mempty Env
forall a. Monoid a => a
mempty
        update :: Value -> Value -> (Int64, Value) -> EvalM Value
update Value
fun Value
arr (Int64
i, Value
v) =
          Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
arr (Maybe Value -> Value) -> EvalM (Maybe Value) -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> Value -> EvalM Value)
-> [Indexing] -> Value -> Value -> EvalM (Maybe Value)
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 Text
"hist_2d" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> ((Value
     -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
    -> TermBinding)
-> (Value
    -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> Maybe TermBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 ((Value
  -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
 -> Maybe TermBinding)
-> (Value
    -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
_ Value
arr Value
fun Value
_ Value
is Value
vs ->
      (Value -> (Maybe [Value], Value) -> EvalM Value)
-> Value -> [(Maybe [Value], Value)] -> EvalM Value
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
        ([Maybe [Value]] -> [Value] -> [(Maybe [Value], Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Value -> Maybe [Value]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Maybe [Value]
forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple ([Value] -> [Maybe [Value]]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
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 SrcLoc
forall a. Monoid a => a
mempty Env
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) =
          Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
arr
            (Maybe Value -> Value) -> EvalM (Maybe Value) -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> Value -> EvalM Value)
-> [Indexing] -> Value -> Value -> EvalM (Maybe Value)
forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (Value -> Value -> Value -> EvalM Value
op Value
fun) ((Value -> Indexing) -> [Value] -> [Indexing]
forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix (Int64 -> Indexing) -> (Value -> Int64) -> Value -> Indexing
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)
_ =
          [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error [Char]
"hist_2d: bad index value"
    def Text
"hist_3d" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> ((Value
     -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
    -> TermBinding)
-> (Value
    -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> Maybe TermBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> TermBinding
fun6 ((Value
  -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
 -> Maybe TermBinding)
-> (Value
    -> Value -> Value -> Value -> Value -> Value -> EvalM Value)
-> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
_ Value
arr Value
fun Value
_ Value
is Value
vs ->
      (Value -> (Maybe [Value], Value) -> EvalM Value)
-> Value -> [(Maybe [Value], Value)] -> EvalM Value
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
        ([Maybe [Value]] -> [Value] -> [(Maybe [Value], Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Value -> Maybe [Value]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Maybe [Value]
forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple ([Value] -> [Maybe [Value]]) -> [Value] -> [Maybe [Value]]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
is) ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
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 SrcLoc
forall a. Monoid a => a
mempty Env
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) =
          Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
arr
            (Maybe Value -> Value) -> EvalM (Maybe Value) -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> Value -> EvalM Value)
-> [Indexing] -> Value -> Value -> EvalM (Maybe Value)
forall (m :: * -> *).
Monad m =>
(Value -> Value -> m Value)
-> [Indexing] -> Value -> Value -> m (Maybe Value)
updateArray (Value -> Value -> Value -> EvalM Value
op Value
fun) ((Value -> Indexing) -> [Value] -> [Indexing]
forall a b. (a -> b) -> [a] -> [b]
map (Int64 -> Indexing
IndexingFix (Int64 -> Indexing) -> (Value -> Int64) -> Value -> Indexing
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)
_ =
          [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error [Char]
"hist_2d: bad index value"
    def Text
"partition" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
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 (Value -> Int) -> EvalM Value -> EvalM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f Value
x
              [[Value]] -> EvalM [[Value]]
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[Value]] -> EvalM [[Value]]) -> [[Value]] -> EvalM [[Value]]
forall a b. (a -> b) -> a -> b
$ Int -> Value -> [[Value]] -> [[Value]]
forall {t} {t}. (Eq t, Num t) => t -> t -> [[t]] -> [[t]]
insertAt Int
i Value
x [[Value]]
outs
            pack :: [[Value m]] -> Value m
pack [[Value m]]
parts =
              [Value m] -> Value m
forall (m :: * -> *). [Value m] -> Value m
toTuple
                [ ValueShape -> [Value m] -> Value m
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape ([Value m] -> Value m) -> [Value m] -> Value m
forall a b. (a -> b) -> a -> b
$ [[Value m]] -> [Value m]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Value m]]
parts,
                  ValueShape -> [Value m] -> Value m
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape ([Value m] -> Value m) -> [Value m] -> Value m
forall a b. (a -> b) -> a -> b
$
                    ([Value m] -> Value m) -> [[Value m]] -> [Value m]
forall a b. (a -> b) -> [a] -> [b]
map (PrimValue -> Value m
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value m)
-> ([Value m] -> PrimValue) -> [Value m] -> Value m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntValue -> PrimValue
SignedValue (IntValue -> PrimValue)
-> ([Value m] -> IntValue) -> [Value m] -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> IntValue
Int64Value (Int64 -> IntValue)
-> ([Value m] -> Int64) -> [Value m] -> IntValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Value m] -> Int64
forall i a. Num i => [a] -> i
genericLength) [[Value m]]
parts
                ]

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

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

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

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

        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Value] -> Value
forall (m :: * -> *). [Value m] -> Value m
toTuple ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$ ([Value], [Value]) -> [Value]
forall {m :: * -> *}. ([Value m], [Value m]) -> [Value m]
listPair (([Value], [Value]) -> [Value]) -> ([Value], [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ [(Value, Value)] -> ([Value], [Value])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Value, Value)] -> ([Value], [Value]))
-> [(Value, Value)] -> ([Value], [Value])
forall a b. (a -> b) -> a -> b
$ (Value -> (Value, Value)) -> [Value] -> [(Value, Value)]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe [Value] -> (Value, Value)
forall {b}. Maybe [b] -> (b, b)
fromPair (Maybe [Value] -> (Value, Value))
-> (Value -> Maybe [Value]) -> Value -> (Value, Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Maybe [Value]
forall (m :: * -> *). Value m -> Maybe [Value m]
fromTuple) ([Value] -> [(Value, Value)]) -> [Value] -> [(Value, Value)]
forall a b. (a -> b) -> a -> b
$ (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
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]
_ = [Char] -> (b, b)
forall a. HasCallStack => [Char] -> a
error [Char]
"Not a pair"
    def Text
"zip" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 ((Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
xs Value
ys -> do
        let ShapeDim Int64
_ ValueShape
xs_rowshape = Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
xs
            ShapeDim Int64
_ ValueShape
ys_rowshape = Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
ys
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
          ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' (Map Name ValueShape -> ValueShape
forall d. Map Name (Shape d) -> Shape d
ShapeRecord ([ValueShape] -> Map Name ValueShape
forall a. [a] -> Map Name a
tupleFields [ValueShape
xs_rowshape, ValueShape
ys_rowshape])) ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$
            ([Value] -> Value) -> [[Value]] -> [Value]
forall a b. (a -> b) -> [a] -> [b]
map [Value] -> Value
forall (m :: * -> *). [Value m] -> Value m
toTuple ([[Value]] -> [Value]) -> [[Value]] -> [Value]
forall a b. (a -> b) -> a -> b
$
              [[Value]] -> [[Value]]
forall a. [[a]] -> [[a]]
transpose [(ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
xs, (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value]) -> [Value]
forall a b. (a -> b) -> a -> b
$ Value -> (ValueShape, [Value])
fromArray Value
ys]
    def Text
"concat" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> EvalM Value) -> TermBinding
fun2 ((Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> EvalM Value) -> TermBinding
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
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray' ValueShape
rowshape ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$ [Value]
xs' [Value] -> [Value] -> [Value]
forall a. [a] -> [a] -> [a]
++ [Value]
ys'
    def Text
"transpose" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> EvalM Value) -> TermBinding
fun1 ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
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
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$
          ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim Int64
m (Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim Int64
n ValueShape
shape)) ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$
            ([Value] -> Value) -> [[Value]] -> [Value]
forall a b. (a -> b) -> [a] -> [b]
map (ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim Int64
n ValueShape
shape)) ([[Value]] -> [Value]) -> [[Value]] -> [Value]
forall a b. (a -> b) -> a -> b
$
              -- Slight hack to work around empty dimensions.
              Int64 -> [[Value]] -> [[Value]]
forall i a. Integral i => i -> [a] -> [a]
genericTake Int64
m ([[Value]] -> [[Value]]) -> [[Value]] -> [[Value]]
forall a b. (a -> b) -> a -> b
$
                [[Value]] -> [[Value]]
forall a. [[a]] -> [[a]]
transpose ((Value -> [Value]) -> [Value] -> [[Value]]
forall a b. (a -> b) -> [a] -> [b]
map ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (Value -> (ValueShape, [Value])) -> Value -> [Value]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> (ValueShape, [Value])
fromArray) [Value]
xs') [[Value]] -> [[Value]] -> [[Value]]
forall a. [a] -> [a] -> [a]
++ [Value] -> [[Value]]
forall a. a -> [a]
repeat []
    def Text
"flatten" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> EvalM Value) -> TermBinding
fun1 ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
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
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray (Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim (Int64
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
* Int64
m) ValueShape
shape) ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$ (Value -> [Value]) -> [Value] -> [Value]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> (Value -> (ValueShape, [Value])) -> Value -> [Value]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> (ValueShape, [Value])
fromArray) [Value]
xs'
    def Text
"unflatten" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
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 = Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim (Value -> Int64
asInt64 Value
m) ValueShape
innershape
            shape :: ValueShape
shape = Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim (Value -> Int64
asInt64 Value
n) ValueShape
rowshape
        if Value -> Int64
asInt64 Value
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
* Value -> Int64
asInt64 Value
m Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Int64
xs_size Bool -> Bool -> Bool
|| Value -> Int64
asInt64 Value
n Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
0 Bool -> Bool -> Bool
|| Value -> Int64
asInt64 Value
m Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
0
          then
            SrcLoc -> Env -> Text -> EvalM Value
forall a. SrcLoc -> Env -> Text -> EvalM a
bad SrcLoc
forall a. Monoid a => a
mempty Env
forall a. Monoid a => a
mempty (Text -> EvalM Value) -> Text -> EvalM Value
forall a b. (a -> b) -> a -> b
$
              Text
"Cannot unflatten array of shape ["
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int64 -> Text
forall a. Pretty a => a -> Text
prettyText Int64
xs_size
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"] to array of shape ["
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int64 -> Text
forall a. Pretty a => a -> Text
prettyText (Value -> Int64
asInt64 Value
n)
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"]["
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int64 -> Text
forall a. Pretty a => a -> Text
prettyText (Value -> Int64
asInt64 Value
m)
                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"]"
          else Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
shape ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$ ([Value] -> Value) -> [[Value]] -> [Value]
forall a b. (a -> b) -> [a] -> [b]
map (ValueShape -> [Value] -> Value
forall (m :: * -> *). ValueShape -> [Value m] -> Value m
toArray ValueShape
rowshape) ([[Value]] -> [Value]) -> [[Value]] -> [Value]
forall a b. (a -> b) -> a -> b
$ Int -> [Value] -> [[Value]]
forall a. Int -> [a] -> [[a]]
chunk (Value -> Int
asInt Value
m) [Value]
xs'
    def Text
"manifest" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ (Value -> EvalM Value) -> TermBinding
fun1 Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    def Text
"vjp2" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      -- TODO: This could be much better. Currently, it is very inefficient
      -- Perhaps creating VJPValues could be abstracted into a function
      -- exposed by the AD module?
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
f Value
v Value
s -> do
        -- Get the depth
        Int
depth <- [Loc] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Loc] -> Int) -> EvalM [Loc] -> EvalM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM [Loc]
stacktrace

        -- Augment the values
        let v' :: Value
v' =
              Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> Value) -> [Char] -> Value
forall a b. (a -> b) -> a -> b
$ [Char]
"vjp: invalid values " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v) (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$
                (Int -> Value -> Maybe Value) -> Value -> Maybe Value
forall {f :: * -> *} {t} {m :: * -> *}.
(Monad f, Num t) =>
(t -> Value m -> f (Value m)) -> Value m -> f (Value m)
modifyValueM (\Int
i Value
lv -> Int -> ADVariable -> Value
forall (m :: * -> *). Int -> ADVariable -> Value m
ValueAD Int
depth (ADVariable -> Value)
-> (ADValue -> ADVariable) -> ADValue -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VJPValue -> ADVariable
AD.VJP (VJPValue -> ADVariable)
-> (ADValue -> VJPValue) -> ADValue -> ADVariable
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tape -> VJPValue
AD.VJPValue (Tape -> VJPValue) -> (ADValue -> Tape) -> ADValue -> VJPValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ADValue -> Tape
AD.TapeID Int
i (ADValue -> Value) -> Maybe ADValue -> Maybe Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Maybe ADValue
forall {m :: * -> *}. Value m -> Maybe ADValue
getAD Value
lv) Value
v
        -- Turn the seeds into a list of ADValues
        let s' :: [ADValue]
s' =
              [ADValue] -> Maybe [ADValue] -> [ADValue]
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> [ADValue]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [ADValue]) -> [Char] -> [ADValue]
forall a b. (a -> b) -> a -> b
$ [Char]
"vjp: invalid seeds " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Value -> [Char]
forall a. Show a => a -> [Char]
show Value
s) (Maybe [ADValue] -> [ADValue]) -> Maybe [ADValue] -> [ADValue]
forall a b. (a -> b) -> a -> b
$
                (Value -> Maybe ADValue) -> [Value] -> Maybe [ADValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Value -> Maybe ADValue
forall {m :: * -> *}. Value m -> Maybe ADValue
getAD ([Value] -> Maybe [ADValue]) -> [Value] -> Maybe [ADValue]
forall a b. (a -> b) -> a -> b
$
                  ([Value], Value) -> [Value]
forall a b. (a, b) -> a
fst (([Value], Value) -> [Value]) -> ([Value], Value) -> [Value]
forall a b. (a -> b) -> a -> b
$
                    ([Value] -> Value -> ([Value], Value))
-> [Value] -> Value -> ([Value], Value)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\[Value]
a Value
b -> (Value
b Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
a, Value
b)) [] Value
s

        -- Run the function, and turn its outputs into a list of Values
        Value
o <- SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f Value
v'
        let o' :: [Value]
o' = ([Value], Value) -> [Value]
forall a b. (a, b) -> a
fst (([Value], Value) -> [Value]) -> ([Value], Value) -> [Value]
forall a b. (a -> b) -> a -> b
$ ([Value] -> Value -> ([Value], Value))
-> [Value] -> Value -> ([Value], Value)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\[Value]
a Value
b -> (Value
b Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
a, Value
b)) [] Value
o

        -- For each output..
        let m :: [(Value, Map Int ADValue)]
m = (((Value, ADValue) -> (Value, Map Int ADValue))
 -> [(Value, ADValue)] -> [(Value, Map Int ADValue)])
-> [(Value, ADValue)]
-> ((Value, ADValue) -> (Value, Map Int ADValue))
-> [(Value, Map Int ADValue)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Value, ADValue) -> (Value, Map Int ADValue))
-> [(Value, ADValue)] -> [(Value, Map Int ADValue)]
forall a b. (a -> b) -> [a] -> [b]
map ([Value] -> [ADValue] -> [(Value, ADValue)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Value]
o' [ADValue]
s') (((Value, ADValue) -> (Value, Map Int ADValue))
 -> [(Value, Map Int ADValue)])
-> ((Value, ADValue) -> (Value, Map Int ADValue))
-> [(Value, Map Int ADValue)]
forall a b. (a -> b) -> a -> b
$ \(Value
on, ADValue
sn) -> case Value
on of
              -- If it is a VJP variable of the correct depth, run
              -- deriveTapqe on it- and its corresponding seed
              (ValueAD Int
d (AD.VJP (AD.VJPValue Tape
t)))
                | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
depth ->
                    (ADValue -> Value
forall {m :: * -> *}. ADValue -> Value m
putAD (ADValue -> Value) -> ADValue -> Value
forall a b. (a -> b) -> a -> b
$ Tape -> ADValue
AD.tapePrimal Tape
t, Tape -> ADValue -> Map Int ADValue
AD.deriveTape Tape
t ADValue
sn)
              -- Otherwise, its partial derivatives are all 0
              Value
_ -> (Value
on, Map Int ADValue
forall k a. Map k a
M.empty)

        -- Add together every derivative
        let drvs :: Map Int (Maybe (Value m))
drvs = (ADValue -> Maybe (Value m))
-> Map Int ADValue -> Map Int (Maybe (Value m))
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Value m -> Maybe (Value m)
forall a. a -> Maybe a
Just (Value m -> Maybe (Value m))
-> (ADValue -> Value m) -> ADValue -> Maybe (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADValue -> Value m
forall {m :: * -> *}. ADValue -> Value m
putAD) (Map Int ADValue -> Map Int (Maybe (Value m)))
-> Map Int ADValue -> Map Int (Maybe (Value m))
forall a b. (a -> b) -> a -> b
$ (ADValue -> ADValue -> ADValue)
-> [Map Int ADValue] -> Map Int ADValue
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith ADValue -> ADValue -> ADValue
add ([Map Int ADValue] -> Map Int ADValue)
-> [Map Int ADValue] -> Map Int ADValue
forall a b. (a -> b) -> a -> b
$ ((Value, Map Int ADValue) -> Map Int ADValue)
-> [(Value, Map Int ADValue)] -> [Map Int ADValue]
forall a b. (a -> b) -> [a] -> [b]
map (Value, Map Int ADValue) -> Map Int ADValue
forall a b. (a, b) -> b
snd [(Value, Map Int ADValue)]
m

        -- Extract the output values, and the partial derivatives
        let ov :: Value
ov = (Int -> Value -> Value) -> Value -> Value
forall {t} {m :: * -> *}.
Num t =>
(t -> Value m -> Value m) -> Value m -> Value m
modifyValue (\Int
i Value
_ -> (Value, Map Int ADValue) -> Value
forall a b. (a, b) -> a
fst ((Value, Map Int ADValue) -> Value)
-> (Value, Map Int ADValue) -> Value
forall a b. (a -> b) -> a -> b
$ [(Value, Map Int ADValue)]
m [(Value, Map Int ADValue)] -> Int -> (Value, Map Int ADValue)
forall a. HasCallStack => [a] -> Int -> a
!! ([(Value, Map Int ADValue)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Value, Map Int ADValue)]
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)) Value
o
        let od :: Value
od =
              Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Value
forall a. HasCallStack => [Char] -> a
error [Char]
"vjp: differentiation failed") (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$
                (Int -> Value -> Maybe Value) -> Value -> Maybe Value
forall {f :: * -> *} {t} {m :: * -> *}.
(Monad f, Num t) =>
(t -> Value m -> f (Value m)) -> Value m -> f (Value m)
modifyValueM (\Int
i Value
vo -> Maybe Value -> Int -> Map Int (Maybe Value) -> Maybe Value
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault (PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> (ADValue -> PrimValue) -> ADValue -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV (PrimValue -> PrimValue)
-> (ADValue -> PrimValue) -> ADValue -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> PrimValue
P.blankPrimValue (PrimType -> PrimValue)
-> (ADValue -> PrimType) -> ADValue -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimType
P.primValueType (PrimValue -> PrimType)
-> (ADValue -> PrimValue) -> ADValue -> PrimType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADValue -> PrimValue
AD.primitive (ADValue -> Value) -> Maybe ADValue -> Maybe Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Maybe ADValue
forall {m :: * -> *}. Value m -> Maybe ADValue
getAD Value
vo) Int
i Map Int (Maybe Value)
forall {m :: * -> *}. Map Int (Maybe (Value m))
drvs) Value
v

        -- Return a tuple of the output values, and partial derivatives
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Value] -> Value
forall (m :: * -> *). [Value m] -> Value m
toTuple [Value
ov, Value
od]
      where
        modifyValue :: (t -> Value m -> Value m) -> Value m -> Value m
modifyValue t -> Value m -> Value m
f Value m
v = (t, Value m) -> Value m
forall a b. (a, b) -> b
snd ((t, Value m) -> Value m) -> (t, Value m) -> Value m
forall a b. (a -> b) -> a -> b
$ (t -> Value m -> (t, Value m)) -> t -> Value m -> (t, Value m)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\t
a Value m
b -> (t
a t -> t -> t
forall a. Num a => a -> a -> a
+ t
1, t -> Value m -> Value m
f t
a Value m
b)) t
0 Value m
v
        modifyValueM :: (t -> Value m -> f (Value m)) -> Value m -> f (Value m)
modifyValueM t -> Value m -> f (Value m)
f Value m
v =
          (t, Value m) -> Value m
forall a b. (a, b) -> b
snd
            ((t, Value m) -> Value m) -> f (t, Value m) -> f (Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (t -> Value m -> f (t, Value m)) -> t -> Value m -> f (t, Value m)
forall (f :: * -> *) a (m :: * -> *).
Monad f =>
(a -> Value m -> f (a, Value m)) -> a -> Value m -> f (a, Value m)
valueAccumLM
              ( \t
a Value m
b -> do
                  Value m
b' <- t -> Value m -> f (Value m)
f t
a Value m
b
                  (t, Value m) -> f (t, Value m)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (t
a t -> t -> t
forall a. Num a => a -> a -> a
+ t
1, Value m
b')
              )
              t
0
              Value m
v

        -- TODO: Perhaps this could be fully abstracted by AD?
        -- Making addFor private would be nice..
        add :: ADValue -> ADValue -> ADValue
add ADValue
x ADValue
y =
          ADValue -> Maybe ADValue -> ADValue
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> ADValue
forall a. HasCallStack => [Char] -> a
error [Char]
"jvp: illtyped add") (Maybe ADValue -> ADValue) -> Maybe ADValue -> ADValue
forall a b. (a -> b) -> a -> b
$
            Op -> [ADValue] -> Maybe ADValue
AD.doOp (BinOp -> Op
AD.OpBin (BinOp -> Op) -> BinOp -> Op
forall a b. (a -> b) -> a -> b
$ PrimType -> BinOp
AD.addFor (PrimType -> BinOp) -> PrimType -> BinOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
P.primValueType (PrimValue -> PrimType) -> PrimValue -> PrimType
forall a b. (a -> b) -> a -> b
$ ADValue -> PrimValue
AD.primitive ADValue
x) [ADValue
x, ADValue
y]
    def Text
"jvp2" = TermBinding -> Maybe TermBinding
forall a. a -> Maybe a
Just (TermBinding -> Maybe TermBinding)
-> TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$
      -- TODO: This could be much better. Currently, it is very inefficient
      -- Perhaps creating JVPValues could be abstracted into a function
      -- exposed by the AD module?
      (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 ((Value -> Value -> Value -> EvalM Value) -> TermBinding)
-> (Value -> Value -> Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
f Value
v Value
s -> do
        -- Get the depth
        Int
depth <- [Loc] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Loc] -> Int) -> EvalM [Loc] -> EvalM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM [Loc]
stacktrace

        -- Turn the seeds into a list of ADValues
        let s' :: [ADValue]
s' =
              [Char] -> Maybe [ADValue] -> [ADValue]
forall {a}. [Char] -> Maybe a -> a
expectJust ([Char]
"jvp: invalid seeds " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Value -> [Char]
forall a. Show a => a -> [Char]
show Value
s) (Maybe [ADValue] -> [ADValue]) -> Maybe [ADValue] -> [ADValue]
forall a b. (a -> b) -> a -> b
$
                (Value -> Maybe ADValue) -> [Value] -> Maybe [ADValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Value -> Maybe ADValue
forall {m :: * -> *}. Value m -> Maybe ADValue
getAD ([Value] -> Maybe [ADValue]) -> [Value] -> Maybe [ADValue]
forall a b. (a -> b) -> a -> b
$
                  ([Value], Value) -> [Value]
forall a b. (a, b) -> a
fst (([Value], Value) -> [Value]) -> ([Value], Value) -> [Value]
forall a b. (a -> b) -> a -> b
$
                    ([Value] -> Value -> ([Value], Value))
-> [Value] -> Value -> ([Value], Value)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\[Value]
a Value
b -> (Value
b Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
a, Value
b)) [] Value
s
        -- Augment the values
        let v' :: Value
v' =
              [Char] -> Maybe Value -> Value
forall {a}. [Char] -> Maybe a -> a
expectJust ([Char]
"jvp: invalid values " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Value -> [Char]
forall a. Show a => a -> [Char]
show Value
v) (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$
                (Int -> Value -> Maybe Value) -> Value -> Maybe Value
forall {f :: * -> *} {t} {m :: * -> *}.
(Monad f, Num t) =>
(t -> Value m -> f (Value m)) -> Value m -> f (Value m)
modifyValueM
                  ( \Int
i Value
lv -> do
                      ADValue
lv' <- Value -> Maybe ADValue
forall {m :: * -> *}. Value m -> Maybe ADValue
getAD Value
lv
                      Value -> Maybe Value
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Int -> ADVariable -> Value
forall (m :: * -> *). Int -> ADVariable -> Value m
ValueAD Int
depth (ADVariable -> Value)
-> (ADValue -> ADVariable) -> ADValue -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JVPValue -> ADVariable
AD.JVP (JVPValue -> ADVariable)
-> (ADValue -> JVPValue) -> ADValue -> ADVariable
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADValue -> ADValue -> JVPValue
AD.JVPValue ADValue
lv' (ADValue -> Value) -> ADValue -> Value
forall a b. (a -> b) -> a -> b
$ [ADValue]
s' [ADValue] -> Int -> ADValue
forall a. HasCallStack => [a] -> Int -> a
!! ([ADValue] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ADValue]
s' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)
                  )
                  Value
v

        -- Run the function, and turn its outputs into a list of Values
        Value
o <- SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f Value
v'
        let o' :: [Value]
o' = ([Value], Value) -> [Value]
forall a b. (a, b) -> a
fst (([Value], Value) -> [Value]) -> ([Value], Value) -> [Value]
forall a b. (a -> b) -> a -> b
$ ([Value] -> Value -> ([Value], Value))
-> [Value] -> Value -> ([Value], Value)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\[Value]
a Value
b -> (Value
b Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
: [Value]
a, Value
b)) [] Value
o

        -- For each output..
        let m :: [(Value, Value m)]
m =
              [Char] -> Maybe [(Value, Value m)] -> [(Value, Value m)]
forall {a}. [Char] -> Maybe a -> a
expectJust [Char]
"jvp: differentiation failed" (Maybe [(Value, Value m)] -> [(Value, Value m)])
-> Maybe [(Value, Value m)] -> [(Value, Value m)]
forall a b. (a -> b) -> a -> b
$
                (Value -> Maybe (Value, Value m))
-> [Value] -> Maybe [(Value, Value m)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                  ( \Value
on -> case Value
on of
                      -- If it is a JVP variable of the correct depth, return its primal and derivative
                      (ValueAD Int
d (AD.JVP (AD.JVPValue ADValue
pv ADValue
dv))) | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
depth -> (Value, Value m) -> Maybe (Value, Value m)
forall a. a -> Maybe a
Just (ADValue -> Value
forall {m :: * -> *}. ADValue -> Value m
putAD ADValue
pv, ADValue -> Value m
forall {m :: * -> *}. ADValue -> Value m
putAD ADValue
dv)
                      -- Otherwise, its partial derivatives are all 0
                      Value
_ -> (Value
on,) (Value m -> (Value, Value m))
-> (ADValue -> Value m) -> ADValue -> (Value, Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> Value m
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value m)
-> (ADValue -> PrimValue) -> ADValue -> Value m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimValue
putV (PrimValue -> PrimValue)
-> (ADValue -> PrimValue) -> ADValue -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> PrimValue
P.blankPrimValue (PrimType -> PrimValue)
-> (ADValue -> PrimType) -> ADValue -> PrimValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> PrimType
P.primValueType (PrimValue -> PrimType)
-> (ADValue -> PrimValue) -> ADValue -> PrimType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADValue -> PrimValue
AD.primitive (ADValue -> (Value, Value m))
-> Maybe ADValue -> Maybe (Value, Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Maybe ADValue
forall {m :: * -> *}. Value m -> Maybe ADValue
getAD Value
on
                  )
                  [Value]
o'

        -- Extract the output values, and the partial derivatives
        let ov :: Value
ov = (Int -> Value -> Value) -> Value -> Value
forall {t} {m :: * -> *}.
Num t =>
(t -> Value m -> Value m) -> Value m -> Value m
modifyValue (\Int
i Value
_ -> (Value, Value Any) -> Value
forall a b. (a, b) -> a
fst ((Value, Value Any) -> Value) -> (Value, Value Any) -> Value
forall a b. (a -> b) -> a -> b
$ [(Value, Value Any)]
forall {m :: * -> *}. [(Value, Value m)]
m [(Value, Value Any)] -> Int -> (Value, Value Any)
forall a. HasCallStack => [a] -> Int -> a
!! ([(Value, Value Any)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Value, Value Any)]
forall {m :: * -> *}. [(Value, Value m)]
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)) Value
o
            od :: Value
od = (Int -> Value -> Value) -> Value -> Value
forall {t} {m :: * -> *}.
Num t =>
(t -> Value m -> Value m) -> Value m -> Value m
modifyValue (\Int
i Value
_ -> (Value, Value) -> Value
forall a b. (a, b) -> b
snd ((Value, Value) -> Value) -> (Value, Value) -> Value
forall a b. (a -> b) -> a -> b
$ [(Value, Value)]
forall {m :: * -> *}. [(Value, Value m)]
m [(Value, Value)] -> Int -> (Value, Value)
forall a. HasCallStack => [a] -> Int -> a
!! ([(Value, Value Any)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Value, Value Any)]
forall {m :: * -> *}. [(Value, Value m)]
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)) Value
o

        -- Return a tuple of the output values, and partial derivatives
        Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> EvalM Value) -> Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Value] -> Value
forall (m :: * -> *). [Value m] -> Value m
toTuple [Value
ov, Value
od]
      where
        modifyValue :: (t -> Value m -> Value m) -> Value m -> Value m
modifyValue t -> Value m -> Value m
f Value m
v = (t, Value m) -> Value m
forall a b. (a, b) -> b
snd ((t, Value m) -> Value m) -> (t, Value m) -> Value m
forall a b. (a -> b) -> a -> b
$ (t -> Value m -> (t, Value m)) -> t -> Value m -> (t, Value m)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\t
a Value m
b -> (t
a t -> t -> t
forall a. Num a => a -> a -> a
+ t
1, t -> Value m -> Value m
f t
a Value m
b)) t
0 Value m
v
        modifyValueM :: (t -> Value m -> f (Value m)) -> Value m -> f (Value m)
modifyValueM t -> Value m -> f (Value m)
f Value m
v =
          (t, Value m) -> Value m
forall a b. (a, b) -> b
snd
            ((t, Value m) -> Value m) -> f (t, Value m) -> f (Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (t -> Value m -> f (t, Value m)) -> t -> Value m -> f (t, Value m)
forall (f :: * -> *) a (m :: * -> *).
Monad f =>
(a -> Value m -> f (a, Value m)) -> a -> Value m -> f (a, Value m)
valueAccumLM
              ( \t
a Value m
b -> do
                  Value m
b' <- t -> Value m -> f (Value m)
f t
a Value m
b
                  (t, Value m) -> f (t, Value m)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (t
a t -> t -> t
forall a. Num a => a -> a -> a
+ t
1, Value m
b')
              )
              t
0
              Value m
v

        expectJust :: [Char] -> Maybe a -> a
expectJust [Char]
_ (Just a
v) = a
v
        expectJust [Char]
s Maybe a
Nothing = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
s
    def Text
"acc" = Maybe TermBinding
forall a. Maybe a
Nothing
    def Text
s | Text -> Name
nameFromText Text
s Name -> Map Name PrimType -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map Name PrimType
namesToPrimTypes = Maybe TermBinding
forall a. Maybe a
Nothing
    def Text
s = [Char] -> Maybe TermBinding
forall a. HasCallStack => [Char] -> a
error ([Char] -> Maybe TermBinding) -> [Char] -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ [Char]
"Missing intrinsic: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
T.unpack Text
s

    tdef :: Name -> Maybe TypeBinding
    tdef :: Name -> Maybe TypeBinding
tdef Name
s = do
      PrimType
t <- Name
s Name -> Map Name PrimType -> Maybe PrimType
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map Name PrimType
namesToPrimTypes
      TypeBinding -> Maybe TypeBinding
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBinding -> Maybe TypeBinding)
-> TypeBinding -> Maybe TypeBinding
forall a b. (a -> b) -> a -> b
$ Env -> [TypeParam] -> StructRetType -> TypeBinding
TypeBinding Env
forall a. Monoid a => a
mempty [] (StructRetType -> TypeBinding) -> StructRetType -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (StructType -> StructRetType) -> StructType -> StructRetType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (ExpBase Info VName) NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase (ExpBase Info VName) NoUniqueness -> StructType)
-> ScalarTypeBase (ExpBase Info VName) NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (ExpBase Info VName) NoUniqueness
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
t

intrinsicVal :: Name -> Value
intrinsicVal :: Name -> Value
intrinsicVal Name
name =
  case VName -> Map VName TermBinding -> Maybe TermBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Name -> VName
intrinsicVar Name
name) (Map VName TermBinding -> Maybe TermBinding)
-> Map VName TermBinding -> Maybe TermBinding
forall a b. (a -> b) -> a -> b
$ Env -> Map VName TermBinding
envTerm (Env -> Map VName TermBinding) -> Env -> Map VName TermBinding
forall a b. (a -> b) -> a -> b
$ Ctx -> Env
ctxEnv Ctx
initialCtx of
    Just (TermValue Maybe BoundV
_ Value
v) -> Value
v
    Maybe TermBinding
_ -> [Char] -> Value
forall a. HasCallStack => [Char] -> a
error ([Char] -> Value) -> [Char] -> Value
forall a b. (a -> b) -> a -> b
$ [Char]
"intrinsicVal: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
name

intrinsicsNeg :: Value
intrinsicsNeg :: Value
intrinsicsNeg = Name -> Value
intrinsicVal Name
"neg"

intrinsicsNot :: Value
intrinsicsNot :: Value
intrinsicsNot = Name -> Value
intrinsicVal Name
"!"

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

interpretDecs :: Ctx -> [Dec] -> F ExtOp Env
interpretDecs :: Ctx -> [DecBase Info VName] -> F ExtOp Env
interpretDecs Ctx
ctx [DecBase Info VName]
decs =
  Map ImportName Env -> EvalM Env -> F ExtOp Env
forall a. Map ImportName Env -> EvalM a -> F ExtOp a
runEvalM (Ctx -> Map ImportName Env
ctxImports Ctx
ctx) (EvalM Env -> F ExtOp Env) -> EvalM Env -> F ExtOp Env
forall a b. (a -> b) -> a -> b
$ do
    Env
env <- (Env -> DecBase Info VName -> EvalM Env)
-> Env -> [DecBase Info VName] -> EvalM Env
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) [DecBase Info VName]
decs
    -- 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
extEnv
    Env -> EvalM Env
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env -> EvalM Env) -> Env -> EvalM Env
forall a b. (a -> b) -> a -> b
$ Env
env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
sizes

interpretDec :: Ctx -> Dec -> F ExtOp Ctx
interpretDec :: Ctx -> DecBase Info VName -> F ExtOp Ctx
interpretDec Ctx
ctx DecBase Info VName
d = do
  Env
env <- Ctx -> [DecBase Info VName] -> F ExtOp Env
interpretDecs Ctx
ctx [DecBase Info VName
d]
  Ctx -> F ExtOp Ctx
forall a. a -> F ExtOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx
ctx {ctxEnv = env}

interpretImport :: Ctx -> (ImportName, Prog) -> F ExtOp Ctx
interpretImport :: Ctx -> (ImportName, Prog) -> F ExtOp Ctx
interpretImport Ctx
ctx (ImportName
fp, Prog
prog) = do
  Env
env <- Ctx -> [DecBase Info VName] -> F ExtOp Env
interpretDecs Ctx
ctx ([DecBase Info VName] -> F ExtOp Env)
-> [DecBase Info VName] -> F ExtOp Env
forall a b. (a -> b) -> a -> b
$ Prog -> [DecBase Info VName]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs Prog
prog
  Ctx -> F ExtOp Ctx
forall a. a -> F ExtOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx
ctx {ctxImports = M.insert fp env $ ctxImports 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 = mconcat (reverse envs) <> ctxEnv ctx}

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

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

-- | Execute the named function on the given arguments; may fail
-- horribly if these are ill-typed.
interpretFunction :: Ctx -> VName -> [V.Value] -> Either T.Text (F ExtOp Value)
interpretFunction :: Ctx -> VName -> [Value] -> Either Text (F ExtOp Value)
interpretFunction Ctx
ctx VName
fname [Value]
vs = do
  let env :: Env
env = Ctx -> Env
ctxEnv Ctx
ctx

  (StructType
ft, EvalM Value
mkf) <- case QualName VName -> Env -> Maybe TermBinding
lookupVar (VName -> QualName VName
forall v. v -> QualName v
qualName VName
fname) Env
env of
    Just (TermValue (Just (T.BoundV [TypeParam]
_ StructType
t)) Value
v) -> do
      StructType
ft <- [ValueType] -> StructType -> Either Text StructType
updateType ((Value -> ValueType) -> [Value] -> [ValueType]
forall a b. (a -> b) -> [a] -> [b]
map Value -> ValueType
valueType [Value]
vs) StructType
t
      (StructType, EvalM Value) -> Either Text (StructType, EvalM Value)
forall a. a -> Either Text a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType
ft, Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v)
    Just (TermPoly (Just (T.BoundV [TypeParam]
_ StructType
t)) TypeBase SizeClosure NoUniqueness -> EvalM Value
v) -> do
      StructType
ft <- [ValueType] -> StructType -> Either Text StructType
updateType ((Value -> ValueType) -> [Value] -> [ValueType]
forall a b. (a -> b) -> [a] -> [b]
map Value -> ValueType
valueType [Value]
vs) StructType
t
      (StructType, EvalM Value) -> Either Text (StructType, EvalM Value)
forall a. a -> Either Text a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType
ft, TypeBase SizeClosure NoUniqueness -> EvalM Value
v (Env -> StructType -> TypeBase SizeClosure NoUniqueness
structToEval Env
env StructType
ft))
    Maybe TermBinding
_ ->
      Text -> Either Text (StructType, EvalM Value)
forall a b. a -> Either a b
Left (Text -> Either Text (StructType, EvalM Value))
-> Text -> Either Text (StructType, EvalM Value)
forall a b. (a -> b) -> a -> b
$ Text
"Unknown function `" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Name -> Text
nameToText (VName -> Name
forall v. IsName v => v -> Name
toName VName
fname) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"`."

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

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

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

    checkInput :: ValueType -> StructType -> Either T.Text ()
    checkInput :: ValueType -> StructType -> Either Text ()
checkInput (Scalar (Prim PrimType
vt)) (Scalar (Prim PrimType
pt))
      | PrimType
vt PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimType
pt = PrimType -> PrimType -> Either Text ()
forall {a} {a} {b}. (Pretty a, Pretty a) => a -> a -> Either Text b
badPrim PrimType
vt PrimType
pt
    checkInput (Array NoUniqueness
_ Shape Int64
_ (Prim PrimType
vt)) (Array NoUniqueness
_ Shape (ExpBase Info VName)
_ (Prim PrimType
pt))
      | PrimType
vt PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimType
pt = PrimType -> PrimType -> Either Text ()
forall {a} {a} {b}. (Pretty a, Pretty a) => a -> a -> Either Text b
badPrim PrimType
vt PrimType
pt
    checkInput vArr :: ValueType
vArr@(Array NoUniqueness
_ (F.Shape [Int64]
vd) ScalarTypeBase Int64 NoUniqueness
_) pArr :: StructType
pArr@(Array NoUniqueness
_ (F.Shape [ExpBase Info VName]
pd) ScalarTypeBase (ExpBase Info VName) NoUniqueness
_)
      | [Int64] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int64]
vd Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [ExpBase Info VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info VName]
pd = ValueType -> StructType -> Either Text ()
forall {a} {a} {b}. (Pretty a, Pretty a) => a -> a -> Either Text b
badDim ValueType
vArr StructType
pArr
      | Bool -> Bool
not (Bool -> Bool) -> ([Bool] -> Bool) -> [Bool] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Int64 -> ExpBase Info VName -> Bool)
-> [Int64] -> [ExpBase Info VName] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int64 -> ExpBase Info VName -> Bool
sameShape [Int64]
vd [ExpBase Info VName]
pd = ValueType -> StructType -> Either Text ()
forall {a} {a} {b}. (Pretty a, Pretty a) => a -> a -> Either Text b
badDim ValueType
vArr StructType
pArr
      where
        sameShape :: Int64 -> Size -> Bool
        sameShape :: Int64 -> ExpBase Info VName -> Bool
sameShape Int64
shape0 (IntLit Integer
shape1 Info StructType
_ SrcLoc
_) = Int64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
shape0 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
shape1
        sameShape Int64
_ ExpBase Info VName
_ = Bool
True
    checkInput ValueType
_ StructType
_ =
      () -> Either Text ()
forall a b. b -> Either a b
Right ()

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

    badDim :: a -> a -> Either Text b
badDim a
vd a
pd =
      Text -> Either Text b
forall a b. a -> Either a b
Left (Text -> Either Text b)
-> (Doc Any -> Text) -> Doc Any -> Either Text b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc Any -> Text
forall a. Doc a -> Text
docText (Doc Any -> Either Text b) -> Doc Any -> Either Text b
forall a b. (a -> b) -> a -> b
$
        Doc Any
"Invalid argument dimensions."
          Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc Any
"Expected:"
          Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann
align (a -> Doc Any
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
pd)
          Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc Any
"Got:     "
          Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann
align (a -> Doc Any
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
vd)