{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}

-- | Generic parsing infrastructure
--
-- TODO: Perhaps we could move this (or some form of this) to ghc-tcplugin-api?
-- (The @typelet@ package could then use it, too.)
module Data.Record.Anon.Internal.Plugin.TC.Parsing (
    -- * Basic infrastructure
    ParseResult(..)
  , parseAll
  , parseAll'
  , withOrig
    -- * Parsers for specific (but not @large-anon@ specific) constructs
  , parseConstraint
  , parseConstraint'
  , parseCons
  , parseNil
  , parseInjTyConApp
  ) where

import Data.Bifunctor
import Data.Foldable (toList)
import Data.Void
import GHC.Stack

import Data.Record.Anon.Internal.Plugin.TC.GhcTcPluginAPI
import Data.Record.Anon.Internal.Plugin.TC.TyConSubst

{-------------------------------------------------------------------------------
  Basic infrastructure
-------------------------------------------------------------------------------}

data ParseResult e a =
    -- | Parse successful
    ParseOk a

    -- | Different constraint than we're looking for (does not imply an error)
  | ParseNoMatch

    -- | Constraint of the shape we're looking for, but something is wrong
  | ParseError e
  deriving (forall a b. a -> ParseResult e b -> ParseResult e a
forall a b. (a -> b) -> ParseResult e a -> ParseResult e b
forall e a b. a -> ParseResult e b -> ParseResult e a
forall e a b. (a -> b) -> ParseResult e a -> ParseResult e b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> ParseResult e b -> ParseResult e a
$c<$ :: forall e a b. a -> ParseResult e b -> ParseResult e a
fmap :: forall a b. (a -> b) -> ParseResult e a -> ParseResult e b
$cfmap :: forall e a b. (a -> b) -> ParseResult e a -> ParseResult e b
Functor)

instance Bifunctor ParseResult where
  bimap :: forall a b c d.
(a -> b) -> (c -> d) -> ParseResult a c -> ParseResult b d
bimap a -> b
_ c -> d
g (ParseOk c
a)    = forall e a. a -> ParseResult e a
ParseOk (c -> d
g c
a)
  bimap a -> b
_ c -> d
_ ParseResult a c
ParseNoMatch   = forall e a. ParseResult e a
ParseNoMatch
  bimap a -> b
f c -> d
_ (ParseError a
e) = forall e a. e -> ParseResult e a
ParseError (a -> b
f a
e)

-- | Apply parser to each value in turn, bailing at the first error
parseAll :: forall e a b. (a -> ParseResult e b) -> [a] -> Either e [b]
parseAll :: forall e a b. (a -> ParseResult e b) -> [a] -> Either e [b]
parseAll a -> ParseResult e b
f = [b] -> [a] -> Either e [b]
go []
  where
    go :: [b] -> [a] -> Either e [b]
    go :: [b] -> [a] -> Either e [b]
go [b]
acc []     = forall a b. b -> Either a b
Right (forall a. [a] -> [a]
reverse [b]
acc)
    go [b]
acc (a
a:[a]
as) = case a -> ParseResult e b
f a
a of
                      ParseOk b
b    -> [b] -> [a] -> Either e [b]
go (b
bforall a. a -> [a] -> [a]
:[b]
acc) [a]
as
                      ParseResult e b
ParseNoMatch -> [b] -> [a] -> Either e [b]
go    [b]
acc  [a]
as
                      ParseError e
e -> forall a b. a -> Either a b
Left e
e

-- | Variation on 'parseAll' which rules out the error case
parseAll' :: (a -> ParseResult Void b) -> [a] -> [b]
parseAll' :: forall a b. (a -> ParseResult Void b) -> [a] -> [b]
parseAll' a -> ParseResult Void b
f = forall b. Either Void [b] -> [b]
aux forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a b. (a -> ParseResult e b) -> [a] -> Either e [b]
parseAll a -> ParseResult Void b
f
  where
    aux :: Either Void [b] -> [b]
    aux :: forall b. Either Void [b] -> [b]
aux (Left  Void
v)  = forall a. Void -> a
absurd Void
v
    aux (Right [b]
bs) = [b]
bs

-- | Bundle the parse result with the original value
withOrig :: (a -> ParseResult e b) -> (a -> ParseResult e (a, b))
withOrig :: forall a e b. (a -> ParseResult e b) -> a -> ParseResult e (a, b)
withOrig a -> ParseResult e b
f a
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a
x, ) forall a b. (a -> b) -> a -> b
$ a -> ParseResult e b
f a
x

{-------------------------------------------------------------------------------
  Parsers for specific (but not @large-anon@ specific) constructs
-------------------------------------------------------------------------------}

-- | Generic constraint parser
--
-- TODO: If we add some parsing infra to ghc-tcplugin-api, maybe a (form of)
-- this function could live there too.
parseConstraint ::
     HasCallStack
  => (Class -> [Type] -> Maybe a) -- ^ Do we want to try and match against this?
  -> (a -> Maybe b)               -- ^ Parser for the class arguments
  -> Ct                           -- ^ Constraint to parse
  -> ParseResult e (GenLocated CtLoc b)
