ProjectionExt.hs

{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}
{-# LANGUAGE TypeApplications       #-}

-- ---------------------------------------------------------------------------
-- Module      :  Data.Comp.ProjectionExt
--
-- Extends Data.Comp.Projection with:
--
--   1. Support for (:*:) products alongside tuples (,).
--      Because Elem is a closed type family in the original module,
--      we re-implement it here with the extra cases, and re-derive
--      everything on top of it.  The original (:<) / pr are NOT
--      re-exported; use (:<~) and pr~ from this module instead.
--
--   2. uncons — projects a component AND returns the remainder
--      of the product with that component removed, as an abstract type.
--
--   3. recons — rebuilds the original product from a remainder
--      and a (possibly modified) projected component.
--
-- Supported product constructors:
--   (,)   — Haskell built-in pair
--   (:*:) — GHC.Generics / Data.Functor.Product style functor product
--           (used at kind Type by applying both sides to a common index)
--
-- Adding further product constructors (e.g. Data.Strict.Pair) requires
-- one new equation in Elem, one in Remainder2, and two new instances
-- each in Proj, ProjR, and Reconstruct.
-- ---------------------------------------------------------------------------

module Data.Comp.ProjectionExt
  ( -- * Constraint and plain projection (replaces :< / pr)
    (:<~), (:<|), (:*:)(..)
  , pr
    -- * Projection with remainder
  , RemainderOf
  , uncons
  , recons
  , modify
    -- * Re-exported position machinery (for advanced use)
  , Elem
  ) where

import Data.Comp.SubsumeCommon
import Data.Kind    (Type)

data a :*: b = a :*: b
infixr 6 :*:

-- ---------------------------------------------------------------------------
-- 1.  Elem — closed type family recognising BOTH (,) and (:*:)
--
-- The structure exactly mirrors the original Elem, with every
-- tuple case duplicated for (:*:).
-- ---------------------------------------------------------------------------

type family Elem (f :: Type) (g :: Type) :: Emb where

  -- Base: exact match
  Elem f f = 'Found 'Here

  -- LHS is a tuple-pair: look for each component independently in g
  Elem (f1, f2) g = Sum' (Elem f1 g) (Elem f2 g)

  -- LHS is a (:*:)-pair: same idea
  Elem (f1 :*: f2) g = Sum' (Elem f1 g) (Elem f2 g)

  -- RHS is a tuple-pair: choose which side f lives in
  Elem f (g1, g2) = Choose (Elem f g1) (Elem f g2)

  -- RHS is a (:*:)-pair: same idea
  Elem f (g1 :*: g2) = Choose (Elem f g1) (Elem f g2)

  -- No match
  Elem f g = 'NotFound

-- ---------------------------------------------------------------------------
-- 2.  Proj — plain projection, mirrors the original Proj class
--     but dispatches on both (,) and (:*:) at the value level.
-- ---------------------------------------------------------------------------

class Proj (e :: Emb) (p :: Type) (q :: Type) where
  proj' :: Proxy e -> q -> p

-- Base: p is q
instance Proj ('Found 'Here) f f where
  proj' _ x = x

-- p is in the left branch of a tuple
instance Proj ('Found pos) f g
    => Proj ('Found ('Le pos)) f (g, g') where
  proj' _ (l, _) = proj' (P :: Proxy ('Found pos)) l

-- p is in the right branch of a tuple
instance Proj ('Found pos) f g
    => Proj ('Found ('Ri pos)) f (g', g) where
  proj' _ (_, r) = proj' (P :: Proxy ('Found pos)) r

-- p is in the left branch of a (:*:)
instance Proj ('Found pos) f g
    => Proj ('Found ('Le pos)) f (g :*: g') where
  proj' _ (l :*: _) = proj' (P :: Proxy ('Found pos)) l

-- p is in the right branch of a (:*:)
instance Proj ('Found pos) f g
    => Proj ('Found ('Ri pos)) f (g' :*: g) where
  proj' _ (_ :*: r) = proj' (P :: Proxy ('Found pos)) r

-- p is a pair (f1, f2) scattered across q — tuple variant
instance ( Proj ('Found p1) f1 g
         , Proj ('Found p2) f2 g
         )
    => Proj ('Found ('Sum p1 p2)) (f1, f2) g where
  proj' _ x =
    ( proj' (P :: Proxy ('Found p1)) x
    , proj' (P :: Proxy ('Found p2)) x
    )

-- p is a pair (f1 :*: f2) scattered across q — (:*:) variant
instance ( Proj ('Found p1) f1 g
         , Proj ('Found p2) f2 g
         )
    => Proj ('Found ('Sum p1 p2)) (f1 :*: f2) g where
  proj' _ x =
    proj' (P :: Proxy ('Found p1)) x
    :*:
    proj' (P :: Proxy ('Found p2)) x

-- ---------------------------------------------------------------------------
-- 3.  Public constraint and projection function (replaces :< / pr)
-- ---------------------------------------------------------------------------

-- | @p :\<~ q@ means @p@ is a uniquely-occurring component of @q@,
-- where @q@ may be built from either @(,)@ or @(':*:')@ products.
type p :<~ q = Proj (ComprEmb (Elem p q)) p q

-- | Project component @p@ out of product @q@.
--   Generalises 'Data.Comp.Projection.pr' to also handle @(':*:')@.
pr :: forall p q. (p :<~ q) => q -> p
pr = proj' (P :: Proxy (ComprEmb (Elem p q)))

-- ---------------------------------------------------------------------------
-- 4.  Remainder type family
--
-- Remainder e p q  is the type of q with the component at position e
-- removed.  The result type depends on which product constructor was
-- used at each node, so we need cases for both (,) and (:*:).
-- ---------------------------------------------------------------------------

type family Remainder (e :: Emb) (p :: Type) (q :: Type) :: Type where

  -- Base: p was the whole product
  Remainder ('Found 'Here) p p = ()

  -- p was in the left branch of a tuple
  Remainder ('Found ('Le pos)) p (g, g') =
      (Remainder ('Found pos) p g, g')

  -- p was in the right branch of a tuple
  Remainder ('Found ('Ri pos)) p (g', g) =
      (g', Remainder ('Found pos) p g)

  -- p was in the left branch of a (:*:)
  Remainder ('Found ('Le pos)) p (g :*: g') =
      Remainder ('Found pos) p g :*: g'

  -- p was in the right branch of a (:*:)
  Remainder ('Found ('Ri pos)) p (g' :*: g) =
      g' :*: Remainder ('Found pos) p g

  -- p is a tuple-pair (f1, f2) scattered through q
  Remainder ('Found ('Sum pos1 pos2)) (f1, f2) q =
      Remainder ('Found pos2) f2 (Remainder ('Found pos1) f1 q)

  -- p is a (:*:)-pair (f1 :*: f2) scattered through q
  Remainder ('Found ('Sum pos1 pos2)) (f1 :*: f2) q =
      Remainder ('Found pos2) f2 (Remainder ('Found pos1) f1 q)

-- Convenience alias hiding the position evidence.
type RemainderOf p q = Remainder (ComprEmb (Elem p q)) p q

-- ---------------------------------------------------------------------------
-- 5.  ProjR — projection that also returns the remainder
-- ---------------------------------------------------------------------------

class ProjR (e :: Emb) (p :: Type) (q :: Type) where
  prR' :: Proxy e -> q -> (p, Remainder e p q)

-- Base
instance ProjR ('Found 'Here) f f where
  prR' _ x = (x, ())

-- Left branch of a tuple
instance ProjR ('Found pos) f g
    => ProjR ('Found ('Le pos)) f (g, g') where
  prR' _ (l, r) =
    let (p, remL) = prR' (P :: Proxy ('Found pos)) l
    in  (p, (remL, r))

-- Right branch of a tuple
instance ProjR ('Found pos) f g
    => ProjR ('Found ('Ri pos)) f (g', g) where
  prR' _ (l, r) =
    let (p, remR) = prR' (P :: Proxy ('Found pos)) r
    in  (p, (l, remR))

-- Left branch of a (:*:)
instance ProjR ('Found pos) f g
    => ProjR ('Found ('Le pos)) f (g :*: g') where
  prR' _ (l :*: r) =
    let (p, remL) = prR' (P :: Proxy ('Found pos)) l
    in  (p, remL :*: r)

