module Language.Futhark.Interpreter.FFI.Push
  ( put,
    putArgs,
    get,
    hasLazy,
    getLazy,
    lazyGet,
    ResShape (..),
    resultShape,
  )
where

import Control.Monad (zipWithM)
import Control.Monad.Except (throwError)
import Data.Array qualified as A
import Data.Map qualified as M
import Data.Text qualified as T
import Language.Futhark.Interpreter.FFI.ServerM
import Language.Futhark.Interpreter.Values qualified as I
import Language.Futhark.Syntax (Int64, Name, nameToString)

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

get :: I.ValueShape -> ValueRef -> ServerM (I.Value m)
get :: forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
get ValueShape
I.ShapeLeaf ValueRef
vr = PrimValue -> Value m
forall (m :: * -> *). PrimValue -> Value m
I.ValuePrim (PrimValue -> Value m) -> ServerM PrimValue -> ServerM (Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ValueRef -> ServerM PrimValue
getPrim ValueRef
vr
get rshp :: ValueShape
rshp@(I.ShapeDim {}) ValueRef
vr
  | ValueShape -> Bool
forall {d}. Shape d -> Bool
primArray ValueShape
rshp =
      -- Retrieving the array in one go is much faster than element by
      -- element.
      ServerM (Value m)
-> (Value -> ServerM (Value m)) -> Maybe Value -> ServerM (Value m)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (ValueShape -> ServerM (Value m)
forall (m :: * -> *). ValueShape -> ServerM (Value m)
getArray ValueShape
rshp) (Value m -> ServerM (Value m)
forall a. a -> ServerM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value m -> ServerM (Value m))
-> (Value -> Value m) -> Value -> ServerM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Value m
forall (m :: * -> *). Value -> Value m
I.fromDataValue) (Maybe Value -> ServerM (Value m))
-> ServerM (Maybe Value) -> ServerM (Value m)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ValueRef -> ServerM (Maybe Value)
getData ValueRef
vr
  | Bool
otherwise = ValueShape -> ServerM (Value m)
forall (m :: * -> *). ValueShape -> ServerM (Value m)
getArray ValueShape
rshp
  where
    primArray :: Shape d -> Bool
primArray (I.ShapeDim d
_ Shape d
shp) = Shape d -> Bool
primArray Shape d
shp
    primArray Shape d
I.ShapeLeaf = Bool
True
    primArray Shape d
_ = Bool
False
    getArray :: I.ValueShape -> ServerM (I.Value m)
    getArray :: forall (m :: * -> *). ValueShape -> ServerM (Value m)
getArray ValueShape
shp = ValueShape -> [Int64] -> ServerM (Value m)
forall (m :: * -> *). ValueShape -> [Int64] -> ServerM (Value m)
getArray' ValueShape
shp []
    getArray' :: I.ValueShape -> [Int64] -> ServerM (I.Value m)
    getArray' :: forall (m :: * -> *). ValueShape -> [Int64] -> ServerM (Value m)
getArray' shp :: ValueShape
shp@(I.ShapeDim Int64
n ValueShape
cshp) [Int64]
is = ValueShape -> Array Int (Value m) -> Value m
forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
I.ValueArray ValueShape
shp (Array Int (Value m) -> Value m)
-> ([Value m] -> Array Int (Value m)) -> [Value m] -> Value m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Value m] -> Array Int (Value m)
forall a. [a] -> Array Int a
toArray ([Value m] -> Value m) -> ServerM [Value m] -> ServerM (Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int64 -> ServerM (Value m)) -> [Int64] -> ServerM [Value m]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ValueShape -> [Int64] -> ServerM (Value m)
forall (m :: * -> *). ValueShape -> [Int64] -> ServerM (Value m)
getArray' ValueShape
cshp ([Int64] -> ServerM (Value m))
-> (Int64 -> [Int64]) -> Int64 -> ServerM (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
: [Int64]
is)) [Int64
0 .. Int64
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1]
    getArray' ValueShape
cshp [Int64]
is = [Int64] -> ValueRef -> ServerM ValueRef
index ([Int64] -> [Int64]
forall a. [a] -> [a]
reverse [Int64]
is) ValueRef
vr ServerM ValueRef
-> (ValueRef -> ServerM (Value m)) -> ServerM (Value m)
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 m)
forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
get ValueShape
cshp
get (I.ShapeRecord Map Name ValueShape
sm) ValueRef
vr =
  Map Name (Value m) -> Value m
forall (m :: * -> *). Map Name (Value m) -> Value m
I.ValueRecord
    (Map Name (Value m) -> Value m)
