{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Strict #-}

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

    -- * Values
    Value,
    fromTuple,
    isEmptyArray,
    asByteString,
    prettyEmptyArray,
    prettyValue,
    valueText,

    -- * IO
    IOOp,
    determineIO,
    doIOOp,
    ioRelativeTo,
  )
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.Functor (($>), (<&>))
import Data.List
  ( find,
    genericLength,
    genericTake,
    transpose,
  )
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict 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.FFI.Push qualified as FFI
import Language.Futhark.Interpreter.FFI.ServerM qualified as FFI
import Language.Futhark.Interpreter.IO
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
  | -- | An arithmetic operation produced a fresh NaN.
    BreakNaN

data ExtOp a
  = ExtOpTrace T.Text (Doc ()) a
  | ExtOpBreak Loc BreakReason (NE.NonEmpty StackFrame) a
  | ExtOpError InterpreterError
  | ExtOpFFI (FFI.ServerM Value) (Value -> a)
  | -- | Perform an IO operation.
    ExtOpIO IOOp (Value -> a)

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
  fmap a -> b
f (ExtOpFFI ServerM Value
vr Value -> a
c) = ServerM Value -> (Value -> b) -> ExtOp b
forall a. ServerM Value -> (Value -> a) -> ExtOp a
ExtOpFFI ServerM Value
vr ((Value -> b) -> ExtOp b) -> (Value -> b) -> ExtOp b
forall a b. (a -> b) -> a -> b
$ a -> b
f (a -> b) -> (Value -> a) -> Value -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> a
c
  fmap a -> b
f (ExtOpIO IOOp
v Value -> a
c) = IOOp -> (Value -> b) -> ExtOp b
forall a. IOOp -> (Value -> a) -> ExtOp a
ExtOpIO IOOp
v ((Value -> b) -> ExtOp b) -> (Value -> b) -> ExtOp b
forall a b. (a -> b) -> a -> b
$ a -> b
f (a -> b) -> (Value -> a) -> Value -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> a
c

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, AD.Counter) (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, AD.Counter)
    )

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, Counter) (F ExtOp))
  a
m) = StateT (Exts, Counter) (F ExtOp) a -> (Exts, Counter) -> F ExtOp a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (ReaderT
  ([StackFrame], Map ImportName Env)
  (StateT (Exts, Counter) (F ExtOp))
  a
-> ([StackFrame], Map ImportName Env)
-> StateT (Exts, Counter) (F ExtOp) a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT
  ([StackFrame], Map ImportName Env)
  (StateT (Exts, Counter) (F ExtOp))
  a
m ([StackFrame]
forall a. Monoid a => a
mempty, Map ImportName Env
imports)) (Exts
forall a. Monoid a => a
mempty, Int -> Counter
AD.Counter Int
0)

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

-- | Instead of tracking the actual depth of AD, we just use the size
-- of the stack as a proxy.
adDepth :: EvalM AD.Depth
adDepth :: EvalM Depth
adDepth = Int -> Depth
AD.Depth (Int -> Depth) -> ([Loc] -> Int) -> [Loc] -> Depth
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Loc] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Loc] -> Depth) -> EvalM [Loc] -> EvalM Depth
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM [Loc]
stacktrace

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, Counter) -> (Exts, Counter)) -> EvalM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((Exts, Counter) -> (Exts, Counter)) -> EvalM ())
-> ((Exts, Counter) -> (Exts, Counter)) -> EvalM ()
forall a b. (a -> b) -> a -> b
$ (Exts -> Exts) -> (Exts, Counter) -> (Exts, Counter)
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 ((Exts -> Exts) -> (Exts, Counter) -> (Exts, Counter))
-> (Exts -> Exts) -> (Exts, Counter) -> (Exts, Counter)
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 = ((Exts, Counter) -> Exts) -> EvalM Exts
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Exts, Counter) -> Exts
forall a b. (a, b) -> a
fst

putCounter :: AD.Counter -> EvalM ()
putCounter :: Counter -> EvalM ()
putCounter Counter
i = ((Exts, Counter) -> (Exts, Counter)) -> EvalM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((Exts, Counter) -> (Exts, Counter)) -> EvalM ())
-> ((Exts, Counter) -> (Exts, Counter)) -> EvalM ()
forall a b. (a -> b) -> a -> b
$ (Counter -> Counter) -> (Exts, Counter) -> (Exts, Counter)
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 ((Counter -> Counter) -> (Exts, Counter) -> (Exts, Counter))
-> (Counter -> Counter) -> (Exts, Counter) -> (Exts, Counter)
forall a b. (a -> b) -> a -> b
$ Counter -> Counter -> Counter
forall a b. a -> b -> a
const Counter
i

getCounter :: EvalM AD.Counter
getCounter :: EvalM Counter
getCounter = ((Exts, Counter) -> Counter) -> EvalM Counter
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Exts, Counter) -> Counter
forall a b. (a, b) -> b
snd

-- | 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
  e <- EvalM Exts
getExts
  x <- m
  i <- getCounter
  put (e, i)
  pure 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 (Maybe BoundV
forall a. Maybe a
Nothing,) (Exts -> Env) -> EvalM Exts -> EvalM Env
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EvalM Exts
getExts

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
brackets (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 NoMode

structToEval :: Env -> StructType -> EvalType
structToEval :: Env -> StructType -> TypeBase SizeClosure NoMode
structToEval Env
env = (ExpBase Info VName -> SizeClosure)
-> StructType -> TypeBase SizeClosure NoMode
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 NoMode -> StructType
evalToStruct = (SizeClosure -> ExpBase Info VName)
-> TypeBase SizeClosure NoMode -> 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 NoMode
-> ([(VName, ([VName], TypeBase SizeClosure NoMode))],
    [(VName, SizeClosure)])
resolveTypeParams [VName]
names StructType
orig_t1 TypeBase SizeClosure NoMode
orig_t2 =
  State
  ([(VName, ([VName], TypeBase SizeClosure NoMode))],
   [(VName, SizeClosure)])
  ()
-> ([(VName, ([VName], TypeBase SizeClosure NoMode))],
    [(VName, SizeClosure)])
-> ([(VName, ([VName], TypeBase SizeClosure NoMode))],
    [(VName, SizeClosure)])
forall s a. State s a -> s -> s
execState ([VName]
-> StructType
-> TypeBase SizeClosure NoMode
-> State
     ([(VName, ([VName], TypeBase SizeClosure NoMode))],
      [(VName, SizeClosure)])
     ()
forall {m :: * -> *} {p :: * -> * -> *}.
(MonadState
   (p [(VName, ([VName], TypeBase SizeClosure NoMode))]
      [(VName, SizeClosure)])
   m,
 Bifunctor p) =>
[VName] -> StructType -> TypeBase SizeClosure NoMode -> m ()
match [VName]
forall a. Monoid a => a
mempty StructType
orig_t1 TypeBase SizeClosure NoMode
orig_t2) ([(VName, ([VName], TypeBase SizeClosure NoMode))],
 [(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 NoMode -> m ()
match [VName]
bound (Scalar (TypeVar NoMode
_ QualName VName
tn [TypeArg (ExpBase Info VName)]
_)) TypeBase SizeClosure NoMode
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 NoMode) -> 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 NoMode
t)
    match [VName]
bound (Scalar (Record Map Name StructType
poly_fields)) (Scalar (Record Map Name (TypeBase SizeClosure NoMode)
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 NoMode -> m ())
-> Map Name StructType
-> Map Name (TypeBase SizeClosure NoMode)
-> 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 NoMode -> m ()
match [VName]
bound) Map Name StructType
poly_fields Map Name (TypeBase SizeClosure NoMode)
fields
    match [VName]
bound (Scalar (Sum Map Name [StructType]
poly_fields)) (Scalar (Sum Map Name [TypeBase SizeClosure NoMode]
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 NoMode] -> [m ()])
-> Map Name [StructType]
-> Map Name [TypeBase SizeClosure NoMode]
-> 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 NoMode -> m ())
-> [StructType] -> [TypeBase SizeClosure NoMode] -> [m ()]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((StructType -> TypeBase SizeClosure NoMode -> m ())
 -> [StructType] -> [TypeBase SizeClosure NoMode] -> [m ()])
-> (StructType -> TypeBase SizeClosure NoMode -> m ())
-> [StructType]
-> [TypeBase SizeClosure NoMode]
-> [m ()]
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> TypeBase SizeClosure NoMode -> m ()
match [VName]
bound) Map Name [StructType]
poly_fields Map Name [TypeBase SizeClosure NoMode]
fields
    match
      [VName]
bound
      (Scalar (Arrow NoMode
_ PName
p1 Diet
_ StructType
poly_t1 (RetType [VName]
dims1 TypeBase (ExpBase Info VName) Freshness
poly_t2)))
      (Scalar (Arrow NoMode
_ PName
p2 Diet
_ TypeBase SizeClosure NoMode
t1 (RetType [VName]
dims2 TypeBase SizeClosure Freshness
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 NoMode -> m ()
match [VName]
bound' StructType
poly_t1 TypeBase SizeClosure NoMode
t1
        [VName] -> StructType -> TypeBase SizeClosure NoMode -> m ()
match [VName]
bound' (TypeBase (ExpBase Info VName) Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) Freshness
poly_t2) (TypeBase SizeClosure Freshness -> TypeBase SizeClosure NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase SizeClosure Freshness
t2)
    match [VName]
bound StructType
poly_t TypeBase SizeClosure NoMode
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 o. TypeBase dim o -> Shape dim
arrayShape StructType
poly_t),
        SizeClosure
d2 : [SizeClosure]
_ <- Shape SizeClosure -> [SizeClosure]
forall dim. Shape dim -> [dim]
shapeDims (TypeBase SizeClosure NoMode -> Shape SizeClosure
forall dim o. TypeBase dim o -> Shape dim
arrayShape TypeBase SizeClosure NoMode
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 NoMode -> m ()
match [VName]
bound (Int -> StructType -> StructType
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 StructType
poly_t) (Int -> TypeBase SizeClosure NoMode -> TypeBase SizeClosure NoMode
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 TypeBase SizeClosure NoMode
t)
    match [VName]
bound StructType
t1 TypeBase SizeClosure NoMode
t2
      | Just StructType
t1' <- StructType -> Maybe StructType
forall d o. TypeBase d o -> Maybe (TypeBase d NoMode)
isAccType StructType
t1,
        Just TypeBase SizeClosure NoMode
t2' <- TypeBase SizeClosure NoMode -> Maybe (TypeBase SizeClosure NoMode)
forall d o. TypeBase d o -> Maybe (TypeBase d NoMode)
isAccType TypeBase SizeClosure NoMode
t2 =
          [VName] -> StructType -> TypeBase SizeClosure NoMode -> m ()
match [VName]
bound StructType
t1' TypeBase SizeClosure NoMode
t2'
    match [VName]
_ StructType
_ TypeBase SizeClosure NoMode
_ = () -> 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)
      | Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (ExpBase Info VName -> Maybe Int
isAnySize ExpBase Info VName
e1) Bool -> Bool -> Bool
|| Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (ExpBase Info VName -> Maybe Int
isAnySize ExpBase Info VName
e2) = () -> 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
  size_env <- EvalM Env
extEnv
  eval (size_env <> env) e

evalResolved ::
  ([(VName, ([VName], EvalType))], [(VName, SizeClosure)]) ->
  EvalM Env
evalResolved :: ([(VName, ([VName], TypeBase SizeClosure NoMode))],
 [(VName, SizeClosure)])
-> EvalM Env
evalResolved ([(VName, ([VName], TypeBase SizeClosure NoMode))]
ts, [(VName, SizeClosure)]
ds) = do
  ts' <- ((VName, ([VName], TypeBase SizeClosure NoMode))
 -> EvalM (VName, TypeBase SizeClosure NoMode))
-> [(VName, ([VName], TypeBase SizeClosure NoMode))]
-> EvalM [(VName, TypeBase SizeClosure NoMode)]
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 NoMode)
 -> EvalM (TypeBase SizeClosure NoMode))
-> (VName, ([VName], TypeBase SizeClosure NoMode))
-> EvalM (VName, TypeBase SizeClosure NoMode)
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 NoMode)
  -> EvalM (TypeBase SizeClosure NoMode))
 -> (VName, ([VName], TypeBase SizeClosure NoMode))
 -> EvalM (VName, TypeBase SizeClosure NoMode))
-> (([VName], TypeBase SizeClosure NoMode)
    -> EvalM (TypeBase SizeClosure NoMode))
-> (VName, ([VName], TypeBase SizeClosure NoMode))
-> EvalM (VName, TypeBase SizeClosure NoMode)
forall a b. (a -> b) -> a -> b
$ \([VName]
bound, TypeBase SizeClosure NoMode
t) -> (Either Int64 SizeClosure -> SizeClosure)
-> TypeBase (Either Int64 SizeClosure) NoMode
-> TypeBase SizeClosure NoMode
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 -> SizeClosure
forall {a}. Integral a => Either a SizeClosure -> SizeClosure
onDim (TypeBase (Either Int64 SizeClosure) NoMode
 -> TypeBase SizeClosure NoMode)
-> EvalM (TypeBase (Either Int64 SizeClosure) NoMode)
-> EvalM (TypeBase SizeClosure NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName
-> TypeBase SizeClosure NoMode
-> EvalM (TypeBase (Either Int64 SizeClosure) NoMode)
evalType ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
bound) TypeBase SizeClosure NoMode
t) [(VName, ([VName], TypeBase SizeClosure NoMode))]
ts
  ds' <- mapM (traverse $ \(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) ds
  pure $ typeEnv (M.fromList ts') <> i64Env (M.fromList ds')
  where
    onDim :: Either a SizeClosure -> SizeClosure
onDim (Left a
x) = Env -> ExpBase Info VName -> SizeClosure
SizeClosure Env
forall a. Monoid a => a
mempty (ExpBase Info VName -> SizeClosure)
-> ExpBase Info VName -> SizeClosure
forall a b. (a -> b) -> a -> b
$ 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
e) = SizeClosure
e

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 :: * -> *} {o} {a}.
TypeBase (ExpBase f VName) o -> Shape a -> Map VName a
match
  where
    match :: TypeBase (ExpBase f VName) o -> Shape a -> Map VName a
match (Scalar (Record Map Name (TypeBase (ExpBase f VName) o)
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) o -> Shape a -> Map VName a)
-> Map Name (TypeBase (ExpBase f VName) o)
-> 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) o -> Shape a -> Map VName a
match Map Name (TypeBase (ExpBase f VName) o)
poly_fields Map Name (Shape a)
fields
    match (Scalar (Sum Map Name [TypeBase (ExpBase f VName) o]
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) o] -> [Shape a] -> [Map VName a])
-> Map Name [TypeBase (ExpBase f VName) o]
-> 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) o -> Shape a -> Map VName a)
-> [TypeBase (ExpBase f VName) o] -> [Shape a] -> [Map VName a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase (ExpBase f VName) o -> Shape a -> Map VName a
match) Map Name [TypeBase (ExpBase f VName) o]
poly_fields Map Name [Shape a]
fields
    match TypeBase (ExpBase f VName) o
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) o -> Shape (ExpBase f VName)
forall dim o. TypeBase dim o -> Shape dim
arrayShape TypeBase (ExpBase f VName) o
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) o -> Shape a -> Map VName a
match (Int -> TypeBase (ExpBase f VName) o -> TypeBase (ExpBase f VName) o
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 TypeBase (ExpBase f VName) o
poly_t) Shape a
rowshape
    match TypeBase (ExpBase f VName) o
_ 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 Depth
_ 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 Depth
_ 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 Depth
_ 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 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, 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 x
x, Maybe TermBinding
_) -> x -> Maybe x
forall a. a -> Maybe a
Just x
x
        (Maybe x
Nothing, Just (TermModule (Module Env
mod))) -> Env -> [VName] -> Maybe x
f Env
mod [VName]
qs
        (Maybe x, 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 'EvalType' provided
    --  is the type of the instantiation.
    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 NoMode -> 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
  = TypeConBinding
      Env
      [TypeParam]
      (RetTypeBase Size NoMode)
  | TypeBinding
      EvalType
  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 EvalType -> Env
typeEnv :: Map VName (TypeBase SizeClosure NoMode) -> Env
typeEnv Map VName (TypeBase SizeClosure NoMode)
m =
  Env
    { envTerm :: Map VName TermBinding
envTerm = Map VName TermBinding
forall a. Monoid a => a
mempty,
      envType :: Map VName TypeBinding
envType = (TypeBase SizeClosure NoMode -> TypeBinding)
-> Map VName (TypeBase SizeClosure NoMode) -> Map VName TypeBinding
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase SizeClosure NoMode -> TypeBinding
TypeBinding Map VName (TypeBase SizeClosure NoMode)
m
    }

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) NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase (ExpBase Info VName) NoMode -> StructType)
-> ScalarTypeBase (ExpBase Info VName) NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (ExpBase Info VName) NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase (ExpBase Info VName) NoMode)
-> PrimType -> ScalarTypeBase (ExpBase Info VName) NoMode
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
  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
  liftF . ExtOpError . InterpreterError $
    "Error at\n" <> prettyStacktrace 0 ss <> s

trace :: T.Text -> Value -> EvalM ()
trace :: Text -> Value -> EvalM ()
trace Text
w Value
v = do
  -- Printing a value requires having it in full.
  v' <- Value -> EvalM Value
