{-# LANGUAGE CPP #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DerivingVia #-}
{-# OPTIONS_GHC -Wno-x-partial #-}              -- 'prodCon' head: guarded by the product contract
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} -- 'classDict' always returns @EvExpr@

-- | The Stock extension SDK.
--
-- A /synthesizer/ for a class @Cls@ is a function @Class -> Datatype
-- -> Synth EvTerm@: given a structural view of the wrapped type,
-- build the class dictionary as Core — the same static, zero-cost
-- evidence the built-in synthesizers produce (no @Generic@, no
-- runtime @Rep@).
--
-- The only non-trivial primitive a synthesizer needs beyond the
-- structure is 'field': request the dictionary for a field's type and
-- continue with its evidence.  This is the \"continuation\" that lets
-- a synthesizer pause, have GHC solve a sub-constraint, and resume —
-- so @Eq@ (consumer), @Functor@ (transformer) and @Arbitrary@
-- (producer) are all just monadic programs over the same structure.
--
-- Companion packages register support for a new class by writing an
-- instance of 'DeriveStock'; the plugin discovers and runs it.  This
-- module deliberately depends only on @ghc@ + @base@, so companions
-- stay light.
module Stock.Derive
  ( -- * Structural view
    Datatype(..)
  , Constructor(..)
    -- * The synthesis monad
  , Synth
  , runSynth
  , synthTc
  , liftTc
  , field
  , fresh
  , castInto
  , classDict
  , classDictWith
  , classMethod
    -- * SOP-style sum-of-products combinators (generics-sop flavour)
  , productCon
  , matchSOP
  , injectSOP
  , fromProduct
  , toProduct
  , pureFields
  , cpureFields
  , mapFields
  , cmapFields
  , zipFields
  , czipFields
  , foldlFields
  , cfoldlFields
  , traverseFields
  , ctraverseFields
    -- * The witness interface
  , Deriver(..)
  , DeriveStock(..)
  , Deriver1(..)
  , DeriveStock1(..)
  , Deriver2(..)
  , DeriveStock2(..)
  ) where

import GHC.Plugins
import GHC.Tc.Plugin (TcPluginM, newWanted, unsafeTcPluginTcM, tcLookupId)
import GHC.Tc.Types.Constraint (Ct, mkNonCanonical, ctEvExpr)
#if MIN_VERSION_ghc(9,12,0)
import GHC.Tc.Types.CtLoc (CtLoc)
#else
import GHC.Tc.Types.Constraint (CtLoc)
#endif
import GHC.Tc.Types.Evidence (EvTerm(EvExpr))
import GHC.Core.Predicate (mkClassPred)
import GHC.Core.Class (Class, classMethods, classOpItems)
import GHC.Types.Fixity (Fixity)
import Control.Monad (forM, foldM)
import Stock.Trans (ReaderT(..), WriterT(..))
-- The kinds for the witness-class indices.  (GHC.Plugins' unqualified @Type@ is
-- the compiler's AST type; here we need the /kind/ @Data.Kind.Type@.)
import qualified Data.Kind as K

-- ---------------------------------------------------------------------------
-- Structural view of the wrapped type
-- ---------------------------------------------------------------------------

-- | What a synthesizer sees when solving @C (Stock T)@: the via-target it is
-- building the instance for (@Stock T@), the underlying analysed type (@T@),
-- the newtype-unwrap coercion between them, and @T@'s constructors.  Field
-- types are already instantiated at @T@'s actual arguments.
data Datatype = Datatype
  { Datatype -> Type
dtVia    :: Type            -- ^ the via-target, e.g. @Stock T@ — what the instance is /for/
  , Datatype -> Coercion
dtUnwrap :: Coercion        -- ^ @dtVia ~R dtType@ (newtype unwrap)
  , Datatype -> Type
dtType   :: Type            -- ^ the underlying type, e.g. @T@ (or @T a@)
  , Datatype -> [Constructor]
dtCons   :: [Constructor]
  }

-- | One constructor of the analysed type: its 'DataCon', field types
-- (instantiated at @T@'s arguments), fixity, and record labels if any.
data Constructor = Constructor
  { Constructor -> DataCon
conDataCon :: DataCon
  , Constructor -> [Type]
conFields  :: [Type]        -- ^ field types the synthesizer sees, instantiated at
                                --   @T@'s arguments — the /modifier/ types where a field
                                --   is overridden (see "Stock.Override"), else the real ones
  , Constructor -> Fixity
conFixity  :: Fixity        -- ^ for infix constructors (default @defaultFixity@)
  , Constructor -> Maybe [FieldLabel]
conLabels  :: Maybe [FieldLabel]  -- ^ record selectors, if a record
  , Constructor -> [Coercion]
conFieldCos :: [Coercion]   -- ^ per field, @realFieldType ~R conFields!!i@; 'Refl' when
                                --   the field is not overridden.  'matchSOP'\/'injectSOP'
                                --   apply these so synthesizers see only 'conFields'.
  }

-- ---------------------------------------------------------------------------
-- The synthesis monad: a writer of emitted wanteds over a reader of the CtLoc
-- ---------------------------------------------------------------------------

-- | A Core-building computation that may request sub-instances (emitting
-- wanted constraints) and allocate fresh binders.  Structurally this is a
-- reader of the 'CtLoc' over a writer of emitted 'Ct's over 'TcPluginM', so we
-- derive the monad instances straight from that stack (the representations are
-- coercible).
newtype Synth a = Synth (CtLoc -> TcPluginM (a, [Ct]))
  deriving ((forall a b. (a -> b) -> Synth a -> Synth b)
-> (forall a b. a -> Synth b -> Synth a) -> Functor Synth
forall a b. a -> Synth b -> Synth a
forall a b. (a -> b) -> Synth a -> Synth 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) -> Synth a -> Synth b
fmap :: forall a b. (a -> b) -> Synth a -> Synth b
$c<$ :: forall a b. a -> Synth b -> Synth a
<$ :: forall a b. a -> Synth b -> Synth a
Functor, Functor Synth
Functor Synth =>
(forall a. a -> Synth a)
-> (forall a b. Synth (a -> b) -> Synth a -> Synth b)
-> (forall a b c. (a -> b -> c) -> Synth a -> Synth b -> Synth c)
-> (forall a b. Synth a -> Synth b -> Synth b)
-> (forall a b. Synth a -> Synth b -> Synth a)
-> Applicative Synth
forall a. a -> Synth a
forall a b. Synth a -> Synth b -> Synth a
forall a b. Synth a -> Synth b -> Synth b
forall a b. Synth (a -> b) -> Synth a -> Synth b
forall a b c. (a -> b -> c) -> Synth a -> Synth b -> Synth c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> Synth a
pure :: forall a. a -> Synth a
$c<*> :: forall a b. Synth (a -> b) -> Synth a -> Synth b
<*> :: forall a b. Synth (a -> b) -> Synth a -> Synth b
$cliftA2 :: forall a b c. (a -> b -> c) -> Synth a -> Synth b -> Synth c
liftA2 :: forall a b c. (a -> b -> c) -> Synth a -> Synth b -> Synth c
$c*> :: forall a b. Synth a -> Synth b -> Synth b
*> :: forall a b. Synth a -> Synth b -> Synth b
$c<* :: forall a b. Synth a -> Synth b -> Synth a
<* :: forall a b. Synth a -> Synth b -> Synth a
Applicative, Applicative Synth
Applicative Synth =>
(forall a b. Synth a -> (a -> Synth b) -> Synth b)
-> (forall a b. Synth a -> Synth b -> Synth b)
-> (forall a. a -> Synth a)
-> Monad Synth
forall a. a -> Synth a
forall a b. Synth a -> Synth b -> Synth b
forall a b. Synth a -> (a -> Synth b) -> Synth b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. Synth a -> (a -> Synth b) -> Synth b
>>= :: forall a b. Synth a -> (a -> Synth b) -> Synth b
$c>> :: forall a b. Synth a -> Synth b -> Synth b
>> :: forall a b. Synth a -> Synth b -> Synth b
$creturn :: forall a. a -> Synth a
return :: forall a. a -> Synth a
Monad)
    via ReaderT CtLoc (WriterT [Ct] TcPluginM)

-- | Run a synthesizer at a constraint location, collecting the wanteds it
-- emitted (to be returned to GHC alongside the solution).
runSynth :: CtLoc -> Synth a -> TcPluginM (a, [Ct])
runSynth :: forall a. CtLoc -> Synth a -> TcPluginM (a, [Ct])
runSynth CtLoc
loc (Synth CtLoc -> TcPluginM (a, [Ct])
g) = CtLoc -> TcPluginM (a, [Ct])
g CtLoc
loc

-- | Build a @Synth@ from a location-dependent, wanted-emitting action — the
-- inverse of 'runSynth'.  Lets a raw @CtLoc -> TcPluginM (EvTerm, [Ct])@
-- synthesizer be presented as a @Deriver@ (see @viaSynth@ in the plugin).
synthTc :: (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
synthTc :: forall a. (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
synthTc = (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
forall a. (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
Synth

-- | Lift a plugin action.
liftTc :: TcPluginM a -> Synth a
liftTc :: forall a. TcPluginM a -> Synth a
liftTc TcPluginM a
m = (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
forall a. (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
Synth \CtLoc
_ -> do a
a <- TcPluginM a
m; (a, [Ct]) -> TcPluginM (a, [Ct])
forall a. a -> TcPluginM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, [])

-- | The continuation: request the dictionary for @C ty@ and resume with its
-- evidence.  Emits the wanted so GHC solves it (possibly via this very plugin,
-- enabling recursion into the field's own instance).
field :: Class -> Type -> Synth CoreExpr
field :: Class -> Type -> Synth CoreExpr
field Class
cls Type
ty = (CtLoc -> TcPluginM (CoreExpr, [Ct])) -> Synth CoreExpr
forall a. (CtLoc -> TcPluginM (a, [Ct])) -> Synth a
Synth \CtLoc
loc -> do
  CtEvidence
ev <- CtLoc -> Type -> TcPluginM CtEvidence
newWanted CtLoc
loc (Class -> [Type] -> Type
mkClassPred Class
cls [Type
ty])
  (CoreExpr, [Ct]) -> TcPluginM (CoreExpr, [Ct])
forall a. a -> TcPluginM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HasDebugCallStack => CtEvidence -> CoreExpr
CtEvidence -> CoreExpr
ctEvExpr CtEvidence
ev, [CtEvidence -> Ct
mkNonCanonical CtEvidence
ev])

-- | A fresh local binder of the given type.
fresh :: Type -> String -> Synth Id
fresh :: Type -> String -> Synth Id
fresh Type
ty String
s = TcPluginM Id -> Synth Id
forall a. TcPluginM a -> Synth a
liftTc (TcPluginM Id -> Synth Id) -> TcPluginM Id -> Synth Id
forall a b. (a -> b) -> a -> b
$ do
  Unique
u <- TcM Unique -> TcPluginM Unique
forall a. TcM a -> TcPluginM a
unsafeTcPluginTcM TcM Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
  Id -> TcPluginM Id
forall a. a -> TcPluginM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HasDebugCallStack => Name -> Type -> Type -> Id
Name -> Type -> Type -> Id
mkLocalId (Unique -> OccName -> Name
mkSystemName Unique
u (String -> OccName
mkVarOcc String
s)) Type
manyDataConTy Type
ty)

-- | A class method selected by its source name — order-independent, unlike
-- indexing @classMethods@ positionally (whose order can differ across GHC
-- versions).  Panics if the class has no such method (a plugin bug).
classMethod :: String -> Class -> Id
classMethod :: String -> Class -> Id
classMethod String
name Class
cls =
  case (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter ((String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
name) (String -> Bool) -> (Id -> String) -> Id -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> String
occNameString (OccName -> String) -> (Id -> OccName) -> Id -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> OccName
forall name. HasOccName name => name -> OccName
occName) (Class -> [Id]
classMethods Class
cls) of
    (Id
m : [Id]
_) -> Id
m
    []      -> String -> SDoc -> Id
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"stock: classMethod: no method" (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
name SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> Class -> SDoc
forall a. Outputable a => a -> SDoc
ppr Class
cls)

-- | Apply a class's dictionary constructor: @C:Cls \@ty m1 .. mn@.
classDict :: Class -> Type -> [CoreExpr] -> EvTerm
classDict :: Class -> Type -> [CoreExpr] -> EvTerm
classDict Class
cls Type
ty [CoreExpr]
methods =
  CoreExpr -> EvTerm
EvExpr (CoreExpr -> [CoreExpr] -> CoreExpr
forall b. Expr b -> [Expr b] -> Expr b
mkApps (Id -> CoreExpr
forall b. Id -> Expr b
Var (DataCon -> Id
dataConWorkId (Class -> DataCon
classDataCon Class
cls))) (Type -> CoreExpr
forall b. Type -> Expr b
Type Type
ty CoreExpr -> [CoreExpr] -> [CoreExpr]
forall a. a -> [a] -> [a]
: [CoreExpr]
methods))

-- | Build a (recursive) dictionary giving explicit superclass dictionaries and
-- implementations for the listed method indices; every other method is taken
-- from the class's own default (applied to the recursive dictionary).  Lets a
-- synthesizer fill a many-method class from a few key methods — e.g.
-- @Hashable@ from just @hashWithSalt@ (its @hash@ has a default), with its
-- @Eq@ superclass supplied as @[field eqCls ty]@.
classDictWith :: Class -> Type -> [CoreExpr] -> [(Int, CoreExpr)] -> Synth EvTerm
classDictWith :: Class -> Type -> [CoreExpr] -> [(Int, CoreExpr)] -> Synth EvTerm
classDictWith Class
cls Type
ty [CoreExpr]
supers [(Int, CoreExpr)]
overrides = do
  Id
dvar    <- Type -> String -> Synth Id
fresh (Class -> [Type] -> Type
mkClassPred Class
cls [Type
ty]) String
"dict"
  [CoreExpr]
methods <- [(Int, Id)] -> ((Int, Id) -> Synth CoreExpr) -> Synth [CoreExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Int] -> [Id] -> [(Int, Id)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Class -> [Id]
classMethods Class
cls)) \(Int
i, Id
_) ->
    case Int -> [(Int, CoreExpr)] -> Maybe CoreExpr
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
i [(Int, CoreExpr)]
overrides of
      Just CoreExpr
e  -> CoreExpr -> Synth CoreExpr
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CoreExpr
e
      Maybe CoreExpr
Nothing -> case (Id, DefMethInfo) -> DefMethInfo
forall a b. (a, b) -> b
snd (Class -> [(Id, DefMethInfo)]
classOpItems Class
cls [(Id, DefMethInfo)] -> Int -> (Id, DefMethInfo)
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) of
        Just (Name
nm, DefMethSpec Type
_) -> do Id
dm <- TcPluginM Id -> Synth Id
forall a. TcPluginM a -> Synth a
liftTc (Name -> TcPluginM Id
tcLookupId Name
nm)
                           CoreExpr -> Synth CoreExpr
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CoreExpr -> [CoreExpr] -> CoreExpr
forall b. Expr b -> [Expr b] -> Expr b
mkApps (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
dm) [Type -> CoreExpr
forall b. Type -> Expr b
Type Type
ty, Id -> CoreExpr
forall b. Id -> Expr b
Var Id
dvar])
        DefMethInfo
Nothing      -> String -> Synth CoreExpr
forall a. HasCallStack => String -> a
error String
"stock: classDictWith: method has no default and no override"
  let EvExpr CoreExpr
con = Class -> Type -> [CoreExpr] -> EvTerm
classDict Class
cls Type
ty ([CoreExpr]
supers [CoreExpr] -> [CoreExpr] -> [CoreExpr]
forall a. [a] -> [a] -> [a]
++ [CoreExpr]
methods)
  EvTerm -> Synth EvTerm
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CoreExpr -> EvTerm
EvExpr (Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let ([(Id, CoreExpr)] -> Bind Id
forall b. [(b, Expr b)] -> Bind b
Rec [(Id
dvar, CoreExpr
con)]) (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
dvar)))

-- ---------------------------------------------------------------------------
-- SOP-style sum-of-products combinators
-- ---------------------------------------------------------------------------
--
-- A thin @generics-sop@ flavour over the structural view.  The correspondence:
--
-- @
--   generics-sop          role             Stock.Derive
--   ────────────          ────             ────────────
--   from x  (NS elim, ∃)  sum   dispatch    matchSOP   dt r x (\\i con fs -> …)
--   to . inj (NS intro)   sum   build       injectSOP  dt con fs
--   cpure_NP  (NP, Π)     field tabulate    cpureFields C k con
--   hcmap     (NP)        field map          cmapFields  C k con xs
--   cliftA2_NP (NP)       field zip          czipFields  C k con xs ys
--   hcfoldl   (NP)        field collapse     cfoldlFields C step z con xs
-- @
--
-- The split mirrors @SOP f = NS (NP f)@: 'matchSOP'\/'injectSOP' handle the
-- /sum/ (the @NS@, an existential over constructors); the @cpure@\/@cmap@\/
-- @czip@\/@cfoldl@ family handle one constructor's /product/ (its @NP@, the
-- representable @Fin n -> f@).  Because the @NP@ combinators take a
-- @Constructor@, they compose inside 'matchSOP' for any constructor of a sum,
-- not just the sole product one.  The @All c xs@ constraint is implicit: the
-- @c@-prefixed combinators call 'field' per field for the @cls@ dictionary.
-- So @eqDeriver@ (sum), @semigroupDeriver@\/@monoidDeriver@ (product) and
-- companions like @stock-hashable@\/@NFData@ read like their generic kin.

-- | The constructor of a product (the sole one).  Exported so product
-- synthesizers can feed it to the per-@Constructor@ NP combinators below.
productCon :: Datatype -> Constructor
productCon :: Datatype -> Constructor
productCon = [Constructor] -> Constructor
forall a. HasCallStack => [a] -> a
head ([Constructor] -> Constructor)
-> (Datatype -> [Constructor]) -> Datatype -> Constructor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Datatype -> [Constructor]
dtCons

-- | The SOP eliminator (@from@ + @case@): scrutinise a value of the via-type
-- and dispatch on its constructor.  @k idx con fields@ builds the @resTy@-typed
-- branch body for constructor @con@ (index @idx@ in 'dtCons') with its bound
-- field expressions.  One alternative per constructor — so this is the
-- sum-of-products generalisation of 'fromProduct'.
matchSOP :: Datatype -> Type -> CoreExpr
         -> (Int -> Constructor -> [CoreExpr] -> Synth CoreExpr) -> Synth CoreExpr
matchSOP :: Datatype
-> Type
-> CoreExpr
-> (Int -> Constructor -> [CoreExpr] -> Synth CoreExpr)
-> Synth CoreExpr
matchSOP Datatype
dt Type
resTy CoreExpr
v Int -> Constructor -> [CoreExpr] -> Synth CoreExpr
k = do
  Id
cb   <- Type -> String -> Synth Id
fresh (Datatype -> Type
dtType Datatype
dt) String
"s"
  [Alt Id]
alts <- [(Int, Constructor)]
-> ((Int, Constructor) -> Synth (Alt Id)) -> Synth [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Int] -> [Constructor] -> [(Int, Constructor)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Datatype -> [Constructor]
dtCons Datatype
dt)) \(Int
i, Constructor
c) -> do
    -- bind each pattern var at the /real/ field type (the coercion's LHS),
    -- then present the continuation the value coerced to the modifier type.
    [Id]
xs   <- (Coercion -> Synth Id) -> [Coercion] -> Synth [Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Coercion
co -> Type -> String -> Synth Id
fresh (Coercion -> Type
coercionLKind Coercion
co) String
"x") (Constructor -> [Coercion]
conFieldCos Constructor
c)
    CoreExpr
body <- Int -> Constructor -> [CoreExpr] -> Synth CoreExpr
k Int
i Constructor
c ((CoreExpr -> Coercion -> CoreExpr)
-> [CoreExpr] -> [Coercion] -> [CoreExpr]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith CoreExpr -> Coercion -> CoreExpr
castInto ((Id -> CoreExpr) -> [Id] -> [CoreExpr]
forall a b. (a -> b) -> [a] -> [b]
map Id -> CoreExpr
forall b. Id -> Expr b
Var [Id]
xs) (Constructor -> [Coercion]
conFieldCos Constructor
c))
    Alt Id -> Synth (Alt Id)
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AltCon -> [Id] -> CoreExpr -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt (DataCon -> AltCon
DataAlt (Constructor -> DataCon
conDataCon Constructor
c)) [Id]
xs CoreExpr
body)
  CoreExpr -> Synth CoreExpr
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast CoreExpr
v (Datatype -> Coercion
dtUnwrap Datatype
dt)) Id
cb Type
resTy [Alt Id]
alts)

