{-# LANGUAGE FlexibleContexts, DeriveGeneric, StandaloneDeriving, StrictData #-}


module OutputGeneration.Retrieval where

import Base

import Data.SValue
import Data.SymbolicExpression

import Analysis.Context
import Analysis.Pointers 
import Analysis.ControlFlow

import Instantiation.SymbolicPropagation (get_invariant)

import Generic.SymbolicConstituents
import Generic.HasSize (HasSize(sizeof))
import Generic.Instruction (GenericInstruction(..))
import Generic.Operand
import X86.Instruction


import qualified Data.Map as M
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Set as S
import Data.Maybe (fromJust,catMaybes,mapMaybe)
import Data.List 
import Data.Foldable
import Data.Word
import Control.Monad.State.Strict
import System.Exit (die)
import Debug.Trace

-- | Retrieves a set of instructions.
ctxt_get_instructions :: Context -> S.Set Instruction
ctxt_get_instructions :: Context -> Set Instruction
ctxt_get_instructions Context
ctxt =
  let entries :: Set Int
entries = Context -> Set Int
ctxt_get_function_entries Context
ctxt
      cfgs :: [Maybe CFG]
cfgs    = (Int -> Maybe CFG) -> [Int] -> [Maybe CFG]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
entry -> Int -> IntMap CFG -> Maybe CFG
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
entry (IntMap CFG -> Maybe CFG) -> IntMap CFG -> Maybe CFG
forall a b. (a -> b) -> a -> b
$ Context -> IntMap CFG
ctxt_cfgs Context
ctxt) ([Int] -> [Maybe CFG]) -> [Int] -> [Maybe CFG]
forall a b. (a -> b) -> a -> b
$ Set Int -> [Int]
forall a. Set a -> [a]
S.toList Set Int
entries
      instrs :: [[Instruction]]
instrs  = (Maybe CFG -> [Instruction]) -> [Maybe CFG] -> [[Instruction]]
forall a b. (a -> b) -> [a] -> [b]
map (\(Just CFG
cfg) -> [[Instruction]] -> [Instruction]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Instruction]] -> [Instruction])
-> [[Instruction]] -> [Instruction]
forall a b. (a -> b) -> a -> b
$ IntMap [Instruction] -> [[Instruction]]
forall a. IntMap a -> [a]
IM.elems (IntMap [Instruction] -> [[Instruction]])
-> IntMap [Instruction] -> [[Instruction]]
forall a b. (a -> b) -> a -> b
$ CFG -> IntMap [Instruction]
cfg_instrs CFG
cfg) [Maybe CFG]
cfgs in
    if Set Int -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set Int
entries Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= IntMap CFG -> Int
forall a. IntMap a -> Int
IM.size (Context -> IntMap CFG
ctxt_cfgs Context
ctxt) then
      [Char] -> Set Instruction
forall a. HasCallStack => [Char] -> a
error ([Char] -> Set Instruction) -> [Char] -> Set Instruction
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Set Int -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set Int
entries, [Maybe CFG] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe CFG]
cfgs)
    else
      [Set Instruction] -> Set Instruction
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ([Set Instruction] -> Set Instruction)
-> [Set Instruction] -> Set Instruction
forall a b. (a -> b) -> a -> b
$ ([Instruction] -> Set Instruction)
-> [[Instruction]] -> [Set Instruction]
forall a b. (a -> b) -> [a] -> [b]
map [Instruction] -> Set Instruction
forall a. Ord a => [a] -> Set a
S.fromList [[Instruction]]
instrs



-- | Retrieves a set of funtion entries.
ctxt_get_function_entries :: Context -> S.Set Int
ctxt_get_function_entries :: Context -> Set Int
ctxt_get_function_entries = [Int] -> Set Int
forall a. Ord a => [a] -> Set a
S.fromList ([Int] -> Set Int) -> (Context -> [Int]) -> Context -> Set Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap FReturnBehavior -> [Int]
forall a. IntMap a -> [Int]
IM.keys (IntMap FReturnBehavior -> [Int])
-> (Context -> IntMap FReturnBehavior) -> Context -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Context -> IntMap FReturnBehavior
ctxt_calls


-- | Retrieves function summaries (preconditions, postconditions)
ctxt_get_function_summary :: Context -> Int -> (a, (FInit, Maybe FReturnBehavior))
ctxt_get_function_summary Context
ctxt Int
entry =
  let finit :: FInit
finit = Context -> IntMap FInit
ctxt_finits Context
ctxt IntMap FInit -> Int -> FInit
forall a. IntMap a -> Int -> a
IM.! Int
entry
      post :: Maybe FReturnBehavior
post  = Int -> IntMap FReturnBehavior -> Maybe FReturnBehavior
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
entry (IntMap FReturnBehavior -> Maybe FReturnBehavior)
-> IntMap FReturnBehavior -> Maybe FReturnBehavior
forall a b. (a -> b) -> a -> b
$ Context -> IntMap FReturnBehavior
ctxt_calls Context
ctxt in
    (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
entry,(FInit
finit,Maybe FReturnBehavior
post))



