{- |

This module exports simple paths and instructions. These instructions wrap
LLVM isntructions, but are designed to be easier to symbolically execute.

-}
module InternalIR.SimplePath (
  -- * Simple instructions
    SimpleInstruction(..)
  , isSimpleInstr, isPathEq, isParEqs, isCallMarker
  , getInstr
  , getNamedInstr
  , setEq
  , setEq'
  , setPathEq
  , setPathNEq
  , phisToPar
  , switchToPar
  -- * a simple utility useful elsewhere
  , withoutName
  -- * Module info
  , module InternalIR.ModuleInfo
  -- * Simple paths
  , SimplePath(..)
  , emptySimplePath
  , SimplePathBBEntry(..)
  , asNames
  ) where

import           Data.ByteString.Short (fromShort)
import           Data.ByteString.UTF8  (toString)
import           Data.List             (intercalate)
import           Data.Text.Lazy        (unpack)
import           InternalIR.ModuleInfo
import           LLVM.AST
import qualified LLVM.AST.Constant     as C
import           LLVM.Pretty
import           LLVMAST.Interface

-- | A simple instruction is either a named instruction or a list of operands
-- we should track.
data SimpleInstruction = Instr (Named Instruction)
                       -- ^ Instruction
                       | PathEq (Named Instruction)
                       -- ^ Equality setting for conditional branches (instruction is always a bitcast)
                       | PathNEq (Named Instruction)
                       -- ^ False equality setting for conditional branches (instruction is always a bitcast)
                       | ParEqs [[Named Instruction]]
                       -- ^ Instructions that should be handled in parallel as
                       -- a disjunction. Other code assumes parEqs to be
                       -- bitcasts as generated by phisToPar or switchToPar, so
                       -- only safe to use as such.
                       | TrackOps [Operand]
                       -- ^ Operand to track (usually branch condition)
                       | EnterCall Name
                       -- ^ Entering call
                       | ExitCall (Maybe Name)
                       -- ^ Exiting call w/ optional ret val
                       deriving (Eq, Ord)

instance HasOps SimpleInstruction where
  getOperands (Instr ni)     = getOperands ni
  getOperands (TrackOps ops) = ops
  getOperands (PathEq ni)    = getOperands ni
  getOperands (PathNEq ni)   = getOperands ni
  getOperands EnterCall{}    = []
  getOperands ExitCall{}     = []
  getOperands si             = error $ "BUG: getOperands for " ++ show si ++ "does not make sense"
  getResultName (Instr ni)   = getResultName ni
  getResultName TrackOps{}   = Nothing
  getResultName (PathEq ni)  = getResultName ni
  getResultName (PathNEq ni) = getResultName ni
  getResultName EnterCall{}  = Nothing
  getResultName ExitCall{}   = Nothing
  getResultName si           = error $ "BUG: getResultName for " ++ show si ++ "does not make sense"


getNamedInstr :: SimpleInstruction -> Maybe (Named Instruction)
getNamedInstr (Instr ni) = Just ni
getNamedInstr _          = Nothing

-- | Like 'setEq', but these variables will NOT be renamed when convering to SSA. This is purely for adding path conditions.
setPathEq :: Name -> Type -> Operand -> InstructionMetadata -> SimpleInstruction
setPathEq nvar ty op meta = PathEq $ setEq' nvar ty op meta

-- | Like 'setPathEq' but not equals.
setPathNEq :: Name -> Type -> Operand -> InstructionMetadata -> SimpleInstruction
setPathNEq nvar ty op meta = PathNEq $ setEq' nvar ty op meta

-- | Set a variable equal to the operand.
setEq :: Name -> Type -> Operand -> InstructionMetadata -> SimpleInstruction
setEq nvar ty op meta = Instr $ setEq' nvar ty op meta

-- | Set a variable equal to the operand.
setEq' :: Name -> Type -> Operand -> InstructionMetadata -> Named Instruction
setEq' nvar ty op meta = nvar := BitCast { operand0 = op, type' = ty, metadata = meta }