force Value
v
  liftF $ ExtOpTrace w (prettyValue v') ()

doIO :: IOOp -> EvalM Value
doIO :: IOOp -> EvalM Value
doIO IOOp
io = ExtOp Value -> EvalM Value
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (ExtOp Value -> EvalM Value) -> ExtOp Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ IOOp -> (Value -> Value) -> ExtOp Value
forall a. IOOp -> (Value -> a) -> ExtOp a
ExtOpIO IOOp
io Value -> Value
forall a. a -> a
id

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
  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
  backtrace <- asks ((StackFrame loc (Ctx env imports) NE.:|) . fst)
  liftF $ ExtOpBreak loc BreakPoint backtrace ()

-- | Fetch in full a value that may reside on a server, such that it can be
-- handled by code that knows only about ordinary values. Only indexing (see
-- 'indexArray') is able to work directly on a server value, so anything else
-- must force it first.
force :: Value -> EvalM Value
force :: Value -> EvalM Value
force Value
v
  | Value -> Bool
forall (m :: * -> *). Value m -> Bool
FFI.hasLazy Value
v = ExtOp Value -> EvalM Value
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (ExtOp Value -> EvalM Value) -> ExtOp Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ServerM Value -> (Value -> Value) -> ExtOp Value
forall a. ServerM Value -> (Value -> a) -> ExtOp a
ExtOpFFI (Value -> ServerM Value
forall (a :: * -> *). Value a -> ServerM (Value a)
FFI.getLazy Value
v) Value -> Value
forall a. a -> a
id
  | Bool
otherwise = Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v

fromArray :: Value -> EvalM (ValueShape, [Value])
fromArray :: Value -> EvalM (ValueShape, [Value])
fromArray Value
v =
  Value -> EvalM Value
force Value
v EvalM Value
-> (Value -> EvalM (ValueShape, [Value]))
-> EvalM (ValueShape, [Value])
forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    ValueArray ValueShape
shape Array Int Value
as -> (ValueShape, [Value]) -> EvalM (ValueShape, [Value])
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ValueShape
shape, Array Int Value -> [Value]
forall i e. Array i e -> [e]
elems Array Int Value
as)
    Value
v' -> [Char] -> EvalM (ValueShape, [Value])
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM (ValueShape, [Value]))
-> [Char] -> EvalM (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'

project :: Name -> Value -> Value
project :: Name -> Value -> Value
project Name
f (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
v'
project Name
_ Value
_ = [Char] -> Value
forall a. HasCallStack => [Char] -> a
error [Char]
"Value does not have expected field."

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
  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
  apply noLoc mempty f' y

matchPat :: Env -> Pat (TypeBase Size o) -> Value -> EvalM Env
matchPat :: forall o.
Env -> Pat (TypeBase (ExpBase Info VName) o) -> Value -> EvalM Env
matchPat Env
env Pat (TypeBase (ExpBase Info VName) o)
p Value
v = do
  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) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env Pat (TypeBase (ExpBase Info VName) o)
p Value
v
  case 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) o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (TypeBase (ExpBase Info VName) o -> StructType)
-> Pat (TypeBase (ExpBase Info VName) o)
-> PatBase Info VName StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat (TypeBase (ExpBase Info VName) o)
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 o) -> Value -> MaybeT EvalM Env
patternMatch :: forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env (PatAttr AttrInfo VName
_ PatBase Info VName (TypeBase (ExpBase Info VName) o)
p SrcLoc
_) Value
val =
  Env
-> PatBase Info VName (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env PatBase Info VName (TypeBase (ExpBase Info VName) o)
p Value
val
patternMatch Env
env (Id VName
v (Info TypeBase (ExpBase Info VName) o
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) o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) o
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) o)]
ps SrcLoc
_) (ValueRecord Map Name Value
vs) =
  (Env
 -> (PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)
 -> MaybeT EvalM Env)
-> Env
-> [(PatBase Info VName (TypeBase (ExpBase Info VName) o), 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) o)
p, Value
v) -> Env
-> PatBase Info VName (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env' PatBase Info VName (TypeBase (ExpBase Info VName) o)
p Value
v) Env
env ([(PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)]
 -> MaybeT EvalM Env)
-> [(PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)]
-> MaybeT EvalM Env
forall a b. (a -> b) -> a -> b
$
    [PatBase Info VName (TypeBase (ExpBase Info VName) o)]
-> [Value]
-> [(PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatBase Info VName (TypeBase (ExpBase Info VName) o)]
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) o))]
ps SrcLoc
_) (ValueRecord Map Name Value
vs) =
  (Env
 -> (PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)
 -> MaybeT EvalM Env)
-> Env
-> Map
     Name (PatBase Info VName (TypeBase (ExpBase Info VName) o), 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) o)
p, Value
v) -> Env
-> PatBase Info VName (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env' PatBase Info VName (TypeBase (ExpBase Info VName) o)
p Value
v) Env
env (Map
   Name (PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)
 -> MaybeT EvalM Env)
-> Map
     Name (PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)
-> MaybeT EvalM Env
forall a b. (a -> b) -> a -> b
$
    (PatBase Info VName (TypeBase (ExpBase Info VName) o)
 -> Value
 -> (PatBase Info VName (TypeBase (ExpBase Info VName) o), Value))
-> Map Name (PatBase Info VName (TypeBase (ExpBase Info VName) o))
-> Map Name Value
-> Map
     Name (PatBase Info VName (TypeBase (ExpBase Info VName) o), 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) o))]
-> Map Name (PatBase Info VName (TypeBase (ExpBase Info VName) o))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))]
 -> Map Name (PatBase Info VName (TypeBase (ExpBase Info VName) o)))
-> [(Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))]
-> Map Name (PatBase Info VName (TypeBase (ExpBase Info VName) o))
forall a b. (a -> b) -> a -> b
$ ((L Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))
 -> (Name, PatBase Info VName (TypeBase (ExpBase Info VName) o)))
-> [(L Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))]
-> [(Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))
-> (Name, PatBase Info VName (TypeBase (ExpBase Info VName) o))
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) o))]
ps) Map Name Value
vs
patternMatch Env
env (PatParens PatBase Info VName (TypeBase (ExpBase Info VName) o)
p SrcLoc
_) Value
v = Env
-> PatBase Info VName (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env PatBase Info VName (TypeBase (ExpBase Info VName) o)
p Value
v
patternMatch Env
env (PatAscription PatBase Info VName (TypeBase (ExpBase Info VName) o)
p TypeExp (ExpBase Info VName) VName
_ SrcLoc
_) Value
v =
  Env
-> PatBase Info VName (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env PatBase Info VName (TypeBase (ExpBase Info VName) o)
p Value
v
patternMatch Env
env (PatLit PatLit
l Info (TypeBase (ExpBase Info VName) o)
t SrcLoc
_) Value
v = do
  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) o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (TypeBase (ExpBase Info VName) o -> StructType)
-> Info (TypeBase (ExpBase Info VName) o) -> Info StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info (TypeBase (ExpBase Info VName) o)
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) o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (TypeBase (ExpBase Info VName) o -> StructType)
-> Info (TypeBase (ExpBase Info VName) o) -> Info StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info (TypeBase (ExpBase Info VName) o)
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 v == l'
    then pure env
    else mzero
patternMatch Env
env (PatConstr Name
n Info (TypeBase (ExpBase Info VName) o)
_ [PatBase Info VName (TypeBase (ExpBase Info VName) o)]
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) o), Value)
 -> MaybeT EvalM Env)
-> Env
-> [(PatBase Info VName (TypeBase (ExpBase Info VName) o), 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) o)
p, Value
v) -> Env
-> PatBase Info VName (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env' PatBase Info VName (TypeBase (ExpBase Info VName) o)
p Value
v) Env
env ([(PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)]
 -> MaybeT EvalM Env)
-> [(PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)]
-> MaybeT EvalM Env
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName (TypeBase (ExpBase Info VName) o)]
-> [Value]
-> [(PatBase Info VName (TypeBase (ExpBase Info VName) o), Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatBase Info VName (TypeBase (ExpBase Info VName) o)]
ps [Value]
vs
patternMatch Env
_ PatBase Info VName (TypeBase (ExpBase Info VName) o)
_ 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 -> Doc ann
forall a. Semigroup a => a -> a -> a
<>) (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

iaArrayLength :: Value -> Int64
iaArrayLength :: Value -> Int64
iaArrayLength Value
v = case Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v of
  ShapeDim Int64
n ValueShape
_ -> Int64
n
  ValueShape
_ -> [Char] -> Int64
forall a. HasCallStack => [Char] -> a
error [Char]
"Expected array."

iaRowShape :: Value -> ValueShape
iaRowShape :: Value -> ValueShape
iaRowShape Value
v = case Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v of
  ShapeDim Int64
_ ValueShape
rowshape -> ValueShape
rowshape
  ValueShape
_ -> [Char] -> ValueShape
forall a. HasCallStack => [Char] -> a
error [Char]
"Expected array."

iaIndexOnce :: Value -> Int64 -> EvalM Value
iaIndexOnce :: Value -> Int64 -> EvalM Value
iaIndexOnce (ValueArray ValueShape
_ Array Int Value
arr) Int64
i = 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
$ 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
iaIndexOnce v :: Value
v@(ValueLazyFFI ValueShape
shp ValueRef
vr [Int64]
os) Int64
i
  -- Indexing an array of arrays merely notes the index; nothing is
  -- fetched until we reach an element.
  | ShapeDim Int64
_ ShapeDim {} <- 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
forall {m :: * -> *}. Value m
v'
  | Bool
otherwise = ExtOp Value -> EvalM Value
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (ExtOp Value -> EvalM Value) -> ExtOp Value -> EvalM Value
forall a b. (a -> b) -> a -> b
$ ServerM Value -> (Value -> Value) -> ExtOp Value
forall a. ServerM Value -> (Value -> a) -> ExtOp a
ExtOpFFI ([Int64] -> ValueRef -> ServerM ValueRef
FFI.index [Int64]
is ValueRef
vr ServerM ValueRef -> (ValueRef -> ServerM Value) -> ServerM Value
forall a b. ServerM a -> (a -> ServerM b) -> ServerM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ValueShape -> ValueRef -> ServerM Value
forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
FFI.get (Value (ZonkAny 2) -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value (ZonkAny 2)
forall {m :: * -> *}. Value m
v')) Value -> Value
forall a. a -> a
id
  where
    v' :: Value m
v' = ValueShape -> ValueRef -> [Int64] -> Value m
forall (m :: * -> *). ValueShape -> ValueRef -> [Int64] -> Value m
ValueLazyFFI ValueShape
shp ValueRef
vr ([Int64] -> Value m) -> [Int64] -> Value m
forall a b. (a -> b) -> a -> b
$ Int64
i Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
: [Int64]
os
    is :: [Int64]
is = (Int64 -> Int64) -> [Int64] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map Int64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int64] -> [Int64]) -> [Int64] -> [Int64]
forall a b. (a -> b) -> a -> b
$ [Int64] -> [Int64]
forall a. [a] -> [a]
reverse ([Int64] -> [Int64]) -> [Int64] -> [Int64]
forall a b. (a -> b) -> a -> b
$ Int64
i Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
: [Int64]
os
iaIndexOnce Value
_ Int64
_ = [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error [Char]
"Expected array."

indexArray :: [Indexing] -> Value -> MaybeT EvalM Value
indexArray :: [Indexing] -> Value -> MaybeT EvalM Value
indexArray (IndexingFix Int64
i : [Indexing]
is) 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
< Value -> Int64
iaArrayLength 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 (Value -> Int64 -> EvalM Value
iaIndexOnce Value
v Int64
i) MaybeT EvalM Value
-> (Value -> MaybeT EvalM Value) -> MaybeT EvalM Value
forall a b.
MaybeT EvalM a -> (a -> MaybeT EvalM b) -> MaybeT EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Indexing] -> Value -> MaybeT EvalM Value
indexArray [Indexing]
is
  | Bool
otherwise =
      EvalM (Maybe Value) -> MaybeT EvalM Value
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (EvalM (Maybe Value) -> MaybeT EvalM Value)
-> EvalM (Maybe Value) -> MaybeT EvalM Value
forall a b. (a -> b) -> a -> b
$ Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Value
forall a. Maybe a
Nothing
indexArray (IndexingSlice Maybe Int64
start Maybe Int64
end Maybe Int64
stride : [Indexing]
is) Value
v = do
  js <- EvalM (Maybe [Int]) -> MaybeT EvalM [Int]
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (EvalM (Maybe [Int]) -> MaybeT EvalM [Int])
-> EvalM (Maybe [Int]) -> MaybeT EvalM [Int]
forall a b. (a -> b) -> a -> b
$ Maybe [Int] -> EvalM (Maybe [Int])
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [Int] -> EvalM (Maybe [Int]))
-> (Int64 -> Maybe [Int]) -> Int64 -> EvalM (Maybe [Int])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Int64 -> Maybe Int64 -> Maybe Int64 -> Int64 -> Maybe [Int]
indexesFor Maybe Int64
start Maybe Int64
end Maybe Int64
stride (Int64 -> EvalM (Maybe [Int])) -> Int64 -> EvalM (Maybe [Int])
forall a b. (a -> b) -> a -> b
$ Value -> Int64
iaArrayLength Value
v
  toArray' (indexShape is $ iaRowShape v) <$> mapM (\Int
i -> EvalM (Maybe Value) -> MaybeT EvalM Value
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> EvalM Value -> EvalM (Maybe Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Int64 -> EvalM Value
iaIndexOnce Value
v (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)) MaybeT EvalM Value
-> (Value -> MaybeT EvalM Value) -> MaybeT EvalM Value
forall a b.
MaybeT EvalM a -> (a -> MaybeT EvalM b) -> MaybeT EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Indexing] -> Value -> MaybeT EvalM Value
indexArray [Indexing]
is) js
indexArray [Indexing]
_ Value
v = Value -> MaybeT EvalM Value
forall a. a -> MaybeT EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 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
      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
      pure $ do
        v'' <- v'
        Just $ ValueArray shape $ arr // [(i', 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
            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'
            pure $ do
              x' <- x
              Just $ arr' // [(i, 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
arrayValueShape Value
arr)
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"."
  v <- 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
$ [Indexing] -> Value -> MaybeT EvalM Value
indexArray [Indexing]
is Value
arr
  maybe oob pure v

-- | Expand type based on information that was not available at
-- type-checking time (the structure of abstract types).
expandType :: (Pretty o) => Env -> TypeBase Size o -> TypeBase SizeClosure o
expandType :: forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
_ (Scalar (Prim PrimType
pt)) = ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o)
-> ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase SizeClosure o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
pt
expandType Env
env (Scalar (Record Map Name (TypeBase (ExpBase Info VName) o)
fs)) = ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o)
-> ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase SizeClosure o) -> ScalarTypeBase SizeClosure o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase SizeClosure o) -> ScalarTypeBase SizeClosure o)
-> Map Name (TypeBase SizeClosure o)
-> ScalarTypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ (TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o)
-> Map Name (TypeBase (ExpBase Info VName) o)
-> Map Name (TypeBase SizeClosure o)
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) o -> TypeBase SizeClosure o
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env) Map Name (TypeBase (ExpBase Info VName) o)
fs
expandType Env
env (Scalar (Arrow o
o PName
p Diet
d StructType
t1 (RetType [VName]
dims TypeBase (ExpBase Info VName) Freshness
t2))) =
  ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o)
-> ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ o
-> PName
-> Diet
-> TypeBase SizeClosure NoMode
-> RetTypeBase SizeClosure Freshness
-> ScalarTypeBase SizeClosure o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
o PName
p Diet
d (Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env StructType
t1) ([VName]
-> TypeBase SizeClosure Freshness
-> RetTypeBase SizeClosure Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims (Env
-> TypeBase (ExpBase Info VName) Freshness
-> TypeBase SizeClosure Freshness
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env TypeBase (ExpBase Info VName) Freshness
t2))
expandType Env
env t :: TypeBase (ExpBase Info VName) o
t@(Array o
o Shape (ExpBase Info VName)
shape ScalarTypeBase (ExpBase Info VName) NoMode
_) =
  let et :: TypeBase (ExpBase Info VName) o
et = Int
-> TypeBase (ExpBase Info VName) o
-> TypeBase (ExpBase Info VName) o
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray (Shape (ExpBase Info VName) -> Int
forall dim. Shape dim -> Int
shapeRank Shape (ExpBase Info VName)
shape) TypeBase (ExpBase Info VName) o
t
      et' :: TypeBase SizeClosure o
et' = Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env TypeBase (ExpBase Info VName) o
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 (NoMode -> o)
-> TypeBase SizeClosure NoMode -> TypeBase SizeClosure o
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 (o -> NoMode -> o
forall a b. a -> b -> a
const o
o) (Shape SizeClosure
-> TypeBase SizeClosure NoMode -> TypeBase SizeClosure NoMode
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
arrayOf Shape SizeClosure
shape' (TypeBase SizeClosure NoMode -> TypeBase SizeClosure NoMode)
-> TypeBase SizeClosure NoMode -> TypeBase SizeClosure NoMode
forall a b. (a -> b) -> a -> b
$ TypeBase SizeClosure o -> TypeBase SizeClosure NoMode
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase SizeClosure o
et')
expandType Env
env (Scalar (TypeVar o
o QualName VName
tn [TypeArg (ExpBase Info VName)]
args)) =
  case QualName VName -> Env -> Maybe TypeBinding
lookupType QualName VName
tn Env
env of
    Just (TypeBinding TypeBase SizeClosure NoMode
t') ->
      (NoMode -> o)
-> TypeBase SizeClosure NoMode -> TypeBase SizeClosure o
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 (o -> NoMode -> o
forall a b. a -> b -> a
const o
o) TypeBase SizeClosure NoMode
t'
    Just (TypeConBinding 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 -> SizeClosure)
-> ExpBase Info VName -> SizeClosure
forall a b. (a -> b) -> a -> b
$ Int -> ExpBase Info VName
anySize Int
0
            | 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 (RetTypeBase (ExpBase Info VName) NoMode))
-> Maybe (Subst (RetTypeBase (ExpBase Info VName) NoMode))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (RetTypeBase (ExpBase Info VName) NoMode))
forall {t}. Map VName (Subst t)
substs) ExpBase Info VName
dim
       in (SizeClosure -> SizeClosure)
-> (NoMode -> o)
-> TypeBase SizeClosure NoMode
-> TypeBase SizeClosure o
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 (o -> NoMode -> o
forall a b. a -> b -> a
const o
o) (TypeBase SizeClosure NoMode -> TypeBase SizeClosure o)
-> TypeBase SizeClosure NoMode -> TypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
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 o -> TypeBase SizeClosure o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (o
-> QualName VName
-> [TypeArg SizeClosure]
-> ScalarTypeBase SizeClosure o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
tn ([TypeArg SizeClosure] -> ScalarTypeBase SizeClosure o)
-> [TypeArg SizeClosure] -> ScalarTypeBase SizeClosure o
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'' :: TypeBase SizeClosure NoMode
t'' = Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env StructType
t'
       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 (TypeBase SizeClosure NoMode -> TypeBinding
TypeBinding TypeBase SizeClosure NoMode
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 NoMode -> TypeArg SizeClosure
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (TypeBase SizeClosure NoMode -> TypeArg SizeClosure)
-> TypeBase SizeClosure NoMode -> TypeArg SizeClosure
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env StructType
t
expandType Env
env (Scalar (Sum Map Name [TypeBase (ExpBase Info VName) o]
cs)) = ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o)
-> ScalarTypeBase SizeClosure o -> TypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase SizeClosure o] -> ScalarTypeBase SizeClosure o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase SizeClosure o] -> ScalarTypeBase SizeClosure o)
-> Map Name [TypeBase SizeClosure o]
-> ScalarTypeBase SizeClosure o
forall a b. (a -> b) -> a -> b
$ (([TypeBase (ExpBase Info VName) o] -> [TypeBase SizeClosure o])
-> Map Name [TypeBase (ExpBase Info VName) o]
-> Map Name [TypeBase SizeClosure o]
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) o] -> [TypeBase SizeClosure o])
 -> Map Name [TypeBase (ExpBase Info VName) o]
 -> Map Name [TypeBase SizeClosure o])
-> ((TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o)
    -> [TypeBase (ExpBase Info VName) o] -> [TypeBase SizeClosure o])
-> (TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o)
-> Map Name [TypeBase (ExpBase Info VName) o]
-> Map Name [TypeBase SizeClosure o]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o)
-> [TypeBase (ExpBase Info VName) o] -> [TypeBase SizeClosure o]
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) o -> TypeBase SizeClosure o
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env) Map Name [TypeBase (ExpBase Info VName) o]
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) NoMode)
evalType :: Set VName
-> TypeBase SizeClosure NoMode
-> EvalM (TypeBase (Either Int64 SizeClosure) NoMode)
evalType Set VName
outer_bound TypeBase SizeClosure NoMode
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 NoMode
-> EvalM (TypeBase (Either Int64 SizeClosure) NoMode)
forall (f :: * -> *) fdim tdim o.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim o -> f (TypeBase tdim o)
traverseDims Set VName
-> DimPos -> SizeClosure -> EvalM (Either Int64 SizeClosure)
forall {p}.
Set VName -> p -> SizeClosure -> EvalM (Either Int64 SizeClosure)
evalDim TypeBase SizeClosure NoMode
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 not be a
-- size-dependent function type, or one that has existentials.
evalTypeFully :: EvalType -> EvalM ValueType
evalTypeFully :: TypeBase SizeClosure NoMode -> EvalM ValueType
evalTypeFully TypeBase SizeClosure NoMode
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)
-> (NoMode -> EvalM NoMode)
-> TypeBase SizeClosure NoMode
-> 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 NoMode -> EvalM NoMode
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase SizeClosure NoMode
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 NoMode -> EvalM Value
v) -> TypeBase SizeClosure NoMode -> EvalM Value
v (TypeBase SizeClosure NoMode -> EvalM Value)
-> TypeBase SizeClosure NoMode -> EvalM Value
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
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
      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
      error $
        prettyString qv
          <> " is not bound to a value.\n"
          <> T.unpack (prettyStacktrace 0 ss)
          <> "Bound to\n"
          <> show x

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

-- | Compute the shape of a sum value with the given type, constructor
-- and payload values. The shapes of the payload values take
-- precedence over the type annotation, which may contain existential
-- sizes that cannot be evaluated. The shapes of the other
-- constructors are computed from the type, after resolving
-- existential sizes against the concrete payload where possible; the
-- type system guarantees that dimensions not determined by the
-- payload can be evaluated.
sumValueShape :: Env -> StructType -> Name -> [Value] -> EvalM ValueShape
sumValueShape :: Env -> StructType -> Name -> [Value] -> EvalM ValueShape
sumValueShape Env
env StructType
t Name
c [Value]
vs =
  case Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env StructType
t of
    Scalar (Sum Map Name [TypeBase SizeClosure NoMode]
cs) -> do
      let payload_shapes :: [ValueShape]
payload_shapes = (Value -> ValueShape) -> [Value] -> [ValueShape]
forall a b. (a -> b) -> [a] -> [b]
map Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape [Value]
vs
          learned :: Env
learned =
            Map VName Int64 -> Env
i64Env (Map VName Int64 -> Env) -> Map VName Int64 -> Env
forall a b. (a -> b) -> a -> b
$ [Map VName Int64] -> Map VName Int64
forall a. Monoid a => [a] -> a
mconcat ([Map VName Int64] -> Map VName Int64)
-> [Map VName Int64] -> Map VName Int64
forall a b. (a -> b) -> a -> b
$ [Map VName Int64]
-> ([TypeBase SizeClosure NoMode] -> [Map VName Int64])
-> Maybe [TypeBase SizeClosure NoMode]
-> [Map VName Int64]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] ((ValueShape -> TypeBase SizeClosure NoMode -> Map VName Int64)
-> [ValueShape]
-> [TypeBase SizeClosure NoMode]
-> [Map VName Int64]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ValueShape -> TypeBase SizeClosure NoMode -> Map VName Int64
forall {a} {o}. Shape a -> TypeBase SizeClosure o -> Map VName a
learn [ValueShape]
payload_shapes) (Maybe [TypeBase SizeClosure NoMode] -> [Map VName Int64])
-> Maybe [TypeBase SizeClosure NoMode] -> [Map VName Int64]
forall a b. (a -> b) -> a -> b
$ Name
-> Map Name [TypeBase SizeClosure NoMode]
-> Maybe [TypeBase SizeClosure NoMode]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name [TypeBase SizeClosure NoMode]
cs
          evalDim :: SizeClosure -> EvalM Int64
evalDim (SizeClosure Env
denv 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
learned Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
denv) ExpBase Info VName
e
          onConstr :: Name -> [TypeBase SizeClosure o] -> EvalM [ValueShape]
onConstr Name
c' [TypeBase SizeClosure o]
fts
            | Name
c' Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
c = [ValueShape] -> EvalM [ValueShape]
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ValueShape]
payload_shapes
            | Bool
otherwise = (TypeBase Int64 o -> ValueShape)
-> [TypeBase Int64 o] -> [ValueShape]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Int64 o -> ValueShape
forall d o. TypeBase d o -> Shape d
typeShape ([TypeBase Int64 o] -> [ValueShape])
-> EvalM [TypeBase Int64 o] -> EvalM [ValueShape]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase SizeClosure o -> EvalM (TypeBase Int64 o))
-> [TypeBase SizeClosure o] -> EvalM [TypeBase Int64 o]
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)
-> (o -> EvalM o)
-> TypeBase SizeClosure o
-> EvalM (TypeBase Int64 o)
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 o -> EvalM o
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [TypeBase SizeClosure o]
fts
      Map Name [ValueShape] -> ValueShape
forall d. Map Name [Shape d] -> Shape d
ShapeSum (Map Name [ValueShape] -> ValueShape)
-> EvalM (Map Name [ValueShape]) -> EvalM ValueShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> [TypeBase SizeClosure NoMode] -> EvalM [ValueShape])
-> Map Name [TypeBase SizeClosure NoMode]
-> EvalM (Map Name [ValueShape])
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
M.traverseWithKey Name -> [TypeBase SizeClosure NoMode] -> EvalM [ValueShape]
forall {o}. Name -> [TypeBase SizeClosure o] -> EvalM [ValueShape]
onConstr Map Name [TypeBase SizeClosure NoMode]
cs
    TypeBase SizeClosure NoMode
t' -> ValueType -> ValueShape
forall d o. TypeBase d o -> Shape d
typeShape (ValueType -> ValueShape) -> EvalM ValueType -> EvalM ValueShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase SizeClosure NoMode -> EvalM ValueType
evalTypeFully TypeBase SizeClosure NoMode
t'
  where
    learn :: Shape a -> TypeBase SizeClosure o -> Map VName a
learn (ShapeDim a
d Shape a
s) TypeBase SizeClosure o
ft
      | SizeClosure Env
_ (Var QualName VName
v Info StructType
_ SrcLoc
_) : [SizeClosure]
_ <- Shape SizeClosure -> [SizeClosure]
forall dim. Shape dim -> [dim]
shapeDims (TypeBase SizeClosure o -> Shape SizeClosure
forall dim o. TypeBase dim o -> Shape dim
arrayShape TypeBase SizeClosure o
ft) =
          VName -> a -> Map VName a -> Map VName a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) a
d (Map VName a -> Map VName a) -> Map VName a -> Map VName a
forall a b. (a -> b) -> a -> b
$ Shape a -> TypeBase SizeClosure o -> Map VName a
learn Shape a
s (Int -> TypeBase SizeClosure o -> TypeBase SizeClosure o
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 TypeBase SizeClosure o
ft)
      | Bool
