{-# LANGUAGE PatternSynonyms #-} -- For Comb
{-# LANGUAGE TemplateHaskell #-} -- For branch
{-# LANGUAGE ViewPatterns #-} -- For unSomeComb
{-# OPTIONS_GHC -fno-warn-orphans #-} -- For MakeLetName TH.Name
-- | Bottom-up optimization of 'Comb'inators,
-- reexamining downward as needed after each optimization.
module Symantic.Parser.Grammar.Optimize where

import Data.Bool (Bool(..))
import Data.Either (Either(..), either)
import Data.Eq (Eq(..))
import Data.Function (($), (.))
import Data.Kind (Constraint)
import Data.Maybe (Maybe(..))
import Data.Set (Set)
import Data.Functor.Identity (Identity(..))
import Data.Functor.Product (Product(..))
import Type.Reflection (Typeable, typeRep, eqTypeRep, (:~~:)(..))
import qualified Data.Foldable as Foldable
import qualified Data.Functor as Functor
import qualified Data.List as List

import Symantic.Parser.Grammar.Combinators hiding (code)
import Symantic.Parser.Grammar.Production
import Symantic.Univariant.Letable
import Symantic.Univariant.Trans
import qualified Symantic.Parser.Grammar.Production as Prod
import qualified Symantic.Univariant.Data as H
import qualified Symantic.Univariant.Lang as H

{-
import Data.Function (($), flip)
import Debug.Trace (trace)

(&) = flip ($)
infix 0 &
-}

-- * Type 'OptimizeGrammar'
type OptimizeGrammar = SomeComb

optimizeGrammar ::
  Trans (SomeComb repr) repr =>
  SomeComb repr a -> repr a
optimizeGrammar = trans

-- * Data family 'Comb'
-- | 'Comb'inators of the 'Grammar'.
-- This is an extensible data-type.
data family Comb
  (comb :: ReprComb -> Constraint)
  :: ReprComb -> ReprComb

-- | Convenient utility to pattern-match a 'SomeComb'.
pattern Comb :: Typeable comb => Comb comb repr a -> SomeComb repr a
pattern Comb x <- (unSomeComb -> Just x)

-- ** Type 'SomeComb'
-- | Some 'Comb'inator existentialized over the actual combinator symantic class.
-- Useful to handle a list of 'Comb'inators
-- without requiring impredicative quantification.
-- Must be used by pattern-matching
-- on the 'SomeComb' data-constructor,
-- to bring the constraints in scope.
--
-- The optimizations are directly applied within it,
-- to avoid introducing an extra newtype,
-- this also give a more understandable code.
data SomeComb repr a =
  forall comb.
  (Trans (Comb comb repr) repr, Typeable comb) =>
  SomeComb (Comb comb repr a)

instance Trans (SomeComb repr) repr where
  trans (SomeComb x) = trans x

-- | @(unSomeComb c :: 'Maybe' ('Comb' comb repr a))@
-- extract the data-constructor from the given 'SomeComb'
-- iif. it belongs to the @('Comb' comb repr a)@ data-instance.
unSomeComb ::
  forall comb repr a.
  Typeable comb =>
  SomeComb repr a -> Maybe (Comb comb repr a)
unSomeComb (SomeComb (c::Comb c repr a)) =
  case typeRep @comb `eqTypeRep` typeRep @c of
    Just HRefl -> Just c
    Nothing -> Nothing

-- CombAlternable
data instance Comb CombAlternable repr a where
  Alt :: Exception -> SomeComb repr a -> SomeComb repr a -> Comb CombAlternable repr a
  Empty :: Comb CombAlternable repr a
  Failure :: SomeFailure -> Comb CombAlternable repr a
  Throw :: ExceptionLabel -> Comb CombAlternable repr a
  Try :: SomeComb repr a -> Comb CombAlternable repr a
instance CombAlternable repr => Trans (Comb CombAlternable repr) repr where
  trans = \case
    Alt exn x y -> alt exn (trans x) (trans y)
    Empty -> empty
    Failure sf -> failure sf
    Throw exn -> throw exn
    Try x -> try (trans x)
instance
  ( CombAlternable repr
  , CombApplicable repr
  , CombLookable repr
  , CombMatchable repr
  , CombSelectable repr
  ) => CombAlternable (SomeComb repr) where
  empty = SomeComb Empty
  failure sf = SomeComb (Failure sf)

  alt _exn p@(Comb Pure{}) _ = p
    -- & trace "Left Catch Law"
  alt _exn (Comb Empty) u = u
    -- & trace "Left Neutral Law"
  alt _exn u (Comb Empty) = u
    -- & trace "Right Neutral Law"
  alt exn (Comb (Alt exn' u v)) w | exn' == exn = u <|> (v <|> w)
    -- See Lemma 1 (Associativity of choice for labeled PEGs)
    -- in https://doi.org/10.1145/2851613.2851750
    -- & trace "Associativity Law"
  alt exn (Comb (Look p)) (Comb (Look q)) = look (alt exn (try p) q)
    -- & trace "Distributivity Law"
  alt exn x y = SomeComb (Alt exn x y)

  throw exn = SomeComb (Throw exn)

  try (Comb (p :$>: x)) = try p $> x
    -- & trace "Try Interchange Law"
  try (Comb (f :<$>: p)) = f <$> try p
    -- & trace "Try Interchange Law"
  try x = SomeComb (Try x)

-- CombApplicable
data instance Comb CombApplicable repr a where
  Pure :: Production a -> Comb CombApplicable repr a
  (:<*>:) :: SomeComb repr (a -> b) -> SomeComb repr a -> Comb CombApplicable repr b
  (:<*:) :: SomeComb repr a -> SomeComb repr b -> Comb CombApplicable repr a
  (:*>:) :: SomeComb repr a -> SomeComb repr b -> Comb CombApplicable repr b
infixl 4 :<*>:, :<*:, :*>:
pattern (:<$>:) :: Production (a -> b) -> SomeComb repr a -> Comb CombApplicable repr b
pattern t :<$>: x <- Comb (Pure t) :<*>: x
pattern (:$>:) :: SomeComb repr a -> Production b -> Comb CombApplicable repr b
pattern x :$>: t <- x :*>: Comb (Pure t)
instance CombApplicable repr => Trans (Comb CombApplicable repr) repr where
  trans = \case
    Pure x -> pure (optimizeProduction x)
    f :<*>: x -> trans f <*> trans x
    x :<*: y -> trans x <* trans y
    x :*>: y -> trans x *> trans y
instance
  ( CombApplicable repr
  , CombAlternable repr
  , CombLookable repr
  , CombMatchable repr
  , CombSelectable repr
  ) => CombApplicable (SomeComb repr) where
  pure = SomeComb . Pure
  f <$> Comb (Branch b l r) =
    branch b
      ((H..) H..@ f <$> l)
      ((H..) H..@ f <$> r)
    -- & trace "Branch Distributivity Law"
  f <$> Comb (Conditional a ps bs d) =
    conditional a ps
      ((f <$>) Functor.<$> bs)
      (f <$> d)
    -- & trace "Conditional Distributivity Law"
  -- Being careful here to use (<*>),
  -- instead of SomeComb (f <$> unOptComb x),
  -- in order to apply the optimizations of (<*>).
  f <$> x = pure f <*> x

  x <$ u = u $> x
    -- & trace "Commutativity Law"

  Comb Empty <*> _ = empty
    -- & trace "App Right Absorption Law"
  u <*> Comb Empty = u *> empty
    -- & trace "App Failure Weakening Law"
  Comb (Pure f) <*> Comb (Pure x) = pure (f H..@ x)
    -- & trace "Homomorphism Law"
  {-
  Comb (Pure f) <*> Comb (g :<$>: p) =
    -- This is basically a shortcut,
    -- it can be caught by one Composition Law
    -- and two Homomorphism Law.
    (H..) H..@ f H..@ g <$> p
    -- & trace "Functor Composition Law"
  -}
  u <*> Comb (Pure x) = H.flip H..@ (H.$) H..@ x <$> u
    -- & trace "Interchange Law"
  u <*> Comb (v :<*>: w) = (((H..) <$> u) <*> v) <*> w
    -- & trace "Composition Law"
  Comb (u :*>: v) <*> w = u *> (v <*> w)
    -- & trace "Reassociation Law 1"
  u <*> Comb (v :<*: w) = (u <*> v) <* w
    -- & trace "Reassociation Law 2"
  u <*> Comb (v :$>: x) = (u <*> pure x) <* v
    -- & trace "Reassociation Law 3"
  p <*> Comb (NegLook q) =
    (p <*> pure H.unit) <* negLook q
    -- & trace "Absorption Law"
  x <*> y = SomeComb (x :<*>: y)

  Comb Empty *> _ = empty
    -- & trace "App Right Absorption Law"
  Comb (_ :<$>: p) *> q = p *> q
    -- & trace "Right Absorption Law"
  Comb Pure{} *> u = u
    -- & trace "Identity Law"
  Comb (u :$>: _) *> v = u *> v
    -- & trace "Identity Law"
  u *> Comb (v :*>: w) = (u *> v) *> w
    -- & trace "Associativity Law"
  x *> y = SomeComb (x :*>: y)

  Comb Empty <* _ = empty
    -- & trace "App Right Absorption Law"
  u <* Comb Empty = u *> empty
    -- & trace "App Failure Weakening Law"
  p <* Comb (_ :<$>: q) = p <* q
    -- & trace "Left Absorption Law"
  u <* Comb Pure{} = u
    -- & trace "Identity Law"
  u <* Comb (v :$>: _) = u <* v
    -- & trace "Identity Law"
  Comb (u :<*: v) <* w = u <* (v <* w)
    -- & trace "Associativity Law"
  x <* y = SomeComb (x :<*: y)

-- CombFoldable
data instance Comb CombFoldable repr a where
  ChainPreC :: SomeComb repr (a -> a) -> SomeComb repr a -> Comb CombFoldable repr a
  ChainPostC :: SomeComb repr a -> SomeComb repr (a -> a) -> Comb CombFoldable repr a
instance CombFoldable repr => Trans (Comb CombFoldable repr) repr where
  trans = \case
    ChainPreC x y -> chainPre (trans x) (trans y)
    ChainPostC x y -> chainPost (trans x) (trans y)
instance CombFoldable repr => CombFoldable (SomeComb repr) where
  chainPre x = SomeComb . ChainPreC x
  chainPost x = SomeComb . ChainPostC x

-- Letable
data instance Comb (Letable letName) repr a where
  Shareable :: letName -> SomeComb repr a -> Comb (Letable letName) repr a
  Ref :: Bool -> letName -> Comb (Letable letName) repr a
instance
  Letable letName repr =>
  Trans (Comb (Letable letName) repr) repr where
  trans = \case
    Shareable n x -> shareable n (trans x)
    Ref isRec n -> ref isRec n
instance
  (Letable letName repr, Typeable letName) =>
  Letable letName (SomeComb repr) where
  shareable n = SomeComb . Shareable n
  ref isRec = SomeComb . Ref isRec

-- Letsable
data instance Comb (Letsable letName) repr a where
  Lets :: LetBindings letName (SomeComb repr) ->
          SomeComb repr a -> Comb (Letsable letName) repr a
instance
  Letsable letName repr =>
  Trans (Comb (Letsable letName) repr) repr where
  trans = \case
    Lets defs x -> lets ((\(SomeLet sub) -> SomeLet (trans sub)) Functor.<$> defs) (trans x)
instance
  (Letsable letName repr, Typeable letName) =>
  Letsable letName (SomeComb repr) where
  lets defs = SomeComb . Lets defs

-- CombLookable
data instance Comb CombLookable repr a where
  Look :: SomeComb repr a -> Comb CombLookable repr a
  NegLook :: SomeComb repr a -> Comb CombLookable repr ()
  Eof :: Comb CombLookable repr ()
instance CombLookable repr => Trans (Comb CombLookable repr) repr where
  trans = \case
    Look x -> look (trans x)
    NegLook x -> negLook (trans x)
    Eof -> eof
instance
  ( CombAlternable repr
  , CombApplicable repr
  , CombLookable repr
  , CombSelectable repr
  , CombMatchable repr
  ) => CombLookable (SomeComb repr) where
  look p@(Comb Pure{}) = p
    -- & trace "Pure Look Law"
  look p@(Comb Empty) = p
    -- & trace "Dead Look Law"
  look (Comb (Look x)) = look x
    -- & trace "Idempotence Law"
  look (Comb (NegLook x)) = negLook x
    -- & trace "Left Identity Law"
  look (Comb (p :$>: x)) = look p $> x
    -- & trace "Interchange Law"
  look (Comb (f :<$>: p)) = f <$> look p
    -- & trace "Interchange Law"
  look x = SomeComb (Look x)

  negLook (Comb Pure{}) = empty
    -- & trace "Pure Negative-Look"
  negLook (Comb Empty) = pure H.unit
    -- & trace "Dead Negative-Look"
  negLook (Comb (NegLook x)) = look (try x *> pure H.unit)
    -- & trace "Double Negation Law"
  negLook (Comb (Try x)) = negLook x
    -- & trace "Zero Consumption Law"
  negLook (Comb (Look x)) = negLook x
    -- & trace "Right Identity Law"
  negLook (Comb (Alt _exn (Comb (Try p)) q)) = negLook p *> negLook q
    -- FIXME: see if this really holds for all exceptions.
    -- & trace "Transparency Law"
  negLook (Comb (p :$>: _)) = negLook p
    -- & trace "NegLook Idempotence Law"
  negLook x = SomeComb (NegLook x)

  eof = SomeComb Eof

-- CombMatchable
data instance Comb CombMatchable repr a where
  Conditional :: Eq a =>
    SomeComb repr a ->
    [Production (a -> Bool)] ->
    [SomeComb repr b] ->
    SomeComb repr b ->
    Comb CombMatchable repr b
instance CombMatchable repr => Trans (Comb CombMatchable repr) repr where
  trans = \case
    Conditional a ps bs b ->
      conditional (trans a)
        (optimizeProduction Functor.<$> ps)
        (trans Functor.<$> bs) (trans b)
instance
  ( CombApplicable repr
  , CombAlternable repr
  , CombLookable repr
  , CombSelectable repr
  , CombMatchable repr
  ) => CombMatchable (SomeComb repr) where
  conditional (Comb Empty) _ _ d = d
    -- & trace "Conditional Absorption Law"
  conditional p _ qs (Comb Empty)
    | Foldable.all (\case { Comb Empty -> True; _ -> False }) qs = p *> empty
      -- & trace "Conditional Weakening Law"
  conditional a _ps bs (Comb Empty)
    | Foldable.all (\case { Comb Empty -> True; _ -> False }) bs = a *> empty
      -- & trace "Conditional Weakening Law"
  conditional (Comb (Pure a)) ps bs d =
    Foldable.foldr (\(p, b) next ->
      if runValue (p H..@ a) then b else next
    ) d (List.zip ps bs)
    -- & trace "Conditional Pure Law"
  conditional a ps bs d = SomeComb (Conditional a ps bs d)

-- CombSatisfiable
data instance Comb (CombSatisfiable tok) repr a where
  -- | To include the @('Set' 'SomeFailure')@ is a little kludge
  -- it would be cleaner to be able to pattern-match
  -- on @(alt exn (Comb 'Satisfy'{}) (Failure{}))@
  -- when generating 'Program', but this is not easily possible then
  -- because data types have already been converted back to class methods,
  -- hence pattern-matching is no longer possible
  -- (at least not without reintroducing data types).
  SatisfyOrFail ::
    CombSatisfiable tok repr =>
    Set SomeFailure ->
    Production (tok -> Bool) ->
    Comb (CombSatisfiable tok) repr tok
instance
  CombSatisfiable tok repr =>
  Trans (Comb (CombSatisfiable tok) repr) repr where
  trans = \case
    SatisfyOrFail fs p -> satisfyOrFail fs (optimizeProduction p)
instance
  (CombSatisfiable tok repr, Typeable tok) =>
  CombSatisfiable tok (SomeComb repr) where
  satisfyOrFail fs = SomeComb . SatisfyOrFail fs

-- CombSelectable
data instance Comb CombSelectable repr a where
  Branch ::
    SomeComb repr (Either a b) ->
    SomeComb repr (a -> c) ->
    SomeComb repr (b -> c) ->
    Comb CombSelectable repr c
instance CombSelectable repr => Trans (Comb CombSelectable repr) repr where
  trans = \case
    Branch lr l r -> branch (trans lr) (trans l) (trans r)
instance
  ( CombApplicable repr
  , CombAlternable repr
  , CombLookable repr
  , CombSelectable repr
  , CombMatchable repr
  ) => CombSelectable (SomeComb repr) where
  branch (Comb Empty) _ _ = empty
    -- & trace "Branch Absorption Law"
  branch b (Comb Empty) (Comb Empty) = b *> empty
    -- & trace "Branch Weakening Law"
  branch (Comb (Pure lr)) l r =
    case runValue lr of
      Left value -> l <*> pure (Pair v c)
        where
        v = H.SomeData $ H.Var $ Identity value
        c = H.SomeData $ H.Var
          [|| case $$(runCode lr) of Left x -> x ||]
      Right value -> r <*> pure (Pair v c)
        where
        v = H.SomeData $ H.Var $ Identity value
        c = H.SomeData $ H.Var
          [|| case $$(runCode lr) of Right x -> x ||]
    -- & trace "Branch Pure Either Law"
  branch b (Comb (Pure l)) (Comb (Pure r)) =
    Pair v c <$> b
    -- & trace "Branch Generalised Identity Law"
    where
    v = H.SomeData $ H.Var $ Identity $ either (runValue l) (runValue r)
    c = H.SomeData $ H.Var [|| either $$(runCode l) $$(runCode r) ||]
  branch (Comb (x :*>: y)) p q = x *> branch y p q
    -- & trace "Interchange Law"
  branch b l (Comb Empty) =
    branch (pure (Pair v c) <*> b) empty l
    -- & trace "Negated Branch Law"
    where
    v = H.SomeData $ H.Var $ Identity $ either Right Left
    c = H.SomeData $ H.Var $ [||either Right Left||]
  branch (Comb (Branch b (Comb Empty) (Comb (Pure lr)))) (Comb Empty) br =
    branch (pure (Pair v c) <*> b) empty br
    -- & trace "Branch Fusion Law"
    where
    v = H.SomeData $ H.Var $ Identity $ \case
      Left{} -> Left ()
      Right r ->
        case runValue lr r of
          Left{} -> Left ()
          Right rr -> Right rr
    c = H.SomeData $ H.Var
      [|| \case Left{} -> Left ()
                Right r -> case $$(runCode lr) r of
                             Left{} -> Left ()
                             Right rr -> Right rr ||]
  branch b l r = SomeComb (Branch b l r)