{-# LANGUAGE CPP, PatternGuards, BangPatterns #-}
module Djinn.GHC (Environment, MaxSolutions(..), djinn) where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad (forM)
import Data.Maybe (isJust)
import Data.Set (Set, insert, union, unions, empty, toList)

import qualified Djinn.HTypes as D
import qualified Djinn.LJT as D

import MonadUtils
import qualified DataCon as G
import qualified GHC as G
import qualified Name as G
import qualified TyCon as G
import qualified Type as G

data NoExtraInfo = NoExtraInfo
type HEnvironment1 a = [(D.HSymbol, ([D.HSymbol], D.HType, a))]
type HEnvironment = HEnvironment1 NoExtraInfo

getConTs :: G.Type -> Set G.Name
getConTs :: Type -> Set Name
getConTs Type
t | Just (TyCoVar
_, Type
i)  <- Type -> Maybe (TyCoVar, Type)
G.splitForAllTy_maybe Type
t = Type -> Set Name
getConTs Type
i
getConTs Type
t | Just (Type
t1,Type
t2) <- Type -> Maybe (Type, Type)
G.splitFunTy_maybe Type
t    = Type -> Set Name
getConTs Type
t1 Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
`union` Type -> Set Name
getConTs Type
t2
getConTs Type
t | Just (TyCon
c, [Type]
ts) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
G.splitTyConApp_maybe Type
t = 
  let args :: Set Name
args = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
unions ([Set Name] -> Set Name) -> [Set Name] -> Set Name
forall a b. (a -> b) -> a -> b
$ (Type -> Set Name) -> [Type] -> [Set Name]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Set Name
getConTs [Type]
ts
   in if TyCon -> Bool
G.isTupleTyCon TyCon
c then Set Name
args else Name -> Set Name -> Set Name
forall a. Ord a => a -> Set a -> Set a
insert (TyCon -> Name
forall a. NamedThing a => a -> Name
G.getName TyCon
c) Set Name
args
getConTs Type
t | Just (Type
t1,Type
t2) <- Type -> Maybe (Type, Type)
G.splitAppTy_maybe Type
t    = Type -> Set Name
getConTs Type
t1 Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
`union` Type -> Set Name
getConTs Type
t2
getConTs Type
t | Just TyCoVar
_       <- Type -> Maybe TyCoVar
G.getTyVar_maybe Type
t      = Set Name
forall a. Set a
empty
getConTs Type
_                                           = Set Name
forall a. Set a
empty

hType :: G.Type -> D.HType
hType :: Type -> HType
hType Type
t | Just (TyCoVar
_, Type
i)  <- Type -> Maybe (TyCoVar, Type)
G.splitForAllTy_maybe Type
t = Type -> HType
hType Type
i
hType Type
t | Just (Type
t1,Type
t2) <- Type -> Maybe (Type, Type)
G.splitFunTy_maybe Type
t    = HType -> HType -> HType
D.HTArrow (Type -> HType
hType Type
t1) (Type -> HType
hType Type
t2)
hType Type
t | Just (TyCon
c, [Type]
ts) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
G.splitTyConApp_maybe Type
t =
  let args :: [HType]
args = (Type -> HType) -> [Type] -> [HType]
forall a b. (a -> b) -> [a] -> [b]
map Type -> HType
hType [Type]
ts
   in if TyCon -> Bool
G.isTupleTyCon TyCon
c  -- Check if we have a tuple
         then [HType] -> HType
D.HTTuple [HType]
args
         else HSymbol -> [HType] -> HType
createHTApp (TyCon -> HSymbol
forall a. NamedThing a => a -> HSymbol
G.getOccString TyCon
c) ([HType] -> [HType]
forall a. [a] -> [a]
reverse [HType]
args)
  where createHTApp :: HSymbol -> [HType] -> HType
createHTApp HSymbol
n []     = HSymbol -> HType
D.HTCon HSymbol
n
        createHTApp HSymbol
n (HType
x:[HType]
xs) = HType -> HType -> HType
D.HTApp (HSymbol -> [HType] -> HType
createHTApp HSymbol
n [HType]
xs) HType
x
hType Type
t | Just (Type
t1,Type
t2) <- Type -> Maybe (Type, Type)
G.splitAppTy_maybe Type
t    = HType -> HType -> HType
D.HTApp (Type -> HType
hType Type
t1) (Type -> HType
hType Type
t2)
hType Type
t | Just TyCoVar
var <- Type -> Maybe TyCoVar
G.getTyVar_maybe Type
t          = HSymbol -> HType
D.HTVar (TyCoVar -> HSymbol
forall a. NamedThing a => a -> HSymbol
toHSymbol TyCoVar
var)
hType Type
_                                           = HSymbol -> HType
forall a. HasCallStack => HSymbol -> a
error HSymbol
"Unimplemented"

environment :: G.GhcMonad m => Maybe G.ModuleInfo -> G.Type -> m HEnvironment
environment :: Maybe ModuleInfo -> Type -> m HEnvironment
environment Maybe ModuleInfo
minfo = ([HEnvironment] -> HEnvironment)
-> m [HEnvironment] -> m HEnvironment
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [HEnvironment] -> HEnvironment
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (m [HEnvironment] -> m HEnvironment)
-> (Type -> m [HEnvironment]) -> Type -> m HEnvironment
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name -> m HEnvironment) -> [Name] -> m [HEnvironment]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Maybe ModuleInfo -> Name -> m HEnvironment
forall (m :: * -> *).
GhcMonad m =>
Maybe ModuleInfo -> Name -> m HEnvironment
environment1 Maybe ModuleInfo
minfo) ([Name] -> m [HEnvironment])
-> (Type -> [Name]) -> Type -> m [HEnvironment]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set Name -> [Name]
forall a. Set a -> [a]
toList (Set Name -> [Name]) -> (Type -> Set Name) -> Type -> [Name]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Set Name
getConTs

