<commit_before>{-# LANGUAGE TypeFamilies, TypeSynonymInstances, FlexibleInstances,
    PatternGuards, RecordWildCards #-}
-----------------------------------------------------------------------------
-- |
-- Module      : Main
-- Copyright   : (c) 2012, 2013 <NAME>
-- License     : BSD-style
-- Maintainer  : <EMAIL>
-- Stability   : experimental
-- Portability : ghc
--
-- Compilation of Python 3 source code into bytecode.
--
-----------------------------------------------------------------------------

module Compile (compileFile) where

import Prelude hiding (mapM)
import Utils (isPureExpr, isPyObjectExpr)
import StackDepth (maxStackDepth)
import ProgName (progName)
import State
   ( setBlockState, getBlockState, initBlockState, initState
   , emitCodeNoArg, emitCodeArg, compileConstantEmit
   , compileConstant, getFileName, newLabel, labelNextInstruction
   , getObjectName, setObjectName, getLastInstruction, getGlobals
   , getNestedScope, ifDump, emptyDefinitionScope, lookupNestedScope
   , indexedVarSetKeys, lookupVar, emitReadVar, emitWriteVar )
import Assemble (assemble)
import Monad (Compile (..), runCompileMonad)
import Types
   (Identifier, CompileConfig (..)
   , ConstantID, CompileState (..), BlockState (..)
   , AnnotatedCode (..), Dumpable (..), IndexedVarSet, VarInfo (..) )
import Scope (topScope, renderScope)
import Blip.Marshal as Blip (writePyc, PycFile (..), PyObject (..))
import Blip.Bytecode (Bytecode (..), BytecodeArg (..), Opcode (..), encode)
import Language.Python.Version3.Parser (parseModule)
import Language.Python.Common.AST as AST
   ( ModuleSpan (..), Module (..), StatementSpan (..), Statement (..)
   , ExprSpan (..), Expr (..), Ident (..), ArgumentSpan (..), Argument (..)
   , OpSpan, Op (..))
import Language.Python.Common (prettyText)
import System.FilePath ((<.>), takeBaseName)
import System.Directory (doesFileExist, getModificationTime, canonicalizePath)
import System.Time (ClockTime (..))
import System.IO (openFile, IOMode(..), Handle, hClose, hFileSize, hGetContents)
import Data.Word (Word32, Word16)
import Data.Traversable as Traversable (mapM)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.ByteString.Lazy as B (empty)
import Data.List (sort)
import Control.Monad (unless, forM_, when)
import Control.Exception (try)
import System.IO.Error (IOError, userError, ioError)
import Control.Monad.Trans (liftIO)

compiler :: Compilable a => a -> CompileState -> IO (CompileResult a)
compiler = runCompileMonad . compile

class Compilable a where
   type CompileResult a :: *
   compile :: a -> Compile (CompileResult a)

instance Compilable a => Compilable [a] where
   type CompileResult [a] = [CompileResult a]
   compile = mapM compile

compileFile :: CompileConfig -> FilePath -> IO ()
compileFile config path = do
   r <- try $ do
      pyHandle <- openFile path ReadMode
      sizeInBytes <- hFileSize pyHandle
      fileContents <- hGetContents pyHandle
      modifiedTime <- getModificationTime path
      let modSeconds = case modifiedTime of TOD secs _picoSecs -> secs
      pyModule <- parseAndCheckErrors fileContents path
      let (globals, nestedScope) = topScope pyModule
      canonicalPath <- canonicalizePath path 
      let state = initState globals emptyDefinitionScope
                     nestedScope config canonicalPath
      pyc <- compileModule state (fromIntegral modSeconds)
                (fromIntegral sizeInBytes) pyModule
      let pycFilePath = takeBaseName path <.> ".pyc"
      pycHandle <- openFile pycFilePath WriteMode 
      writePyc pycHandle pyc
      hClose pycHandle
   -- XXX maybe we want more customised error messages for different kinds of
   -- IOErrors?
   case r of
      Left e -> putStrLn $ progName ++ ": " ++ show (e :: IOError)
      Right () -> return ()

parseAndCheckErrors :: String -> FilePath -> IO ModuleSpan
parseAndCheckErrors fileContents sourceName =
   case parseModule fileContents sourceName of
      Left e -> error $ "parse error: " ++ prettyText e
      Right (pyModule, _comments) -> return pyModule

compileModule :: CompileState
              -> Word32         -- modification time
              -> Word32         -- size in bytes
              -> ModuleSpan     -- AST
              -> IO PycFile
compileModule state pyFileModifiedTime pyFileSizeBytes mod = do
   obj <- compiler mod state
   return $ PycFile
      { magic = compileConfig_magic $ state_config state
      , modified_time = pyFileModifiedTime 
      , size = pyFileSizeBytes
      , object = obj }

instance Compilable ModuleSpan where
   type CompileResult ModuleSpan = PyObject
   compile (Module stmts) = do
      maybeDumpScope 
      setObjectName "<module>"
      compile $ Body stmts

nestedBlock :: Identifier -> Compile a -> Compile a
nestedBlock objectName comp = do
   -- save the current block state
   oldBlockState <- getBlockState id
   -- set the new block state to initial values, and the
   -- scope of the current definition
   (definitionScope, nestedScope) <- lookupNestedScope objectName
   setBlockState $ initBlockState definitionScope nestedScope
   -- set the new object name
   setObjectName objectName
   -- run the nested computation
   result <- comp
   -- restore the original block state
   setBlockState oldBlockState
   return result

-- body of module, function and class
newtype Body = Body [StatementSpan]

instance Compilable Body where
   type CompileResult Body = PyObject
   compile (Body stmts) = do
      compile stmts
      -- if the last instruction is not a return, then return None
      maybeLastInstruction <- getLastInstruction
      case maybeLastInstruction of
         Nothing -> returnNone
         Just (Bytecode { opcode = RETURN_VALUE }) -> return ()
         other -> returnNone
      assemble
      makeObject

instance Compilable StatementSpan where
   type CompileResult StatementSpan = ()
   compile (Assign [Var ident _] e _) = do
      compile e
      varInfo <- lookupVar $ ident_string ident
      emitWriteVar varInfo
   compile (Return { return_expr = Nothing }) = returnNone
   compile (Return { return_expr = Just expr }) =  
      compile expr >> emitCodeNoArg RETURN_VALUE
   compile (Pass {}) = return ()
   compile (StmtExpr {..}) = 
      unless (isPureExpr stmt_expr) $ 
         compile stmt_expr >> emitCodeNoArg POP_TOP
   compile (Conditional {..}) = do
      restLabel <- newLabel
      mapM_ (compileGuard restLabel) cond_guards 
      compile cond_else
      labelNextInstruction restLabel
   compile (While {..}) = do
      startLoop <- newLabel
      endLoop <- newLabel
      anchor <- newLabel
      emitCodeArg SETUP_LOOP endLoop
      labelNextInstruction startLoop
      compile while_cond
      emitCodeArg POP_JUMP_IF_FALSE anchor
      compile while_body
      emitCodeArg JUMP_ABSOLUTE startLoop
      labelNextInstruction anchor 
      emitCodeNoArg POP_BLOCK
      compile while_else
      labelNextInstruction endLoop
   compile (Fun {..}) = do
      let funName = ident_string $ fun_name
      varInfo <- lookupVar funName
      funBodyObj <- nestedBlock funName $ do
         compileFunDocString fun_body
         compile $ Body fun_body
      compileConstantEmit funBodyObj
      compileConstantEmit $ Unicode funName
      emitCodeArg MAKE_FUNCTION 0 -- XXX need to figure out this arg
                                  -- appears to be related to keyword args, and defaults
      emitWriteVar varInfo
   compile (NonLocal {}) = return ()
   compile (Global {}) = return ()
   compile other = error ("Unsupported statement " ++ show other) 

-- Check for a docstring in the first statement of a function body.
-- The first constant in the corresponding code object is inspected
-- by the interpreter for the docstring. If their is no docstring
-- then the first constant must be None
compileFunDocString :: [StatementSpan] -> Compile ()
compileFunDocString (firstStmt:_stmts)
   | StmtExpr {..} <- firstStmt,
     Strings {} <- stmt_expr
        = do compileConstant $ constantToPyObject stmt_expr
             return ()
   | otherwise = compileConstant Blip.None >> return ()

compileGuard :: Word16 -> (ExprSpan, [StatementSpan]) -> Compile ()<commit_msg>Added support for assert statements<commit_after>{-# LANGUAGE TypeFamilies, TypeSynonymInstances, FlexibleInstances,
    PatternGuards, RecordWildCards #-}
-----------------------------------------------------------------------------
-- |
-- Module      : Main
-- Copyright   : (c) 2012, 2013 <NAME>
-- License     : BSD-style
-- Maintainer  : <EMAIL>
-- Stability   : experimental
-- Portability : ghc
--
-- Compilation of Python 3 source code into bytecode.
--
-----------------------------------------------------------------------------

module Compile (compileFile) where

import Prelude hiding (mapM)
import Utils (isPureExpr, isPyObjectExpr)
import StackDepth (maxStackDepth)
import ProgName (progName)
import State
   ( setBlockState, getBlockState, initBlockState, initState
   , emitCodeNoArg, emitCodeArg, compileConstantEmit
   , compileConstant, getFileName, newLabel, labelNextInstruction
   , getObjectName, setObjectName, getLastInstruction, getGlobals
   , getNestedScope, ifDump, emptyDefinitionScope, lookupNestedScope
   , indexedVarSetKeys, lookupVar, emitReadVar, emitWriteVar
   , lookupGlobalVar )
import Assemble (assemble)
import Monad (Compile (..), runCompileMonad)
import Types
   (Identifier, CompileConfig (..)
   , ConstantID, CompileState (..), BlockState (..)
   , AnnotatedCode (..), Dumpable (..), IndexedVarSet, VarInfo (..) )
import Scope (topScope, renderScope)
import Blip.Marshal as Blip (writePyc, PycFile (..), PyObject (..))
import Blip.Bytecode (Bytecode (..), BytecodeArg (..), Opcode (..), encode)
import Language.Python.Version3.Parser (parseModule)
import Language.Python.Common.AST as AST
   ( ModuleSpan (..), Module (..), StatementSpan (..), Statement (..)
   , ExprSpan (..), Expr (..), Ident (..), ArgumentSpan (..), Argument (..)
   , OpSpan, Op (..))
import Language.Python.Common (prettyText)
import System.FilePath ((<.>), takeBaseName)
import System.Directory (doesFileExist, getModificationTime, canonicalizePath)
import System.Time (ClockTime (..))
import System.IO (openFile, IOMode(..), Handle, hClose, hFileSize, hGetContents)
import Data.Word (Word32, Word16)
import Data.Traversable as Traversable (mapM)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.ByteString.Lazy as B (empty)
import Data.List (sort)
import Control.Monad (unless, forM_, when)
import Control.Exception (try)
import System.IO.Error (IOError, userError, ioError)
import Control.Monad.Trans (liftIO)

compiler :: Compilable a => a -> CompileState -> IO (CompileResult a)
compiler = runCompileMonad . compile

class Compilable a where
   type CompileResult a :: *
   compile :: a -> Compile (CompileResult a)

instance Compilable a => Compilable [a] where
   type CompileResult [a] = [CompileResult a]
   compile = mapM compile

compileFile :: CompileConfig -> FilePath -> IO ()
compileFile config path = do
   r <- try $ do
      pyHandle <- openFile path ReadMode
      sizeInBytes <- hFileSize pyHandle
      fileContents <- hGetContents pyHandle
      modifiedTime <- getModificationTime path
      let modSeconds = case modifiedTime of TOD secs _picoSecs -> secs
      pyModule <- parseAndCheckErrors fileContents path
      let (globals, nestedScope) = topScope pyModule
      canonicalPath <- canonicalizePath path 
      let state = initState globals emptyDefinitionScope
                     nestedScope config canonicalPath
      pyc <- compileModule state (fromIntegral modSeconds)
                (fromIntegral sizeInBytes) pyModule
      let pycFilePath = takeBaseName path <.> ".pyc"
      pycHandle <- openFile pycFilePath WriteMode 
      writePyc pycHandle pyc
      hClose pycHandle
   -- XXX maybe we want more customised error messages for different kinds of
   -- IOErrors?
   case r of
      Left e -> putStrLn $ progName ++ ": " ++ show (e :: IOError)
      Right () -> return ()

parseAndCheckErrors :: String -> FilePath -> IO ModuleSpan
parseAndCheckErrors fileContents sourceName =
   case parseModule fileContents sourceName of
      Left e -> error $ "parse error: " ++ prettyText e
      Right (pyModule, _comments) -> return pyModule

compileModule :: CompileState
              -> Word32         -- modification time
              -> Word32         -- size in bytes
              -> ModuleSpan     -- AST
              -> IO PycFile
compileModule state pyFileModifiedTime pyFileSizeBytes mod = do
   obj <- compiler mod state
   return $ PycFile
      { magic = compileConfig_magic $ state_config state
      , modified_time = pyFileModifiedTime 
      , size = pyFileSizeBytes
      , object = obj }

instance Compilable ModuleSpan where
   type CompileResult ModuleSpan = PyObject
   compile (Module stmts) = do
      maybeDumpScope 
      setObjectName "<module>"
      compile $ Body stmts

nestedBlock :: Identifier -> Compile a -> Compile a
nestedBlock objectName comp = do
   -- save the current block state
   oldBlockState <- getBlockState id
   -- set the new block state to initial values, and the
   -- scope of the current definition
   (definitionScope, nestedScope) <- lookupNestedScope objectName
   setBlockState $ initBlockState definitionScope nestedScope
   -- set the new object name
   setObjectName objectName
   -- run the nested computation
   result <- comp
   -- restore the original block state
   setBlockState oldBlockState
   return result

-- body of module, function and class
newtype Body = Body [StatementSpan]

instance Compilable Body where
   type CompileResult Body = PyObject
   compile (Body stmts) = do
      compile stmts
      -- if the last instruction is not a return, then return None
      maybeLastInstruction <- getLastInstruction
      case maybeLastInstruction of
         Nothing -> returnNone
         Just (Bytecode { opcode = RETURN_VALUE }) -> return ()
         other -> returnNone
      assemble
      makeObject

instance Compilable StatementSpan where
   type CompileResult StatementSpan = ()
   compile (Assign [Var ident _] e _) = do
      compile e
      varInfo <- lookupVar $ ident_string ident
      emitWriteVar varInfo
   compile (Return { return_expr = Nothing }) = returnNone
   compile (Return { return_expr = Just expr }) =  
      compile expr >> emitCodeNoArg RETURN_VALUE
   compile (Pass {}) = return ()
   compile (StmtExpr {..}) = 
      unless (isPureExpr stmt_expr) $ 
         compile stmt_expr >> emitCodeNoArg POP_TOP
   compile (Conditional {..}) = do
      restLabel <- newLabel
      mapM_ (compileGuard restLabel) cond_guards 
      compile cond_else
      labelNextInstruction restLabel
   compile (While {..}) = do
      startLoop <- newLabel
      endLoop <- newLabel
      anchor <- newLabel
      emitCodeArg SETUP_LOOP endLoop
      labelNextInstruction startLoop
      compile while_cond
      emitCodeArg POP_JUMP_IF_FALSE anchor
      compile while_body
      emitCodeArg JUMP_ABSOLUTE startLoop
      labelNextInstruction anchor 
      emitCodeNoArg POP_BLOCK
      compile while_else
      labelNextInstruction endLoop
   compile (Fun {..}) = do
      let funName = ident_string $ fun_name
      varInfo <- lookupVar funName
      funBodyObj <- nestedBlock funName $ do
         compileFunDocString fun_body
         compile $ Body fun_body
      compileConstantEmit funBodyObj
      compileConstantEmit $ Unicode funName
      emitCodeArg MAKE_FUNCTION 0 -- XXX need to figure out this arg
                                  -- appears to be related to keyword args, and defaults
      emitWriteVar varInfo
   -- XXX assertions appear to be turned off if the code is compiled
   -- for optimisation
   -- If the assertion expression is a tuple of non-zero length, then
   -- it is always True: CPython warns about this
   compile (Assert {..}) = do
      case assert_exprs of
         test_expr:restAssertExprs -> do
            compile test_expr
            end <- newLabel
            emitCodeArg POP_JUMP_IF_TRUE end
            GlobalVar assertionErrorVar <- lookupGlobalVar "AssertionError"
            emitCodeArg LOAD_GLOBAL assertionErrorVar
            case restAssertExprs of
               assertMsgExpr:_ -> do
                  compile assertMsgExpr
                  emitCodeArg CALL_FUNCTION 1
               _other -> return ()
            emitCodeArg RAISE_VARARGS 1
            labelNextInstruction end
         _other -> error "assert with no test"
              
   compile (NonLocal {}) = return ()
   compile (Global {}) = return ()
   compile other = error ("Unsupported statement " ++ show other) 

-- Check for a docstring in the first statement of a function body.
-- The first constant in the corresponding code object is inspected
-- by the interpreter for the docstring. If their is no docstring
-- then the first constant must be None
compileFunDocString :: [StatementSpan] -> Compile ()
compileFunDocString (firstStmt:_stmts)
   | StmtExpr {..} <- firstStmt,
     Strings {} <- stmt_expr
        = do compileConstant $ constantToPyObject stmt_expr
             return ()
   | otherwise = compileConstant Blip.None >> return ()

compileGuard :: Word16 -> (ExprSpan, [StatementSpan]) -> Compile ()
{'fraction_non_alphanumeric': 0.09831055757366085, 'fraction_numerical': 0.0025370466470622153, 'mean_word_length': 2.7951859956236325, 'pattern_counts': {'":': 2, '<': 61, '<?xml version=': 0, '>': 81, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6308217', 'n_tokens_mistral': 4767, 'n_tokens_neox': 4633, 'n_words': 1551}