{-# 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@
    forall term c (m :: * -> *). ExploreFuncs term c m -> c -> c -> c
union           :: c -> c -> c
    -- | @c0 `subsumes` c1@ if @c0@ permits all orderings permited by @c1@
  , forall term c (m :: * -> *).
ExploreFuncs term c m -> c -> c -> m Bool
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.
  , forall term c (m :: * -> *).
ExploreFuncs term c m -> c -> term -> term -> c
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 :: forall term c (m :: * -> *). ExploredTerms term c m -> Int
size (ET HashMap term (c, HashSet term)
m ExploreFuncs term c m
_ ExploreStrategy
_) = HashMap term (c, HashSet term) -> Int
forall k v. HashMap k v -> Int
M.size HashMap term (c, HashSet term)
m

empty :: ExploreFuncs term c m -> ExploreStrategy -> ExploredTerms term c m
empty :: forall term c (m :: * -> *).
ExploreFuncs term c m -> ExploreStrategy -> ExploredTerms term c m
empty = HashMap term (c, HashSet term)
-> ExploreFuncs term c m
-> ExploreStrategy
-> ExploredTerms term c m
forall term c (m :: * -> *).
HashMap term (c, HashSet term)
-> ExploreFuncs term c m
-> ExploreStrategy
-> ExploredTerms term c m
ET HashMap term (c, HashSet term)
forall k v. HashMap k v
M.empty

visited :: (Eq term, Hashable term) => term -> ExploredTerms term c m -> Bool
visited :: forall term c (m :: * -> *).
(Eq term, Hashable term) =>
term -> ExploredTerms term c m -> Bool
visited term
t (ET HashMap term (c, HashSet term)
m ExploreFuncs term c m
_ ExploreStrategy
_) = term -> HashMap term (c, HashSet term) -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
M.member term
t HashMap term (c, HashSet term)
m

insert :: (Eq term, Hashable term) => term -> c -> S.HashSet term -> ExploredTerms term c m -> ExploredTerms term c m
insert :: forall term c (m :: * -> *).
(Eq term, Hashable term) =>
term
-> c
-> HashSet term
-> ExploredTerms term c m
-> ExploredTerms term c m
insert term
t c
oc HashSet term
s (ET HashMap term (c, HashSet term)
etMap ef :: ExploreFuncs term c m
ef@(EF c -> c -> c
union c -> c -> m Bool
_ c -> term -> term -> c
_) ExploreStrategy
strategy) = HashMap term (c, HashSet term)
-> ExploreFuncs term c m
-> ExploreStrategy
-> ExploredTerms term c m
forall term c (m :: * -> *).
HashMap term (c, HashSet term)
-> ExploreFuncs term c m
-> ExploreStrategy
-> ExploredTerms term c m
ET (((c, HashSet term) -> (c, HashSet term) -> (c, HashSet term))
-> term
-> (c, HashSet term)
-> HashMap term (c, HashSet term)
-> HashMap term (c, HashSet term)
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
M.insertWith (c, HashSet term) -> (c, HashSet term) -> (c, HashSet term)
forall {a}.
Eq a =>
(c, HashSet a) -> (c, HashSet a) -> (c, HashSet a)
go term
t (c
oc, HashSet term
s) HashMap term (c, HashSet term)
etMap) ExploreFuncs term c m
ef ExploreStrategy
strategy
  where
    go :: (c, HashSet a) -> (c, HashSet a) -> (c, HashSet a)
go (c
oc1, HashSet a
s1) (c
oc2, HashSet a
s2) = (c
oc1 c -> c -> c
`union` c
oc2, HashSet a -> HashSet a -> HashSet a
forall a. Eq a => HashSet a -> HashSet a -> HashSet a
S.union HashSet a
s1 HashSet a
s2)

lookup :: (Eq term, Hashable term) => term -> ExploredTerms term c m -> Maybe (c, S.HashSet term)
lookup :: forall term c (m :: * -> *).
(Eq term, Hashable term) =>
term -> ExploredTerms term c m -> Maybe (c, HashSet term)
lookup term
t (ET HashMap term (c, HashSet term)
etMap ExploreFuncs term c m
_ ExploreStrategy
_) = term -> HashMap term (c, HashSet term) -> Maybe (c, HashSet term)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup term
t HashMap term (c, HashSet term)
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 :: 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 term
t0 c
oc0 et :: ExploredTerms term c m
et@(ET HashMap term (c, HashSet term)
_ (EF{c -> c -> m Bool
subsumes :: forall term c (m :: * -> *).
ExploreFuncs term c m -> c -> c -> m Bool
subsumes :: c -> c -> m Bool
subsumes,c -> term -> term -> c
exRefine :: forall term c (m :: * -> *).
ExploreFuncs term c m -> c -> term -> term -> c
exRefine :: c -> term -> term -> c
exRefine}) ExploreStrategy
_) = m Bool
result where

  result :: m Bool
result = HashSet (term, c) -> [(term, c)] -> m Bool
go HashSet (term, c)
forall a. HashSet a
S.empty [(term
t0, c
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 :: HashSet (term, c) -> [(term, c)] -> m Bool
go HashSet (term, c)
_ []       = Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

  -- Term `h` has been seen before at constraints `oc`
  go HashSet (term, c)
seen ((term
h, c
oc'):[(term, c)]
rest) | Just (c
oc, HashSet term
trms) <- term -> ExploredTerms term c m -> Maybe (c, HashSet term)
forall term c (m :: * -> *).
(Eq term, Hashable term) =>
term -> ExploredTerms term c m -> Maybe (c, HashSet term)
lookup term
h ExploredTerms term c m
et
                = do
                    Bool
exploringPathWouldNotPermitDifferentSteps <- c
oc c -> c -> m Bool
`subsumes` c
oc'
                    if Bool
exploringPathWouldNotPermitDifferentSteps
                      then HashSet (term, c) -> [(term, c)] -> m Bool
go HashSet (term, c)
seen' [(term, c)]
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' :: HashSet (term, c)
trms' = (term -> (term, c)) -> HashSet term -> HashSet (term, c)
forall b a.
(Hashable b, Eq b) =>
(a -> b) -> HashSet a -> HashSet b
S.map (\term
t -> (term
t, c -> term -> term -> c
exRefine c
oc' term
h term
t)) HashSet term
trms
                          ts :: HashSet (term, c)
ts    = HashSet (term, c) -> HashSet (term, c) -> HashSet (term, c)
forall a. Eq a => HashSet a -> HashSet a -> HashSet a
S.union HashSet (term, c)
trms' ([(term, c)] -> HashSet (term, c)
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList [(term, c)]
rest) HashSet (term, c) -> HashSet (term, c) -> HashSet (term, c)
forall a. (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a
`S.difference` HashSet (term, c)
seen'
                        in
                          HashSet (term, c) -> [(term, c)] -> m Bool
go HashSet (term, c)
seen' (HashSet (term, c) -> [(term, c)]
forall a. HashSet a -> [a]
S.toList HashSet (term, c)
ts)
                  where
                    seen' :: HashSet (term, c)
seen' = (term, c) -> HashSet (term, c) -> HashSet (term, c)
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
S.insert (term
h, c
oc') HashSet (term, c)
seen

  -- There exists a reachable term that has never previously been seen; not fully explored
  go HashSet (term, c)
_ [(term, c)]
_         = Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
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 :: 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 term
t c
oc et :: ExploredTerms term c m
et@(ET HashMap term (c, HashSet term)
_ EF{c -> c -> m Bool
subsumes :: forall term c (m :: * -> *).
ExploreFuncs term c m -> c -> c -> m Bool
subsumes :: c -> c -> m Bool
subsumes} ExploreStrategy
strategy) =
  case ExploreStrategy
strategy of
    ExploreStrategy
ExploreWhenNeeded      -> Bool -> Bool
not (Bool -> Bool) -> m Bool -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> term -> c -> ExploredTerms term c m -> m Bool
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 term
t c
oc ExploredTerms term c m
et
    ExploreStrategy
ExploreOnce            -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ term -> ExploredTerms term c m -> Bool
forall term c (m :: * -> *).
(Eq term, Hashable term) =>
term -> ExploredTerms term c m -> Bool
visited term
t ExploredTerms term c m
et
    ExploreStrategy
ExploreAlways          -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
    ExploreStrategy
ExploreLessConstrained ->
      case term -> ExploredTerms term c m -> Maybe (c, HashSet term)
forall term c (m :: * -> *).
(Eq term, Hashable term) =>
term -> ExploredTerms term c m -> Maybe (c, HashSet term)
lookup term
t ExploredTerms term c m
et of
        Just (c
oc', HashSet term
_) -> do
          Bool
s <- c
oc' c -> c -> m Bool
`subsumes` c
oc
          Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not Bool
s
        Maybe (c, HashSet term)
Nothing       -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True