environment1 :: G.GhcMonad m => Maybe G.ModuleInfo -> G.Name -> m HEnvironment
environment1 :: Maybe ModuleInfo -> Name -> m HEnvironment
environment1 Maybe ModuleInfo
minfo Name
name = do
  Maybe TyThing
thing <- Maybe ModuleInfo -> Name -> m (Maybe TyThing)
forall (m :: * -> *).
GhcMonad m =>
Maybe ModuleInfo -> Name -> m (Maybe TyThing)
lookupNameEverywhere Maybe ModuleInfo
minfo Name
name
  case Maybe TyThing
thing of
    Just (G.ATyCon TyCon
tycon) | TyCon -> Bool
G.isAlgTyCon TyCon
tycon -> do
      let tyconName :: HSymbol
tyconName = Name -> HSymbol
forall a. NamedThing a => a -> HSymbol
toHSymbol (Name -> HSymbol) -> Name -> HSymbol
forall a b. (a -> b) -> a -> b
$ TyCon -> Name
G.tyConName TyCon
tycon
          varsH :: [HSymbol]
varsH = (TyCoVar -> HSymbol) -> [TyCoVar] -> [HSymbol]
forall a b. (a -> b) -> [a] -> [b]
map TyCoVar -> HSymbol
forall a. NamedThing a => a -> HSymbol
toHSymbol ([TyCoVar] -> [HSymbol]) -> [TyCoVar] -> [HSymbol]
forall a b. (a -> b) -> a -> b
$ TyCon -> [TyCoVar]
G.tyConTyVars TyCon
tycon
          Just [DataCon]
datacons = TyCon -> Maybe [DataCon]
G.tyConDataCons_maybe TyCon
tycon
      [((HSymbol, [HType]), [HEnvironment])]
