{-|
Module      : Foreign.Storable.Generic.Plugin.Internal.GroupTypes
Copyright   : (c) Mateusz Kłoczko, 2016
License     : MIT
Maintainer  : mateusz.p.kloczko@gmail.com
Stability   : experimental
Portability : GHC-only

Grouping methods, both for types and core bindings.
-}
module Foreign.Storable.Generic.Plugin.Internal.GroupTypes 
    (
    -- Type ordering
      calcGroupOrder
    , substituteTyCon
    , getDataConArgs
    -- CoreBind ordering
    , groupBinds
    )
where

-- Management of Core.
import CoreSyn (Bind(..),Expr(..), CoreExpr, CoreBind, CoreProgram, Alt)
import Literal (Literal(..))
import Id  (isLocalId, isGlobalId,Id)
import Var (Var(..))
import Name (getOccName,mkOccName)
import OccName (OccName(..), occNameString)
import qualified Name as N (varName)
import SrcLoc (noSrcSpan)
import Unique (getUnique)
-- Compilation pipeline stuff
import HscMain (hscCompileCoreExpr)
import HscTypes (HscEnv,ModGuts(..))
import CoreMonad (CoreM,CoreToDo(..), getHscEnv)
import BasicTypes (CompilerPhase(..))
-- Haskell types 
import Type (isAlgType, splitTyConApp_maybe)
import TyCon (algTyConRhs, visibleDataCons)
import TyCoRep (Type(..), TyBinder(..))
import TysWiredIn (intDataCon)
import DataCon    (dataConWorkId,dataConOrigArgTys) 

import MkCore (mkWildValBinder)
-- Printing
import Outputable (cat, ppr, SDoc, showSDocUnsafe)
import Outputable (text, (<+>), ($$), nest)
import CoreMonad (putMsg, putMsgS)

-- Used to get to compiled values
import GHCi.RemoteTypes

import TyCon
import Type hiding (eqType)

import Unsafe.Coerce

import Data.List
import Data.Maybe
import Data.Either
import Debug.Trace
import Control.Monad.IO.Class

import Foreign.Storable.Generic.Plugin.Internal.Error
import Foreign.Storable.Generic.Plugin.Internal.Helpers
import Foreign.Storable.Generic.Plugin.Internal.Predicates
import Foreign.Storable.Generic.Plugin.Internal.Types


-- | Calculate the order of types. 
calcGroupOrder :: [Type] -> ([[Type]], Maybe Error)
calcGroupOrder :: [Type] -> ([[Type]], Maybe Error)
calcGroupOrder [Type]
types = [Type] -> [[Type]] -> ([[Type]], Maybe Error)
calcGroupOrder_rec [Type]
types []

calcGroupOrder_rec :: [Type]
                   -> [[Type]]
                   -> ([[Type]], Maybe Error)
calcGroupOrder_rec :: [Type] -> [[Type]] -> ([[Type]], Maybe Error)
calcGroupOrder_rec []    [[Type]]
acc = ([[Type]] -> [[Type]]
forall a. [a] -> [a]
reverse [[Type]]
acc, Maybe Error
forall a. Maybe a
Nothing)
calcGroupOrder_rec [Type]
types [[Type]]
acc = do
    let ([Type]
layer, [Type]
rest) = [Type] -> [Type] -> [Type] -> [Type] -> ([Type], [Type])
calcGroupOrder_iteration [Type]
types [] [] []
        layer' :: [Type]
layer'       = (Type -> Type -> Bool) -> [Type] -> [Type]
forall a. (a -> a -> Bool) -> [a] -> [a]
nubBy Type -> Type -> Bool
eqType [Type]
layer
    if [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
layer' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then ([[Type]] -> [[Type]]
forall a. [a] -> [a]
reverse [[Type]]
acc, Error -> Maybe Error
forall a. a -> Maybe a
Just (Error -> Maybe Error) -> Error -> Maybe Error
forall a b. (a -> b) -> a -> b
$ Int -> [Type] -> Error
OrderingFailedTypes ([[Type]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Type]]
acc) [Type]
rest)
        else [Type] -> [[Type]] -> ([[Type]], Maybe Error)
calcGroupOrder_rec [Type]
rest ([Type]
layer'[Type] -> [[Type]] -> [[Type]]
forall a. a -> [a] -> [a]
:[[Type]]
acc)

