{-
  Copyright 2025 Arista Networks

  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 DataKinds #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}

-- | Presents right-associative folds as 'Foldable' sequences.
module Proto3.Wire.Encode.Repeated
  ( Repeated(..)
  , nullRepeated
  , ToRepeated(..)
  , mapRepeated
  ) where

import Data.Functor.Identity (Identity(..))
import Data.IntMap.Lazy qualified
import Data.IntSet qualified
import Data.Kind (Type)
import Data.List.NonEmpty qualified
import Data.Map.Lazy qualified
import Data.Sequence qualified
import Data.Set qualified
import Data.Vector qualified
import Data.Vector.Storable qualified
import Data.Vector.Unboxed qualified
import Foreign (Storable)
import GHC.Exts (Constraint, TYPE)
import GHC.Generics (Generic)
import Proto3.Wire.FoldR (FoldR(..), fromFoldR)

-- | Expresses a sequence of values /in reverse order/ for encoding as a repeated field.
type Repeated :: forall er . TYPE er -> Type
data Repeated e = ReverseRepeated
  { forall e. Repeated e -> Maybe Int
countRepeated :: Maybe Int
      -- ^ Optionally predicts the number of elements in the sequence.  Predict
      -- the count only when it is practical to do so accurately and quickly.
      --
      -- A prediction that is too low causes undefined behavior--possibly
      -- a crash.  A length prediction that is too high overallocates
      -- output space, as if the sequence really were that length.
  , forall e. Repeated e -> FoldR e
reverseRepeated :: FoldR e
      -- ^ A lazy right-associative fold over the /reverse/
      -- of the desired sequence of field values.
      --
      -- Design Note: We could have used a lazy left-associative fold, but
      -- vectors perform such folds using a left-to-right iteration, instead
      -- of the right-to-left iteration that would yield best performance.
      --
      -- Therefore in order to avoid accidental misuse of 'foldl', we ask
      -- for sequence reversal explicitly.  Thanks to vector fusion rules,
      -- it is fast to 'foldr' on the result of reversing a vector.
  }
  deriving stock ((forall a b. (a -> b) -> Repeated a -> Repeated b)
-> (forall a b. a -> Repeated b -> Repeated a) -> Functor Repeated
forall a b. a -> Repeated b -> Repeated a
forall a b. (a -> b) -> Repeated a -> Repeated 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) -> Repeated a -> Repeated b
fmap :: forall a b. (a -> b) -> Repeated a -> Repeated b
$c<$ :: forall a b. a -> Repeated b -> Repeated a
<$ :: forall a b. a -> Repeated b -> Repeated a
Functor, (forall x. Repeated e -> Rep (Repeated e) x)
-> (forall x. Rep (Repeated e) x -> Repeated e)
-> Generic (Repeated e)
forall x. Rep (Repeated e) x -> Repeated e
forall x. Repeated e -> Rep (Repeated e) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall e x. Rep (Repeated e) x -> Repeated e
forall e x. Repeated e -> Rep (Repeated e) x
$cfrom :: forall e x. Repeated e -> Rep (Repeated e) x
from :: forall x. Repeated e -> Rep (Repeated e) x
$cto :: forall e x. Rep (Repeated e) x -> Repeated e
to :: forall x. Rep (Repeated e) x -> Repeated e
Generic)

deriving stock instance Eq e => Eq (Repeated e)
deriving stock instance Read e => Read (Repeated e)
deriving stock instance Show e => Show (Repeated e)