-- | @x |> co@, skipping the cast entirely when @co@ is reflexive (the
-- not-overridden case) so the generated Core stays byte-identical.
castInto :: CoreExpr -> Coercion -> CoreExpr
castInto :: CoreExpr -> Coercion -> CoreExpr
castInto CoreExpr
e Coercion
co = if Coercion -> Bool
isReflCo Coercion
co then CoreExpr
e else CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast CoreExpr
e Coercion
co

-- | The SOP introducer (@inj@ + @to@): build a value of the via-type from a
-- chosen constructor and one expression per its fields.
injectSOP :: Datatype -> Constructor -> [CoreExpr] -> CoreExpr
injectSOP :: Datatype -> Constructor -> [CoreExpr] -> CoreExpr
injectSOP Datatype
dt Constructor
c [CoreExpr]
es =
  CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast (DataCon -> [CoreExpr] -> CoreExpr
mkCoreConApps (Constructor -> DataCon
conDataCon Constructor
c) ((Type -> CoreExpr) -> [Type] -> [CoreExpr]
forall a b. (a -> b) -> [a] -> [b]
map Type -> CoreExpr
forall b. Type -> Expr b
Type (HasDebugCallStack => Type -> [Type]
Type -> [Type]
tyConAppArgs (Datatype -> Type
dtType Datatype
dt)) [CoreExpr] -> [CoreExpr] -> [CoreExpr]
forall a. [a] -> [a] -> [a]
++ [CoreExpr]
es'))
       (Coercion -> Coercion
mkSymCo (Datatype -> Coercion
dtUnwrap Datatype
dt))
  where -- each result comes back at the modifier type; coerce it to the real
        -- field type before reapplying the constructor (no-op when reflexive).
        es' :: [CoreExpr]
es' = (CoreExpr -> Coercion -> CoreExpr)
-> [CoreExpr] -> [Coercion] -> [CoreExpr]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\CoreExpr
e Coercion
co -> CoreExpr -> Coercion -> CoreExpr
castInto CoreExpr
e (Coercion -> Coercion
mkSymCo Coercion
co)) [CoreExpr]
es (Constructor -> [Coercion]
conFieldCos Constructor
c)

-- | @productTypeFrom@ + a continuation: the single-constructor case of
-- 'matchSOP'.
fromProduct :: Datatype -> Type -> CoreExpr -> ([CoreExpr] -> Synth CoreExpr)
            -> Synth CoreExpr
fromProduct :: Datatype
-> Type
-> CoreExpr
-> ([CoreExpr] -> Synth CoreExpr)
-> Synth CoreExpr
fromProduct Datatype
dt Type
resTy CoreExpr
v [CoreExpr] -> Synth CoreExpr
k = Datatype
-> Type
-> CoreExpr
-> (Int -> Constructor -> [CoreExpr] -> Synth CoreExpr)
-> Synth CoreExpr
matchSOP Datatype
dt Type
resTy CoreExpr
v \Int
_ Constructor
_ [CoreExpr]
fields -> [CoreExpr] -> Synth CoreExpr
k [CoreExpr]
fields

-- | @productTypeTo@: the single-constructor case of 'injectSOP'.
toProduct :: Datatype -> [CoreExpr] -> CoreExpr
toProduct :: Datatype -> [CoreExpr] -> CoreExpr
toProduct Datatype
dt = Datatype -> Constructor -> [CoreExpr] -> CoreExpr
injectSOP Datatype
dt (Datatype -> Constructor
productCon Datatype
dt)

-- The NP combinators below all operate on ONE @Constructor@ (≅ one @NP@), so
-- they compose directly inside 'matchSOP' for /any/ constructor of a sum — not
-- just the sole product one.  An @NP@ is the representable functor @Fin n ->
-- f@: 'pureFields' tabulates it, 'mapFields'\/'zipFields' act positionwise, and
-- 'foldlFields' collapses it.  Each has a @c@onstrained sibling that requests
-- the field's @cls@ dictionary (the implicit @All c xs@).

-- | @pure_NP@ \/ @cpure_NP@: one result per field, produced from its type
-- (@cpure@ also from its @cls@ dictionary, e.g. @k ft d = mempty \@ft d@).
pureFields :: (Type -> Synth a) -> Constructor -> Synth [a]
pureFields :: forall a. (Type -> Synth a) -> Constructor -> Synth [a]
pureFields Type -> Synth a
k Constructor
con = (Type -> Synth a) -> [Type] -> Synth [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Type -> Synth a
k (Constructor -> [Type]
conFields Constructor
con)

-- | 'pureFields' that also hands each field its own @cls@ dictionary.
cpureFields :: Class -> (Type -> CoreExpr -> CoreExpr) -> Constructor -> Synth [CoreExpr]
cpureFields :: Class
-> (Type -> CoreExpr -> CoreExpr)
-> Constructor
-> Synth [CoreExpr]
cpureFields Class
cls Type -> CoreExpr -> CoreExpr
k = (Type -> Synth CoreExpr) -> Constructor -> Synth [CoreExpr]
forall a. (Type -> Synth a) -> Constructor -> Synth [a]
pureFields \Type
ft -> do CoreExpr
d <- Class -> Type -> Synth CoreExpr
field Class
cls Type
ft; CoreExpr -> Synth CoreExpr
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> CoreExpr -> CoreExpr
k Type
ft CoreExpr
d)

-- | @hmap@ \/ @hcmap@: map over the fields positionwise (the basic @NP@ action;
-- @cmap@ also hands each field's @cls@ dictionary to the step).
mapFields :: (Type -> CoreExpr -> Synth a) -> Constructor -> [CoreExpr] -> Synth [a]
mapFields :: forall a.
(Type -> CoreExpr -> Synth a)
-> Constructor -> [CoreExpr] -> Synth [a]
mapFields Type -> CoreExpr -> Synth a
k Constructor
con [CoreExpr]
xs = [Synth a] -> Synth [a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ((Type -> CoreExpr -> Synth a) -> [Type] -> [CoreExpr] -> [Synth a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> CoreExpr -> Synth a
k (Constructor -> [Type]
conFields Constructor
con) [CoreExpr]
xs)

-- | 'mapFields' that also hands each field its own @cls@ dictionary.
cmapFields :: Class -> (Type -> CoreExpr -> CoreExpr -> CoreExpr) -> Constructor -> [CoreExpr] -> Synth [CoreExpr]
cmapFields :: Class
-> (Type -> CoreExpr -> CoreExpr -> CoreExpr)
-> Constructor
-> [CoreExpr]
-> Synth [CoreExpr]
cmapFields Class
cls Type -> CoreExpr -> CoreExpr -> CoreExpr
k = (Type -> CoreExpr -> Synth CoreExpr)
-> Constructor -> [CoreExpr] -> Synth [CoreExpr]
forall a.
(Type -> CoreExpr -> Synth a)
-> Constructor -> [CoreExpr] -> Synth [a]
mapFields \Type
ft CoreExpr
x -> do CoreExpr
d <- Class -> Type -> Synth CoreExpr
field Class
cls Type
ft; CoreExpr -> Synth CoreExpr
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> CoreExpr -> CoreExpr -> CoreExpr
k Type
ft CoreExpr
d CoreExpr
x)

-- | @liftA2_NP@ \/ @cliftA2_NP@: combine two field-lists positionwise (@czip@
-- via each field's @cls@ dictionary, e.g. @k ft d x y = (\<>) \@ft d x y@).
zipFields :: (Type -> CoreExpr -> CoreExpr -> Synth a)
          -> Constructor -> [CoreExpr] -> [CoreExpr] -> Synth [a]
zipFields :: forall a.
(Type -> CoreExpr -> CoreExpr -> Synth a)
-> Constructor -> [CoreExpr] -> [CoreExpr] -> Synth [a]
zipFields Type -> CoreExpr -> CoreExpr -> Synth a
k Constructor
con [CoreExpr]
xs [CoreExpr]
ys = [Synth a] -> Synth [a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ((Type -> CoreExpr -> CoreExpr -> Synth a)
-> [Type] -> [CoreExpr] -> [CoreExpr] -> [Synth a]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Type -> CoreExpr -> CoreExpr -> Synth a
k (Constructor -> [Type]
conFields Constructor
con) [CoreExpr]
xs [CoreExpr]
ys)

-- | 'zipFields' that also hands each field its own @cls@ dictionary.
czipFields :: Class -> (Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr)
           -> Constructor -> [CoreExpr] -> [CoreExpr] -> Synth [CoreExpr]
czipFields :: Class
-> (Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr)
-> Constructor
-> [CoreExpr]
-> [CoreExpr]
-> Synth [CoreExpr]
czipFields Class
cls Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr
k = (Type -> CoreExpr -> CoreExpr -> Synth CoreExpr)
-> Constructor -> [CoreExpr] -> [CoreExpr] -> Synth [CoreExpr]
forall a.
(Type -> CoreExpr -> CoreExpr -> Synth a)
-> Constructor -> [CoreExpr] -> [CoreExpr] -> Synth [a]
zipFields \Type
ft CoreExpr
x CoreExpr
y -> do CoreExpr
d <- Class -> Type -> Synth CoreExpr
field Class
cls Type
ft; CoreExpr -> Synth CoreExpr
forall a. a -> Synth a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr
k Type
ft CoreExpr
d CoreExpr
x CoreExpr
y)

-- | @hfoldl@ \/ @cfoldl_NP@: collapse the fields left-to-right through an
-- accumulator @a@ (any host value — a 'CoreExpr', a list, …).  The step of the
-- unconstrained 'foldlFields' has full @Synth@ access (request any/no/several
-- dictionaries); 'cfoldlFields' hands it each field's @cls@ dictionary.  Needed
-- by accumulating classes, e.g. @Hashable@'s @hashWithSalt@ threading the salt.
foldlFields :: (a -> Type -> CoreExpr -> Synth a) -> a -> Constructor -> [CoreExpr] -> Synth a
foldlFields :: forall a.
(a -> Type -> CoreExpr -> Synth a)
-> a -> Constructor -> [CoreExpr] -> Synth a
foldlFields a -> Type -> CoreExpr -> Synth a
step a
z Constructor
con [CoreExpr]
fields =
  (a -> (Type, CoreExpr) -> Synth a)
-> a -> [(Type, CoreExpr)] -> Synth a
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\a
acc (Type
ft, CoreExpr
e) -> a -> Type -> CoreExpr -> Synth a
step a
acc Type
ft CoreExpr
e) a
z ([Type] -> [CoreExpr] -> [(Type, CoreExpr)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Constructor -> [Type]
conFields Constructor
con) [CoreExpr]
fields)

-- | 'foldlFields' that also hands each field its own @cls@ dictionary.
cfoldlFields :: Class
             -> (CoreExpr -> Type -> CoreExpr -> CoreExpr -> Synth CoreExpr)  -- ^ @acc ft dict field@
             -> CoreExpr            -- ^ initial accumulator
             -> Constructor -> [CoreExpr] -> Synth CoreExpr
cfoldlFields :: Class
-> (CoreExpr -> Type -> CoreExpr -> CoreExpr -> Synth CoreExpr)
-> CoreExpr
-> Constructor
-> [CoreExpr]
-> Synth CoreExpr
cfoldlFields Class
cls CoreExpr -> Type -> CoreExpr -> CoreExpr -> Synth CoreExpr
step =
  (CoreExpr -> Type -> CoreExpr -> Synth CoreExpr)
-> CoreExpr -> Constructor -> [CoreExpr] -> Synth CoreExpr
forall a.
(a -> Type -> CoreExpr -> Synth a)
-> a -> Constructor -> [CoreExpr] -> Synth a
foldlFields \CoreExpr
acc Type
ft CoreExpr
e -> do CoreExpr
d <- Class -> Type -> Synth CoreExpr
field Class
cls Type
ft; CoreExpr -> Type -> CoreExpr -> CoreExpr -> Synth CoreExpr
step CoreExpr
acc Type
ft CoreExpr
d CoreExpr
e

-- | @htraverse@ over a list of fields: produce one @a@ per field in @Synth@.
-- The most general traversal-shaped combinator — used for applicative-effectful
-- work like generating `Gen a` values (for @Arbitrary@).  'traverseFields' has
-- full @Synth@ access; 'ctraverseFields' hands each field's @cls@ dictionary
-- to the step.
traverseFields :: (Type -> CoreExpr -> Synth a) -> Constructor -> [CoreExpr] -> Synth [a]
traverseFields :: forall a.
(Type -> CoreExpr -> Synth a)
-> Constructor -> [CoreExpr] -> Synth [a]
traverseFields Type -> CoreExpr -> Synth a
k Constructor
con [CoreExpr]
xs = [Synth a] -> Synth [a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ((Type -> CoreExpr -> Synth a) -> [Type] -> [CoreExpr] -> [Synth a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> CoreExpr -> Synth a
k (Constructor -> [Type]
conFields Constructor
con) [CoreExpr]
xs)

-- | @hctraverse@: the @c@onstrained 'traverseFields' — requests each field's
-- @cls@ dictionary and hands it to the step (alongside the field value).
-- Used by @Arbitrary@ (request `Arbitrary ft` per field, emit `Gen ft`),
-- @CoArbitrary@ (request `CoArbitrary ft`, emit function generator),
-- and @Shrink@ (request `Shrink ft`, emit shrink list).
ctraverseFields :: Class
                -> (Type -> CoreExpr -> CoreExpr -> Synth CoreExpr)  -- ^ @ft dict field@
                -> Constructor -> [CoreExpr] -> Synth [CoreExpr]
ctraverseFields :: Class
-> (Type -> CoreExpr -> CoreExpr -> Synth CoreExpr)
-> Constructor
-> [CoreExpr]
-> Synth [CoreExpr]
ctraverseFields Class
cls Type -> CoreExpr -> CoreExpr -> Synth CoreExpr
k = (Type -> CoreExpr -> Synth CoreExpr)
-> Constructor -> [CoreExpr] -> Synth [CoreExpr]
forall a.
(Type -> CoreExpr -> Synth a)
-> Constructor -> [CoreExpr] -> Synth [a]
traverseFields \Type
ft CoreExpr
e -> do CoreExpr
d <- Class -> Type -> Synth CoreExpr
field Class
cls Type
ft; Type -> CoreExpr -> CoreExpr -> Synth CoreExpr
k Type
ft CoreExpr
d CoreExpr
e

-- ---------------------------------------------------------------------------
-- The witness interface
-- ---------------------------------------------------------------------------

-- | A class's synthesizer, keyed by the wrapper arity it works through.
newtype Deriver = Deriver { Deriver -> Class -> Datatype -> Synth EvTerm
runDeriver :: Class -> Datatype -> Synth EvTerm }

-- | Register synthesis of a class @cls@ derived @via Stock@.  The method does
-- not mention @cls@, so the plugin selects the instance by looking it up in the
-- instance environment rather than by ordinary dispatch.
class DeriveStock (cls :: K.Type -> K.Constraint) where
  deriveStock :: Deriver

-- | A @Stock1@ synthesizer for a @(Type -> Type)@ class: given the class, the
-- constraint location, the via-target @Stock1 F@ and the inner @F@, build the
-- dictionary — or 'Nothing' if a field shape is unsupported.  (Lifted classes
-- need the parameter-variance walk, so they get the raw form rather than the
-- 'Datatype'-based 'Deriver'; the @Stock1@ 'TyCon' is recoverable as
-- @tyConAppTyCon@ of the via-target.)
newtype Deriver1 = Deriver1
  { Deriver1
-> Class
-> CtLoc
-> Type
-> Type
-> TcPluginM (Maybe (EvTerm, [Ct]))
runDeriver1 :: Class -> CtLoc -> Type -> Type -> TcPluginM (Maybe (EvTerm, [Ct])) }

-- | Register synthesis of a @(Type -> Type)@ class derived @via Stock1@ (the
-- lifted counterpart of 'DeriveStock' — e.g. @NFData1@, @Hashable1@).
class DeriveStock1 (cls :: (K.Type -> K.Type) -> K.Constraint) where
  deriveStock1 :: Deriver1

-- | The @Stock2@ analogue of 'Deriver1': given the class, the location, the
-- via-target @Stock2 P@ and the inner @P@, build the dictionary.
newtype Deriver2 = Deriver2
  { Deriver2
-> Class
-> CtLoc
-> Type
-> Type
-> TcPluginM (Maybe (EvTerm, [Ct]))
runDeriver2 :: Class -> CtLoc -> Type -> Type -> TcPluginM (Maybe (EvTerm, [Ct])) }

-- | Register synthesis of a @(Type -> Type -> Type)@ class derived @via Stock2@
-- (e.g. @NFData2@).
class DeriveStock2 (cls :: (K.Type -> K.Type -> K.Type) -> K.Constraint) where
  deriveStock2 :: Deriver2