-- Copyright 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.

-- | Extends 'Representable' with support for modifying elements.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Data.Functor.Update (Update(..), updateRep, ixRep, GUpdate(..)) where

import Data.Coerce (coerce)
import Data.Functor ((<&>))
import GHC.Generics
         ( Generic1(..)
         , (:*:)(..), (:.:)(..)
         , M1(..), Rec1(..), U1(..), Par1(..)
         )

import Data.Functor.Rep (Representable(..))

import Data.Functor.Field (Field(..), GTabulate(..), FieldRep(..))
import Data.Ten.Internal (mapStarFst, mapStarSnd)

-- | Extends 'Representable' with support for modifying elements.
--
-- If @'Eq' ('Rep' f)@ is available, this is already possible by roundabout
-- means, but this class lets instances provide a more direct method that
-- doesn't require per-field equality tests.
class Representable f => Update f where
  overRep :: Rep f -> (a -> a) -> f a -> f a

-- | Update an @f a@ at a given index.
updateRep :: Update f => Rep f -> a -> f a -> f a
updateRep :: Rep f -> a -> f a -> f a
updateRep Rep f
i = Rep f -> (a -> a) -> f a -> f a
forall (f :: * -> *) a. Update f => Rep f -> (a -> a) -> f a -> f a
overRep Rep f
i ((a -> a) -> f a -> f a) -> (a -> a -> a) -> a -> f a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> a
forall a b. a -> b -> a
const

