content: {-
  Copyright   : (c) Jean-Christophe Mincke, 2016-2017
  
  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
  -}
  
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  
  module Control.Distributed.Blast.Slave.Analyser
  (
    SExp (..)
    , NodeTypeInfo (..)
    , LExpInfo (..)
    , RConstInfo (..)
    , RMapInfo (..)
    , InfoMap
    , analyseLocal
  )
  where
  
  --import Debug.Trace
  import            Control.Bool (unlessM)
  import            Control.Monad.Logger
  import            Control.Monad.IO.Class
  import            Control.Monad.Trans.Either
  import            Control.Monad.Trans.State
  import qualified  Data.ByteString as BS
  import qualified  Data.Map as M
  import qualified  Data.Set as S
  import qualified  Data.Serialize as S
  import qualified  Data.Text as T
  import qualified  Data.Vault.Strict as V
  
  import            Control.Distributed.Blast.Types
  import            Control.Distributed.Blast.Common.Analyser
  
  
  
  data SExp (k::Kind) a where
    SRApply :: Int -> V.Key (Data b) -> ExpClosure SExp a b -> SExp 'Remote a -> SExp 'Remote b
    SRConst ::  (S.Serialize b) => Int -> V.Key (Data b) -> a -> SExp 'Remote b
    SLConst :: Int -> V.Key a -> IO a -> SExp 'Local a
    SCollect :: (S.Serialize b) => Int -> V.Key a -> SExp 'Remote b -> SExp 'Local a
    SLApply :: Int -> V.Key b -> SExp 'Local (a -> b) -> SExp 'Local a -> SExp 'Local b
  
  
  
  instance (MonadLoggerIO m) => Builder m SExp where
    makeRApply i f a = do
      k <- liftIO V.newKey
      return $ SRApply i k f a
    makeRConst i _ a = do
      k <- liftIO V.newKey
      return $ SRConst i k a
    makeLConst i a = do
      k <- liftIO V.newKey
      return $ SLConst i k a
    makeCollect i _ a = do
      k <- liftIO V.newKey
      return $ SCollect i k a
    makeLApply i f a = do
      k <- liftIO V.newKey
      return $ SLApply i k f a
  
  instance Indexable SExp where
    getIndex (SRApply n _ _ _) = n
    getIndex (SRConst n _ _) = n
    getIndex (SLConst n _ _) = n
    getIndex (SCollect n _ _) = n
    getIndex (SLApply n _ _ _) = n
  
  
  
  type RemoteCacher = Data BS.ByteString -> V.Vault -> V.Vault
  type LocalCacher = BS.ByteString -> V.Vault -> V.Vault
  
  type RemoteCacheReader = V.Vault -> Maybe (Data BS.ByteString)
  
  type UnCacher = V.Vault -> V.Vault
  
  data NodeTypeInfo =
    NtRMap RMapInfo
    |NtRConst RConstInfo
    |NtLExp LExpInfo
    |NtLExpNoCache
  
  data RMapInfo = MkRMapInfo {
    _rmRemoteClosure :: RemoteClosureImpl
    , _rmUnCacher :: UnCacher
    , _rmCacheReader :: Maybe RemoteCacheReader
    }
  
  data RConstInfo = MkRConstInfo {
    _rcstCacher :: RemoteCacher
    , _rcstUnCacher :: UnCacher
    , _rcstCacheReader :: Maybe RemoteCacheReader
    }
  
  data LExpInfo = MkLExpInfo {
    _lexpCacher :: LocalCacher
    , _lexpUnCacher :: UnCacher
    }
  
  type InfoMap = GenericInfoMap NodeTypeInfo
  
  
  
  getVal :: (Monad m) =>  CachedValType -> V.Vault -> V.Key a -> EitherT RemoteClosureResult m a
  getVal cvt vault key =
    case V.lookup key vault of
    Just v -> right v
    Nothing -> left $ RcRespCacheMiss cvt
  
  getLocalVal :: (Monad m) =>  CachedValType -> V.Vault -> V.Key a -> EitherT RemoteClosureResult m a
  getLocalVal cvt vault key  =
    case V.lookup key vault of
    Just v -> right v
    Nothing -> left $ RcRespCacheMiss cvt
  
  makeUnCacher :: V.Key a -> V.Vault -> V.Vault
  makeUnCacher k vault = V.delete k vault
  
  mkRemoteClosure :: forall a b m . (MonadLoggerIO m) =>
    V.Key (Data a) -> V.Key (Data b) -> ExpClosure SExp a b -> StateT InfoMap m RemoteClosureImpl
  mkRemoteClosure keya keyb (ExpClosure e f) = do
    analyseLocal e
    addLocalExpCacheM e
    let keyc = getLocalVaultKey e
    return $ wrapClosure keyc keya keyb f
  
  
  wrapClosure :: forall a b c .
              V.Key c -> V.Key (Data a) -> V.Key (Data b) -> (c -> a -> IO b) -> RemoteClosureImpl
  wrapClosure keyc keya keyb f =
      proc
      where
      proc vault = do
        r' <- runEitherT r
        return $ either (\l -> (l, vault)) id r'
        where
        r = do
          c <- getLocalVal CachedFreeVar vault keyc
          av <- getVal CachedArg vault keya
          brdd <- liftIO $ f' c av
          let vault' = V.insert keyb brdd vault
          return (RcRespOk, vault')
        f' _ NoData = return NoData
        f' c (Data a) = do
          x <-  f c a
          return $ Data x
  
  visitLocalExp :: Int -> InfoMap -> InfoMap
  visitLocalExp n m =
    case M.lookup n m of
      Just (GenericInfo _ _ ) -> m
      Nothing -> M.insert n (GenericInfo S.empty NtLExpNoCache) m
  
  
  
  visitLocalExpM :: (MonadLoggerIO m) => Int -> StateT InfoMap m ()
  visitLocalExpM n = do
    $(logInfo) $ T.pack  ("Visiting local exp node: " ++ show n)
    m <- get
    put $ visitLocalExp n m
  
  
  addLocalExpCache :: (S.Serialize a) => Int -> V.Key a -> InfoMap -> InfoMap
  addLocalExpCache n key m =
    case M.lookup n m of
    Just (GenericInfo c NtLExpNoCache) -> M.insert n (GenericInfo c (NtLExp (MkLExpInfo (makeCacher key) (makeUnCacher key)))) m
    Nothing -> M.insert n (GenericInfo S.empty (NtLExp (MkLExpInfo (makeCacher key) (makeUnCacher key)))) m
    Just (GenericInfo _ (NtLExp _)) -> m
    _ ->  error ("Node " ++ show n ++ " cannot add local exp cache")
    where
    makeCacher :: (S.Serialize a) => V.Key a -> BS.ByteString -> V.Vault -> V.Vault
    makeCacher k bs vault =
      case S.decode bs of
      Left e -> error $ ("Cannot deserialize value: " ++ e)
      Right a -> V.insert k a vault
  
  addLocalExpCacheM :: forall a m. (MonadLoggerIO m, S.Serialize a) =>
    SExp 'Local a -> StateT InfoMap m ()
  addLocalExpCacheM e = do
    let n = getLocalIndex e
    $(logInfo) $ T.pack  ("Adding cache to local exp node: " ++ show n)
    let key = getLocalVaultKey e
    m <- get
    put $ addLocalExpCache n key m
  
  addRemoteExpCacheReader :: (S.Serialize a) => Int -> V.Key (Data a) -> InfoMap -> InfoMap
  addRemoteExpCacheReader n key m =
    case M.lookup n m of
    Just (GenericInfo _ (NtRMap (MkRMapInfo _ _ (Just _)))) -> m
    Just (GenericInfo c (NtRMap (MkRMapInfo cs uncacher Nothing))) ->
      M.insert n (GenericInfo c (NtRMap (MkRMapInfo cs uncacher (Just cacheReader)))) m
    Just (GenericInfo _ (NtRConst (MkRConstInfo _ _ (Just _)))) -> m
    Just (GenericInfo c (NtRConst (MkRConstInfo cacher uncacher Nothing))) ->
      M.insert n (GenericInfo c (NtRConst (MkRConstInfo cacher uncacher (Just cacheReader)))) m
    _ ->  error ("Node " ++ show n ++ " cannot add remote exp cache reader")
    where
    cacheReader :: V.Vault -> Maybe (Data BS.ByteString)
    cacheReader vault =
      case V.lookup key vault of
        Nothing -> Nothing
        Just NoData -> Just NoData
        Just (Data b) -> Just $ Data $ S.encode b
  
  
  addRemoteExpCacheReaderM ::
    forall a m. (MonadLoggerIO m, S.Serialize a)
    => SExp 'Remote a -> StateT InfoMap m ()
  addRemoteExpCacheReaderM e = do
    let n = getRemoteIndex e
    $(logInfo) $ T.pack  ("Adding cache reader to remote exp node: " ++ show n)
    let key = getRemoteVaultKey e
    m <- get
    put $ addRemoteExpCacheReader n key m
  
  
  
  getRemoteIndex :: SExp 'Remote a -> Int
  getRemoteIndex (SRApply i _ _ _) = i
  getRemoteIndex (SRConst i _ _) = i
  
  getRemoteVaultKey :: SExp 'Remote a -> V.Key (Data a)
  getRemoteVaultKey (SRApply _ k _ _) = k
  getRemoteVaultKey (SRConst _ k _) = k
  
  getLocalIndex :: SExp 'Local a -> Int
  getLocalIndex (SLConst i _ _) = i
  getLocalIndex (SCollect i _ _) = i
  getLocalIndex (SLApply i _ _ _) = i
  
  getLocalVaultKey :: SExp 'Local a -> V.Key a
  getLocalVaultKey (SLConst _ k _) = k
  getLocalVaultKey (SCollect _ k _) = k
  getLocalVaultKey (SLApply _ k _ _) = k
  
  
  analyseRemote :: (MonadLoggerIO m) => SExp 'Remote a -> StateT InfoMap m ()
  analyseRemote (SRApply n keyb cs@(ExpClosure ce _) a) = do
    unlessM (wasVisitedM n) $ do
      analyseRemote a
      referenceM n (getRemoteIndex a)
      analyseLocal ce
      referenceM n (getLocalIndex ce)
      $(logInfo) $ T.pack ("create closure for RApply node " ++ show n)
      let keya = getRemoteVaultKey a
      rcs <- mkRemoteClosure keya keyb cs
      visitRApplyM rcs
    where
    visitRApplyM cs' = do
      $(logInfo) $ T.pack  ("Visiting RMap node: " ++ show n)
      m <- get
      put $ visitRApply cs' m
    visitRApply cs' m =
      case M.lookup n m of
      Just (GenericInfo _ _) -> m
      Nothing -> M.insert n (GenericInfo S.empty (NtRMap (MkRMapInfo cs' (makeUnCacher keyb) Nothing))) m
  
  analyseRemote (SRConst n key _) = do
    unlessM (wasVisitedM n) $ visitRConstExpM
    where
    visitRConstExpM = do
      $(logInfo) $ T.pack  ("Visiting RConst node: " ++ show n)
      m <- get
      put $ visitRConst m
    visitRConst m =
      case M.lookup n m of
      Just (GenericInfo _ _) -> error ("RConst Node " ++ show n ++ " has already been visited")
      Nothing -> M.insert n (GenericInfo S.empty (NtRConst (MkRConstInfo (makeCacher key) (makeUnCacher key) Nothing))) m
      where
      makeCacher :: (S.Serialize a) => V.Key (Data a) -> (Data BS.ByteString) -> V.Vault -> V.Vault
      makeCacher k (NoData) vault = V.insert k NoData vault
      makeCacher k (Data bs) vault =
        case S.decode bs of
        Left e -> error $ ("Cannot deserialize value: " ++ e)
        Right a -> V.insert k (Data a) vault
  
  
  
  analyseLocal :: (MonadLoggerIO m) => SExp 'Local a -> StateT InfoMap m ()
  
  analyseLocal(SLConst n _ _) = do
    unlessM (wasVisitedM n) $ visitLocalExpM n
  
  analyseLocal (SCollect n _ e) = do
    unlessM (wasVisitedM n) $ do
      analyseRemote e
      addRemoteExpCacheReaderM e
      referenceM n (getRemoteIndex e)
      visitLocalExpM n
  
  analyseLocal (SLApply n _ f e) = do
    unlessM (wasVisitedM n) $ do
      analyseLocal f
      analyseLocal e
      visitLocalExpM n
  
  
  
  
  
  
  
  