dtypes <- [DataCon]
-> (DataCon -> m ((HSymbol, [HType]), [HEnvironment]))
-> m [((HSymbol, [HType]), [HEnvironment])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [DataCon]
datacons ((DataCon -> m ((HSymbol, [HType]), [HEnvironment]))
 -> m [((HSymbol, [HType]), [HEnvironment])])
-> (DataCon -> m ((HSymbol, [HType]), [HEnvironment]))
-> m [((HSymbol, [HType]), [HEnvironment])]
forall a b. (a -> b) -> a -> b
$ \DataCon
dcon -> do
        let dconN :: HSymbol
dconN = Name -> HSymbol
forall a. NamedThing a => a -> HSymbol
toHSymbol (Name -> HSymbol) -> Name -> HSymbol
forall a b. (a -> b) -> a -> b
$ DataCon -> Name
G.dataConName DataCon
dcon
            ([TyCoVar]
_,[Type]
_,[Type]
dconT,Type
_) = DataCon -> ([TyCoVar], [Type], [Type], Type)
G.dataConSig DataCon
dcon
        [HEnvironment]
dconE <- (Type -> m HEnvironment) -> [Type] -> m [HEnvironment]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Maybe ModuleInfo -> Type -> m HEnvironment
forall (m :: * -> *).
GhcMonad m =>
Maybe ModuleInfo -> Type -> m HEnvironment
environment Maybe ModuleInfo
minfo) [Type]
dconT
        ((HSymbol, [HType]), [HEnvironment])
-> m ((HSymbol, [HType]), [HEnvironment])
forall (m :: * -> *) a. Monad m => a -> m a
return ((HSymbol
dconN, (Type -> HType) -> [Type] -> [HType]
forall a b. (a -> b) -> [a] -> [b]
map Type -> HType
hType [Type]
dconT), [HEnvironment]
dconE)
      let dtypesT :: [(HSymbol, [HType])]
dtypesT = (((HSymbol, [HType]), [HEnvironment]) -> (HSymbol, [HType]))
-> [((HSymbol, [HType]), [HEnvironment])] -> [(HSymbol, [HType])]
forall a b. (a -> b) -> [a] -> [b]
map ((HSymbol, [HType]), [HEnvironment]) -> (HSymbol, [HType])
forall a b. (a, b) -> a
fst [((HSymbol, [HType]), [HEnvironment])]
dtypes
          dtypesE :: [HEnvironment]
dtypesE = (((HSymbol, [HType]), [HEnvironment]) -> [HEnvironment])
-> [((HSymbol, [HType]), [HEnvironment])] -> [HEnvironment]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((HSymbol, [HType]), [HEnvironment]) -> [HEnvironment]
forall a b. (a, b) -> b
snd [((HSymbol, [HType]), [HEnvironment])]
dtypes
      HEnvironment -> m HEnvironment
forall (m :: * -> *) a. Monad m => a -> m a
return (HEnvironment -> m HEnvironment) -> HEnvironment -> m HEnvironment
forall a b. (a -> b) -> a -> b
$ (HSymbol
tyconName, ([HSymbol]
varsH, [(HSymbol, [HType])] -> HType
D.HTUnion [(HSymbol, [HType])]
dtypesT, NoExtraInfo
NoExtraInfo)) (HSymbol, ([HSymbol], HType, NoExtraInfo))
-> HEnvironment -> HEnvironment
forall a. a -> [a] -> [a]
: [HEnvironment] -> HEnvironment
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [HEnvironment]
dtypesE
#if __GLASGOW_HASKELL__ >= 710
    Just (G.ATyCon TyCon
tycon) | TyCon -> Bool
G.isTypeSynonymTyCon TyCon
tycon -> do
#else
    Just (G.ATyCon tycon) | G.isSynTyCon tycon -> do
#endif
      -- Get information for this type synonym
      let tyconName :: HSymbol
tyconName = Name -> HSymbol
forall a. NamedThing a => a -> HSymbol
toHSymbol (Name -> HSymbol) -> Name -> HSymbol
forall a b. (a -> b) -> a -> b
$ TyCon -> Name
G.tyConName TyCon
tycon
#if __GLASGOW_HASKELL__ >= 708
          Just ([TyCoVar]
vars, Type
defn) = TyCon -> Maybe ([TyCoVar], Type)
G.synTyConDefn_maybe TyCon
tycon
#else
          (vars, defn) = G.synTyConDefn tycon
#endif
          varsH :: [HSymbol]
varsH = (TyCoVar -> HSymbol) -> [TyCoVar] -> [HSymbol]
forall a b. (a -> b) -> [a] -> [b]
map TyCoVar -> HSymbol
forall a. NamedThing a => a -> HSymbol
toHSymbol [TyCoVar]
vars
          htype :: HType
