-- 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
  { _vbContents :: [a]
  }

nil :: VecBuilder 0 a
nil = VecBuilder []

cons :: a -> VecBuilder n a -> VecBuilder (1+n) a
cons x (VecBuilder xs) = VecBuilder (x:xs)

vbToVec' :: SInt n -> VecBuilder n a -> Vec n a
vbToVec' n (VecBuilder c) = reifySInt n $ Vec.fromList 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 { _unVecView :: [a] }

uncons :: VecView (1+n) a -> (a, VecView n a)
uncons (VecView (x:xs)) = (x, VecView xs)
uncons _ = error "Internal error: invalid VecView."

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

vvFromVec :: Vec n a -> VecView n a
vvFromVec = VecView . 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.
  { _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 f (Batching go) = Batching $ \k m g rqs -> go k m (\vv x -> g vv (f x)) rqs
  {-# INLINE fmap #-}

instance Applicative (Batching rq rs) where
  pure x = Batching $ \k m g rqs -> k m (`g` x) rqs
  {-# INLINE pure #-}
  Batching f <*> Batching x =
    Batching $ \k m g rqs -> f (x k) m (\vv h -> g vv . h) rqs
  {-# INLINE (<*>) #-}
  liftA2 f (Batching x) (Batching y) = Batching $ \k m g rqs ->
    x (y k) m (\vv a b -> g vv (f a b)) rqs
  {-# INLINE liftA2 #-}

-- | Issue one request and retrieve its response.
request :: rq -> Batching rq rs rs
request rq = Batching $ \k n g !rqs -> k
  (sintVal `addSInt` n)
  -- Consume our response off the front of the VecView and pass it to the
  -- result function.
  (\rss -> let !(rs, rss') = uncons rss in g rss' rs)
  -- Add our request to the front of the VecBuilder.
  (rq `cons` 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
  { _bcCont :: A.Array rs -> Int -> (Int, a)
  }

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

instance Applicative (BulkCont rq rs) where
  pure x = BulkCont (\_ i -> (i, x))
  {-# INLINE pure #-}

  BulkCont kf <*> BulkCont kx = BulkCont (\a i ->
    let !(i', f) = kf a i
        !(i'', x) = kx a i'
    in  (i'', f x))
  {-# INLINE (<*>) #-}

  liftA2 f (BulkCont kx) (BulkCont ky) = BulkCont (\a i ->
    let !(i', x) = kx a i
        !(i'', y) = ky a i'
    in  (i'', f x y))
  {-# INLINE liftA2 #-}

rqBulkCont :: rq -> BulkCont rq rs rs
rqBulkCont _ = BulkCont $ \a i ->
  case A.indexArrayM a i of ((), x) -> (1 + i, x)
{-# INLINE rqBulkCont #-}

vvToArray :: forall n a. SInt n -> VecView n a -> A.Array a
vvToArray n (VecView v) = runST $
  A.newArray (unSInt n) undefined >>= \arr -> do
    sequenceA_ $ zipWith (A.writeArray arr) [0..] v
    A.unsafeFreezeArray 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 rqs0 =
  let !rqs = \rest -> F.foldr (:) rest rqs0
      !n = length rqs0
  in  withSInt n $ \sn ->
        Batching $ \k m g (VecBuilder rqs1) -> k
          (sn `addSInt` m)
          (\rss ->
            let !(rss0, rss1) = split sn rss
                !(_, rssT) =
                  _bcCont (traverse rqBulkCont rqs0) (vvToArray sn rss0) 0
            in  g rss1 rssT)
          (VecBuilder (rqs 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 f (Batching go) = go
  (\n q rqs -> q . vvFromVec <$> f (vbToVec' n rqs))
  sintVal
  (const id)
  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_ f = runIdentity . runBatching (Identity . 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