nullRepeated :: Repeated e -> Bool
nullRepeated :: forall e. Repeated e -> Bool
nullRepeated Repeated e
c = FoldR e -> Bool
forall a. FoldR a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Repeated e -> FoldR e
forall e. Repeated e -> FoldR e
reverseRepeated Repeated e
c)
{-# INLINE nullRepeated #-}

-- | For each container type, specifies the optimal method for reverse iteration.
type ToRepeated :: forall cr . TYPE cr -> forall er . TYPE er -> Constraint
class ToRepeated c e | c -> e
  where
    -- | Converts to a reverse iteration over the elements.
    toRepeated :: c -> Repeated e

instance forall er (e :: TYPE er) .
         ToRepeated (Repeated e) e
  where
    toRepeated :: Repeated e -> Repeated e
toRepeated = Repeated e -> Repeated e
forall a. a -> a
id
    {-# INLINE toRepeated #-}

instance ToRepeated (Identity a) a
  where
    toRepeated :: Identity a -> Repeated a
toRepeated Identity a
x = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
1) ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> a -> b -> b
f (Identity a -> a
forall a. Identity a -> a
runIdentity Identity a
x) b
z))
    {-# INLINE toRepeated #-}

instance ToRepeated [a] a
  where
    toRepeated :: [a] -> Repeated a
toRepeated [a]
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated Maybe Int
forall a. Maybe a
Nothing ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> (b -> a -> b) -> b -> [a] -> b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> b -> b) -> b -> a -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
f) b
z [a]
xs))
      -- Unavoidably reads to the end of the list before presenting the last element.
    {-# INLINE toRepeated #-}

instance ToRepeated (Data.List.NonEmpty.NonEmpty a) a
  where
    toRepeated :: NonEmpty a -> Repeated a
toRepeated NonEmpty a
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated Maybe Int
forall a. Maybe a
Nothing ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> (b -> a -> b) -> b -> NonEmpty a -> b
forall b a. (b -> a -> b) -> b -> NonEmpty a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> b -> b) -> b -> a -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
f) b
z NonEmpty a
xs))
      -- Unavoidably reads to the end of the list before presenting the last element.
    {-# INLINE toRepeated #-}

instance ToRepeated (Data.Vector.Vector a) a
  where
    toRepeated :: Vector a -> Repeated a
toRepeated Vector a
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      (Int -> Maybe Int
forall a. a -> Maybe a
Just (Vector a -> Int
forall a. Vector a -> Int
Data.Vector.length Vector a
xs))
      (Vector a -> FoldR a
forall (t :: * -> *) a. Foldable t => t a -> FoldR a
fromFoldR (Vector a -> Vector a
forall a. Vector a -> Vector a
Data.Vector.reverse Vector a
xs))
      -- Vector fusion should convert this to right-to-left iteration.
    {-# INLINE toRepeated #-}

instance Storable a =>
         ToRepeated (Data.Vector.Storable.Vector a) a
  where
    toRepeated :: Vector a -> Repeated a
toRepeated Vector a
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      (Int -> Maybe Int
forall a. a -> Maybe a
Just (Vector a -> Int
forall a. Storable a => Vector a -> Int
Data.Vector.Storable.length Vector a
xs))
      ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> (a -> b -> b) -> b -> Vector a -> b
forall a b. Storable a => (a -> b -> b) -> b -> Vector a -> b
Data.Vector.Storable.foldr a -> b -> b
f b
z (Vector a -> Vector a
forall a. Storable a => Vector a -> Vector a
Data.Vector.Storable.reverse Vector a
xs)))
      -- Vector fusion should convert this to right-to-left iteration.
    {-# INLINE toRepeated #-}

instance Data.Vector.Unboxed.Unbox a =>
         ToRepeated (Data.Vector.Unboxed.Vector a) a
  where
    toRepeated :: Vector a -> Repeated a
toRepeated Vector a
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      (Int -> Maybe Int
forall a. a -> Maybe a
Just (Vector a -> Int
forall a. Unbox a => Vector a -> Int
Data.Vector.Unboxed.length Vector a
xs))
      ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> (a -> b -> b) -> b -> Vector a -> b
forall a b. Unbox a => (a -> b -> b) -> b -> Vector a -> b
Data.Vector.Unboxed.foldr a -> b -> b
f b
z (Vector a -> Vector a
forall a. Unbox a => Vector a -> Vector a
Data.Vector.Unboxed.reverse Vector a
xs)))
      -- Vector fusion should convert this to right-to-left iteration.
    {-# INLINE toRepeated #-}

instance ToRepeated (Data.Sequence.Seq a) a
  where
    toRepeated :: Seq a -> Repeated a
toRepeated Seq a
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      (Int -> Maybe Int
forall a. a -> Maybe a
Just (Seq a -> Int
forall a. Seq a -> Int
Data.Sequence.length Seq a
xs))
      ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> (b -> a -> b) -> b -> Seq a -> b
forall b a. (b -> a -> b) -> b -> Seq a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> b -> b) -> b -> a -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
f) b
z Seq a
xs))
      -- Should present the last element without having to read through the whole sequence.
    {-# INLINE toRepeated #-}

instance ToRepeated (Data.Set.Set a) a
  where
    toRepeated :: Set a -> Repeated a
toRepeated Set a
xs = Maybe Int -> FoldR a -> Repeated a
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      (Int -> Maybe Int
forall a. a -> Maybe a
Just (Set a -> Int
forall a. Set a -> Int
Data.Set.size Set a
xs))
      ((forall b. (a -> b -> b) -> b -> b) -> FoldR a
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\a -> b -> b
f b
z -> (b -> a -> b) -> b -> Set a -> b
forall b a. (b -> a -> b) -> b -> Set a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> b -> b) -> b -> a -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
f) b
z Set a
xs))
      -- Should present the last element without having to read through the whole sequence.
    {-# INLINE toRepeated #-}

instance ToRepeated Data.IntSet.IntSet Int
  where
    toRepeated :: IntSet -> Repeated Int
toRepeated IntSet
xs = Maybe Int -> FoldR Int -> Repeated Int
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated Maybe Int
forall a. Maybe a
Nothing ((forall b. (Int -> b -> b) -> b -> b) -> FoldR Int
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\Int -> b -> b
f b
z -> (b -> Int -> b) -> b -> IntSet -> b
forall a. (a -> Int -> a) -> a -> IntSet -> a
Data.IntSet.foldl ((Int -> b -> b) -> b -> Int -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> b -> b
f) b
z IntSet
xs))
      -- Should present the last element without having to read through the whole sequence.
    {-# INLINE toRepeated #-}

instance ToRepeated (Data.Map.Lazy.Map k a) (k, a)
  where
    toRepeated :: Map k a -> Repeated (k, a)
toRepeated Map k a
xs = Maybe Int -> FoldR (k, a) -> Repeated (k, a)
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      (Int -> Maybe Int
forall a. a -> Maybe a
Just (Map k a -> Int
forall k a. Map k a -> Int
Data.Map.Lazy.size Map k a
xs))
      ((forall b. ((k, a) -> b -> b) -> b -> b) -> FoldR (k, a)
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\(k, a) -> b -> b
f b
z -> (b -> k -> a -> b) -> b -> Map k a -> b
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
Data.Map.Lazy.foldlWithKey (\b
a k
k a
v -> (k, a) -> b -> b
f (k
k, a
v) b
a) b
z Map k a
xs))
      -- Should present the last key-value pair without having to read through the whole map.
    {-# INLINE toRepeated #-}

instance ToRepeated (Data.IntMap.Lazy.IntMap a) (Int, a)
  where
    toRepeated :: IntMap a -> Repeated (Int, a)
toRepeated IntMap a
xs = Maybe Int -> FoldR (Int, a) -> Repeated (Int, a)
forall e. Maybe Int -> FoldR e -> Repeated e
ReverseRepeated
      Maybe Int
forall a. Maybe a
Nothing
      ((forall b. ((Int, a) -> b -> b) -> b -> b) -> FoldR (Int, a)
forall a. (forall b. (a -> b -> b) -> b -> b) -> FoldR a
FoldR (\(Int, a) -> b -> b
f b
z -> (b -> Int -> a -> b) -> b -> IntMap a -> b
forall a b. (a -> Int -> b -> a) -> a -> IntMap b -> a
Data.IntMap.Lazy.foldlWithKey (\b
a Int
k a
v -> (Int, a) -> b -> b
f (Int
k, a
v) b
a) b
z IntMap a
xs))
      -- Should present the last key-value pair without having to read through the whole map.
    {-# INLINE toRepeated #-}

-- | A convenience function that maps a function over a sequence,
-- provided that the relevant types are all lifted.
mapRepeated ::
  forall (c :: Type) (e :: Type) (a :: Type) . ToRepeated c e => (e -> a) -> c -> Repeated a
mapRepeated :: forall c e a. ToRepeated c e => (e -> a) -> c -> Repeated a
mapRepeated e -> a
f c
xs = (e -> a) -> Repeated e -> Repeated a
forall a b. (a -> b) -> Repeated a -> Repeated b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap e -> a
f (c -> Repeated e
forall c e. ToRepeated c e => c -> Repeated e
toRepeated c
xs)
{-# INLINE mapRepeated #-}