{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module AlgebraCheckers.Suggestions where

import AlgebraCheckers.Patterns
import AlgebraCheckers.Ppr
import AlgebraCheckers.Unification
import Control.Monad
import Data.Char
import Data.Data
import Data.Generics.Schemes (listify)
import Data.Group
import Data.List
import Data.Maybe
import Data.Semigroup
import Data.Traversable
import Language.Haskell.TH hiding (ppr)
import Language.Haskell.TH.Syntax
import Prelude hiding (exp)
import THInstanceReification


data Suggestion
  = HomoSuggestion Name Name Int Type Type Exp
  deriving (Suggestion -> Suggestion -> Bool
(Suggestion -> Suggestion -> Bool)
-> (Suggestion -> Suggestion -> Bool) -> Eq Suggestion
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Suggestion -> Suggestion -> Bool
$c/= :: Suggestion -> Suggestion -> Bool
== :: Suggestion -> Suggestion -> Bool
$c== :: Suggestion -> Suggestion -> Bool
Eq, Eq Suggestion
Eq Suggestion
-> (Suggestion -> Suggestion -> Ordering)
-> (Suggestion -> Suggestion -> Bool)
-> (Suggestion -> Suggestion -> Bool)
-> (Suggestion -> Suggestion -> Bool)
-> (Suggestion -> Suggestion -> Bool)
-> (Suggestion -> Suggestion -> Suggestion)
-> (Suggestion -> Suggestion -> Suggestion)
-> Ord Suggestion
Suggestion -> Suggestion -> Bool
Suggestion -> Suggestion -> Ordering
Suggestion -> Suggestion -> Suggestion
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Suggestion -> Suggestion -> Suggestion
$cmin :: Suggestion -> Suggestion -> Suggestion
max :: Suggestion -> Suggestion -> Suggestion
$cmax :: Suggestion -> Suggestion -> Suggestion
>= :: Suggestion -> Suggestion -> Bool
$c>= :: Suggestion -> Suggestion -> Bool
> :: Suggestion -> Suggestion -> Bool
$c> :: Suggestion -> Suggestion -> Bool
<= :: Suggestion -> Suggestion -> Bool
$c<= :: Suggestion -> Suggestion -> Bool
< :: Suggestion -> Suggestion -> Bool
$c< :: Suggestion -> Suggestion -> Bool
compare :: Suggestion -> Suggestion -> Ordering
$ccompare :: Suggestion -> Suggestion -> Ordering
$cp1Ord :: Eq Suggestion
Ord, Int -> Suggestion -> ShowS
[Suggestion] -> ShowS
Suggestion -> String
(Int -> Suggestion -> ShowS)
-> (Suggestion -> String)
-> ([Suggestion] -> ShowS)
-> Show Suggestion
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Suggestion] -> ShowS
$cshowList :: [Suggestion] -> ShowS
show :: Suggestion -> String
$cshow :: Suggestion -> String
showsPrec :: Int -> Suggestion -> ShowS
$cshowsPrec :: Int -> Suggestion -> ShowS
Show)

homoSuggestionEq :: Suggestion -> Suggestion -> Bool
homoSuggestionEq :: Suggestion -> Suggestion -> Bool
homoSuggestionEq (HomoSuggestion Name
_ Name
fn1 Int
ix1 Type
_ Type
_ Exp
_)
                 (HomoSuggestion Name
_ Name
fn2 Int
ix2 Type
_ Type
_ Exp
_) = Name
fn1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
fn2
                                                 Bool -> Bool -> Bool
&& Int
ix1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ix2


pprSuggestion :: Suggestion -> Doc
pprSuggestion :: Suggestion -> Doc
pprSuggestion (HomoSuggestion Name
nm Name
_ Int
_ Type
arg_ty Type
res_ty (LamE [VarP Name
var] Exp
exp)) =
  Exp -> Doc