-> ServerM (Map Name (Value m)) -> ServerM (Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (ServerM (Value m)) -> ServerM (Map Name (Value m))
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 ((Name -> ValueShape -> ServerM (Value m))
-> Map Name ValueShape -> Map Name (ServerM (Value m))
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (\Name
fn ValueShape
cshp -> ValueRef -> Name -> ServerM ValueRef
project ValueRef
vr Name
fn ServerM ValueRef
-> (ValueRef -> ServerM (Value m)) -> ServerM (Value m)
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 m)
forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
get ValueShape
cshp) Map Name ValueShape
sm)
get shp :: ValueShape
shp@(I.ShapeSum Map Name [ValueShape]
sm) ValueRef
vr = do
  (vn, vrs) <- ValueRef -> ServerM (Name, [ValueRef])
destruct ValueRef
vr
  shps <- throwNothing ("Invalid variant " ++ nameToString vn ++ " in shape " ++ show sm ++ ".") $ M.lookup vn sm
  I.ValueSum shp vn <$> zipWithM get shps vrs

-- | Does this value contain any references to values residing on a server?
hasLazy :: I.Value m -> Bool
hasLazy :: forall (m :: * -> *). Value m -> Bool
hasLazy I.ValueLazyFFI {} = Bool
True
hasLazy (I.ValueArray ValueShape
_ Array Int (Value m)
arr) = (Value m -> Bool) -> [Value m] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Value m -> Bool
forall (m :: * -> *). Value m -> Bool
hasLazy ([Value m] -> Bool) -> [Value m] -> Bool
forall a b. (a -> b) -> a -> b
$ Array Int (Value m) -> [Value m]
forall i e. Array i e -> [e]
A.elems Array Int (Value m)
arr
hasLazy (I.ValueRecord Map Name (Value m)
fs) = (Value m -> Bool) -> Map Name (Value m) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Value m -> Bool
forall (m :: * -> *). Value m -> Bool
hasLazy Map Name (Value m)
fs
hasLazy (I.ValueSum ValueShape
_ Name
_ [Value m]
vs) = (Value m -> Bool) -> [Value m] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Value m -> Bool
forall (m :: * -> *). Value m -> Bool
hasLazy [Value m]
vs
hasLazy Value m
_ = Bool
False

-- | Get all the values residing on the server. The resulting 'I.Value' has no
-- 'I.ValueLazyFFI' in it.
getLazy :: I.Value a -> ServerM (I.Value a)
getLazy :: forall (a :: * -> *). Value a -> ServerM (Value a)
getLazy (I.ValueArray ValueShape
shp Array Int (Value a)
arr) = ValueShape -> Array Int (Value a) -> Value a
forall (m :: * -> *). ValueShape -> Array Int (Value m) -> Value m
I.ValueArray ValueShape
shp (Array Int (Value a) -> Value a)
-> ServerM (Array Int (Value a)) -> ServerM (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value a -> ServerM (Value a))
-> Array Int (Value a) -> ServerM (Array Int (Value a))
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) -> Array Int a -> m (Array Int b)
mapM Value a -> ServerM (Value a)
forall (a :: * -> *). Value a -> ServerM (Value a)
getLazy Array Int (Value a)
arr
getLazy (I.ValueRecord Map Name (Value a)
m) = Map Name (Value a) -> Value a
forall (m :: * -> *). Map Name (Value m) -> Value m
I.ValueRecord (Map Name (Value a) -> Value a)
-> ServerM (Map Name (Value a)) -> ServerM (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value a -> ServerM (Value a))
-> Map Name (Value a) -> ServerM (Map Name (Value a))
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) -> Map Name a -> m (Map Name b)
mapM Value a -> ServerM (Value a)
forall (a :: * -> *). Value a -> ServerM (Value a)
getLazy Map Name (Value a)
m
getLazy (I.ValueSum ValueShape
shp Name
vn [Value a]
vs) = ValueShape -> Name -> [Value a] -> Value a
forall (m :: * -> *). ValueShape -> Name -> [Value m] -> Value m
I.ValueSum ValueShape
shp Name
vn ([Value a] -> Value a) -> ServerM [Value a] -> ServerM (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value a -> ServerM (Value a)) -> [Value a] -> ServerM [Value a]
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 a -> ServerM (Value a)
forall (a :: * -> *). Value a -> ServerM (Value a)
getLazy [Value a]
vs
getLazy (I.ValueLazyFFI ValueShape
shp ValueRef
r [Int64]
os) = do
  v <- ValueShape -> ValueRef -> ServerM (Value a)
forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
get ValueShape
shp ValueRef
r
  pure $ foldl (\(I.ValueArray ValueShape
_ Array Int (Value a)
a) Int64
i -> Array Int (Value a)
a Array Int (Value a) -> Int -> Value a
forall i e. Ix i => Array i e -> i -> e
A.! Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i) v $ reverse os
getLazy Value a
v = Value a -> ServerM (Value a)
forall a. a -> ServerM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Value a
v

-- | As 'get', but arrays, including those in records, are left on the server.
lazyGet :: I.ValueShape -> ValueRef -> ServerM (I.Value m)
lazyGet :: forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
lazyGet shp :: ValueShape
shp@(I.ShapeDim {}) ValueRef
vr = Value m -> ServerM (Value m)
forall a. a -> ServerM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value m -> ServerM (Value m)) -> Value m -> ServerM (Value m)
forall a b. (a -> b) -> a -> b
$ ValueShape -> ValueRef -> [Int64] -> Value m
forall (m :: * -> *). ValueShape -> ValueRef -> [Int64] -> Value m
I.ValueLazyFFI ValueShape
shp ValueRef
vr []
lazyGet (I.ShapeRecord Map Name ValueShape
sm) ValueRef
vr =
  Map Name (Value m) -> Value m
forall (m :: * -> *). Map Name (Value m) -> Value m
I.ValueRecord
    (Map Name (Value m) -> Value m)
-> ServerM (Map Name (Value m)) -> ServerM (Value m)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (ServerM (Value m)) -> ServerM (Map Name (Value m))
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 ((Name -> ValueShape -> ServerM (Value m))
-> Map Name ValueShape -> Map Name (ServerM (Value m))
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (\Name
fn ValueShape
cshp -> ValueRef -> Name -> ServerM ValueRef
project ValueRef
vr Name
fn ServerM ValueRef
-> (ValueRef -> ServerM (Value m)) -> ServerM (Value m)
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 m)
forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
lazyGet ValueShape
cshp) Map Name ValueShape
sm)
lazyGet ValueShape
shp ValueRef
vr = ValueShape -> ValueRef -> ServerM (Value m)
forall (m :: * -> *). ValueShape -> ValueRef -> ServerM (Value m)
get ValueShape
shp ValueRef
vr

-- | Put these values on the server as the arguments of the given entry point.
putArgs :: Name -> [I.Value m] -> ServerM [ValueRef]
putArgs :: forall (m :: * -> *). Name -> [Value m] -> ServerM [ValueRef]
putArgs Name
fn [Value m]
vs = do
  ts <- Name -> ServerM [TypeName]
inputs Name
fn
  zipWithM put ts vs

put :: TypeName -> I.Value m -> ServerM ValueRef
put :: forall (m :: * -> *). TypeName -> Value m -> ServerM ValueRef
put TypeName
_ (I.ValuePrim PrimValue
p) = PrimValue -> ServerM ValueRef
putPrim PrimValue
p
put TypeName
tn pv :: Value m
pv@(I.ValueArray ValueShape
shp Array Int (Value m)
_)
  -- TODO: array elements residing on the server currently have to be fetched
  -- first, but this could be optimised further.
  | Value m -> Bool
forall (m :: * -> *). Value m -> Bool
hasLazy Value m
pv = TypeName -> Value m -> ServerM ValueRef
forall (m :: * -> *). TypeName -> Value m -> ServerM ValueRef
put TypeName
tn (Value m -> ServerM ValueRef)
-> ServerM (Value m) -> ServerM ValueRef
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value m -> ServerM (Value m)
forall (a :: * -> *). Value a -> ServerM (Value a)
getLazy Value m
pv
  | Bool
otherwise = do
      et <- TypeName -> ServerM TypeName
elemType TypeName
tn
      mapM (put et) (flatten pv) >>= mkArray tn (dims shp)
  where
    flatten :: I.Value m -> [I.Value m]
    flatten :: forall (m :: * -> *). Value m -> [Value m]
flatten (I.ValueArray ValueShape
_ Array Int (Value m)
a) = ([Value m] -> Value m -> [Value m])
-> [Value m] -> [Value m] -> [Value m]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[Value m]
o Value m
n -> [Value m]
o [Value m] -> [Value m] -> [Value m]
forall a. [a] -> [a] -> [a]
++ Value m -> [Value m]
forall (m :: * -> *). Value m -> [Value m]
flatten Value m
n) [] ([Value m] -> [Value m]) -> [Value m] -> [Value m]
forall a b. (a -> b) -> a -> b
$ Array Int (Value m) -> [Value m]
forall i e. Array i e -> [e]
A.elems Array Int (Value m)
a
    flatten Value m