{--
ctxt_disassemble_address :: Context -> Word64 -> IO (Instruction,String)
ctxt_disassemble_address ctxt a = do
  let entries = ctxt_get_function_entries ctxt
      cfgs    = map (\entry -> (entry,IM.lookup entry (C.ctxt_cfgs ctxt))) $ S.toList entries
      instrs  = map (\(entry,Just cfg) -> (entry,concat $ IM.elems $ C.cfg_instrs cfg)) cfgs


  i <- ctxt_get_instruction ctxt $ fromIntegral a
  return $ (mk_json_instruction i, C.pp_instruction ctxt i)
--}



-- | Retrieves a textual representation of the function boundaries
ctxt_mk_function_boundary ctxt entry =
  let cfg = ctxt_cfgs ctxt IM.! entry
      addresses = concat $ IM.elems $ cfg_blocks cfg in
    (fromIntegral entry, intercalate "\n" $ map show_chunk $ mk_consecutive_chunks addresses)
 where
  show_chunk [i]   = showHex i ++ " (single instruction)"
  show_chunk chunk = showHex (head chunk) ++ "-->" ++ showHex (last chunk)

  mk_consecutive_chunks :: [Int] -> [[Int]]
  mk_consecutive_chunks = split_consecutives . sort

  split_consecutives :: [Int] -> [[Int]]
  split_consecutives []         = []
  split_consecutives [i]        = [[i]]
  split_consecutives (i0:i1:is) = if i1 < i0 + 16 then add_to_hd i0 $ split_consecutives (i1:is) else [i0] : split_consecutives (i1:is)


  add_to_hd :: Int -> [[Int]] -> [[Int]]
  add_to_hd i []       = [[i]]
  add_to_hd i (is:iss) = (i:is) : iss



-- | Given an address @a@, retrieve all invariants (the address may occur in multiple functions).
ctxt_get_inv :: Context -> Word64 -> (Word64,Maybe [(Word64,Predicate)])
ctxt_get_inv ctxt a = do
  let entries = ctxt_get_function_entries ctxt
      invs    = map get_invariant_per_entry $ S.toList entries in
    if invs == [] || all ((==) Nothing) invs then do
      (a,Nothing)
    else
      (a,Just $ catMaybes invs)
 where
  get_invariant_per_entry entry =
    let fctxt = mk_fcontext ctxt entry in
     case get_invariant fctxt $ fromIntegral a of
       Nothing -> Nothing
       Just inv -> Just (fromIntegral entry,inv)






-- | Retrieve an overview of memory operands
-- Returns per instruction a tuple (entry, addr, ptrs).
-- Here:
--   entry is the entry address of the function of the instruction,
--   addr is the address of the instruction
--   ptrs is a list of symbolic pointers for each operand (or Nothing if not a memory-operand)
ctxt_resolve_mem_operands :: Context -> [(Word64,Word64, [Maybe SValue])]
ctxt_resolve_mem_operands ctxt = 
  let entries = ctxt_get_function_entries ctxt
      cfgs    = map (\entry -> (entry,IM.lookup entry $ ctxt_cfgs ctxt)) $ S.toList entries
      instrs  = map (\(entry,Just cfg) -> (entry,concat $ IM.elems $ cfg_instrs cfg)) cfgs in
    concatMap resolve_mem_operands instrs
 where
  resolve_mem_operands (entry,instrs) = map (resolve_mem_operand entry) instrs

  resolve_mem_operand entry i = (fromIntegral entry, addressof i, map (resolve entry i) $ get_operands i)

  resolve entry i (Memory a si) = 
    let fctxt   = mk_fcontext ctxt entry
        Just p  = get_invariant fctxt (fromIntegral $ addressof i)
        ptr     = evalState (sset_rip fctxt i >> sresolve_address fctxt a) (p,S.empty) in
      --(if ptr == Top then trace ("TOP(U): " ++ show i ++ "\n" ++ show p) else id) $ Just ptr
      Just ptr
  resolve entry i _ = Nothing

  get_operands i = 
    case dest $ head $ canonicalize i of
      Nothing  -> []--srcs i
      Just dst -> [dst]-- : srcs i




-- | Retrieve instruction for a given instruction address, both as datastructure and pretty-printed
ctxt_get_instruction :: Context -> Int -> IO Instruction
ctxt_get_instruction ctxt a = do
  i <- fetch_instruction ctxt $ fromIntegral a
  case i of
    Nothing -> die $ "Could not disassemble instruction at address: " ++ showHex a
    Just i  -> return i


-- | Retrieve all instruction addresses.
ctxt_get_instruction_addresses :: Context -> S.Set Word64
ctxt_get_instruction_addresses ctxt =
  S.map fromIntegral $ S.unions $ map cfg_to_addresses $ IM.elems $ ctxt_cfgs ctxt
 where
  cfg_to_addresses g = S.fromList $ concat $ IM.elems $ cfg_blocks g

-- | Given an address @a@, retrieve the set of next addresses.
ctxt_get_controlflow :: Context -> Word64 -> IO (Word64,[Word64])
ctxt_get_controlflow ctxt a = do
  let entries = ctxt_get_function_entries ctxt
  posts   <- mapM get_post_per_entry $ S.toList entries
  return $ (fromIntegral a,map fromIntegral $ IS.toList $ IS.unions posts)
 where
  get_post_per_entry entry = do
    post <- stepA ctxt entry $ fromIntegral a
    case post of
      Left _     -> return $ IS.empty
      Right nxts -> return $ IS.fromList $ map fst nxts