-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
module Retrie.PatternMap.Class where

import Control.Monad
import Data.Kind
import Data.Maybe

import Retrie.AlphaEnv
import Retrie.ExactPrint
import Retrie.GHC
import Retrie.Quantifiers
import Retrie.Substitution

------------------------------------------------------------------------

data MatchEnv = ME
  { MatchEnv -> AlphaEnv
meAlphaEnv :: AlphaEnv
  , MatchEnv -> forall a. a -> Annotated a
mePruneA :: forall a. a -> Annotated a
  }

extendMatchEnv :: MatchEnv -> [RdrName] -> MatchEnv
extendMatchEnv :: MatchEnv -> [RdrName] -> MatchEnv
extendMatchEnv MatchEnv
me [RdrName]
bs =
  MatchEnv
me { meAlphaEnv = foldr extendAlphaEnvInternal (meAlphaEnv me) bs }

pruneMatchEnv :: Int -> MatchEnv -> MatchEnv
pruneMatchEnv :: Int -> MatchEnv -> MatchEnv
pruneMatchEnv Int
i MatchEnv
me = MatchEnv
me { meAlphaEnv = pruneAlphaEnv i (meAlphaEnv me) }

------------------------------------------------------------------------

-- TODO: Maybe a -> a ??? -- we never need to delete
type A a = Maybe a -> Maybe a