-- Right branch of a (:*:)
instance ProjR ('Found pos) f g
    => ProjR ('Found ('Ri pos)) f (g' :*: g) where
  prR' _ (l :*: r) =
    let (p, remR) = prR' (P :: Proxy ('Found pos)) r
    in  (p, l :*: remR)

-- Scattered tuple-pair (f1, f2)
instance ( ProjR ('Found pos1) f1 q
         , ProjR ('Found pos2) f2 (Remainder ('Found pos1) f1 q)
         )
    => ProjR ('Found ('Sum pos1 pos2)) (f1, f2) q where
  prR' _ x =
    let (f1, rem1) = prR' (P :: Proxy ('Found pos1)) x
        (f2, rem2) = prR' (P :: Proxy ('Found pos2)) rem1
    in  ((f1, f2), rem2)

-- Scattered (:*:)-pair (f1 :*: f2)
instance ( ProjR ('Found pos1) f1 q
         , ProjR ('Found pos2) f2 (Remainder ('Found pos1) f1 q)
         )
    => ProjR ('Found ('Sum pos1 pos2)) (f1 :*: f2) q where
  prR' _ x =
    let (f1, rem1) = prR' (P :: Proxy ('Found pos1)) x
        (f2, rem2) = prR' (P :: Proxy ('Found pos2)) rem1
    in  (f1 :*: f2, rem2)

-- ---------------------------------------------------------------------------
-- 6.  Reconstruct — inverse of ProjR
-- ---------------------------------------------------------------------------

class Reconstruct (e :: Emb) (p :: Type) (q :: Type) where
  recons' :: Proxy e -> Remainder e p q -> p -> q

-- Base
instance Reconstruct ('Found 'Here) f f where
  recons' _ () x = x