-- | A 'Control.Lens.Lens' to the field identified by a given 'Rep'.
--
-- @
--     ixRep :: Update f => Rep f -> Lens' (f a) a
-- @
ixRep :: (Update f, Functor m) => Rep f -> (a -> m a) -> f a -> m (f a)
ixRep :: Rep f -> (a -> m a) -> f a -> m (f a)
ixRep Rep f
i a -> m a
f = \f a
fa -> a -> m a
f (f a -> Rep f -> a
forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f a
fa Rep f
i) m a -> (a -> f a) -> m (f a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
ma -> Rep f -> a -> f a -> f a
forall (f :: * -> *) a. Update f => Rep f -> a -> f a -> f a
updateRep Rep f
i a
ma f a
fa

instance (Generic1 f, GTabulate (Rep1 f), GUpdate (Rep1 f), Functor f)
      => Update (FieldRep f) where
  overRep :: Rep (FieldRep f) -> (a -> a) -> FieldRep f a -> FieldRep f a
overRep =
    \Rep (FieldRep f)
i a -> a
f (FieldRep f a
fa) -> f a -> FieldRep f a
forall (f :: * -> *) a. f a -> FieldRep f a
FieldRep (f a -> FieldRep f a) -> f a -> FieldRep f a
forall a b. (a -> b) -> a -> b
$ FieldSetter f -> (a -> a) -> f a -> f a
forall (f :: * -> *).
FieldSetter f -> forall a. (a -> a) -> f a -> f a
runFS (Field f -> f (FieldSetter f) -> FieldSetter f
forall (f :: * -> *). Field f -> forall a. f a -> a
getField Rep (FieldRep f)
Field f
i f (FieldSetter f)
setters_) a -> a
f f a
fa
   where
    setters_ :: f (FieldSetter f)
    setters_ :: f (FieldSetter f)
setters_ = f (FieldSetter f)
forall (f :: * -> *).
(Generic1 f, GUpdate (Rep1 f)) =>
f (FieldSetter f)
setters

-- | The 'Generic1' implementation of 'Update'.
class GUpdate rec where
  gsetters :: ((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r

instance GUpdate U1 where
  gsetters :: ((forall a. (a -> a) -> U1 a -> U1 a) -> r) -> U1 r
gsetters (forall a. (a -> a) -> U1 a -> U1 a) -> r
_r = U1 r
forall k (p :: k). U1 p
U1
  {-# INLINE gsetters #-}

instance GUpdate rec => GUpdate (Rec1 rec) where
  gsetters :: ((forall a. (a -> a) -> Rec1 rec a -> Rec1 rec a) -> r)
-> Rec1 rec r
gsetters (forall a. (a -> a) -> Rec1 rec a -> Rec1 rec a) -> r
r = rec r -> Rec1 rec r
forall k (f :: k -> *) (p :: k). f p -> Rec1 f p
Rec1 (rec r -> Rec1 rec r) -> rec r -> Rec1 rec r
forall a b. (a -> b) -> a -> b
$ ((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (\forall a. (a -> a) -> rec a -> rec a
s -> (forall a. (a -> a) -> Rec1 rec a -> Rec1 rec a) -> r
r ((forall a. (a -> a) -> Rec1 rec a -> Rec1 rec a) -> r)
-> (forall a. (a -> a) -> Rec1 rec a -> Rec1 rec a) -> r
forall a b. (a -> b) -> a -> b
$ \a -> a
f -> rec a -> Rec1 rec a
forall k (f :: k -> *) (p :: k). f p -> Rec1 f p
Rec1 (rec a -> Rec1 rec a)
-> (Rec1 rec a -> rec a) -> Rec1 rec a -> Rec1 rec a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> rec a -> rec a
forall a. (a -> a) -> rec a -> rec a
s a -> a
f (rec a -> rec a) -> (Rec1 rec a -> rec a) -> Rec1 rec a -> rec a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rec1 rec a -> rec a
forall k (f :: k -> *) (p :: k). Rec1 f p -> f p
unRec1)
  {-# INLINE gsetters #-}

instance GUpdate rec => GUpdate (M1 k i rec) where
  gsetters :: ((forall a. (a -> a) -> M1 k i rec a -> M1 k i rec a) -> r)
-> M1 k i rec r
gsetters (forall a. (a -> a) -> M1 k i rec a -> M1 k i rec a) -> r
r = rec r -> M1 k i rec r
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 (rec r -> M1 k i rec r) -> rec r -> M1 k i rec r
forall a b. (a -> b) -> a -> b
$ ((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (\forall a. (a -> a) -> rec a -> rec a
s -> (forall a. (a -> a) -> M1 k i rec a -> M1 k i rec a) -> r
r ((forall a. (a -> a) -> M1 k i rec a -> M1 k i rec a) -> r)
-> (forall a. (a -> a) -> M1 k i rec a -> M1 k i rec a) -> r
forall a b. (a -> b) -> a -> b
$ \a -> a
f -> rec a -> M1 k i rec a
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 (rec a -> M1 k i rec a)
-> (M1 k i rec a -> rec a) -> M1 k i rec a -> M1 k i rec a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> rec a -> rec a
forall a. (a -> a) -> rec a -> rec a
s a -> a
f (rec a -> rec a)
-> (M1 k i rec a -> rec a) -> M1 k i rec a -> rec a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. M1 k i rec a -> rec a
forall i (c :: Meta) k (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1)
  {-# INLINE gsetters #-}

instance GUpdate Par1 where
  gsetters :: ((forall a. (a -> a) -> Par1 a -> Par1 a) -> r) -> Par1 r
gsetters (forall a. (a -> a) -> Par1 a -> Par1 a) -> r
r = r -> Par1 r
forall p. p -> Par1 p
Par1 (r -> Par1 r) -> r -> Par1 r
forall a b. (a -> b) -> a -> b
$ (forall a. (a -> a) -> Par1 a -> Par1 a) -> r
r ((forall a. (a -> a) -> Par1 a -> Par1 a) -> r)
-> (forall a. (a -> a) -> Par1 a -> Par1 a) -> r
forall a b. (a -> b) -> a -> b
$ \a -> a
f -> a -> Par1 a
forall p. p -> Par1 p
Par1 (a -> Par1 a) -> (Par1 a -> a) -> Par1 a -> Par1 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f (a -> a) -> (Par1 a -> a) -> Par1 a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Par1 a -> a
forall p. Par1 p -> p
unPar1
  {-# INLINE gsetters #-}

instance (GUpdate f, GUpdate g) => GUpdate (f :*: g) where
  gsetters :: ((forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r)
-> (:*:) f g r
gsetters (forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r
r = f r
fs f r -> g r -> (:*:) f g r
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: g r
gs
   where
    fs :: f r
fs = ((forall a. (a -> a) -> f a -> f a) -> r) -> f r
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (((forall a. (a -> a) -> f a -> f a) -> r) -> f r)
-> ((forall a. (a -> a) -> f a -> f a) -> r) -> f r
forall a b. (a -> b) -> a -> b
$ \forall a. (a -> a) -> f a -> f a
s -> (forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r
r ((forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r)
-> (forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r
forall a b. (a -> b) -> a -> b
$ (f a -> f a) -> (:*:) f g a -> (:*:) f g a
forall k (f :: k -> *) (m :: k) (g :: k -> *).
(f m -> f m) -> (:*:) f g m -> (:*:) f g m
mapStarFst ((f a -> f a) -> (:*:) f g a -> (:*:) f g a)
-> ((a -> a) -> f a -> f a)
-> (a -> a)
-> (:*:) f g a
-> (:*:) f g a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> f a -> f a
forall a. (a -> a) -> f a -> f a
s
    gs :: g r
gs = ((forall a. (a -> a) -> g a -> g a) -> r) -> g r
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (((forall a. (a -> a) -> g a -> g a) -> r) -> g r)
-> ((forall a. (a -> a) -> g a -> g a) -> r) -> g r
forall a b. (a -> b) -> a -> b
$ \forall a. (a -> a) -> g a -> g a
s -> (forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r
r ((forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r)
-> (forall a. (a -> a) -> (:*:) f g a -> (:*:) f g a) -> r
forall a b. (a -> b) -> a -> b
$ (g a -> g a) -> (:*:) f g a -> (:*:) f g a
forall k (g :: k -> *) (m :: k) (f :: k -> *).
(g m -> g m) -> (:*:) f g m -> (:*:) f g m
mapStarSnd ((g a -> g a) -> (:*:) f g a -> (:*:) f g a)
-> ((a -> a) -> g a -> g a)
-> (a -> a)
-> (:*:) f g a
-> (:*:) f g a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> g a -> g a
forall a. (a -> a) -> g a -> g a
s
  {-# INLINE gsetters #-}

instance (GUpdate f, GUpdate g) => GUpdate (f :.: g) where
  gsetters :: ((forall a. (a -> a) -> (:.:) f g a -> (:.:) f g a) -> r)
-> (:.:) f g r
gsetters (forall a. (a -> a) -> (:.:) f g a -> (:.:) f g a) -> r
r = f (g r) -> (:.:) f g r
forall k2 k1 (f :: k2 -> *) (g :: k1 -> k2) (p :: k1).
f (g p) -> (:.:) f g p
Comp1 (f (g r) -> (:.:) f g r) -> f (g r) -> (:.:) f g r
forall a b. (a -> b) -> a -> b
$
    ((forall a. (a -> a) -> f a -> f a) -> g r) -> f (g r)
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (((forall a. (a -> a) -> f a -> f a) -> g r) -> f (g r))
-> ((forall a. (a -> a) -> f a -> f a) -> g r) -> f (g r)
forall a b. (a -> b) -> a -> b
$ \ forall a. (a -> a) -> f a -> f a
s0 ->
    ((forall a. (a -> a) -> g a -> g a) -> r) -> g r
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (((forall a. (a -> a) -> g a -> g a) -> r) -> g r)
-> ((forall a. (a -> a) -> g a -> g a) -> r) -> g r
forall a b. (a -> b) -> a -> b
$ \ forall a. (a -> a) -> g a -> g a
s1 ->
    (forall a. (a -> a) -> (:.:) f g a -> (:.:) f g a) -> r
r ((forall a. (a -> a) -> (:.:) f g a -> (:.:) f g a) -> r)
-> (forall a. (a -> a) -> (:.:) f g a -> (:.:) f g a) -> r
forall a b. (a -> b) -> a -> b
$ \a -> a
f -> (f (g a) -> f (g a)) -> (:.:) f g a -> (:.:) f g a
coerce ((g a -> g a) -> f (g a) -> f (g a)
forall a. (a -> a) -> f a -> f a
s0 ((a -> a) -> g a -> g a
forall a. (a -> a) -> g a -> g a
s1 a -> a
f))
  {-# INLINE gsetters #-}

newtype FieldSetter f = FS { FieldSetter f -> forall a. (a -> a) -> f a -> f a
runFS :: forall a. (a -> a) -> f a -> f a }

setters :: (Generic1 f, GUpdate (Rep1 f)) => f (FieldSetter f)
setters :: f (FieldSetter f)
setters = Rep1 f (FieldSetter f) -> f (FieldSetter f)
forall k (f :: k -> *) (a :: k). Generic1 f => Rep1 f a -> f a
to1 (Rep1 f (FieldSetter f) -> f (FieldSetter f))
-> Rep1 f (FieldSetter f) -> f (FieldSetter f)
forall a b. (a -> b) -> a -> b
$ ((forall a. (a -> a) -> Rep1 f a -> Rep1 f a) -> FieldSetter f)
-> Rep1 f (FieldSetter f)
forall (rec :: * -> *) r.
GUpdate rec =>
((forall a. (a -> a) -> rec a -> rec a) -> r) -> rec r
gsetters (\forall a. (a -> a) -> Rep1 f a -> Rep1 f a
overI -> (forall a. (a -> a) -> f a -> f a) -> FieldSetter f
forall (f :: * -> *).
(forall a. (a -> a) -> f a -> f a) -> FieldSetter f
FS ((forall a. (a -> a) -> f a -> f a) -> FieldSetter f)
-> (forall a. (a -> a) -> f a -> f a) -> FieldSetter f
forall a b. (a -> b) -> a -> b
$ \a -> a
f -> Rep1 f a -> f a
forall k (f :: k -> *) (a :: k). Generic1 f => Rep1 f a -> f a
to1 (Rep1 f a -> f a) -> (f a -> Rep1 f a) -> f a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> Rep1 f a -> Rep1 f a
forall a. (a -> a) -> Rep1 f a -> Rep1 f a
overI a -> a
f (Rep1 f a -> Rep1 f a) -> (f a -> Rep1 f a) -> f a -> Rep1 f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> Rep1 f a
forall k (f :: k -> *) (a :: k). Generic1 f => f a -> Rep1 f a
from1)