forall a. Ppr a => a -> Doc
ppr (Exp -> Doc) -> Exp -> Doc
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
forall a. Data a => a -> a
deModuleName (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$
    Name -> Exp
VarE 'law Exp -> Type -> Exp
`AppTypeE` Name -> Type
ConT Name
nm Exp -> Exp -> Exp
`AppE` ([Pat] -> Exp -> Exp
LamE [Pat -> Type -> Pat
SigP (Name -> Pat
VarP Name
var) Type
arg_ty] (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Type -> Exp
SigE Exp
exp Type
res_ty)
pprSuggestion (HomoSuggestion Name
nm Name
_ Int
_ Type
_ Type
_ Exp
exp) =
  Exp -> Doc
forall a. Ppr a => a -> Doc
ppr (Exp -> Doc) -> Exp -> Doc
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
forall a. Data a => a -> a
deModuleName (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$
    Name -> Exp
VarE 'law Exp -> Type -> Exp
`AppTypeE` Name -> Type
ConT Name
nm Exp -> Exp -> Exp
`AppE` Exp
exp


knownSuggestionHierarchies :: [[Name]]
knownSuggestionHierarchies :: [[Name]]
knownSuggestionHierarchies =
  [ [ ''Group, ''Monoid, ''Semigroup ]
  ]

suggest :: Data a => Module -> a -> Q [Suggestion]
suggest :: Module -> a -> Q [Suggestion]
suggest Module
md a
a = do
  let surface :: [Name]
surface = Module -> a -> [Name]
forall a. Data a => Module -> a -> [Name]
getSurface Module
md a
a
  ([[[Suggestion]]] -> [Suggestion])
-> Q [[[Suggestion]]] -> Q [Suggestion]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([[Suggestion]] -> [Suggestion]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([[Suggestion]] -> [Suggestion])
-> ([[[Suggestion]]] -> [[Suggestion]])
-> [[[Suggestion]]]
-> [Suggestion]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[[Suggestion]]] -> [[Suggestion]]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join) (Q [[[Suggestion]]] -> Q [Suggestion])
-> Q [[[Suggestion]]] -> Q [Suggestion]
forall a b. (a -> b) -> a -> b
$
    [Name] -> (Name -> Q [[Suggestion]]) -> Q [[[Suggestion]]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Name]
surface ((Name -> Q [[Suggestion]]) -> Q [[[Suggestion]]])
-> (Name -> Q [[Suggestion]]) -> Q [[[Suggestion]]]
forall a b. (a -> b) -> a -> b
$ \Name
nm ->
      [[Name]] -> ([Name] -> Q [Suggestion]) -> Q [[Suggestion]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [[Name]]
knownSuggestionHierarchies (([Name] -> Q [Suggestion]) -> Q [[Suggestion]])
-> ([Name] -> Q [Suggestion]) -> Q [[Suggestion]]
forall a b. (a -> b) -> a -> b
$ \[Name]
hierarchy -> do
        [Suggestion]
zs <- ([[Suggestion]] -> [Suggestion])
-> Q [[Suggestion]] -> Q [Suggestion]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Suggestion]] -> [Suggestion]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Q [[Suggestion]] -> Q [Suggestion])
-> Q [[Suggestion]] -> Q [Suggestion]
forall a b. (a -> b) -> a -> b
$ [Name] -> (Name -> Q [Suggestion]) -> Q [[Suggestion]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Name]
hierarchy ((Name -> Q [Suggestion]) -> Q [[Suggestion]])
-> (Name -> Q [Suggestion]) -> Q [[Suggestion]]
forall a b. (a -> b) -> a -> b
$ \Name
tc_name -> do
          VarI Name
_ Type
ty Maybe Dec
_ <- Name -> Q Info
reify Name
nm
          Name -> Name -> Type -> Q [Suggestion]
possibleHomos Name
tc_name Name
nm Type
ty
        [Suggestion] -> Q [Suggestion]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Suggestion] -> Q [Suggestion]) -> [Suggestion] -> Q [Suggestion]
forall a b. (a -> b) -> a -> b
$ (Suggestion -> Suggestion -> Bool) -> [Suggestion] -> [Suggestion]
forall a. (a -> a -> Bool) -> [a] -> [a]
nubBy Suggestion -> Suggestion -> Bool
homoSuggestionEq [Suggestion]
zs


suggest' :: Data a => a -> Q [Suggestion]
suggest' :: a -> Q [Suggestion]
suggest' a
a = do
  Module
md <- Q Module
thisModule
  Module -> a -> Q [Suggestion]
forall a. Data a => Module -> a -> Q [Suggestion]
suggest Module
md a
a



possibleHomos :: Name -> Name -> Type -> Q [Suggestion]
possibleHomos :: Name -> Name -> Type -> Q [Suggestion]
possibleHomos Name
tc_name Name
fn Type
ty = do
  let ([Type]
args, Type
res) = Type -> ([Type], Type)
unrollTyArr Type
ty
  Name -> Type -> Q Bool
hasInstance Name
tc_name Type
res Q Bool -> (Bool -> Q [Suggestion]) -> Q [Suggestion]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Bool
False -> [Suggestion] -> Q [Suggestion]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    Bool
True  -> do
      [Name]
names <- [Type] -> (Type -> Q Name) -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Type]
args ((Type -> Q Name) -> Q [Name]) -> (Type -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName (String -> Q Name) -> (Type -> String) -> Type -> Q Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> String
goodTyName
      ([Maybe Suggestion] -> [Suggestion])
-> Q [Maybe Suggestion] -> Q [Suggestion]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Maybe Suggestion] -> [Suggestion]
forall a. [Maybe a] -> [a]
catMaybes (Q [Maybe Suggestion] -> Q [Suggestion])
-> Q [Maybe Suggestion] -> Q [Suggestion]
forall a b. (a -> b) -> a -> b
$ [(Name, Type, Int)]
-> ((Name, Type, Int) -> Q (Maybe Suggestion))
-> Q [Maybe Suggestion]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for ([Name] -> [Type] -> [Int] -> [(Name, Type, Int)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Name]
names [Type]
args [Int
0..]) (((Name, Type, Int) -> Q (Maybe Suggestion))
 -> Q [Maybe Suggestion])