v = [Value m
v]
    dims :: I.ValueShape -> [Int64]
    dims :: ValueShape -> [Int64]
dims (I.ShapeDim Int64
n ValueShape
cshp) = Int64
n Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
: ValueShape -> [Int64]
dims ValueShape
cshp
    dims ValueShape
_ = []
put TypeName
tn (I.ValueRecord Map Name (Value m)
vm) = do
  fm <- [(Name, TypeName)] -> Map Name TypeName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeName)] -> Map Name TypeName)
-> ServerM [(Name, TypeName)] -> ServerM (Map Name TypeName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeName -> ServerM [(Name, TypeName)]
fieldOrder TypeName
tn
  vrm <- sequence $ M.intersectionWith put fm vm
  mkRecord tn vrm
put TypeName
tn (I.ValueSum ValueShape
_ Name
vn [Value m]
vs) = do
  vts <-
    String -> Maybe [TypeName] -> ServerM [TypeName]
forall (m :: * -> *) a.
MonadError String m =>
String -> Maybe a -> m a
throwNothing (String
"Invalid variant " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
vn String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" in type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ TypeName -> String
T.unpack TypeName
tn String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
".") (Maybe [TypeName] -> ServerM [TypeName])
-> (Map Name [TypeName] -> Maybe [TypeName])
-> Map Name [TypeName]
-> ServerM [TypeName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Map Name [TypeName] -> Maybe [TypeName]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
vn
      (Map Name [TypeName] -> ServerM [TypeName])
-> ServerM (Map Name [TypeName]) -> ServerM [TypeName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TypeName -> ServerM (Map Name [TypeName])
variants TypeName
tn
  vrs <- zipWithM put vts vs
  mkSum tn vn vrs
-- The value already resides on the server.
put TypeName
_ (I.ValueLazyFFI ValueShape
_ ValueRef
r []) = ValueRef -> ServerM ValueRef
forall a. a -> ServerM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ValueRef
r
put TypeName
tn v :: Value m
v@(I.ValueLazyFFI {}) = do
  -- A partially indexed array cannot be constructed on the server, as the
  -- 'index' command only indexes all the way down to an element, so this
  -- one does have to go through the interpreter.
  iv <- Value m -> ServerM (Value m)
forall (a :: * -> *). Value a -> ServerM (Value a)
getLazy Value m
v
  put tn iv
put TypeName
_ Value m
v = String -> ServerM ValueRef
forall a. HasCallStack => String -> a
error (String -> ServerM ValueRef) -> String -> ServerM ValueRef
forall a b. (a -> b) -> a -> b
$ String
"Values of type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Value m -> String
forall a. Show a => a -> String
show Value m
v String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" are unsupported in FFI."

-- | How much of the shape of a value its type determines. Array dimensions may
-- be unknown at the type level, so they must be extracted from the actual
-- server-side value.
data ResShape
  = -- | An array of any rank; the argument describes its elements (after
    -- stripping array dimensions).
    ResArray ResShape
  | -- | A record, whose fields are described individually.
    ResRecord (M.Map Name ResShape)
  | -- | A sum type, whose constructor payloads are described individually.
    -- Only the payload of the constructor that is actually present can be
    -- inspected, so the resulting shape describes only that constructor.
    ResSum (M.Map Name [ResShape])
  | -- | Determined by the type alone.
    ResKnown I.ValueShape

-- | Determine the shape of a value residing on the server, consulting the
-- server for whatever the type did not settle.
resultShape :: ResShape -> ValueRef -> ServerM I.ValueShape
resultShape :: ResShape -> ValueRef -> ServerM ValueShape
resultShape (ResKnown ValueShape
shp) ValueRef
_ = ValueShape -> ServerM ValueShape
forall a. a -> ServerM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ValueShape
shp
resultShape (ResRecord Map Name ResShape
fs) ValueRef
vr =
  Map Name ValueShape -> ValueShape
forall d. Map Name (Shape d) -> Shape d
I.ShapeRecord (Map Name ValueShape -> ValueShape)
-> ServerM (Map Name ValueShape) -> ServerM ValueShape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> ResShape -> ServerM ValueShape)
-> Map Name ResShape -> ServerM (Map Name ValueShape)
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
M.traverseWithKey (\Name
f ResShape
shp -> ResShape -> ValueRef -> ServerM ValueShape
resultShape ResShape
shp (ValueRef -> ServerM ValueShape)
-> ServerM ValueRef -> ServerM ValueShape
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ValueRef -> Name -> ServerM ValueRef
project ValueRef
vr Name
f) Map Name ResShape
fs
resultShape (ResSum Map Name [ResShape]
cs) ValueRef
vr = do
  (c, vrs) <- ValueRef -> ServerM (Name, [ValueRef])
destruct ValueRef
vr
  payload <- throwNothing ("Unexpected variant " ++ nameToString c ++ ".") $ M.lookup c cs
  I.ShapeSum . M.singleton c <$> zipWithM resultShape payload vrs
resultShape (ResArray ResShape
eshp) ValueRef
vr = do
  dims <- ValueRef -> ServerM [Int64]
shape ValueRef
vr
  foldr I.ShapeDim <$> elemShape (length dims) eshp vr <*> pure dims

-- | The shape of the elements of an array. The elements cannot be inspected
-- one at a time, as there may not be any, so a record is instead unzipped
-- into one array per field - which has a shape even when it is empty. The
-- outer dimensions of those arrays are the ones we started with, and are
-- dropped again.
elemShape :: Int -> ResShape -> ValueRef -> ServerM I.ValueShape
elemShape :: Int -> ResShape -> ValueRef -> ServerM ValueShape
elemShape Int
_ (ResKnown ValueShape
shp) ValueRef
_ = ValueShape -> ServerM ValueShape
forall a. a -> ServerM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ValueShape
shp
elemShape Int
k (ResRecord Map Name ResShape
fs) ValueRef
arr = do
  etn <- TypeName -> ServerM TypeName
elemType (TypeName -> ServerM TypeName)
-> ServerM TypeName -> ServerM TypeName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ValueRef -> ServerM TypeName
vtype ValueRef
arr
  order <- fieldOrder etn
  refs <- unzipArray arr $ length order
  I.ShapeRecord . M.fromList <$> zipWithM onField order refs
  where
    onField :: (Name, b) -> ValueRef -> ServerM (Name, ValueShape)
onField (Name
f, b
_) ValueRef
ref = do
      shp <- ServerM ValueShape
-> (ResShape -> ServerM ValueShape)
-> Maybe ResShape
-> ServerM ValueShape
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Name -> ServerM ValueShape
forall {m :: * -> *} {a}. MonadError String m => Name -> m a
unknownField Name
f) (\ResShape
e -> ResShape -> ValueRef -> ServerM ValueShape
resultShape (ResShape -> ResShape
arrayOf ResShape
e) ValueRef
ref) (Maybe ResShape -> ServerM ValueShape)
-> Maybe ResShape -> ServerM ValueShape
forall a b. (a -> b) -> a -> b
$ Name -> Map Name ResShape -> Maybe ResShape
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name ResShape
fs
      pure (f, dropDims k shp)
    unknownField :: Name -> m a
unknownField Name
f =
      String -> m a
forall a. String -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> m a) -> String -> m a
forall a b. (a -> b) -> a -> b
$ String
"Unzipping produced unexpected field " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
nameToString Name
f String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"."
elemShape Int
_ (ResSum Map Name [ResShape]
_) ValueRef
_ =
  -- Unlike records, an array of sums cannot be split into arrays that have a
  -- shape even when empty.
  String -> ServerM ValueShape
forall a. String -> ServerM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError String
"Cannot determine the sizes of an array of sum types."
elemShape Int
_ (ResArray ResShape
_) ValueRef
_ =
  -- 'ResArray' covers every dimension at once, so it never describes the
  -- elements of an array.
  String -> ServerM ValueShape
forall a. String -> ServerM a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError String
"Array element is itself an array."

-- | The shape of an array whose elements are described by the argument.
arrayOf :: ResShape -> ResShape
arrayOf :: ResShape -> ResShape
arrayOf (ResArray ResShape
eshp) = ResShape -> ResShape
ResArray ResShape
eshp
arrayOf ResShape
shp = ResShape -> ResShape
ResArray ResShape
shp

dropDims :: Int -> I.ValueShape -> I.ValueShape
dropDims :: Int -> ValueShape -> ValueShape
dropDims Int
0 ValueShape
shp = ValueShape
shp
dropDims Int
k (I.ShapeDim Int64
_ ValueShape
shp) = Int -> ValueShape -> ValueShape
dropDims (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ValueShape
shp
dropDims Int
_ ValueShape
_ = String -> ValueShape
forall a. HasCallStack => String -> a
error String
"Unzipped field has too few dimensions."