htype = Type -> HType
hType Type
defn
      -- Recursively obtain it for the environment of the type
      HEnvironment
defnEnv <- Maybe ModuleInfo -> Type -> m HEnvironment
forall (m :: * -> *).
GhcMonad m =>
Maybe ModuleInfo -> Type -> m HEnvironment
environment Maybe ModuleInfo
minfo Type
defn
      HEnvironment -> m HEnvironment
forall (m :: * -> *) a. Monad m => a -> m a
return (HEnvironment -> m HEnvironment) -> HEnvironment -> m HEnvironment
forall a b. (a -> b) -> a -> b
$ (HSymbol
tyconName, ([HSymbol]
varsH, HType
htype, NoExtraInfo
NoExtraInfo)) (HSymbol, ([HSymbol], HType, NoExtraInfo))
-> HEnvironment -> HEnvironment
forall a. a -> [a] -> [a]
: HEnvironment
defnEnv
    Maybe TyThing
_ -> HEnvironment -> m HEnvironment
forall (m :: * -> *) a. Monad m => a -> m a
return []

lookupNameEverywhere :: G.GhcMonad m => Maybe G.ModuleInfo -> G.Name -> m (Maybe G.TyThing)
lookupNameEverywhere :: Maybe ModuleInfo -> Name -> m (Maybe TyThing)
lookupNameEverywhere (Just ModuleInfo
minfo) Name
name = do
  Maybe TyThing
thing <- ModuleInfo -> Name -> m (Maybe TyThing)
forall (m :: * -> *).
GhcMonad m =>
ModuleInfo -> Name -> m (Maybe TyThing)
G.modInfoLookupName ModuleInfo
minfo Name
name
  if Maybe TyThing -> Bool
forall a. Maybe a -> Bool
isJust Maybe TyThing
thing then Maybe TyThing -> m (Maybe TyThing)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe TyThing
thing else Name -> m (Maybe TyThing)
forall (m :: * -> *). GhcMonad m => Name -> m (Maybe TyThing)
G.lookupGlobalName Name
name
lookupNameEverywhere Maybe ModuleInfo
Nothing    Name
name = Name -> m (Maybe TyThing)
forall (m :: * -> *). GhcMonad m => Name -> m (Maybe TyThing)
G.lookupGlobalName Name
name

toHSymbol :: G.NamedThing a => a -> D.HSymbol
toHSymbol :: a -> HSymbol
toHSymbol = a -> HSymbol
forall a. NamedThing a => a -> HSymbol
G.getOccString

toLJTSymbol :: G.NamedThing a => a -> D.Symbol
toLJTSymbol :: a -> Symbol
toLJTSymbol = HSymbol -> Symbol
D.Symbol (HSymbol -> Symbol) -> (a -> HSymbol) -> a -> Symbol
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> HSymbol
forall a. NamedThing a => a -> HSymbol
G.getOccString

-- |Bindings which are in scope at a specific point.
type Environment = [(G.Name, G.Type)]

-- |Obtain a maximum number of solutions.
newtype MaxSolutions = Max Int

-- |Obtain the list of expressions which could fill
-- something with the given type.
-- The first flag specifies whether to return one
-- or more solutions to the problem.
djinn :: G.GhcMonad m => Bool -> Maybe G.ModuleInfo -> Environment -> G.Type -> MaxSolutions -> Int -> m [String]
djinn :: Bool
-> Maybe ModuleInfo
-> Environment
-> Type
-> MaxSolutions
-> Int
-> m [HSymbol]
djinn Bool
multi Maybe ModuleInfo
minfo Environment
env Type
ty (Max Int
mx) Int
microsec = do
  HEnvironment
tyEnv <- Maybe ModuleInfo -> Type -> m HEnvironment
forall (m :: * -> *).
GhcMonad m =>
Maybe ModuleInfo -> Type -> m HEnvironment
environment Maybe ModuleInfo
minfo Type
ty
  let form :: Formula
form = HEnvironment -> HType -> Formula
forall a. [(HSymbol, ([HSymbol], HType, a))] -> HType -> Formula
D.hTypeToFormula HEnvironment
tyEnv (Type -> HType
hType Type
ty)
      envF :: [(Symbol, Formula)]
envF = ((Name, Type) -> (Symbol, Formula))
-> Environment -> [(Symbol, Formula)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n,Type
t) -> (Name -> Symbol
forall a. NamedThing a => a -> Symbol
toLJTSymbol Name
n, HEnvironment -> HType -> Formula
forall a. [(HSymbol, ([HSymbol], HType, a))] -> HType -> Formula
D.hTypeToFormula HEnvironment
tyEnv (Type -> HType
hType Type
t))) Environment
env
      prfs :: [Proof]
prfs = Bool -> [(Symbol, Formula)] -> Formula -> [Proof]
D.prove Bool
multi [(Symbol, Formula)]
envF Formula
form
      trms :: [HSymbol]
trms = (Proof -> HSymbol) -> [Proof] -> [HSymbol]
forall a b. (a -> b) -> [a] -> [b]
map (HExpr -> HSymbol
D.hPrExpr (HExpr -> HSymbol) -> (Proof -> HExpr) -> Proof -> HSymbol
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proof -> HExpr
D.termToHExpr) [Proof]
prfs
  IO [HSymbol] -> m [HSymbol]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [HSymbol] -> m [HSymbol]) -> IO [HSymbol] -> m [HSymbol]
forall a b. (a -> b) -> a -> b
$ [HSymbol] -> Int -> Int -> (HSymbol -> Bool) -> IO [HSymbol]
forall a. [a] -> Int -> Int -> (a -> Bool) -> IO [a]
cropList [HSymbol]
trms Int
microsec Int
mx (\HSymbol
x -> HSymbol -> Int -> Bool
forall a. [a] -> Int -> Bool
lengthLessThan HSymbol
x Int
1000)

cropList :: [a] -> Int -> Int -> (a -> Bool) -> IO [a]
cropList :: [a] -> Int -> Int -> (a -> Bool) -> IO [a]
cropList [a]
_   Int
_  Int
0 a -> Bool
_ = [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
cropList [a]
lst Int
ms Int
n a -> Bool
chk =
  IO [a] -> (Async [a] -> IO [a]) -> IO [a]
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync (let !l :: [a]
l = [a]
lst in [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
l) ((Async [a] -> IO [a]) -> IO [a])
-> (Async [a] -> IO [a]) -> IO [a]
forall a b. (a -> b) -> a -> b
$ \Async [a]
a -> do
    Int -> IO ()
threadDelay Int
ms
    Maybe (Either SomeException [a])
res <- Async [a] -> IO (Maybe (Either SomeException [a]))
forall a. Async a -> IO (Maybe (Either SomeException a))
poll Async [a]
a
    case Maybe (Either SomeException [a])
res of
      Just Either SomeException [a]
r -> case Either SomeException [a]
r of
        Right (a
x:[a]
xs) -> if a -> Bool
chk a
x then do [a]
ys <- [a] -> Int -> Int -> (a -> Bool) -> IO [a]
forall a. [a] -> Int -> Int -> (a -> Bool) -> IO [a]
cropList [a]
xs Int
ms (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a -> Bool
chk
                                         [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> IO [a]) -> [a] -> IO [a]
forall a b. (a -> b) -> a -> b
$ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys
                                 else [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
        Either SomeException [a]
_            -> [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []
      Maybe (Either SomeException [a])
Nothing -> do Async [a] -> IO ()
forall a. Async a -> IO ()
cancel Async [a]
a
                    [a] -> IO [a]
forall (m :: * -> *) a. Monad m => a -> m a
return []

lengthLessThan :: [a] -> Int -> Bool
lengthLessThan :: [a] -> Int -> Bool
lengthLessThan []     Int
_ = Bool
True
lengthLessThan (a
_:[a]
_)  Int
0 = Bool
False
lengthLessThan (a
x:[a]
xs) Int
n = [a] -> Int -> Bool
forall a. [a] -> Int -> Bool
lengthLessThan [a]
xs (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)