toA :: PatternMap m => (m a -> m a) -> A (m a)
toA :: forall (m :: * -> *) a. PatternMap m => (m a -> m a) -> A (m a)
toA m a -> m a
f = m a -> Maybe (m a)
forall a. a -> Maybe a
Just (m a -> Maybe (m a))
-> (Maybe (m a) -> m a) -> Maybe (m a) -> Maybe (m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> m a
f (m a -> m a) -> (Maybe (m a) -> m a) -> Maybe (m a) -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> Maybe (m a) -> m a
forall a. a -> Maybe a -> a
fromMaybe m a
forall a. m a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty

toAList :: A a -> A [a]
toAList :: forall a. A a -> A [a]
toAList A a
f Maybe [a]
Nothing = (a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[]) (a -> [a]) -> Maybe a -> Maybe [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> A a
f Maybe a
forall a. Maybe a
Nothing
toAList A a
f (Just [a]
xs) = [a] -> Maybe [a]
forall a. a -> Maybe a
Just ([a] -> Maybe [a]) -> [a] -> Maybe [a]
forall a b. (a -> b) -> a -> b
$ (a -> Maybe a) -> [a] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (A a
f A a -> (a -> Maybe a) -> a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just) [a]
xs

unionOn :: PatternMap m => (a -> m b) -> a -> a -> m b
unionOn :: forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn a -> m b
f a
m1 a
m2 = m b -> m b -> m b
forall a. m a -> m a -> m a
forall (m :: * -> *) a. PatternMap m => m a -> m a -> m a
mUnion (a -> m b
f a
m1) (a -> m b
f a
m2)

------------------------------------------------------------------------

class PatternMap m where
  type Key m :: Type

  mEmpty :: m a
  mUnion :: m a -> m a -> m a

  mAlter :: AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
  mMatch :: MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]

-- Useful to get the chain started in mMatch
mapFor :: (b -> c) -> (a, b) -> [(a, c)]
mapFor :: forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor b -> c
f (a
hs,b
m) = [(a
hs, b -> c
f b
m)]

-- Useful for using existing lookup functions in mMatch
maybeMap :: (b -> Maybe c) -> (a, b) -> [(a, c)]
maybeMap :: forall b c a. (b -> Maybe c) -> (a, b) -> [(a, c)]
maybeMap b -> Maybe c
f (a
hs,b
m) = Maybe (a, c) -> [(a, c)]
forall a. Maybe a -> [a]
maybeToList (Maybe (a, c) -> [(a, c)]) -> Maybe (a, c) -> [(a, c)]
forall a b. (a -> b) -> a -> b
$ (a
hs,) (c -> (a, c)) -> Maybe c -> Maybe (a, c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> b -> Maybe c
f b
m

maybeListMap :: (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap :: forall b c a. (b -> Maybe [c]) -> (a, b) -> [(a, c)]
maybeListMap b -> Maybe [c]
f (a
hs, b
m) = [ (a
a, c
c) | (a
a, [c]
cs) <- (b -> Maybe [c]) -> (a, b) -> [(a, [c])]
forall b c a. (b -> Maybe c) -> (a, b) -> [(a, c)]
maybeMap b -> Maybe [c]
f (a
hs, b
m), c
c <- [c]
cs ]

------------------------------------------------------------------------

newtype MaybeMap a = MaybeMap [a]
  deriving ((forall a b. (a -> b) -> MaybeMap a -> MaybeMap b)
-> (forall a b. a -> MaybeMap b -> MaybeMap a) -> Functor MaybeMap
forall a b. a -> MaybeMap b -> MaybeMap a
forall a b. (a -> b) -> MaybeMap a -> MaybeMap b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> MaybeMap a -> MaybeMap b
fmap :: forall a b. (a -> b) -> MaybeMap a -> MaybeMap b
$c<$ :: forall a b. a -> MaybeMap b -> MaybeMap a
<$ :: forall a b. a -> MaybeMap b -> MaybeMap a
Functor)

instance PatternMap MaybeMap where
  type Key MaybeMap = ()

  mEmpty :: MaybeMap a
  mEmpty :: forall a. MaybeMap a
mEmpty = [a] -> MaybeMap a
forall a. [a] -> MaybeMap a
MaybeMap []

  mUnion :: MaybeMap a -> MaybeMap a -> MaybeMap a
  mUnion :: forall a. MaybeMap a -> MaybeMap a -> MaybeMap a
mUnion (MaybeMap [a]
m1) (MaybeMap [a]
m2) = [a] -> MaybeMap a
forall a. [a] -> MaybeMap a
MaybeMap ([a] -> MaybeMap a) -> [a] -> MaybeMap a
forall a b. (a -> b) -> a -> b
$ [a]
m1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
m2

  mAlter :: AlphaEnv -> Quantifiers -> Key MaybeMap -> A a -> MaybeMap a -> MaybeMap a
  mAlter :: forall a.
AlphaEnv
-> Quantifiers -> Key MaybeMap -> A a -> MaybeMap a -> MaybeMap a
mAlter AlphaEnv
_ Quantifiers
_ () A a
f (MaybeMap []) = [a] -> MaybeMap a
forall a. [a] -> MaybeMap a
MaybeMap ([a] -> MaybeMap a) -> [a] -> MaybeMap a
forall a b. (a -> b) -> a -> b
$ Maybe a -> [a]
forall a. Maybe a -> [a]
maybeToList (Maybe a -> [a]) -> Maybe a -> [a]
forall a b. (a -> b) -> a -> b
$ A a
f Maybe a
forall a. Maybe a
Nothing
  mAlter AlphaEnv
_ Quantifiers
_ () A a
f (MaybeMap [a]
xs) = [a] -> MaybeMap a
forall a. [a] -> MaybeMap a
MaybeMap ([a] -> MaybeMap a) -> [a] -> MaybeMap a
forall a b. (a -> b) -> a -> b
$ (a -> Maybe a) -> [a] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (A a
f A a -> (a -> Maybe a) -> a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just) [a]
xs

  mMatch
    :: MatchEnv
    -> Key MaybeMap
    -> (Substitution, MaybeMap a)
    -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
mMatch MatchEnv
_ () (Substitution
hs, MaybeMap [a]
xs) = (a -> (Substitution, a)) -> [a] -> [(Substitution, a)]
forall a b. (a -> b) -> [a] -> [b]
map (Substitution
hs,) [a]
xs

------------------------------------------------------------------------

data ListMap m a = ListMap
  { forall (m :: * -> *) a. ListMap m a -> MaybeMap a
lmNil  :: MaybeMap a
  , forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lmCons :: m (ListMap m a)
  }
  deriving ((forall a b. (a -> b) -> ListMap m a -> ListMap m b)
-> (forall a b. a -> ListMap m b -> ListMap m a)
-> Functor (ListMap m)
forall a b. a -> ListMap m b -> ListMap m a
forall a b. (a -> b) -> ListMap m a -> ListMap m b
forall (m :: * -> *) a b.
Functor m =>
a -> ListMap m b -> ListMap m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ListMap m a -> ListMap m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ListMap m a -> ListMap m b
fmap :: forall a b. (a -> b) -> ListMap m a -> ListMap m b
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> ListMap m b -> ListMap m a
<$ :: forall a b. a -> ListMap m b -> ListMap m a
Functor)

instance PatternMap m => PatternMap (ListMap m) where
  type Key (ListMap m) = [Key m]

  mEmpty :: ListMap m a
  mEmpty :: forall a. ListMap m a
mEmpty = MaybeMap a -> m (ListMap m a) -> ListMap m a
forall (m :: * -> *) a.
MaybeMap a -> m (ListMap m a) -> ListMap m a
ListMap MaybeMap a
forall a. MaybeMap a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty m (ListMap m a)
forall a. m a
forall (m :: * -> *) a. PatternMap m => m a
mEmpty

  mUnion :: ListMap m a -> ListMap m a -> ListMap m a
  mUnion :: forall a. ListMap m a -> ListMap m a -> ListMap m a
mUnion ListMap m a
m1 ListMap m a
m2 = ListMap
    { lmNil :: MaybeMap a
lmNil = (ListMap m a -> MaybeMap a)
-> ListMap m a -> ListMap m a -> MaybeMap a
forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn ListMap m a -> MaybeMap a
forall (m :: * -> *) a. ListMap m a -> MaybeMap a
lmNil ListMap m a
m1 ListMap m a
m2
    , lmCons :: m (ListMap m a)
lmCons = (ListMap m a -> m (ListMap m a))
-> ListMap m a -> ListMap m a -> m (ListMap m a)
forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn ListMap m a -> m (ListMap m a)
forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lmCons ListMap m a
m1 ListMap m a
m2
    }

  mAlter :: AlphaEnv -> Quantifiers -> Key (ListMap m) -> A a -> ListMap m a -> ListMap m a
  mAlter :: forall a.
AlphaEnv
-> Quantifiers
-> Key (ListMap m)
-> A a
-> ListMap m a
-> ListMap m a
mAlter AlphaEnv
env Quantifiers
vs []     A a
f ListMap m a
m = ListMap m a
m { lmNil  = mAlter env vs () f (lmNil m) }
  mAlter AlphaEnv
env Quantifiers
vs (Key m
x:[Key m]
xs) A a
f ListMap m a
m = ListMap m a
m { lmCons = mAlter env vs x (toA (mAlter env vs xs f)) (lmCons m) }

  mMatch :: MatchEnv -> Key (ListMap m) -> (Substitution, ListMap m a) -> [(Substitution, a)]
  mMatch :: forall a.
MatchEnv
-> Key (ListMap m)
-> (Substitution, ListMap m a)
-> [(Substitution, a)]
mMatch MatchEnv
env []     = (ListMap m a -> MaybeMap a)
-> (Substitution, ListMap m a) -> [(Substitution, MaybeMap a)]
forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor ListMap m a -> MaybeMap a
forall (m :: * -> *) a. ListMap m a -> MaybeMap a
lmNil ((Substitution, ListMap m a) -> [(Substitution, MaybeMap a)])
-> ((Substitution, MaybeMap a) -> [(Substitution, a)])
-> (Substitution, ListMap m a)
-> [(Substitution, a)]
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
forall a.
MatchEnv
-> Key MaybeMap
-> (Substitution, MaybeMap a)
-> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env ()
  mMatch MatchEnv
env (Key m
x:[Key m]
xs) = (ListMap m a -> m (ListMap m a))
-> (Substitution, ListMap m a) -> [(Substitution, m (ListMap m a))]
forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor ListMap m a -> m (ListMap m a)
forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lmCons ((Substitution, ListMap m a) -> [(Substitution, m (ListMap m a))])
-> ((Substitution, m (ListMap m a)) -> [(Substitution, a)])
-> (Substitution, ListMap m a)
-> [(Substitution, a)]
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> MatchEnv
-> Key m
-> (Substitution, m (ListMap m a))
-> [(Substitution, ListMap m a)]
forall a.
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env Key m
x ((Substitution, m (ListMap m a)) -> [(Substitution, ListMap m a)])
-> ((Substitution, ListMap m a) -> [(Substitution, a)])
-> (Substitution, m (ListMap m a))
-> [(Substitution, a)]
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> MatchEnv
-> Key (ListMap m)
-> (Substitution, ListMap m a)
-> [(Substitution, a)]
forall a.
MatchEnv
-> Key (ListMap m)
-> (Substitution, ListMap m a)
-> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env [Key m]
Key (ListMap m)
xs

------------------------------------------------------------------------

findMatch :: PatternMap m => MatchEnv -> Key m -> m a -> [(Substitution, a)]
findMatch :: forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> m a -> [(Substitution, a)]
findMatch MatchEnv
env Key m
k m a
m = MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
forall a.
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env Key m
k (Substitution
emptySubst, m a
m)

insertMatch :: PatternMap m => AlphaEnv -> Quantifiers -> Key m -> a -> m a -> m a
insertMatch :: forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> a -> m a -> m a
insertMatch AlphaEnv
env Quantifiers
vs Key m
k a
x = AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
forall a. AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
vs Key m
k (Maybe a -> A a
forall a b. a -> b -> a
const (a -> Maybe a
forall a. a -> Maybe a
Just a
x))