parseConstraint :: forall a b e.
HasCallStack =>
(Class -> [Type] -> Maybe a)
-> (a -> Maybe b) -> Ct -> ParseResult e (GenLocated CtLoc b)
parseConstraint Class -> [Type] -> Maybe a
p a -> Maybe b
f Ct
ct = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall l e. l -> e -> GenLocated l e
L forall a b. (a -> b) -> a -> b
$ Ct -> CtLoc
ctLoc Ct
ct) forall a b. (a -> b) -> a -> b
$
    -- TODO: classify up to equalities..?
    case Type -> Pred
classifyPredType (Ct -> Type
ctPred Ct
ct) of
      ClassPred Class
cls [Type]
args | Just a
a <- Class -> [Type] -> Maybe a
p Class
cls [Type]
args ->
        case a -> Maybe b
f a
a of
          Just b
parsed ->
            forall e a. a -> ParseResult e a
ParseOk b
parsed
          Maybe b
Nothing ->
            forall a. String -> a
panic forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
                String
"Unexpected "
              , SDoc -> String
showSDocUnsafe (forall a. Outputable a => a -> SDoc
ppr Class
cls)
              , String
" constraint with arguments:\n"
              , [String] -> String
unlines (forall a b. (a -> b) -> [a] -> [b]
map (SDoc -> String
showSDocUnsafe forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Outputable a => a -> SDoc
ppr) [Type]
args)
              , String
"\nat\n"
              , CallStack -> String
prettyCallStack HasCallStack => CallStack
callStack
              ]
      Pred
_otherwise ->
        forall e a. ParseResult e a
ParseNoMatch

-- | Specialization of 'parseConstraint', just checking the class name
parseConstraint' ::
     HasCallStack
  => Class                        -- ^ Predicate we want to match against
  -> ([Type] -> Maybe a)          -- ^ Parser for the class arguments
  -> Ct                           -- ^ Constraint to parse
  -> ParseResult e (GenLocated CtLoc a)
parseConstraint' :: forall a e.
HasCallStack =>
Class
-> ([Type] -> Maybe a) -> Ct -> ParseResult e (GenLocated CtLoc a)
parseConstraint' Class
cls = forall a b e.
HasCallStack =>
(Class -> [Type] -> Maybe a)
-> (a -> Maybe b) -> Ct -> ParseResult e (GenLocated CtLoc b)
parseConstraint Class -> [Type] -> Maybe [Type]
p
  where
    p :: Class -> [Type] -> Maybe [Type]
    p :: Class -> [Type] -> Maybe [Type]
p Class
cls' [Type]
args = if Class
cls forall a. Eq a => a -> a -> Bool
== Class
cls' then forall a. a -> Maybe a
Just [Type]
args else forall a. Maybe a
Nothing

-- | Parse @x ': xs == (':) x xs == ((':) x) xs@
parseCons :: TyConSubst -> Type -> Maybe (Type, Type)
parseCons :: TyConSubst -> Type -> Maybe (Type, Type)
parseCons TyConSubst
tcs Type
t = do
    [Type]
args <- TyConSubst -> TyCon -> Type -> Maybe [Type]
parseInjTyConApp TyConSubst
tcs TyCon
promotedConsDataCon Type
t
    case [Type]
args of
      [Type
_k, Type
x, Type
xs] -> forall a. a -> Maybe a
Just (Type
x, Type
xs)
      [Type]
_otherwise  -> forall a. Maybe a
Nothing

-- | Parse @'[]@
parseNil :: TyConSubst -> Type -> Maybe ()
parseNil :: TyConSubst -> Type -> Maybe ()
parseNil TyConSubst
tcs Type
t = do
    [Type]
args <- TyConSubst -> TyCon -> Type -> Maybe [Type]
parseInjTyConApp TyConSubst
tcs TyCon
promotedNilDataCon Type
t
    case [Type]
args of
      [Type
_k]       -> forall a. a -> Maybe a
Just ()
      [Type]
_otherwise -> forall a. Maybe a
Nothing

-- | Parse application of an injective type constructor
parseInjTyConApp :: TyConSubst -> TyCon -> Type -> Maybe [Type]
parseInjTyConApp :: TyConSubst -> TyCon -> Type -> Maybe [Type]
parseInjTyConApp TyConSubst
tcs TyCon
tyCon Type
t = do
    NonEmpty (TyCon, [Type])
splits <- TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type]))
splitTyConApp_upTo TyConSubst
tcs Type
t

    -- At this point we might have multiple matches
    --
    -- > t ~ TyCon1 args1
    -- > t ~ TyCon1 args1'
    -- > t ~ TyCon2 args2
    -- > ..
    --
    -- We are only interested in the equalities with @tyCon@ at the head, but
    -- this may still leave us with multiple equalities
    --
    -- > t ~ tyCon args1
    -- > t ~ tyCon args1'
    --
    -- When this is the case, however, by injectivity of 'tyCon' we know that
    -- @args1 ~ args1'@, so we can just return /any/ of the matches; we will
    -- return the first.
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup TyCon
tyCon (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (TyCon, [Type])
splits)