-- Left branch of a tuple
instance Reconstruct ('Found pos) f g
    => Reconstruct ('Found ('Le pos)) f (g, g') where
  recons' _ (remL, r) p =
    (recons' (P :: Proxy ('Found pos)) remL p, r)

-- Right branch of a tuple
instance Reconstruct ('Found pos) f g
    => Reconstruct ('Found ('Ri pos)) f (g', g) where
  recons' _ (l, remR) p =
    (l, recons' (P :: Proxy ('Found pos)) remR p)

-- Left branch of a (:*:)
instance Reconstruct ('Found pos) f g
    => Reconstruct ('Found ('Le pos)) f (g :*: g') where
  recons' _ (remL :*: r) p =
    recons' (P :: Proxy ('Found pos)) remL p :*: r

-- Right branch of a (:*:)
instance Reconstruct ('Found pos) f g
    => Reconstruct ('Found ('Ri pos)) f (g' :*: g) where
  recons' _ (l :*: remR) p =
    l :*: recons' (P :: Proxy ('Found pos)) remR p

-- Scattered tuple-pair (f1, f2)
instance ( Reconstruct ('Found pos2) f2 (Remainder ('Found pos1) f1 q)
         , Reconstruct ('Found pos1) f1 q
         )
    => Reconstruct ('Found ('Sum pos1 pos2)) (f1, f2) q where
  recons' _ rem2 (f1, f2) =
    let rem1 = recons' (P :: Proxy ('Found pos2)) rem2 f2
    in  recons' (P :: Proxy ('Found pos1)) rem1 f1

-- Scattered (:*:)-pair (f1 :*: f2)
instance ( Reconstruct ('Found pos2) f2 (Remainder ('Found pos1) f1 q)
         , Reconstruct ('Found pos1) f1 q
         )
    => Reconstruct ('Found ('Sum pos1 pos2)) (f1 :*: f2) q where
  recons' _ rem2 (f1 :*: f2) =
    let rem1 = recons' (P :: Proxy ('Found pos2)) rem2 f2
    in  recons' (P :: Proxy ('Found pos1)) rem1 f1

-- ---------------------------------------------------------------------------
-- 7.  Public API
-- ---------------------------------------------------------------------------

-- | Bundles all constraints needed for 'uncons' and 'recons'.
type p :<| q =
  ( ProjR       (ComprEmb (Elem p q)) p q
  , Reconstruct (ComprEmb (Elem p q)) p q
  , p :<~ q
  )

-- | Project component @p@ from product @q@, also returning the remainder.
--
-- The remainder type is abstract — it should only be passed to 'recons'.
--
-- Law:  @uncurry (flip recons) (uncons x) == x@
--
-- Examples:
--   uncons @Int  (True, (42 :: Int, 'x'))  == (42,  (True, 'x'))
--   uncons @Bool (True :*: (42 :: Int))     == (True, () :*: 42)
uncons
  :: forall p q. (p :<| q)
  => q
  -> (p, RemainderOf p q)
uncons = prR' (P :: Proxy (ComprEmb (Elem p q)))

-- | Reconstruct the original product from a remainder and a
-- (possibly modified) projected component.
--
-- Examples:
--   let (n, rem) = uncons @Int (True, (42 :: Int, 'x'))
--   recons rem (n + 1)                    == (True, (43, 'x'))
--
--   let (b, rem) = uncons @Bool (True :*: (42 :: Int))
--   recons rem (not b)                    == (False :*: 42)
recons
  :: forall p q. (p :<| q)
  => RemainderOf p q
  -> p
  -> q
recons = recons' (P :: Proxy (ComprEmb (Elem p q)))

-- | Deconstruct and reconstruct a product, modifying a field
-- determined by the type of the given transformation function.
modify :: forall p q. (p :<| q)
  => (p -> p) -> q -> q
modify f q = recons rem (f p)
  where (p, rem) = uncons @p q

-- ---------------------------------------------------------------------------
-- 8.  Sanity checks (unexported; remove before packaging)
-- ---------------------------------------------------------------------------

-- Tuple-only product
_t1 :: (Bool, (Int, Char))
_t1 = let (n, rem) = uncons @Int (True, (42 :: Int, 'x'))
      in  recons rem (n + 1)           -- (True, (43, 'x'))

-- (:*:)-only product
_t2 :: Bool :*: Int
_t2 = let (b, rem) = uncons @Bool (True :*: (42 :: Int))
      in  recons rem (not b)           -- (False :*: 42)

-- Mixed: tuple on the outside, (:*:) on the inside
_t3 :: (Bool :*: Int, Char)
_t3 = let (n, rem) = uncons @Int ((True :*: (42 :: Int)), 'x')
      in  recons rem (n + 1)           -- ((True :*: 43), 'x')

-- Mixed: (:*:) on the outside, tuple on the inside
_t4 :: Bool :*: Int :*: Char
_t4 = let (b, rem) = uncons @Bool (True :*: (42 :: Int) :*: 'x')
      in  recons rem (not b)           -- ((False, 42) :*: 'x')