Internal.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-- UndecidableInstances is needed because finding the nth constructor computes
-- with type families in instance contexts: GCount (f :+: g) adds the counts of
-- its branches, and the instances for :+: compare and subtract that count.
-- Resolution terminates: GCount recurses only on the smaller branches f and g
-- of f :+: g, and every GInjection, GInjectionSum and GProduct constraint in an
-- instance context is on a component of the instance head (the inside of an
-- M1, a branch of a :+:, or a factor of a :*:), so each step of resolution is
-- on a smaller representation type. Rep, CmpNat, + and - are evaluated by GHC,
-- and IsLT and AssertZero are not recursive.
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wall #-}

{- | The generic machinery behind 'Data.Lens.Injection.Generic.injection'.

This module is internal: its contents may change in any release, without a
major version bump. Use "Data.Lens.Injection.Generic" instead.
-}
module Data.Lens.Injection.Generic.Internal (
  GInjection (..),
  GInjectionSum (..),
  GProduct (..),
  GCount,
  IsLT,
  AssertZero,
) where

import Control.Lens (Prism, iso, prism, withPrism)
import Data.Bifunctor (first)
import Data.Kind (Constraint, Type)
import Data.Proxy (Proxy (..))
import GHC.Generics (C, C1, D, D1, Generic (..), K1 (..), M1 (..), S, U1 (..), (:*:) (..), (:+:) (..))
import GHC.TypeLits (CmpNat, ErrorMessage (..), Nat, TypeError, type (+), type (-))

-- | The number of constructors in a generic sum.
type family GCount (f :: Type -> Type) :: Nat where
  GCount (f :+: g) = GCount f + GCount g
  GCount _ = 1

-- | Whether an 'Ordering' is 'LT'.
type family IsLT (o :: Ordering) :: Bool where
  IsLT 'LT = 'True
  IsLT _ = 'False

-- | The index has reached a constructor, so it must be 0, otherwise it was out of range.
type family AssertZero (n :: Nat) :: Constraint where
  AssertZero 0 = ()
  AssertZero n = TypeError ('Text "Injection index out of range: it is " ':<>: 'ShowType n ':<>: 'Text " past the last constructor")

-- | A prism to the @n@th constructor of a generic representation.
class GInjection (n :: Nat) (s :: Type -> Type) (t :: Type -> Type) a b where
  gInjection :: proxy n -> Prism (s x) (t x) a b

instance (GInjection n s t a b) => GInjection n (M1 D m s) (M1 D m' t) a b where
  gInjection n =
    iso unM1 M1 . gInjection n
  {-# INLINE gInjection #-}

instance (GInjectionSum (IsLT (CmpNat n (GCount f))) n f g f' g' a b) => GInjection n (f :+: g) (f' :+: g') a b where
  gInjection =
    gInjectionSum (Proxy :: Proxy (IsLT (CmpNat n (GCount f))))
  {-# INLINE gInjection #-}

-- | A constructor with no fields.
instance (AssertZero n, a ~ (), b ~ ()) => GInjection n (M1 C m U1) (M1 C m' U1) a b where
  gInjection _ =
    iso (const ()) (const (M1 U1))
  {-# INLINE gInjection #-}

-- | A constructor with one field.
instance (AssertZero n, x ~ a, y ~ b) => GInjection n (M1 C m (M1 S ms (K1 i x))) (M1 C m' (M1 S ms' (K1 i' y))) a b where
  gInjection _ =
    iso (unK1 . unM1 . unM1) (M1 . M1 . K1)
  {-# INLINE gInjection #-}

-- | A constructor with more than one field.
instance
  ( AssertZero n
  , Generic a
  , Generic b
  , Rep a ~ D1 da (C1 ca p)
  , Rep b ~ D1 db (C1 cb q)
  , GProduct (f :*: g) p
  , GProduct (f' :*: g') q
  ) =>
  GInjection n (M1 C m (f :*: g)) (M1 C m' (f' :*: g')) a b
  where
  gInjection _ =
    iso
      (\(M1 x) -> to (M1 (M1 (gProductTo x))))
      (\b -> case from b of M1 (M1 y) -> M1 (gProductFrom y))
  {-# INLINE gInjection #-}

{- | Descend into the left (@'True@) or right (@'False@) of a generic sum.
The other side of the sum is unchanged.
-}
class GInjectionSum (left :: Bool) (n :: Nat) f g f' g' a b where
  gInjectionSum :: proxy left -> proxy' n -> Prism ((f :+: g) x) ((f' :+: g') x) a b

instance (GInjection n f f' a b, g ~ g') => GInjectionSum 'True n f g f' g' a b where
  gInjectionSum _ n =
    withPrism (gInjection n :: Prism (f x) (f' x) a b) $ \bt sta ->
      prism (L1 . bt) $ \case
        L1 x -> first L1 (sta x)
        R1 y -> Left (R1 y)
  {-# INLINE gInjectionSum #-}

instance (GInjection (n - GCount f) g g' a b, f ~ f') => GInjectionSum 'False n f g f' g' a b where
  gInjectionSum _ _ =
    withPrism (gInjection (Proxy :: Proxy (n - GCount f)) :: Prism (g x) (g' x) a b) $ \bt sta ->
      prism (R1 . bt) $ \case
        L1 x -> Left (L1 x)
        R1 y -> first R1 (sta y)
  {-# INLINE gInjectionSum #-}

-- | Convert between generic products with the same fields, ignoring metadata.
class GProduct p q where
  gProductTo :: p x -> q x
  gProductFrom :: q x -> p x

instance (GProduct f f', GProduct g g') => GProduct (f :*: g) (f' :*: g') where
  gProductTo (f :*: g) =
    gProductTo f :*: gProductTo g
  {-# INLINE gProductTo #-}
  gProductFrom (f :*: g) =
    gProductFrom f :*: gProductFrom g
  {-# INLINE gProductFrom #-}

instance (x ~ y) => GProduct (M1 S m (K1 i x)) (M1 S m' (K1 i' y)) where
  gProductTo (M1 (K1 x)) =
    M1 (K1 x)
  {-# INLINE gProductTo #-}
  gProductFrom (M1 (K1 y)) =
    M1 (K1 y)
  {-# INLINE gProductFrom #-}