{-|
Module      : HsToCoq.Coq.Gallina.Rewrite
Description : Rewrite rule enginge for the Gallina AST
Copyright   : Copyright © 2018 Joachim Breitner
License     : MIT
Maintainer  : antal.b.sz@gmail.com
Stability   : experimental

This module implements rewrite rules. So far, this is straight forward and
not very efficiently.

Efficiency of finding a rewrite rule (candidate) could be greatly improved by
 * sorting the rewrite rules into a map, based on the outermost functions.
 * implementing an actual trie
This can be relevant when the numer of rewrite rules increases.

Another improvemet would be to detect bad patterns (non-linear ones,
unsupported AST nodes) and report such errors upon parsing.
-}

{-# LANGUAGE OverloadedLists #-}
module HsToCoq.Coq.Gallina.Rewrite (Rewrite(..), Rewrites, rewrite) where

import qualified Data.Map as M
import qualified Data.Set as S

import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Writer.Strict
import Data.Foldable

import HsToCoq.Coq.Subst
import HsToCoq.Coq.Gallina
import HsToCoq.Coq.Gallina.Util

data Rewrite = Rewrite
    { patternVars :: [Ident]
    , lhs :: Term
    , rhs :: Term
    }
  deriving (Eq, Ord, Show)

-- | This could be replaced by something more efficient, if needed.
-- See comments in HsToCoq.Coq.Gallina.Rewrite.
type Rewrites = [Rewrite]

-- | Applies all rewrite rules that match at the root of the given term.
rewrite :: Rewrites -> Term -> Term
rewrite rws t = foldl' (flip rewrite1) t rws

rewrite1 :: Rewrite -> Term -> Term
rewrite1 (Rewrite patVars lhs rhs) term
    | Just s <- match patVars lhs term
    = subst s rhs
    | otherwise
    = term

-- | Normalizes the outermost constructor
-- (Maybe we should drop InFix completely)
norm :: Term -> Term
norm (HsString s) = String s
norm t | Right (f, args) <- collectArgs t = appList (Qualid f) (map PosArg args)
norm t = t

match :: [Ident] -> Term -> Term -> Maybe (M.Map Qualid Term)
match patVars lhs term = execWriterT (go lhs term)
  where
    isPatVar = (`S.member` S.fromList patVars)

    go :: Term -> Term -> WriterT (M.Map Qualid Term) Maybe ()
    go lhs term = go' (norm lhs) (norm term)

    go' :: Term -> Term -> WriterT (M.Map Qualid Term) Maybe ()
    go' (String s1) (String s2) = guard (s1 == s2)
    go' (Num n1) (Num n2)       = guard (n1 == n2)
    go' (Qualid qid@(Bare v)) t | isPatVar v = do
        tell (M.singleton qid t)
        return ()
    go' (Qualid pqid) (Qualid qid) = guard (pqid == qid)
    go' (App pf pa) (App f a) = do
        go pf f
        zipWithSameLengthM_ goA (toList pa) (toList a)
    -- Simple matching of simple `match`es only – no dependent match clauses, no
    -- return types, no reordering match clauses.  We are VERY DUMB ABOUT BINDING.
    go' (Match scruts1 Nothing eqs1) (Match scruts2 Nothing eqs2) = do
        zipWithSameLengthM_ goMI scruts1 scruts2
        zipWithSameLengthM_ goE  eqs1    eqs2
    go' _lhs _term = mzero

    goA :: Arg -> Arg -> WriterT (M.Map Qualid Term) Maybe ()
    goA (PosArg pt) (PosArg p) = go pt p
    goA (NamedArg pi pt) (NamedArg i p) = do
        guard (pi == i)
        go pt p
    goA _lhs _term = mzero

    -- Handle simple `MatchItem`s only – terms without `as` or `in` clauses
    goMI (MatchItem tm1 Nothing Nothing) (MatchItem tm2 Nothing Nothing) =
      go tm1 tm2
    goMI _ _ =
      mzero

    goE (Equation mps1 tm1) (Equation mps2 tm2) = do
      zipWithSameLengthM_ goMP mps1 mps2
      go tm1 tm2

    goMP (MultPattern mps1) (MultPattern mps2) =
      zipWithSameLengthM_ goP mps1 mps2

    -- We are VERY DUMB ABOUT BINDING.
    goP UnderscorePat    UnderscorePat    = pure ()
    goP (NumPat num1)    (NumPat num2)    = guard $ num1 == num2
    goP (StringPat str1) (StringPat str2) = guard $ str1 == str2
    goP (InScopePat pat1 scope1) (InScopePat pat2 scope2) = do
      guard $ scope1 == scope2
      goP pat1 pat2
    goP (ArgsPat con1 args1) (ArgsPat con2 args2) = do
      goPQid con1 con2
      zipWithSameLengthM_ goP args1 args2
    goP (ExplicitArgsPat con1 args1) (ExplicitArgsPat con2 args2) = do
      goPQid con1 con2
      zipWithSameLengthM_ goP args1 args2
    goP (InfixPat lhs1 op1 rhs1) (InfixPat lhs2 op2 rhs2) = do
      goPQid (Bare op1) (Bare op2)
      goP lhs1 lhs2
      goP rhs1 rhs2
    goP (AsPat pat1 var1) (AsPat pat2 var2) = do
      goPQid var1 var2
      goP pat1 pat2
    goP (QualidPat qid1@(Bare v)) p | isPatVar v, Just t <- patToTerm p =
      tell $ M.singleton qid1 t
    goP (QualidPat qid1) (QualidPat qid2) =
      guard $ qid1 == qid2
    goP (OrPats pats1) (OrPats pats2) =
      zipWithSameLengthM_ goOP pats1 pats2
    goP _ _ = mzero

    patToTerm (ArgsPat con args)         = appList (Qualid con) <$> traverse (fmap PosArg . patToTerm) args
    patToTerm (ExplicitArgsPat con args) = ExplicitApp con . toList <$> traverse patToTerm args
    patToTerm (InfixPat lhs op rhs)      = App2 (Qualid $ Bare op) <$> patToTerm lhs <*> patToTerm rhs
    patToTerm (InScopePat p scope)       = InScope <$> patToTerm p <*> pure scope
    patToTerm (QualidPat qid)            = Just $ Qualid qid
    patToTerm (NumPat n)                 = Just $ Num n
    patToTerm (StringPat s)              = Just $ String s
    patToTerm (AsPat _ _)                = Nothing
    patToTerm UnderscorePat              = Nothing
    patToTerm (OrPats _)                 = Nothing
              
    goPQid lhs@(Bare v) rhs | isPatVar v = tell $ M.singleton lhs (Qualid rhs)
    goPQid lhs          rhs              = guard $ lhs == rhs

    goOP (OrPattern pats1) (OrPattern pats2) = zipWithSameLengthM_ goP pats1 pats2

    zipWithSameLengthM_ :: (Alternative f, Foldable g, Foldable h) => (a -> b -> f c) -> g a -> h b -> f ()
    zipWithSameLengthM_ f as0 bs0 =
      let zipGo (a:as) (b:bs) = f a b *> zipGo as bs
          zipGo []     []     = pure ()
          zipGo _      _      = empty
      in zipGo (toList as0) (toList bs0)
