<commit_before>{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE NoMonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ParallelListComp #-}

-- | High-level API to K3 toolchain stages.
module Language.K3.Stages where

import Control.Arrow hiding ( left )
import Control.DeepSeq
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State
import Control.Monad.Trans.Except

import Control.Concurrent

import Criterion.Measurement
import Criterion.Types

import Data.Binary ( Binary )
import Data.Serialize ( Serialize )

import Data.Function
import Data.Functor.Identity
import Data.List
import Data.List.Split
import Data.Map ( Map )
import Data.Monoid
import qualified Data.Map as Map
import qualified Data.Text as T
import Data.Tree
import Data.Tuple

import Debug.Trace

import GHC.Generics ( Generic )

import qualified Text.Printf as TPF

import Language.K3.Core.Common
import Language.K3.Core.Annotation
import Language.K3.Core.Declaration
import Language.K3.Core.Utils

import Language.K3.Analysis.Core
import Language.K3.Analysis.HMTypes.Inference hiding ( tenv, inferDeclTypes )

import qualified Language.K3.Analysis.Properties           as Properties
import qualified Language.K3.Analysis.CArgs                as CArgs
import qualified Language.K3.Analysis.Provenance.Inference as Provenance
import qualified Language.K3.Analysis.SEffects.Inference   as SEffects

import qualified Language.K3.Analysis.SendGraph as SG

import qualified Language.K3.Codegen.CPP.Materialization.Inference as Mat

import Language.K3.Transform.Simplification
import Language.K3.Transform.TriggerSymbols (triggerSymbols)

import Language.K3.Utils.Pretty
import qualified Language.K3.Utils.PrettyText as PT

import Language.K3.Codegen.CPP.Materialization.Inference
import Language.K3.Codegen.CPP.Preprocessing

-- | Snapshot specifications are a list of pass names to capture per declaration.
type SnapshotSpec = Map String [String]

-- | Stage specifications are a pair of pass identifiers, and a snapshot spec.
data StageSpec = StageSpec { passesToRun    :: Maybe [Identifier]
                           , passesToFilter :: Maybe [Identifier]
                           , snapshotSpec   :: SnapshotSpec }
                 deriving (Eq, Ord, Read, Show, Generic)

-- | Configuration metadata for compiler stages
data CompilerSpec = CompilerSpec { blockSize :: Int
                                 , stageSpec :: StageSpec }
                    deriving (Eq, Ord, Read, Show, Generic)

-- | Compilation profiling
data TransformReport = TransformReport { statistics :: Map String [Measured]
                                       , snapshots  :: Map String [K3 Declaration] }
                      deriving (Eq, Read, Show, Generic)

-- | The program transformation composition monad
data TransformSt = TransformSt { nextuid    :: ParGenSymS
                               , cseCnt     :: ParGenSymS
                               , tenv       :: TIEnv
                               , prenv      :: Properties.PIEnv
                               , penv       :: Provenance.PIEnv
                               , fenv       :: SEffects.FIEnv
                               , report     :: TransformReport
                               , mzFlags    :: MZFlags }
                  deriving (Eq, Read, Show, Generic)

-- | Monoid instance for transform reports.
instance Monoid TransformReport where
  mempty = TransformReport Map.empty Map.empty

  mappend (TransformReport ast asn) (TransformReport bst bsn) =
    TransformReport (ast <> bst) (asn <> bsn)

data TransformStSymS = TransformStSymS { uidSym :: ParGenSymS, cseSym :: ParGenSymS
                                       , prvSym :: ParGenSymS, effSym :: ParGenSymS }
                      deriving (Eq, Ord, Read, Show, Generic)

-- | Stage-based transformation monad.
type TransformM = ExceptT String (StateT TransformSt IO)

{- Stage-based transform instances -}
instance Binary StageSpec
instance Binary CompilerSpec
instance Binary TransformReport
instance Binary TransformSt
instance Binary TransformStSymS

instance Serialize StageSpec
instance Serialize CompilerSpec
instance Serialize TransformReport
instance Serialize TransformSt
instance Serialize TransformStSymS

ss0 :: StageSpec
ss0 = StageSpec Nothing Nothing Map.empty

cs0 :: CompilerSpec
cs0 = CompilerSpec 16 ss0

rp0 :: TransformReport
rp0 = TransformReport Map.empty Map.empty

mz0 :: MZFlags
mz0 = MZFlags

st0 :: Maybe ParGenSymS -> K3 Declaration -> IO (Either String TransformSt)
st0 symSOpt prog =
  return $ mkEnv >>= \(stpe, stfe) ->
    return $ TransformSt uidSymS cseSymS tienv0 Properties.pienv0 stpe stfe rp0 mz0

  where puid = let UID i = maxProgramUID prog in i + 1
        uidSymS = maybe (contigsymAtS puid) (lowerboundsymS puid) symSOpt
        [cseSymS, prvSymS, effSymS] = replicate 3 $ resetsymS uidSymS
        mkEnv = do
          vpenv <- variablePositions prog
          let pe = Provenance.pienv0 (Just prvSymS) vpenv
          return (pe, SEffects.fienv0 (Just effSymS) (Provenance.ppenv pe) $ lcenv vpenv)

runTransformM :: TransformSt -> TransformM a -> IO (Either String (a, TransformSt))
runTransformM st m = do
  (a, s) <- runStateT (runExceptT m) st
  return $ either Left (Right . (,s)) a

evalTransformM :: TransformSt -> TransformM a -> IO (Either String a)
evalTransformM st m = do
  e <- runTransformM st m
  return $ either Left (Right . fst) e

liftEitherM :: Either String a -> TransformM a
liftEitherM = either throwE return

{- TransformSt utilities -}<commit_msg>Set mz0 in stages.<commit_after>{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE NoMonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ParallelListComp #-}

-- | High-level API to K3 toolchain stages.
module Language.K3.Stages where

import Control.Arrow hiding ( left )
import Control.DeepSeq
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State
import Control.Monad.Trans.Except

import Control.Concurrent

import Criterion.Measurement
import Criterion.Types

import Data.Binary ( Binary )
import Data.Serialize ( Serialize )

import Data.Function
import Data.Functor.Identity
import Data.List
import Data.List.Split
import Data.Map ( Map )
import Data.Monoid
import qualified Data.Map as Map
import qualified Data.Text as T
import Data.Tree
import Data.Tuple

import Debug.Trace

import GHC.Generics ( Generic )

import qualified Text.Printf as TPF

import Language.K3.Core.Common
import Language.K3.Core.Annotation
import Language.K3.Core.Declaration
import Language.K3.Core.Utils

import Language.K3.Analysis.Core
import Language.K3.Analysis.HMTypes.Inference hiding ( tenv, inferDeclTypes )

import qualified Language.K3.Analysis.Properties           as Properties
import qualified Language.K3.Analysis.CArgs                as CArgs
import qualified Language.K3.Analysis.Provenance.Inference as Provenance
import qualified Language.K3.Analysis.SEffects.Inference   as SEffects

import qualified Language.K3.Analysis.SendGraph as SG

import qualified Language.K3.Codegen.CPP.Materialization.Inference as Mat
import qualified Language.K3.Codegen.CPP.Materialization.Inference as Mat

import Language.K3.Transform.Simplification
import Language.K3.Transform.TriggerSymbols (triggerSymbols)

import Language.K3.Utils.Pretty
import qualified Language.K3.Utils.PrettyText as PT

import Language.K3.Codegen.CPP.Materialization.Inference
import Language.K3.Codegen.CPP.Preprocessing

-- | Snapshot specifications are a list of pass names to capture per declaration.
type SnapshotSpec = Map String [String]

-- | Stage specifications are a pair of pass identifiers, and a snapshot spec.
data StageSpec = StageSpec { passesToRun    :: Maybe [Identifier]
                           , passesToFilter :: Maybe [Identifier]
                           , snapshotSpec   :: SnapshotSpec }
                 deriving (Eq, Ord, Read, Show, Generic)

-- | Configuration metadata for compiler stages
data CompilerSpec = CompilerSpec { blockSize :: Int
                                 , stageSpec :: StageSpec }
                    deriving (Eq, Ord, Read, Show, Generic)

-- | Compilation profiling
data TransformReport = TransformReport { statistics :: Map String [Measured]
                                       , snapshots  :: Map String [K3 Declaration] }
                      deriving (Eq, Read, Show, Generic)

-- | The program transformation composition monad
data TransformSt = TransformSt { nextuid    :: ParGenSymS
                               , cseCnt     :: ParGenSymS
                               , tenv       :: TIEnv
                               , prenv      :: Properties.PIEnv
                               , penv       :: Provenance.PIEnv
                               , fenv       :: SEffects.FIEnv
                               , report     :: TransformReport
                               , mzFlags    :: MZFlags }
                  deriving (Eq, Read, Show, Generic)

-- | Monoid instance for transform reports.
instance Monoid TransformReport where
  mempty = TransformReport Map.empty Map.empty

  mappend (TransformReport ast asn) (TransformReport bst bsn) =
    TransformReport (ast <> bst) (asn <> bsn)

data TransformStSymS = TransformStSymS { uidSym :: ParGenSymS, cseSym :: ParGenSymS
                                       , prvSym :: ParGenSymS, effSym :: ParGenSymS }
                      deriving (Eq, Ord, Read, Show, Generic)

-- | Stage-based transformation monad.
type TransformM = ExceptT String (StateT TransformSt IO)

{- Stage-based transform instances -}
instance Binary StageSpec
instance Binary CompilerSpec
instance Binary TransformReport
instance Binary TransformSt
instance Binary TransformStSymS

instance Serialize StageSpec
instance Serialize CompilerSpec
instance Serialize TransformReport
instance Serialize TransformSt
instance Serialize TransformStSymS

ss0 :: StageSpec
ss0 = StageSpec Nothing Nothing Map.empty

cs0 :: CompilerSpec
cs0 = CompilerSpec 16 ss0

rp0 :: TransformReport
rp0 = TransformReport Map.empty Map.empty

mz0 :: MZFlags
mz0 = Mat.defaultMZFlags

st0 :: Maybe ParGenSymS -> K3 Declaration -> IO (Either String TransformSt)
st0 symSOpt prog =
  return $ mkEnv >>= \(stpe, stfe) ->
    return $ TransformSt uidSymS cseSymS tienv0 Properties.pienv0 stpe stfe rp0 mz0

  where puid = let UID i = maxProgramUID prog in i + 1
        uidSymS = maybe (contigsymAtS puid) (lowerboundsymS puid) symSOpt
        [cseSymS, prvSymS, effSymS] = replicate 3 $ resetsymS uidSymS
        mkEnv = do
          vpenv <- variablePositions prog
          let pe = Provenance.pienv0 (Just prvSymS) vpenv
          return (pe, SEffects.fienv0 (Just effSymS) (Provenance.ppenv pe) $ lcenv vpenv)

runTransformM :: TransformSt -> TransformM a -> IO (Either String (a, TransformSt))
runTransformM st m = do
  (a, s) <- runStateT (runExceptT m) st
  return $ either Left (Right . (,s)) a

evalTransformM :: TransformSt -> TransformM a -> IO (Either String a)
evalTransformM st m = do
  e <- runTransformM st m
  return $ either Left (Right . fst) e

liftEitherM :: Either String a -> TransformM a
liftEitherM = either throwE return

{- TransformSt utilities -}
{'fraction_non_alphanumeric': 0.07847378439176328, 'fraction_numerical': 0.007613774009344178, 'mean_word_length': 2.881464069845534, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '208568', 'n_tokens_mistral': 3417, 'n_tokens_neox': 3269, 'n_words': 1014}