{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | This module implements the optimizations to prune the
-- exploration of rewrites of terms that have been already considered
-- (section 6.4 of the REST paper).
module Language.REST.ExploredTerms
   (
     ExploredTerms
   , empty
   , insert
   , shouldExplore
   , size
   , visited
   , ExploreFuncs(..)
   , ExploreStrategy(..)
   )  where

import qualified Data.HashMap.Strict as M
import qualified Data.HashSet as S
import Data.Hashable

import Prelude hiding (lookup)

-- | 'ExploreStrategy' defines how 'shouldExplore' should decide whether or not
-- | to consider rewrites from a given term
data ExploreStrategy =
    ExploreAlways -- ^ Always explore, even when it's not necessary.
  | ExploreLessConstrained -- ^ Explore terms unless the constraints are stricter.
                           -- This may stil explore unnecessary paths, the terms
                           -- were already fully explored with the different constraints.
  | ExploreWhenNeeded -- ^ Explore terms unless the constraints are stricter OR if all
                      --   terms reachable via transitive rewrites were already explored.
  | ExploreOnce -- ^ Explore each term only once. This may cause some terms not to be
                --   explored if the terms leading to them were initially visited at
                --   strict constraints.


data ExploreFuncs term c m = EF
  { -- | When a term @t@ is visited at constraints @c0@, and then at constraints
    --   @c1@, the constraints for term @t@ is set to @c0 `union` c1@
    union           :: c -> c -> c
    -- | @c0 `subsumes` c1@ if @c0@ permits all orderings permited by @c1@
  , subsumes        :: c -> c -> m Bool
    -- | @'exRefine' c t u@ strengthens constraints @c@ to permit the rewrite step
    --   from @t@ to @u@. This is used to determine if considering term @u@ by rewriting
    --   from @t@ would permit more rewrite applications.
  , exRefine        :: c -> term -> term -> c
  }

-- | A mapping of terms, to the rewritten terms that need to be fully explored
-- | in order for this term to be fully explored
data ExploredTerms term c m =
  ET (M.HashMap term (c, S.HashSet term)) (ExploreFuncs term c m) ExploreStrategy

size :: ExploredTerms term c m -> Int
size (ET m _ _) = M.size m

empty :: ExploreFuncs term c m -> ExploreStrategy -> ExploredTerms term c m
empty = ET M.empty

visited :: (Eq term, Hashable term) => term -> ExploredTerms term c m -> Bool
visited t (ET m _ _) = M.member t m

insert :: (Eq term, Hashable term) => term -> c -> S.HashSet term -> ExploredTerms term c m -> ExploredTerms term c m
insert t oc s (ET etMap ef@(EF union _ _) strategy) = ET (M.insertWith go t (oc, s) etMap) ef strategy
  where
    go (oc1, s1) (oc2, s2) = (oc1 `union` oc2, S.union s1 s2)

lookup :: (Eq term, Hashable term) => term -> ExploredTerms term c m -> Maybe (c, S.HashSet term)
lookup t (ET etMap _ _) = M.lookup t etMap

-- | @isFullyExplored t c M = not explorable(t, c)@ where @explorable@ is
-- defined as in the REST paper. Also incorporates an optimization described
-- here: https://github.com/zgrannan/rest/issues/9
isFullyExplored :: forall term c m . (Monad m, Eq term, Hashable term, Hashable c, Eq c, Show c) =>
  term -> c -> ExploredTerms term c m -> m Bool
isFullyExplored t0 oc0 et@(ET _ (EF{subsumes,exRefine}) _) = result where

  result = go S.empty [(t0, oc0)]

  -- Arg 1: Steps that have already been seen
  -- Arg 2: Steps to consider
  go :: S.HashSet (term, c) -> [(term, c)] -> m Bool

  -- Completed worklist, this term is fully explored at these constraints
  go _ []       = return True

  -- Term `h` has been seen before at constraints `oc`
  go seen ((h, oc'):rest) | Just (oc, trms) <- lookup h et
                = do
                    exploringPathWouldNotPermitDifferentSteps <- oc `subsumes` oc'
                    if exploringPathWouldNotPermitDifferentSteps
                      then go seen' rest
                      else
                        let
                          -- Exploring `h` at these constraints
                          -- would allow exploration of each t in trms,
                          -- at the constraints generated by the step from h to t
                          trms' = S.map (\t -> (t, exRefine oc' h t)) trms
                          ts    = S.union trms' (S.fromList rest) `S.difference` seen'
                        in
                          go seen' (S.toList ts)
                  where
                    seen' = S.insert (h, oc') seen

  -- There exists a reachable term that has never previously been seen; not fully explored
  go _ _         = return False

-- | @'shouldExplore' t c et@ determines if rewrites originating from term @t@ at
--   constraints @c@ should be considered, given the already explored terms of @et@
--   and the associated 'ExploreStrategy'
shouldExplore :: forall term c m . (Monad m, Eq term, Hashable term, Eq c, Show c, Hashable c) =>
  term -> c -> ExploredTerms term c m -> m Bool
shouldExplore t oc et@(ET _ EF{subsumes} strategy) =
  case strategy of
    ExploreWhenNeeded      -> not <$> isFullyExplored t oc et
    ExploreOnce            -> return $ not $ visited t et
    ExploreAlways          -> return True
    ExploreLessConstrained ->
      case lookup t et of
        Just (oc', _) -> do
          s <- oc' `subsumes` oc
          return $ not s
        Nothing       -> return True