-- Copyright 2019-2021 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

-- | An 'Applicative' for deferring "requests" to handle them all in bulk.

module Control.Batching
         ( Batching, request, batchRequest, runBatching, runBatching_
         ) where

import Control.Applicative (Applicative(..))
import Control.Monad.ST (runST)
import Data.Foldable (sequenceA_, toList)
import qualified Data.Foldable as F
import Data.Functor.Identity (Identity(..))
import GHC.TypeNats (type (+), Nat)

import qualified Data.Primitive.Array as A
import Data.SInt (SInt(unSInt), reifySInt, withSInt, sintVal, addSInt)
import Data.Vec.Short (Vec)
import qualified Data.Vec.Short as Vec

-- Quick-and-dirty Vec builder standin: a list.
--
-- Supports O(1) cons, O(n) conversion to Vec, and O(m) prepend of length-m to
-- length-n.
--
-- Note we won't ever left-associate appends, because we design the 'Batching'
-- type to avoid it, so O(m) prepend won't lead to any super-linear behavior.
newtype VecBuilder (n :: Nat) a = VecBuilder
  { VecBuilder n a -> [a]
_vbContents :: [a]
  }

nil :: VecBuilder 0 a
nil :: VecBuilder 0 a
nil = [a] -> VecBuilder 0 a
forall (n :: Nat) a. [a] -> VecBuilder n a
VecBuilder []

cons :: a -> VecBuilder n a -> VecBuilder (1+n) a
cons :: a -> VecBuilder n a -> VecBuilder (1 + n) a
cons a
x (VecBuilder [a]
xs) = [a] -> VecBuilder (1 + n) a
forall (n :: Nat) a. [a] -> VecBuilder n a
VecBuilder (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)

vbToVec' :: SInt n -> VecBuilder n a -> Vec n a
vbToVec' :: SInt n -> VecBuilder n a -> Vec n a
vbToVec' SInt n
n (VecBuilder [a]
c) = SInt n -> (KnownNat n => Vec n a) -> Vec n a
forall (n :: Nat) r. SInt n -> (KnownNat n => r) -> r
reifySInt SInt n
n ((KnownNat n => Vec n a) -> Vec n a)
-> (KnownNat n => Vec n a) -> Vec n a
forall a b. (a -> b) -> a -> b
$ [a] -> Vec n a
forall (n :: Nat) a. (HasCallStack, KnownNat n) => [a] -> Vec n a
Vec.fromList [a]
c

-- Quick-and-dirty Vec iterator standin: a list.
--
-- Supports O(1) uncons, O(m) bulk-uncons (split) of length m, and O(n)
-- conversion from Vec.
--
-- Note we won't ever recursively split, since splits are only generated by
-- 'batchRequest', and not by ('<*>'); so O(m) split won't lead to any
-- super-linear behavior.
newtype VecView (n :: Nat) a = VecView { VecView n a -> [a]
_unVecView :: [a] }

uncons :: VecView (1+n) a -> (a, VecView n a)
uncons :: VecView (1 + n) a -> (a, VecView n a)
uncons (VecView (a
x:[a]
xs)) = (a
x, [a] -> VecView n a
forall (n :: Nat) a. [a] -> VecView n a
VecView [a]
xs)
uncons VecView (1 + n) a
_ = [Char] -> (a, VecView n a)
forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error: invalid VecView."

split
  :: forall m n a. SInt m -> VecView (m+n) a -> (VecView m a, VecView n a)
split :: SInt m -> VecView (m + n) a -> (VecView m a, VecView n a)
split SInt m
m (VecView [a]
xs) =
  let ([a]
ma, [a]
na) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt (SInt m -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt m
m) [a]
xs
  in  ([a] -> VecView m a
forall (n :: Nat) a. [a] -> VecView n a
VecView [a]
ma, [a] -> VecView n a
forall (n :: Nat) a. [a] -> VecView n a
VecView [a]
na)

vvFromVec :: Vec n a -> VecView n a
vvFromVec :: Vec n a -> VecView n a
vvFromVec = [a] -> VecView n a
forall (n :: Nat) a. [a] -> VecView n a
VecView ([a] -> VecView n a) -> (Vec n a -> [a]) -> Vec n a -> VecView n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