-- | This could be done more efficently if we'd 
-- represent the problem as a graph problem.
calcGroupOrder_iteration :: [Type] -- ^ Type to check 
                         -> [Type] -- ^ Type that are checked
                         -> [Type] -- ^ Type that are in this layer
                         -> [Type] -- ^ Type that are not.
                         -> ([Type], [Type]) -- Returning types in this layer and the next ones.
calcGroupOrder_iteration :: [Type] -> [Type] -> [Type] -> [Type] -> ([Type], [Type])
calcGroupOrder_iteration []     [Type]
checked [Type]
accepted [Type]
rejected = ([Type]
accepted, [Type]
rejected)
calcGroupOrder_iteration (Type
t:[Type]
ts) [Type]
checked [Type]
accepted [Type]
rejected = do
    let args :: [Type]
args = Type -> [Type]
getDataConArgs Type
t
        -- Are the t's arguments equal to some other ?
        is_arg_somewhere :: Bool
is_arg_somewhere = (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Type
t -> Type -> [Type] -> Bool
elemType Type
t [Type]
args) [Type]
checked Bool -> Bool -> Bool
|| (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Type
t -> Type -> [Type] -> Bool
elemType Type
t [Type]
args) [Type]
ts

    if Bool
is_arg_somewhere
        then [Type] -> [Type] -> [Type] -> [Type] -> ([Type], [Type])
calcGroupOrder_iteration [Type]
ts (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
checked)  [Type]
accepted    (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
rejected)
        else [Type] -> [Type] -> [Type] -> [Type] -> ([Type], [Type])
calcGroupOrder_iteration [Type]
ts (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
checked) (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
accepted)  [Type]
rejected

-- | Used for type substitution. 
-- Whether a TyVar appears, replace it with a Type.
type TypeScope = (TyVar, Type)

-- | Functions doing the type substitutions.

-- Examples
-- 
-- substituteTyCon [(a,Int)]           a          = Int
-- substituteTyCon [(a,Int),(b,Char)] (AType b a) = AType Char Int
substituteTyCon :: [TypeScope] -> Type -> Type
substituteTyCon :: [TypeScope] -> Type -> Type
substituteTyCon []         Type
tc_app             = Type
tc_app
substituteTyCon [TypeScope]
type_scope old :: Type
old@(TyVarTy  Var
ty_var) 
-- Substitute simple type variables
    = case (TypeScope -> Bool) -> [TypeScope] -> Maybe TypeScope
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(Var
av,Type
_) -> Var
av Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
ty_var) [TypeScope]
type_scope of
          Just (Var
_, Type
new_type) -> Type
new_type
          Maybe TypeScope
Nothing            -> Type
old
substituteTyCon [TypeScope]
type_scope (TyConApp TyCon
tc [Type]
args)
-- Substitute type constructors
    = TyCon -> [Type] -> Type
TyConApp TyCon
tc ([Type] -> Type) -> [Type] -> Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ([TypeScope] -> Type -> Type
substituteTyCon [TypeScope]
type_scope) [Type]
args
substituteTyCon [TypeScope]
type_scope Type
t = Type
t 

-- | Get data constructor arguments from an algebraic type.
getDataConArgs :: Type -> [Type]
getDataConArgs :: Type -> [Type]
getDataConArgs Type
t 
    | Type -> Bool
isAlgType Type
t
    , Just (TyCon
tc, [Type]
ty_args) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
t
    , [Var]
ty_vars <- TyCon -> [Var]
tyConTyVars TyCon
tc
    = do
    -- Substitute data_cons args with type args,
    -- using ty_vars as keys.
    let type_scope :: [TypeScope]
type_scope = [Var] -> [Type] -> [TypeScope]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
ty_vars [Type]
ty_args
        data_cons :: [Type]
data_cons  = (DataCon -> [Type]) -> [DataCon] -> [Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DataCon -> [Type]
dataConOrigArgTys ([DataCon] -> [Type]) -> [DataCon] -> [Type]
forall a b. (a -> b) -> a -> b
$ (AlgTyConRhs -> [DataCon]
visibleDataCons(AlgTyConRhs -> [DataCon])
-> (TyCon -> AlgTyConRhs) -> TyCon -> [DataCon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.TyCon -> AlgTyConRhs
algTyConRhs) TyCon
tc
    (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ([TypeScope] -> Type -> Type
substituteTyCon [TypeScope]
type_scope) [Type]
data_cons  
    | Bool
otherwise = []



-- | Group bindings according to type groups.
groupBinds :: [[Type]]   -- ^ Type groups.
           -> [CoreBind] -- ^ Should be only NonRecs. 
           -> ([[CoreBind]], Maybe Error)
-- perhaps add some safety so non-recs won't get here.
groupBinds :: [[Type]] -> [CoreBind] -> ([[CoreBind]], Maybe Error)
groupBinds [[Type]]
type_groups [CoreBind]
binds = [[Type]]
-> [CoreBind] -> [[CoreBind]] -> ([[CoreBind]], Maybe Error)
groupBinds_rec [[Type]]
type_groups [CoreBind]
binds [] 

-- | Iteration for groupBinds
groupBinds_rec :: [[Type]]      -- ^ Group of types
               -> [CoreBind]    -- ^ Ungrouped bindings
               -> [[CoreBind]]  -- ^ Grouped bindings
               -> ([[CoreBind]], Maybe Error) -- ^ Grouped bindings, and perhaps an error)
groupBinds_rec :: [[Type]]
-> [CoreBind] -> [[CoreBind]] -> ([[CoreBind]], Maybe Error)
groupBinds_rec []       []    [[CoreBind]]
acc = ([[CoreBind]] -> [[CoreBind]]
forall a. [a] -> [a]
reverse [[CoreBind]]
acc,Maybe Error
forall a. Maybe a
Nothing)
groupBinds_rec ([Type]
a:[[Type]]
as)   []    [[CoreBind]]
acc = ([[CoreBind]] -> [[CoreBind]]
forall a. [a] -> [a]
reverse [[CoreBind]]
acc,Error -> Maybe Error
forall a. a -> Maybe a
Just (Error -> Maybe Error) -> Error -> Maybe Error
forall a b. (a -> b) -> a -> b
$ SDoc -> Error
OtherError SDoc
msg)
    where msg :: SDoc
msg =    String -> SDoc
text String
"Could not find any bindings." 
                SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"Is the second pass placed after main simplifier phases ?" 
groupBinds_rec []       [CoreBind]
binds [[CoreBind]]
acc = ([[CoreBind]] -> [[CoreBind]]
forall a. [a] -> [a]
reverse [[CoreBind]]
acc,Error -> Maybe Error
forall a. a -> Maybe a
Just (Error -> Maybe Error) -> Error -> Maybe Error
forall a b. (a -> b) -> a -> b
$ Int -> [CoreBind] -> Error
OrderingFailedBinds ([[CoreBind]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[CoreBind]]
acc) [CoreBind]
binds)
groupBinds_rec ([Type]
tg:[[Type]]
tgs) [CoreBind]
binds [[CoreBind]]
acc = do
    let predicate :: CoreBind -> Bool
predicate (NonRec Var
id Expr Var
_) = case Type -> Maybe Type
getGStorableType (Type -> Maybe Type) -> Type -> Maybe Type
forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
id of
            Just Type
t -> Type
t Type -> [Type] -> Bool
`elemType` [Type]
tg
            Maybe Type
Nothing -> Bool
False
        predicate (Rec [(Var, Expr Var)]
_) = Bool
False
    let ([CoreBind]
layer, [CoreBind]
rest) = (CoreBind -> Bool) -> [CoreBind] -> ([CoreBind], [CoreBind])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition CoreBind -> Bool
predicate [CoreBind]
binds
    if [CoreBind] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [CoreBind]
layer Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 
        then ([[CoreBind]] -> [[CoreBind]]
forall a. [a] -> [a]
reverse [[CoreBind]]
acc, Error -> Maybe Error
forall a. a -> Maybe a
Just (Error -> Maybe Error) -> Error -> Maybe Error
forall a b. (a -> b) -> a -> b
$ Int -> [CoreBind] -> Error
OrderingFailedBinds ([[CoreBind]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[CoreBind]]
acc) [CoreBind]
rest)
        else [[Type]]
-> [CoreBind] -> [[CoreBind]] -> ([[CoreBind]], Maybe Error)
groupBinds_rec [[Type]]
tgs [CoreBind]
rest ([CoreBind] -> [CoreBind]
forall a. [a] -> [a]
reverse [CoreBind]
layer[CoreBind] -> [[CoreBind]] -> [[CoreBind]]
forall a. a -> [a] -> [a]
:[[CoreBind]]
acc)