-- 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 :: AlphaEnv
meAlphaEnv = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr RdrName -> AlphaEnv -> AlphaEnv
extendAlphaEnvInternal (MatchEnv -> AlphaEnv
meAlphaEnv MatchEnv
me) [RdrName]
bs }

pruneMatchEnv :: Int -> MatchEnv -> MatchEnv
pruneMatchEnv :: Int -> MatchEnv -> MatchEnv
pruneMatchEnv Int
i MatchEnv
me = MatchEnv
me { meAlphaEnv :: AlphaEnv
meAlphaEnv = Int -> AlphaEnv -> AlphaEnv
pruneAlphaEnv Int
i (MatchEnv -> AlphaEnv
meAlphaEnv MatchEnv
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 = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> m a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a -> a
fromMaybe 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 = (forall a. a -> [a] -> [a]
:[]) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> A a
f forall a. Maybe a
Nothing
toAList A a
f (Just [a]
xs) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (A a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = 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) = forall a. Maybe a -> [a]
maybeToList forall a b. (a -> b) -> a -> b
$ (a
hs,) 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) <- 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 -> 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
<$ :: forall a b. a -> MaybeMap b -> MaybeMap a
$c<$ :: forall a b. a -> MaybeMap b -> MaybeMap a
fmap :: forall a b. (a -> b) -> MaybeMap a -> MaybeMap b
$cfmap :: forall a b. (a -> b) -> MaybeMap a -> MaybeMap b
Functor)

instance PatternMap MaybeMap where
  type Key MaybeMap = ()

  mEmpty :: MaybeMap a
  mEmpty :: forall a. MaybeMap a
mEmpty = 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) = forall a. [a] -> MaybeMap a
MaybeMap forall a b. (a -> b) -> a -> b
$ [a]
m1 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 []) = forall a. [a] -> MaybeMap a
MaybeMap forall a b. (a -> b) -> a -> b
$ forall a. Maybe a -> [a]
maybeToList forall a b. (a -> b) -> a -> b
$ A a
f forall a. Maybe a
Nothing
  mAlter AlphaEnv
_ Quantifiers
_ () A a
f (MaybeMap [a]
xs) = forall a. [a] -> MaybeMap a
MaybeMap forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (A a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = 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 -> 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
<$ :: forall a b. a -> ListMap m b -> ListMap m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> ListMap m b -> ListMap m a
fmap :: forall a b. (a -> b) -> ListMap m a -> ListMap m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ListMap m a -> ListMap m b
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 = forall (m :: * -> *) a.
MaybeMap a -> m (ListMap m a) -> ListMap m a
ListMap forall (m :: * -> *) a. PatternMap m => m a
mEmpty 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 = forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn forall (m :: * -> *) a. ListMap m a -> MaybeMap a
lmNil ListMap m a
m1 ListMap m a
m2
    , lmCons :: m (ListMap m a)
lmCons = forall (m :: * -> *) a b.
PatternMap m =>
(a -> m b) -> a -> a -> m b
unionOn 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 :: MaybeMap a
lmNil  = forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
vs () A a
f (forall (m :: * -> *) a. ListMap m a -> MaybeMap a
lmNil ListMap m a
m) }
  mAlter AlphaEnv
env Quantifiers
vs (Key m
x:[Key m]
xs) A a
f ListMap m a
m = ListMap m a
m { lmCons :: m (ListMap m a)
lmCons = forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
vs Key m
x (forall (m :: * -> *) a. PatternMap m => (m a -> m a) -> A (m a)
toA (forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
vs [Key m]
xs A a
f)) (forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lmCons ListMap m a
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 []     = forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor forall (m :: * -> *) a. ListMap m a -> MaybeMap a
lmNil forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> 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) = forall b c a. (b -> c) -> (a, b) -> [(a, c)]
mapFor forall (m :: * -> *) a. ListMap m a -> m (ListMap m a)
lmCons forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env Key m
x forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (m :: * -> *) a.
PatternMap m =>
MatchEnv -> Key m -> (Substitution, m a) -> [(Substitution, a)]
mMatch MatchEnv
env [Key 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 = 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 = forall (m :: * -> *) a.
PatternMap m =>
AlphaEnv -> Quantifiers -> Key m -> A a -> m a -> m a
mAlter AlphaEnv
env Quantifiers
vs Key m
k (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just a
x))