-- TODO(awpr): consider factoring out a @SizedBatch n rq rs a@ type and
-- providing a scoped API for accessing it, e.g.:
--
-- withSizedBatch
--   :: Batching rq rs a
--   -> (forall n. KnownNat n => SizedBatch n rq rs a -> r)
--   -> r
-- getBatchRequests  :: SizedBatch n rq rs a -> Vec n rq
-- putBatchResponses :: SizedBatch n rq rs a -> Vec n rs -> a

-- -- Start with a simple version: contains the requests and a function to --
-- consume the responses.
-- data Batching0 rq rs a = forall n. Batching0
--   { n        :: SInt n
--   , requests :: VecBuilder n rq
--   , cont     :: VecView n rs -> a
--   }
--
-- -- First transformation: actually use the curried (<**>) type internally so
-- -- that we can right-associate all of the request prepends.
-- newtype Batching1 rq rs a = Batching1
--   { _unBatching1 :: forall r. Batching0 rq rs (a -> r) -> Batching0 rq rs r
--   }
--
-- -- Second transformation: turn the above function inside out, to enable the
-- -- third transformation.
-- newtype Batching2 rq rs a = Batching2
--   { _unBatching2
--       :: forall z r. (Batching0 rq rs r -> z) -> (Batching0 rq rs (a -> r) -> z)
--   }
--
-- -- Third transformation: curry both of the function types above to get rid
-- -- of a GADT constructor.
-- newtype Batching rq rs a = Batching
--   { _unBatching
--       :: forall z r
--        . (forall n. SInt n -> (VecView n rs -> r)      -> VecBuilder n rq -> z)
--       -> (forall m. SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
--   }
--
-- Now we have a newtype Batching that:
--   - takes a continuation for what to do with a VecBuilder of requests and a
--   VecView-consuming function returning @r@, with matching lengths.
--   - takes a VecView-consuming function for a (smaller) VecView returning
--   @a -> r@.
--   - takes a (smaller) VecBuilder of requests.
--   - wraps the VecView-consuming function with code to consume a prefix and
--   apply away the @a@ parameter.
--   - prepends some of its own requests to the VecBuilder.
--   - passes them on to the continuation and returns its result.

-- | The bulk request-response Applicative.
--
-- A value of type @Batching rq rs a@ describes a computation that gathers some
-- number of @rq@ request values, expects the same number of @rs@ response
-- values, and ultimately returns an @a@ result value derived from the
-- responses.
--
-- This can be used to apply an offline resource allocation algorithm to code
-- written as if allocation requests were satisfied incrementally.
--
-- This synergizes well with @-XApplicativeDo@, which allows using do-notation
-- for this type, as long as requests do not depend on earlier responses.
newtype Batching rq rs a = Batching
-- This is essentially the same as the Ap type from
-- Control.Applicative.Free.Fast, but specialized to lists of requests and
-- responses.
  { Batching rq rs a
-> forall z r.
   (forall (n :: Nat).
    SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
   -> forall (m :: Nat).
      SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
_unBulk
      :: forall z r
       . (forall n. SInt n -> (VecView n rs -> r)      -> VecBuilder n rq -> z)
      -> (forall m. SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
  }

instance Functor (Batching rq rs) where
  fmap :: (a -> b) -> Batching rq rs a -> Batching rq rs b
fmap a -> b
f (Batching forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
go) = (forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z)
-> Batching rq rs b
forall rq rs a.
(forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
Batching ((forall z r.
  (forall (n :: Nat).
   SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
  -> forall (m :: Nat).
     SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z)
 -> Batching rq rs b)
-> (forall z r.
    (forall (n :: Nat).
     SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
    -> forall (m :: Nat).
       SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z)
-> Batching rq rs b
forall a b. (a -> b) -> a -> b
$ \forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m VecView m rs -> b -> r
g VecBuilder m rq
rqs -> (forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
go forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m (\VecView m rs
vv a
x -> VecView m rs -> b -> r
g VecView m rs
vv (a -> b
f a
x)) VecBuilder m rq
rqs
  {-# INLINE fmap #-}

instance Applicative (Batching rq rs) where
  pure :: a -> Batching rq rs a
pure a
x = (forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
forall rq rs a.
(forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
Batching ((forall z r.
  (forall (n :: Nat).
   SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
  -> forall (m :: Nat).
     SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
 -> Batching rq rs a)
-> (forall z r.
    (forall (n :: Nat).
     SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
    -> forall (m :: Nat).
       SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
forall a b. (a -> b) -> a -> b
$ \forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m VecView m rs -> a -> r
g VecBuilder m rq
rqs -> SInt m -> (VecView m rs -> r) -> VecBuilder m rq -> z
forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m (VecView m rs -> a -> r
`g` a
x) VecBuilder m rq
rqs
  {-# INLINE pure #-}
  Batching forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> (a -> b) -> r) -> VecBuilder m rq -> z
f <*> :: Batching rq rs (a -> b) -> Batching rq rs a -> Batching rq rs b
<*> Batching forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
x =
    (forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z)
-> Batching rq rs b
forall rq rs a.
(forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
Batching ((forall z r.
  (forall (n :: Nat).
   SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
  -> forall (m :: Nat).
     SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z)
 -> Batching rq rs b)
-> (forall z r.
    (forall (n :: Nat).
     SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
    -> forall (m :: Nat).
       SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z)
-> Batching rq rs b
forall a b. (a -> b) -> a -> b
$ \forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m VecView m rs -> b -> r
g VecBuilder m rq
rqs -> (forall (n :: Nat).
 SInt n -> (VecView n rs -> a -> r) -> VecBuilder n rq -> z)
-> SInt m
-> (VecView m rs -> (a -> b) -> a -> r)
-> VecBuilder m rq
-> z
forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> (a -> b) -> r) -> VecBuilder m rq -> z
f ((forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (n :: Nat).
   SInt n -> (VecView n rs -> a -> r) -> VecBuilder n rq -> z
forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
x forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k) SInt m
m (\VecView m rs
vv a -> b
h -> VecView m rs -> b -> r
g VecView m rs
vv (b -> r) -> (a -> b) -> a -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
h) VecBuilder m rq
rqs
  {-# INLINE (<*>) #-}
  liftA2 :: (a -> b -> c)
-> Batching rq rs a -> Batching rq rs b -> Batching rq rs c
liftA2 a -> b -> c
f (Batching forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
x) (Batching forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z
y) = (forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> c -> r) -> VecBuilder m rq -> z)
-> Batching rq rs c
forall rq rs a.
(forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
Batching ((forall z r.
  (forall (n :: Nat).
   SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
  -> forall (m :: Nat).
     SInt m -> (VecView m rs -> c -> r) -> VecBuilder m rq -> z)
 -> Batching rq rs c)
-> (forall z r.
    (forall (n :: Nat).
     SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
    -> forall (m :: Nat).
       SInt m -> (VecView m rs -> c -> r) -> VecBuilder m rq -> z)
-> Batching rq rs c
forall a b. (a -> b) -> a -> b
$ \forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m VecView m rs -> c -> r
g VecBuilder m rq
rqs ->
    (forall (n :: Nat).
 SInt n -> (VecView n rs -> b -> r) -> VecBuilder n rq -> z)
-> SInt m -> (VecView m rs -> a -> b -> r) -> VecBuilder m rq -> z
forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
x ((forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (n :: Nat).
   SInt n -> (VecView n rs -> b -> r) -> VecBuilder n rq -> z
forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> b -> r) -> VecBuilder m rq -> z
y forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k) SInt m
m (\VecView m rs
vv a
a b
b -> VecView m rs -> c -> r
g VecView m rs
vv (a -> b -> c
f a
a b
b)) VecBuilder m rq
rqs
  {-# INLINE liftA2 #-}

-- | Issue one request and retrieve its response.
request :: rq -> Batching rq rs rs
request :: rq -> Batching rq rs rs
request rq
rq = (forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> rs -> r) -> VecBuilder m rq -> z)
-> Batching rq rs rs
forall rq rs a.
(forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
Batching ((forall z r.
  (forall (n :: Nat).
   SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
  -> forall (m :: Nat).
     SInt m -> (VecView m rs -> rs -> r) -> VecBuilder m rq -> z)
 -> Batching rq rs rs)
-> (forall z r.
    (forall (n :: Nat).
     SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
    -> forall (m :: Nat).
       SInt m -> (VecView m rs -> rs -> r) -> VecBuilder m rq -> z)
-> Batching rq rs rs
forall a b. (a -> b) -> a -> b
$ \forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
n VecView m rs -> rs -> r
g !VecBuilder m rq
rqs -> SInt (1 + m)
-> (VecView (1 + m) rs -> r) -> VecBuilder (1 + m) rq -> z
forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k
  (SInt 1
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal SInt 1 -> SInt m -> SInt (1 + m)
forall (m :: Nat) (n :: Nat).
HasCallStack =>
SInt m -> SInt n -> SInt (m + n)
`addSInt` SInt m
n)
  -- Consume our response off the front of the VecView and pass it to the
  -- result function.
  (\VecView (1 + m) rs
rss -> let !(rs
rs, VecView m rs
rss') = VecView (1 + m) rs -> (rs, VecView m rs)
forall (n :: Nat) a. VecView (1 + n) a -> (a, VecView n a)
uncons VecView (1 + m) rs
rss in VecView m rs -> rs -> r
g VecView m rs
rss' rs
rs)
  -- Add our request to the front of the VecBuilder.
  (rq
rq rq -> VecBuilder m rq -> VecBuilder (1 + m) rq
forall a (n :: Nat). a -> VecBuilder n a -> VecBuilder (1 + n) a
`cons` VecBuilder m rq
rqs)
{-# INLINE request #-}

-- Specialized implementation for issuing Traversables of requests at once.
-- The goal here is to separate the traversal gathering requests from the
-- traversal building up the result, so we don't have to allocate a bunch of
-- memory to hold a closure that will build the resulting data structure while
-- we're handling the requests.
newtype BulkCont rq rs a = BulkCont
  { BulkCont rq rs a -> Array rs -> Int -> (Int, a)
_bcCont :: A.Array rs -> Int -> (Int, a)
  }

instance Functor (BulkCont rq rs) where
  fmap :: (a -> b) -> BulkCont rq rs a -> BulkCont rq rs b
fmap a -> b
f (BulkCont Array rs -> Int -> (Int, a)
k) = (Array rs -> Int -> (Int, b)) -> BulkCont rq rs b
forall rq rs a. (Array rs -> Int -> (Int, a)) -> BulkCont rq rs a
BulkCont (\Array rs
a Int
i ->
    let !(Int
i', a
x) = Array rs -> Int -> (Int, a)
k Array rs
a Int
i
    in  (Int
i', a -> b
f a
x))
  {-# INLINE fmap #-}

instance Applicative (BulkCont rq rs) where
  pure :: a -> BulkCont rq rs a
pure a
x = (Array rs -> Int -> (Int, a)) -> BulkCont rq rs a
forall rq rs a. (Array rs -> Int -> (Int, a)) -> BulkCont rq rs a
BulkCont (\Array rs
_ Int
i -> (Int
i, a
x))
  {-# INLINE pure #-}

  BulkCont Array rs -> Int -> (Int, a -> b)
kf <*> :: BulkCont rq rs (a -> b) -> BulkCont rq rs a -> BulkCont rq rs b
<*> BulkCont Array rs -> Int -> (Int, a)
kx = (Array rs -> Int -> (Int, b)) -> BulkCont rq rs b
forall rq rs a. (Array rs -> Int -> (Int, a)) -> BulkCont rq rs a
BulkCont (\Array rs
a Int
i ->
    let !(Int
i', a -> b
f) = Array rs -> Int -> (Int, a -> b)
kf Array rs
a Int
i
        !(Int
i'', a
x) = Array rs -> Int -> (Int, a)
kx Array rs
a Int
i'
    in  (Int
i'', a -> b
f a
x))
  {-# INLINE (<*>) #-}

  liftA2 :: (a -> b -> c)
-> BulkCont rq rs a -> BulkCont rq rs b -> BulkCont rq rs c
liftA2 a -> b -> c
f (BulkCont Array rs -> Int -> (Int, a)
kx) (BulkCont Array rs -> Int -> (Int, b)
ky) = (Array rs -> Int -> (Int, c)) -> BulkCont rq rs c
forall rq rs a. (Array rs -> Int -> (Int, a)) -> BulkCont rq rs a
BulkCont (\Array rs
a Int
i ->
    let !(Int
i', a
x) = Array rs -> Int -> (Int, a)
kx Array rs
a Int
i
        !(Int
i'', b
y) = Array rs -> Int -> (Int, b)
ky Array rs
a Int
i'
    in  (Int
i'', a -> b -> c
f a
x b
y))
  {-# INLINE liftA2 #-}

rqBulkCont :: rq -> BulkCont rq rs rs
rqBulkCont :: rq -> BulkCont rq rs rs
rqBulkCont rq
_ = (Array rs -> Int -> (Int, rs)) -> BulkCont rq rs rs
forall rq rs a. (Array rs -> Int -> (Int, a)) -> BulkCont rq rs a
BulkCont ((Array rs -> Int -> (Int, rs)) -> BulkCont rq rs rs)
-> (Array rs -> Int -> (Int, rs)) -> BulkCont rq rs rs
forall a b. (a -> b) -> a -> b
$ \Array rs
a Int
i ->
  case Array rs -> Int -> ((), rs)
forall (m :: * -> *) a. Monad m => Array a -> Int -> m a
A.indexArrayM Array rs
a Int
i of ((), rs
x) -> (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i, rs
x)
{-# INLINE rqBulkCont #-}

vvToArray :: forall n a. SInt n -> VecView n a -> A.Array a
vvToArray :: SInt n -> VecView n a -> Array a
vvToArray SInt n
n (VecView [a]
v) = (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array a)) -> Array a)
-> (forall s. ST s (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$
  Int -> a -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
A.newArray (SInt n -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt n
n) a
forall a. HasCallStack => a
undefined ST s (MutableArray s a)
-> (MutableArray s a -> ST s (Array a)) -> ST s (Array a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MutableArray s a
arr -> do
    [ST s ()] -> ST s ()
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Applicative f) =>
t (f a) -> f ()
sequenceA_ ([ST s ()] -> ST s ()) -> [ST s ()] -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Int -> a -> ST s ()) -> [Int] -> [a] -> [ST s ()]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
A.writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr) [Int
0..] [a]
v
    MutableArray (PrimState (ST s)) a -> ST s (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
A.unsafeFreezeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr

-- | Issue a Traversable of requests and retrieve their responses.
batchRequest :: forall t rq rs. Traversable t => t rq -> Batching rq rs (t rs)
batchRequest :: t rq -> Batching rq rs (t rs)
batchRequest t rq
rqs0 =
  let !rqs :: [rq] -> [rq]
rqs = \[rq]
rest -> (rq -> [rq] -> [rq]) -> [rq] -> t rq -> [rq]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr (:) [rq]
rest t rq
rqs0
      !n :: Int
n = t rq -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t rq
rqs0
  in  Int
-> (forall (n :: Nat). SInt n -> Batching rq rs (t rs))
-> Batching rq rs (t rs)
forall r.
HasCallStack =>
Int -> (forall (n :: Nat). SInt n -> r) -> r
withSInt Int
n ((forall (n :: Nat). SInt n -> Batching rq rs (t rs))
 -> Batching rq rs (t rs))
-> (forall (n :: Nat). SInt n -> Batching rq rs (t rs))
-> Batching rq rs (t rs)
forall a b. (a -> b) -> a -> b
$ \SInt n
sn ->
        (forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> t rs -> r) -> VecBuilder m rq -> z)
-> Batching rq rs (t rs)
forall rq rs a.
(forall z r.
 (forall (n :: Nat).
  SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
 -> forall (m :: Nat).
    SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z)
-> Batching rq rs a
Batching ((forall z r.
  (forall (n :: Nat).
   SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
  -> forall (m :: Nat).
     SInt m -> (VecView m rs -> t rs -> r) -> VecBuilder m rq -> z)
 -> Batching rq rs (t rs))
-> (forall z r.
    (forall (n :: Nat).
     SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
    -> forall (m :: Nat).
       SInt m -> (VecView m rs -> t rs -> r) -> VecBuilder m rq -> z)
-> Batching rq rs (t rs)
forall a b. (a -> b) -> a -> b
$ \forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k SInt m
m VecView m rs -> t rs -> r
g (VecBuilder [rq]
rqs1) -> SInt (n + m)
-> (VecView (n + m) rs -> r) -> VecBuilder (n + m) rq -> z
forall (n :: Nat).
SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z
k
          (SInt n
sn SInt n -> SInt m -> SInt (n + m)
forall (m :: Nat) (n :: Nat).
HasCallStack =>
SInt m -> SInt n -> SInt (m + n)
`addSInt` SInt m
m)
          (\VecView (n + m) rs
rss ->
            let !(VecView n rs
rss0, VecView m rs
rss1) = SInt n -> VecView (n + m) rs -> (VecView n rs, VecView m rs)
forall (m :: Nat) (n :: Nat) a.
SInt m -> VecView (m + n) a -> (VecView m a, VecView n a)
split SInt n
sn VecView (n + m) rs
rss
                !(Int
_, t rs
rssT) =
                  BulkCont rq rs (t rs) -> Array rs -> Int -> (Int, t rs)
forall rq rs a. BulkCont rq rs a -> Array rs -> Int -> (Int, a)
_bcCont ((rq -> BulkCont rq rs rs) -> t rq -> BulkCont rq rs (t rs)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse rq -> BulkCont rq rs rs
forall rq rs. rq -> BulkCont rq rs rs
rqBulkCont t rq
rqs0) (SInt n -> VecView n rs -> Array rs
forall (n :: Nat) a. SInt n -> VecView n a -> Array a
vvToArray SInt n
sn VecView n rs
rss0) Int
0
            in  VecView m rs -> t rs -> r
g VecView m rs
rss1 t rs
rssT)
          ([rq] -> VecBuilder (n + m) rq
forall (n :: Nat) a. [a] -> VecBuilder n a
VecBuilder ([rq] -> [rq]
rqs [rq]
rqs1))
{-# INLINE batchRequest #-}

-- | Given an allocator function in any 'Functor', run a 'Batching' computation.
runBatching
  :: Functor f
  => (forall n. Vec n rq -> f (Vec n rs))
  -> Batching rq rs a -> f a
runBatching :: (forall (n :: Nat). Vec n rq -> f (Vec n rs))
-> Batching rq rs a -> f a
runBatching forall (n :: Nat). Vec n rq -> f (Vec n rs)
f (Batching forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
go) = (forall (n :: Nat).
 SInt n -> (VecView n rs -> a) -> VecBuilder n rq -> f a)
-> SInt 0 -> (VecView 0 rs -> a -> a) -> VecBuilder 0 rq -> f a
forall z r.
(forall (n :: Nat).
 SInt n -> (VecView n rs -> r) -> VecBuilder n rq -> z)
-> forall (m :: Nat).
   SInt m -> (VecView m rs -> a -> r) -> VecBuilder m rq -> z
go
  (\SInt n
n VecView n rs -> a
q VecBuilder n rq
rqs -> VecView n rs -> a
q (VecView n rs -> a) -> (Vec n rs -> VecView n rs) -> Vec n rs -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n rs -> VecView n rs
forall (n :: Nat) a. Vec n a -> VecView n a
vvFromVec (Vec n rs -> a) -> f (Vec n rs) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vec n rq -> f (Vec n rs)
forall (n :: Nat). Vec n rq -> f (Vec n rs)
f (SInt n -> VecBuilder n rq -> Vec n rq
forall (n :: Nat) a. SInt n -> VecBuilder n a -> Vec n a
vbToVec' SInt n
n VecBuilder n rq
rqs))
  SInt 0
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal
  ((a -> a) -> VecView 0 rs -> a -> a
forall a b. a -> b -> a
const a -> a
forall a. a -> a
id)
  VecBuilder 0 rq
forall a. VecBuilder 0 a
nil
{-# INLINE runBatching #-}

-- | Like 'runBatching', but without a 'Functor' (or implicitly in 'Identity').
runBatching_ :: (forall n. Vec n rq -> Vec n rs) -> Batching rq rs a -> a
runBatching_ :: (forall (n :: Nat). Vec n rq -> Vec n rs) -> Batching rq rs a -> a
runBatching_ forall (n :: Nat). Vec n rq -> Vec n rs
f = Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a)
-> (Batching rq rs a -> Identity a) -> Batching rq rs a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (n :: Nat). Vec n rq -> Identity (Vec n rs))
-> Batching rq rs a -> Identity a
forall (f :: * -> *) rq rs a.
Functor f =>
(forall (n :: Nat). Vec n rq -> f (Vec n rs))
-> Batching rq rs a -> f a
runBatching (Vec n rs -> Identity (Vec n rs)
forall a. a -> Identity a
Identity (Vec n rs -> Identity (Vec n rs))
-> (Vec n rq -> Vec n rs) -> Vec n rq -> Identity (Vec n rs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n rq -> Vec n rs
forall (n :: Nat). Vec n rq -> Vec n rs
f)
{-# INLINE runBatching_ #-}

-- TODO(awpr): consider adding a Batched monad that supports many batches of
-- requests:
--
-- type Batched rq rs a = Free (Batching rq rs) a
--
-- -- Use ApplicativeDo inside this to group requests into a batch.
-- batch :: Batching rq rs a -> Batched rq rs a
--
-- -- Note this requires Monad rather than just Functor, since it can have
-- multiple batches.
-- runBatched
--   :: Monad f
--   => (forall n. KnownNat n => Vec n rq -> f (Vec n rs))
--   -> Batched rq rs a -> f a