otherwise = Shape a -> TypeBase SizeClosure o -> Map VName a
learn Shape a
s (Int -> TypeBase SizeClosure o -> TypeBase SizeClosure o
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 TypeBase SizeClosure o
ft)
    learn (ShapeRecord Map Name (Shape a)
fs) (Scalar (Record Map Name (TypeBase SizeClosure o)
fts)) =
      [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
$ (Shape a -> TypeBase SizeClosure o -> Map VName a)
-> Map Name (Shape a)
-> Map Name (TypeBase SizeClosure o)
-> 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 Shape a -> TypeBase SizeClosure o -> Map VName a
learn Map Name (Shape a)
fs Map Name (TypeBase SizeClosure o)
fts
    learn (ShapeSum Map Name [Shape a]
fs) (Scalar (Sum Map Name [TypeBase SizeClosure o]
fts)) =
      [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
$ ([Shape a] -> [TypeBase SizeClosure o] -> [Map VName a])
-> Map Name [Shape a]
-> Map Name [TypeBase SizeClosure o]
-> 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 ((Shape a -> TypeBase SizeClosure o -> Map VName a)
-> [Shape a] -> [TypeBase SizeClosure o] -> [Map VName a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Shape a -> TypeBase SizeClosure o -> Map VName a
learn) Map Name [Shape a]
fs Map Name [TypeBase SizeClosure o]
fts
    learn Shape a
_ TypeBase SizeClosure o
_ = Map VName a
forall a. Monoid a => a
mempty

-- 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 o) -> Value -> Env -> Env
linkMissingSizes :: forall o.
[VName]
-> Pat (TypeBase (ExpBase Info VName) o) -> Value -> Env -> Env
linkMissingSizes [] Pat (TypeBase (ExpBase Info VName) o)
_ Value
_ Env
env = Env
env
linkMissingSizes [VName]
missing_sizes Pat (TypeBase (ExpBase Info VName) o)
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 NoMode -> StructType
evalToStruct (TypeBase SizeClosure NoMode -> StructType)
-> TypeBase SizeClosure NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env (StructType -> TypeBase SizeClosure NoMode)
-> StructType -> TypeBase SizeClosure NoMode
forall a b. (a -> b) -> a -> b
$ Pat (TypeBase (ExpBase Info VName) o) -> StructType
forall o. Pat (TypeBase (ExpBase Info VName) o) -> StructType
patternStructType Pat (TypeBase (ExpBase Info VName) o)
p

evalBinding :: 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.
evalBinding :: Env
-> [VName]
-> [Pat ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Freshness
-> EvalM Value
evalBinding Env
env [VName]
missing_sizes [] ExpBase Info VName
body TypeBase (ExpBase Info VName) Freshness
rettype =
  -- Eta-expand the rest to make any sizes visible.
  [Value]
-> Env -> TypeBase (ExpBase Info VName) Freshness -> EvalM Value
etaExpand [] Env
env TypeBase (ExpBase Info VName) Freshness
rettype
  where
    etaExpand :: [Value]
-> Env -> TypeBase (ExpBase Info VName) Freshness -> EvalM Value
etaExpand [Value]
vs Env
env' (Scalar (Arrow Freshness
_ PName
_ Diet
_ StructType
p_t (RetType [VName]
_ TypeBase (ExpBase Info VName) Freshness
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'' <- [VName] -> PatBase Info VName StructType -> Value -> Env -> Env
forall o.
[VName]
-> Pat (TypeBase (ExpBase Info VName) o) -> 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 o.
Env -> Pat (TypeBase (ExpBase Info VName) o) -> Value -> EvalM Env
matchPat Env
env' PatBase Info VName StructType
forall {vn}. PatBase Info vn StructType
p Value
v
        etaExpand (v : vs) env'' rt
    etaExpand [Value]
vs Env
env' TypeBase (ExpBase Info VName) Freshness
_ = do
      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
      foldM (apply noLoc mempty) f $ reverse vs
evalBinding Env
env [VName]
missing_sizes (Pat ParamType
p : [Pat ParamType]
ps) ExpBase Info VName
body TypeBase (ExpBase Info VName) Freshness
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' <- [VName] -> Pat ParamType -> Value -> Env -> Env
forall o.
[VName]
-> Pat (TypeBase (ExpBase Info VName) o) -> 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 o.
Env -> Pat (TypeBase (ExpBase Info VName) o) -> Value -> EvalM Env
matchPat Env
env Pat ParamType
p Value
v
    evalBinding env' missing_sizes ps body rettype

evalValBinding ::
  Env ->
  VName ->
  [TypeParam] ->
  [Pat ParamType] ->
  ResRetType ->
  Exp ->
  EvalM TermBinding
evalValBinding :: Env
-> VName
-> [TypeParam]
-> [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness
-> ExpBase Info VName
-> EvalM TermBinding
evalValBinding Env
env VName
name [TypeParam]
tparams [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret ExpBase Info VName
fbody = do
  let ftype :: StructType
ftype = TypeBase SizeClosure NoMode -> StructType
evalToStruct (TypeBase SizeClosure NoMode -> StructType)
-> TypeBase SizeClosure NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env (StructType -> TypeBase SizeClosure NoMode)
-> StructType -> TypeBase SizeClosure NoMode
forall a b. (a -> b) -> a -> b
$ [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret
      retext :: [VName]
retext = case [Pat ParamType]
ps of
        [] -> RetTypeBase (ExpBase Info VName) Freshness -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims RetTypeBase (ExpBase Info VName) Freshness
ret
        [Pat ParamType]
_ -> []

  -- A function binding may refer to itself recursively, which means it muse be
  -- in scope of itself. We tie the knot by evaluating the body in an
  -- environment ('recenv') that maps the 'name' to itself. Crucially, in each
  -- case below the binding is a lazy, self-referential 'let' whose right-hand
  -- side is a value constructor (a 'ValueFun' closure or a 'TermPoly' closure).
  --
  -- Distinguish polymorphic, monomorphic-function, and plain-value
  -- bindings here.
  case ([TypeParam] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParam]
tparams, [Pat ParamType]
ps) of
    (Bool
True, []) ->
      -- Not a syntactic function, so cannot be recursive.
      (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) Freshness
-> [VName]
-> Value
-> EvalM Value
forall o.
Env
-> TypeBase (ExpBase Info VName) o
-> [VName]
-> Value
-> EvalM Value
returned Env
env (RetTypeBase (ExpBase Info VName) Freshness
-> TypeBase (ExpBase Info VName) Freshness
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType RetTypeBase (ExpBase Info VName) Freshness
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) Freshness
-> EvalM Value
evalBinding Env
env [] [Pat ParamType]
ps ExpBase Info VName
fbody (RetTypeBase (ExpBase Info VName) Freshness
-> TypeBase (ExpBase Info VName) Freshness
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType RetTypeBase (ExpBase Info VName) Freshness
ret)
    (Bool
True, Pat ParamType
p : [Pat ParamType]
ps') ->
      -- A monomorphic function; "unfold" the first step of what evalBinding
      -- would do.
      let recenv :: Env
recenv = Env
env {envTerm = M.insert name binding $ envTerm env}
          binding :: TermBinding
binding =
            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) (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
v -> do
              env' <- Env -> Pat ParamType -> Value -> EvalM Env
forall o.
Env -> Pat (TypeBase (ExpBase Info VName) o) -> Value -> EvalM Env
matchPat Env
recenv Pat ParamType
p Value
v
              evalBinding env' [] ps' fbody (retType ret)
       in TermBinding -> EvalM TermBinding
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TermBinding
binding
    (Bool
False, [Pat ParamType]
_) ->
      -- A polymorphic function.
      let binding :: TermBinding
binding = Maybe BoundV
-> (TypeBase SizeClosure NoMode -> 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 NoMode -> EvalM Value) -> TermBinding)
-> (TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \TypeBase SizeClosure NoMode
ftype' -> do
            let resolved :: ([(VName, ([VName], TypeBase SizeClosure NoMode))],
 [(VName, SizeClosure)])
resolved =
                  [VName]
-> StructType
-> TypeBase SizeClosure NoMode
-> ([(VName, ([VName], TypeBase SizeClosure NoMode))],
    [(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 NoMode
ftype'
            tparam_env <- ([(VName, ([VName], TypeBase SizeClosure NoMode))],
 [(VName, SizeClosure)])
-> EvalM Env
evalResolved ([(VName, ([VName], TypeBase SizeClosure NoMode))],
 [(VName, SizeClosure)])
resolved
            let recenv = Env
env {envTerm = M.insert name binding $ envTerm env}
                env' = Env
tparam_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
recenv
                -- 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 -> 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)
            returned env (retType ret) retext
              =<< evalBinding env' missing_sizes ps fbody (retType ret)
       in TermBinding -> EvalM TermBinding
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TermBinding
binding

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
  v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  case 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 ()
  pure v

returned :: Env -> TypeBase Size o -> [VName] -> Value -> EvalM Value
returned :: forall o.
Env
-> TypeBase (ExpBase Info VName) o
-> [VName]
-> Value
-> EvalM Value
returned Env
_ TypeBase (ExpBase Info VName) o
_ [] 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) o
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 NoMode -> StructType
evalToStruct (TypeBase SizeClosure NoMode -> StructType)
-> TypeBase SizeClosure NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env (StructType -> TypeBase SizeClosure NoMode)
-> StructType -> TypeBase SizeClosure NoMode
forall a b. (a -> b) -> a -> b
$ TypeBase (ExpBase Info VName) o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) o
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
  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_second' <- traverse (fmap asInteger . eval env) maybe_second
  end' <- traverse (fmap asInteger . eval env) end

  let (end_adj, step, ok) =
        case (end', 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 ok
    then pure $ toArray' ShapeLeaf $ map toInt [start', start' + step .. end_adj]
    else bad loc env $ badRange start' maybe_second' 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
  v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  env' <- matchPat env p v
  let p_t = TypeBase SizeClosure NoMode -> StructType
evalToStruct (TypeBase SizeClosure NoMode -> StructType)
-> TypeBase SizeClosure NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env (StructType -> TypeBase SizeClosure NoMode)
-> StructType -> TypeBase SizeClosure NoMode
forall a b. (a -> b) -> a -> b
$ PatBase Info VName StructType -> StructType
forall o. Pat (TypeBase (ExpBase Info VName) o) -> StructType
patternStructType PatBase Info VName StructType
p
      v_s = Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v
      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)
  eval env'' body
evalAppExp Env
env (LetFun (VName
f, SrcLoc
_) ([TypeParam]
tparams, [Pat ParamType]
ps, Maybe (TypeExp (ExpBase Info VName) VName)
_, Info RetTypeBase (ExpBase Info VName) Freshness
ret, ExpBase Info VName
fbody) ExpBase Info VName
body SrcLoc
_) = do
  binding <- Env
-> VName
-> [TypeParam]
-> [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness
-> ExpBase Info VName
-> EvalM TermBinding
evalValBinding Env
env VName
f [TypeParam]
tparams [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret ExpBase Info VName
fbody
  eval (env {envTerm = M.insert f binding $ envTerm env}) 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 -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"&&" = do
      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 x'
        then eval env y
        else pure $ ValuePrim $ BoolValue False
  | VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"||" = do
      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 x'
        then pure $ ValuePrim $ BoolValue True
        else eval env y
  | Bool
otherwise = do
      x' <- Env -> ExpBase Info VName -> Maybe VName -> EvalM Value
evalArg Env
env ExpBase Info VName
x Maybe VName
xext
      y' <- evalArg env y yext
      op' <- evalTermVar env op op_t
      op'' <- apply loc env op' x'
      apply loc env op'' y'
evalAppExp Env
env (If ExpBase Info VName
cond ExpBase Info VName
e1 ExpBase Info VName
e2 SrcLoc
_) = do
  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 cond' then eval env e1 else eval env 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.
  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)
  f' <- eval env f
  foldM (apply loc env) f' 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
  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
  arr <- eval env e
  evalIndex loc env is' arr
evalAppExp Env
env (LetWith IdentBase Info VName StructType
dest IdentBase Info VName StructType
src [UpdateStep Info VName]
steps 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
  src_v <- 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 o. TypeBase dim o -> TypeBase dim NoMode
toStruct StructType
src_t)
  v' <- eval env v
  dest' <- maybe oob pure =<< evalUpdateSteps env steps src_v v'
  let t = [TypeParam] -> StructType -> BoundV
T.BoundV [] (StructType -> BoundV) -> StructType -> BoundV
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
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
  eval (valEnv (M.singleton (identName dest) (Just t, dest')) <> env) 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
  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 form of
    For IdentBase Info VName StructType
iv ExpBase Info VName
bound -> do
      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
      forLoop (identName iv) bound' (zero bound') init_v
    ForIn PatBase Info VName StructType
in_pat ExpBase Info VName
in_e -> do
      (_, in_vs) <- Value -> EvalM (ValueShape, [Value])
fromArray (Value -> EvalM (ValueShape, [Value]))
-> EvalM Value -> EvalM (ValueShape, [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
in_e
      foldM (forInLoop in_pat) init_v 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 o. Pat (TypeBase (ExpBase Info VName) o) -> StructType
patternStructType Pat ParamType
pat) (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v)
       in Env -> Pat ParamType -> Value -> EvalM Env
forall o.
Env -> Pat (TypeBase (ExpBase Info VName) o) -> 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) NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase (ExpBase Info VName) NoMode -> StructType)
-> ScalarTypeBase (ExpBase Info VName) NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (ExpBase Info VName) NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase (ExpBase Info VName) NoMode)
-> PrimType -> ScalarTypeBase (ExpBase Info VName) NoMode
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' <- Value -> EvalM Env
withLoopParams Value
v
          forLoop iv bound (inc i) =<< evalBody (forLoopEnv iv i <> env')

    whileLoop :: ExpBase Info VName -> Value -> EvalM Value
whileLoop ExpBase Info VName
cond Value
v = do
      env' <- Value -> EvalM Env
withLoopParams Value
v
      continue <- asBool <$> eval env' cond
      if continue
        then whileLoop cond =<< evalBody env'
        else pure v

    forInLoop :: Pat (TypeBase (ExpBase Info VName) o)
-> Value -> Value -> EvalM Value
forInLoop Pat (TypeBase (ExpBase Info VName) o)
in_pat Value
v Value
in_v = do
      env' <- Value -> EvalM Env
withLoopParams Value
v
      env'' <- matchPat env' in_pat in_v
      evalBody env''
evalAppExp Env
env (Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs SrcLoc
_) = do
  v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  match v (NE.toList 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
      c' <- Value -> Env -> CaseBase Info VName -> EvalM (Maybe Value)
evalCase Value
v Env
env CaseBase Info VName
c
      case 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, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_) =
  -- Already handled by the names added by the type checker.
  Env -> ExpBase Info VName -> EvalM Value
eval 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
      v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
      pure (k, v)
    evalField (RecordFieldImplicit (L Loc
_ VName
k) Info StructType
t SrcLoc
loc) = do
      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
      pure (baseName k, 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
  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 o. TypeBase dim o -> TypeBase dim NoMode
toStruct StructType
t
  pure $ toArray t' []
eval Env
env (ArrayLit (ExpBase Info VName
v : [ExpBase Info VName]
vs) Info StructType
_ SrcLoc
_) = do
  v' <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
v
  vs' <- mapM (eval env) vs
  pure $ toArray' (valueShape v') (v' : 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
  v <- Env -> AppExp -> EvalM Value
evalAppExp Env
env AppExp
e
  returned env (toStruct t) retext 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 o. TypeBase dim o -> TypeBase dim NoMode
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
  v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  t' <- evalTypeFully $ expandType env t
  case checkShape (typeShape t') (valueShape 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 (ZonkAny 1) -> Text) -> Doc (ZonkAny 1) -> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 1) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 1) -> EvalM Value) -> Doc (ZonkAny 1) -> EvalM Value
forall a b. (a -> b) -> a -> b
$
        Doc (ZonkAny 1)
"Value"
          Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann
dquotes (Value -> Doc (ZonkAny 1)
forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
v)
          Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 1)
"of shape"
          Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann
dquotes (ValueShape -> Doc (ZonkAny 1)
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 (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 1)
"cannot match shape of type"
          Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann
dquotes (TypeExp (ExpBase Info VName) VName -> Doc (ZonkAny 1)
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 (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann
parens (Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann
dquotes (ValueType -> Doc (ZonkAny 1)
forall a ann. Pretty a => a -> Doc ann
forall ann. ValueType -> Doc ann
pretty ValueType
t'))
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
  ev <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  apply loc env intrinsicsNeg 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 [UpdateStep Info VName]
steps ExpBase Info VName
v Info StructType
_ SrcLoc
loc) = do
  src' <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
src
  v' <- eval env v
  res <- evalUpdateSteps env steps src' v'
  maybe oob pure res
  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"
-- 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) Freshness
rt)) SrcLoc
_) =
  Env
-> [VName]
-> [Pat ParamType]
-> ExpBase Info VName
-> TypeBase (ExpBase Info VName) Freshness
-> EvalM Value
evalBinding Env
env [] [Pat ParamType]
ps ExpBase Info VName
body TypeBase (ExpBase Info VName) Freshness
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 o. TypeBase dim o -> TypeBase dim NoMode
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) Freshness
t), Info [VName]
_) SrcLoc
loc) = do
  v <- Env -> ExpBase Info VName -> Maybe VName -> EvalM Value
evalArg Env
env ExpBase Info VName
e Maybe VName
argext
  f <- evalTermVar env qv t'
  apply loc env f v
  where
    t' :: StructType
t' = TypeBase (ExpBase Info VName) Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) Freshness
t
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) Freshness
t)) SrcLoc
loc) = do
  y <- Env -> ExpBase Info VName -> Maybe VName -> EvalM Value
evalArg Env
env ExpBase Info VName
e Maybe VName
argext
  pure $
    ValueFun $ \Value
x -> do
      f <- Env -> QualName VName -> StructType -> EvalM Value
evalTermVar Env
env QualName VName
qv StructType
t'
      f' <- apply loc env f x
      apply loc env f' y
  where
    t' :: StructType
t' = TypeBase (ExpBase Info VName) Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) Freshness
t
eval Env
env (UpdateSection [UpdateStep Info VName]
steps Info StructType
_ SrcLoc
loc) =
  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
$ [UpdateStep Info VName] -> Value -> EvalM Value
evalSection [UpdateStep Info VName]
steps
  where
    evalSection :: [UpdateStep Info VName] -> Value -> EvalM Value
evalSection [] Value
v = Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v
    evalSection (UpdateStepField Name
f : [UpdateStep Info VName]
rest) (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 =
          [UpdateStep Info VName] -> Value -> EvalM Value
evalSection [UpdateStep Info VName]
rest Value
v'
    evalSection (UpdateStepField Name
_ : [UpdateStep Info VName]
_) Value
_ =
      [Char] -> EvalM Value
forall a. HasCallStack => [Char] -> a
error [Char]
"Value does not have expected field."
    evalSection (UpdateStepSlice SliceBase Info VName
is : [UpdateStep Info VName]
rest) Value
arr = do
      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
      evalIndex loc env is' arr >>= evalSection rest
eval Env
env (Project Name
f ExpBase Info VName
e Info StructType
_ SrcLoc
_) = do
  Name -> Value -> Value
project Name
f (Value -> Value) -> EvalM Value -> EvalM 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
e
eval Env
env (Assert ExpBase Info VName
what ExpBase Info VName
e (Info Text
s) SrcLoc
loc) = do
  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
  unless cond $ bad loc env s
  eval env e
eval Env
env (Constr Name
c [ExpBase Info VName]
es (Info StructType
t) SrcLoc
_) = do
  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
  shape <- sumValueShape env (toStruct t) c vs
  pure $ ValueSum shape c 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
  v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  trace (locText (locOf loc)) v
  pure v
eval Env
env (Attr (AttrComp Name
"trace" [AttrAtom (AtomName Name
tag) SrcLoc
_] SrcLoc
_) ExpBase Info VName
e SrcLoc
_) = do
  v <- Env -> ExpBase Info VName -> EvalM Value
eval Env
env ExpBase Info VName
e
  trace (nameToText tag) v
  pure 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

evalUpdateSteps :: Env -> [UpdateStep Info VName] -> Value -> Value -> EvalM (Maybe Value)
evalUpdateSteps :: Env
-> [UpdateStep Info VName] -> Value -> Value -> EvalM (Maybe Value)
evalUpdateSteps Env
env = [UpdateStep Info VName] -> Value -> Value -> EvalM (Maybe Value)
go
  where
    -- The value we update with must be in full, as 'writeArray' will store
    -- it directly in the destination.
    go :: [UpdateStep Info VName] -> Value -> Value -> EvalM (Maybe Value)
go [] Value
_ Value
newv = Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> EvalM Value -> EvalM (Maybe Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM Value
force Value
newv
    go (UpdateStepField Name
f : [UpdateStep Info VName]
rest) (ValueRecord Map Name Value
fs) Value
newv
      | Just Value
old <- 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 = do
          newf <- [UpdateStep Info VName] -> Value -> Value -> EvalM (Maybe Value)
go [UpdateStep Info VName]
rest Value
old Value
newv
          pure $ fmap (\Value
v' -> Map Name Value -> Value
forall (m :: * -> *). Map Name (Value m) -> Value m
ValueRecord (Map Name Value -> Value) -> Map Name Value -> Value
forall a b. (a -> b) -> a -> b
$ Name -> Value -> Map Name Value -> Map Name Value
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
f Value
v' Map Name Value
fs) newf
    go (UpdateStepField Name
_ : [UpdateStep Info VName]
_) Value
_ Value
_ =
      [Char] -> EvalM (Maybe Value)
forall a. HasCallStack => [Char] -> a
error [Char]
"eval update: invalid field update."
    go (UpdateStepSlice SliceBase Info VName
is : [UpdateStep Info VName]
rest) Value
arr Value
newv = do
      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
      -- 'writeArray' cannot update an array residing on a server.
      arr' <- force arr
      v <- runMaybeT $ indexArray is' arr'
      case v of
        Maybe Value
Nothing -> Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Value
forall a. Maybe a
Nothing
        Just Value
old -> do
          newsub <- [UpdateStep Info VName] -> Value -> Value -> EvalM (Maybe Value)
go [UpdateStep Info VName]
rest Value
old Value
newv
          case newsub of
            Maybe Value
Nothing -> Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Value
forall a. Maybe a
Nothing
            Just Value
vsub -> Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Value -> EvalM (Maybe Value))
-> Maybe Value -> EvalM (Maybe Value)
forall a b. (a -> b) -> a -> b
$ [Indexing] -> Value -> Value -> Maybe Value
writeArray [Indexing]
is' Value
arr' Value
vsub

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 -> PatBase Info VName StructType -> Value -> MaybeT EvalM Env
forall o.
Env
-> Pat (TypeBase (ExpBase Info VName) o)
-> Value
-> MaybeT EvalM Env
patternMatch Env
env PatBase Info VName StructType
p Value
v
  lift $ eval env' 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
      (k, v) <- Map VName t -> [(VName, t)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName t
m
      k' <- replace k
      pure (k', f 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
forall a. a -> a
id 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 NoMode -> EvalM Value
v) = Maybe BoundV
-> (TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding
TermPoly Maybe BoundV
t TypeBase SizeClosure NoMode -> 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

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
  f' <- ImportName -> EvalM (Maybe Env)
lookupImport ImportName
f
  known <- asks snd
  case 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
  (_, built_env) <- Env -> [DecBase Info VName] -> EvalM (Env, Env)
evalDecs Env
env [DecBase Info VName]
ds
  pure (built_env, Module built_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
  (f_env, f') <- Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env ModExp
f
  (e_env, e') <- evalModExp env e
  case f' of
    ModuleFun Module -> EvalM Module
f'' -> do
      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 = case Module
res_mod of
            Module Env
x -> Env
x
            Module
_ -> Env
forall a. Monoid a => a
mempty
      -- The following environment handles the case where rsubst refers to names
      -- that are not actually defined in the module itself, but merely
      -- inherited from an outer environment (see #2273).
      let env_substs = (Map VName TermBinding -> Map VName TypeBinding -> Env
`Env` Map VName TypeBinding
forall a. Monoid a => a
mempty) (Map VName TermBinding -> Env) -> Map VName TermBinding -> Env
forall a b. (a -> b) -> a -> b
$ [(VName, TermBinding)] -> Map VName TermBinding
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, TermBinding)] -> Map VName TermBinding)
-> [(VName, TermBinding)] -> Map VName TermBinding
forall a b. (a -> b) -> a -> b
$ do
            (to, from) <- Map VName VName -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName VName
rsubst
            x <- maybeToList $ M.lookup from $ envTerm env
            pure (to, x)
      pure (f_env <> e_env <> res_env <> env_substs, res_mod)
    Module
_ -> [Char] -> EvalM (Env, Module)
forall a. HasCallStack => [Char] -> a
error [Char]
"Expected ModuleFun."

-- | Call an external function with the given return type.
extCall :: Name -> ResRetType -> Env -> [Value] -> EvalM Value
extCall :: Name
-> RetTypeBase (ExpBase Info VName) Freshness
-> Env
-> [Value]
-> EvalM Value
extCall Name
n RetTypeBase (ExpBase Info VName) Freshness
ret Env
env [Value]
ps = do
  v <- ExtOp Value -> EvalM Value
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (ExtOp Value -> EvalM Value)
-> (ResShape -> ExtOp Value) -> ResShape -> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ServerM Value -> (Value -> Value) -> ExtOp Value
forall a. ServerM Value -> (Value -> a) -> ExtOp a
`ExtOpFFI` Value -> Value
forall a. a -> a
id) (ServerM Value -> ExtOp Value)
-> (ResShape -> ServerM Value) -> ResShape -> ExtOp Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ResShape -> ServerM Value
forall {m :: * -> *}. ResShape -> ServerM (Value m)
call (ResShape -> EvalM Value) -> EvalM ResShape -> EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [VName] -> Env -> StructType -> EvalM ResShape
resShapeOf [VName]
ext Env
env StructType
t
  -- An existential size that the server cannot reveal is one that occurs only
  -- in the payload of an absent constructor. It is observable by binding it in
  -- a size pattern, so this differs from the interpreter, which knows the
  -- actual size. Like the C API, we use 0 for the arrays of an absent
  -- constructor.
  let ext_env =
        Map VName Int64 -> Env
i64Env (Map VName Int64 -> Env) -> Map VName Int64 -> Env
forall a b. (a -> b) -> a -> b
$
          [VName] -> StructType -> ValueShape -> Map VName Int64
resolveExistentials [VName]
ext StructType
t (Value -> ValueShape
forall (m :: * -> *). Value m -> ValueShape
valueShape Value
v) Map VName Int64 -> Map VName Int64 -> Map VName Int64
forall a. Semigroup a => a -> a -> a
<> [(VName, Int64)] -> Map VName Int64
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ((VName -> (VName, Int64)) -> [VName] -> [(VName, Int64)]
forall a b. (a -> b) -> [a] -> [b]
map (,Int64
0) [VName]
ext)
  completeSums (ext_env <> env) t v
  where
    ext :: [VName]
ext = RetTypeBase (ExpBase Info VName) Freshness -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims RetTypeBase (ExpBase Info VName) Freshness
ret
    t :: StructType
t = StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ ([(PName, ParamType)], StructType) -> StructType
forall a b. (a, b) -> b
snd (([(PName, ParamType)], StructType) -> StructType)
-> ([(PName, ParamType)], StructType) -> StructType
forall a b. (a -> b) -> a -> b
$ TypeBase (ExpBase Info VName) Freshness
-> ([(PName, ParamType)], StructType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType (TypeBase (ExpBase Info VName) Freshness
 -> ([(PName, ParamType)], StructType))
-> TypeBase (ExpBase Info VName) Freshness
-> ([(PName, ParamType)], StructType)
forall a b. (a -> b) -> a -> b
$ RetTypeBase (ExpBase Info VName) Freshness
-> TypeBase (ExpBase Info VName) Freshness
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType RetTypeBase (ExpBase Info VName) Freshness
ret
    call :: ResShape -> ServerM (Value m)
call ResShape
resshp = do
      ServerM ()
FFI.gc
      vr <- Name -> [ValueRef] -> ServerM ValueRef
FFI.call Name
n ([ValueRef] -> ServerM ValueRef)
-> ServerM [ValueRef] -> ServerM ValueRef
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> [Value] -> ServerM [ValueRef]
forall (m :: * -> *). Name -> [Value m] -> ServerM [ValueRef]
FFI.putArgs Name
n [Value]
ps
      shp <- FFI.resultShape resshp vr
      FFI.lazyGet shp vr

-- | Describe a result type for 'FFI.resultShape'. Arrays, records, and sums
-- with existential sizes are merely described, never evaluated; only what is
-- none of these has its shape taken from the type.
resShapeOf :: [VName] -> Env -> StructType -> EvalM FFI.ResShape
resShapeOf :: [VName] -> Env -> StructType -> EvalM ResShape
resShapeOf [VName]
ext Env
env StructType
t
  | Int
rank Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 =
      ResShape -> ResShape
FFI.ResArray (ResShape -> ResShape) -> EvalM ResShape -> EvalM ResShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VName] -> Env -> StructType -> EvalM ResShape
resShapeOf [VName]
ext Env
env (Int -> StructType -> StructType
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
rank StructType
t)
  | Scalar (Record Map Name StructType
fs) <- StructType
t =
      Map Name ResShape -> ResShape
FFI.ResRecord (Map Name ResShape -> ResShape)
-> EvalM (Map Name ResShape) -> EvalM ResShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (StructType -> EvalM ResShape)
-> Map Name StructType -> EvalM (Map Name ResShape)
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) -> Map Name a -> f (Map Name b)
traverse ([VName] -> Env -> StructType -> EvalM ResShape
resShapeOf [VName]
ext Env
env) Map Name StructType
fs
  | Scalar (Sum Map Name [StructType]
cs) <- StructType
t,
    (VName -> Bool) -> [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` FV -> Set VName
fvVars (StructType -> FV
forall o. TypeBase (ExpBase Info VName) o -> FV
freeInType StructType
t)) [VName]
ext =
      Map Name [ResShape] -> ResShape
FFI.ResSum (Map Name [ResShape] -> ResShape)
-> EvalM (Map Name [ResShape]) -> EvalM ResShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([StructType] -> EvalM [ResShape])
-> Map Name [StructType] -> EvalM (Map Name [ResShape])
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) -> Map Name a -> f (Map Name b)
traverse ((StructType -> EvalM ResShape) -> [StructType] -> EvalM [ResShape]
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) -> [a] -> f [b]
traverse ([VName] -> Env -> StructType -> EvalM ResShape
resShapeOf [VName]
ext Env
env)) Map Name [StructType]
cs
  | Bool
otherwise =
      ValueShape -> ResShape
FFI.ResKnown (ValueShape -> ResShape)
-> (ValueType -> ValueShape) -> ValueType -> ResShape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueType -> ValueShape
forall d o. TypeBase d o -> Shape d
typeShape (ValueType -> ResShape) -> EvalM ValueType -> EvalM ResShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase SizeClosure NoMode -> EvalM ValueType
evalTypeFully (Env -> StructType -> TypeBase SizeClosure NoMode
structToEval Env
env StructType
t)
  where
    rank :: Int
rank = StructType -> Int
forall d o. TypeBase d o -> Int
arrayRank StructType
t

-- | The shape of a sum produced by 'FFI.resultShape' describes only the
-- constructor that is present. Add the others, whose shapes are determined by
-- the type in the given environment.
completeSums :: Env -> StructType -> Value -> EvalM Value
completeSums :: Env -> StructType -> Value -> EvalM Value
completeSums Env
env (Scalar (Record Map Name StructType
ts)) (ValueRecord Map Name Value
vs) =
  Map Name Value -> Value
forall (m :: * -> *). Map Name (Value m) -> Value m
ValueRecord (Map Name Value -> Value) -> EvalM (Map Name Value) -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (EvalM Value) -> EvalM (Map Name Value)
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 ((StructType -> Value -> EvalM Value)
-> Map Name StructType -> Map Name Value -> Map Name (EvalM Value)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (Env -> StructType -> Value -> EvalM Value
completeSums Env
env) Map Name StructType
ts Map Name Value
vs)
completeSums Env
env t :: StructType
t@(Scalar (Sum Map Name [StructType]
cs)) (ValueSum ValueShape
_ Name
c [Value]
vs) = do
  vs' <- EvalM [Value]
-> ([StructType] -> EvalM [Value])
-> Maybe [StructType]
-> EvalM [Value]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Value] -> EvalM [Value]
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Value]
vs) (\[StructType]
ts -> (StructType -> Value -> EvalM Value)
-> [StructType] -> [Value] -> EvalM [Value]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (Env -> StructType -> Value -> EvalM Value
completeSums Env
env) [StructType]
ts [Value]
vs) (Maybe [StructType] -> EvalM [Value])
-> Maybe [StructType] -> EvalM [Value]
forall a b. (a -> b) -> a -> b
$ Name -> Map Name [StructType] -> Maybe [StructType]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name [StructType]
cs
  shape <- sumValueShape env t c vs'
  pure $ ValueSum shape c vs'
completeSums Env
_ StructType
_ Value
v = Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
v

-- | The parameters of an external binding, represented as functions that
-- bind an argument in the environment.
extParams :: [Pat ParamType] -> ResRetType -> [Env -> Value -> EvalM Env]
extParams :: [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness
-> [Env -> Value -> EvalM Env]
extParams [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret = (Pat ParamType -> Env -> Value -> EvalM Env)
-> [Pat ParamType] -> [Env -> Value -> EvalM Env]
forall a b. (a -> b) -> [a] -> [b]
map Pat ParamType -> Env -> Value -> EvalM Env
forall {o}.
Pat (TypeBase (ExpBase Info VName) o) -> Env -> Value -> EvalM Env
fromPat [Pat ParamType]
ps [Env -> Value -> EvalM Env]
-> [Env -> Value -> EvalM Env] -> [Env -> Value -> EvalM Env]
forall a. Semigroup a => a -> a -> a
<> ((PName, ParamType) -> Env -> Value -> EvalM Env)
-> [(PName, ParamType)] -> [Env -> Value -> EvalM Env]
forall a b. (a -> b) -> [a] -> [b]
map (PName, ParamType) -> Env -> Value -> EvalM Env
forall {f :: * -> *} {o}.
Applicative f =>
(PName, TypeBase (ExpBase Info VName) o) -> Env -> Value -> f Env
fromArrow (([(PName, ParamType)], StructType) -> [(PName, ParamType)]
forall a b. (a, b) -> a
fst (([(PName, ParamType)], StructType) -> [(PName, ParamType)])
-> ([(PName, ParamType)], StructType) -> [(PName, ParamType)]
forall a b. (a -> b) -> a -> b
$ TypeBase (ExpBase Info VName) Freshness
-> ([(PName, ParamType)], StructType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType (TypeBase (ExpBase Info VName) Freshness
 -> ([(PName, ParamType)], StructType))
-> TypeBase (ExpBase Info VName) Freshness
-> ([(PName, ParamType)], StructType)
forall a b. (a -> b) -> a -> b
$ RetTypeBase (ExpBase Info VName) Freshness
-> TypeBase (ExpBase Info VName) Freshness
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType RetTypeBase (ExpBase Info VName) Freshness
ret)
  where
    fromPat :: Pat (TypeBase (ExpBase Info VName) o) -> Env -> Value -> EvalM Env
fromPat Pat (TypeBase (ExpBase Info VName) o)
p Env
env = Env -> Pat (TypeBase (ExpBase Info VName) o) -> Value -> EvalM Env
forall o.
Env -> Pat (TypeBase (ExpBase Info VName) o) -> Value -> EvalM Env
matchPat Env
env Pat (TypeBase (ExpBase Info VName) o)
p
    fromArrow :: (PName, TypeBase (ExpBase Info VName) o) -> Env -> Value -> f Env
fromArrow (Named VName
v, TypeBase (ExpBase Info VName) o
t) Env
env Value
val =
      Env -> f Env
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env -> f Env) -> Env -> f 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) o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) o
t, Value
val)) Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env
    fromArrow (PName
Unnamed, TypeBase (ExpBase Info VName) o
_) Env
env Value
_ = Env -> f Env
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
env

-- | Construct the value of an external binding with the given parameters. The
-- shape of the result may depend on the values of the parameters, so it is
-- computed only once every argument has been supplied and bound in the
-- environment. A binding of arity zero is not a function, and is called
-- immediately.
extFun ::
  Name ->
  Env ->
  [Env -> Value -> EvalM Env] ->
  ResRetType ->
  EvalM Value
extFun :: Name
-> Env
-> [Env -> Value -> EvalM Env]
-> RetTypeBase (ExpBase Info VName) Freshness
-> EvalM Value
extFun Name
n = [Value]
-> Env
-> [Env -> Value -> EvalM Env]
-> RetTypeBase (ExpBase Info VName) Freshness
-> EvalM Value
extFun' []
  where
    extFun' :: [Value]
-> Env
-> [Env -> Value -> EvalM Env]
-> RetTypeBase (ExpBase Info VName) Freshness
-> EvalM Value
extFun' [Value]
vs Env
env [] RetTypeBase (ExpBase Info VName) Freshness
ret = Name
-> RetTypeBase (ExpBase Info VName) Freshness
-> Env
-> [Value]
-> EvalM Value
extCall Name
n RetTypeBase (ExpBase Info VName) Freshness
ret Env
env ([Value] -> EvalM Value) -> [Value] -> EvalM Value
forall a b. (a -> b) -> a -> b
$ [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
vs
    extFun' [Value]
vs Env
env (Env -> Value -> EvalM Env
bind : [Env -> Value -> EvalM Env]
binds) RetTypeBase (ExpBase Info VName) Freshness
ret = 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 -> Value -> EvalM Env
bind Env
env Value
v
      extFun' (v : vs) env' binds ret

evalDec :: Env -> Dec -> EvalM Env
evalDec :: Env -> DecBase Info VName -> EvalM Env
evalDec Env
env (ValDec vb :: ValBindBase Info VName
vb@(ValBind (Just Info EntryPoint
_) vn :: VName
vn@(VName Name
n Int
_) SrcLoc
_ Maybe (TypeExp (ExpBase Info VName) VName)
_ (Info RetTypeBase (ExpBase Info VName) Freshness
ret) [TypeParam]
tparams [Pat ParamType]
ps ExpBase Info VName
_ Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_)) | AttrInfo VName
"$external" AttrInfo VName -> [AttrInfo VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ValBindBase Info VName -> [AttrInfo VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [AttrInfo vn]
valBindAttrs ValBindBase Info VName
vb = 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
  let ftype :: StructType
ftype = TypeBase SizeClosure NoMode -> StructType
evalToStruct (TypeBase SizeClosure NoMode -> StructType)
-> TypeBase SizeClosure NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Env -> StructType -> TypeBase SizeClosure NoMode
forall o.
Pretty o =>
Env -> TypeBase (ExpBase Info VName) o -> TypeBase SizeClosure o
expandType Env
env (StructType -> TypeBase SizeClosure NoMode)
-> StructType -> TypeBase SizeClosure NoMode
forall a b. (a -> b) -> a -> b
$ [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret
      bv :: Maybe BoundV
bv = 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
      params :: [Env -> Value -> EvalM Env]
params = [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness
-> [Env -> Value -> EvalM Env]
extParams [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret
  sizes <- EvalM Env
extEnv
  if null tparams
    then do
      f <- extFun n env params ret
      pure $ mempty {envTerm = M.singleton vn $ TermValue bv f} <> sizes
    else
      -- TODO: Add missing sizes?
      let pfn TypeBase SizeClosure NoMode
ftype' = do
            tparam_env <-
              ([(VName, ([VName], TypeBase SizeClosure NoMode))],
 [(VName, SizeClosure)])
-> EvalM Env
evalResolved (([(VName, ([VName], TypeBase SizeClosure NoMode))],
  [(VName, SizeClosure)])
 -> EvalM Env)
-> ([(VName, ([VName], TypeBase SizeClosure NoMode))],
    [(VName, SizeClosure)])
-> EvalM Env
forall a b. (a -> b) -> a -> b
$ [VName]
-> StructType
-> TypeBase SizeClosure NoMode
-> ([(VName, ([VName], TypeBase SizeClosure NoMode))],
    [(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 NoMode
ftype'
            extFun n (tparam_env <> env) params ret
       in pure $ mempty {envTerm = M.singleton vn $ TermPoly bv pfn} <> sizes
evalDec Env
env (ValDec (ValBind Maybe (Info EntryPoint)
_ VName
v SrcLoc
_ Maybe (TypeExp (ExpBase Info VName) VName)
_ (Info RetTypeBase (ExpBase Info VName) Freshness
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
  binding <- Env
-> VName
-> [TypeParam]
-> [Pat ParamType]
-> RetTypeBase (ExpBase Info VName) Freshness
-> ExpBase Info VName
-> EvalM TermBinding
evalValBinding Env
env VName
v [TypeParam]
tparams [Pat ParamType]
ps RetTypeBase (ExpBase Info VName) Freshness
ret ExpBase Info VName
fbody
  sizes <- extEnv
  pure $ mempty {envTerm = M.singleton v binding} <> sizes
evalDec Env
env (OpenDec ModExp
me SrcLoc
_) = do
  (me_env, me') <- Env -> ModExp -> EvalM (Env, Module)
evalModExp Env
env ModExp
me
  case 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
    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
forall a. Monoid a => a
mempty
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]
-> RetTypeBase (ExpBase Info VName) NoMode
-> TypeBinding
TypeConBinding Env
env [TypeParam]
ps (RetTypeBase (ExpBase Info VName) NoMode -> TypeBinding)
-> RetTypeBase (ExpBase Info VName) NoMode -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> RetTypeBase (ExpBase Info VName) NoMode
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims StructType
t
  Env -> EvalM Env
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
forall a. Monoid a => a
mempty {envType = M.singleton v abbr}
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
  (mod_env, 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
  pure $ modEnv (M.singleton v mod) <> mod_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

evalDecs :: Env -> [Dec] -> EvalM (Env, Env)
evalDecs :: Env -> [DecBase Info VName] -> EvalM (Env, Env)
evalDecs Env
env = ((Env, Env) -> DecBase Info VName -> EvalM (Env, Env))
-> (Env, Env) -> [DecBase Info VName] -> EvalM (Env, Env)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Env, Env) -> DecBase Info VName -> EvalM (Env, Env)
evalDec' (Env
env, Env
forall a. Monoid a => a
mempty)
  where
    evalDec' :: (Env, Env) -> DecBase Info VName -> EvalM (Env, Env)
evalDec' (Env
env', Env
built_env) DecBase Info VName
dec = do
      dec_env <- Env -> DecBase Info VName -> EvalM Env
evalDec Env
env' DecBase Info VName
dec
      pure (dec_env <> env', dec_env <> built_env)

-- | 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
      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 NE.nonEmpty 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 ()

getV :: PrimValue -> Maybe P.PrimValue
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 :: P.PrimValue -> PrimValue
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

getAD :: Value -> Maybe AD.ADValue
getAD :: Value -> 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 Depth
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
$ Depth -> ADVariable -> ADValue
AD.Variable Depth
d ADVariable
v
getAD Value
_ = Maybe ADValue
forall a. Maybe a
Nothing

putAD :: AD.ADValue -> Value
putAD :: ADValue -> Value
putAD (AD.Variable Depth
d ADVariable
s) = Depth -> ADVariable -> Value
forall (m :: * -> *). Depth -> ADVariable -> Value m
ValueAD Depth
d ADVariable
s
putAD (AD.Constant PrimValue
v) = PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> PrimValue -> Value
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
putV PrimValue
v

modifyValue :: (Num t) => (t -> Value -> Value) -> Value -> Value
modifyValue :: forall t. Num t => (t -> Value -> Value) -> Value -> Value
modifyValue t -> Value -> Value
f Value
v = (t, Value) -> Value
forall a b. (a, b) -> b
snd ((t, Value) -> Value) -> (t, Value) -> Value
forall a b. (a -> b) -> a -> b
$ (t -> Value -> (t, Value)) -> t -> Value -> (t, Value)
forall a (m :: * -> *).
(a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)
valueAccum (\t
a Value
b -> (t
a t -> t -> t
forall a. Num a => a -> a -> a
+ t
1, t -> Value -> Value
f t
a Value
b)) t
0 Value
v

modifyValueM ::
  (Num t, Monad m) =>
  (t -> Value -> m Value) ->
  Value ->
  m Value
modifyValueM :: forall t (m :: * -> *).
(Num t, Monad m) =>
(t -> Value -> m Value) -> Value -> m Value
modifyValueM t -> Value -> m Value
f Value
v = (t, Value) -> Value
forall a b. (a, b) -> b
snd ((t, Value) -> Value) -> m (t, Value) -> m Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (t -> Value -> m (t, Value)) -> t -> Value -> m (t, Value)
forall (f :: * -> *) a (m :: * -> *).
Monad f =>
(a -> Value m -> f (a, Value m)) -> a -> Value m -> f (a, Value m)
valueAccumLM t -> Value -> m (t, Value)
g t
0 Value
v
  where
    g :: t -> Value -> m (t, Value)
g t
a Value
b = do
      b' <- t -> Value -> m Value
f t
a Value
b
      pure (a + 1, 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?
doJVP2 :: Value -> Value -> Value -> EvalM Value
doJVP2 :: Value -> Value -> Value -> EvalM Value
doJVP2 Value
f Value
v0 Value
s0 = do
  -- Differentiation traverses the values themselves, so anything
  -- residing on a server must be fetched first.
  v <- Value -> EvalM Value
force Value
v0
  s <- force s0

  depth <- adDepth

  -- Turn the seeds into a list of ADValues
  let 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]
"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
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 -> 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]
"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 t (m :: * -> *).
(Num t, Monad m) =>
(t -> Value -> m Value) -> Value -> m Value
modifyValueM
            ( \Int
i Value
lv -> do
                lv' <- Value -> Maybe ADValue
getAD Value
lv
                pure $ ValueAD depth . AD.JVP . AD.JVPValue lv' $ s' !! (length s' - 1 - i)
            )
            Value
v

  -- Run the function, and turn its outputs into a list of Values
  o <- apply noLoc mempty f v'
  let 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)] -> Maybe [(Value, Value)] -> [(Value, Value)]
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> [(Value, Value)]
forall a. HasCallStack => [Char] -> a
error [Char]
"jvp: differentiation failed") (Maybe [(Value, Value)] -> [(Value, Value)])
-> Maybe [(Value, Value)] -> [(Value, Value)]
forall a b. (a -> b) -> a -> b
$
          [Value]
-> (Value -> Maybe (Value, Value)) -> Maybe [(Value, Value)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Value]
o' ((Value -> Maybe (Value, Value)) -> Maybe [(Value, Value)])
-> (Value -> Maybe (Value, Value)) -> Maybe [(Value, Value)]
forall a b. (a -> b) -> a -> b
$ \Value
on -> case Value
on of
            -- If it is a JVP variable of the correct depth, return
            -- its primal and derivative
            (ValueAD Depth
d (AD.JVP (AD.JVPValue ADValue
pv ADValue
dv)))
              | Depth
d Depth -> Depth -> Bool
forall a. Eq a => a -> a -> Bool
== Depth
depth -> (Value, Value) -> Maybe (Value, Value)
forall a. a -> Maybe a
Just (ADValue -> Value
putAD ADValue
pv, ADValue -> Value
putAD ADValue
dv)
            -- Otherwise, its partial derivatives are all 0
            Value
_ ->
              (Value
on,)
                (Value -> (Value, Value))
-> (ADValue -> Value) -> ADValue -> (Value, Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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, Value))
-> Maybe ADValue -> Maybe (Value, Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Maybe ADValue
getAD Value
on

  -- Extract the output values, and the partial derivatives
  let ov = (Int -> Value -> Value) -> Value -> Value
forall t. Num t => (t -> Value -> Value) -> Value -> Value
modifyValue (\Int
i Value
_ -> (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)]
m [(Value, Value)] -> Int -> (Value, Value)
forall a. HasCallStack => [a] -> Int -> a
!! ([(Value, Value)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Value, Value)]
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 = (Int -> Value -> Value) -> Value -> Value
forall t. Num t => (t -> Value -> Value) -> Value -> Value
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)]
m [(Value, Value)] -> Int -> (Value, Value)
forall a. HasCallStack => [a] -> Int -> a
!! ([(Value, Value)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Value, Value)]
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
  pure $ toTuple [ov, od]

-- 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?
doVJP2 :: Value -> Value -> Value -> EvalM Value
doVJP2 :: Value -> Value -> Value -> EvalM Value
doVJP2 Value
f Value
v0 Value
s0 = do
  -- Differentiation traverses the values themselves, so anything
  -- residing on a server must be fetched first.
  v <- Value -> EvalM Value
force Value
v0
  s <- force s0

  -- Get the depth
  depth <- adDepth

  -- Augment the values
  let 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
$
          (Counter -> Value -> Maybe Value) -> Value -> Maybe Value
forall t (m :: * -> *).
(Num t, Monad m) =>
(t -> Value -> m Value) -> Value -> m Value
modifyValueM (\Counter
i Value
lv -> Depth -> ADVariable -> Value
forall (m :: * -> *). Depth -> ADVariable -> Value m
ValueAD Depth
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
. Counter -> ADValue -> Tape
AD.TapeID Counter
i (ADValue -> Value) -> Maybe ADValue -> Maybe Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Maybe ADValue
getAD Value
lv) Value
v
  -- Turn the seeds into a list of ADValues
  let 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
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
  o <- apply noLoc mempty f v'
  let 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..
  m <-
    forM
      (zip o' s')
      ( \(Value
on, ADValue
sn) -> case Value
on of
          -- If it is a VJP variable of the correct depth, run
          -- deriveTape on it- and its corresponding seed
          (ValueAD Depth
d (AD.VJP (AD.VJPValue Tape
t)))
            | Depth
d Depth -> Depth -> Bool
forall a. Eq a => a -> a -> Bool
== Depth
depth ->
                EvalM Counter
getCounter
                  EvalM Counter
-> (Counter -> EvalM (Either [Char] (Value, Map Counter ADValue)))
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ([Char] -> EvalM (Either [Char] (Value, Map Counter ADValue)))
-> ((Map Counter ADValue, Counter)
    -> EvalM (Either [Char] (Value, Map Counter ADValue)))
-> Either [Char] (Map Counter ADValue, Counter)
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
                    (Either [Char] (Value, Map Counter ADValue)
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either [Char] (Value, Map Counter ADValue)
 -> EvalM (Either [Char] (Value, Map Counter ADValue)))
-> ([Char] -> Either [Char] (Value, Map Counter ADValue))
-> [Char]
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Either [Char] (Value, Map Counter ADValue)
forall a b. a -> Either a b
Left)
                    (\(Map Counter ADValue
m', Counter
i) -> Counter -> EvalM ()
putCounter Counter
i EvalM ()
-> Either [Char] (Value, Map Counter ADValue)
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> (Value, Map Counter ADValue)
-> Either [Char] (Value, Map Counter ADValue)
forall a b. b -> Either a b
Right (ADValue -> Value
putAD (ADValue -> Value) -> ADValue -> Value
forall a b. (a -> b) -> a -> b
$ Tape -> ADValue
AD.tapePrimal Tape
t, Map Counter ADValue
m'))
                    (Either [Char] (Map Counter ADValue, Counter)
 -> EvalM (Either [Char] (Value, Map Counter ADValue)))
-> (Counter -> Either [Char] (Map Counter ADValue, Counter))
-> Counter
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tape
-> ADValue
-> Counter
-> Either [Char] (Map Counter ADValue, Counter)
AD.deriveTape Tape
t ADValue
sn
          -- Otherwise, its partial derivatives are all 0
          Value
_ -> Either [Char] (Value, Map Counter ADValue)
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either [Char] (Value, Map Counter ADValue)
 -> EvalM (Either [Char] (Value, Map Counter ADValue)))
-> Either [Char] (Value, Map Counter ADValue)
-> EvalM (Either [Char] (Value, Map Counter ADValue))
forall a b. (a -> b) -> a -> b
$ (Value, Map Counter ADValue)
-> Either [Char] (Value, Map Counter ADValue)
forall a b. b -> Either a b
Right (Value
on, Map Counter ADValue
forall k a. Map k a
M.empty)
      )
      <&> either (error . show) id . sequence

  -- Add together every derivative
  drvs' <- AD.unionsWithM add (map snd m)
  let drvs = (ADValue -> Maybe Value)
-> Map Counter ADValue -> Map Counter (Maybe Value)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value)
-> (ADValue -> Value) -> ADValue -> Maybe Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ADValue -> Value
putAD) Map Counter ADValue
drvs'

  -- Extract the output values, and the partial derivatives
  let ov = (Int -> Value -> Value) -> Value -> Value
forall t. Num t => (t -> Value -> Value) -> Value -> Value
modifyValue (\Int
i Value
_ -> (Value, Map Counter ADValue) -> Value
forall a b. (a, b) -> a
fst ((Value, Map Counter ADValue) -> Value)
-> (Value, Map Counter ADValue) -> Value
forall a b. (a -> b) -> a -> b
$ [(Value, Map Counter ADValue)]
m [(Value, Map Counter ADValue)]
-> Int -> (Value, Map Counter ADValue)
forall a. HasCallStack => [a] -> Int -> a
!! ([(Value, Map Counter ADValue)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Value, Map Counter 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 -> 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
$
          (Counter -> Value -> Maybe Value) -> Value -> Maybe Value
forall t (m :: * -> *).
(Num t, Monad m) =>
(t -> Value -> m Value) -> Value -> m Value
modifyValueM (\Counter
i Value
vo -> Maybe Value -> Counter -> Map Counter (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
getAD Value
vo) Counter
i Map Counter (Maybe Value)
drvs) Value
v

  -- Return a tuple of the output values, and partial derivatives
  pure $ toTuple [ov, od]
  where
    -- TODO: Perhaps this could be fully abstracted by AD?
    -- Making addFor private would be nice..
    add :: ADValue -> ADValue -> EvalM ADValue
add ADValue
x ADValue
y =
      EvalM Counter
getCounter
        EvalM Counter -> (Counter -> EvalM ADValue) -> EvalM ADValue
forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ([Char] -> EvalM ADValue)
-> ((ADValue, Counter) -> EvalM ADValue)
-> Either [Char] (ADValue, Counter)
-> EvalM ADValue
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
          ([Char] -> EvalM ADValue
forall a. HasCallStack => [Char] -> a
error ([Char] -> EvalM ADValue) -> ShowS -> [Char] -> EvalM ADValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
forall a. Show a => a -> [Char]
show)
          (\(ADValue
a, Counter
b) -> Counter -> EvalM ()
putCounter Counter
b EvalM () -> EvalM ADValue -> EvalM ADValue
forall a b. EvalM a -> EvalM b -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ADValue -> EvalM ADValue
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ADValue
a)
          (Either [Char] (ADValue, Counter) -> EvalM ADValue)
-> (Counter -> Either [Char] (ADValue, Counter))
-> Counter
-> EvalM ADValue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op -> [ADValue] -> Counter -> Either [Char] (ADValue, Counter)
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]

-- | 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 -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
intOp IntType -> BinOp
f = (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp IntType -> BinOp
f [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintOp IntType -> BinOp
f
    floatOp :: (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
intOp IntType -> BinOp
f [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter)
adBinOp (Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter)
forall a b. (a -> b) -> a -> b
$ CmpOp -> Op
AD.OpCmp CmpOp
f)]

    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

    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 -> Counter -> Maybe (ADValue, Counter)
adBinOp Op
op ADValue
x ADValue
y Counter
i =
      ([Char] -> Maybe (ADValue, Counter))
-> ((ADValue, Counter) -> Maybe (ADValue, Counter))
-> Either [Char] (ADValue, Counter)
-> Maybe (ADValue, Counter)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (ADValue, Counter) -> [Char] -> Maybe (ADValue, Counter)
forall a b. a -> b -> a
const Maybe (ADValue, Counter)
forall a. Maybe a
Nothing) (ADValue, Counter) -> Maybe (ADValue, Counter)
forall a. a -> Maybe a
Just (Either [Char] (ADValue, Counter) -> Maybe (ADValue, Counter))
-> Either [Char] (ADValue, Counter) -> Maybe (ADValue, Counter)
forall a b. (a -> b) -> a -> b
$ Op -> [ADValue] -> Counter -> Either [Char] (ADValue, Counter)
AD.doOp Op
op [ADValue
x, ADValue
y] Counter
i
    adUnOp :: Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
adUnOp Op
op ADValue
x Counter
i = ([Char] -> Maybe (ADValue, Counter))
-> ((ADValue, Counter) -> Maybe (ADValue, Counter))
-> Either [Char] (ADValue, Counter)
-> Maybe (ADValue, Counter)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (ADValue, Counter) -> [Char] -> Maybe (ADValue, Counter)
forall a b. a -> b -> a
const Maybe (ADValue, Counter)
forall a. Maybe a
Nothing) (ADValue, Counter) -> Maybe (ADValue, Counter)
forall a. a -> Maybe a
Just (Either [Char] (ADValue, Counter) -> Maybe (ADValue, Counter))
-> Either [Char] (ADValue, Counter) -> Maybe (ADValue, Counter)
forall a b. (a -> b) -> a -> b
$ Op -> [ADValue] -> Counter -> Either [Char] (ADValue, Counter)
AD.doOp Op
op [ADValue
x] Counter
i

    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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> do
      i <- EvalM Counter
getCounter
      case (x, 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 -> Counter -> Maybe (ADValue, Counter))
 -> Maybe PrimValue)
-> [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [Maybe PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map ((PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
 ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> (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 -> Counter -> Maybe (ADValue, Counter))]
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
getAD Value
x,
            Just ADValue
y' <- Value -> Maybe ADValue
getAD Value
y,
            Just (ADValue
z, Counter
i') <- [Maybe (ADValue, Counter)] -> Maybe (ADValue, Counter)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum ([Maybe (ADValue, Counter)] -> Maybe (ADValue, Counter))
-> [Maybe (ADValue, Counter)] -> Maybe (ADValue, Counter)
forall a b. (a -> b) -> a -> b
$ ((PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
 -> Maybe (ADValue, Counter))
-> [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [Maybe (ADValue, Counter)]
forall a b. (a -> b) -> [a] -> [b]
map ((PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
 ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> (ADValue, ADValue, Counter) -> Maybe (ADValue, Counter)
forall {a} {b} {c} {t} {t} {t} {t}.
(a, b, c, t -> t -> t -> t) -> (t, t, t) -> t
`bopDefAD` (ADValue
x', ADValue
y', Counter
i)) [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
fs -> do
              Counter -> EvalM ()
putCounter Counter
i'
              [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
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 (ZonkAny 4) -> Text) -> Doc (ZonkAny 4) -> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 4) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 4) -> EvalM Value) -> Doc (ZonkAny 4) -> EvalM Value
forall a b. (a -> b) -> a -> b
$
            Doc (ZonkAny 4)
"Cannot apply operator to arguments"
              Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann
dquotes (Value -> Doc (ZonkAny 4)
forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
x)
              Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 4)
"and"
              Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann
dquotes (Value -> Doc (ZonkAny 4)
forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
y)
              Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall a. Semigroup a => a -> a -> a
<> Doc (ZonkAny 4)
"."
      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
          x' <- t -> m t
valf t
x
          y' <- valf y
          retf =<< op x' y'
        bopDefAD :: (a, b, c, t -> t -> t -> t) -> (t, t, t) -> t
bopDefAD (a
_, b
_, c
_, t -> t -> t -> t
dop) (t
x, t
y, t
i) = t -> t -> t -> t
dop t
x t
y t
i

    unopDef :: [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
unopDef [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
fs = (Value -> EvalM Value) -> TermBinding
fun1 ((Value -> EvalM Value) -> TermBinding)
-> (Value -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \Value
x -> do
      i <- EvalM Counter
getCounter
      case 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 -> Counter -> Maybe (ADValue, Counter))
 -> Maybe PrimValue)
-> [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
     ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [Maybe PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map ((PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
 ADValue -> Counter -> Maybe (ADValue, Counter))
-> 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 -> Counter -> Maybe (ADValue, Counter))]
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
getAD Value
x,
            Just (ADValue
r, Counter
i') <- [Maybe (ADValue, Counter)] -> Maybe (ADValue, Counter)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum ([Maybe (ADValue, Counter)] -> Maybe (ADValue, Counter))
-> [Maybe (ADValue, Counter)] -> Maybe (ADValue, Counter)
forall a b. (a -> b) -> a -> b
$ ((PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))
 -> Maybe (ADValue, Counter))
-> [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
     ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [Maybe (ADValue, Counter)]
forall a b. (a -> b) -> [a] -> [b]
map ((PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
 ADValue -> Counter -> Maybe (ADValue, Counter))
-> (ADValue, Counter) -> Maybe (ADValue, Counter)
forall {a} {b} {c} {t} {t} {t}.
(a, b, c, t -> t -> t) -> (t, t) -> t
`unopDefAD'` (ADValue
x', Counter
i)) [(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
fs -> do
              Counter -> EvalM ()
putCounter Counter
i'
              [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
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 (ZonkAny 5) -> Text) -> Doc (ZonkAny 5) -> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 5) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 5) -> EvalM Value) -> Doc (ZonkAny 5) -> EvalM Value
forall a b. (a -> b) -> a -> b
$
            Doc (ZonkAny 5)
"Cannot apply function to argument"
              Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann
dquotes (Value -> Doc (ZonkAny 5)
forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
x)
              Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall a. Semigroup a => a -> a -> a
<> Doc (ZonkAny 5)
"."
      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
          x' <- t -> m t
valf t
x
          retf =<< op x'
        unopDefAD' :: (a, b, c, t -> t -> t) -> (t, t) -> t
unopDefAD' (a
_, b
_, c
_, t -> t -> t
dop) (t
x, t
i) = t -> t -> t
dop t
x t
i

    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 -> do
      i <- EvalM Counter
getCounter
      case fromTuple 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
getAD Value
x,
            Just ADValue
y' <- Value -> Maybe ADValue
getAD Value
y,
            Right (ADValue
z, Counter
i') <- Op -> [ADValue] -> Counter -> Either [Char] (ADValue, Counter)
AD.doOp Op
op [ADValue
x', ADValue
y'] Counter
i -> do
              Counter -> EvalM ()
putCounter Counter
i'
              [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
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 (ZonkAny 6) -> Text) -> Doc (ZonkAny 6) -> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 6) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 6) -> EvalM Value) -> Doc (ZonkAny 6) -> EvalM Value
forall a b. (a -> b) -> a -> b
$
            Doc (ZonkAny 6)
