content: {-# LANGUAGE RecordWildCards #-}
  
  -----------------------------------------------------------------------------
  -- |
  -- Module      : Blip.Compiler.State
  -- Copyright   : (c) 2012, 2013, 2014 Bernie Pope
  -- License     : BSD-style
  -- Maintainer  : florbitous@gmail.com
  -- Stability   : experimental
  -- Portability : ghc
  --
  -- Management of state for the compiler. There is global state which
  -- persists through the whole compilation (such as command line flags), and
  -- there is block state, which is local for the compilation of a block
  -- of code.
  --
  -----------------------------------------------------------------------------
  module Blip.Compiler.State
     ( setBlockState, getBlockState, initBlockState, initState, modifyBlockState
     , emitCode, emitCodeNoArg, emitCodeArg, compileConstant
     , getFileName, newLabel, compileConstantEmit, labelNextInstruction
     , getObjectName, setObjectName, getLabelMap
     , getNestedScope, ifDump, emptyVarSet, emptyLocalScope
     , getLocalScope, indexedVarSetKeys, lookupNameVar
     , emitReadVar, emitWriteVar, emitDeleteVar, lookupClosureVar, setFlag
     , pushFrameBlock, popFrameBlock, peekFrameBlock, withFrameBlock 
     , setFastLocals, setArgCount, emptyParameterTypes, setLineNumber
     , setFirstLineNumber )
     where
  
  import Blip.Compiler.Monad (Compile (..))
  import Blip.Compiler.Types
     ( Identifier, CompileConfig (..), VarIndex, IndexedVarSet
     , ConstantID, CompileState (..), BlockState (..)
     , AnnotatedCode (..), LabelMap, Dumpable, VarSet, NestedScope (..)
     , LocalScope (..), VarInfo (..), ScopeIdentifier
     , FrameBlockInfo (..), Context (..), ParameterTypes (..) )
  import Blip.Bytecode
     (Bytecode (..), Opcode (..), BytecodeArg (..), bytecodeSize)
  import Blip.Marshal (PyObject (..), CodeObjectFlagMask, co_varargs, co_varkeywords)
  import Blip.Compiler.Utils (identsFromParameters, countPosParameters, getSpanLine)
  import Data.Word (Word16, Word32)
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  import Control.Monad.State.Strict as State hiding (State)
  import Data.List (sort)
  import Data.Bits ((.|.))
  import Language.Python.Common.SrcLocation (SrcSpan (..))
  
  emptyVarSet :: VarSet
  emptyVarSet = Set.empty
  
  emptyParameterTypes :: ParameterTypes
  emptyParameterTypes =
     ParameterTypes { parameterTypes_pos = []
                    , parameterTypes_varPos = Nothing
                    , parameterTypes_varKeyword = Nothing
                    }
  
  emptyLocalScope :: LocalScope
  emptyLocalScope =
     LocalScope
     { localScope_params = emptyParameterTypes 
     , localScope_locals = emptyVarSet
     , localScope_freeVars = emptyVarSet
     , localScope_cellVars = emptyVarSet
     , localScope_explicitGlobals = emptyVarSet
     }
  
  initBlockState :: Context -> LocalScope -> BlockState
  initBlockState context (LocalScope {..}) = BlockState
     { state_label = 0
     , state_instructions = []
     , state_labelNextInstruction = [] 
     , state_constants = [] 
     , state_constantCache = Map.empty
     , state_nextConstantID = 0
     , state_names = []
     , state_nameCache = Map.empty
     , state_nextNameID = 0
     , state_objectName = ""
     , state_instruction_index = 0
     , state_labelMap = Map.empty
     , state_locals = localScope_locals
     , state_fastLocals =
          if context == FunctionContext
             then makeLocalsIndexedSet (identsFromParameters localScope_params)
                     localScope_locals 
             else Map.empty
     -- the indices for cellvars and freevars are used as offsets into an array
     -- within the code object. The cellvars come first, followed by the
     -- freevars. These indices are used by the LOAD_DEREF and STORE_DEREF
     -- bytecode instructions.
     -- cellvars are indexed from 0 upwards
     , state_cellVars = indexedVarSet 0 $ localScope_cellVars 
     -- freevars are indexed from (length cellvars) 
     , state_freeVars = indexedVarSet
                           (fromIntegral $ Set.size localScope_cellVars) 
                           localScope_freeVars 
     , state_explicitGlobals = localScope_explicitGlobals
     , state_argcount = fromIntegral $ countPosParameters localScope_params
     , state_flags = varArgsFlags localScope_params 0
     , state_frameBlockStack = []
     , state_context = context
     , state_lineNumber = 0
     , state_lineNumberTable = []
     , state_firstLineNumber = 0
     }
  
  varArgsFlags :: ParameterTypes -> Word32 -> Word32
  varArgsFlags (ParameterTypes {..}) flags =
     flags .|. posVarArgsMask .|. keywordVarArgsMask
     where
     posVarArgsMask :: CodeObjectFlagMask
     posVarArgsMask = 
        maybe 0 (const co_varargs) parameterTypes_varPos
     keywordVarArgsMask :: CodeObjectFlagMask
     keywordVarArgsMask =
        maybe 0 (const co_varkeywords) parameterTypes_varKeyword 
  
  -- Local variables are indexed starting with parameters first, in the order
  -- that they appear in the function head, followed by the other
  -- locally defined variables, which can appear in any order.
  makeLocalsIndexedSet :: [Identifier] -> VarSet -> IndexedVarSet
  makeLocalsIndexedSet params locals =
     Map.fromList $ zip (params ++ Set.toList localsNotParams) [0..]
     where
     localsNotParams = locals `Set.difference` Set.fromList params
  
  indexedVarSet :: VarIndex -> VarSet -> IndexedVarSet
  indexedVarSet from set =
     Map.fromList $ sort $ zip (Set.toList set) [from..]
  
  -- Return the keys of an IndexedVarSet in ascending order of the indices
  indexedVarSetKeys :: IndexedVarSet -> [Identifier]
  indexedVarSetKeys varset =
     map snd $ sort [ (index, name) | (name, index) <- Map.assocs varset ]
  
  incInstructionIndex :: Bytecode -> Compile Word16
  incInstructionIndex bytecode = do
     currentIndex <- getBlockState state_instruction_index
     let nextIndex = currentIndex + (fromIntegral $ bytecodeSize bytecode)
     modifyBlockState $ \s -> s { state_instruction_index = nextIndex }
     return currentIndex
  
  setFastLocals :: [Identifier] -> Compile ()
  setFastLocals idents = do
     let localsVarSet = Set.fromList idents
     modifyBlockState $ \s -> s { state_fastLocals = indexedVarSet 0 localsVarSet }
  
  setArgCount :: Word32 -> Compile ()
  setArgCount n = modifyBlockState $ \s -> s { state_argcount = n }
  
  initState :: Context       -- module, class or function?
            -> LocalScope    -- local scope of the top-level of the module
            -> NestedScope   -- nested scope of the rest of the module (not at the top-level)
            -> CompileConfig -- configuration options
            -> FilePath      -- file path of the Python source
            -> CompileState
  initState context localScope nestedScope config pyFilename = CompileState
     { state_config = config
     , state_blockState = initBlockState context localScope
     , state_filename = pyFilename
     , state_nestedScope = nestedScope
     }
  
  ifDump :: Dumpable -> Compile () -> Compile ()
  ifDump dumpable action = do
     state <- get
     if dumpable `Set.member` (compileConfig_dumps $ state_config state)
        then action
        else return () 
  
  -- get the nested scope for the current block
  getNestedScope :: Compile NestedScope
  getNestedScope = gets state_nestedScope
  
  getLocalScope :: ScopeIdentifier -> Compile (String, LocalScope)
  getLocalScope scopeIdent = do
     NestedScope nestedScope <- getNestedScope
     case Map.lookup scopeIdent nestedScope of
        Just scope -> return scope
        -- this case should never happen
        Nothing -> error $ "no scope found for: " ++ show scopeIdent
  
  getFileName :: Compile FilePath
  getFileName = gets state_filename
  
  getObjectName :: Compile String
  getObjectName = getBlockState state_objectName
  
  setObjectName :: String -> Compile ()
  setObjectName str = modifyBlockState $ \s -> s { state_objectName = str }
  
  setBlockState :: BlockState -> Compile ()
  setBlockState blockState = do
     oldState <- get
     put $ oldState { state_blockState = blockState }
  
  getBlockState :: (BlockState -> a) -> Compile a 
  getBlockState f = gets (f . state_blockState)
  
  modifyBlockState :: (BlockState -> BlockState) -> Compile ()
  modifyBlockState f = do
     state <- getBlockState id
     setBlockState $! f state
  
  newLabel :: Compile Word16
  newLabel = do
     currentLabel <- getBlockState state_label
     let newLabel = currentLabel + 1
     modifyBlockState $ \s -> s { state_label = newLabel }
     return currentLabel
  
  -- prefix this new label onto the existing ones
  labelNextInstruction :: Word16 -> Compile ()
  labelNextInstruction newLabel = do
     currentLabels <- getBlockState state_labelNextInstruction
     modifyBlockState $ \ s -> s { state_labelNextInstruction = newLabel : currentLabels }
  
  {-
            | Free  | Cell  | Local | Explicit Global | Implicit Global
  ---------------------------------------------------------------------
  Class     | Deref | X     | Name  | Global          | Name          
  Module    | X     | X     | Name  | X               | Name
  Funcition | Deref | Deref | Fast  | Global          | Global
  
  -}
  
  data VarOpcodeType = Deref | Name | Global | Fast
  
  emitReadVar :: Identifier -> Compile ()
  emitReadVar ident = do
     (opcodeType, index) <- getVarOpcodeType ident
     case opcodeType of
        Deref -> emitCodeArg LOAD_DEREF index
        Name -> emitCodeArg LOAD_NAME index
        Global -> emitCodeArg LOAD_GLOBAL index
        Fast -> emitCodeArg LOAD_FAST index
  
  emitWriteVar :: Identifier -> Compile ()
  emitWriteVar ident = do
     (opcodeType, index) <- getVarOpcodeType ident
     case opcodeType of
        Deref -> emitCodeArg STORE_DEREF index
        Name -> emitCodeArg STORE_NAME index
        Global -> emitCodeArg STORE_GLOBAL index
        Fast -> emitCodeArg STORE_FAST index
  
  emitDeleteVar :: Identifier -> Compile ()
  emitDeleteVar ident = do
     (opcodeType, index) <- getVarOpcodeType ident
     case opcodeType of
        Deref -> emitCodeArg DELETE_DEREF index
        Name -> emitCodeArg DELETE_NAME index
        Global -> emitCodeArg DELETE_GLOBAL index
        Fast -> emitCodeArg DELETE_FAST index
  
  getVarOpcodeType :: Identifier -> Compile (VarOpcodeType, VarIndex)
  getVarOpcodeType ident = do
     context <- getBlockState state_context
     varInfo <- lookupVar ident
     getVarInContext context varInfo
     where
     getVarInContext :: Context -> VarInfo -> Compile (VarOpcodeType, VarIndex)
     getVarInContext ClassContext info =
        case info of
           FreeVar index -> return (Deref, index)
           LocalVar -> do
              index <- lookupNameVar ident
              return (Name, index)
           ExplicitGlobal -> do
              index <- lookupNameVar ident
              return (Global, index)
           ImplicitGlobal -> do
              index <- lookupNameVar ident
              return (Name, index)
           CellVar _index -> error $ "class with a cell variable: " ++ ident
     getVarInContext ModuleContext info =
        case info of
           LocalVar -> do
              index <- lookupNameVar ident
              return (Name, index)
           ImplicitGlobal -> do
              index <- lookupNameVar ident
              return (Name, index)
           FreeVar _index ->
              error $ "module with a free variable: " ++ ident
           CellVar _index ->
              error $ "module with a cell variable: " ++ ident
           ExplicitGlobal ->
              error $ "module with an explicit global variable: " ++ ident
     getVarInContext FunctionContext info =
        case info of
           FreeVar index -> return (Deref, index)
           CellVar index -> return (Deref, index)
           LocalVar -> do
              fastLocals <- getBlockState state_fastLocals 
              case Map.lookup ident fastLocals of
                 Just index -> return (Fast, index)
                 Nothing -> error $ "local function variable not in fast locals: " ++ ident
           ExplicitGlobal -> do
              index <- lookupNameVar ident
              return (Global, index)
           ImplicitGlobal -> do
              index <- lookupNameVar ident
              return (Global, index)
  
  emitCodeArg :: Opcode -> Word16 -> Compile ()
  emitCodeArg opCode arg = emitCode $ Bytecode opCode (Just $ Arg16 arg)
  
  emitCodeNoArg :: Opcode -> Compile ()
  emitCodeNoArg opCode = emitCode $ Bytecode opCode Nothing
  
  emitCode :: Bytecode -> Compile ()
  emitCode instruction = do
     -- Attach a label to the instruction if necesary.
     labels <- getBlockState state_labelNextInstruction
     -- Ensure current labels are used only once.
     modifyBlockState $ \s -> s { state_labelNextInstruction = [] }
     instructionIndex <- incInstructionIndex instruction
     -- add a mapping from instruction offset to source code line number
     updateLineNumberTable instructionIndex
     -- Map each label to its instruction index
     forM_ labels $ \label -> updateLabelMap label instructionIndex
     let annotatedInstruction =
            AnnotatedCode { annotatedCode_bytecode = instruction
                          , annotatedCode_labels = labels
                          , annotatedCode_index = instructionIndex } 
     oldInstructions <- getBlockState state_instructions
     modifyBlockState $
        \s -> s { state_instructions = annotatedInstruction : oldInstructions }
  
  getLabelMap :: Compile LabelMap
  getLabelMap = getBlockState state_labelMap
  
  updateLabelMap :: Word16 -> Word16 -> Compile ()
  updateLabelMap label index = do
     oldLabelMap <- getBlockState state_labelMap
     let newLabelMap = Map.insert label index oldLabelMap
     modifyBlockState $ \s -> s { state_labelMap = newLabelMap }
  
  updateLineNumberTable :: Word16 -> Compile ()
  updateLineNumberTable offset = do
     lineNumber <- getBlockState state_lineNumber
     oldTable <- getBlockState state_lineNumberTable
     let updateTable =
            modifyBlockState $ \s -> s { state_lineNumberTable = (offset, lineNumber) : oldTable }
     case oldTable of
        [] -> updateTable 
        (_prevOffset, prevLineNumber):_rest
           -- don't update the table if the current line number is not less
           -- than the previously stored line number
           | prevLineNumber >= lineNumber -> return ()
           | otherwise -> updateTable
  
  compileConstant :: PyObject -> Compile ConstantID
  -- Code objects are not cached to avoid complex equality comparisons
  compileConstant obj@(Code {}) = do
     oldConstants <- getBlockState state_constants
     constantID <- getBlockState state_nextConstantID 
     modifyBlockState $
        \s -> s { state_constants = obj : oldConstants
                , state_nextConstantID = constantID + 1 }
     return constantID
  compileConstant obj = do
     blockState <- getBlockState id
     let constantCache = state_constantCache blockState
     case Map.lookup obj constantCache of
        -- We haven't seen this (non-code) constant before
        Nothing -> do
           let constantID = state_nextConstantID blockState
               newConstantCache = Map.insert obj constantID constantCache 
               oldConstants = state_constants blockState
           setBlockState $ blockState
              { state_nextConstantID = constantID + 1
              , state_constantCache = newConstantCache
              , state_constants = obj : oldConstants }
           return constantID
        Just constantID -> return constantID
  
  compileConstantEmit :: PyObject -> Compile ()
  compileConstantEmit obj = do
     constantID <- compileConstant obj
     emitCodeArg LOAD_CONST constantID
  
  {-
  
  check if var is:
  
    cellvar
    localvar
    freevar
    explicit global
    implicit global 
  
  We check local vars first before free vars because classes can
  have a variable with the same name being local and also free.
  If a local version of the variable is defined, that is the
  one we want to see (not the free variable). If we need to
  see the free variable, then we can look it up specially.
  
  If we can't find it defined anywhere then we presume it
  to be an implicit global variable.
  
  -}
  
  lookupVar :: Identifier -> Compile VarInfo
  lookupVar identifier = do
     -- cell
     cellvars <- getBlockState state_cellVars
     case Map.lookup identifier cellvars of
        Just index -> return $ CellVar index
        Nothing -> do
           -- local
           locals <- getBlockState state_locals
           if identifier `Set.member` locals
              then return LocalVar
              else do
                 -- free
                 freevars <- getBlockState state_freeVars
                 case Map.lookup identifier freevars of
                    Just index -> return $ FreeVar index
                    Nothing -> do
                       -- explicit global
                       explicitGlobals <- getBlockState state_explicitGlobals
                       if identifier `Set.member` explicitGlobals
                          then return ExplicitGlobal
                          -- implicit global 
                          else return ImplicitGlobal
  
  {- lookup a variable in cell vars or free vars only.
     We avoid looking in other places because, for example,
     classes can have free variables with the same name as
     locally defined variables, and we don't want to get them
     confused.
  -}
  
  lookupClosureVar :: Identifier -> Compile (Maybe VarInfo)
  lookupClosureVar identifier = do
     cellvars <- getBlockState state_cellVars
     case Map.lookup identifier cellvars of
        Just index -> return $ Just $ CellVar index
        Nothing -> do
           freevars <- getBlockState state_freeVars
           case Map.lookup identifier freevars of
              Just index -> return $ Just $ FreeVar index
              Nothing -> return Nothing
  
  -- look up a variable in the "names" vector. Add it if it is not there.
  -- return the index of the variable in the vector.
  lookupNameVar :: Identifier -> Compile VarIndex
  lookupNameVar ident = do
     blockState <- getBlockState id
     let nameCache = state_nameCache blockState
     case Map.lookup ident nameCache of
        -- We haven't seen this name before
        Nothing -> do
           let index = state_nextNameID blockState
               newNameCache = Map.insert ident index nameCache
               oldNames = state_names blockState 
           setBlockState $
              blockState { state_nextNameID = index + 1
                         , state_nameCache = newNameCache 
                         , state_names = ident : oldNames }
           return index 
        Just index -> return index 
  
  -- set a flag in the code object by applying a mask 
  setFlag :: CodeObjectFlagMask -> Compile ()
  setFlag mask = do
      oldFlags <- getBlockState state_flags 
      let newFlags = oldFlags .|. mask
      modifyBlockState $ \state -> state { state_flags = newFlags }
  
  pushFrameBlock :: FrameBlockInfo -> Compile ()
  pushFrameBlock info = do
     oldFrameStack <- getBlockState state_frameBlockStack
     let newFrameStack = info : oldFrameStack
     modifyBlockState $ \state -> state { state_frameBlockStack = newFrameStack }
  
  popFrameBlock :: Compile FrameBlockInfo
  popFrameBlock = do
     oldFrameStack <- getBlockState state_frameBlockStack
     case oldFrameStack of
        [] -> error "attempt to pop from an empty frame block stack"
        top:rest -> do
           modifyBlockState $ \state -> state { state_frameBlockStack = rest }
           return top
  
  peekFrameBlock :: Compile (Maybe FrameBlockInfo)
  peekFrameBlock = do
     oldFrameStack <- getBlockState state_frameBlockStack
     case oldFrameStack of
        [] -> return Nothing
        top:_rest -> return $ Just top
  
  withFrameBlock :: FrameBlockInfo -> Compile a -> Compile a
  withFrameBlock pushedInfo comp = do 
      pushFrameBlock pushedInfo
      result <- comp
      poppedInfo <- popFrameBlock
      if pushedInfo /= poppedInfo
         then error $ "pushed frame block not equal to popped frame block"
         else return result
  
  {- 
  
     From Python/compile.c
  
     The line number is reset in the following cases:
     - when entering a new scope
     - on each statement
     - on each expression that start a new line
     - before the "except" clause
     - before the "for" and "while" expressions
  
     Our own remarks:
  
     - the CPython compiler does not seem to follow the above comment
       for "for" and "while" statements (not expressions).
     - I'm not sure why they do something special for the "except" clause
       and why not the "else" and "finally"?
  -}
  
  setLineNumber :: SrcSpan -> Compile ()
  setLineNumber span =
     case getSpanLine span of
        Nothing -> return ()
        Just line -> do
           let lineWord32 = fromIntegral line
           oldLineNumber <- getBlockState state_lineNumber
           -- We ensure that line numbers are monotonically increasing.
           if lineWord32 > oldLineNumber
              then modifyBlockState $ \s -> s { state_lineNumber = lineWord32 }
              else return ()
  
  setFirstLineNumber :: SrcSpan -> Compile ()
  setFirstLineNumber span = 
     case getSpanLine span of
        Nothing -> return ()
        Just line ->
           modifyBlockState $ \state -> state { state_firstLineNumber = fromIntegral line }
  