instance Show SimpleInstruction where
  show (Instr ni)          = unpack . ppll $ ni
  show (PathEq ni)         = unwords [ "path-eq", unpack . ppll $ ni ]
  show (PathNEq ni)        = unwords [ "path-neq", unpack . ppll $ ni ]
  show (ParEqs niss)       = let listStr xs = unwords ["[", intercalate "," xs, "]"]
                                 str_niss = map (listStr . map (unpack . ppll)) niss
                             in "par-eqs " ++ listStr str_niss
  show (TrackOps ops)      = unwords $ "track-ops" : map show ops
  show (EnterCall n)       = unwords ["enter-call", unpack . ppll $ n]
  show (ExitCall Nothing)  = "exit-call"
  show (ExitCall (Just n)) = unwords ["exit-call", unpack . ppll $ n]

-- | Is simple instruction?
isSimpleInstr :: SimpleInstruction -> Bool
isSimpleInstr (Instr _) = True
isSimpleInstr _         = False

-- | Is path equation?
isPathEq :: SimpleInstruction -> Bool
isPathEq (PathEq _)  = True
isPathEq (PathNEq _) = True
isPathEq _           = False

-- | Is parallel instruction equation?
isParEqs :: SimpleInstruction -> Bool
isParEqs (ParEqs _) = True
isParEqs _          = False

-- | Is an enter- or exit-call marker?
isCallMarker :: SimpleInstruction -> Bool
isCallMarker EnterCall{} = True
isCallMarker ExitCall{}  = True
isCallMarker _           = False

-- | Get LLVM instruction from simple instruction.
getInstr :: SimpleInstruction -> Maybe Instruction
getInstr (Instr i) = Just $ withoutName i
getInstr _         = Nothing

-- | Convert a set of phis to a prallel instruction.
phisToPar :: [Named Instruction] -> SimpleInstruction
phisToPar nis = ParEqs $ map extract branchNames
  where branchNames = let ((_ := Phi _ ivs _):_) = nis
                      in map snd ivs
        -- Convert all the nis phis into eqs
        extract bName = map (\(nv := Phi ty ivs meta) ->
                              setEq' nv ty (fst . head $ filter ((== bName) . snd) ivs) meta) nis

-- | Used to turn a switch statement into a simple instruction.
-- Equivalent to lhs = rhs0 OR lhs = rhs1 OR .. OR lhs = rhsN
switchToPar :: Operand             -- ^ lhs
            -> [C.Constant]        -- ^ multiple rhs
            -> InstructionMetadata
            -> SimpleInstruction
switchToPar op' cs meta = ParEqs $ [map (mkPathCond op' meta) cs]
  where mkPathCond op meta' c = case op of
          LocalReference ty nvar -> setEq' nvar ty (ConstantOperand c) meta'
          _                      -> error $ "mkPathCond: not yet supported: " ++ show op'

-- | A simple path is simply a list of straightline instructions. All control
-- flow is collapsed.
data SimplePath = SimplePath { instrs :: [SimpleInstruction]
                             , bbs    :: [SimplePathBBEntry]
                             } deriving (Eq, Ord)

-- | Empty simple path
emptySimplePath :: SimplePath
emptySimplePath = SimplePath [] []

instance Semigroup SimplePath where
  (SimplePath s0 bb0) <> (SimplePath s1 bb1) = SimplePath (s0 ++ s1) (bb0 ++ bb1)

instance Monoid SimplePath where
  mempty = emptySimplePath

instance Show SimplePath where
  show (SimplePath nis _) = "\n<" ++ (intercalate ";\n" $ map show nis) ++ ">\n"

-- | Basic blocks traversed in the SimplePath
data SimplePathBBEntry = SimplePathBBEntry { spbeMod  :: ModuleInfo
                                           , spbeFunc :: Name
                                           -- ^ Function name. (In the spbeMod module of course)
                                           , spbeBB   :: BasicBlock
                                           -- ^ Basic block itself. (In the spbeFunc function of course)
                                           }
                                            deriving (Ord)

-- | We assume that module names are unique, so this makes comparison much faster
instance Eq SimplePathBBEntry where
  spbe1 == spbe2 = asNames spbe1 == asNames spbe2

instance Show SimplePathBBEntry where
  show spbe = "{SPBE " ++ modName (spbeMod spbe) ++
              " " ++ showStem (spbeFunc spbe) ++
              " " ++ showStem (getBlockName (spbeBB spbe)) ++
              "}"
    where showStem :: Name -> String
          showStem (Name x)   = toString $ fromShort x
          showStem (UnName x) = show x

asNames :: SimplePathBBEntry -> (String, Name, Name)
asNames spbe = (modName $ spbeMod spbe, spbeFunc spbe, getBlockName $ spbeBB spbe)