"Cannot apply operator to argument"
              Doc (ZonkAny 6) -> Doc (ZonkAny 6) -> Doc (ZonkAny 6)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 6) -> Doc (ZonkAny 6)
forall ann. Doc ann -> Doc ann
dquotes (Value -> Doc (ZonkAny 6)
forall (m :: * -> *) a. Value m -> Doc a
prettyValue Value
v)
              Doc (ZonkAny 6) -> Doc (ZonkAny 6) -> Doc (ZonkAny 6)
forall a. Semigroup a => a -> a -> a
<> Doc (ZonkAny 6)
"."

    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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
-> 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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
-> 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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp (IntType -> Safety -> BinOp
`P.SDiv` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintOp (IntType -> Safety -> BinOp
`P.UDiv` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp (IntType -> Safety -> BinOp
`P.SMod` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintOp (IntType -> Safety -> BinOp
`P.UMod` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp (IntType -> Safety -> BinOp
`P.SQuot` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp (IntType -> Safety -> BinOp
`P.SRem` Safety
P.Unsafe)
            [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp IntType -> BinOp
P.AShr [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe PrimValue,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintOp IntType -> BinOp
P.LShr [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe PrimValue,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> BinOp)
-> [(PrimValue -> Maybe PrimValue, PrimValue -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe PrimValue,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintOp IntType -> BinOp
P.LShr
    -- Equality inspects the values themselves, so anything residing on a
    -- server must be fetched first.
    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 -> PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> (Bool -> PrimValue) -> Bool -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue (Bool -> Value) -> EvalM Bool -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> Value -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Value -> Value -> Bool) -> EvalM Value -> EvalM (Value -> Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM Value
force Value
xs EvalM (Value -> Bool) -> EvalM Value -> EvalM Bool
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Value -> EvalM Value
force 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 -> PrimValue -> Value
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value) -> (Bool -> PrimValue) -> Bool -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue (Bool -> Value) -> EvalM Bool -> EvalM Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> Value -> Bool
forall a. Eq a => a -> a -> Bool
(/=) (Value -> Value -> Bool) -> EvalM Value -> EvalM (Value -> Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM Value
force Value
xs EvalM (Value -> Bool) -> EvalM Value -> EvalM Bool
forall a b. EvalM (a -> b) -> EvalM a -> EvalM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Value -> EvalM Value
force 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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintCmp IntType -> CmpOp
P.CmpSlt
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintCmp IntType -> CmpOp
P.CmpUlt
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
floatCmp FloatType -> CmpOp
P.FCmpLt
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
 -> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
      PrimValue -> PrimValue -> Maybe Bool,
      ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))])
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a b. (a -> b) -> a -> b
$
            (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintCmp IntType -> CmpOp
P.CmpSlt
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintCmp IntType -> CmpOp
P.CmpUlt
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
floatCmp FloatType -> CmpOp
P.FCmpLt
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintCmp IntType -> CmpOp
P.CmpSle
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintCmp IntType -> CmpOp
P.CmpUle
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
floatCmp FloatType -> CmpOp
P.FCmpLe
            [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> t -> Maybe a,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
bopDef ([(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
   PrimValue -> PrimValue -> Maybe Bool,
   ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
 -> TermBinding)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall a b. (a -> b) -> a -> b
$
          [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Counter -> Maybe (ADValue, Counter))]
 -> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
      PrimValue -> PrimValue -> Maybe Bool,
      ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))])
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a b. (a -> b) -> a -> b
$
            (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
sintCmp IntType -> CmpOp
P.CmpSle
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (IntType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
uintCmp IntType -> CmpOp
P.CmpUle
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ (FloatType -> CmpOp)
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
floatCmp FloatType -> CmpOp
P.FCmpLe
              [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
  PrimValue -> PrimValue -> Maybe Bool,
  ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
forall a. [a] -> [a] -> [a]
++ CmpOp
-> [(PrimValue -> Maybe PrimValue, Bool -> Maybe PrimValue,
     PrimValue -> PrimValue -> Maybe Bool,
     ADValue -> ADValue -> Counter -> Maybe (ADValue, Counter))]
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 -> Text -> Bool
forall a. Eq a => a -> a -> Bool
==) (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 -> Text -> Bool
forall a. Eq a => a -> a -> Bool
==) (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 -> Text -> Bool
forall a. Eq a => a -> a -> Bool
==) (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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
-> 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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Text -> Bool
forall a. Eq a => a -> a -> Bool
==) (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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
-> 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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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 -> Counter -> Maybe (ADValue, Counter))]
-> TermBinding
forall {t} {a}.
[(PrimValue -> Maybe t, a -> Maybe PrimValue, t -> Maybe a,
  ADValue -> Counter -> Maybe (ADValue, Counter))]
-> 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 -> Counter -> Maybe (ADValue, Counter)
adUnOp (Op -> ADValue -> Counter -> Maybe (ADValue, Counter))
-> Op -> ADValue -> Counter -> Maybe (ADValue, Counter)
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'
                ValueAD {} -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
x -- FIXME: these do not carry signs.
                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'
                ValueAD {} -> Value -> EvalM Value
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value
x -- FIXME: these do not carry signs.
                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 NoMode -> EvalM Value) -> TermBinding
TermPoly Maybe BoundV
forall a. Maybe a
Nothing ((TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding)
-> (TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \TypeBase SizeClosure NoMode
t -> do
        t' <- TypeBase SizeClosure NoMode -> EvalM ValueType
evalTypeFully TypeBase SizeClosure NoMode
t
        pure $ ValueFun $ \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 -> ([(PName, TypeBase Int64 Diet)], ValueType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType ValueType
t' of
            ([(PName, TypeBase Int64 Diet)
_, (PName, TypeBase Int64 Diet)
_], ValueType
ret_t)
              | ValueShape
rowshape <- ValueType -> ValueShape
forall d o. TypeBase d o -> Shape d
typeShape (ValueType -> ValueShape) -> ValueType -> ValueShape
forall a b. (a -> b) -> a -> b
$ Int -> ValueType -> ValueType
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 ValueType
ret_t -> do
                  xs' <- (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> EvalM (ValueShape, [Value]) -> EvalM [Value]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
                  toArray' rowshape <$> mapM (apply noLoc mempty f) xs'
            ([(PName, 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 NoMode -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase SizeClosure NoMode
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
"flatmap" = 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 NoMode -> EvalM Value) -> TermBinding
TermPoly Maybe BoundV
forall a. Maybe a
Nothing ((TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding)
-> (TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \TypeBase SizeClosure NoMode
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 TypeBase SizeClosure NoMode
-> ([(PName, TypeBase SizeClosure Diet)],
    TypeBase SizeClosure NoMode)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType TypeBase SizeClosure NoMode
t of
            ([(PName, TypeBase SizeClosure Diet)
_, (PName, TypeBase SizeClosure Diet)
_], TypeBase SizeClosure NoMode
ret_t)
              | Just [TypeBase SizeClosure NoMode
_, TypeBase SizeClosure NoMode
_, TypeBase SizeClosure NoMode
_, TypeBase SizeClosure NoMode
irreg_t, TypeBase SizeClosure NoMode
reg_t] <- TypeBase SizeClosure NoMode -> Maybe [TypeBase SizeClosure NoMode]
forall dim o. TypeBase dim o -> Maybe [TypeBase dim o]
isTupleRecord TypeBase SizeClosure NoMode
ret_t -> do
                  irreg_rowshape <- ValueType -> ValueShape
forall d o. TypeBase d o -> Shape d
typeShape (ValueType -> ValueShape) -> EvalM ValueType -> EvalM ValueShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase SizeClosure NoMode -> EvalM ValueType
evalTypeFully (Int -> TypeBase SizeClosure NoMode -> TypeBase SizeClosure NoMode
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 TypeBase SizeClosure NoMode
irreg_t)
                  reg_rowshape <- typeShape <$> evalTypeFully (stripArray 1 reg_t)
                  yss <-
                    mapM (apply noLoc mempty f) . snd
                      =<< fromArray xs
                  -- Each application produces a segment, which is concatenated
                  -- with the others, and a value that is merely collected.
                  let (segs, regs) = unzip $ map (fromPair . fromTuple) yss
                  segs' <- mapM (fmap snd . fromArray) segs
                  let seg_sizes = ([Value] -> Int64) -> [[Value]] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map [Value] -> Int64
forall i a. Num i => [a] -> i
genericLength [[Value]]
segs' :: [Int64]
                      offsets = [Int64] -> [Int64]
forall a. HasCallStack => [a] -> [a]
init ([Int64] -> [Int64]) -> [Int64] -> [Int64]
forall a b. (a -> b) -> a -> b
$ (Int64 -> Int64 -> Int64) -> Int64 -> [Int64] -> [Int64]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
(+) Int64
0 [Int64]
seg_sizes
                      flag a
s = if a
s a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then [] else Bool
True Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Bool
False
                      mkI64 = PrimValue -> Value m
forall (m :: * -> *). PrimValue -> Value m
ValuePrim (PrimValue -> Value m) -> (Int64 -> PrimValue) -> Int64 -> Value m
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
                  pure $
                    toTuple
                      [ toArray' ShapeLeaf $ map mkI64 seg_sizes,
                        toArray' ShapeLeaf $ map (ValuePrim . BoolValue) $ concatMap flag seg_sizes,
                        toArray' ShapeLeaf $ map mkI64 offsets,
                        toArray' irreg_rowshape $ concat segs',
                        toArray' reg_rowshape regs
                      ]
            ([(PName, TypeBase SizeClosure Diet)], TypeBase SizeClosure NoMode)
_ ->
              [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 flatmap intrinsic:\n"
                  [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unlines [Value -> [Char]
forall a. Show a => a -> [Char]
show Value
f, Value -> [Char]
forall a. Show a => a -> [Char]
show Value
xs]
      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]
"flatmap: lambda did not return a pair"
    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)
-> ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value])
-> EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> EvalM Value)
-> EvalM (ValueShape, [Value]) -> EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value -> EvalM (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
              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
              pure (x' : out, x')
        xs' <- (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> EvalM (ValueShape, [Value]) -> EvalM [Value]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
        toArray' (valueShape ne) . reverse . fst
          <$> foldM next ([], ne) 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 ->
        Value -> EvalM Value
force Value
arr EvalM Value -> (Value -> EvalM Value) -> EvalM Value
forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          ValueArray ValueShape
shape Array Int Value
arr' -> do
            is' <- (Value -> Int) -> [Value] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Int
asInt ([Value] -> [Int])
-> ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value])
-> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Int])
-> EvalM (ValueShape, [Value]) -> EvalM [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
is
            vs' <- snd <$> fromArray vs
            pure $ ValueArray shape $ foldl' update arr' $ zip is' vs'
          Value
arr' ->
            [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 ->
        Value -> EvalM Value
force Value
arr EvalM Value -> (Value -> EvalM Value) -> EvalM Value
forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          arr' :: Value
arr'@ValueArray {} -> do
            is' <- (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]])
-> ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value])
-> [Maybe [Value]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Maybe [Value]])
-> EvalM (ValueShape, [Value]) -> EvalM [Maybe [Value]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
is
            vs' <- snd <$> fromArray vs
            pure $ foldl' update arr' $ zip is' vs'
          Value
arr' ->
            [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 ->
        Value -> EvalM Value
force Value
arr EvalM Value -> (Value -> EvalM Value) -> EvalM Value
forall a b. EvalM a -> (a -> EvalM b) -> EvalM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          arr' :: Value
arr'@ValueArray {} -> do
            is' <- (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]])
-> ((ValueShape, [Value]) -> [Value])
-> (ValueShape, [Value])
-> [Maybe [Value]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Maybe [Value]])
-> EvalM (ValueShape, [Value]) -> EvalM [Maybe [Value]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
is
            vs' <- snd <$> fromArray vs
            pure $ foldl' update arr' $ zip is' vs'
          Value
arr' ->
            [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 -> do
      arr' <- Value -> EvalM Value
force Value
arr
      is' <- map asInt64 . snd <$> fromArray is
      vs' <- snd <$> fromArray vs
      foldM (update fun) arr' $ zip is' 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 -> do
      arr' <- Value -> EvalM Value
force Value
arr
      is' <- map fromTuple . snd <$> fromArray is
      vs' <- snd <$> fromArray vs
      foldM (update fun) arr' $ zip is' 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 -> do
      arr' <- Value -> EvalM Value
force Value
arr
      is' <- map fromTuple . snd <$> fromArray is
      vs' <- snd <$> fromArray vs
      foldM (update fun) arr' $ zip is' 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
        (xs_shape, xs') <- Value -> EvalM (ValueShape, [Value])
fromArray Value
xs

        let ShapeDim _ rowshape = xs_shape

            next [[Value]]
outs Value
x = do
              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
              pure $ insertAt i x outs
            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
                ]

        pack . map reverse
          <$> foldM next (replicate (asInt k) []) 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 -> do
        dest' <- Value -> EvalM Value
force Value
dest
        vs' <- force vs
        case (dest', 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
              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 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
dest', 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 -> do
        dest' <- Value -> EvalM Value
force Value
dest
        vs' <- force vs
        case (dest', 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
              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 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
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'
              res <- Value m -> Value m -> m (Value m)
op Value m
x Value m
v
              pure $ ValueAcc shape op $ acc_arr // [(fromIntegral i', 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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk1 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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'])) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk2 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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']) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 -> MaybeT EvalM Value
f Int64
i Int64
j =
            [Indexing] -> Value -> MaybeT EvalM 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

      vs <- 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
$ [MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk1 [[MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk2 [Int64 -> Int64 -> MaybeT EvalM 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']
      case vs 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) -> MaybeT EvalM Value
f Value
arr' (Int64
i, Int64
j) =
                EvalM (Maybe Value) -> MaybeT EvalM Value
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (EvalM (Maybe Value) -> MaybeT EvalM Value)
-> (Value -> EvalM (Maybe Value)) -> Value -> MaybeT EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Value -> EvalM (Maybe Value))
-> (Value -> Maybe Value) -> Value -> EvalM (Maybe Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [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 -> MaybeT EvalM Value)
-> MaybeT EvalM Value -> MaybeT EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> MaybeT EvalM Value
indexArray [Int64 -> Indexing
IndexingFix Int64
i, Int64 -> Indexing
IndexingFix Int64
j] Value
v
          dest <- Value -> EvalM Value
force Value
arr
          o <- runMaybeT $ foldM f dest [(i, j) | i <- iota n1, j <- iota n2]
          case o 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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk1 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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'])) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk2 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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']) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk3 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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']) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 -> MaybeT EvalM Value
f Int64
i Int64
j Int64
l =
            [Indexing] -> Value -> MaybeT EvalM 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

      vs <- 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
$ [MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk1 [[MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk2 [[MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk3 [Int64 -> Int64 -> Int64 -> MaybeT EvalM 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']
      case vs 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) -> MaybeT EvalM Value
f Value
arr' (Int64
i, Int64
j, Int64
l) =
                EvalM (Maybe Value) -> MaybeT EvalM Value
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (EvalM (Maybe Value) -> MaybeT EvalM Value)
-> (Value -> EvalM (Maybe Value)) -> Value -> MaybeT EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Value -> EvalM (Maybe Value))
-> (Value -> Maybe Value) -> Value -> EvalM (Maybe Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [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 -> MaybeT EvalM Value)
-> MaybeT EvalM Value -> MaybeT EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> MaybeT EvalM Value
indexArray [Int64 -> Indexing
IndexingFix Int64
i, Int64 -> Indexing
IndexingFix Int64
j, Int64 -> Indexing
IndexingFix Int64
l] Value
v
          dest <- Value -> EvalM Value
force Value
arr
          o <- runMaybeT $ foldM f dest [(i, j, l) | i <- iota n1, j <- iota n2, l <- iota n3]
          case o 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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk1 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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'])) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk2 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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']) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk3 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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']) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 :: [MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk4 = ([Value m] -> Value m)
-> MaybeT EvalM [Value m] -> MaybeT EvalM (Value m)
forall a b. (a -> b) -> MaybeT EvalM a -> MaybeT EvalM 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']) (MaybeT EvalM [Value m] -> MaybeT EvalM (Value m))
-> ([MaybeT EvalM (Value m)] -> MaybeT EvalM [Value m])
-> [MaybeT EvalM (Value m)]
-> MaybeT EvalM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [MaybeT EvalM (Value m)] -> MaybeT EvalM [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 -> MaybeT EvalM Value
f Int64
i Int64
j Int64
l Int64
m =
            [Indexing] -> Value -> MaybeT EvalM 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

      vs <- 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
$ [MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk1 [[MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk2 [[MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk3 [[MaybeT EvalM Value] -> MaybeT EvalM Value
forall {m :: * -> *}.
[MaybeT EvalM (Value m)] -> MaybeT EvalM (Value m)
mk4 [Int64 -> Int64 -> Int64 -> Int64 -> MaybeT EvalM 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']
      case vs 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) -> MaybeT EvalM Value
f Value
arr' (Int64
i, Int64
j, Int64
l, Int64
m) =
                EvalM (Maybe Value) -> MaybeT EvalM Value
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (EvalM (Maybe Value) -> MaybeT EvalM Value)
-> (Value -> EvalM (Maybe Value)) -> Value -> MaybeT EvalM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Value -> EvalM (Maybe Value)
forall a. a -> EvalM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Value -> EvalM (Maybe Value))
-> (Value -> Maybe Value) -> Value -> EvalM (Maybe Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [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 -> MaybeT EvalM Value)
-> MaybeT EvalM Value -> MaybeT EvalM Value
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Indexing] -> Value -> MaybeT EvalM Value
indexArray [Int64 -> Indexing
IndexingFix Int64
i, Int64 -> Indexing
IndexingFix Int64
j, Int64 -> Indexing
IndexingFix Int64
l, Int64 -> Indexing
IndexingFix Int64
m] Value
v
          dest <- Value -> EvalM Value
force Value
arr
          o <- runMaybeT $ foldM f dest [(i, j, l, m) | i <- iota n1, j <- iota n2, l <- iota n3, m <- iota n4]
          case o 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]

        x' <- (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> EvalM (ValueShape, [Value]) -> EvalM [Value]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
x
        pure $ toTuple $ listPair $ unzip $ map (fromPair . fromTuple) 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
        xs' <- (ValueShape, [Value]) -> [Value]
forall a b. (a, b) -> b
snd ((ValueShape, [Value]) -> [Value])
-> EvalM (ValueShape, [Value]) -> EvalM [Value]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
        ys' <- snd <$> fromArray ys
        pure $
          toArray' (ShapeRecord (tupleFields [xs_rowshape, ys_rowshape])) $
            map toTuple $
              transpose [xs', 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
        (xs_shape, xs') <- Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
        (_, ys') <- fromArray ys
        let ShapeDim _ rowshape = xs_shape
        pure $ toArray' rowshape $ xs' ++ 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
        (xs_shape, xs') <- Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
        let ShapeDim n (ShapeDim m shape) = xs_shape
        rows <- mapM (fmap snd . fromArray) xs'
        pure $
          toArray (ShapeDim m (ShapeDim n shape)) $
            map (toArray (ShapeDim n shape)) $
              -- Slight hack to work around empty dimensions.
              genericTake m $
                transpose rows ++ 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
        (xs_shape, xs') <- Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
        let ShapeDim n (ShapeDim m shape) = xs_shape
        rows <- mapM (fmap snd . fromArray) xs'
        pure $ toArray (ShapeDim (n * m) shape) $ concat rows
    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
        (xs_shape, xs') <- Value -> EvalM (ValueShape, [Value])
fromArray Value
xs
        let ShapeDim xs_size innershape = xs_shape
            rowshape = Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim (Value -> Int64
asInt64 Value
m) ValueShape
innershape
            shape = Int64 -> ValueShape -> ValueShape
forall d. d -> Shape d -> Shape d
ShapeDim (Value -> Int64
asInt64 Value
n) ValueShape
rowshape
        if asInt64 n * asInt64 m /= xs_size || asInt64 n < 0 || asInt64 m < 0
          then
            bad mempty mempty $
              "Cannot unflatten array of shape ["
                <> prettyText xs_size
                <> "] to array of shape ["
                <> prettyText (asInt64 n)
                <> "]["
                <> prettyText (asInt64 m)
                <> "]"
          else pure $ toArray shape $ map (toArray rowshape) $ chunk (asInt m) 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
"jvp2" = 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
doJVP2
    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
$ (Value -> Value -> Value -> EvalM Value) -> TermBinding
fun3 Value -> Value -> Value -> EvalM Value
doVJP2
    def Text
"jmp2" = 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
x Value
seeds -> do
      v <- 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
      seeds' <- snd <$> fromArray seeds
      dvs <-
        toArray' (valueShape v) . map (project "1")
          <$> mapM (doJVP2 f x) seeds'
      pure $ toTuple [v, dvs]
    def Text
"mjp2" = 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
x Value
seeds -> do
      v <- 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
      seeds' <- snd <$> fromArray seeds
      dvs <-
        toArray' (valueShape x) . map (project "1")
          <$> mapM (doVJP2 f x) seeds'
      pure $ toTuple [v, dvs]
    def Text
"with_vjp" = 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
_ Value
arg ->
      -- XXX? We simply ignore the custom derivative. This is correct, but makes
      -- it more of a hassle to test them.
      SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. IsLocation a => a
noLoc Env
forall a. Monoid a => a
mempty Value
f Value
arg
    def Text
"acc" = Maybe TermBinding
forall a. Maybe a
Nothing
    def Text
fname
      | Text -> Name
nameFromText Text
fname 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
      | Just EvalM ValueType -> Value -> EvalM IOOp
op <- Text -> Maybe (EvalM ValueType -> Value -> EvalM IOOp)
forall (m :: * -> *).
Monad m =>
Text -> Maybe (m ValueType -> Value m -> m IOOp)
determineIO Text
fname =
          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 NoMode -> EvalM Value) -> TermBinding
TermPoly Maybe BoundV
forall a. Maybe a
Nothing ((TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding)
-> (TypeBase SizeClosure NoMode -> EvalM Value) -> TermBinding
forall a b. (a -> b) -> a -> b
$ \TypeBase SizeClosure NoMode
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
$ IOOp -> EvalM Value
doIO (IOOp -> EvalM Value)
-> (Value -> EvalM IOOp) -> Value -> EvalM Value
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< EvalM ValueType -> Value -> EvalM IOOp
op (TypeBase SizeClosure NoMode -> EvalM ValueType
evalTypeFully TypeBase SizeClosure NoMode
t) (Value -> EvalM IOOp)
-> (Value -> EvalM Value) -> Value -> EvalM IOOp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Value -> EvalM Value
force
    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
      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
      pure $ TypeConBinding mempty [] $ RetType [] $ Scalar $ Prim 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
"!"

-- | Interpret a single expression.
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

-- | Interpret an application of a function value to an argument value. It
-- better be well-typed!
interpretApply :: Ctx -> Value -> Value -> F ExtOp Value
interpretApply :: Ctx -> Value -> Value -> F ExtOp Value
interpretApply Ctx
ctx Value
f Value
arg =
  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
$ SrcLoc -> Env -> Value -> Value -> EvalM Value
apply SrcLoc
forall a. Monoid a => a
mempty (Ctx -> Env
ctxEnv Ctx
ctx) Value
f Value
arg

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 -> [DecBase Info VName] -> EvalM (Env, Env)
evalDecs (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.
    sizes <- extEnv
    pure $ 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 <- Ctx -> [DecBase Info VName] -> F ExtOp Env
interpretDecs Ctx
ctx [DecBase Info VName
d]
  pure 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 <- 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
  pure 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 NoMode -> TypeBase dim NoMode
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 NoMode -> ValueType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (PrimType -> ScalarTypeBase Int64 NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
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 o. TypeBase dim o -> TypeBase dim NoMode
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 (ZonkAny 0) -> Text) -> Doc (ZonkAny 0) -> Either Text ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 0) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 0) -> Either Text ())
-> Doc (ZonkAny 0) -> Either Text ()
forall a b. (a -> b) -> a -> b
$
        Doc (ZonkAny 0)
forall {ann}. Doc ann
expected
          Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 0)
"Got input of types"
          Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 0)] -> Doc (ZonkAny 0)
forall a. [Doc a] -> Doc a
stack ((StructType -> Doc (ZonkAny 0))
-> [StructType] -> [Doc (ZonkAny 0)]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> Doc (ZonkAny 0)
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty [StructType]
args_ts))
  where
    param_ts :: [ParamType]
param_ts = ((PName, ParamType) -> ParamType)
-> [(PName, ParamType)] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (PName, ParamType) -> ParamType
forall a b. (a, b) -> b
snd ([(PName, ParamType)] -> [ParamType])
-> [(PName, ParamType)] -> [ParamType]
forall a b. (a -> b) -> a -> b
$ ([(PName, ParamType)], StructType) -> [(PName, ParamType)]
forall a b. (a, b) -> a
fst (([(PName, ParamType)], StructType) -> [(PName, ParamType)])
-> ([(PName, ParamType)], StructType) -> [(PName, ParamType)]
forall a b. (a -> b) -> a -> b
$ StructType -> ([(PName, ParamType)], StructType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
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

  (ft, 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
      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
      pure (ft, pure v)
    Just (TermPoly (Just (T.BoundV [TypeParam]
_ StructType
t)) TypeBase SizeClosure NoMode -> EvalM Value
v) -> do
      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
      pure (ft, v (structToEval env 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 -> Value m) -> [Value] -> [Value m]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Value m
forall (m :: * -> *). Value -> Value m
fromDataValue [Value]
vs

  checkEntryArgs fname vs ft

  Right $ runEvalM (ctxImports ctx) $ do
    f <- mkf
    foldM (apply noLoc mempty) f vs'
  where
    updateType :: [ValueType] -> StructType -> Either Text StructType
updateType (ValueType
vt : [ValueType]
vts) (Scalar (Arrow NoMode
als PName
pn Diet
d StructType
pt (RetType [VName]
dims TypeBase (ExpBase Info VName) Freshness
rt))) = do
      ValueType -> StructType -> Either Text ()
checkInput ValueType
vt StructType
pt
      ScalarTypeBase (ExpBase Info VName) NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase (ExpBase Info VName) NoMode -> StructType)
-> (StructType -> ScalarTypeBase (ExpBase Info VName) NoMode)
-> StructType
-> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NoMode
-> PName
-> Diet
-> StructType
-> RetTypeBase (ExpBase Info VName) Freshness
-> ScalarTypeBase (ExpBase Info VName) NoMode
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow NoMode
als PName
pn Diet
d (ValueType -> StructType
valueStructType ValueType
vt) (RetTypeBase (ExpBase Info VName) Freshness
 -> ScalarTypeBase (ExpBase Info VName) NoMode)
-> (StructType -> RetTypeBase (ExpBase Info VName) Freshness)
-> StructType
-> ScalarTypeBase (ExpBase Info VName) NoMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName]
-> TypeBase (ExpBase Info VName) Freshness
-> RetTypeBase (ExpBase Info VName) Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims (TypeBase (ExpBase Info VName) Freshness
 -> RetTypeBase (ExpBase Info VName) Freshness)
-> (StructType -> TypeBase (ExpBase Info VName) Freshness)
-> StructType
-> RetTypeBase (ExpBase Info VName) Freshness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Freshness -> StructType -> TypeBase (ExpBase Info VName) Freshness
forall o.
Freshness
-> TypeBase (ExpBase Info VName) o
-> TypeBase (ExpBase Info VName) Freshness
toRes Freshness
Nonfresh
        (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) Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase (ExpBase Info VName) Freshness
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 NoMode
_ Shape Int64
_ (Prim PrimType
vt)) (Array NoMode
_ 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 NoMode
_ (F.Shape [Int64]
vd) ScalarTypeBase Int64 NoMode
_) pArr :: StructType
pArr@(Array NoMode
_ (F.Shape [ExpBase Info VName]
pd) ScalarTypeBase (ExpBase Info VName) NoMode
_)
      | [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 (ZonkAny 8) -> Text) -> Doc (ZonkAny 8) -> Either Text b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 8) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 8) -> Either Text b)
-> Doc (ZonkAny 8) -> Either Text b
forall a b. (a -> b) -> a -> b
$
        Doc (ZonkAny 8)
"Invalid argument type."
          Doc (ZonkAny 8) -> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 8)
"Expected:"
          Doc (ZonkAny 8) -> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Doc ann -> Doc ann
align (a -> Doc (ZonkAny 8)
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
pt)
          Doc (ZonkAny 8) -> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 8)
"Got:     "
          Doc (ZonkAny 8) -> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 8) -> Doc (ZonkAny 8)
forall ann. Doc ann -> Doc ann
align (a -> Doc (ZonkAny 8)
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 (ZonkAny 7) -> Text) -> Doc (ZonkAny 7) -> Either Text b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc (ZonkAny 7) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 7) -> Either Text b)
-> Doc (ZonkAny 7) -> Either Text b
forall a b. (a -> b) -> a -> b
$
        Doc (ZonkAny 7)
"Invalid argument dimensions."
          Doc (ZonkAny 7) -> Doc (ZonkAny 7) -> Doc (ZonkAny 7)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 7)
"Expected:"
          Doc (ZonkAny 7) -> Doc (ZonkAny 7) -> Doc (ZonkAny 7)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 7) -> Doc (ZonkAny 7)
forall ann. Doc ann -> Doc ann
align (a -> Doc (ZonkAny 7)
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
pd)
          Doc (ZonkAny 7) -> Doc (ZonkAny 7) -> Doc (ZonkAny 7)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 7)
"Got:     "
          Doc (ZonkAny 7) -> Doc (ZonkAny 7) -> Doc (ZonkAny 7)
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc (ZonkAny 7) -> Doc (ZonkAny 7)
forall ann. Doc ann -> Doc ann
align (a -> Doc (ZonkAny 7)
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
vd)