{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Language.REST.Internal.Rewrite
  ( Rewrite(..)
  , Subst
  , named
  , subst
  , unify
  ) where

import GHC.Generics (Generic)

import           Data.Maybe (isNothing)
import           Data.Hashable
import qualified Data.HashMap.Strict as M
import qualified Data.HashSet as S
import           Text.Printf

import Language.REST.RewriteRule
import Language.REST.MetaTerm as MT
import Language.REST.RuntimeTerm


-- | @Rewrite t u s@ defines a rewrite rule \( t \rightarrow u \), with
--   an optional name @s@.
data Rewrite = Rewrite MetaTerm MetaTerm (Maybe String)
  deriving (Eq, Ord, Generic, Hashable, Show)

-- | 'Subst' is a mapping from variable names to 'RuntimeTerm's.
--   Normally this would be generated by unifying the left-hand-side of
--   a 'Rewrite' with a term.
type Subst = M.HashMap String RuntimeTerm

-- | @named r n@ assigns the name @n@ to rule @r@, replacing any
--   existing name
named :: Rewrite -> String -> Rewrite
named (Rewrite t u _) n = Rewrite t u (Just n)

-- | @subst s m@ replaces the variables in the 'MetaTerm' @m@ with 'RuntimeTerm's
--   in the substitution 's'. This function returns an error if any variables in 'm'
--   do not have a substituion
subst :: Subst -> MetaTerm -> RuntimeTerm
subst s (MT.Var v)  | Just t <- M.lookup v s = t
                    | otherwise
                    = error $ printf "No value for metavar %s during subst %s" (show v) (show s)
subst s (MT.RWApp op xs) = App op (map (subst s) xs)

unifyAll :: Subst -> [(MetaTerm, RuntimeTerm)] -> Maybe Subst
unifyAll su [] = Just su
unifyAll su ((x, y) : ts)
  | Just s <- unify x y su
  = unifyAll s ts
  | otherwise
  = Nothing

-- | @unify m r su@ extends the substitution @su@ to generate a new
--   substitution that unifies @m@ and @r@. Returns 'Nothing' if su
--   cannot be extended to unify the terms.
unify :: MetaTerm -> RuntimeTerm -> Subst -> Maybe Subst
unify (MT.Var s) term su | M.lookup s su == Just term
  = Just su
unify (MT.Var s) term su | isNothing (M.lookup s su)
  = Just $ M.insert s term su
unify (MT.RWApp o1 xs) (App o2 ys) su | o1 == o2 && length xs == length ys =
  unifyAll su (zip xs ys)
unify _ _ _ = Nothing

instance Monad m => RewriteRule m Rewrite RuntimeTerm where
  apply t (Rewrite left right _) = return $ S.unions $ map go (subTerms t)
    where
      go (t', tf) | Just su <- unify left t' M.empty = S.singleton (tf $ subst su right)
      go _                                = S.empty