content: {-# LANGUAGE CPP, LambdaCase, TupleSections, ViewPatterns #-}
  {-# OPTIONS -Wno-name-shadowing #-}
  module TypeLevel.Rewrite.Internal.ApplyRules where
  
  import Control.Applicative
  import Control.Monad
  import Control.Monad.Trans.Class
  import Control.Monad.Trans.State
  import Data.Foldable (asum, for_)
  import Data.Map (Map)
  import Data.Maybe (listToMaybe, maybeToList)
  import Data.Traversable
  import qualified Data.Map as Map
  
  -- GHC API
  #if MIN_VERSION_ghc(9,0,0)
  import GHC.Plugins (TyVar)
  #else
  import Type (TyVar)
  #endif
  
  -- term-rewriting API
  import Data.Rewriting.Rule (Rule(..))
  import Data.Rewriting.Substitution (gApply)
  import Data.Rewriting.Term (Term(..))
  import qualified Data.Rewriting.Substitution.Type as Substitution
  
  import TypeLevel.Rewrite.Internal.TypeEq
  import TypeLevel.Rewrite.Internal.TypeNode
  import TypeLevel.Rewrite.Internal.TypeRule
  import TypeLevel.Rewrite.Internal.TypeSubst
  import TypeLevel.Rewrite.Internal.TypeTerm
  
  
  type Subst = Map TyVar (Term TypeNode TypeEq)
  
  applyRules
    :: Traversable t
    => TypeSubst
    -> [TypeRule]
    -> t TypeTerm
    -> Maybe (TypeRule,t TypeTerm)
  applyRules typeSubst rules inputs
    = annotatedTraverseFirst (multiRewrite typeSubst rules) inputs
  
  multiRewrite
    :: TypeSubst
    -> [TypeRule]
    -> TypeTerm
    -> Maybe (TypeRule, TypeTerm)
  multiRewrite typeSubst rules input
    = asum
      [ (rule,) <$> singleRewrite typeSubst rule input
      | rule <- rules
      ]
  
  -- >>> singleRewrite (F x (F x y) ~ F x y) [F a (F a b)]
  -- Just [F a b]
  singleRewrite
    :: TypeSubst
    -> TypeRule
    -> TypeTerm
    -> Maybe TypeTerm
  singleRewrite typeSubst rule input@(Fun inputF inputXS)
      = topLevelRewrite typeSubst rule input
    <|> (Fun inputF <$> traverseFirst (singleRewrite typeSubst rule) inputXS)
  singleRewrite typeSubst rule input
    = topLevelRewrite typeSubst rule input
  
  
  -- >>> topLevelRewrite (F x (F x y) ~ F x y) (F a (F a b))
  -- Just (F a b)
  topLevelRewrite
    :: TypeSubst
    -> TypeRule
    -> TypeTerm
    -> Maybe TypeTerm
  topLevelRewrite typeSubst (Rule pattern0 pattern') input0 = do
    subst <- execStateT (go pattern0 input0) Map.empty
    gApply (Substitution.fromMap subst) pattern'
    where
      go
        :: Term TypeNode TyVar
        -> TypeTerm
        -> StateT Subst Maybe ()
      go (Var var) input = do
        subst <- get
        case Map.lookup var subst of
          Nothing -> do
            modify (Map.insert var input)
          Just term -> do
            guard (input == term)
      go (Fun patternF patternXS)
         (Fun inputF inputXS)
         = do
        guard (patternF == inputF)
        guard (length patternXS == length inputXS)
        for_ (zip patternXS inputXS) $ \(pattern, input) -> do
          go pattern input
      go pattern (Var var) = do
        let possibleReplacements = fmap snd
                                 . filter ((== var) . fst)
                                 $ typeSubst
        asum $ fmap (go pattern) possibleReplacements
  
  -- >>> traverseFirst (\x -> if even x then Just (10 + x) else Nothing) [1,3,5]
  -- Nothing
  -- >>> traverseFirst (\x -> if even x then Just (10 + x) else Nothing) [1,2,4]
  -- Just [1,12,4]
  traverseFirst
    :: Traversable t
    => (a -> Maybe a)
    -> t a
    -> Maybe (t a)
  traverseFirst f = listToMaybe . traverseAll f
  
  annotatedTraverseFirst
    :: Traversable t
    => (a -> Maybe (annotation, a))
    -> t a
    -> Maybe (annotation, t a)
  annotatedTraverseFirst f = listToMaybe . annotatedTraverseAll f
  
  -- >>> traverseAll (\x -> if even x then Just (10 + x) else Nothing) [1,3,5]
  -- []
  -- >>> traverseAll (\x -> if even x then Just (10 + x) else Nothing) [1,2,4]
  -- [[1,12,4], [1,2,14]]
  traverseAll
    :: Traversable t
    => (a -> Maybe a)
    -> t a
    -> [t a]
  traverseAll f
    = fmap snd
    . annotatedTraverseAll (fmap ((),) . f)
  
  annotatedTraverseAll
    :: Traversable t
    => (a -> Maybe (annotation, a))
    -> t a
    -> [(annotation, t a)]
  annotatedTraverseAll f ta = flip evalStateT Nothing $ do
    ta' <- for ta $ \a -> do
      get >>= \case
        Just _ -> do
          -- already picked one
          pure a
        Nothing -> do
          pickIt <- lift [True,False]
          if pickIt
            then do
              (annotation, a) <- lift $ maybeToList $ f a
              put (Just annotation)
              pure a
            else do
              pure a
    maybeAnnotation <- get
    annotation <- lift $ maybeToList maybeAnnotation
    pure (annotation, ta')
  