-> ((Name, Type, Int) -> Q (Maybe Suggestion))
-> Q [Maybe Suggestion]
forall a b. (a -> b) -> a -> b
$ \(Name
name, Type
arg, Int
ix) ->
        Name -> Type -> Q Bool
hasInstance Name
tc_name Type
arg Q Bool -> (Bool -> Q (Maybe Suggestion)) -> Q (Maybe Suggestion)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Bool
False -> Maybe Suggestion -> Q (Maybe Suggestion)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Suggestion
forall a. Maybe a
Nothing
          Bool
True  -> do
            Exp
exp <- [PatQ] -> ExpQ -> ExpQ
lamE [Name -> PatQ
varP Name
name] (ExpQ -> ExpQ) -> ExpQ -> ExpQ
forall a b. (a -> b) -> a -> b
$ [ExpQ] -> ExpQ
appsE ([ExpQ] -> ExpQ) -> [ExpQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$ Name -> ExpQ
varE Name
fn ExpQ -> [ExpQ] -> [ExpQ]
forall a. a -> [a] -> [a]
: (Name -> ExpQ) -> [Name] -> [ExpQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> ExpQ
varE [Name]
names
            Maybe Suggestion -> Q (Maybe Suggestion)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Suggestion -> Q (Maybe Suggestion))
-> Maybe Suggestion -> Q (Maybe Suggestion)
forall a b. (a -> b) -> a -> b
$ Suggestion -> Maybe Suggestion
forall a. a -> Maybe a
Just (Suggestion -> Maybe Suggestion) -> Suggestion -> Maybe Suggestion
forall a b. (a -> b) -> a -> b
$ Name -> Name -> Int -> Type -> Type -> Exp -> Suggestion
HomoSuggestion Name
tc_name Name
fn Int
ix Type
arg Type
res Exp
exp


goodTyName :: Type -> String
goodTyName :: Type -> String
goodTyName = (Char -> Char) -> ShowS
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Char -> Char
toLower ShowS -> (Type -> String) -> Type -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShowS
forall a. Int -> [a] -> [a]
take Int
1 ShowS -> (Type -> String) -> Type -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Bool -> Bool
not (Bool -> Bool) -> (Char -> Bool) -> Char -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Bool
isAlpha) ShowS -> (Type -> String) -> Type -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> String
render (Doc -> String) -> (Type -> Doc) -> Type -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Doc
forall a. Ppr a => a -> Doc
ppr (Type -> Doc) -> (Type -> Type) -> Type -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
forall a. Data a => a -> a
deModuleName

getSurface :: Data a => Module -> a -> [Name]
getSurface :: Module -> a -> [Name]
getSurface Module
m = (Name -> Bool) -> GenericQ [Name]
forall r. Typeable r => (r -> Bool) -> GenericQ [r]
listify (Module -> Name -> Bool
sameModule Module
m)


sameModule :: Module -> Name -> Bool
sameModule :: Module -> Name -> Bool
sameModule (Module (PkgName String
pkg) (ModName String
md)) Name
n =
  Name -> Maybe String
nameModule Name
n Maybe String -> Maybe String -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Maybe String
forall a. a -> Maybe a
Just String
md Bool -> Bool -> Bool
&& Name -> Maybe String
namePackage Name
n Maybe String -> Maybe String -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Maybe String
forall a. a -> Maybe a
Just String
pkg


unrollTyArr :: Type -> ([Type], Type)
unrollTyArr :: Type -> ([Type], Type)
unrollTyArr Type
ty =
  let tys :: [Type]
tys = Type -> [Type]
unloopTyArrs Type
ty
   in ([Type] -> [Type]
forall a. [a] -> [a]
init [Type]
tys, [Type] -> Type
forall a. [a] -> a
last [Type]
tys)
  where
    unloopTyArrs :: Type -> [Type]
    unloopTyArrs :: Type -> [Type]
unloopTyArrs (Type
ArrowT `AppT` Type
a `AppT` Type
b) =  Type
a Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: Type -> [Type]
unloopTyArrs Type
b
    unloopTyArrs Type
t =  [Type
t]

hasInstance :: Name -> Type -> Q Bool
hasInstance :: Name -> Type -> Q Bool
hasInstance Name
tc_name = Name -> [Type] -> Q Bool
isProperInstance Name
tc_name ([Type] -> Q Bool) -> (Type -> [Type]) -> Type -> Q Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure