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

Contains methods for calculating type ordering and performing the compile-substitution optimisation.

-}
{-#LANGUAGE CPP #-}

module Foreign.Storable.Generic.Plugin.Internal 
    ( groupTypes
    , gstorableSubstitution)
where

-- Management of Core.

import Prelude hiding ((<>))

#if MIN_VERSION_GLASGOW_HASKELL(9,0,1,0)
import GHC.Core          (Bind(..),Expr(..), CoreExpr, CoreBind, CoreProgram, Alt)
import GHC.Types.Literal (Literal(..))
import GHC.Types.Id      (isLocalId, isGlobalId,Id, modifyInlinePragma, setInlinePragma, idInfo)
import GHC.Types.Id.Info
import GHC.Types.Var             (Var(..))
import GHC.Types.Name            (getOccName,mkOccName)
import GHC.Types.Name.Occurrence (OccName(..), occNameString)
import qualified GHC.Types.Name as N (varName)
import GHC.Types.SrcLoc (noSrcSpan)
import GHC.Types.Unique (getUnique)
import GHC.Driver.Main (hscCompileCoreExpr, getHscEnv)
#if MIN_VERSION_GLASGOW_HASKELL(9,2,0,0)
import GHC.Driver.Env.Types (HscEnv)
import GHC.Unit.Module.ModGuts (ModGuts(..))
#else
import GHC.Driver.Types (HscEnv,ModGuts(..))
#endif
import GHC.Core.Opt.Monad
    (CoreM, CoreToDo(..), 
     getHscEnv, getDynFlags, putMsg, putMsgS)
import GHC.Types.Basic (CompilerPhase(..))
import GHC.Core.Type (isAlgType, splitTyConApp_maybe)
import GHC.Core.TyCon (tyConKind, algTyConRhs, visibleDataCons)
import GHC.Core.TyCo.Rep (Type(..), TyBinder(..))
import GHC.Builtin.Types   (intDataCon)
import GHC.Core.DataCon    (dataConWorkId,dataConOrigArgTys) 
import GHC.Core.Make       (mkWildValBinder)
import GHC.Utils.Outputable 
    (cat, ppr, SDoc, showSDocUnsafe,
     ($$), ($+$), hsep, vcat, empty,text, 
     (<>), (<+>), nest, int, colon,hcat, comma, 
     punctuate, fsep) 
import GHC.Core.Opt.Monad (putMsg, putMsgS)
#elif MIN_VERSION_GLASGOW_HASKELL(8,2,1,0)
import CoreSyn (Bind(..),Expr(..), CoreExpr, CoreBind, CoreProgram, Alt)
import Literal (Literal(..))
import Id  (isLocalId, isGlobalId,Id, modifyInlinePragma, setInlinePragma, idInfo)
import IdInfo
import Var (Var(..))
import Name (getOccName,mkOccName)
import OccName (OccName(..), occNameString)
import qualified Name as N (varName)
import SrcLoc (noSrcSpan)
import Unique (getUnique)
import HscMain (hscCompileCoreExpr)
import HscTypes (HscEnv,ModGuts(..))
import CoreMonad 
    (CoreM, CoreToDo(..), 
     getHscEnv, getDynFlags, putMsg, putMsgS)
import BasicTypes (CompilerPhase(..))
import Type (isAlgType, splitTyConApp_maybe)
import TyCon (tyConKind, algTyConRhs, visibleDataCons)
import TyCoRep (Type(..), TyBinder(..))
import TysWiredIn (intDataCon)
import DataCon    (dataConWorkId,dataConOrigArgTys) 
import MkCore (mkWildValBinder)
import Outputable 
    (cat, ppr, SDoc, showSDocUnsafe,
     ($$), ($+$), hsep, vcat, empty,text, 
     (<>), (<+>), nest, int, colon,hcat, comma, 
     punctuate, fsep) 
import CoreMonad (putMsg, putMsgS)
#endif



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

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



--------------------
-- Grouping types --
--------------------


groupTypes_errors :: Flags -> [Error] -> CoreM ()
groupTypes_errors :: Flags -> [Error] -> CoreM ()
groupTypes_errors Flags
flags [Error]
errors = do
    let (Flags Verbosity
verb CrashOnWarning
to_crash) = Flags
flags
        crasher :: [a] -> m ()
crasher [a]
errs = case [a]
errs of
            [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
            [a]
_  -> forall a. HasCallStack => [Char] -> a
error [Char]
"Crashing..."
        print_header :: SDoc -> SDoc
print_header SDoc
txt = case Verbosity
verb of
            Verbosity
None  -> SDoc
empty
            Verbosity
other ->    [Char] -> SDoc
text [Char]
"Errors while grouping types - types not found for: "
                     SDoc -> SDoc -> SDoc
$$ Int -> SDoc -> SDoc
nest Int
4 SDoc
txt
        print_tyNotF :: Verbosity -> Var -> SDoc
print_tyNotF Verbosity
verb Var
id = case Verbosity
verb of
            Verbosity
None  -> SDoc
empty
            Verbosity
other -> forall a. Outputable a => a -> SDoc
ppr Var
id SDoc -> SDoc -> SDoc
$$ Int -> SDoc -> SDoc
nest Int
13 ([Char] -> SDoc
text [Char]
"::") SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr (Var -> Type
varType Var
id)
        print_err :: Error -> SDoc
print_err    Error
err = case Error
err of
            TypeNotFound Var
id -> Verbosity -> Var -> SDoc
print_tyNotF Verbosity
verb Var
id
            Error
other           -> Verbosity -> Error -> SDoc
pprError Verbosity
verb Error
other
        printer :: [Error] -> CoreM ()
printer [Error]
errs = case [Error]
errs of
            [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
            [Error]
ls ->  SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ SDoc -> SDoc
print_header ([SDoc] -> SDoc
vcat (forall a b. (a -> b) -> [a] -> [b]
map Error -> SDoc
print_err [Error]
errs)) 
    -- Do printing
    -- Eventually crash.
    [Error] -> CoreM ()
printer [Error]
errors
    forall (f :: * -> *).
Applicative f =>
CrashOnWarning -> f () -> f ()
when CrashOnWarning
to_crash forall a b. (a -> b) -> a -> b
$ forall {m :: * -> *} {a}. Monad m => [a] -> m ()
crasher [Error]
errors

groupTypes_info :: Flags -> [[Type]] -> CoreM ()
groupTypes_info :: Flags -> [[Type]] -> CoreM ()
groupTypes_info Flags
flags [[Type]]
types = do
    let (Flags Verbosity
verb CrashOnWarning
_) = Flags
flags
        -- If verbosity is set, do the printing
        print_header :: SDoc -> SDoc
print_header SDoc
txt = case Verbosity
verb of
            Verbosity
None  -> SDoc
empty
            Verbosity
other ->    [Char] -> SDoc
text [Char]
"GStorable instances will be optimised in the following order"
                    SDoc -> SDoc -> SDoc
$+$ Int -> SDoc -> SDoc
nest Int
4 SDoc
txt
                    SDoc -> SDoc -> SDoc
$+$ [Char] -> SDoc
text [Char]
""
        print_layer :: [a] -> Int -> SDoc
print_layer [a]
layer Int
ix = Int -> SDoc
int Int
ix SDoc -> SDoc -> SDoc
<> [Char] -> SDoc
text [Char]
":" SDoc -> SDoc -> SDoc
<+> [SDoc] -> SDoc
fsep (SDoc -> [SDoc] -> [SDoc]
punctuate SDoc
comma forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a. Outputable a => a -> SDoc
ppr [a]
layer)
        -- Print groups of types
        printer :: [[a]] -> CoreM ()
printer [[a]]
groups = case [[a]]
groups of
            [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
            [[a]]
_  -> SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ SDoc -> SDoc
print_header ([SDoc] -> SDoc
vcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a}. Outputable a => [a] -> Int -> SDoc
print_layer [[a]]
groups [Int
1..])
    -- Do the printing
    forall {a}. Outputable a => [[a]] -> CoreM ()
printer [[Type]]
types


-- | Find GStorable identifiers, obtain their types, and calculate
-- the order of compilation.
groupTypes :: Flags -> IORef [[Type]] -> ModGuts -> CoreM ModGuts
groupTypes :: Flags -> IORef [[Type]] -> ModGuts -> CoreM ModGuts
groupTypes Flags
flags IORef [[Type]]
type_order_ref ModGuts
guts = do
    let binds :: CoreProgram
binds = ModGuts -> CoreProgram
mg_binds ModGuts
guts
        -- Get GStorable ids that are fully defined.
        all_ids :: [Var]
all_ids = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap CoreBind -> [Var]
getIdsBind CoreProgram
binds
        with_typecheck :: Var -> CrashOnWarning
with_typecheck = (Type -> Maybe Type)
-> (Var -> CrashOnWarning) -> Var -> CrashOnWarning
withTypeCheck Type -> Maybe Type
getGStorableType Var -> CrashOnWarning
isGStorableId
        predicate :: Var -> CrashOnWarning
predicate Var
id = forall (t :: * -> *).
Foldable t =>
t CrashOnWarning -> CrashOnWarning
and [ Var -> CrashOnWarning
with_typecheck Var
id
                           , CrashOnWarning -> CrashOnWarning
not (Type -> CrashOnWarning
hasGStorableConstraints forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
id)
                           ]
        gstorable_ids :: [Var]
gstorable_ids = forall a. (a -> CrashOnWarning) -> [a] -> [a]
filter Var -> CrashOnWarning
predicate [Var]
all_ids
        -- Now process them - different ids
        -- will have different type signatures.
        -- It is possible to fetch the types from them.
        m_gstorable_types :: [Maybe Type]
m_gstorable_types = forall a b. (a -> b) -> [a] -> [b]
map (Type -> Maybe Type
getGStorableTypeforall b c a. (b -> c) -> (a -> b) -> a -> c
.Var -> Type
varType) [Var]
gstorable_ids
        -- Grab any errors related to types not found.
        bad_types_zip :: Var -> Maybe a -> Maybe Error
bad_types_zip Var
id Maybe a
m_t = case Maybe a
m_t of
            Maybe a
Nothing -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Var -> Error
TypeNotFound Var
id
            Just a
_  -> forall a. Maybe a
Nothing
        bad_types :: [Error]
bad_types     =    forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a}. Var -> Maybe a -> Maybe Error
bad_types_zip [Var]
gstorable_ids [Maybe Type]
m_gstorable_types 
        -- type_list is used instead of type_set because Type has no uniquable instance.
        type_list :: [Type]
type_list = [ Type
t | Just Type
t <- [Maybe Type]
m_gstorable_types]
        -- Calculate the type ordering.
        ([[Type]]
type_order,Maybe Error
m_error) = [Type] -> ([[Type]], Maybe Error)
calcGroupOrder [Type]
type_list

    Flags -> [[Type]] -> CoreM ()
groupTypes_info Flags
flags [[Type]]
type_order
    Flags -> [Error] -> CoreM ()
groupTypes_errors Flags
flags [Error]
bad_types
    
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef [[Type]]
type_order_ref [[Type]]
type_order
    forall (m :: * -> *) a. Monad m => a -> m a
return ModGuts
guts



------------------------------------------------
-- Grouping and compiling GStorable CoreBinds --
------------------------------------------------

-- | Print errors related to CoreBind grouping.
-- Return the badly grouped bindings, and perhaps crash
-- the compiler.
grouping_errors :: Flags            -- ^ Verbosity and ToCrash options 
                -> Maybe Error      -- ^ The error
                -> CoreM [CoreBind] -- ^ Recovered bindings.
grouping_errors :: Flags -> Maybe Error -> CoreM CoreProgram
grouping_errors Flags
flags Maybe Error
m_err = do
   let (Flags Verbosity
_ CrashOnWarning
to_crash) = Flags
flags
       verb :: Verbosity
verb = Verbosity
Some
       crasher :: Maybe a -> m ()
crasher Maybe a
m_e = case Maybe a
m_e of
           Maybe a
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
           Just a
_  -> forall a. HasCallStack => [Char] -> a
error [Char]
"Crashing..."
       print_header :: SDoc -> SDoc
print_header SDoc
txt = case Verbosity
verb of
           Verbosity
None  -> SDoc
empty
           Verbosity
other ->    [Char] -> SDoc
text [Char]
"Errors while grouping bindings: "
                    SDoc -> SDoc -> SDoc
$$ Int -> SDoc -> SDoc
nest Int
4 SDoc
txt 
       printer :: Maybe Error -> CoreM ()
printer Maybe Error
m_err = case Maybe Error
m_err of
           Maybe Error
Nothing  -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
           Just Error
err ->  SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ SDoc -> SDoc
print_header (Verbosity -> Error -> SDoc
pprError Verbosity
verb Error
err) 
       ungroup :: Maybe Error -> CoreProgram
ungroup Maybe Error
m_e = case Maybe Error
m_e of
           Just (OrderingFailedBinds Int
_ CoreProgram
rest) -> CoreProgram
rest
           Maybe Error
_                                 -> []
   Maybe Error -> CoreM ()
printer Maybe Error
m_err
   forall (f :: * -> *).
Applicative f =>
CrashOnWarning -> f () -> f ()
when CrashOnWarning
to_crash forall a b. (a -> b) -> a -> b
$ forall {m :: * -> *} {a}. Monad m => Maybe a -> m ()
crasher Maybe Error
m_err
   forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Maybe Error -> CoreProgram
ungroup Maybe Error
m_err


-- | Print the information related to found GStorable ids.
foundBinds_info :: Flags    -- ^ Verbosity and ToCrash options 
                -> [Id]     -- ^ GStorable ids.
                -> CoreM ()
foundBinds_info :: Flags -> [Var] -> CoreM ()
foundBinds_info Flags
flags [Var]
ids = do
    -- For Pretty printing
    DynFlags
dyn_flags <- forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
    let (Flags Verbosity
verb CrashOnWarning
_) = Flags
flags
        -- If verbosity is set, do the printing
        print_header :: SDoc -> SDoc
print_header SDoc
txt = case Verbosity
verb of
            Verbosity
None  -> SDoc
empty
            Verbosity
other ->    [Char] -> SDoc
text [Char]
"The following bindings are to be optimised:"
                    SDoc -> SDoc -> SDoc
$+$ Int -> SDoc -> SDoc
nest Int
4 SDoc
txt
        print_binding :: a -> SDoc
print_binding a
id = forall a. Outputable a => a -> SDoc
ppr a
id
        -- Print groups of types
        printer :: [[Var]] -> CoreM ()
printer [[Var]]
the_groups = case [[Var]]
the_groups of
            [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
            [[Var]]
_  -> SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ SDoc -> SDoc
print_header forall a b. (a -> b) -> a -> b
$ [SDoc] -> SDoc
vcat (forall a b. (a -> b) -> [a] -> [b]
map [Var] -> SDoc
print_group [[Var]]
the_groups)
        -- Use eqType for maybes
        eqType_maybe :: Maybe Type -> Maybe Type -> CrashOnWarning
eqType_maybe (Just Type
t1) (Just Type
t2) = Type
t1 Type -> Type -> CrashOnWarning
`eqType` Type
t2
        eqType_maybe Maybe Type
_         Maybe Type
_         = CrashOnWarning
False
        -- group and sort the bindings 
        grouped :: [[Var]]
grouped = forall a. (a -> a -> CrashOnWarning) -> [a] -> [[a]]
groupBy (\Var
i1 Var
i2 -> (Type -> Maybe Type
getGStorableType forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
i1) Maybe Type -> Maybe Type -> CrashOnWarning
`eqType_maybe` (Type -> Maybe Type
getGStorableType forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
i2) ) [Var]
ids
        sorting :: [Var] -> [Var]
sorting = forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (\Var
i1 Var
i2 -> Var -> Name
varName Var
i1 forall a. Ord a => a -> a -> Ordering
`compare` Var -> Name
varName Var
i2)
        sorted :: [[Var]]
sorted  = forall a b. (a -> b) -> [a] -> [b]
map [Var] -> [Var]
sorting [[Var]]
grouped
        -- print groups of bindings
        print_group :: [Var] -> SDoc
print_group [Var]
the_group = case [Var]
the_group of
            [] -> SDoc
empty
            (Var
h:[Var]
_) -> case Type -> Maybe Type
getGStorableType forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
h of
                Just Type
gtype ->     forall a. Outputable a => a -> SDoc
ppr  Type
gtype
                              SDoc -> SDoc -> SDoc
$+$ ([SDoc] -> SDoc
fsep forall a b. (a -> b) -> a -> b
$ SDoc -> [SDoc] -> [SDoc]
punctuate SDoc
comma (forall a b. (a -> b) -> [a] -> [b]
map forall a. Outputable a => a -> SDoc
print_binding [Var]
the_group))
                Maybe Type
Nothing    -> forall a. Outputable a => a -> SDoc
ppr [Char]
"Could not get the type of a binding:" 
                              SDoc -> SDoc -> SDoc
$+$ Int -> SDoc -> SDoc
nest Int
4 (forall a. Outputable a => a -> SDoc
ppr Var
h SDoc -> SDoc -> SDoc
<+> [Char] -> SDoc
text [Char]
"::" SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr (Var -> Type
varType Var
h))
    -- Print the ids
    [[Var]] -> CoreM ()
printer [[Var]]
sorted

-- | Do the optimisation for GStorable bindings.
gstorableSubstitution :: Flags          -- ^ Verbosity and ToCrash options.
                      -> IORef [[Type]] -- ^ Reference to grouped types.
                      -> ModGuts        -- ^ Information about compiled module.
                      -> CoreM ModGuts  -- ^ Information about compiled module, with GStorable optimisations.
gstorableSubstitution :: Flags -> IORef [[Type]] -> ModGuts -> CoreM ModGuts
gstorableSubstitution Flags
flags IORef [[Type]]
type_order_ref ModGuts
guts = do 
    [[Type]]
type_hierarchy <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef [[Type]]
type_order_ref 
    let binds :: CoreProgram
binds  = ModGuts -> CoreProgram
mg_binds ModGuts
guts
        -- Get all GStorable binds.
        -- Check whether the type has no constraints.
        typeCheck :: Type -> Maybe Type
typeCheck Type
t = if Type -> CrashOnWarning
hasGStorableConstraints Type
t
            then forall a. Maybe a
Nothing
            else Type -> Maybe Type
getGStorableMethodType Type
t
        -- predicate = toIsBind (withTypeCheck typeCheck isGStorableMethodId)
        predicate :: CoreBind -> CrashOnWarning
predicate = (Var -> CrashOnWarning) -> CoreBind -> CrashOnWarning
toIsBind (Var -> CrashOnWarning
isGStorableMethodId)
        
        (CoreProgram
gstorable_binds,CoreProgram
rest) = forall a. (a -> CrashOnWarning) -> [a] -> ([a], [a])
partition CoreBind -> CrashOnWarning
predicate CoreProgram
binds
        -- Check if there are any recursives somehow
        -- The plugin won't be able to handle them.
        (CoreProgram
nonrecs, CoreProgram
recs) = forall a. (a -> CrashOnWarning) -> [a] -> ([a], [a])
partition CoreBind -> CrashOnWarning
isNonRecBind CoreProgram
gstorable_binds
        -- Group the gstorables by nestedness
        ([CoreProgram]
grouped_binds, Maybe Error
m_err_group) = [[Type]] -> CoreProgram -> ([CoreProgram], Maybe Error)
groupBinds [[Type]]
type_hierarchy CoreProgram
nonrecs
   
    Flags -> [Var] -> CoreM ()
foundBinds_info Flags
flags forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap CoreBind -> [Var]
getIdsBind forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [CoreProgram]
grouped_binds 
    -- Check for errors
    CoreProgram
not_grouped <- Flags -> Maybe Error -> CoreM CoreProgram
grouping_errors Flags
flags Maybe Error
m_err_group
    -- Compile and replace gstorable bindings
    CoreProgram
new_gstorables <- Flags -> [CoreProgram] -> CoreProgram -> CoreM CoreProgram
compileGroups Flags
flags [CoreProgram]
grouped_binds CoreProgram
rest -- perhaps return errors here ?
    
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ ModGuts
guts {mg_binds :: CoreProgram
mg_binds = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [CoreProgram
new_gstorables, CoreProgram
not_grouped,CoreProgram
recs,CoreProgram
rest]}